From e5245ea34e3ae556117ea6f1a109a58523fde000 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Sun, 31 May 2026 06:03:13 +0000 Subject: [PATCH] =?UTF-8?q?fix:=20V=20tensor=20must=20be=20(B,=20n=5Fh,=20?= =?UTF-8?q?hd,=20N)=20for=20FMHA=20=E2=80=94=20was=20transposed=20wrong?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- single_shot_inference.py | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/single_shot_inference.py b/single_shot_inference.py index b84f5f72..e3503ea2 100644 --- a/single_shot_inference.py +++ b/single_shot_inference.py @@ -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