From 94dec5922d1039a6e976c39cb5bd4095c14e612a Mon Sep 17 00:00:00 2001 From: biondizzle Date: Sun, 17 May 2026 16:56:52 +0000 Subject: [PATCH] Scale assembly Phase 2: use CPU-computed offsets for Python slicing 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. --- vllm/nvfp4_cutedsl.py | 28 ++++++++++++++++++---------- 1 file changed, 18 insertions(+), 10 deletions(-) diff --git a/vllm/nvfp4_cutedsl.py b/vllm/nvfp4_cutedsl.py index cdee4255..e08ab8da 100644 --- a/vllm/nvfp4_cutedsl.py +++ b/vllm/nvfp4_cutedsl.py @@ -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):