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

Reduce hot-path .item() syncs: gate li>=58 diagnostics behind VERBOSE>=2, topk on float

Browse Source
This commit is contained in:
biondizzle
2026-06-01 22:33:03 +00:00
parent 6f4bbc997a
commit e53645654d
1 changed files with 5 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
single_shot_inference.py
View File

@@ -499,7 +499,7 @@ def moe_forward(x, li, moe_runner, se_runner, router, token_id):
if VERBOSE >= 2 and (li < 3 or li >= 58):
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 >= 58:
if VERBOSE >= 2 and li >= 58:
print(f" L{li} MoE DIAG: topk_ids={topk_ids[0].tolist()} topk_w=[{','.join(f'{w:.3f}' for w in topk_w[0].tolist())}]", flush=True)
# Also print gate logits for debugging
if hasattr(router, '_gate_lin') and router._gate_lin is not None:
@@ -509,7 +509,7 @@ def moe_forward(x, li, moe_runner, se_runner, router, token_id):
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)
shared_out = se_runner.run(x)
if li >= 58:
if VERBOSE >= 2 and li >= 58:
print(f" L{li} MoE DIAG: |routed|={routed_out.abs().max().item():.1f} |shared|={shared_out.abs().max().item():.1f} |x|={x.abs().max().item():.1f}", flush=True)
if VERBOSE >= 2 and li < 3:
has_nan = torch.isnan(shared_out).any().item()
@@ -541,11 +541,11 @@ def forward_layer(X_l, w, li, cfg, rope_cos, rope_sin,
x_in_f, ctx_f = ffn_mhc.pre_block(X_mid); x_ffn = rmsnorm(x_in_f, ffn_norm_w)
F_ffn = moe_forward(x_ffn, li, moe_runner, se_runner, router, token_id)
X_next = ffn_mhc.post_block(X_mid, F_ffn, ctx_f)
if VERBOSE >= 1 and (li < 3 or li >= 58):
if VERBOSE >= 2 and (li < 3 or li >= 58):
print(f" L{li}: |X|={X_l.abs().max().item():.1f}->{X_next.abs().max().item():.1f} "
f"|Fa|={F_attn.abs().max().item():.1f} |Ff|={F_ffn.abs().max().item():.1f}", flush=True)
# Detailed diagnostics for last 3 layers or any layer with explosive growth
if li >= 58 or (li > 0 and X_next.abs().max().item() > 200):
# Detailed diagnostics — only with VERBOSE >= 2 to avoid .item() syncs on hot path
if VERBOSE >= 2 and (li >= 58 or (li > 0 and X_next.abs().max().item() > 200)):
A_a, B_a, C_a = attn_mhc._dynamic_params(X_l)
A_f, B_f, C_f = ffn_mhc._dynamic_params(X_mid)
print(f" L{li} DIAG: A_attn=[{A_a.min().item():.4f},{A_a.max().item():.4f}] "