From 7fbbdc52047775528ff9a23040279159bf7b59f3 Mon Sep 17 00:00:00 2001
From: biondizzle <biondizzle@gmail.com>
Date: Mon, 1 Jun 2026 01:27:16 +0000
Subject: [PATCH] diag: validate router output before MoE

---
 single_shot_inference.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/single_shot_inference.py b/single_shot_inference.py
index 146ab1ee..e1b24a25 100644
--- a/single_shot_inference.py
+++ b/single_shot_inference.py
@@ -380,6 +380,9 @@ def forward_attention(x_normed, w, li, cfg, rope_cos, rope_sin,
 # =====================================================================
 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
+    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)
     shared_out = se_runner.run(x)
     return routed_out + shared_out