minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 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):