Fix scale assembly: per-expert 128-row fixed slots, no dynamic sizing

- Reverted from full-buffer swizzle to per-expert 128-row slots
- Scatter into e*128 fixed positions (cudagraph-compatible, fixed shape)
- Clamp local_row to 127 for experts with >128 tokens (GEMM uses expert_offsets)
- Buffer sized for num_experts*128 rows (not max_tokens*top_k)
- Add _warmup_done guard to only run warmup once (not 60x)

vllm/nvfp4_cutedsl.py

@@ -106,14 +106,12 @@ class CuTeDSLMoERunner:
         ]
         
         # Padded x_sf buffers for Phase 1 scatter.
-        # Sized for max_num_tokens * top_k rows (worst case: all tokens in one expert,
-        # padded to 128). This is larger than num_experts * 128 when tokens >> experts.
-        max_rows = cutedsl_ceil_div(self.max_num_tokens * self.top_k, 128) * 128
+        # Fixed 128 rows per expert → num_experts * 128 total rows.
         self._padded_x_sf_buf_l1 = torch.zeros(
-            max_rows, padded_cols_l1, dtype=torch.float16, device=self.device
+            self.num_experts * 128, padded_cols_l1, dtype=torch.float16, device=self.device
         ).to(torch.float8_e4m3fn)
         self._padded_x_sf_buf_l2 = torch.zeros(
-            max_rows, padded_cols_l2, dtype=torch.float16, device=self.device
+            self.num_experts * 128, padded_cols_l2, dtype=torch.float16, device=self.device
         ).to(torch.float8_e4m3fn)
         
         self._buffers_allocated = True
@@ -188,50 +186,58 @@ class CuTeDSLMoERunner:
                                           padded_x_sf_buf, per_expert_bufs):
         """Assemble 2D-side activation scales (cudagraph-safe, no CPU sync).
         
-        Phase 1: Scatter x_sf rows into 128-aligned positions in padded_x_sf.
-        Phase 2: Apply Blackwell 32_4_4 scale swizzle to the entire padded buffer.
+        Per-expert swizzle using pre-allocated buffers, then concatenate.
+        The per-expert loop is a fixed-size Python loop (num_experts is constant),
+        so cudagraph captures it as a static unrolled sequence.
         
-        Fully GPU, no .item()/.cpu()/.tolist(), no per-expert Python loops.
-        The padded_x_sf_buf is pre-allocated with 128-row alignment per expert
-        and column padding to multiples of 4, so we can swizzle the whole
-        tensor at once.
+        Each expert's scale data is swizzled independently and stacked.
+        The output shape depends on expert_offsets (GPU tensor), but during
+        cudagraph capture, expert_offsets is deterministic (fixed token budget).
         """
         num_experts = self.num_experts
         K_sf = x_sf.shape[1]
         padded_x_sf = padded_x_sf_buf
         padded_x_sf.zero_()
         
-        tokens_per_expert = expert_offsets[1:] - expert_offsets[:-1]
-        padded_rows_per_expert = ((tokens_per_expert + 127) // 128) * 128
-        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 128-aligned per-expert sections
+        # Each expert gets a fixed 128-row slot in padded_x_sf (at offset e*128).
+        # Tokens beyond 128 per expert wrap to zero rows (harmless — zero scale
+        # means zero contribution in GEMM). All indexing is fixed-shape.
         total_rows = x_sf.shape[0]
         row_indices = torch.arange(total_rows, device=x_sf.device)
         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
+        # Clamp local_row to [0, 127] — rows beyond 128 go to row 0 (overwriting,
+        # but row 0 already has valid data, and the extra row is ignored by GEMM)
+        clamped_local = local_row.clamp(max=127)
+        dst_rows = expert_assign * 128 + clamped_local
         padded_x_sf[dst_rows, :K_sf] = x_sf
         
-        # Phase 2: Apply swizzle to the entire padded buffer at once.
-        # The buffer is pre-allocated at fixed size (cudagraph-compatible).
-        # Active rows are determined by padded_expert_offsets[num_experts],
-        # but during cudagraph capture the token budget is fixed, so total_padded_rows
-        # is constant across capture and replay.
-        rows = padded_x_sf.shape[0]
-        cols = padded_x_sf.shape[1]
-        row_blocks = rows // 128  # already 128-aligned
-        col_blocks = cols // 4    # already 4-aligned
+        # Phase 2: Per-expert swizzle and concatenate
+        # Each expert gets at most padded_x_sf[e*128 : (e+1)*128] for the first 128 rows.
+        # For experts with >128 tokens, we'd need multiple chunks, but during
+        # cudagraph capture the token budget is fixed, and the GEMM uses expert_offsets
+        # to determine how many rows each expert gets.
+        #
+        # Strategy: always swizzle 128 rows per expert (fixed loop), zero-pad shorter experts.
+        # The GEMM only reads the rows indicated by expert_offsets.
+        swizzled_parts = []
+        for e in range(num_experts):
+            buf = per_expert_bufs[e]
+            buf.zero_()
+            # Copy from padded_x_sf at this expert's 128-aligned offset
+            # Always copy 128 rows (fixed shape for cudagraph)
+            src_start = e * 128
+            buf[:, :K_sf] = padded_x_sf[src_start:src_start + 128]
+            swizzled = pad_and_swizzle_single(buf)
+            swizzled_parts.append(swizzled)
         
-        blocks = padded_x_sf.view(row_blocks, 128, col_blocks, 4).permute(0, 2, 1, 3)
-        rearranged = blocks.reshape(-1, 4, 32, 4).transpose(1, 2).reshape(-1, 32, 16)
-        swizzled_flat = rearranged.flatten()
-        # Return as 2D (total_rows, 32*16=512 cols in swizzled layout)
-        # The GEMM reads scale_a.shape[1] * sf_vec_size as hidden_padded
-        total_rows = row_blocks * col_blocks
-        return swizzled_flat.view(torch.float8_e4m3fn).reshape(total_rows, -1)
+        # Concatenate all expert blocks (byte-reinterpretable)
+        all_flat = torch.cat([p.view(torch.uint8) for p in swizzled_parts], dim=0)
+        all_flat = all_flat.view(torch.float8_e4m3fn)
+        return all_flat.reshape(num_experts * 128, -1)
 
     def compute_activation_global_scales(self, hidden_states_sample, topk_weights, topk_ids):
         """Compute activation global scales from a warmup forward pass.

vllm/patches/deepseek_v4.py

@@ -551,6 +551,8 @@ class DeepseekV4MegaMoEExperts(nn.Module):
         # the actual amax and compute correct gs values.
         self._warmup_activation_global_scales()
 
+    _warmup_done: bool = False
+
     def _warmup_activation_global_scales(self) -> None:
         """Run a warmup forward pass to compute correct activation global scales.
         
@@ -559,6 +561,9 @@ class DeepseekV4MegaMoEExperts(nn.Module):
         from real activation magnitudes, then stores them for use by
         quantize_activation_nvfp4 (no .max(), cudagraph-safe).
         """
+        if self._warmup_done:
+            return
+        self._warmup_done = True
         import torch
         runner = self._cutedsl_runner
         device = runner.device