From b497b35a103a28796b249275896818a7cbb75bc4 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Sun, 17 May 2026 06:06:10 +0000 Subject: [PATCH] fix: dynamic activation quantization (quantize_to_nvfp4) + per-expert scale assembly The checkpoint input_scale was a calibration value that doesn't match runtime input magnitudes. This caused all block scales to saturate at float8_e4m3fn max (448.0), producing garbage output. Fix: use quantize_to_nvfp4 (computes global_scale from input amax) and assemble_scales_2d_side (per-expert split, matches working layertest). This breaks cudagraph (has .max() and .tolist() in forward). Will re-enable cudagraph later with a warmup-based caching approach. --- vllm/nvfp4_cutedsl.py | 175 ++++++++++-------------------------- vllm/patches/deepseek_v4.py | 20 +---- 2 files changed, 53 insertions(+), 142 deletions(-) diff --git a/vllm/nvfp4_cutedsl.py b/vllm/nvfp4_cutedsl.py index c9542e3a..12f3ee9c 100644 --- a/vllm/nvfp4_cutedsl.py +++ b/vllm/nvfp4_cutedsl.py @@ -1,38 +1,37 @@ """ vLLM integration for the CuTeDSL NVFP4 MoE kernel. -CUDA-graph-compatible design: -- All intermediate buffers pre-allocated at max_num_tokens * top_k size -- No .item(), .tolist(), .cpu() — zero CPU-GPU syncs -- No dynamic slicing with GPU scalars — always operate on full pre-allocated buffers -- Extra slots (beyond real tokens) are zero and contribute nothing to output -- Fixed-shape tensors throughout the forward pass +Eager-mode design (cudagraph disabled for now): +- Uses quantize_to_nvfp4 for dynamic activation global_scale +- Uses assemble_scales_2d_side for correct per-expert scale assembly +- Expert offsets computed on GPU, converted to Python list for scale assembly +- .item() calls in scale assembly are acceptable in eager mode -vLLM cudagraph captures at fixed token budgets (1,2,4,8,...,8192). -During capture, num_tokens equals the budget — all shapes are fixed. -During replay, inputs are padded to the budget size. Our runner always -processes max_slots = budget * top_k rows; padding rows are zeros. +The previous cudagraph-safe design used quantize_activation_nvfp4 with +a fixed global_scale from checkpoint input_scale. This caused garbage +output because the calibration input_scale doesn't match runtime input +magnitudes, leading to saturated block scales (all 448.0). + +Future: re-enable cudagraph by computing global_scale during warmup +and caching it, or by using a custom CUDA kernel for dynamic quantization. """ import torch from cutedsl.bridge import ( - quantize_activation_nvfp4, + quantize_to_nvfp4, quantize_weight_to_nvfp4, make_b_k_major, + assemble_scales_2d_side, assemble_scales_3d_side, + compute_expert_offsets, run_nvfp4_grouped_gemm, ) -from cutedsl.kernel.moe.torch_scaled_grouped_mm import ( - ceil_div as cutedsl_ceil_div, - pad_and_swizzle_single, -) class CuTeDSLMoERunner: """Manages NVFP4 MoE execution via the CuTeDSL kernel. - CUDA-graph-compatible: all buffers pre-allocated, no CPU-GPU syncs, - no dynamic shapes. Always computes at max_num_tokens * top_k capacity. + Eager-mode: uses dynamic activation quantization for correct output. """ def __init__(self, num_experts, hidden_size, intermediate_size, @@ -60,26 +59,14 @@ class CuTeDSLMoERunner: self._l1_gsb = None self._l2_gsb = None - self._l1_activation_global_scale = None # set from checkpoint input_scale - self._l2_activation_global_scale = None - - # Pre-allocated cudagraph buffers (set in _allocate_buffers) + # Pre-allocated buffers self._token_indices = None self._expert_id_range = None self._expert_offsets_buf = None - self._padded_scales_buf = None - self._padded_expert_offsets_buf = None self._buffers_allocated = False def _allocate_buffers(self): - """Pre-allocate all buffers at max size for cudagraph compatibility.""" - max_slots = self.max_num_tokens * self.top_k - K_sf = cutedsl_ceil_div(self.hidden_size, 16) - padded_cols = cutedsl_ceil_div(K_sf, 4) * 4 - # Worst case: 1 token per expert, each padded to 128 rows - max_padded_rows = self.num_experts * 128 - - # Slot -> token mapping: [0,0,...,0, 1,1,...,1, ...] (top_k repeats) + """Pre-allocate routing buffers.""" self._token_indices = torch.arange( self.max_num_tokens, device=self.device ).unsqueeze(1).expand(-1, self.top_k).reshape(-1) @@ -89,12 +76,6 @@ class CuTeDSLMoERunner: self._expert_offsets_buf = torch.zeros( self.num_experts + 1, dtype=torch.int32, device=self.device ) - self._padded_expert_offsets_buf = torch.zeros( - self.num_experts + 1, dtype=torch.int32, device=self.device - ) - self._padded_scales_buf = torch.zeros( - max_padded_rows, padded_cols, dtype=torch.float16, device=self.device - ).to(torch.float8_e4m3fn) self._buffers_allocated = True @@ -140,71 +121,15 @@ class CuTeDSLMoERunner: self.l2_gs.append(w_gs) self._l1_mat_b = None - def _assemble_scales_cudagraph_safe(self, x_sf, expert_offsets): - """Assemble 2D-side activation scales (cudagraph-safe, no CPU sync). - - Uses GPU-computed indices to scatter scale data into padded positions, - then applies the swizzle. Returns 2D tensor. - No .item(), no .tolist(), no Python control flow on GPU data. - """ - num_experts = self.num_experts - K_sf = x_sf.shape[1] - padded_cols = cutedsl_ceil_div(K_sf, 4) * 4 - - # Compute tokens per expert (GPU) - tokens_per_expert = expert_offsets[1:] - expert_offsets[:-1] - - # Compute padded rows per expert (round up to 128) - padded_rows_per_expert = ((tokens_per_expert + 127) // 128) * 128 - - # Compute padded offsets - padded_expert_offsets = self._padded_expert_offsets_buf - padded_expert_offsets.zero_() - padded_expert_offsets[1:] = padded_rows_per_expert.cumsum(0) - - # Use the FULL pre-allocated scales buffer (no GPU scalar slicing) - padded_scales = self._padded_scales_buf - padded_scales.zero_() - - # Build index mapping: for each row in x_sf, which expert does it belong to? - total_rows = x_sf.shape[0] - row_indices = self._token_indices[:total_rows] - expert_assign = torch.searchsorted( - expert_offsets[1:], row_indices, right=False - ).clamp(max=num_experts - 1) - - # Destination row in padded buffer - local_row = row_indices - expert_offsets[expert_assign] - dst_rows = padded_expert_offsets[expert_assign] + local_row - - # Scatter x_sf into padded_scales - padded_scales[dst_rows, :K_sf] = x_sf - - # Apply swizzle, reshape to 2D (element count preserved by swizzle) - swizzled = pad_and_swizzle_single(padded_scales) - return swizzled.reshape(padded_scales.shape[0], -1) - def run(self, hidden_states, topk_weights, topk_ids, expert_indices=None): """Run the NVFP4 MoE forward pass. - Fully cudagraph-safe: no CPU-GPU syncs, no dynamic shapes. + Uses dynamic activation quantization (quantize_to_nvfp4) for correct + global_scale computation. This means .max() is called during forward, + which is a CPU-GPU sync — acceptable in eager mode but not cudagraph. - expert_offsets are computed from the actual token distribution - via GPU-only ops (argsort, broadcast ==, cumsum). These offsets - are passed to the GEMM as a GPU tensor, never converted to Python. - - The GEMM and quantize functions see the full slot buffer. - Padding rows are zeros that produce zero output, contributing - nothing to the final scatter_add. - - Args: - hidden_states: (num_tokens, hidden_size) bf16 - topk_weights: (num_tokens, top_k) float32 - topk_ids: (num_tokens, top_k) int - expert_indices: ignored (uses all experts) - - Returns: - (num_tokens, hidden_size) bf16 - MoE output + The scale assembly splits activation scales per-expert (same as the + working moe_pipeline), ensuring correct swizzle layout. """ num_tokens = hidden_states.shape[0] top_k = topk_ids.shape[1] @@ -212,7 +137,7 @@ class CuTeDSLMoERunner: self._ensure_stacked() - # -- Build slot mapping -- + # -- Build slot mapping (GPU) -- flat_ids = topk_ids.reshape(-1) flat_weights = topk_weights.reshape(-1) num_slots = num_tokens * top_k @@ -223,33 +148,32 @@ class CuTeDSLMoERunner: sorted_weights = flat_weights[sort_idx] sorted_token_ids = token_indices[sort_idx] - # Expert offsets (GPU-only, never touches CPU) + # Expert offsets 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) + tokens_per_expert_gpu = (sorted_ids.unsqueeze(1) == expert_id_range.unsqueeze(0)).sum(dim=0).int() + tokens_per_expert = tokens_per_expert_gpu.tolist() + expert_offsets = compute_expert_offsets(tokens_per_expert, self.num_experts) # -- Gather hidden states into slot order -- slot_hidden = hidden_states[sorted_token_ids] # === L1: gate + up === - x_fp4, x_sf = quantize_activation_nvfp4( - slot_hidden, self._l1_activation_global_scale - ) + x_fp4, x_sf, x_igs = quantize_to_nvfp4(slot_hidden) - l1_scale_a = self._assemble_scales_cudagraph_safe( - x_sf, expert_offsets[:self.num_experts + 1] - ) - l1_gsa = torch.full( - (self.num_experts,), self._l1_activation_global_scale, - dtype=torch.float32, device=device - ) + # Split scales per-expert for correct swizzle layout + x_sf_parts = [] + offset = 0 + for tpe in tokens_per_expert: + x_sf_parts.append(x_sf[offset:offset + tpe]) + offset += tpe + + l1_scale_a = assemble_scales_2d_side(x_sf_parts) + l1_gsa = torch.tensor([x_igs] * self.num_experts, 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[:self.num_experts + 1], + expert_offsets=expert_offsets, global_scale_a=l1_gsa, global_scale_b=self._l1_gsb, ) @@ -259,22 +183,21 @@ class CuTeDSLMoERunner: activated = torch.nn.functional.silu(gate) * up # === L2: down === - l2_x_fp4, l2_x_sf = quantize_activation_nvfp4( - activated, self._l2_activation_global_scale - ) + l2_x_fp4, l2_x_sf, l2_igs = quantize_to_nvfp4(activated) - l2_scale_a = self._assemble_scales_cudagraph_safe( - l2_x_sf, expert_offsets[:self.num_experts + 1] - ) - l2_gsa = torch.full( - (self.num_experts,), self._l2_activation_global_scale, - dtype=torch.float32, device=device - ) + l2_sf_parts = [] + offset = 0 + for tpe in tokens_per_expert: + l2_sf_parts.append(l2_x_sf[offset:offset + tpe]) + offset += tpe + + l2_scale_a = assemble_scales_2d_side(l2_sf_parts) + l2_gsa = torch.tensor([l2_igs] * self.num_experts, dtype=torch.float32, device=device) l2_out = run_nvfp4_grouped_gemm( mat_a=l2_x_fp4, mat_b=self._l2_mat_b, scale_a=l2_scale_a, scale_b=self._l2_scale_b, - expert_offsets=expert_offsets[:self.num_experts + 1], + expert_offsets=expert_offsets, global_scale_a=l2_gsa, global_scale_b=self._l2_gsb, ) diff --git a/vllm/patches/deepseek_v4.py b/vllm/patches/deepseek_v4.py index 61625034..0532a199 100644 --- a/vllm/patches/deepseek_v4.py +++ b/vllm/patches/deepseek_v4.py @@ -512,22 +512,10 @@ class DeepseekV4MegaMoEExperts(nn.Module): self._cutedsl_runner.l2_sf = l2_sf self._cutedsl_runner.l2_gs = l2_gs - # Set activation global scales from checkpoint input_scale - # The input_scale is the pre-computed activation normalization factor. - # w13_input_scale shape: (num_experts, 2) for gate+up, but may be (num_experts,) after EP split - # w2_input_scale shape: (num_experts, 1) or (num_experts,) - w13_igs = self.w13_input_scale.data - w2_igs = self.w2_input_scale.data - if w13_igs.dim() == 2: - l1_igs = w13_igs[:, 0] # gate input_scale - else: - l1_igs = w13_igs # already 1D per expert - if w2_igs.dim() == 2: - l2_igs = w2_igs[:, 0] - else: - l2_igs = w2_igs - self._cutedsl_runner._l1_activation_global_scale = l1_igs.mean().item() - self._cutedsl_runner._l2_activation_global_scale = l2_igs.mean().item() + # Activation global scales are now computed dynamically in the runner + # (quantize_to_nvfp4 computes global_scale from actual input magnitude). + # The checkpoint input_scale is a calibration value that doesn't match + # runtime input magnitudes, causing saturated block scales. # Drop the original loader-side parameters self._w13_input_scale = self.w13_input_scale.data.clone()