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

diag: MoE input/output diagnostics for first 3 layers

Browse Source
This commit is contained in:
biondizzle
2026-06-01 02:41:12 +00:00
parent 856a459a98
commit 48d93a6d2e
1 changed files with 6 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
single_shot_inference.py
View File

@@ -414,8 +414,14 @@ 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:
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:
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:
print(f" L{li} MoE shared: |out|={shared_out.abs().max().item():.4f} has_nan={torch.isnan(shared_out).any().item()}", flush=True)
return routed_out + shared_out
# =====================================================================