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

fix: move token_ids to same GPU as router (was cuda:0 but router on cuda:N)

Browse Source
This commit is contained in:
biondizzle
2026-06-01 01:49:40 +00:00
parent 7804b779ce
commit 905623793b
1 changed files with 4 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
single_shot_inference.py
View File

@@ -385,8 +385,10 @@ def forward_attention(x_normed, w, li, cfg, rope_cos, rope_sin,
# MoE — production kernels
# =====================================================================
def moe_forward(x, li, moe_runner, se_runner, router, token_id):
topk_w, topk_ids = router(x, token_ids=token_id)
# Diag: validate router output before MoE
# Ensure token_id is on same GPU as router
token_id_dev = token_id.to(x.device) if token_id.device != x.device else token_id
topk_w, topk_ids = router(x, token_ids=token_id_dev)
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)
routed_out = moe_runner.run(x, topk_w, topk_ids)