xinyun/vllm
1
0
Fork 0
mirror of https://git.datalinker.icu/vllm-project/vllm.git synced 2025-12-10 00:45:01 +08:00
Code Issues Packages Projects Releases Wiki Activity

[Bug] Fix DeepGEMM Attention Test (#26423)

Browse Source 
Signed-off-by: yewentao256 <zhyanwentao@126.com>
This commit is contained in:
Wentao Ye 2025-10-08 12:23:41 -04:00 committed by GitHub
parent 76afe4edf8
commit 9fb3ae4e6f
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 2 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
pyproject.toml
View File

@ -184,6 +184,7 @@ ba = "ba"
[tool.typos.type.py.extend-words] [tool.typos.type.py.extend-words]
ba = "ba" ba = "ba"
nd = "nd"
[tool.typos.type.cpp] [tool.typos.type.cpp]
extend-glob = ["*.cu"] extend-glob = ["*.cu"]

3
tests/kernels/attention/test_deepgemm_attention.py
View File

@ -82,8 +82,7 @@ def _ref_fp8_mqa_logits(
torch.arange(0, seq_len_kv, device="cuda")[None, :] < cu_seqlen_ke[:, None] torch.arange(0, seq_len_kv, device="cuda")[None, :] < cu_seqlen_ke[:, None]
) )
mask = mask_lo & mask_hi mask = mask_lo & mask_hi
score = torch.einsum("mhd,nd->hmn", q, k)
score = torch.einsum("mhd,and->hmn", q, k)
logits = (score.relu() * weights.unsqueeze(-1).transpose(0, 1)).sum(dim=0) logits = (score.relu() * weights.unsqueeze(-1).transpose(0, 1)).sum(dim=0)
logits = logits.masked_fill(~mask, float("-inf")) logits = logits.masked_fill(~mask, float("-inf"))