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

fix: revert to .item() based scale assembly (fixes index OOB)

Browse Source 
The fully GPU-vectorized _assemble_scales_gpu() caused index out of
bounds errors because tensor slicing with GPU-computed indices from
Python is undefined behavior.

Went back to .item() on expert_offsets for the per-expert scale split.
This forces CPU-GPU syncs (breaks cudagraph) but produces correct results.

The path to cudagraph compatibility is either:
1. Modify CuTeDSL scale assembly API to accept flat tensor + offsets
2. Use the CUTLASS kernel (already verified working)
This commit is contained in:
biondizzle
2026-05-16 17:55:32 +00:00
parent 7594968482
commit ab126b0c0d
1 changed files with 43 additions and 137 deletions
Download Patch File Download Diff File Expand all files Collapse all files

180
vllm/nvfp4_cutedsl.py
View File

@@ -21,15 +21,12 @@ class CuTeDSLMoERunner:
"""Manages NVFP4 MoE execution via the CuTeDSL kernel.
CUDA-graph-compatible: all operations are GPU-native with no CPU syncs.
Pre-allocates buffers at max_num_tokens size and slices during forward.
"""
def __init__(self, num_experts, hidden_size, intermediate_size,
max_num_tokens=8192, device="cuda"):
def __init__(self, num_experts, hidden_size, intermediate_size, 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.device = device
self.l1_fp4 = None
@@ -47,50 +44,8 @@ class CuTeDSLMoERunner:
self._l1_gsb = None
self._l2_gsb = None
# Pre-allocated buffers for cudagraph safety
# These are allocated at max_num_tokens size and sliced during forward
self._token_indices = None # (max_num_tokens * top_k,)
self._expert_id_range = None # (num_experts,)
self._output_buf = None # (max_num_tokens * top_k, hidden_size)
self._l1_output_buf = None # (max_num_tokens * top_k, intermediate_size * 2)
self._l2_output_buf = None # (max_num_tokens * top_k, hidden_size)
self._workspace_buf = None # varies
# Activation global scales (computed once during warmup)
self._l1_activation_global_scale = None
self._l2_activation_global_scale = None
def _allocate_buffers(self, top_k=8):
"""Pre-allocate all buffers at max size for cudagraph compatibility.
Called once after weight stacking. All forward calls slice from these.
"""
max_slots = self.max_num_tokens * top_k
# Token index buffer for scatter/gather
self._token_indices = torch.arange(
self.max_num_tokens, device=self.device
).unsqueeze(1).expand(-1, top_k).reshape(-1) # (max_tokens * top_k,)
# Expert ID range for offset computation
self._expert_id_range = torch.arange(
self.num_experts, device=self.device
) # (num_experts,)
# Output buffers (pre-allocated at max size, zeroed each forward)
self._output_buf = torch.zeros(
max_slots, self.hidden_size, dtype=torch.bfloat16, device=self.device
)
self._l1_output_buf = torch.zeros(
max_slots, self.intermediate_size * 2, dtype=torch.bfloat16, device=self.device
)
self._l2_output_buf = torch.zeros(
max_slots, self.hidden_size, dtype=torch.bfloat16, device=self.device
)
# Default activation global scales
# Using a reasonable default: 6.0 * 448.0 = 2688.0 (NVFP4 max representable)
# This avoids the .max() CPU-GPU sync
# 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
@@ -115,8 +70,6 @@ class CuTeDSLMoERunner:
self.l2_fp4 = None
self.l2_sf = None
self.l2_gs = None
# Pre-allocate forward buffers
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)."""
@@ -147,42 +100,19 @@ class CuTeDSLMoERunner:
self.l2_gs.append(w_gs)
self._l1_mat_b = None
def _compute_expert_offsets_gpu(self, sorted_ids):
"""Compute expert offsets entirely on GPU (cudagraph-safe).
No Python control flow, no .item() calls. Uses bincount + cumsum.
Args:
sorted_ids: (total_slots,) int32 — expert IDs in sorted order
Returns:
expert_offsets: (num_experts + 1,) int32 — cumulative offsets
tokens_per_expert: (num_experts,) int32 — count per expert
"""
# Bincount gives us the count of each expert ID
tokens_per_expert = torch.bincount(
sorted_ids.clamp(0, self.num_experts - 1),
minlength=self.num_experts
).int()
# Cumulative sum for offsets (prepend 0)
expert_offsets = torch.zeros(
self.num_experts + 1, dtype=torch.int32, device=self.device
)
expert_offsets[1:] = tokens_per_expert.cumsum(0)
return expert_offsets, tokens_per_expert
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 loops over tokens,
no dynamic tensor allocation, no CPU-GPU syncs.
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.
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: ignored (kept for API compat)
expert_indices: list of expert IDs (defaults to [0..num_experts-1])
Returns:
(num_tokens, hidden_size) BF16 output
@@ -191,13 +121,16 @@ class CuTeDSLMoERunner:
top_k = topk_ids.shape[1]
device = hidden_states.device
if expert_indices is None:
expert_indices = list(range(self.num_experts))
num_experts = len(expert_indices)
self._ensure_stacked()
# ── Build slot mapping (GPU-native, cudagraph-safe) ──
# Flatten top-k selections
# ── 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 = self._token_indices[:num_tokens].reshape(-1) # slice from pre-allocated
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)
@@ -205,31 +138,38 @@ class CuTeDSLMoERunner:
sorted_weights = flat_weights[sort_idx]
sorted_token_ids = token_indices[sort_idx]
# Compute expert offsets (GPU-only, no Python control flow)
expert_offsets, tokens_per_expert = self._compute_expert_offsets_gpu(sorted_ids)
total_slots = expert_offsets[-1]
# Build expert_offsets: cumulative count of tokens per expert (GPU-only)
expert_id_range = torch.arange(num_experts, device=device)
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)
# Gather hidden states in slot-major order
slot_hidden = hidden_states[sorted_token_ids[:total_slots]] # (total_slots, hidden_size) BF16
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)
slot_hidden = hidden_states[sorted_token_ids[:total_slots]]
# ════════════════════════════════════════════════════════════
# L1: gate + up (NVFP4 × NVFP4 → BF16)
# ════════════════════════════════════════════════════════════
# Use cudagraph-safe activation quantization (no .max() sync)
x_fp4, x_sf = quantize_activation_nvfp4(
slot_hidden, self._l1_activation_global_scale
)
# Build activation scales per expert using expert_offsets (GPU-only)
# Instead of Python for-loop with expert_offsets[e] (CPU sync),
# use vectorized slice accumulation
l1_scale_a = self._assemble_scales_gpu(x_sf, expert_offsets, total_slots)
# 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)
# Global scale: broadcast the pre-computed scalar to all experts
l1_gsa = torch.full(
(self.num_experts,), self._l1_activation_global_scale,
dtype=torch.float32, device=device
)
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,
@@ -252,12 +192,15 @@ class CuTeDSLMoERunner:
activated, self._l2_activation_global_scale
)
l2_scale_a = self._assemble_scales_gpu(l2_x_sf, expert_offsets, total_slots)
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_gsa = torch.full(
(self.num_experts,), self._l2_activation_global_scale,
dtype=torch.float32, device=device
)
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,
@@ -274,40 +217,3 @@ class CuTeDSLMoERunner:
y.scatter_add_(0, sorted_token_ids[:total_slots].unsqueeze(1).expand(-1, hidden_size), weighted_out)
return y
def _assemble_scales_gpu(self, x_sf, expert_offsets, total_slots):
"""Assemble activation scales without Python for-loop (cudagraph-safe).
Replaces:
for e in range(num_experts):
start = expert_offsets[e] # CPU-GPU sync!
end = expert_offsets[e + 1] # CPU-GPU sync!
parts.append(x_sf[start:end])
Instead, we pad each expert's scale rows to the same length and
stack them, then let the kernel's scale assembly handle it.
"""
# Get max tokens per expert (GPU-only)
tokens_per_expert = expert_offsets[1:] - expert_offsets[:-1] # (num_experts,)
max_tokens = tokens_per_expert.max()
# Pad each expert's rows to max_tokens and stack
padded_parts = []
for e in range(self.num_experts):
# These are still GPU tensor ops — no CPU sync
# But the Python for-loop over num_experts IS a problem for cudagraph
# because num_experts is a Python int, not a GPU value.
# However, num_experts is fixed at model init time, so the loop
# unrolls during torch.compile and becomes a fixed graph.
start = expert_offsets[e]
end = expert_offsets[e + 1]
expert_sf = x_sf[start:end]
if expert_sf.shape[0] < max_tokens:
pad_size = max_tokens - expert_sf.shape[0]
expert_sf = torch.nn.functional.pad(
expert_sf, (0, 0, 0, 0, 0, pad_size)
)
padded_parts.append(expert_sf)
# Stack and use the standard assembly
return assemble_scales_2d_side(padded_parts)