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

tune: reduce verbose diagnostics, print every decode step

Browse Source
This commit is contained in:
biondizzle
2026-06-01 05:40:48 +00:00
parent df8acae66b
commit 311fae490f
1 changed files with 9 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

18
single_shot_inference.py
View File

@@ -350,7 +350,7 @@ 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:
if VERBOSE >= 2 and 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:
@@ -399,7 +399,7 @@ 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:
if VERBOSE >= 2 and li < 3:
# Compare with PyTorch reference
k_exp = all_kv.unsqueeze(0).expand(n_h, -1, -1).contiguous()
v_exp = k_exp.clone()
@@ -419,7 +419,7 @@ def forward_attention(x_normed, w, li, cfg, rope_cos, rope_sin,
a_grp = a_flat.reshape(T, o_groups, gid); oa_3d = oa_bf.reshape(o_groups, o_rank, gid)
g_out = torch.bmm(a_grp.permute(1, 0, 2), oa_3d.transpose(1, 2))
g_flat = g_out.permute(1, 0, 2).reshape(T, o_groups * o_rank)
if li < 3:
if VERBOSE >= 2 and li < 3:
print(f" L{li} wo_a: |g_flat|={g_flat.abs().max().item():.6f} shape={g_flat.shape}", flush=True)
F_attn = prod_lin['o_b'](g_flat)
else:
@@ -430,7 +430,7 @@ def forward_attention(x_normed, w, li, cfg, rope_cos, rope_sin,
else:
log.warning(f"L{li}: No o_a_proj weight, zero attention output")
F_attn = torch.zeros(T, cfg["hidden_size"], dtype=torch.bfloat16, device=dev)
if li < 3:
if VERBOSE >= 2 and li < 3:
print(f" L{li} F_attn: |F_attn|={F_attn.abs().max().item():.6f}", flush=True)
return F_attn, q_a
@@ -444,13 +444,13 @@ def moe_forward(x, li, moe_runner, se_runner, router, token_id):
torch.cuda.synchronize(x.device)
if topk_ids.max().item() >= 384 or topk_ids.min().item() < 0:
print(f" L{li} BAD topk_ids: min={topk_ids.min().item()} max={topk_ids.max().item()}", flush=True)
if li < 3:
if VERBOSE >= 2 and li < 3:
print(f" L{li} MoE input: |x|={x.abs().max().item():.4f} has_nan={torch.isnan(x).any().item()}", flush=True)
routed_out = moe_runner.run(x, topk_w, topk_ids)
if li < 3:
if VERBOSE >= 2 and li < 3:
print(f" L{li} MoE routed: |out|={routed_out.abs().max().item():.4f} has_nan={torch.isnan(routed_out).any().item()}", flush=True)
shared_out = se_runner.run(x)
if li < 3:
if VERBOSE >= 2 and li < 3:
has_nan = torch.isnan(shared_out).any().item()
out_max = shared_out.abs().max().item() if not has_nan else float('nan')
print(f" L{li} MoE shared: |out|={out_max:.4f} has_nan={has_nan}", flush=True)
@@ -794,7 +794,7 @@ def main():
err = torch.cuda.current_stream(gpu).query()
print(f" CRASH at token {pi} layer {li} gpu {gpu}: {e}", flush=True)
raise
if pi == 0 and li < 3:
if VERBOSE >= 2 and pi == 0 and li < 3:
torch.cuda.synchronize(gpu)
print(f" Token {pi} L{li}: OK |X|={X.abs().max().item():.1f}", flush=True)
X = X.to('cuda:0'); torch.cuda.set_device(0)
@@ -829,7 +829,7 @@ def main():
next_id = torch.argmax(logits, -1).item(); all_tokens.append(next_id)
dt = time.time() - t1
has_nan = torch.isnan(logits.float()).any().item()
if step % 5 == 0 or has_nan:
if step % 1 == 0 or has_nan:
tv, ti = torch.topk(logits[0], 5)
top5 = ' '.join(f'{tokenizer.decode([t.item()])}({v.item():.1f})' for t, v in zip(ti[:5], tv[:5]))
print(f" Step {step}: {next_id} '{tokenizer.decode([next_id])}' ({dt:.2f}s) "