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

PART A: add detailed blowup diagnostics — capture mHC intermediate values when |X| > 1e6

Browse Source
This commit is contained in:
biondizzle
2026-06-03 06:10:33 +00:00
parent 6459fbca9a
commit 262f844e2e
1 changed files with 39 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

49
tests/unit/test_part_a_decode_diagnostics.py
View File

@@ -240,6 +240,7 @@ def main():
if X.device != torch.device(dev): X = X.to(dev)
torch.cuda.set_device(gpu)
X_prev = X.clone() # Save for blowup diagnostics
X_in_mag = X.abs().max().item()
X = forward_layer(X, layer_w[li], li, cfg, *rope_caches[gpu],
attn_mhcs.get(li), ffn_mhcs.get(li), attn_norms.get(li), ffn_norms.get(li),
@@ -254,16 +255,44 @@ def main():
if pi < 3 or pi == len(input_ids) - 1:
print(f" {pi:>3} {li:>3} {X_in_mag:>12.2f} {X_out_mag:>12.2f} {ratio:>5} {kc.n_comp:>6} {kc.swa_len:>4}", flush=True)
# Early abort if |X| blows up
if X_out_mag > 1e10:
print(f" *** BLOWUP at token {pi} layer {li}: |X|={X_out_mag:.2e} — ABORTING ***", flush=True)
print(f" This means the production pipeline is numerically unstable.", flush=True)
print(f" Check: mHC residual growth, NVFP4 quantization, MoE scaling.", flush=True)
# Print KV cache state at this point
for l2 in range(li + 1):
kc2 = kv_caches[l2]
r2 = cr[l2] if l2 < len(cr) else 128
print(f" L{l2} (ratio={r2}): n_comp={kc2.n_comp} swa_len={kc2.swa_len}", flush=True)
# Early abort if |X| blows up — run detailed diagnostics on THIS layer
if X_out_mag > 1e6:
print(f" *** BLOWUP at token {pi} layer {li}: |X|={X_out_mag:.2e} ***", flush=True)
print(f" Re-running layer {li} with detailed diagnostics...", flush=True)
# Re-run the SAME input through forward_layer but capture intermediates
X_diag = X_prev.clone() # X before this layer
attn_mhc_d = attn_mhcs.get(li)
ffn_mhc_d = ffn_mhcs.get(li)
A_l_a, B_l_a, C_l_a = attn_mhc_d._dynamic_params(X_diag)
ctx_a_d = mHCContext(B_l=B_l_a, C_l=C_l_a)
x_quant_attn = mhc_rmsnorm_quantize_nvfp4(
X_diag, A_l_a, attn_norms.get(li).to(dev, torch.float32))
x_normed = dequantize_nvfp4(x_quant_attn.x_fp4, x_quant_attn.x_sf, x_quant_attn.gsa)
print(f" |x_normed|={x_normed.abs().max().item():.2f} gsa={x_quant_attn.gsa}", flush=True)
F_attn_d, q_a_d = forward_attention(
x_normed, layer_w[li], li, cfg, *rope_caches[gpu],
kv_caches[li], pos, compressors.get(li), indexers.get(li), prod_lins.get(li),
x_quant=x_quant_attn)
print(f" |F_attn|={F_attn_d.abs().max().item():.2f}", flush=True)
X_mid_d = attn_mhc_d.post_block(X_diag, F_attn_d, ctx_a_d)
print(f" |X_mid|={X_mid_d.abs().max().item():.2f} B_l_row=[{B_l_a.sum(-1).min().item():.4f},{B_l_a.sum(-1).max().item():.4f}] C_l=[{C_l_a.min().item():.4f},{C_l_a.max().item():.4f}]", flush=True)
A_l_f, B_l_f, C_l_f = ffn_mhc_d._dynamic_params(X_mid_d)
ctx_f_d = mHCContext(B_l=B_l_f, C_l=C_l_f)
x_quant_ffn = mhc_rmsnorm_quantize_nvfp4(
X_mid_d, A_l_f, ffn_norms.get(li).to(dev, torch.float32))
x_ffn = dequantize_nvfp4(x_quant_ffn.x_fp4, x_quant_ffn.x_sf, x_quant_ffn.gsa)
F_ffn_d = moe_forward(x_ffn, li, moe_runners.get(li), se_runners.get(li),
routers.get(li), tid32.to(dev))
print(f" |F_ffn|={F_ffn_d.abs().max().item():.2f}", flush=True)
X_next_d = ffn_mhc_d.post_block(X_mid_d, F_ffn_d, ctx_f_d)
print(f" |X_next|={X_next_d.abs().max().item():.2e}", flush=True)
# Check per-component magnitudes
BX = torch.bmm(ctx_a_d.B_l.transpose(-1, -2), X_diag.float())
CF = ctx_a_d.C_l.unsqueeze(-1) * F_attn_d.unsqueeze(1)
print(f" |B@X|={BX.abs().max().item():.2f} |C*F|={CF.abs().max().item():.2f}", flush=True)
BX_f = torch.bmm(ctx_f_d.B_l.transpose(-1, -2), X_mid_d.float())
CF_f = ctx_f_d.C_l.unsqueeze(-1) * F_ffn_d.unsqueeze(1)
print(f" FFN: |B@X|={BX_f.abs().max().item():.2f} |C*F|={CF_f.abs().max().item():.2f}", flush=True)
return 1
if pi % 5 == 0: