A high-throughput and memory-efficient inference and serving engine for LLMs
[Perf] SM90 cutlass fp8 mm supports odd M by swap_ab, 180~290% kernel performance improvement (#44572)
Signed-off-by: yewentao256 <zhyanwentao@126.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Wentao Ye committed
Jun 13, 2026 at 15:05 UTC
71b961dd356a399150d25738c175c71859aa1301
3 files changed
+87
-159
csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm90_fp8_dispatch.cuh
+87
-39
@@ -25,33 +25,43 @@ using namespace cute;
25
template <class OutType, int ScaleGranularityM,
26
int ScaleGranularityN, int ScaleGranularityK,
27
class MmaTileShape, class ClusterShape,
28
- class EpilogueScheduler, class MainloopScheduler>
28
+ class EpilogueScheduler, class MainloopScheduler,
29
+ bool swap_ab_ = false>
30
struct cutlass_3x_gemm_fp8_blockwise {
31
+ static constexpr bool swap_ab = swap_ab_;
32
using ElementAB = cutlass::float_e4m3_t;
33
34
using ElementA = ElementAB;
35
using LayoutA = cutlass::layout::RowMajor;
36
+ using LayoutA_Transpose = typename cutlass::layout::LayoutTranspose<LayoutA>::type;
37
static constexpr int AlignmentA = 128 / cutlass::sizeof_bits<ElementA>::value;
38
39
using ElementB = ElementAB;
40
using LayoutB = cutlass::layout::ColumnMajor;
41
+ using LayoutB_Transpose = typename cutlass::layout::LayoutTranspose<LayoutB>::type;
42
static constexpr int AlignmentB = 128 / cutlass::sizeof_bits<ElementB>::value;
43
44
using ElementD = OutType;
45
using LayoutD = cutlass::layout::RowMajor;
46
+ using LayoutD_Transpose = typename cutlass::layout::LayoutTranspose<LayoutD>::type;
47
static constexpr int AlignmentD = 128 / cutlass::sizeof_bits<ElementD>::value;
48
49
using ElementC = void; // TODO: support bias
50
using LayoutC = LayoutD;
51
+ using LayoutC_Transpose = LayoutD_Transpose;
52
static constexpr int AlignmentC = AlignmentD;
53
54
using ElementAccumulator = float;
55
using ElementCompute = float;
56
using ElementBlockScale = float;
57
52
- using ScaleConfig = cutlass::detail::Sm90BlockwiseScaleConfig<
58
+ using ScaleConfig = conditional_t<swap_ab,
59
+ cutlass::detail::Sm90BlockwiseScaleConfig<
60
ScaleGranularityM, ScaleGranularityN, ScaleGranularityK,
54
- cute::GMMA::Major::MN, cute::GMMA::Major::K>;
61
+ cute::GMMA::Major::K, cute::GMMA::Major::MN>,
62
+ cutlass::detail::Sm90BlockwiseScaleConfig<
63
+ ScaleGranularityM, ScaleGranularityN, ScaleGranularityK,
64
+ cute::GMMA::Major::MN, cute::GMMA::Major::K>>;
65
66
using LayoutSFA = decltype(ScaleConfig::deduce_layoutSFA());
67
using LayoutSFB = decltype(ScaleConfig::deduce_layoutSFB());
@@ -71,30 +81,46 @@ struct cutlass_3x_gemm_fp8_blockwise {
81
ElementAccumulator,
82
ElementCompute,
83
ElementC,
74
- LayoutC,
84
+ conditional_t<swap_ab, LayoutC_Transpose, LayoutC>,
85
AlignmentC,
86
ElementD,
77
- LayoutD,
87
+ conditional_t<swap_ab, LayoutD_Transpose, LayoutD>,
88
AlignmentD,
89
EpilogueScheduler,
90
DefaultOperation
91
>::CollectiveOp;
92
83
- using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
84
- ArchTag,
85
- OperatorClass,
86
- ElementA,
87
- cute::tuple<LayoutA, LayoutSFA>,
88
- AlignmentA,
89
- ElementB,
90
- cute::tuple<LayoutB, LayoutSFB>,
91
- AlignmentB,
92
- ElementAccumulator,
93
- MmaTileShape,
94
- ClusterShape,
95
- cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
96
- MainloopScheduler
97
- >::CollectiveOp;
93
+ using CollectiveMainloop = conditional_t<swap_ab,
94
+ typename cutlass::gemm::collective::CollectiveBuilder<
95
+ ArchTag,
96
+ OperatorClass,
97
+ ElementB,
98
+ cute::tuple<LayoutB_Transpose, LayoutSFA>,
99
+ AlignmentB,
100
+ ElementA,
101
+ cute::tuple<LayoutA_Transpose, LayoutSFB>,
102
+ AlignmentA,
103
+ ElementAccumulator,
104
+ MmaTileShape,
105
+ ClusterShape,
106
+ cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
107
+ MainloopScheduler
108
+ >::CollectiveOp,
109
+ typename cutlass::gemm::collective::CollectiveBuilder<
110
+ ArchTag,
111
+ OperatorClass,
112
+ ElementA,
113
+ cute::tuple<LayoutA, LayoutSFA>,
114
+ AlignmentA,
115
+ ElementB,
116
+ cute::tuple<LayoutB, LayoutSFB>,
117
+ AlignmentB,
118
+ ElementAccumulator,
119
+ MmaTileShape,
120
+ ClusterShape,
121
+ cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
122
+ MainloopScheduler
123
+ >::CollectiveOp>;
124
125
using KernelType = enable_sm90_or_later<cutlass::gemm::kernel::GemmUniversal<
126
Shape<int, int, int, int>, CollectiveMainloop, CollectiveEpilogue>>;
@@ -107,6 +133,7 @@ void cutlass_gemm_caller_blockwise(torch::stable::Tensor& out, torch::stable::Te
133
torch::stable::Tensor const& b,
134
torch::stable::Tensor const& a_scales,
135
torch::stable::Tensor const& b_scales) {
136
+ static constexpr bool swap_ab = Gemm::swap_ab;
137
using GemmKernel = typename Gemm::GemmKernel;
138
using StrideA = typename Gemm::GemmKernel::StrideA;
139
using StrideB = typename Gemm::GemmKernel::StrideB;
@@ -122,8 +149,6 @@ void cutlass_gemm_caller_blockwise(torch::stable::Tensor& out, torch::stable::Te
149
150
int32_t m = a.size(0), n = b.size(1), k = a.size(1);
151
125
- STD_TORCH_CHECK(m % 4 == 0, "m must be divisible by 4");
126
-
152
StrideA a_stride;
153
StrideB b_stride;
154
StrideC c_stride;
@@ -132,12 +157,16 @@ void cutlass_gemm_caller_blockwise(torch::stable::Tensor& out, torch::stable::Te
157
b_stride =
158
cutlass::make_cute_packed_stride(StrideB{}, cute::make_shape(n, k, 1));
159
c_stride =
135
- cutlass::make_cute_packed_stride(StrideC{}, cute::make_shape(m, n, 1));
160
+ cutlass::make_cute_packed_stride(
161
+ StrideC{}, swap_ab ? cute::make_shape(n, m, 1)
162
+ : cute::make_shape(m, n, 1));
163
137
- LayoutSFA layout_SFA =
138
- ScaleConfig::tile_atom_to_shape_SFA(make_shape(m, n, k, 1));
139
- LayoutSFB layout_SFB =
140
- ScaleConfig::tile_atom_to_shape_SFB(make_shape(m, n, k, 1));
164
+ LayoutSFA layout_SFA = swap_ab
165
+ ? ScaleConfig::tile_atom_to_shape_SFA(make_shape(n, m, k, 1))
166
+ : ScaleConfig::tile_atom_to_shape_SFA(make_shape(m, n, k, 1));
167
+ LayoutSFB layout_SFB = swap_ab
168
+ ? ScaleConfig::tile_atom_to_shape_SFB(make_shape(n, m, k, 1))
169
+ : ScaleConfig::tile_atom_to_shape_SFB(make_shape(m, n, k, 1));
170
171
auto a_ptr = static_cast<ElementAB const*>(a.data_ptr());
172
auto b_ptr = static_cast<ElementAB const*>(b.data_ptr());
@@ -145,15 +174,25 @@ void cutlass_gemm_caller_blockwise(torch::stable::Tensor& out, torch::stable::Te
174
auto b_scales_ptr = static_cast<ElementBlockScale const*>(b_scales.data_ptr());
175
176
typename GemmKernel::MainloopArguments mainloop_args{};
148
- mainloop_args.ptr_A = a_ptr;
149
- mainloop_args.dA = a_stride;
150
- mainloop_args.ptr_B = b_ptr;
151
- mainloop_args.dB = b_stride;
152
- mainloop_args.ptr_SFA = a_scales_ptr;
177
mainloop_args.layout_SFA = layout_SFA;
154
- mainloop_args.ptr_SFB = b_scales_ptr;
178
mainloop_args.layout_SFB = layout_SFB;
156
- auto prob_shape = cute::make_shape(m, n, k, 1);
179
+ if (swap_ab) {
180
+ mainloop_args.ptr_A = b_ptr;
181
+ mainloop_args.dA = b_stride;
182
+ mainloop_args.ptr_B = a_ptr;
183
+ mainloop_args.dB = a_stride;
184
+ mainloop_args.ptr_SFA = b_scales_ptr;
185
+ mainloop_args.ptr_SFB = a_scales_ptr;
186
+ } else {
187
+ mainloop_args.ptr_A = a_ptr;
188
+ mainloop_args.dA = a_stride;
189
+ mainloop_args.ptr_B = b_ptr;
190
+ mainloop_args.dB = b_stride;
191
+ mainloop_args.ptr_SFA = a_scales_ptr;
192
+ mainloop_args.ptr_SFB = b_scales_ptr;
193
+ }
194
+ auto prob_shape = swap_ab ? cute::make_shape(n, m, k, 1)
195
+ : cute::make_shape(m, n, k, 1);
196
197
auto c_ptr = static_cast<ElementD*>(out.data_ptr());
198
typename GemmKernel::EpilogueArguments epilogue_args{
@@ -168,12 +207,21 @@ void cutlass_gemm_blockwise_sm90_fp8_dispatch(torch::stable::Tensor& out,
207
torch::stable::Tensor const& b,
208
torch::stable::Tensor const& a_scales,
209
torch::stable::Tensor const& b_scales) {
171
- // TODO: better heuristics
210
+ bool swap_ab = (a.size(0) % 4) != 0;
211
+ if (!swap_ab) {
212
+ cutlass_gemm_caller_blockwise<cutlass_3x_gemm_fp8_blockwise<
213
+ OutType, 1, 128, 128, Shape<_128, _128, _128>,
214
+ Shape<_1, _2, _1>, cutlass::epilogue::TmaWarpSpecializedCooperative,
215
+ cutlass::gemm::KernelTmaWarpSpecializedCooperativeFP8BlockScaledAccum>>(
216
+ out, a, b, a_scales, b_scales);
217
+ return;
218
+ }
219
+
220
cutlass_gemm_caller_blockwise<cutlass_3x_gemm_fp8_blockwise<
173
- OutType, 1, 128, 128, Shape<_128, _128, _128>,
174
- Shape<_1, _2, _1>, cutlass::epilogue::TmaWarpSpecializedCooperative,
175
- cutlass::gemm::KernelTmaWarpSpecializedCooperativeFP8BlockScaledAccum>>(
176
- out, a, b, a_scales, b_scales);
221
+ OutType, 128, 1, 128, Shape<_128, _16, _128>,
222
+ Shape<_1, _1, _1>, cutlass::epilogue::TmaWarpSpecialized,
223
+ cutlass::gemm::KernelTmaWarpSpecializedPingpongFP8BlockScaledAccum,
224
+ true>>(out, a, b, a_scales, b_scales);
225
}
226
227
} // namespace vllm
\ No newline at end of file
tests/kernels/quantization/test_cutlass_scaled_mm.py
-2
@@ -245,8 +245,6 @@ def test_cutlass_fp8_blockwise_scale_gemm(
245
return
246
if m % a_scale_group_shape[0] != 0 or k % a_scale_group_shape[1] != 0:
247
return
248
- if m % 4 != 0 and current_platform.has_device_capability(100):
249
- return
248
cutlass_fp8_gemm_helper(m, n, k, a_scale_group_shape, b_scale_group_shape, use_bias)
249
250
vllm/model_executor/kernels/linear/scaled_mm/cutlass.py
-118
@@ -20,7 +20,6 @@ from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
20
)
21
from vllm.model_executor.utils import set_weight_attrs
22
from vllm.platforms import current_platform
23
-from vllm.utils.torch_utils import direct_register_custom_op
23
24
from .BlockScaledMMLinearKernel import Fp8BlockScaledMMLinearKernel
25
from .ScaledMMLinearKernel import (
@@ -277,7 +276,6 @@ class CutlassFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel):
276
def __init__(self, config: FP8ScaledMMLinearLayerConfig) -> None:
277
super().__init__(config)
278
act_scale_descriptor = config.activation_quant_key.scale
280
- self.weight_group_shape = config.weight_quant_key.scale.group_shape
279
self.quant_fp8 = QuantFP8(
280
static=act_scale_descriptor.static,
281
group_shape=act_scale_descriptor.group_shape,
@@ -285,7 +283,6 @@ class CutlassFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel):
283
use_ue8m0=False,
284
column_major_scales=True,
285
)
288
- self.is_hopper = current_platform.is_device_capability(90)
286
287
@classmethod
288
def is_supported(cls, compute_capability=None):
@@ -320,16 +317,6 @@ class CutlassFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel):
317
Bs: torch.Tensor,
318
) -> torch.Tensor:
319
out_dtype = self.config.out_dtype
323
- if self.is_hopper:
324
- return torch.ops.vllm.dynamic_padded_cutlass(
325
- A,
326
- B,
327
- As,
328
- Bs,
329
- list(self.weight_group_shape),
330
- out_dtype,
331
- )
332
-
320
return ops.cutlass_scaled_mm(
321
A,
322
B.T,
@@ -354,108 +341,3 @@ def cutlass_scaled_mm(
341
scale_a=As,
342
scale_b=Bs.T,
343
)
357
-
358
-
359
-def _padded_cutlass(
360
- qx: torch.Tensor,
361
- weight: torch.Tensor,
362
- x_scale: torch.Tensor,
363
- weight_scale: torch.Tensor,
364
- block_size: list[int],
365
- output_dtype: torch.dtype,
366
-) -> torch.Tensor:
367
- pad_multiple = 4
368
- dim = qx.shape[0]
369
- padded = (
370
- dim if dim % pad_multiple == 0 else dim + pad_multiple - (dim % pad_multiple)
371
- )
372
-
373
- has_pad = padded > dim
374
-
375
- if has_pad:
376
- padded_shape = [padded, *qx.shape[1:]]
377
- padded_qx = torch.zeros(padded_shape, device=qx.device, dtype=qx.dtype)
378
- padded_qx[0 : qx.shape[0], ...].copy_(qx)
379
-
380
- padded_x_scale_shape = [*x_scale.shape[1:], padded]
381
- padded_x_scale = torch.ones(
382
- padded_x_scale_shape, device=x_scale.device, dtype=x_scale.dtype
383
- ).permute(-1, -2)
384
- padded_x_scale[0 : x_scale.shape[0], ...].copy_(x_scale)
385
-
386
- output = cutlass_scaled_mm(
387
- padded_qx, weight, padded_x_scale, weight_scale, block_size, output_dtype
388
- )
389
- return output[0 : qx.shape[0], ...]
390
- else:
391
- return cutlass_scaled_mm(
392
- qx, weight, x_scale, weight_scale, block_size, output_dtype
393
- )
394
-
395
-
396
-def _padded_cutlass_fake(
397
- qx: torch.Tensor,
398
- weight: torch.Tensor,
399
- x_scale: torch.Tensor,
400
- weight_scale: torch.Tensor,
401
- block_size: list[int],
402
- output_dtype: torch.dtype,
403
-) -> torch.Tensor:
404
- return torch.empty(
405
- (qx.size(0), weight.size(0)), dtype=output_dtype, device=qx.device
406
- )
407
-
408
-
409
-def _dynamic_padded_cutlass(
410
- qx: torch.Tensor,
411
- weight: torch.Tensor,
412
- x_scale: torch.Tensor,
413
- weight_scale: torch.Tensor,
414
- block_size: list[int],
415
- output_dtype: torch.dtype,
416
-) -> torch.Tensor:
417
- def run_padded(
418
- qx: torch.Tensor,
419
- weight: torch.Tensor,
420
- x_scale: torch.Tensor,
421
- weight_scale: torch.Tensor,
422
- ) -> torch.Tensor:
423
- return _padded_cutlass(
424
- qx, weight, x_scale, weight_scale, block_size, output_dtype
425
- )
426
-
427
- def run_direct(
428
- qx: torch.Tensor,
429
- weight: torch.Tensor,
430
- x_scale: torch.Tensor,
431
- weight_scale: torch.Tensor,
432
- ) -> torch.Tensor:
433
- return cutlass_scaled_mm(
434
- qx, weight, x_scale, weight_scale, block_size, output_dtype
435
- )
436
-
437
- if torch.compiler.is_compiling():
438
- return torch.cond(
439
- qx.shape[0] % 4 != 0,
440
- run_padded,
441
- run_direct,
442
- (qx, weight, x_scale, weight_scale),
443
- )
444
-
445
- if qx.shape[0] % 4 != 0:
446
- return run_padded(qx, weight, x_scale, weight_scale)
447
-
448
- return run_direct(qx, weight, x_scale, weight_scale)
449
-
450
-
451
-direct_register_custom_op(
452
- "padded_cutlass",
453
- _padded_cutlass,
454
- fake_impl=_padded_cutlass_fake,
455
-)
456
-
457
-direct_register_custom_op(
458
- "dynamic_padded_cutlass",
459
- _dynamic_padded_cutlass,
460
- fake_impl=_padded_cutlass_fake,
461
-)