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

cudagraph-safe CuTeDSL MoE: searchsorted-based scale assembly

Browse Source 
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
This commit is contained in:
biondizzle
2026-05-16 18:01:47 +00:00
parent ab126b0c0d
commit 5121074782
1 changed files with 119 additions and 68 deletions
Download Patch File Download Diff File Expand all files Collapse all files

187
vllm/nvfp4_cutedsl.py
View File

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