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

fix: dynamic activation quantization (quantize_to_nvfp4) + per-expert scale assembly

Browse Source 
The checkpoint input_scale was a calibration value that doesn't match
runtime input magnitudes. This caused all block scales to saturate at
float8_e4m3fn max (448.0), producing garbage output.

Fix: use quantize_to_nvfp4 (computes global_scale from input amax) and
assemble_scales_2d_side (per-expert split, matches working layertest).

This breaks cudagraph (has .max() and .tolist() in forward). Will
re-enable cudagraph later with a warmup-based caching approach.
This commit is contained in:
biondizzle
2026-05-17 06:06:10 +00:00
parent 78bebff736
commit b497b35a10
2 changed files with 53 additions and 142 deletions
Download Patch File Download Diff File Expand all files Collapse all files

175
vllm/nvfp4_cutedsl.py
View File

@@ -1,38 +1,37 @@
"""
vLLM integration for the CuTeDSL NVFP4 MoE kernel.
CUDA-graph-compatible design:
- All intermediate buffers pre-allocated at max_num_tokens * top_k size
- No .item(), .tolist(), .cpu() — zero CPU-GPU syncs
- No dynamic slicing with GPU scalars — always operate on full pre-allocated buffers
- Extra slots (beyond real tokens) are zero and contribute nothing to output
- Fixed-shape tensors throughout the forward pass
Eager-mode design (cudagraph disabled for now):
- Uses quantize_to_nvfp4 for dynamic activation global_scale
- Uses assemble_scales_2d_side for correct per-expert scale assembly
- Expert offsets computed on GPU, converted to Python list for scale assembly
- .item() calls in scale assembly are acceptable in eager mode
vLLM cudagraph captures at fixed token budgets (1,2,4,8,...,8192).
During capture, num_tokens equals the budget — all shapes are fixed.
During replay, inputs are padded to the budget size. Our runner always
processes max_slots = budget * top_k rows; padding rows are zeros.
The previous cudagraph-safe design used quantize_activation_nvfp4 with
a fixed global_scale from checkpoint input_scale. This caused garbage
output because the calibration input_scale doesn't match runtime input
magnitudes, leading to saturated block scales (all 448.0).
Future: re-enable cudagraph by computing global_scale during warmup
and caching it, or by using a custom CUDA kernel for dynamic quantization.
"""
import torch
from cutedsl.bridge import (
quantize_activation_nvfp4,
quantize_to_nvfp4,
quantize_weight_to_nvfp4,
make_b_k_major,
assemble_scales_2d_side,
assemble_scales_3d_side,
compute_expert_offsets,
run_nvfp4_grouped_gemm,
)
from cutedsl.kernel.moe.torch_scaled_grouped_mm import (
ceil_div as cutedsl_ceil_div,
pad_and_swizzle_single,
)
class CuTeDSLMoERunner:
"""Manages NVFP4 MoE execution via the CuTeDSL kernel.
CUDA-graph-compatible: all buffers pre-allocated, no CPU-GPU syncs,
no dynamic shapes. Always computes at max_num_tokens * top_k capacity.
Eager-mode: uses dynamic activation quantization for correct output.
"""
def __init__(self, num_experts, hidden_size, intermediate_size,
@@ -60,26 +59,14 @@ class CuTeDSLMoERunner:
self._l1_gsb = None
self._l2_gsb = None
self._l1_activation_global_scale = None # set from checkpoint input_scale
self._l2_activation_global_scale = None
# Pre-allocated cudagraph buffers (set in _allocate_buffers)
# Pre-allocated buffers
self._token_indices = None
self._expert_id_range = None
self._expert_offsets_buf = None
self._padded_scales_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
# Worst case: 1 token per expert, each padded to 128 rows
max_padded_rows = self.num_experts * 128
# Slot -> token mapping: [0,0,...,0, 1,1,...,1, ...] (top_k repeats)
"""Pre-allocate routing buffers."""
self._token_indices = torch.arange(
self.max_num_tokens, device=self.device
).unsqueeze(1).expand(-1, self.top_k).reshape(-1)
@@ -89,12 +76,6 @@ class CuTeDSLMoERunner:
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._padded_scales_buf = torch.zeros(
max_padded_rows, padded_cols, dtype=torch.float16, device=self.device
).to(torch.float8_e4m3fn)
self._buffers_allocated = True
@@ -140,71 +121,15 @@ class CuTeDSLMoERunner:
self.l2_gs.append(w_gs)
self._l1_mat_b = None
def _assemble_scales_cudagraph_safe(self, x_sf, expert_offsets):
"""Assemble 2D-side activation scales (cudagraph-safe, no CPU sync).
Uses GPU-computed indices to scatter scale data into padded positions,
then applies the swizzle. Returns 2D tensor.
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)
# Use the FULL pre-allocated scales buffer (no GPU scalar slicing)
padded_scales = self._padded_scales_buf
padded_scales.zero_()
# Build index mapping: for each row in x_sf, which expert does it belong to?
total_rows = x_sf.shape[0]
row_indices = self._token_indices[:total_rows]
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, reshape to 2D (element count preserved by swizzle)
swizzled = pad_and_swizzle_single(padded_scales)
return swizzled.reshape(padded_scales.shape[0], -1)
def run(self, hidden_states, topk_weights, topk_ids, expert_indices=None):
"""Run the NVFP4 MoE forward pass.
Fully cudagraph-safe: no CPU-GPU syncs, no dynamic shapes.
Uses dynamic activation quantization (quantize_to_nvfp4) for correct
global_scale computation. This means .max() is called during forward,
which is a CPU-GPU sync — acceptable in eager mode but not cudagraph.
expert_offsets are computed from the actual token distribution
via GPU-only ops (argsort, broadcast ==, cumsum). These offsets
are passed to the GEMM as a GPU tensor, never converted to Python.
The GEMM and quantize functions see the full slot buffer.
Padding rows are zeros that produce zero output, contributing
nothing to the final scatter_add.
Args:
hidden_states: (num_tokens, hidden_size) bf16
topk_weights: (num_tokens, top_k) float32
topk_ids: (num_tokens, top_k) int
expert_indices: ignored (uses all experts)
Returns:
(num_tokens, hidden_size) bf16 - MoE output
The scale assembly splits activation scales per-expert (same as the
working moe_pipeline), ensuring correct swizzle layout.
"""
num_tokens = hidden_states.shape[0]
top_k = topk_ids.shape[1]
@@ -212,7 +137,7 @@ class CuTeDSLMoERunner:
self._ensure_stacked()
# -- Build slot mapping --
# -- Build slot mapping (GPU) --
flat_ids = topk_ids.reshape(-1)
flat_weights = topk_weights.reshape(-1)
num_slots = num_tokens * top_k
@@ -223,33 +148,32 @@ class CuTeDSLMoERunner:
sorted_weights = flat_weights[sort_idx]
sorted_token_ids = token_indices[sort_idx]
# Expert offsets (GPU-only, never touches CPU)
# Expert offsets
expert_id_range = self._expert_id_range
tokens_per_expert = (sorted_ids.unsqueeze(1) == expert_id_range.unsqueeze(0)).sum(dim=0).int()
expert_offsets = self._expert_offsets_buf
expert_offsets.zero_()
expert_offsets[1:self.num_experts + 1] = tokens_per_expert.cumsum(0)
tokens_per_expert_gpu = (sorted_ids.unsqueeze(1) == expert_id_range.unsqueeze(0)).sum(dim=0).int()
tokens_per_expert = tokens_per_expert_gpu.tolist()
expert_offsets = compute_expert_offsets(tokens_per_expert, self.num_experts)
# -- Gather hidden states into slot order --
slot_hidden = hidden_states[sorted_token_ids]
# === L1: gate + up ===
x_fp4, x_sf = quantize_activation_nvfp4(
slot_hidden, self._l1_activation_global_scale
)
x_fp4, x_sf, x_igs = quantize_to_nvfp4(slot_hidden)
l1_scale_a = self._assemble_scales_cudagraph_safe(
x_sf, expert_offsets[:self.num_experts + 1]
)
l1_gsa = torch.full(
(self.num_experts,), self._l1_activation_global_scale,
dtype=torch.float32, device=device
)
# Split scales per-expert for correct swizzle layout
x_sf_parts = []
offset = 0
for tpe in tokens_per_expert:
x_sf_parts.append(x_sf[offset:offset + tpe])
offset += tpe
l1_scale_a = assemble_scales_2d_side(x_sf_parts)
l1_gsa = torch.tensor([x_igs] * self.num_experts, 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[:self.num_experts + 1],
expert_offsets=expert_offsets,
global_scale_a=l1_gsa, global_scale_b=self._l1_gsb,
)
@@ -259,22 +183,21 @@ class CuTeDSLMoERunner:
activated = torch.nn.functional.silu(gate) * up
# === L2: down ===
l2_x_fp4, l2_x_sf = quantize_activation_nvfp4(
activated, self._l2_activation_global_scale
)
l2_x_fp4, l2_x_sf, l2_igs = quantize_to_nvfp4(activated)
l2_scale_a = self._assemble_scales_cudagraph_safe(
l2_x_sf, expert_offsets[:self.num_experts + 1]
)
l2_gsa = torch.full(
(self.num_experts,), self._l2_activation_global_scale,
dtype=torch.float32, device=device
)
l2_sf_parts = []
offset = 0
for tpe in tokens_per_expert:
l2_sf_parts.append(l2_x_sf[offset:offset + tpe])
offset += tpe
l2_scale_a = assemble_scales_2d_side(l2_sf_parts)
l2_gsa = torch.tensor([l2_igs] * self.num_experts, 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[:self.num_experts + 1],
expert_offsets=expert_offsets,
global_scale_a=l2_gsa, global_scale_b=self._l2_gsb,
)

20
vllm/patches/deepseek_v4.py
View File

@@ -512,22 +512,10 @@ class DeepseekV4MegaMoEExperts(nn.Module):
self._cutedsl_runner.l2_sf = l2_sf
self._cutedsl_runner.l2_gs = l2_gs
# Set activation global scales from checkpoint input_scale
# The input_scale is the pre-computed activation normalization factor.
# w13_input_scale shape: (num_experts, 2) for gate+up, but may be (num_experts,) after EP split
# w2_input_scale shape: (num_experts, 1) or (num_experts,)
w13_igs = self.w13_input_scale.data
w2_igs = self.w2_input_scale.data
if w13_igs.dim() == 2:
l1_igs = w13_igs[:, 0] # gate input_scale
else:
l1_igs = w13_igs # already 1D per expert
if w2_igs.dim() == 2:
l2_igs = w2_igs[:, 0]
else:
l2_igs = w2_igs
self._cutedsl_runner._l1_activation_global_scale = l1_igs.mean().item()
self._cutedsl_runner._l2_activation_global_scale = l2_igs.mean().item()
# Activation global scales are now computed dynamically in the runner
# (quantize_to_nvfp4 computes global_scale from actual input magnitude).
# The checkpoint input_scale is a calibration value that doesn't match
# runtime input magnitudes, causing saturated block scales.
# Drop the original loader-side parameters
self._w13_input_scale = self.w13_input_scale.data.clone()