diff --git a/vllm/patches/deepseek_v4.py b/vllm/patches/deepseek_v4.py index 8d50d383..5c674237 100644 --- a/vllm/patches/deepseek_v4.py +++ b/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