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

diag: compare production FMHA and NVFP4 linear output with PyTorch reference

Browse Source
This commit is contained in:
biondizzle
2026-06-01 02:12:39 +00:00
parent df394f8b40
commit 1eed28dd09
1 changed files with 14 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

15
single_shot_inference.py
View File

@@ -310,6 +310,12 @@ def forward_attention(x_normed, w, li, cfg, rope_cos, rope_sin,
# 1. Q: q_a (NVFP4 GEMM) → q_a_norm → q_b (NVFP4 GEMM) → q_b_norm
q_a = prod_lin['q_a'](x_normed)
if li < 3:
# Compare q_a with PyTorch reference
q_a_ref = do_nvfp4_linear_ref(x_normed, w, pfx, 'q_a_proj')
if q_a_ref is not None:
cos_qa = torch.nn.functional.cosine_similarity(q_a.flatten().float(), q_a_ref.flatten().float(), dim=0).item()
print(f" L{li} q_a: |prod|={q_a.abs().max().item():.6f} |ref|={q_a_ref.abs().max().item():.6f} cos={cos_qa:.6f}", flush=True)
q_norm_w = w.get(f"{pfx}.q_a_norm.weight")
if q_norm_w is not None: q_a = rmsnorm(q_a, q_norm_w.to(dev, torch.float32))
q = prod_lin['q_b'](q_a); q = unweighted_rmsnorm(q).bfloat16()
@@ -354,7 +360,14 @@ def forward_attention(x_normed, w, li, cfg, rope_cos, rope_sin,
# 6. Production FMHA
attn_out = _run_production_fmha(q_heads, all_kv, n_h, hd, T, seq_len, scale, dev, li, w, pfx)
if li < 3:
print(f" L{li} FMHA: |attn_out|={attn_out.abs().max().item():.6f} q_heads_range=[{q_heads.min().item():.3f},{q_heads.max().item():.3f}] all_kv_range=[{all_kv.min().item():.3f},{all_kv.max().item():.3f}] N={seq_len} hd={hd} scale={scale:.6f}", flush=True)
# Compare with PyTorch reference
k_exp = all_kv.unsqueeze(0).expand(n_h, -1, -1).contiguous()
v_exp = k_exp.clone()
q_in = q_heads.permute(1, 0, 2)
ref_scores = torch.matmul(q_in, k_exp.transpose(-1, -2)) * scale
ref_attn = torch.matmul(torch.softmax(ref_scores.float(), -1).bfloat16(), v_exp).permute(1, 0, 2)
cos_sim = torch.nn.functional.cosine_similarity(attn_out.flatten().float(), ref_attn.flatten().float(), dim=0).item()
print(f" L{li} FMHA: |prod|={attn_out.abs().max().item():.6f} |ref|={ref_attn.abs().max().item():.6f} cos={cos_sim:.6f}", flush=True)
# 7. Inverse RoPE
attn_out = _apply_rope(attn_out, positions, rope_cos, rope_sin, rd, inverse=True)