diff --git a/vllm/nvfp4_cutedsl.py b/vllm/nvfp4_cutedsl.py index 81f3a7e7..cb020597 100644 --- a/vllm/nvfp4_cutedsl.py +++ b/vllm/nvfp4_cutedsl.py @@ -83,6 +83,9 @@ class CuTeDSLMoERunner: self._padded_hidden_buf = None self._padded_activated_buf = None self._padded_expert_offsets_buf = None + self._max_chunks_per_expert = cutedsl_ceil_div( + self.max_num_tokens * self.top_k, self.num_experts * 128 + ) self._buffers_allocated = False def _fill_token_indices(self): @@ -249,24 +252,17 @@ class CuTeDSLMoERunner: 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). + # Unused chunks are zero buffers that contribute nothing to the GEMM. + max_chunks = self._max_chunks_per_expert swizzled_parts = [] for e in range(num_experts): - n_padded = padded_rows_per_expert[e] - start = padded_expert_offsets[e] buf = per_expert_bufs[e] - # Process in 128-row chunks - offset = start - remaining = n_padded - while remaining > 0: - buf.zero_() - chunk = min(remaining, 128) - buf[:chunk, :K_sf] = padded_x_sf[offset:offset + chunk] - swizzled = pad_and_swizzle_single(buf) - swizzled_parts.append(swizzled) - offset += 128 - remaining -= 128 - if n_padded == 0: + for c in range(max_chunks): buf.zero_() + src_offset = padded_expert_offsets[e] + c * 128 + # Copy 128 rows — rows beyond n_padded are already zero + buf[:, :K_sf] = padded_x_sf[src_offset:src_offset + 128] swizzled = pad_and_swizzle_single(buf) swizzled_parts.append(swizzled)