From 8fd9579127a43e5ee7b1234daf47639187d6a8a5 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Sat, 16 May 2026 03:36:12 +0000 Subject: [PATCH] =?UTF-8?q?feat:=20vLLM=20integration=20=E2=80=94=20replac?= =?UTF-8?q?e=20C++=20kernel=20with=20CuTeDSL?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit deepseek_v4.py changes: - finalize_weights(): dequantize checkpoint → BF16 → re-quantize to float4_e2m1fn_x2 via CuTeDSLMoERunner (replaces transform_nvfp4_weights_for_mega_moe) - _run_mega_moe(): calls CuTeDSLMoERunner.run() (replaces nvfp4_mega_moe_full) - Removed get_symm_buffer() and SymmBuffer (CuTeDSL manages its own workspace) - Removed _transformed_l1_weights / _transformed_l2_weights - Added _cutedsl_runner class variable - Weight loader unchanged (checkpoint loading is the same) vllm/nvfp4_cutedsl.py: - CuTeDSLMoERunner class handles the full pipeline - prepare_weights_from_dequantized() for weight prep - run() does L1→SiLU→L2→scatter with NVFP4-native GEMMs --- vllm/patches/deepseek_v4.py | 159 +++++++++++++++++------------------- 1 file changed, 74 insertions(+), 85 deletions(-) diff --git a/vllm/patches/deepseek_v4.py b/vllm/patches/deepseek_v4.py index d965dac0..71e43bf5 100644 --- a/vllm/patches/deepseek_v4.py +++ b/vllm/patches/deepseek_v4.py @@ -215,13 +215,13 @@ class DeepseekV4MegaMoEExperts(nn.Module): """MegaMoE experts for DeepSeek V4 with NVFP4 quantization. Loads NVFP4 expert weights (E2M1 packed uint8 + float8_e4m3fn block scales - + float32 global scales) and feeds them natively to the DeepGEMM - fp8_nvfp4_mega_moe kernel (kind::mxf4nvf4.scale_vec::4X). + + float32 global scales) and runs them through the CuTeDSL NVFP4 kernel. - No conversion to MXFP4. Experts stay NVFP4. The global scale (weight_scale_2) - is folded into the block scales before kernel consumption. + The CuTeDSL kernel is a Python-based CUTLASS kernel compiled via MLIR → PTX. + It handles NVFP4 natively with full Blackwell pipeline overlap (TMA → MMA → Epilogue). + This replaces the broken C++ CUTLASS kernel (see README.md for the full story). """ - _symm_buffer_cache: dict[tuple[int, int, int, int, int, int, int], object] = {} + _cutedsl_runner: 'CuTeDSLMoERunner | None' = None # NVFP4 E2M1 lookup table (positive values, sign from bit 3) E2M1_LUT = [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0] @@ -329,8 +329,7 @@ class DeepseekV4MegaMoEExperts(nn.Module): ) set_weight_attrs(self.w2_input_scale, weight_attrs) - self._transformed_l1_weights: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None - self._transformed_l2_weights: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None + self._cutedsl_runner = None # Register in the static forward context so the custom-op wrapper # can look up this module by name from within a torch.compile graph. @@ -413,34 +412,72 @@ class DeepseekV4MegaMoEExperts(nn.Module): ) def finalize_weights(self) -> None: - if self._transformed_l1_weights is not None: + if self._cutedsl_runner is not None and self._cutedsl_runner.l1_fp4 is not None: return self._check_runtime_supported() - from nvfp4_megamoe_kernel import transform_nvfp4_weights_for_mega_moe - # === Native NVFP4 path === - # The CUTLASS nvfp4 mega_moe kernel consumes NVFP4 directly: - # - E2M1 packed int8 (same as checkpoint) - # - UE4M3 block scales (float8_e4m3fn), group_size=16 - # - float32 global scales returned SEPARATELY (NOT folded into float8) - # Previous versions folded global scales into block scales via float8 - # round-trip, which caused ~25% precision loss. Now, global scales - # are applied as per-expert GEMM alpha in float32 (exact). + # Dequantize checkpoint NVFP4 → BF16, then re-quantize to native + # float4_e2m1fn_x2 for the CuTeDSL kernel. + # Future optimization: load checkpoint bytes directly into + # float4_e2m1fn_x2 without the BF16 round-trip. + from vllm.nvfp4_cutedsl import CuTeDSLMoERunner - # Transform weights — returns (w, sf, global_sf) tuples - self._transformed_l1_weights, self._transformed_l2_weights = ( - transform_nvfp4_weights_for_mega_moe( - (self.w13_weight.data.contiguous(), - self.w13_weight_scale.data.contiguous()), - (self.w2_weight.data.contiguous(), - self.w2_weight_scale.data.contiguous()), - l1_weight_scale_2=self.w13_weight_scale_2.data.contiguous(), - l2_weight_scale_2=self.w2_weight_scale_2.data.contiguous(), + E2M1_LUT = torch.tensor([ + 0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, + -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0, + ], dtype=torch.float32, device=self.w13_weight.device) + + def dequant_nvfp4(packed, scale, global_scale): + raw = packed.view(torch.uint8) + lo = E2M1_LUT[(raw & 0x0F).long()] + hi = E2M1_LUT[((raw >> 4) & 0x0F).long()] + out_features = raw.shape[0] + in_features = raw.shape[1] * 2 + unpacked = torch.empty(out_features, in_features, dtype=torch.float32, device=raw.device) + unpacked[:, 0::2] = lo + unpacked[:, 1::2] = hi + bs = scale.float().repeat_interleave(16, dim=1)[:, :in_features] + return (unpacked * bs * global_scale).to(torch.bfloat16) + + l1_weights_bf16 = [] + l2_weights_bf16 = [] + + for e in range(self.num_local_experts): + # L1: gate + up fused + gate_w = dequant_nvfp4( + self.w13_weight.data[e, :self.intermediate_size], + self.w13_weight_scale.data[e, :self.intermediate_size], + self.w13_weight_scale_2.data[e, 0], ) + up_w = dequant_nvfp4( + self.w13_weight.data[e, self.intermediate_size:], + self.w13_weight_scale.data[e, self.intermediate_size:], + self.w13_weight_scale_2.data[e, 1], + ) + fused = torch.cat([gate_w, up_w], dim=0) # (6144, hidden) + l1_weights_bf16.append(fused.T) # (hidden, 6144) — K=hidden packed dim + + # L2: down + l2_w = dequant_nvfp4( + self.w2_weight.data[e], + self.w2_weight_scale.data[e], + self.w2_weight_scale_2.data[e], + ) + l2_weights_bf16.append(l2_w.T) # (intermediate//2, hidden) + + # Create CuTeDSL runner and prepare weights + self._cutedsl_runner = CuTeDSLMoERunner( + num_experts=self.num_local_experts, + hidden_size=self.hidden_size, + intermediate_size=self.intermediate_size, + device=self.w13_weight.device, + ) + self._cutedsl_runner.prepare_weights_from_dequantized( + l1_weights_bf16, l2_weights_bf16, ) - # Drop the original loader-side parameters (preserve input_scales) + # Drop the original loader-side parameters self._w13_input_scale = self.w13_input_scale.data.clone() self._w2_input_scale = self.w2_input_scale.data.clone() self.w13_weight = None @@ -452,35 +489,6 @@ class DeepseekV4MegaMoEExperts(nn.Module): self.w2_weight_scale_2 = None self.w2_input_scale = None - def get_symm_buffer(self): - import nvfp4_megamoe_kernel as deep_gemm - from nvfp4_megamoe_kernel import SymmBuffer, get_symm_buffer_for_nvfp4_mega_moe - - group = get_ep_group().device_group - device = torch.accelerator.current_device_index() - key = ( - id(group), - device, - self.num_experts, - self.max_num_tokens, - self.top_k, - self.hidden_size, - self.intermediate_size, - ) - symm_buffer = self._symm_buffer_cache.get(key) - if symm_buffer is None: - # NVFP4 SymmBuffer: 2x SF size due to group_size=16 - symm_buffer = get_symm_buffer_for_nvfp4_mega_moe( - group, - self.num_experts, - self.max_num_tokens, - self.top_k, - self.hidden_size, - self.intermediate_size, - ) - self._symm_buffer_cache[key] = symm_buffer - return symm_buffer - def forward( self, hidden_states: torch.Tensor, @@ -517,42 +525,23 @@ class DeepseekV4MegaMoEExperts(nn.Module): fast_math: bool, ) -> None: import os - from nvfp4_megamoe_kernel import stage_activation, nvfp4_mega_moe_full - - symm_buffer = self.get_symm_buffer() - symm_buffer.experts_start_idx = self.experts_start_idx - num_tokens = hidden_states.shape[0] - - # Quantize activation using the kernel's PyTorch stage_activation - # Use the checkpoint's input_scale for L1 (w13) activation quantization. - # The checkpoint's input_scale was used during weight calibration — using - # the same scale at runtime ensures the quantized weights are rescaled correctly. - # Dynamic stage_activation computes amax/(6*448) which can be 10x+ off. - w13_input_scale = float(self._w13_input_scale[0]) # same for all experts - x_fp4, x_sf, input_global_scale = stage_activation(hidden_states, input_global_scale=w13_input_scale) - symm_buffer.x[:num_tokens].copy_(x_fp4) - symm_buffer.x_sf[:num_tokens].copy_(x_sf) - symm_buffer.input_global_scale = input_global_scale - symm_buffer.topk_idx[:num_tokens].copy_(topk_ids) - symm_buffer.topk_weights[:num_tokens].copy_(topk_weights) # This method must have been already called during the weight loading phase. # We call it again here to cover the dummy weight loading case. self.finalize_weights() - assert self._transformed_l1_weights is not None - assert self._transformed_l2_weights is not None + assert self._cutedsl_runner is not None + assert self._cutedsl_runner.l1_fp4 is not None - nvfp4_mega_moe_full( - y, - self._transformed_l1_weights, - self._transformed_l2_weights, - symm_buffer, - activation_clamp=activation_clamp, - fast_math=fast_math, - l1_input_scale=self._w13_input_scale, - l2_input_scale=self._w2_input_scale, + # Build expert indices list for this rank + expert_indices = list(range(self.num_local_experts)) + + result = self._cutedsl_runner.run( + hidden_states, topk_weights, topk_ids, + expert_indices=expert_indices, ) + y.copy_(result) + if os.environ.get('NVFP4_DEBUG_SYNC', '') == '1': torch.cuda.synchronize()