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

diag: sync+error check after each layer on first token

Browse Source
This commit is contained in:
biondizzle
2026-06-01 01:26:50 +00:00
parent 91b3929605
commit f5fa84016e
1 changed files with 16 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

23
single_shot_inference.py
View File

@@ -678,13 +678,22 @@ def main():
gpu = li % NUM_GPUS
if X.device != torch.device(f"cuda:{gpu}"): X = X.to(f"cuda:{gpu}")
torch.cuda.set_device(gpu)
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),
kv_caches[li], pos, tid,
compressors.get(li), indexers.get(li),
moe_runners.get(li), se_runners.get(li), routers.get(li),
prod_lin=prod_lins.get(li))
try:
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),
kv_caches[li], pos, tid,
compressors.get(li), indexers.get(li),
moe_runners.get(li), se_runners.get(li), routers.get(li),
prod_lin=prod_lins.get(li))
except Exception as e:
torch.cuda.synchronize()
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:
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)
if pi % 10 == 0: print(f" Token {pi}/{len(generated)}: {time.time()-t1:.2f}s", flush=True)
print(f" Prefill done ({time.time()-t0:.1f}s)")