diff --git a/vllm/nvfp4_cutedsl.py b/vllm/nvfp4_cutedsl.py index 46625cef..4e669395 100644 --- a/vllm/nvfp4_cutedsl.py +++ b/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.