debug more

vllm/patches/deepseek_v4.py

@@ -530,6 +530,21 @@ class DeepseekV4MegaMoEExperts(nn.Module):
         symm_buffer.experts_start_idx = self.experts_start_idx
         num_tokens = hidden_states.shape[0]
 
+        # NaN-trace: check hidden_states BEFORE staging
+        if int(os.environ.get('MEGA_MOE_DEBUG', '0')):
+            hs_f32 = hidden_states.to(torch.float32)
+            nan_frac = torch.isnan(hs_f32).float().mean().item()
+            nan_per_tok = torch.isnan(hs_f32).any(dim=-1)
+            bad_tok_ids = nan_per_tok.nonzero(as_tuple=True)[0][:10].tolist()
+            finite_mask = ~torch.isnan(hs_f32) & ~torch.isinf(hs_f32)
+            finite_max = hs_f32[finite_mask].abs().max().item() if finite_mask.any() else float('nan')
+            print(f"[PRE-stage] hidden_states: "
+                  f"nan_frac={nan_frac:.4f} "
+                  f"inf_any={torch.isinf(hs_f32).any().item()} "
+                  f"finite_max={finite_max:.4e} "
+                  f"first_bad_token_idxs={bad_tok_ids} "
+                  f"experts_start_idx={self.experts_start_idx}")
+
         # Quantize activation using the kernel's PyTorch stage_activation
         # (same code path the kernel uses for L1→L2 requantization).
         # This replaces the broken Triton staging kernel — no more uint32