[Kernel] SM 12.x blockwise FP8: swizzle the CTA raster when the weight exceeds the L2 (#55180)
Jürgen Schmied committed
Sep 7, 2026 at 01:36 UTC
4df80187b5ca8453eb7b47d2914f0555acf9c12f
3 files changed
+56
-10
csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm120_fp8.cu
+30
-3
@@ -4,18 +4,45 @@
4
5
namespace vllm {
6
7
+namespace {
8
+
9
+// CTA swizzle for SM 12.x parts whose L2 does not hold the weight operand.
10
+//
11
+// On a GB10 (24 MiB L2) the blockwise kernel loses most of its throughput once
12
+// the weight is re-streamed from DRAM per row of M tiles: a 16384x2560 FP8
13
+// weight runs at 165 TFLOPS at M=4096 but 90 at M=8192 and 52 at M>=16384;
14
+// 8192x8192 is at 54 from M=6144. With the tile scheduler's max_swizzle_size = 8
15
+// the same launches run at 150-174 TFLOPS at every M, bit-identical to the
16
+// default order (ten N/K shapes, M 2048-16384, all cells identical). The one
17
+// place the default order is better is a narrow band around M=4096 on the
18
+// 2560-wide weights (167 vs 153 at 16384x2560, 160 vs 154 at 12288x2560);
19
+// elsewhere the swizzled order is equal or up to 3.3x faster, so it is used
20
+// whenever the weight exceeds the L2. Parts whose L2 holds the weight (RTX PRO
21
+// 6000 Blackwell / GB202: 96-128 MiB) keep the default order, which is also
22
+// the faster one there (2560x6144 at 15 MiB: 178 vs 163 at M=2048).
23
+constexpr int kBlockwiseFp8SwizzleSize = 8;
24
+
25
+int blockwise_fp8_swizzle_size(int64_t weight_bytes) {
26
+ const int64_t l2_bytes = get_device_prop()->l2CacheSize;
27
+ return (l2_bytes > 0 && weight_bytes > l2_bytes) ? kBlockwiseFp8SwizzleSize
28
+ : 1;
29
+}
30
+
31
+} // namespace
32
+
33
void cutlass_scaled_mm_blockwise_sm120_fp8(
34
torch::stable::Tensor& out, torch::stable::Tensor const& a,
35
torch::stable::Tensor const& b, torch::stable::Tensor const& a_scales,
36
torch::stable::Tensor const& b_scales) {
37
+ // b is [K, N] FP8 (one byte per element).
38
+ const int swizzle = blockwise_fp8_swizzle_size(b.size(1) * b.size(0));
39
if (out.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
40
cutlass_gemm_blockwise_sm120_fp8_dispatch<cutlass::bfloat16_t>(
13
- out, a, b, a_scales, b_scales);
14
-
41
+ out, a, b, a_scales, b_scales, swizzle);
42
} else {
43
STD_TORCH_CHECK(out.scalar_type() == torch::headeronly::ScalarType::Half);
44
cutlass_gemm_blockwise_sm120_fp8_dispatch<cutlass::half_t>(
18
- out, a, b, a_scales, b_scales);
45
+ out, a, b, a_scales, b_scales, swizzle);
46
}
47
}
48
csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm120_fp8_dispatch.cuh
+14
-6
@@ -176,7 +176,8 @@ template <typename Gemm>
176
void cutlass_gemm_caller_blockwise(torch::stable::Tensor& out, torch::stable::Tensor const& a,
177
torch::stable::Tensor const& b,
178
torch::stable::Tensor const& a_scales,
179
- torch::stable::Tensor const& b_scales) {
179
+ torch::stable::Tensor const& b_scales,
180
+ int max_swizzle_size) {
181
static constexpr bool swap_ab = Gemm::swap_ab;
182
using GemmKernel = typename Gemm::GemmKernel;
183
using StrideA = typename Gemm::GemmKernel::StrideA;
@@ -238,8 +239,14 @@ void cutlass_gemm_caller_blockwise(torch::stable::Tensor& out, torch::stable::Te
239
auto c_ptr = static_cast<ElementD*>(out.data_ptr());
240
typename GemmKernel::EpilogueArguments epilogue_args{
241
{}, c_ptr, c_stride, c_ptr, c_stride};
242
+ // CTA rasterization: max_swizzle_size > 1 groups nearby M/N tiles in the
243
+ // persistent scheduler's raster, improving the temporal locality of the
244
+ // shared weight (B) tiles in the L2. Bit-identical to the default order
245
+ // (each output tile's K-reduction is unchanged).
246
+ typename GemmKernel::TileSchedulerArguments scheduler{};
247
+ scheduler.max_swizzle_size = max_swizzle_size;
248
c3x::cutlass_gemm_caller<GemmKernel>(a.device(), prob_shape, mainloop_args,
242
- epilogue_args);
249
+ epilogue_args, scheduler);
250
}
251
252
template <typename OutType>
@@ -247,7 +254,8 @@ void cutlass_gemm_blockwise_sm120_fp8_dispatch(torch::stable::Tensor& out,
254
torch::stable::Tensor const& a,
255
torch::stable::Tensor const& b,
256
torch::stable::Tensor const& a_scales,
250
- torch::stable::Tensor const& b_scales) {
257
+ torch::stable::Tensor const& b_scales,
258
+ int max_swizzle_size) {
259
int M = a.size(0);
260
// more heuristic tuning can be done here by checking N/K dimensions as well
261
bool swap_ab = (M <= 64);
@@ -256,18 +264,18 @@ void cutlass_gemm_blockwise_sm120_fp8_dispatch(torch::stable::Tensor& out,
264
if (M <= 256) {
265
using Gemm = typename sm120_blockwise_fp8_config_pingpong<OutType>::Gemm;
266
return cutlass_gemm_caller_blockwise<Gemm>(
259
- out, a, b, a_scales, b_scales);
267
+ out, a, b, a_scales, b_scales, max_swizzle_size);
268
}
269
// M > 256: use default 128x128x128 config with Cooperative (Auto) schedule
270
using Gemm = typename sm120_blockwise_fp8_config_default<OutType>::Gemm;
271
return cutlass_gemm_caller_blockwise<Gemm>(
264
- out, a, b, a_scales, b_scales);
272
+ out, a, b, a_scales, b_scales, max_swizzle_size);
273
} else {
274
// Swap A/B for small M to improve performance
275
// Use TILE_N=32 as the minimum compatible tile size.
276
using Gemm = typename sm120_blockwise_fp8_config_swapab<OutType>::Gemm;
277
return cutlass_gemm_caller_blockwise<Gemm>(
270
- out, a, b, a_scales, b_scales);
278
+ out, a, b, a_scales, b_scales, max_swizzle_size);
279
}
280
}
281
tests/kernels/quantization/test_cutlass_scaled_mm.py
+12
-1
@@ -229,7 +229,18 @@ def test_cutlass_fp8_gemm_padded(
229
torch.testing.assert_close(out, baseline, rtol=5e-1, atol=1.5e-1)
230
231
232
-@pytest.mark.parametrize("m,n,k", MNK_FACTORS)
232
+# Two prefill-sized cases for the SM 12.x blockwise path, where the op switches
233
+# the tile scheduler to a swizzled CTA order once the weight exceeds the L2
234
+# (16384x2560 = 40 MiB and 5120x5120 = 25 MiB on a 24 MiB L2 part); the odd M
235
+# also covers the non-swap-AB dispatch. Elsewhere they run the default order
236
+# like the rest of the list.
237
+BLOCKWISE_PREFILL_FACTORS = [
238
+ (8193, 16384, 2560),
239
+ (5120, 5120, 5120),
240
+]
241
+
242
+
243
+@pytest.mark.parametrize("m,n,k", MNK_FACTORS + BLOCKWISE_PREFILL_FACTORS)
244
@pytest.mark.parametrize(
245
"a_scale_group_shape,b_scale_group_shape", [((1, 128), (128, 128))]
246
)