biondizzle/nvfp4-megamoe-kernel
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: rewrite CuTeDSLMoERunner for CUDA graph compatibility

Browse Source 
CUDA graphs forbid CPU-GPU syncs (.item()) and Python loops over
tokens during graph capture. The old scatter loop did both.

Changes:
- Slot routing: replaced Python loop with GPU-native argsort + gather
  (sort tokens by expert id, gather hidden states in slot order)
- Scatter: replaced Python loop with torch.scatter_add_ (GPU-native)
- Weight stacking: lazily pre-built once, reused every forward call
- Removed all .item() calls from the forward path
- expert_offsets built from GPU tensor operations

This is required for FULL_AND_PIECEWISE CUDA graph mode which
compiles and captures graphs during startup.
This commit is contained in:
biondizzle
2026-05-16 07:03:08 +00:00
parent dab31b0961
commit 4b0a9557f0
1 changed files with 82 additions and 96 deletions
Download Patch File Download Diff File Expand all files Collapse all files

178
vllm/nvfp4_cutedsl.py
View File

@@ -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