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

diag: cuda sync check after shared expert for first 3 layers

Browse Source
This commit is contained in:
biondizzle
2026-06-01 02:56:28 +00:00
parent b85fcf4d6f
commit 22d7cc9b7a
1 changed files with 7 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
single_shot_inference.py
View File

@@ -421,7 +421,13 @@ def moe_forward(x, li, moe_runner, se_runner, router, token_id):
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)
try:
torch.cuda.synchronize(x.device)
cuda_ok = True
except RuntimeError as e:
cuda_ok = False
print(f" L{li} SE CUDA ERROR: {e}", flush=True)
print(f" L{li} MoE shared: |out|={shared_out.abs().max().item():.4f} has_nan={torch.isnan(shared_out).any().item()} cuda_ok={cuda_ok}", flush=True)
print(f" L{li} SE gsa: l1={se_runner._l1_activation_global_scale:.6f} l2={se_runner._l2_activation_global_scale:.6f} gsb: l1={se_runner._l1_gsb[0].item():.6f} l2={se_runner._l2_gsb[0].item():.6f}", flush=True)
return routed_out + shared_out