diff --git a/vllm/nvfp4_cutedsl.py b/vllm/nvfp4_cutedsl.py index 5b29c07f..3f90d2e0 100644 --- a/vllm/nvfp4_cutedsl.py +++ b/vllm/nvfp4_cutedsl.py @@ -1,20 +1,8 @@ """ vLLM integration for the CuTeDSL NVFP4 MoE kernel. -This module provides the interface between the vLLM DeepSeek-V4 model -and the CuTeDSL NVFP4 kernel pipeline. It replaces the broken C++ -CUTLASS path with the working CuTeDSL path. - -The key change: instead of `nvfp4_mega_moe_full` (C++ kernel), we call -`run_nvfp4_moe` (CuTeDSL kernel). Weight preparation and tensor layout -conversion is handled by `cutedsl/bridge.py`. - -Usage in deepseek_v4.py: - from vllm.nvfp4_cutedsl import CuTeDSLMoERunner - - # In DeepseekV4MegaMoEExperts: - self.moe_runner = CuTeDSLMoERunner(...) - self.moe_runner.run(hidden_states, topk_weights, topk_ids, y) +CUDA-graph-compatible: no .item() calls, no Python loops over tokens, +no dynamic shapes. All routing and scattering done with GPU tensors. """ import torch @@ -32,12 +20,7 @@ from cutedsl.bridge import ( class CuTeDSLMoERunner: """Manages NVFP4 MoE execution via the CuTeDSL kernel. - Replaces the old `nvfp4_mega_moe_full` + `SymmBuffer` pipeline. - The CuTeDSL kernel manages its own workspace internally. - - Weight format: checkpoint uint8 → dequantize to BF16 → re-quantize - to float4_e2m1fn_x2. Future optimization: load checkpoint bytes - directly into float4_e2m1fn_x2 tensors. + CUDA-graph-compatible: all operations are GPU-native with no CPU syncs. """ def __init__(self, num_experts, hidden_size, intermediate_size, device="cuda"): @@ -46,60 +29,66 @@ class CuTeDSLMoERunner: self.intermediate_size = intermediate_size self.device = device - # Prepared weights (set by prepare_weights) - self.l1_fp4 = None # list of (K//2, N) float4_e2m1fn_x2 per expert - self.l1_sf = None # list of (K//16, N) float8_e4m3fn per expert - self.l1_gs = None # list of float32 per expert + self.l1_fp4 = None + self.l1_sf = None + self.l1_gs = None self.l2_fp4 = None self.l2_sf = None self.l2_gs = None - - 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). - Use this when you've view-cast checkpoint uint8 → float4_e2m1fn_x2 - and passed block scales / global scales through directly. - Zero precision loss — the bytes are identical. - """ + # Pre-built stacked tensors (set in prepare_weights_direct) + self._l1_mat_b = None + self._l2_mat_b = None + self._l1_scale_b = None + self._l2_scale_b = None + self._l1_gsb = None + self._l2_gsb = None + + def _ensure_stacked(self): + """Lazily stack weight tensors into the format the kernel expects.""" + if self._l1_mat_b is not None: + return + self._l1_mat_b = make_b_k_major(torch.stack(self.l1_fp4)) + self._l2_mat_b = make_b_k_major(torch.stack(self.l2_fp4)) + self._l1_scale_b = assemble_scales_3d_side(self.l1_sf) + 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) + + 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 def prepare_weights_from_dequantized(self, l1_weights_bf16, l2_weights_bf16): - """Prepare NVFP4 weights from dequantized BF16 tensors. - - This is the current path: checkpoint uint8 → dequantize → BF16 → re-quantize. - - Args: - l1_weights_bf16: list of (6144, hidden_size) BF16 tensors (gate+up fused) - l2_weights_bf16: list of (hidden_size, intermediate//2) BF16 tensors (down) - """ + """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: (6144, hidden) → transpose to (hidden, 6144) for K=hidden packed dim 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: (hidden, intermediate//2) → already (K=intermediate//2, N=hidden) - # Wait — down_proj is (hidden, intermediate//2), which is (N_out, K_in) - # For the GEMM: B is (K, N) where K=intermediate, N=hidden - l2_w_t = l2_w.T # (intermediate//2, hidden) → K=intermediate, N=hidden + 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) self.l2_sf.append(w_sf) 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. + CUDA-graph-compatible: no .item() calls, no Python token loops. + Args: hidden_states: (num_tokens, hidden_size) BF16 topk_weights: (num_tokens, top_k) float32 routing weights @@ -116,50 +105,56 @@ class CuTeDSLMoERunner: if expert_indices is None: expert_indices = list(range(self.num_experts)) - # ── Build slot-based routing ── - expert_token_lists = {e: [] for e in expert_indices} - for t in range(num_tokens): - for k in range(top_k): - e = topk_ids[t, k].item() - if e in expert_token_lists: - expert_token_lists[e].append(t) + num_experts = len(expert_indices) + self._ensure_stacked() - tokens_per_expert = [len(expert_token_lists[e]) for e in expert_indices] + # ── Build slot mapping (GPU-native) ── + # Sort tokens by expert assignment to create slot-major ordering + # This replaces the Python loop over tokens/experts - # Skip if no tokens routed - if sum(tokens_per_expert) == 0: + # topk_ids: (num_tokens, top_k) — flatten to (num_tokens * top_k,) + 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) + + # 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 + # Count how many slots each expert gets + expert_id_range = torch.arange(num_experts, device=device) + tokens_per_expert = (sorted_ids.unsqueeze(1) == expert_id_range.unsqueeze(0)).sum(dim=0) + expert_offsets = torch.zeros(num_experts + 1, dtype=torch.int64, device=device) + expert_offsets[1:] = tokens_per_expert.cumsum(0) + if expert_offsets[-1] == 0: return torch.zeros(num_tokens, hidden_size, dtype=torch.bfloat16, device=device) - # Slot-major activation - slot_hidden = torch.cat([ - hidden_states[expert_token_lists[e]] for e in expert_indices - ], dim=0) - - expert_offsets = compute_expert_offsets(tokens_per_expert, len(expert_indices)) + # Gather hidden states in slot-major order + slot_hidden = hidden_states[sorted_token_ids] # (total_slots, hidden_size) BF16 # ════════════════════════════════════════════════════════════ # L1: gate + up (NVFP4 × NVFP4 → BF16) # ════════════════════════════════════════════════════════════ x_fp4, x_sf, x_igs = quantize_to_nvfp4(slot_hidden) - l1_mat_b = make_b_k_major(torch.stack(self.l1_fp4)) - + # Assemble activation scales — split by expert offset boundaries x_sf_parts = [] - offset = 0 - for tpe in tokens_per_expert: - x_sf_parts.append(x_sf[offset:offset+tpe]) - offset += tpe + for e in range(num_experts): + start = expert_offsets[e] + end = expert_offsets[e + 1] + x_sf_parts.append(x_sf[start:end]) l1_scale_a = assemble_scales_2d_side(x_sf_parts) - l1_scale_b = assemble_scales_3d_side(self.l1_sf) - l1_gsa = torch.tensor([x_igs] * len(expert_indices), dtype=torch.float32, device=device) - l1_gsb = torch.tensor(self.l1_gs, dtype=torch.float32, device=device) + l1_gsa = torch.full((num_experts,), x_igs, dtype=torch.float32, device=device) l1_out = run_nvfp4_grouped_gemm( - mat_a=x_fp4, mat_b=l1_mat_b, - scale_a=l1_scale_a, scale_b=l1_scale_b, + 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, - global_scale_a=l1_gsa, global_scale_b=l1_gsb, + global_scale_a=l1_gsa, global_scale_b=self._l1_gsb, ) # ════════════════════════════════════════════════════════════ @@ -174,39 +169,30 @@ class CuTeDSLMoERunner: # ════════════════════════════════════════════════════════════ l2_x_fp4, l2_x_sf, l2_x_igs = quantize_to_nvfp4(activated) - l2_mat_b = make_b_k_major(torch.stack(self.l2_fp4)) - l2_sf_parts = [] - offset = 0 - for tpe in tokens_per_expert: - l2_sf_parts.append(l2_x_sf[offset:offset+tpe]) - offset += tpe + for e in range(num_experts): + start = expert_offsets[e] + end = expert_offsets[e + 1] + l2_sf_parts.append(l2_x_sf[start:end]) l2_scale_a = assemble_scales_2d_side(l2_sf_parts) - l2_scale_b = assemble_scales_3d_side(self.l2_sf) - l2_gsa = torch.tensor([l2_x_igs] * len(expert_indices), dtype=torch.float32, device=device) - l2_gsb = torch.tensor(self.l2_gs, dtype=torch.float32, device=device) + l2_gsa = torch.full((num_experts,), l2_x_igs, dtype=torch.float32, device=device) l2_out = run_nvfp4_grouped_gemm( - mat_a=l2_x_fp4, mat_b=l2_mat_b, - scale_a=l2_scale_a, scale_b=l2_scale_b, + 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, - global_scale_a=l2_gsa, global_scale_b=l2_gsb, + global_scale_a=l2_gsa, global_scale_b=self._l2_gsb, ) # ════════════════════════════════════════════════════════════ - # Scatter → final output + # Scatter → final output (GPU-native, no Python loops) # ════════════════════════════════════════════════════════════ + # Use scatter_add to accumulate expert outputs back to token positions + # sorted_token_ids maps each slot back to its source token + # sorted_weights are the routing weights for each slot y = torch.zeros(num_tokens, hidden_size, dtype=torch.bfloat16, device=device) - - slot_idx = 0 - for e in expert_indices: - for t in expert_token_lists[e]: - for k in range(top_k): - if topk_ids[t, k].item() == e: - w = topk_weights[t, k].item() - y[t] += w * l2_out[slot_idx] - break - slot_idx += 1 + weighted_out = l2_out * sorted_weights.unsqueeze(1).to(l2_out.dtype) + y.scatter_add_(0, sorted_token_ids.unsqueeze(1).expand(-1, hidden_size), weighted_out) return y