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)