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

diag: add direct segment 0 test to compare with run_segment

Browse Source
This commit is contained in:
biondizzle
2026-05-26 15:37:06 +00:00
parent 2a5f9dc6e3
commit 7f983fb855
1 changed files with 31 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

31
tests/unit/test_d5c_multitile.py
View File

@@ -129,6 +129,37 @@ def test_d5c_multitile():
attn_sink_val = 0.5
attn_sink = torch.tensor([attn_sink_val], dtype=torch.float32, device='cuda')
# Direct test of segment 0 (verify run_segment isn't the issue)
print('--- Direct segment 0 test ---')
k_seg0 = k_combined[0:seg_size]
v_seg0 = v_combined[0:seg_size]
c0 = torch.zeros(m, hd, 1, dtype=torch.bfloat16, device='cuda')
lse0 = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda')
rs0 = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda')
mQ0 = to_cute(q.unsqueeze(-1))
mK0 = to_cute(k_seg0.unsqueeze(-1))
mV0 = to_cute(v_seg0.unsqueeze(-1))
mC0 = to_cute(c0); mLSE0 = to_cute(lse0); mRS0 = to_cute(rs0)
mSB0 = to_cute(attn_sink)
# Compile FRESH for segment 0
k0 = FmhaKernel(head_dim=hd, s_k=seg_size, normalize=False,
apply_swa_mask=True, is_causal=False, n_comp=n_comp, apply_sink_bias=True)
comp0 = cute.compile(k0, mQ0, mK0, mV0, mC0, stream, mLSE0,
swa_len=swa_len, sink_bias=mSB0, row_sums=mRS0)
comp0(mQ0, mK0, mV0, mC0, stream, mLSE0,
swa_len=swa_len, sink_bias=mSB0, row_sums=mRS0)
torch.cuda.synchronize()
ok0 = c0[:, :, 0].float() / rs0[:, 0, 0].float().unsqueeze(1).clamp(min=1e-30)
# Per-segment reference for segment 0
scores0_ref = q.float() @ k_seg0.float().T * scale
scores0_ref[:, n_comp:] += attn_sink_val
if swa_len < (seg_size - n_comp): # 100 < 32? No
scores0_ref[:, n_comp + swa_len:] = float('-inf')
ref0 = (torch.softmax(scores0_ref, dim=-1) @ v_seg0.float()).to(torch.bfloat16)
cos0 = torch.nn.functional.cosine_similarity(ok0.flatten().unsqueeze(0), ref0.flatten().unsqueeze(0).float()).item()
print(f'Seg0 direct: cos {cos0:.6f} LSE_kern={lse0[0,0,0].item():.4f} LSE_ref={torch.logsumexp(scores0_ref, dim=-1)[0].item():.4f}')
print()
# Reference
ref = reference_combined_attention(
q, k_comp, v_comp, k_swa, v_swa,