From 49c28e6562a427b87662d109c24f1217ed61dfe5 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Sun, 17 May 2026 16:55:47 +0000 Subject: [PATCH] Fix: use real padded expert offsets instead of fixed layout MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- vllm/nvfp4_cutedsl.py | 54 ++++++++++++++++++++++++++----------------- 1 file changed, 33 insertions(+), 21 deletions(-) diff --git a/vllm/nvfp4_cutedsl.py b/vllm/nvfp4_cutedsl.py index 20d55caf..cdee4255 100644 --- a/vllm/nvfp4_cutedsl.py +++ b/vllm/nvfp4_cutedsl.py @@ -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)