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

Fix: use real padded expert offsets instead of fixed layout

Browse Source 
Root cause of garbage output: fixed-layout padding with
max_chunks=ceil(avg) was too small for uneven expert assignment.
Tokens beyond max_chunks*128 per expert were silently dropped
(clamped_local overwrote the same row).

Fix: compute padded_expert_offsets from actual tokens_per_expert
(padded to 128). No clamping needed — each expert gets exactly
the space it needs. Pass padded_expert_offsets to scale assembly
and GEMM.
This commit is contained in:
biondizzle
2026-05-17 16:55:47 +00:00
parent 87a223f1ac
commit 49c28e6562
1 changed files with 33 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files

54
vllm/nvfp4_cutedsl.py
View File

@@ -245,52 +245,51 @@ class CuTeDSLMoERunner:
self._l1_mat_b = None
def _assemble_scales_cudagraph_safe(self, x_sf, expert_offsets,
padded_expert_offsets,
padded_x_sf_buf, per_expert_bufs):
"""Assemble 2D-side activation scales (cudagraph-safe, no CPU sync).
Each expert's scale rows are padded to 128, then swizzled independently.
The GEMM reads scale_a according to padded_expert_offsets, which
matches the layout produced here.
Uses real padded_expert_offsets (GPU tensor) matching the GEMM's layout.
"""
num_experts = self.num_experts
K_sf = x_sf.shape[1]
padded_x_sf = padded_x_sf_buf
padded_x_sf.zero_()
# Phase 1: Scatter x_sf into fixed-layout per-expert sections
# Each expert gets max_chunks * 128 rows at offset e * max_chunks * 128.
# This matches Phase 2's fixed reading pattern.
# Phase 1: Scatter x_sf into padded per-expert sections
# Each expert e gets rows [padded_expert_offsets[e], padded_expert_offsets[e+1])
total_rows = x_sf.shape[0]
row_indices = self._row_indices_buf[:total_rows]
expert_assign = torch.searchsorted(
expert_offsets[1:], row_indices, right=True
).clamp(max=num_experts - 1)
local_row = row_indices - expert_offsets[expert_assign]
max_chunks = self._max_chunks_per_expert
# Clamp local_row to max_chunks * 128 - 1 to avoid writing beyond the expert's section
max_rows_per_expert = max_chunks * 128
clamped_local = local_row.clamp(max=max_rows_per_expert - 1)
dst_rows = expert_assign * max_rows_per_expert + clamped_local
dst_rows = padded_expert_offsets[expert_assign] + local_row
padded_x_sf[dst_rows, :K_sf] = x_sf
# Phase 2: Per-expert swizzle and concatenate
# Fixed loop: max_chunks_per_expert iterations per expert (cudagraph-safe).
# CPU-stored offsets from _allocate_buffers ensure no GPU→CPU sync.
# Fixed loop: max_chunks iterations per expert (cudagraph-safe).
# Read from padded_x_sf at padded_expert_offsets[e] + c*128.
max_chunks = self._max_chunks_per_expert
swizzled_parts = []
for e in range(num_experts):
buf = per_expert_bufs[e]
for c in range(max_chunks):
buf.zero_()
src_offset = e * max_chunks * 128 + c * 128
# Read 128 rows starting at this expert's offset + chunk offset
# padded_expert_offsets is a GPU tensor, but we only use it for
# slicing (Python sees a GPU scalar → no sync needed for __getitem__
# on a CUDA tensor with a CUDA index)
src_offset = padded_expert_offsets[e] + c * 128
buf[:, :K_sf] = padded_x_sf[src_offset:src_offset + 128]
swizzled = pad_and_swizzle_single(buf)
swizzled_parts.append(swizzled)
all_flat = torch.cat([p.view(torch.uint8) for p in swizzled_parts], dim=0)
all_flat = all_flat.view(torch.float8_e4m3fn)
# Total rows = num_experts * max_chunks * 128 (fixed)
total_padded = num_experts * max_chunks * 128
# Total padded rows from padded_expert_offsets
total_padded = padded_expert_offsets[num_experts]
return all_flat.reshape(total_padded, -1)
def compute_activation_global_scales(self, hidden_states_sample, topk_weights, topk_ids):
@@ -332,8 +331,15 @@ class CuTeDSLMoERunner:
expert_offsets.zero_()
expert_offsets[1:self.num_experts + 1] = tokens_per_expert.cumsum(0)
# Compute padded expert offsets (same as run())
padded_tokens_per_expert = ((tokens_per_expert + 127) // 128) * 128
padded_expert_offsets = self._padded_expert_offsets_buf
padded_expert_offsets.zero_()
padded_expert_offsets[1:self.num_experts + 1] = padded_tokens_per_expert.cumsum(0)
l1_scale_a = self._assemble_scales_cudagraph_safe(
x_sf, expert_offsets[:self.num_experts + 1],
padded_expert_offsets,
self._padded_x_sf_buf_l1, self._per_expert_scale_bufs_l1
)
l1_gsa = torch.full((self.num_experts,), l1_gs, dtype=torch.float32, device=device)
@@ -397,22 +403,26 @@ class CuTeDSLMoERunner:
expert_offsets.zero_()
expert_offsets[1:self.num_experts + 1] = tokens_per_expert.cumsum(0)
# Fixed padded expert offsets: each expert gets max_chunks * 128 rows
# Pad each expert to 128-row alignment (GPU-only computation)
padded_tokens_per_expert = ((tokens_per_expert + 127) // 128) * 128
padded_expert_offsets = self._padded_expert_offsets_buf
max_rows_per_expert = self._max_chunks_per_expert * 128
padded_expert_offsets.zero_()
padded_expert_offsets[1:self.num_experts + 1] = padded_tokens_per_expert.cumsum(0)
total_padded_slots = padded_expert_offsets[self.num_experts]
# -- Gather hidden states into slot order, pad per expert --
# -- Gather hidden states into slot order, scatter into padded layout --
# Each expert's tokens go at [padded_expert_offsets[e], padded_expert_offsets[e] + tokens_per_expert[e])
# Padding rows between tokens_per_expert and padded_tokens_per_expert are zero.
slot_hidden = hidden_states[sorted_token_ids]
total_padded_slots = self.num_experts * max_rows_per_expert
padded_hidden = self._shared_bufs['hidden'][:total_padded_slots]
padded_hidden.zero_()
# scatter: padded_hidden[padded_expert_offsets[expert_assign] + local_row] = slot_hidden
row_indices = self._row_indices_buf[:num_slots]
expert_assign = torch.searchsorted(
expert_offsets[1:], row_indices, right=True
).clamp(max=self.num_experts - 1)
local_row = row_indices - expert_offsets[expert_assign]
clamped_local = local_row.clamp(max=max_rows_per_expert - 1)
padded_dst = expert_assign * max_rows_per_expert + clamped_local
padded_dst = padded_expert_offsets[expert_assign] + local_row
padded_hidden[padded_dst] = slot_hidden
# === L1: gate + up ===
@@ -422,6 +432,7 @@ class CuTeDSLMoERunner:
l1_scale_a = self._assemble_scales_cudagraph_safe(
x_sf, expert_offsets[:self.num_experts + 1],
padded_expert_offsets,
self._padded_x_sf_buf_l1, self._per_expert_scale_bufs_l1
)
l1_gsa = self._l1_gsa_buf.fill_(self._l1_activation_global_scale)
@@ -452,6 +463,7 @@ class CuTeDSLMoERunner:
l2_scale_a = self._assemble_scales_cudagraph_safe(
l2_x_sf, expert_offsets[:self.num_experts + 1],
padded_expert_offsets,
self._padded_x_sf_buf_l2, self._per_expert_scale_bufs_l2
)
l2_gsa = self._l2_gsa_buf.fill_(self._l2_activation_global_scale)