feat: compute_activation_global_scales warmup method

Uses quantize_to_nvfp4 during warmup to get exact gs values for L1 and L2.
L1 gs comes from slot_hidden, L2 gs from the actual L1 GEMM output.
These values are then used with quantize_activation_nvfp4 (cudagraph-safe)
during inference.

vllm/nvfp4_cutedsl.py

@@ -208,6 +208,70 @@ class CuTeDSLMoERunner:
         all_flat = all_flat.view(torch.float8_e4m3fn)
         return all_flat.reshape(num_experts * 128, -1)
 
+    def compute_activation_global_scales(self, hidden_states_sample, topk_weights, topk_ids):
+        """Compute activation global scales from a warmup forward pass.
+        
+        Called BEFORE cudagraph capture. Runs quantize_to_nvfp4 (which has .max())
+        to observe actual activation magnitudes, then stores the EXACT gs values
+        for use with quantize_activation_nvfp4 (cudagraph-safe, no .max()).
+        
+        L1 and L2 get SEPARATE gs because the L2 activation (after SiLU * up)
+        has a much larger magnitude than the L1 input.
+        
+        The warmup uses the runner's own scale assembly and GEMM to ensure the
+        L2 gs is computed from the actual L1 output (not an approximation).
+        """
+        self._ensure_stacked()
+        device = hidden_states_sample.device
+        num_tokens = hidden_states_sample.shape[0]
+        top_k = topk_ids.shape[1]
+        
+        with torch.no_grad():
+            # Build slot mapping (same as run())
+            flat_ids = topk_ids.reshape(-1)
+            num_slots = num_tokens * top_k
+            token_indices = self._token_indices[:num_slots]
+            sort_idx = flat_ids.argsort(stable=True)
+            sorted_ids = flat_ids[sort_idx]
+            sorted_token_ids = token_indices[sort_idx]
+            slot_hidden = hidden_states_sample[sorted_token_ids]
+            
+            # L1: get exact gs from quantize_to_nvfp4
+            _, _, l1_gs = quantize_to_nvfp4(slot_hidden)
+            
+            # Run L1 GEMM to get L1 output
+            x_fp4, x_sf = quantize_activation_nvfp4(slot_hidden, l1_gs)
+            expert_id_range = self._expert_id_range
+            tokens_per_expert = (sorted_ids.unsqueeze(1) == expert_id_range.unsqueeze(0)).sum(dim=0).int()
+            expert_offsets = self._expert_offsets_buf
+            expert_offsets.zero_()
+            expert_offsets[1:self.num_experts + 1] = tokens_per_expert.cumsum(0)
+            
+            l1_scale_a = self._assemble_scales_cudagraph_safe(
+                x_sf, expert_offsets[:self.num_experts + 1],
+                self._padded_x_sf_buf_l1, self._per_expert_scale_bufs_l1
+            )
+            l1_gsa = torch.full((self.num_experts,), l1_gs, dtype=torch.float32, device=device)
+            
+            l1_out = run_nvfp4_grouped_gemm(
+                mat_a=x_fp4, mat_b=self._l1_mat_b,
+                scale_a=l1_scale_a, scale_b=self._l1_scale_b,
+                expert_offsets=expert_offsets[1:self.num_experts + 1],
+                global_scale_a=l1_gsa, global_scale_b=self._l1_gsb,
+            )
+            
+            # L2: get exact gs from SiLU(gate)*up
+            gate = l1_out[:, :self.intermediate_size]
+            up = l1_out[:, self.intermediate_size:]
+            activated = torch.nn.functional.silu(gate) * up
+            _, _, l2_gs = quantize_to_nvfp4(activated)
+        
+        self._l1_activation_global_scale = l1_gs
+        self._l2_activation_global_scale = l2_gs
+        
+        print(f"  Warmup L1 gs: {l1_gs:.10f}")
+        print(f"  Warmup L2 gs: {l2_gs:.10f}")
+
     def run(self, hidden_states, topk_weights, topk_ids, expert_indices=None):
         """Run the NVFP4 MoE forward pass.