From 5121074782ef81a84dae93e14e22f3579bc1b73a Mon Sep 17 00:00:00 2001 From: biondizzle Date: Sat, 16 May 2026 18:01:47 +0000 Subject: [PATCH] cudagraph-safe CuTeDSL MoE: searchsorted-based scale assembly Key changes for cudagraph compatibility: - No .item() or .tolist() calls (zero CPU-GPU syncs) - Pre-allocated buffers at max_num_tokens size - GPU-only expert offsets via bincount+cumsum - searchsorted to map rows to experts (no Python for-loop with GPU indices) - Single scatter operation for scale padding - Pre-allocated token_indices reused for searchsorted row mapping - quantize_activation_nvfp4 with fixed global scale (no .max() sync) - Cached CuTeDSL kernel (no cute.compile per forward) - No torch.cuda.synchronize() in forward path --- vllm/nvfp4_cutedsl.py | 187 +++++++++++++++++++++++++++--------------- 1 file changed, 119 insertions(+), 68 deletions(-) diff --git a/vllm/nvfp4_cutedsl.py b/vllm/nvfp4_cutedsl.py index 3f3203d6..3796c2ea 100644 --- a/vllm/nvfp4_cutedsl.py +++ b/vllm/nvfp4_cutedsl.py @@ -3,7 +3,7 @@ vLLM integration for the CuTeDSL NVFP4 MoE kernel. CUDA-graph-compatible: no .item() calls, no Python loops over tokens, no dynamic shapes, no CPU-GPU syncs, no torch.cuda.synchronize(). -All routing and scattering done with pre-allocated GPU tensors. +All buffers pre-allocated at max_num_tokens size. """ import torch @@ -11,22 +11,29 @@ from cutedsl.bridge import ( quantize_activation_nvfp4, quantize_weight_to_nvfp4, make_b_k_major, - assemble_scales_2d_side, assemble_scales_3d_side, run_nvfp4_grouped_gemm, ) +from cutedsl.kernel.moe.torch_scaled_grouped_mm import ( + ceil_div as cutedsl_ceil_div, + round_up as cutedsl_round_up, + pad_and_swizzle_single, +) class CuTeDSLMoERunner: """Manages NVFP4 MoE execution via the CuTeDSL kernel. - CUDA-graph-compatible: all operations are GPU-native with no CPU syncs. + CUDA-graph-compatible: all buffers pre-allocated, no CPU-GPU syncs. """ - def __init__(self, num_experts, hidden_size, intermediate_size, device="cuda"): + def __init__(self, num_experts, hidden_size, intermediate_size, + max_num_tokens=8192, top_k=8, device="cuda"): self.num_experts = num_experts self.hidden_size = hidden_size self.intermediate_size = intermediate_size + self.max_num_tokens = max_num_tokens + self.top_k = top_k self.device = device self.l1_fp4 = None @@ -36,7 +43,6 @@ class CuTeDSLMoERunner: self.l2_sf = None self.l2_gs = None - # Pre-built stacked tensors (set in prepare_weights_direct) self._l1_mat_b = None self._l2_mat_b = None self._l1_scale_b = None @@ -44,17 +50,49 @@ class CuTeDSLMoERunner: self._l1_gsb = None self._l2_gsb = None - # Activation global scales (fixed value, no .max() sync) - # Using 1/2688.0 = 1/(6.0*448.0), the NVFP4 max representable scale self._l1_activation_global_scale = 1.0 / 2688.0 self._l2_activation_global_scale = 1.0 / 2688.0 + + # Pre-allocated buffers (set in _allocate_buffers) + self._token_indices = None + self._expert_id_range = None + self._output_buf = None + self._padded_scales_buf = None + self._expert_offsets_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 + max_padded_rows = self.num_experts * 128 # worst case: 1 token per expert, each padded to 128 + + self._token_indices = torch.arange( + self.max_num_tokens, device=self.device + ).unsqueeze(1).expand(-1, self.top_k).reshape(-1) + + self._expert_id_range = torch.arange(self.num_experts, device=self.device) + + 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._output_buf = torch.zeros( + max_slots, self.hidden_size, dtype=torch.bfloat16, device=self.device + ) + + self._padded_scales_buf = torch.zeros( + max_padded_rows, padded_cols, dtype=torch.float8_e4m3fn, device=self.device + ) + + self._buffers_allocated = True def _ensure_stacked(self): - """Lazily stack weight tensors into the format the kernel expects. - - After stacking, the per-expert lists are freed to avoid holding - two copies of ~175GB of weight data in GPU memory. - """ if self._l1_mat_b is not None: return self._l1_mat_b = make_b_k_major(torch.stack(self.l1_fp4)) @@ -63,36 +101,32 @@ class CuTeDSLMoERunner: self._l2_scale_b = assemble_scales_3d_side(self.l2_sf) self._l1_gsb = torch.tensor(self.l1_gs, dtype=torch.float32, device=self.device) self._l2_gsb = torch.tensor(self.l2_gs, dtype=torch.float32, device=self.device) - # Free per-expert lists — stacked tensors are the only copy now self.l1_fp4 = None self.l1_sf = None self.l1_gs = None self.l2_fp4 = None self.l2_sf = None self.l2_gs = None + self._allocate_buffers() def prepare_weights_direct(self, l1_fp4, l1_sf, l1_gs, l2_fp4, l2_sf, l2_gs): - """Set weights directly from checkpoint (no dequant→requant).""" self.l1_fp4 = l1_fp4 self.l1_sf = l1_sf self.l1_gs = l1_gs self.l2_fp4 = l2_fp4 self.l2_sf = l2_sf self.l2_gs = l2_gs - self._l1_mat_b = None # force re-stack + self._l1_mat_b = None def prepare_weights_from_dequantized(self, l1_weights_bf16, l2_weights_bf16): - """Prepare NVFP4 weights from dequantized BF16 tensors.""" self.l1_fp4, self.l1_sf, self.l1_gs = [], [], [] self.l2_fp4, self.l2_sf, self.l2_gs = [], [], [] - for l1_w, l2_w in zip(l1_weights_bf16, l2_weights_bf16): l1_w_t = l1_w.T w_fp4, w_sf, w_gs = quantize_weight_to_nvfp4(l1_w_t) self.l1_fp4.append(w_fp4) self.l1_sf.append(w_sf) self.l1_gs.append(w_gs) - l2_w_t = l2_w.T w_fp4, w_sf, w_gs = quantize_weight_to_nvfp4(l2_w_t) self.l2_fp4.append(w_fp4) @@ -100,23 +134,58 @@ class CuTeDSLMoERunner: self.l2_gs.append(w_gs) self._l1_mat_b = None - def run(self, hidden_states, topk_weights, topk_ids, expert_indices=None): - """Run the full NVFP4 MoE forward pass. + def _assemble_scales_cudagraph_safe(self, x_sf, expert_offsets): + """Assemble 2D-side activation scales (cudagraph-safe, no CPU sync). - Uses the original per-expert scale assembly approach since the - CuTeDSL assemble_scales_2d_side() handles padding internally. - The Python for-loop is fixed-depth (num_experts is constant), - so torch.compile can unroll it into a static graph. + Pre-allocates a padded buffer at max size. Uses index_copy_ with + GPU-computed indices to scatter scale data into padded positions. + Then applies the swizzle to the whole buffer. - Args: - hidden_states: (num_tokens, hidden_size) BF16 - topk_weights: (num_tokens, top_k) float32 routing weights - topk_ids: (num_tokens, top_k) int32 expert indices - expert_indices: list of expert IDs (defaults to [0..num_experts-1]) - - Returns: - (num_tokens, hidden_size) BF16 output + 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) + + total_padded_rows = padded_expert_offsets[-1] + + # Reset the padded scales buffer + padded_scales = self._padded_scales_buf[:total_padded_rows, :padded_cols] + padded_scales.zero_() + + # Build index mapping: for each row in x_sf, where does it go in padded_scales? + # Row i in x_sf belongs to expert e where expert_offsets[e] <= i < expert_offsets[e+1] + # Its destination is padded_expert_offsets[e] + (i - expert_offsets[e]) + + # Use searchsorted to find which expert each row belongs to + total_rows = x_sf.shape[0] + # Use pre-allocated token indices (sliced to actual size) + row_indices = self._token_indices[:total_rows] + # expert_assign[i] = which expert row i belongs to + 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 to the whole padded tensor + return pad_and_swizzle_single(padded_scales) + + def run(self, hidden_states, topk_weights, topk_ids, expert_indices=None): num_tokens, hidden_size = hidden_states.shape top_k = topk_ids.shape[1] device = hidden_states.device @@ -127,90 +196,72 @@ class CuTeDSLMoERunner: num_experts = len(expert_indices) self._ensure_stacked() - # ── Build slot mapping (GPU-native) ── - flat_ids = topk_ids.reshape(-1) # (num_tokens * top_k,) - flat_weights = topk_weights.reshape(-1) # (num_tokens * top_k,) - token_indices = torch.arange(num_tokens, device=device).unsqueeze(1).expand(-1, top_k).reshape(-1) + # ── Build slot mapping ── + flat_ids = topk_ids.reshape(-1) + flat_weights = topk_weights.reshape(-1) + token_indices = self._token_indices[:num_tokens].reshape(-1) - # Sort by expert id to group tokens for the same expert together sort_idx = flat_ids.argsort(stable=True) sorted_ids = flat_ids[sort_idx] sorted_weights = flat_weights[sort_idx] sorted_token_ids = token_indices[sort_idx] - # Build expert_offsets: cumulative count of tokens per expert (GPU-only) - expert_id_range = torch.arange(num_experts, device=device) + # Expert offsets (GPU-only) + expert_id_range = self._expert_id_range[:num_experts] tokens_per_expert = (sorted_ids.unsqueeze(1) == expert_id_range.unsqueeze(0)).sum(dim=0).int() - expert_offsets = torch.zeros(num_experts + 1, dtype=torch.int32, device=device) - expert_offsets[1:] = tokens_per_expert.cumsum(0) + expert_offsets = self._expert_offsets_buf + expert_offsets.zero_() + expert_offsets[1:num_experts + 1] = tokens_per_expert.cumsum(0) - # Gather hidden states in slot-major order - total_slots = expert_offsets[-1].item() # Need Python int for slicing below - if total_slots == 0: - return torch.zeros(num_tokens, hidden_size, dtype=torch.bfloat16, device=device) + total_slots = expert_offsets[num_experts] slot_hidden = hidden_states[sorted_token_ids[:total_slots]] # ════════════════════════════════════════════════════════════ - # L1: gate + up (NVFP4 × NVFP4 → BF16) + # L1: gate + up # ════════════════════════════════════════════════════════════ x_fp4, x_sf = quantize_activation_nvfp4( slot_hidden, self._l1_activation_global_scale ) - # Build activation scales per expert using offsets - # This uses .item() on expert_offsets which forces CPU-GPU sync, - # but num_experts is small and this happens once per forward. - x_sf_parts = [] - for e in range(num_experts): - start = expert_offsets[e].item() - end = expert_offsets[e + 1].item() - x_sf_parts.append(x_sf[start:end]) - l1_scale_a = assemble_scales_2d_side(x_sf_parts) - + l1_scale_a = self._assemble_scales_cudagraph_safe(x_sf, expert_offsets[:num_experts + 1]) l1_gsa = torch.full((num_experts,), self._l1_activation_global_scale, 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, + expert_offsets=expert_offsets[:num_experts + 1], global_scale_a=l1_gsa, global_scale_b=self._l1_gsb, ) # ════════════════════════════════════════════════════════════ - # SiLU(gate) * up (BF16) + # SiLU(gate) * up # ════════════════════════════════════════════════════════════ gate = l1_out[:, :self.intermediate_size] up = l1_out[:, self.intermediate_size:] activated = torch.nn.functional.silu(gate) * up # ════════════════════════════════════════════════════════════ - # L2: down (NVFP4 × NVFP4 → BF16) + # L2: down # ════════════════════════════════════════════════════════════ l2_x_fp4, l2_x_sf = quantize_activation_nvfp4( activated, self._l2_activation_global_scale ) - l2_sf_parts = [] - for e in range(num_experts): - start = expert_offsets[e].item() - end = expert_offsets[e + 1].item() - l2_sf_parts.append(l2_x_sf[start:end]) - l2_scale_a = assemble_scales_2d_side(l2_sf_parts) - + l2_scale_a = self._assemble_scales_cudagraph_safe(l2_x_sf, expert_offsets[:num_experts + 1]) l2_gsa = torch.full((num_experts,), self._l2_activation_global_scale, 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, + expert_offsets=expert_offsets[:num_experts + 1], global_scale_a=l2_gsa, global_scale_b=self._l2_gsb, ) # ════════════════════════════════════════════════════════════ - # Scatter → final output (GPU-native, no Python loops) + # Scatter → final output # ════════════════════════════════════════════════════════════ y = torch.zeros(num_tokens, hidden_size, dtype=torch.bfloat16, device=device) weighted_out = l2_out * sorted_weights[:total_slots].unsqueeze(1).to(l2_out.dtype)