From 3cd910193c03204b4910c643e74cc434fc885eee Mon Sep 17 00:00:00 2001
From: biondizzle <biondizzle@gmail.com>
Date: Sun, 17 May 2026 09:59:12 +0000
Subject: [PATCH] Rewrite scale assembly: no .item() calls, no Python loops,
 fully GPU

Apply to_blocked swizzle on entire padded buffer at once instead of
per-expert loops. No .item()/.cpu() calls. Fully cudagraph-safe.
---
 vllm/nvfp4_cutedsl.py | 52 +++++++++++++++++--------------------------
 1 file changed, 21 insertions(+), 31 deletions(-)

diff --git a/vllm/nvfp4_cutedsl.py b/vllm/nvfp4_cutedsl.py
index ab6dc76a..9845a4b0 100644
--- a/vllm/nvfp4_cutedsl.py
+++ b/vllm/nvfp4_cutedsl.py
@@ -189,10 +189,12 @@ class CuTeDSLMoERunner:
         """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: Per-expert, copy from padded_x_sf at the right offset,
-                 pad_and_swizzle, then concatenate.
+        Phase 2: Apply Blackwell 32_4_4 scale swizzle to the entire padded buffer.
         
-        The output has sum(padded_rows_per_expert) rows (variable per expert).
+        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.
         """
         num_experts = self.num_experts
         K_sf = x_sf.shape[1]
@@ -213,35 +215,23 @@ class CuTeDSLMoERunner:
         dst_rows = padded_expert_offsets[expert_assign] + local_row
         padded_x_sf[dst_rows, :K_sf] = x_sf
         
-        # Phase 2: Per-expert swizzle
-        swizzled_parts = []
-        for e in range(num_experts):
-            n_padded = padded_rows_per_expert[e].item()
-            if n_padded == 0:
-                continue
-            start = padded_expert_offsets[e].item()
-            buf = per_expert_bufs[e]
-            # buf is only 128 rows; process in 128-row chunks
-            offset = start
-            remaining = n_padded
-            while remaining > 0:
-                chunk = min(remaining, 128)
-                buf.zero_()
-                buf[:chunk, :K_sf] = padded_x_sf[offset:offset + chunk]
-                swizzled = pad_and_swizzle_single(buf)
-                swizzled_parts.append(swizzled)
-                offset += chunk
-                remaining -= chunk
+        # 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
         
-        all_flat = torch.cat([p.view(torch.uint8) for p in swizzled_parts], dim=0)
-        all_flat = all_flat.view(torch.float8_e4m3fn)
-        # Reshape to 2D: (total_padded_rows, padded_cols)
-        # padded_cols comes from the swizzle: ceil_div(K_sf, 4) * 4 * 4
-        # (128 rows per row_block, 4 cols per col_block, 32 sub-rows * 16 sub-cols per block)
-        # Simpler: total elements / total_padded_rows
-        total_padded_rows = padded_expert_offsets[num_experts].item()
-        padded_cols = all_flat.shape[0] // total_padded_rows if total_padded_rows > 0 else 0
-        return all_flat.reshape(total_padded_rows, padded_cols)
+        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 = rearranged.flatten().view(torch.float8_e4m3fn)
+        # The GEMM only reads total_padded_rows worth of scale data.
+        # Return the full swizzled buffer; the GEMM uses expert_offsets to
+        # determine how many rows each expert gets.
+        return swizzled
 
     def compute_activation_global_scales(self, hidden_states_sample, topk_weights, topk_ids):
         """Compute activation global scales from a warmup forward pass.