From b0221662e73f323648c5f05aace373b1e52ddfe0 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Sun, 17 May 2026 11:38:19 +0000 Subject: [PATCH] Fix warmup: pass local expert IDs (not global), remove incorrect _warmup_done guard MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit compute_activation_global_scales expects local IDs (0..num_experts-1), not global IDs. EP5/EP7 were getting L2 gs=0 because global IDs (240+, 336+) didn't match expert_id_range (0..47), so no tokens matched any expert → L1 GEMM got zero inputs → L2 gs=0 → NaN/crash. Also removed _warmup_done guard since each layer needs its own warmup (different weights, different gs values). --- vllm/patches/deepseek_v4.py | 12 ++++-------- 1 file changed, 4 insertions(+), 8 deletions(-) diff --git a/vllm/patches/deepseek_v4.py b/vllm/patches/deepseek_v4.py index 8208d449..8070bedb 100644 --- a/vllm/patches/deepseek_v4.py +++ b/vllm/patches/deepseek_v4.py @@ -551,19 +551,14 @@ class DeepseekV4MegaMoEExperts(nn.Module): # the actual amax and compute correct gs values. self._warmup_activation_global_scales() - _warmup_done: bool = False - def _warmup_activation_global_scales(self) -> None: """Run a warmup forward pass to compute correct activation global scales. - Called once during finalize_weights, before cudagraph capture. + Called once per layer during finalize_weights, before cudagraph capture. Uses quantize_to_nvfp4 (which calls .max()) to get the exact gs from real activation magnitudes, then stores them for use by quantize_activation_nvfp4 (no .max(), cudagraph-safe). """ - if self._warmup_done: - return - self._warmup_done = True import torch runner = self._cutedsl_runner device = runner.device @@ -574,11 +569,12 @@ class DeepseekV4MegaMoEExperts(nn.Module): # Sample hidden states: typical BF16 activations have amax ~1-10 hidden_states = torch.randn(num_tokens, runner.hidden_size, dtype=torch.bfloat16, device=device) - # Assign all tokens to the first few experts evenly + # Assign all tokens to local experts (0..num_local_experts-1) + # compute_activation_global_scales expects local IDs, not global topk_ids = torch.zeros(num_tokens, top_k, dtype=torch.int64, device=device) for i in range(num_tokens): for j in range(top_k): - topk_ids[i, j] = (runner.experts_start_idx + j) % (runner.experts_start_idx + runner.num_experts) + topk_ids[i, j] = j % runner.num_experts topk_weights = torch.ones(num_tokens, top_k, dtype=torch.float32, device=device) / top_k runner.compute_activation_global_scales(