biondizzle/nvfp4-megamoe-kernel
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Fix causal mask shape for SDPA: (1,1,T,T) broadcast

Browse Source
This commit is contained in:
biondizzle
2026-05-19 08:00:39 +00:00
parent 57615029a4
commit 1e675ccc9a
1 changed files with 2 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
cutedsl/csa_attention.py
View File

@@ -409,8 +409,8 @@ def full_attention_reference(
# Q: (T, NH, HD) → (T*NH, 1, HD)
q_2d = q.reshape(T * NH, 1, HD)
# Causal mask
causal_mask = torch.tril(torch.ones(T, T, device=q.device, dtype=torch.bool)).unsqueeze(0)
# Causal mask: (1, 1, T, T) broadcast over batch dim T*NH
causal_mask = torch.tril(torch.ones(T, T, device=q.device, dtype=torch.bool)).unsqueeze(0).unsqueeze(0)
out = F.scaled_dot_product_attention(
q_2d, k_2d, v_2d,