minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 13 Packages Projects Releases Wiki Activity

[CPU] Fix issues when running llama3.2-11B vision model with image tasks (#8666)

Browse Source 
Co-authored-by: JieXin Liang <Alcanderian@users.noreply.github.com>
Co-authored-by: Yineng Zhang <me@zhyncs.com>
Co-authored-by: jianan-gu <jianan.gu@intel.com>
This commit is contained in:
blzheng
2026-05-21 13:09:18 +08:00
committed by GitHub
co-authored by JieXin Liang Yineng Zhang jianan-gu
parent 79b937aefb
commit 84ea47eb22
15 changed files with 481 additions and 231 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/cpu/test_decode.py
+29 -10
View File
@@ -21,9 +21,11 @@ class TestDecodeAttention(CustomTestCase):
req_to_token: torch.Tensor,
req_pool_indices: torch.Tensor,
seq_lens: torch.Tensor,
encoder_lens=None,
scaling=None,
enable_gqa=False,
causal=False,
is_cross_attn=False,
):
# [num_tokens, num_heads, head_size] -> [num_heads, num_tokens, head_size]
query = query.movedim(0, query.dim() - 2)
@@ -33,14 +35,21 @@ class TestDecodeAttention(CustomTestCase):
seq_len_q = 1
seq_len_kv = seq_lens[seq_idx]
end_q = start_q + seq_len_q
end_kv = start_kv + seq_len_kv
if encoder_lens is not None:
start_kv = 0 if is_cross_attn else encoder_lens[seq_idx]
end_kv = (
encoder_lens[seq_idx] if is_cross_attn else start_kv + seq_len_kv
)
else:
start_kv = 0
end_kv = start_kv + seq_len_kv
per_req_query = query[:, start_q:end_q, :]
# get key and value from cache. per_req_tokens contains the kv cache
# index for each token in the sequence.
req_pool_idx = req_pool_indices[seq_idx]
per_req_tokens = req_to_token[req_pool_idx, :seq_len_kv]
per_req_tokens = req_to_token[req_pool_idx, start_kv:end_kv]
per_req_key = k_cache[per_req_tokens].movedim(0, query.dim() - 2)
per_req_value = v_cache[per_req_tokens].movedim(0, query.dim() - 2)
@@ -61,10 +70,13 @@ class TestDecodeAttention(CustomTestCase):
return output
def _test_grouped_decode_attention_once(self, B, H_Q, H_KV, D, D_V, dtype, device):
def _test_grouped_decode_attention_once(
self, B, H_Q, H_KV, D, D_V, is_cross_attn, dtype, device
):
# This represents the number of tokens already in the sequence
seq_len = 1024
total_tokens = B * seq_len
encoder_len = 10
total_tokens = B * (seq_len + encoder_len)
sm_scale = 1.0 / (D**0.5)
logit_cap = 0.0
num_kv_splits = 8
@@ -91,11 +103,12 @@ class TestDecodeAttention(CustomTestCase):
req_to_token = (
torch.arange(total_tokens, device=device)
.reshape(B, seq_len)
.reshape(B, seq_len + encoder_len)
.to(torch.int32)
)
b_req_idx = torch.arange(B, device=device).to(torch.int64)
b_seq_len = torch.full((B,), seq_len, device=device).to(torch.int64)
encoder_lens = torch.full((B,), encoder_len, device=device).to(torch.int64)
attn_logits = torch.empty(
(B, H_Q, num_kv_splits, D_V + 1),
@@ -114,8 +127,8 @@ class TestDecodeAttention(CustomTestCase):
k_buffer,
v_buffer,
o,
key,
value,
key if not is_cross_attn else None,
value if not is_cross_attn else None,
loc,
attn_logits,
req_to_token,
@@ -123,6 +136,8 @@ class TestDecodeAttention(CustomTestCase):
b_seq_len,
sm_scale,
logit_cap,
is_cross_attn,
encoder_lens,
)
self._run_sdpa_forward_decode(
@@ -135,15 +150,16 @@ class TestDecodeAttention(CustomTestCase):
b_seq_len,
scaling=sm_scale,
enable_gqa=enable_gqa,
encoder_lens=encoder_lens,
is_cross_attn=is_cross_attn,
)
cos_sim = torch.nn.functional.cosine_similarity(
o.flatten(), o_grouped.flatten(), dim=0
)
self.assertGreater(cos_sim.item(), 0.99)
torch.testing.assert_close(o, o_grouped, atol=3e-2, rtol=1e-6)
def _test_grouped_decode_attention(self, device="cpu"):
def _test_grouped_decode_attention(self, device="cuda"):
configs = [
(2, 16, 16, 64, 64),
(2, 16, 1, 16, 16),
@@ -161,7 +177,10 @@ class TestDecodeAttention(CustomTestCase):
for B, H_Q, H_KV, D, D_V in configs:
for dtype in [torch.bfloat16, torch.float16]:
self._test_grouped_decode_attention_once(
B, H_Q, H_KV, D, D_V, dtype=dtype, device=device
B, H_Q, H_KV, D, D_V, False, dtype=dtype, device=device
)
self._test_grouped_decode_attention_once(
B, H_Q, H_KV, D, D_V, True, dtype=dtype, device=device
)
def test_grouped_decode_attention(self):
test/registered/cpu/test_extend.py
+60 -17
View File
@@ -24,9 +24,11 @@ class TestExtendAttention(CustomTestCase):
seq_lens: torch.Tensor,
extend_prefix_lens: torch.Tensor,
extend_seq_lens: torch.Tensor,
encoder_lens=None,
scaling=None,
enable_gqa=False,
causal=False,
is_cross_attn=False,
):
assert seq_lens.shape[0] == extend_prefix_lens.shape[0]
@@ -43,7 +45,14 @@ class TestExtendAttention(CustomTestCase):
seq_len_kv = seq_lens[seq_idx]
end_q = start_q + extend_seq_len_q
end_kv = start_kv + seq_len_kv
if encoder_lens is not None:
start_kv = 0 if is_cross_attn else encoder_lens[seq_idx]
end_kv = (
encoder_lens[seq_idx] if is_cross_attn else start_kv + seq_len_kv
)
else:
start_kv = 0
end_kv = start_kv + seq_len_kv
per_req_query = query[:, start_q:end_q, :]
per_req_query_redudant = torch.empty(
@@ -57,7 +66,7 @@ class TestExtendAttention(CustomTestCase):
# get key and value from cache. per_req_tokens contains the kv cache
# index for each token in the sequence.
req_pool_idx = req_pool_indices[seq_idx]
per_req_tokens = req_to_token[req_pool_idx, :seq_len_kv]
per_req_tokens = req_to_token[req_pool_idx, start_kv:end_kv]
per_req_key = k_cache[per_req_tokens].movedim(0, query.dim() - 2)
per_req_value = v_cache[per_req_tokens].movedim(0, query.dim() - 2)
@@ -86,6 +95,7 @@ class TestExtendAttention(CustomTestCase):
D,
DV,
mla=False,
is_cross_attn=False,
*,
b_seq_len_prefix=None,
b_seq_len_extend=None,
@@ -94,32 +104,36 @@ class TestExtendAttention(CustomTestCase):
if b_seq_len_prefix is None:
b_seq_len_prefix = torch.randint(1, N_CTX // 2, (B,), dtype=torch.int32)
if mla:
b_seq_len_prefix.zero_()
else:
b_seq_len_prefix = torch.as_tensor(b_seq_len_prefix, dtype=torch.int32)
encoder_lens = torch.randint(1, N_CTX // 2, (B,), dtype=torch.int64)
if mla:
b_seq_len_prefix.zero_()
encoder_lens.zero_()
if b_seq_len_extend is None:
b_seq_len_extend = torch.randint(1, N_CTX // 2, (B,), dtype=torch.int32)
else:
b_seq_len_extend = torch.as_tensor(b_seq_len_extend, dtype=torch.int32)
b_seq_len = b_seq_len_prefix + b_seq_len_extend
max_len_in_batch = torch.max(b_seq_len, 0)[0].item()
max_len_in_batch = (
torch.max(b_seq_len, 0)[0].item() + torch.max(encoder_lens, 0)[0].item()
)
b_req_idx = torch.arange(B, dtype=torch.int32)
req_to_tokens = torch.empty((B, max_len_in_batch), dtype=torch.int32)
b_start_loc = torch.zeros((B,), dtype=torch.int32)
b_start_loc[1:] = torch.cumsum(b_seq_len[:-1], 0)
b_start_loc[1:] = torch.cumsum(b_seq_len[:-1] + encoder_lens[:-1], 0)
b_start_loc_extend = torch.zeros((B,), dtype=torch.int32)
b_start_loc_extend[1:] = torch.cumsum(b_seq_len_extend[:-1], 0)
for i in range(B):
req_to_tokens[i, : b_seq_len[i]] = torch.arange(
b_start_loc[i], b_start_loc[i] + b_seq_len[i]
req_to_tokens[i, : b_seq_len[i] + encoder_lens[i]] = torch.arange(
b_start_loc[i], b_start_loc[i] + b_seq_len[i] + encoder_lens[i]
)
total_token_num = torch.sum(b_seq_len).item()
total_token_num = torch.sum(b_seq_len).item() + torch.sum(encoder_lens).item()
extend_token_num = torch.sum(b_seq_len_extend).item()
H_BUF = 1 if mla else H_KV
@@ -131,8 +145,10 @@ class TestExtendAttention(CustomTestCase):
q_extend = torch.empty((extend_token_num, H_Q, D), dtype=dtype)
for i in range(B):
extend_start_in_buffer = b_start_loc[i] + b_seq_len_prefix[i]
extend_end_in_buffer = b_start_loc[i] + b_seq_len[i]
extend_start_in_buffer = (
b_start_loc[i] + b_seq_len_prefix[i] + encoder_lens[i]
)
extend_end_in_buffer = b_start_loc[i] + b_seq_len[i] + encoder_lens[i]
extend_start = b_start_loc_extend[i]
extend_end = b_start_loc_extend[i] + b_seq_len_extend[i]
k_extend[extend_start:extend_end] = k_buffer[
@@ -178,7 +194,9 @@ class TestExtendAttention(CustomTestCase):
b_seq_len_extend,
scaling=sm_scale,
enable_gqa=enable_gqa,
causal=True,
causal=not is_cross_attn,
is_cross_attn=is_cross_attn,
encoder_lens=encoder_lens,
)
o_extend = torch.empty((extend_token_num, H_Q, DV), dtype=dtype)
@@ -197,16 +215,29 @@ class TestExtendAttention(CustomTestCase):
max_len_extend,
sm_scale,
logit_cap,
is_cross_attn,
encoder_lens,
)
torch.testing.assert_close(o_ref, o_extend, atol=1e-2, rtol=1e-2)
def test_extend_attention(self):
for is_mla in [True, False]:
self._test_extend_attention_once(1, 123, 1, 1, 128, 96, is_mla)
self._test_extend_attention_once(1, 123, 16, 1, 128, 96, is_mla)
self._test_extend_attention_once(4, 1230, 16, 4, 128, 96, is_mla)
self._test_extend_attention_once(1, 9000, 16, 1, 32, 32, is_mla)
for is_cross_attn in [True, False]:
if is_mla and is_cross_attn:
continue
self._test_extend_attention_once(
1, 123, 1, 1, 128, 96, is_mla, is_cross_attn
)
self._test_extend_attention_once(
1, 123, 16, 1, 128, 96, is_mla, is_cross_attn
)
self._test_extend_attention_once(
4, 1230, 16, 4, 128, 96, is_mla, is_cross_attn
)
self._test_extend_attention_once(
1, 9000, 16, 1, 32, 32, is_mla, is_cross_attn
)
def test_extend_attention_large_seq_causal_mask(self):
self._test_extend_attention_once(
@@ -220,6 +251,18 @@ class TestExtendAttention(CustomTestCase):
b_seq_len_extend=[5000],
)
def test_extend_attention_gqa_partial_extend_with_prefix(self):
self._test_extend_attention_once(
B=1,
N_CTX=256,
H_Q=16,
H_KV=4,
D=128,
DV=96,
b_seq_len_prefix=[97],
b_seq_len_extend=[37],
)
if __name__ == "__main__":
unittest.main()
test/registered/cpu/test_gemm.py
+10 -1
View File
@@ -35,6 +35,7 @@ class TestGemm(CustomTestCase):
N = [16, 32 * 13]
K = [32 * 16]
has_bias = [False, True]
dim = [2, 3, 4, 5]
M_int8 = [2, 128]
N_int8 = [32 * 12]
@@ -52,10 +53,16 @@ class TestGemm(CustomTestCase):
N_gptq = [4096]
K_gptq = [4096]
def _bf16_gemm(self, M, N, K, has_bias):
def _bf16_gemm(self, M, N, K, has_bias, dim):
mat1 = torch.randn(M, K, dtype=torch.bfloat16)
mat2 = torch.randn(N, K, dtype=torch.bfloat16)
if dim == 3:
mat1 = mat1.unsqueeze(0).repeat(2, 1, 1)
if dim == 4:
mat1 = mat1.unsqueeze(0).unsqueeze(0).repeat(2, 2, 1, 1)
if dim == 5:
mat1 = mat1.unsqueeze(0).unsqueeze(0).unsqueeze(0).repeat(2, 2, 2, 1, 1)
ref = torch.matmul(mat1.float(), mat2.float().t())
if has_bias:
@@ -83,12 +90,14 @@ class TestGemm(CustomTestCase):
self.N,
self.K,
self.has_bias,
self.dim,
):
with self.subTest(
M=params[0],
N=params[1],
K=params[2],
has_bias=params[3],
dim=params[4],
):
self._bf16_gemm(*params)
test/registered/cpu/test_mla.py
+2
View File
@@ -118,6 +118,8 @@ class TestMLA(CustomTestCase):
b_seq_len,
sm_scale,
logit_cap,
False,
None,
)
self._run_sdpa_forward_decode(