[Kernel] SM 12.x blockwise FP8: swizzle the CTA raster when the weight exceeds the L2

csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm120_fp8.cu

@@ -4,18 +4,45 @@
 
 namespace vllm {
 
+namespace {
+
+// CTA swizzle for SM 12.x parts whose L2 does not hold the weight operand.
+//
+// On a GB10 (24 MiB L2) the blockwise kernel loses most of its throughput once
+// the weight is re-streamed from DRAM per row of M tiles: a 16384x2560 FP8
+// weight runs at 165 TFLOPS at M=4096 but 90 at M=8192 and 52 at M>=16384;
+// 8192x8192 is at 54 from M=6144. With the tile scheduler's max_swizzle_size = 8
+// the same launches run at 150-174 TFLOPS at every M, bit-identical to the
+// default order (ten N/K shapes, M 2048-16384, all cells identical). The one
+// place the default order is better is a narrow band around M=4096 on the
+// 2560-wide weights (167 vs 153 at 16384x2560, 160 vs 154 at 12288x2560);
+// elsewhere the swizzled order is equal or up to 3.3x faster, so it is used
+// whenever the weight exceeds the L2. Parts whose L2 holds the weight (RTX PRO
+// 6000 Blackwell / GB202: 96-128 MiB) keep the default order, which is also
+// the faster one there (2560x6144 at 15 MiB: 178 vs 163 at M=2048).
+constexpr int kBlockwiseFp8SwizzleSize = 8;
+
+int blockwise_fp8_swizzle_size(int64_t weight_bytes) {
+  const int64_t l2_bytes = get_device_prop()->l2CacheSize;
+  return (l2_bytes > 0 && weight_bytes > l2_bytes) ? kBlockwiseFp8SwizzleSize
+                                                   : 1;
+}
+
+}  // namespace
+
 void cutlass_scaled_mm_blockwise_sm120_fp8(
     torch::stable::Tensor& out, torch::stable::Tensor const& a,
     torch::stable::Tensor const& b, torch::stable::Tensor const& a_scales,
     torch::stable::Tensor const& b_scales) {
+  // b is [K, N] FP8 (one byte per element).
+  const int swizzle = blockwise_fp8_swizzle_size(b.size(1) * b.size(0));
   if (out.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
     cutlass_gemm_blockwise_sm120_fp8_dispatch<cutlass::bfloat16_t>(
-        out, a, b, a_scales, b_scales);
-
+        out, a, b, a_scales, b_scales, swizzle);
   } else {
     STD_TORCH_CHECK(out.scalar_type() == torch::headeronly::ScalarType::Half);
     cutlass_gemm_blockwise_sm120_fp8_dispatch<cutlass::half_t>(
-        out, a, b, a_scales, b_scales);
+        out, a, b, a_scales, b_scales, swizzle);
   }
 }
 

csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm120_fp8_dispatch.cuh

@@ -176,7 +176,8 @@ template <typename Gemm>
 void cutlass_gemm_caller_blockwise(torch::stable::Tensor& out, torch::stable::Tensor const& a,
                                    torch::stable::Tensor const& b,
                                    torch::stable::Tensor const& a_scales,
-                                   torch::stable::Tensor const& b_scales) {
+                                   torch::stable::Tensor const& b_scales,
+                                   int max_swizzle_size) {
   static constexpr bool swap_ab = Gemm::swap_ab;
   using GemmKernel = typename Gemm::GemmKernel;
   using StrideA = typename Gemm::GemmKernel::StrideA;
@@ -238,16 +239,23 @@ void cutlass_gemm_caller_blockwise(torch::stable::Tensor& out, torch::stable::Te
   auto c_ptr = static_cast<ElementD*>(out.data_ptr());
   typename GemmKernel::EpilogueArguments epilogue_args{
       {}, c_ptr, c_stride, c_ptr, c_stride};
+  // CTA rasterization: max_swizzle_size > 1 groups nearby M/N tiles in the
+  // persistent scheduler's raster, improving the temporal locality of the
+  // shared weight (B) tiles in the L2. Bit-identical to the default order
+  // (each output tile's K-reduction is unchanged).
+  typename GemmKernel::TileSchedulerArguments scheduler{};
+  scheduler.max_swizzle_size = max_swizzle_size;
   c3x::cutlass_gemm_caller<GemmKernel>(a.device(), prob_shape, mainloop_args,
-                                       epilogue_args);
+                                       epilogue_args, scheduler);
 }
 
 template <typename OutType>
 void cutlass_gemm_blockwise_sm120_fp8_dispatch(torch::stable::Tensor& out,
                                                torch::stable::Tensor const& a,
                                                torch::stable::Tensor const& b,
                                                torch::stable::Tensor const& a_scales,
-                                               torch::stable::Tensor const& b_scales) {
+                                               torch::stable::Tensor const& b_scales,
+                                               int max_swizzle_size) {
   int M = a.size(0);
   // more heuristic tuning can be done here by checking N/K dimensions as well
   bool swap_ab = (M <= 64);
@@ -256,18 +264,18 @@ void cutlass_gemm_blockwise_sm120_fp8_dispatch(torch::stable::Tensor& out,
     if (M <= 256) {
       using Gemm = typename sm120_blockwise_fp8_config_pingpong<OutType>::Gemm;
       return cutlass_gemm_caller_blockwise<Gemm>(
-          out, a, b, a_scales, b_scales);
+          out, a, b, a_scales, b_scales, max_swizzle_size);
     }
     // M > 256: use default 128x128x128 config with Cooperative (Auto) schedule
     using Gemm = typename sm120_blockwise_fp8_config_default<OutType>::Gemm;
     return cutlass_gemm_caller_blockwise<Gemm>(
-        out, a, b, a_scales, b_scales);
+        out, a, b, a_scales, b_scales, max_swizzle_size);
   } else {
     // Swap A/B for small M to improve performance
     // Use TILE_N=32 as the minimum compatible tile size.
     using Gemm = typename sm120_blockwise_fp8_config_swapab<OutType>::Gemm;
     return cutlass_gemm_caller_blockwise<Gemm>(
-        out, a, b, a_scales, b_scales);
+        out, a, b, a_scales, b_scales, max_swizzle_size);
   }
 }
 

tests/kernels/quantization/test_cutlass_scaled_mm.py

@@ -229,7 +229,18 @@ def test_cutlass_fp8_gemm_padded(
     torch.testing.assert_close(out, baseline, rtol=5e-1, atol=1.5e-1)
 
 
-@pytest.mark.parametrize("m,n,k", MNK_FACTORS)
+# Two prefill-sized cases for the SM 12.x blockwise path, where the op switches
+# the tile scheduler to a swizzled CTA order once the weight exceeds the L2
+# (16384x2560 = 40 MiB and 5120x5120 = 25 MiB on a 24 MiB L2 part); the odd M
+# also covers the non-swap-AB dispatch. Elsewhere they run the default order
+# like the rest of the list.
+BLOCKWISE_PREFILL_FACTORS = [
+    (8193, 16384, 2560),
+    (5120, 5120, 5120),
+]
+
+
+@pytest.mark.parametrize("m,n,k", MNK_FACTORS + BLOCKWISE_PREFILL_FACTORS)
 @pytest.mark.parametrize(
     "a_scale_group_shape,b_scale_group_shape", [((1, 128), (128, 128))]
 )