"csrc/vscode:/vscode.git/clone" did not exist on "e8cb0d04950c9e0b39d00f3744806c19dff4b4bc"
Unverified Commit 9fb3ae4e authored by Wentao Ye's avatar Wentao Ye Committed by GitHub
Browse files

[Bug] Fix DeepGEMM Attention Test (#26423)


Signed-off-by: default avataryewentao256 <zhyanwentao@126.com>
parent 76afe4ed
......@@ -184,6 +184,7 @@ ba = "ba"
[tool.typos.type.py.extend-words]
ba = "ba"
nd = "nd"
[tool.typos.type.cpp]
extend-glob = ["*.cu"]
......
......@@ -82,8 +82,7 @@ def _ref_fp8_mqa_logits(
torch.arange(0, seq_len_kv, device="cuda")[None, :] < cu_seqlen_ke[:, None]
)
mask = mask_lo & mask_hi
score = torch.einsum("mhd,and->hmn", q, k)
score = torch.einsum("mhd,nd->hmn", q, k)
logits = (score.relu() * weights.unsqueeze(-1).transpose(0, 1)).sum(dim=0)
logits = logits.masked_fill(~mask, float("-inf"))
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment