From 905623793b0c67b3b2155a5ba01da6e7462fd191 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Mon, 1 Jun 2026 01:49:40 +0000 Subject: [PATCH] fix: move token_ids to same GPU as router (was cuda:0 but router on cuda:N) --- single_shot_inference.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/single_shot_inference.py b/single_shot_inference.py index 8f11ec17..7812a363 100644 --- a/single_shot_inference.py +++ b/single_shot_inference.py @@ -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)