Scale assembly: full-buffer swizzle, zero CPU syncs, no Python loops

Removed .cpu().tolist() and per-expert Python loops. Apply the
Blackwell 32_4_4 swizzle to the entire padded_x_sf buffer at once.
The buffer is already 128-row aligned (padded per expert) and 4-col
aligned, so the full-buffer swizzle produces the correct layout.

The GEMM reads scale_a using padded_expert_offsets, which matches
the scatter layout. Fully GPU, zero CPU syncs, cudagraph-safe.

vllm/nvfp4_cutedsl.py

@@ -247,58 +247,41 @@ class CuTeDSLMoERunner:
     def _assemble_scales_cudagraph_safe(self, x_sf, expert_offsets,
                                           padded_expert_offsets,
                                           padded_x_sf_buf, per_expert_bufs):
-        """Assemble 2D-side activation scales (cudagraph-safe, no CPU sync).
+        """Assemble 2D-side activation scales (cudagraph-safe, NO CPU syncs).
         
-        Each expert's scale rows are padded to 128, then swizzled independently.
-        Uses real padded_expert_offsets (GPU tensor) matching the GEMM's layout.
+        Phase 1: Scatter x_sf into padded per-expert sections (GPU-only).
+        Phase 2: Apply full-buffer Blackwell 32_4_4 swizzle (no Python loops).
+        
+        The buffer is 128-row aligned per expert (from padded_expert_offsets),
+        so the full-buffer swizzle produces the correct layout. The GEMM reads
+        scale_a using padded_expert_offsets, matching the scatter layout.
         """
-        num_experts = self.num_experts
         K_sf = x_sf.shape[1]
         padded_x_sf = padded_x_sf_buf
         padded_x_sf.zero_()
         
-        # Phase 1: Scatter x_sf into padded per-expert sections
-        # Each expert e gets rows [padded_expert_offsets[e], padded_expert_offsets[e+1])
+        # Phase 1: Scatter x_sf into padded per-expert sections (GPU-only)
         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)
+        ).clamp(max=self.num_experts - 1)
         local_row = row_indices - expert_offsets[expert_assign]
         dst_rows = padded_expert_offsets[expert_assign] + local_row
         padded_x_sf[dst_rows, :K_sf] = x_sf
         
-        # Phase 2: Per-expert swizzle and concatenate
-        # Pre-compute padded_expert_offsets on CPU for Python loop indexing.
-        # During cudagraph capture, expert_offsets is deterministic (fixed token budget),
-        # so this CPU computation matches the GPU values.
-        expert_offsets_cpu = expert_offsets[:num_experts + 1].cpu().tolist()
-        padded_offsets_cpu = [0]
-        for e in range(num_experts):
-            n_tokens = expert_offsets_cpu[e + 1] - expert_offsets_cpu[e]
-            padded_n = ((n_tokens + 127) // 128) * 128
-            padded_offsets_cpu.append(padded_offsets_cpu[-1] + padded_n)
-        
-        max_chunks = self._max_chunks_per_expert
-        swizzled_parts = []
-        for e in range(num_experts):
-            buf = per_expert_bufs[e]
-            # Number of 128-row chunks for this expert
-            n_tokens = expert_offsets_cpu[e + 1] - expert_offsets_cpu[e]
-            n_chunks = (n_tokens + 127) // 128
-            for c in range(max_chunks):
-                buf.zero_()
-                if c < n_chunks:
-                    src_offset = padded_offsets_cpu[e] + c * 128
-                    buf[:, :K_sf] = padded_x_sf[src_offset:src_offset + 128]
-                # else: zero buffer (padding chunk)
-                swizzled = pad_and_swizzle_single(buf)
-                swizzled_parts.append(swizzled)
-        
-        all_flat = torch.cat([p.view(torch.uint8) for p in swizzled_parts], dim=0)
-        all_flat = all_flat.view(torch.float8_e4m3fn)
-        total_padded = padded_offsets_cpu[num_experts]
-        return all_flat.reshape(total_padded, -1)
+        # Phase 2: Full-buffer swizzle (no CPU sync, no Python loops)
+        # padded_x_sf is 128-row aligned per expert and 4-col aligned.
+        # to_blocked: (rows, cols) → view(R, 128, C, 4) → permute(0,2,1,3)
+        #   → reshape(-1, 4, 32, 4) → transpose(1,2) → reshape(-1, 32, 16) → flatten
+        rows = padded_x_sf.shape[0]
+        cols = padded_x_sf.shape[1]
+        R = rows // 128
+        C = cols // 4
+        blocks = padded_x_sf.view(R, 128, C, 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)
+        return swizzled.reshape(rows, cols)
 
     def compute_activation_global_scales(self, hidden_states_sample, topk_weights, topk_ids):
         """Compute activation global scales from a warmup forward pass.