add SKIP_ROUTED_MOE debug flag, re-enable sinks

single_shot_inference.py

@@ -48,6 +48,7 @@ CHECKPOINT_DIR = "/root/nvidia-meeting/DeepSeek-V4-Pro-NVFP4"
 MAX_NEW_TOKENS = 10
 PROMPT = "The capital of France is"
 NUM_GPUS = 8
+SKIP_ROUTED_MOE = False  # If True, only use shared expert (debug)
 
 # =====================================================================
 # NVFP4 dequantization — matches checkpoint format exactly
@@ -377,9 +378,8 @@ def forward_layer(X_l, w, li, cfg, rope_cos, rope_sin,
         k_expanded = k_full.expand(n_h, -1, -1).contiguous()
         v_expanded = v_full.expand(n_h, -1, -1).contiguous()
         # Attention sink (paper D5c)
-        # DISABLED for now to check impact
         sink_key = f"{pre}.sinks"
-        if False and sink_key in w and seq_len > 0:
+        if sink_key in w and seq_len > 0:
             sinks = w[sink_key].to(device=device)  # (n_h,) BF16
             sink_k = torch.zeros(n_h, 1, hd, dtype=torch.bfloat16, device=device)
             sink_v = torch.zeros(n_h, 1, hd, dtype=torch.bfloat16, device=device)
@@ -521,8 +521,9 @@ def moe_forward(x, w, li, cfg, token_id, device):
     # ---- Run selected experts ----
     T = x.shape[0]
     expert_outputs = []
-    for i, eid in enumerate(expert_ids):
-        eid_int = eid.item()
+    if not SKIP_ROUTED_MOE:
+        for i, eid in enumerate(expert_ids):
+            eid_int = eid.item()
         epre = f"model.layers.{li}.mlp.experts.{eid_int}"
 
         gate = nvfp4_linear(x,