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

Scale assembly Phase 2: use CPU-computed offsets for Python slicing

Browse Source 
GPU scalars can't be used for Python indexing (requires sync).
Compute padded_expert_offsets on CPU via .cpu().tolist() for
the Python loop. This is OK for cudagraph: Python code only
runs during capture, not replay. The GPU kernel launches
recorded during capture are deterministic.
This commit is contained in:
biondizzle
2026-05-17 16:56:52 +00:00
parent 49c28e6562
commit 94dec5922d
1 changed files with 18 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

28
vllm/nvfp4_cutedsl.py
View File

@@ -269,27 +269,35 @@ class CuTeDSLMoERunner:
padded_x_sf[dst_rows, :K_sf] = x_sf
# Phase 2: Per-expert swizzle and concatenate
# Fixed loop: max_chunks iterations per expert (cudagraph-safe).
# Read from padded_x_sf at padded_expert_offsets[e] + c*128.
# Pre-compute padded_expert_offsets on CPU for Python loop indexing.
# During cudagraph capture, expert_offsets is deterministic (fixed token budget),
# so this CPU computation matches the GPU values.
expert_offsets_cpu = expert_offsets[:num_experts + 1].cpu().tolist()
padded_offsets_cpu = [0]
for e in range(num_experts):
n_tokens = expert_offsets_cpu[e + 1] - expert_offsets_cpu[e]
padded_n = ((n_tokens + 127) // 128) * 128
padded_offsets_cpu.append(padded_offsets_cpu[-1] + padded_n)
max_chunks = self._max_chunks_per_expert
swizzled_parts = []
for e in range(num_experts):
buf = per_expert_bufs[e]
# Number of 128-row chunks for this expert
n_tokens = expert_offsets_cpu[e + 1] - expert_offsets_cpu[e]
n_chunks = (n_tokens + 127) // 128
for c in range(max_chunks):
buf.zero_()
# 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]
if c < n_chunks:
src_offset = padded_offsets_cpu[e] + c * 128
buf[:, :K_sf] = padded_x_sf[src_offset:src_offset + 128]
# else: zero buffer (padding chunk)
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 padded rows from padded_expert_offsets
total_padded = padded_expert_offsets[num_experts]
total_padded = padded_offsets_cpu[num_experts]
return all_flat.reshape(total_padded, -1)
def compute_activation_global_scales(self, hidden_states_sample, topk_weights, topk_ids):