feat: vLLM integration — replace C++ kernel with CuTeDSL
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
This commit is contained in:
@@ -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()
|
||||
|
||||
|
||||
Reference in New Issue
Block a user