diff --git a/vllm/nvfp4_cutedsl.py b/vllm/nvfp4_cutedsl.py index cb020597..63235048 100644 --- a/vllm/nvfp4_cutedsl.py +++ b/vllm/nvfp4_cutedsl.py @@ -115,9 +115,8 @@ class CuTeDSLMoERunner: for _ in range(self.num_experts) ] - # Padded x_sf buffers for Phase 1 scatter. - # Sized for max_num_tokens * top_k rows (worst case: all tokens in one expert). - max_sf_rows = self.max_num_tokens * self.top_k + # Padded x_sf buffers: num_experts * max_chunks * 128 rows (fixed layout) + max_sf_rows = self.num_experts * self._max_chunks_per_expert * 128 self._padded_x_sf_buf_l1 = torch.zeros( max_sf_rows, padded_cols_l1, dtype=torch.float16, device=self.device ).to(torch.float8_e4m3fn) @@ -149,10 +148,14 @@ class CuTeDSLMoERunner: padded_max_slots, self.intermediate_size, dtype=torch.bfloat16, device=self.device ) - # Padded expert offsets buffer (num_experts + 1) + # Padded expert offsets buffer: [0, max_rows, 2*max_rows, ...] (fixed) self._padded_expert_offsets_buf = torch.zeros( self.num_experts + 1, dtype=torch.int32, device=self.device ) + max_rows_per_expert = self._max_chunks_per_expert * 128 + self._padded_expert_offsets_buf[1:] = torch.arange( + 1, self.num_experts + 1, dtype=torch.int32, device=self.device + ) * max_rows_per_expert self._buffers_allocated = True @@ -241,27 +244,32 @@ class CuTeDSLMoERunner: padded_expert_offsets = torch.zeros(num_experts + 1, dtype=torch.int32, device=x_sf.device) padded_expert_offsets[1:] = padded_rows_per_expert.cumsum(0) - # Phase 1: Scatter x_sf into padded per-expert sections + # 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. 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] - dst_rows = padded_expert_offsets[expert_assign] + local_row + 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 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. + # CPU-stored offsets from _allocate_buffers ensure no GPU→CPU sync. 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 = padded_expert_offsets[e] + c * 128 - # Copy 128 rows — rows beyond n_padded are already zero + src_offset = e * max_chunks * 128 + c * 128 buf[:, :K_sf] = padded_x_sf[src_offset:src_offset + 128] swizzled = pad_and_swizzle_single(buf) swizzled_parts.append(swizzled) @@ -375,26 +383,22 @@ class CuTeDSLMoERunner: expert_offsets.zero_() expert_offsets[1:self.num_experts + 1] = tokens_per_expert.cumsum(0) - # Padded expert offsets (each expert padded to 128) - padded_tokens_per_expert = ((tokens_per_expert + 127) // 128) * 128 + # Fixed padded expert offsets: each expert gets max_chunks * 128 rows 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) - total_padded_slots = padded_expert_offsets[self.num_experts] + max_rows_per_expert = self._max_chunks_per_expert * 128 - # -- Gather hidden states into slot order, pad to 128 per expert -- + # -- Gather hidden states into slot order, pad per expert -- slot_hidden = hidden_states[sorted_token_ids] + total_padded_slots = self.num_experts * max_rows_per_expert padded_hidden = self._padded_hidden_buf[:total_padded_slots] padded_hidden.zero_() - # Scatter real tokens into padded positions - # Each expert e: tokens are at [expert_offsets[e], expert_offsets[e+1]) - # In padded buffer: tokens are at [padded_expert_offsets[e], padded_expert_offsets[e]+tokens_per_expert[e]) 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] - padded_dst = padded_expert_offsets[expert_assign] + local_row + clamped_local = local_row.clamp(max=max_rows_per_expert - 1) + padded_dst = expert_assign * max_rows_per_expert + clamped_local padded_hidden[padded_dst] = slot_hidden # === L1: gate + up ===