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

fix: V tensor must be (B, n_h, hd, N) for FMHA — was transposed wrong

Browse Source
This commit is contained in:
biondizzle
2026-05-31 06:03:13 +00:00
parent 91abf0f921
commit e5245ea34e
1 changed files with 5 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
single_shot_inference.py
View File

@@ -368,16 +368,15 @@ def forward_layer(X_l, w, li, cfg, rope_cos, rope_sin,
# -- FMHA with sink bias correction (paper D5c) --
q_input = q_heads.permute(1, 0, 2) # (n_h, T, hd)
from dsv4.kernels.attention.production import dsv4_attention
# Get both output and LSE from the kernel
from dsv4.kernels.attention.fmha_multitile_op import fmha_multitile_decode_raw
scale = 1.0 / math.sqrt(hd)
q_4d = q_input.unsqueeze(0).contiguous() # (1, n_h, T, hd)
k_4d = k_full.unsqueeze(0).contiguous() # (1, 1, seq_len, hd)
v_4d = v_full.unsqueeze(0).contiguous() # (1, 1, seq_len, hd)
# CRITICAL: V must be (B, n_h, hd, N) — transposed!
v_4d = v_full.unsqueeze(0).transpose(-1, -2).contiguous() # (1, 1, hd, seq_len)
o_4d, lse = fmha_multitile_decode_raw(q_4d, k_4d, v_4d, scale) # (1, n_h, T, hd), (1, n_h, T)
o_4d, lse = fmha_multitile_decode_raw(q_4d, k_4d, v_4d, scale)
attn_out = o_4d.squeeze(0).permute(1, 0, 2) # (T, n_h, hd)
# Apply sink bias correction: scale the output by softmax_normalizer / (normalizer + exp(sink))
@@ -387,10 +386,10 @@ def forward_layer(X_l, w, li, cfg, rope_cos, rope_sin,
if sink_key in w and seq_len > 0:
sinks = w[sink_key].to(device=device) # (n_h,) BF16
# lse: (1, n_h, T) — log-sum-exp of attention scores per head
lse_3d = lse.squeeze(0).permute(1, 0) # (T, n_h)
lse_2d = lse.squeeze(0).t() # (T, n_h)
# For each head, compute the correction factor
sink_exp = torch.exp(sinks.float()) # (n_h,)
attn_exp = torch.exp(lse_3d.float()) # (T, n_h)
attn_exp = torch.exp(lse_2d.float()) # (T, n_h)
# Correction: attn_exp / (attn_exp + sink_exp)
correction = attn_exp / (attn_exp + sink_exp.unsqueeze(0) + 1e-10) # (T, n_h)
# Apply per-head correction