From 04245b664b6d6515073074ee71cf841a0355eb3c Mon Sep 17 00:00:00 2001
From: biondizzle <biondizzle@gmail.com>
Date: Sun, 17 May 2026 10:48:24 +0000
Subject: [PATCH] Add warmup-based activation global scale computation in
 finalize_weights
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

The checkpoint input_scale is a calibration value that produces wrong gs
at runtime (too small → block scales saturate → garbage output → EOS).

Now calls compute_activation_global_scales() with sample data during weight
finalization, before cudagraph capture. This observes actual activation
magnitudes and computes correct L1 and L2 gs values.
---
 vllm/patches/deepseek_v4.py | 36 ++++++++++++++++++++++++++++++++++++
 1 file changed, 36 insertions(+)

diff --git a/vllm/patches/deepseek_v4.py b/vllm/patches/deepseek_v4.py
index 6b787da5..ef4d75f2 100644
--- a/vllm/patches/deepseek_v4.py
+++ b/vllm/patches/deepseek_v4.py
@@ -529,6 +529,7 @@ class DeepseekV4MegaMoEExperts(nn.Module):
             l2_igs = w2_igs[:, 0]
         else:
             l2_igs = w2_igs
+        # Use checkpoint input_scale as initial guess, then warmup will override
         self._cutedsl_runner._l1_activation_global_scale = l1_igs.mean().item()
         self._cutedsl_runner._l2_activation_global_scale = l2_igs.mean().item()
 
@@ -544,6 +545,41 @@ class DeepseekV4MegaMoEExperts(nn.Module):
         self.w2_weight_scale_2 = None
         self.w2_input_scale = None
 
+        # Warmup: compute actual activation global scales from sample data.
+        # The checkpoint input_scale is a calibration value that doesn't match
+        # runtime activation magnitudes. We run a small forward pass to observe
+        # the actual amax and compute correct gs values.
+        self._warmup_activation_global_scales()
+
+    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.
+        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).
+        """
+        import torch
+        runner = self._cutedsl_runner
+        device = runner.device
+        num_tokens = min(8, runner.max_num_tokens)
+        top_k = runner.top_k
+        
+        with torch.no_grad():
+            # 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
+            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_weights = torch.ones(num_tokens, top_k, dtype=torch.float32, device=device) / top_k
+            
+            runner.compute_activation_global_scales(
+                hidden_states, topk_weights, topk_ids
+            )
+
         # Note: No explicit CuTeDSL warmup here. With FULL_AND_PIECEWISE
         # CUDA graph mode, the kernel compiles during graph capture (startup).
         # In eager mode, the first inference triggers JIT compilation.