Fix cudagraph: fully fixed-layout per-expert sections, no GPU scalars in Python control flow

- Each expert gets max_chunks*128 rows at fixed offsets (e*max_chunks*128)
- Phase 1 scatters into fixed offsets with clamped local_row
- Phase 2 reads from fixed offsets (pure Python arithmetic, no GPU sync)
- padded_x_sf_buf sized for num_experts * max_chunks * 128
- padded_expert_offsets pre-computed in _allocate_buffers

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 ===