Fix MoE routing: hash layers 0-2 (tid2eid), e_score_correction_bias for layers 3+

- Layers 0-2 use hash routing (tid2eid lookup, uniform weights)
- Layers 3+ use noaux_tc (sqrt(softplus) + e_score_correction_bias for selection only)
- Fixed e_bias key name: e_score_correction_bias (not e_bias)
- Hash routing detection: check tid2eid present AND e_score_correction_bias absent

single_shot_inference.py

@@ -621,39 +621,37 @@ def moe_forward(x, w, li, cfg, token_id, device):
     routed_scaling = cfg.get("routed_scaling_factor", 2.5)
     swiglu_limit = cfg.get("swiglu_limit", 10.0)
     mlp_inter = cfg["moe_intermediate_size"]
-    is_hash = li < 3
-
     # ---- Routing ----
+    # Layers 0-2: hash routing (tid2eid lookup)
+    # Layers 3+: noaux_tc (sqrt(softplus) scoring + e_score_correction_bias for selection only)
+    # Config: topk_method='noaux_tc', scoring_func='sqrtsoftplus'
     expert_ids = None
     expert_weights = None
 
-    if is_hash:
-        tid2eid_key = f"model.layers.{li}.mlp.gate.tid2eid"
-        if tid2eid_key in w:
-            tid2eid = w[tid2eid_key]
-            tid = token_id.item() if token_id.numel() == 1 else token_id[0].item()
-            expert_ids = tid2eid[tid]  # (top_k,) int64
-            expert_weights = torch.ones(top_k, dtype=torch.float32, device=x.device) / top_k
-        else:
-            # Fallback: use dense routing even for hash layers
-            is_hash = False
+    tid2eid_key = f"model.layers.{li}.mlp.gate.tid2eid"
+    e_bias_key = f"model.layers.{li}.mlp.gate.e_score_correction_bias"
+    is_hash = tid2eid_key in w and e_bias_key not in w
 
-    if not is_hash:
-        # Dense routing: sqrt(softplus(X @ W_gate)) + e_bias for selection
+    if is_hash:
+        # Hash routing: deterministic per-token lookup, uniform weights
+        tid2eid = w[tid2eid_key]
+        tid = token_id.item() if token_id.numel() == 1 else token_id[0].item()
+        expert_ids = tid2eid[tid]  # (top_k,) int64
+        expert_weights = torch.ones(top_k, dtype=torch.float32, device=x.device) / top_k
+    else:
+        # Dense routing: sqrt(softplus(logits)) scoring
         gate_w = w[f"model.layers.{li}.mlp.gate.weight"]  # (H, n_experts) BF16
         logits = torch.nn.functional.linear(x, gate_w.bfloat16())  # (T, n_experts)
-        # Activation: sqrt(softplus(logits))
-        activated = torch.sqrt(torch.nn.functional.softplus(logits.float()) + 1e-6)
-        # e_bias: learned per-expert bias for SELECTION ONLY (not in weights)
-        e_bias_key = f"model.layers.{li}.mlp.gate.e_bias"
+        # Scoring: sqrt(softplus(logits))
+        scores = torch.sqrt(torch.nn.functional.softplus(logits.float()) + 1e-6)
+        # e_score_correction_bias: per-expert bias for SELECTION ONLY
+        selection_logits = scores.clone()
         if e_bias_key in w:
-            activated = activated + w[e_bias_key].float().unsqueeze(0)
-        # Top-k
-        scores, indices = activated.topk(top_k, dim=-1)  # (T, top_k)
-        # Renormalize on UNBIASED activation (no e_bias in weights)
-        unbiased = torch.sqrt(torch.nn.functional.softplus(logits.float()) + 1e-6)
-        unbiased_scores = torch.gather(unbiased, -1, indices)
-        expert_weights = unbiased_scores / unbiased_scores.sum(dim=-1, keepdim=True)
+            selection_logits = selection_logits + w[e_bias_key].float().unsqueeze(0)
+        _, indices = selection_logits.topk(top_k, dim=-1)  # (T, top_k)
+        # Weights from UNBIASED scores (no e_bias)
+        expert_weights = torch.gather(scores, -1, indices)
+        expert_weights = expert_weights / expert_weights.sum(dim=-1, keepdim=True)
         # For T=1 decode, squeeze
         if x.shape[0] == 1:
             expert_ids = indices[0]