A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 05:30 UTC
[Bug] Fix DeepGEMM Attention Test (#26423)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Oct 8, 2025 at 12:23 UTC
9fb3ae4e6fec167232367b55ec85065545bf379d
2 files changed
+2
-2
pyproject.toml
+1
@@ -184,6 +184,7 @@ ba = "ba"
184
185
[tool.typos.type.py.extend-words]
186
ba = "ba"
187
+nd = "nd"
188
189
[tool.typos.type.cpp]
190
extend-glob = ["*.cu"]
tests/kernels/attention/test_deepgemm_attention.py
+1
-2
@@ -82,8 +82,7 @@ def _ref_fp8_mqa_logits(
82
torch.arange(0, seq_len_kv, device="cuda")[None, :] < cu_seqlen_ke[:, None]
83
)
84
mask = mask_lo & mask_hi
85
-
86
- score = torch.einsum("mhd,and->hmn", q, k)
85
+ score = torch.einsum("mhd,nd->hmn", q, k)
86
logits = (score.relu() * weights.unsqueeze(-1).transpose(0, 1)).sum(dim=0)
87
logits = logits.masked_fill(~mask, float("-inf"))
88