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

diag: print expected unnorm P@V for comparison with raw kernel output

Browse Source
This commit is contained in:
biondizzle
2026-05-23 01:28:32 +00:00
parent e5b4da7274
commit 9347119ade
1 changed files with 14 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
tests/unit/test_fmha_v3_stage_c.py
View File

@@ -436,14 +436,26 @@ def test():
v = torch.randn(n, hd, dtype=torch.bfloat16, device='cuda')
v_kernel = v.unsqueeze(-1)
c = torch.zeros(m, hd, 1, dtype=torch.bfloat16, device='cuda')
debug = torch.zeros(4, dtype=torch.float32, device='cuda') # [row_sum, row_max, inv_row_sum, 0]
qf = q[:, :, 0].float()
kf = k[:, :, 0].float()
scale = 1.0 / math.sqrt(hd)
attn = qf @ kf.T * scale
attn = torch.softmax(attn, dim=-1)
attn_raw = qf @ kf.T * scale
attn = torch.softmax(attn_raw, dim=-1)
ref = attn @ v.float()
# Expected stats for comparison
print(f' row_sum (should be 1.0): {attn.sum(dim=-1)[:4].tolist()}')
# Unnormalized softmax: exp(S - max)
S_max = attn_raw.max(dim=-1, keepdim=True).values
P_unnorm = torch.exp(attn_raw - S_max)
unnorm_pv = P_unnorm @ v.float()
unnorm_sum = P_unnorm.sum(dim=-1)
print(f' unnorm row_sum: {unnorm_sum[:4].tolist()}')
print(f' unnorm P@V[0,:4]: {unnorm_pv[0,:4].tolist()}')
print(f' kernel out[0,:4] should match unnorm P@V (no normalize)')
mQ = ct.from_dlpack(q).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q))
mK = ct.from_dlpack(k).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k))
mV = ct.from_dlpack(v_kernel).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_kernel))