A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 5, 14:45 UTC
[Kernel] Add swapAB support for SM120 CUTLASS blockwise FP8 GEMM (#38325)
Necofish committed
Apr 3, 2026 at 21:49 UTC
580090db6bcbb95ca9413ccf2d4d6726e6e69799
1 file changed
+82
-26
csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm120_fp8_dispatch.cuh
+82
-26
@@ -26,8 +26,10 @@ using namespace cute;
26
template <class OutType, int ScaleGranularityM,
27
int ScaleGranularityN, int ScaleGranularityK,
28
class MmaTileShape, class ClusterShape,
29
- class EpilogueScheduler, class MainloopScheduler>
29
+ class EpilogueScheduler, class MainloopScheduler,
30
+ bool swap_ab_ = false>
31
struct cutlass_3x_gemm_fp8_blockwise {
32
+ static constexpr bool swap_ab = swap_ab_;
33
using ElementAB = cutlass::float_e4m3_t;
34
35
using ElementA = ElementAB;
@@ -55,9 +57,13 @@ struct cutlass_3x_gemm_fp8_blockwise {
57
using ElementCompute = float;
58
using ElementBlockScale = float;
59
58
- using ScaleConfig = cutlass::detail::Sm120BlockwiseScaleConfig<
60
+ using ScaleConfig = conditional_t<swap_ab,
61
+ cutlass::detail::Sm120BlockwiseScaleConfig<
62
ScaleGranularityM, ScaleGranularityN, ScaleGranularityK,
60
- cute::UMMA::Major::MN, cute::UMMA::Major::K>;
63
+ cute::UMMA::Major::K, cute::UMMA::Major::MN>,
64
+ cutlass::detail::Sm120BlockwiseScaleConfig<
65
+ ScaleGranularityM, ScaleGranularityN, ScaleGranularityK,
66
+ cute::UMMA::Major::MN, cute::UMMA::Major::K>>;
67
68
// layout_SFA and layout_SFB cannot be swapped since they are deduced.
69
using LayoutSFA = decltype(ScaleConfig::deduce_layoutSFA());
@@ -78,17 +84,32 @@ struct cutlass_3x_gemm_fp8_blockwise {
84
ElementAccumulator,
85
ElementCompute,
86
ElementC,
81
- LayoutC,
87
+ conditional_t<swap_ab, LayoutC_Transpose, LayoutC>,
88
AlignmentC,
89
ElementD,
84
- LayoutD,
90
+ conditional_t<swap_ab, LayoutD_Transpose, LayoutD>,
91
AlignmentD,
92
EpilogueScheduler,
93
DefaultOperation
94
>::CollectiveOp;
95
96
using StageCountType = cutlass::gemm::collective::StageCountAuto;
91
- using CollectiveMainloop =
97
+ using CollectiveMainloop = conditional_t<swap_ab,
98
+ typename cutlass::gemm::collective::CollectiveBuilder<
99
+ ArchTag,
100
+ OperatorClass,
101
+ ElementB,
102
+ cute::tuple<LayoutB_Transpose, LayoutSFA>,
103
+ AlignmentB,
104
+ ElementA,
105
+ cute::tuple<LayoutA_Transpose, LayoutSFB>,
106
+ AlignmentA,
107
+ ElementAccumulator,
108
+ MmaTileShape,
109
+ ClusterShape,
110
+ cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
111
+ MainloopScheduler
112
+ >::CollectiveOp,
113
typename cutlass::gemm::collective::CollectiveBuilder<
114
ArchTag,
115
OperatorClass,
@@ -103,7 +124,7 @@ struct cutlass_3x_gemm_fp8_blockwise {
124
ClusterShape,
125
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
126
MainloopScheduler
106
- >::CollectiveOp;
127
+ >::CollectiveOp>;
128
129
// SM12x family to support both SM120 (RTX 5090) and SM121 (DGX Spark)
130
using KernelType = enable_sm120_family<cutlass::gemm::kernel::GemmUniversal<
@@ -115,7 +136,7 @@ struct cutlass_3x_gemm_fp8_blockwise {
136
// Tile configurations for different M ranges
137
template <typename OutType>
138
struct sm120_blockwise_fp8_config_default {
118
- // M > 256: use 128x128x128 tile with Cooperative (Auto) schedule
139
+ // use 128x128x128 tile with Cooperative (Auto) schedule
140
using KernelSchedule = cutlass::gemm::collective::KernelScheduleAuto;
141
using EpilogueSchedule = cutlass::epilogue::collective::EpilogueScheduleAuto;
142
using TileShape = Shape<_128, _128, _128>;
@@ -127,8 +148,8 @@ struct sm120_blockwise_fp8_config_default {
148
};
149
150
template <typename OutType>
130
-struct sm120_blockwise_fp8_config_M64 {
131
- // M in [1, 256]: use 64x128x128 tile with Pingpong schedule
151
+struct sm120_blockwise_fp8_config_pingpong {
152
+ // use 64x128x128 tile with Pingpong schedule
153
using KernelSchedule = cutlass::gemm::KernelTmaWarpSpecializedBlockwisePingpongSm120;
154
using EpilogueSchedule = cutlass::epilogue::collective::EpilogueScheduleAuto;
155
using TileShape = Shape<_64, _128, _128>;
@@ -139,11 +160,24 @@ struct sm120_blockwise_fp8_config_M64 {
160
EpilogueSchedule, KernelSchedule>;
161
};
162
163
+template <typename OutType>
164
+struct sm120_blockwise_fp8_config_swapab {
165
+ // use 128x32x128 tile with Cooperative schedule
166
+ using KernelSchedule = cutlass::gemm::KernelTmaWarpSpecializedBlockwiseCooperativeSm120;
167
+ using EpilogueSchedule = cutlass::epilogue::collective::EpilogueScheduleAuto;
168
+ using TileShape = Shape<_128, _32, _128>;
169
+ using ClusterShape = Shape<_1, _1, _1>;
170
+ using Gemm = cutlass_3x_gemm_fp8_blockwise<
171
+ OutType, 128, 1, 128, TileShape, ClusterShape,
172
+ EpilogueSchedule, KernelSchedule, true>;
173
+};
174
+
175
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) {
180
+ static constexpr bool swap_ab = Gemm::swap_ab;
181
using GemmKernel = typename Gemm::GemmKernel;
182
using StrideA = typename Gemm::GemmKernel::StrideA;
183
using StrideB = typename Gemm::GemmKernel::StrideB;
@@ -167,11 +201,13 @@ void cutlass_gemm_caller_blockwise(torch::stable::Tensor& out, torch::stable::Te
201
b_stride =
202
cutlass::make_cute_packed_stride(StrideB{}, cute::make_shape(n, k, 1));
203
c_stride =
170
- cutlass::make_cute_packed_stride(StrideC{}, cute::make_shape(m, n, 1));
204
+ cutlass::make_cute_packed_stride(StrideC{}, swap_ab ? cute::make_shape(n, m, 1) : cute::make_shape(m, n, 1));
205
172
- LayoutSFA layout_SFA =
206
+ LayoutSFA layout_SFA = swap_ab ?
207
+ ScaleConfig::tile_atom_to_shape_SFA(make_shape(n, m, k, 1)) :
208
ScaleConfig::tile_atom_to_shape_SFA(make_shape(m, n, k, 1));
174
- LayoutSFB layout_SFB =
209
+ LayoutSFB layout_SFB = swap_ab ?
210
+ ScaleConfig::tile_atom_to_shape_SFB(make_shape(n, m, k, 1)) :
211
ScaleConfig::tile_atom_to_shape_SFB(make_shape(m, n, k, 1));
212
213
auto a_ptr = static_cast<ElementAB const*>(a.data_ptr());
@@ -180,15 +216,24 @@ void cutlass_gemm_caller_blockwise(torch::stable::Tensor& out, torch::stable::Te
216
auto b_scales_ptr = static_cast<ElementBlockScale const*>(b_scales.data_ptr());
217
218
typename GemmKernel::MainloopArguments mainloop_args{};
183
- mainloop_args.ptr_A = a_ptr;
184
- mainloop_args.dA = a_stride;
185
- mainloop_args.ptr_B = b_ptr;
186
- mainloop_args.dB = b_stride;
187
- mainloop_args.ptr_SFA = a_scales_ptr;
219
mainloop_args.layout_SFA = layout_SFA;
189
- mainloop_args.ptr_SFB = b_scales_ptr;
220
mainloop_args.layout_SFB = layout_SFB;
191
- auto prob_shape = cute::make_shape(m, n, k, 1);
221
+ if (swap_ab) {
222
+ mainloop_args.ptr_A = b_ptr;
223
+ mainloop_args.dA = b_stride;
224
+ mainloop_args.ptr_B = a_ptr;
225
+ mainloop_args.dB = a_stride;
226
+ mainloop_args.ptr_SFA = b_scales_ptr;
227
+ mainloop_args.ptr_SFB = a_scales_ptr;
228
+ } else {
229
+ mainloop_args.ptr_A = a_ptr;
230
+ mainloop_args.dA = a_stride;
231
+ mainloop_args.ptr_B = b_ptr;
232
+ mainloop_args.dB = b_stride;
233
+ mainloop_args.ptr_SFA = a_scales_ptr;
234
+ mainloop_args.ptr_SFB = b_scales_ptr;
235
+ }
236
+ auto prob_shape = swap_ab ? cute::make_shape(n, m, k, 1) : cute::make_shape(m, n, k, 1);
237
238
auto c_ptr = static_cast<ElementD*>(out.data_ptr());
239
typename GemmKernel::EpilogueArguments epilogue_args{
@@ -204,15 +249,26 @@ void cutlass_gemm_blockwise_sm120_fp8_dispatch(torch::stable::Tensor& out,
249
torch::stable::Tensor const& a_scales,
250
torch::stable::Tensor const& b_scales) {
251
int M = a.size(0);
207
- if (M <= 256) {
208
- using Gemm = typename sm120_blockwise_fp8_config_M64<OutType>::Gemm;
252
+ // more heuristic tuning can be done here by checking N/K dimensions as well
253
+ bool swap_ab = (M <= 64) || (M % 4 != 0);
254
+
255
+ if (!swap_ab) {
256
+ if (M <= 256) {
257
+ using Gemm = typename sm120_blockwise_fp8_config_pingpong<OutType>::Gemm;
258
+ return cutlass_gemm_caller_blockwise<Gemm>(
259
+ out, a, b, a_scales, b_scales);
260
+ }
261
+ // M > 256: use default 128x128x128 config with Cooperative (Auto) schedule
262
+ using Gemm = typename sm120_blockwise_fp8_config_default<OutType>::Gemm;
263
+ return cutlass_gemm_caller_blockwise<Gemm>(
264
+ out, a, b, a_scales, b_scales);
265
+ } else {
266
+ // Swap A/B for small M to improve performance
267
+ // Use TILE_N=32 as the minimum compatible tile size.
268
+ using Gemm = typename sm120_blockwise_fp8_config_swapab<OutType>::Gemm;
269
return cutlass_gemm_caller_blockwise<Gemm>(
270
out, a, b, a_scales, b_scales);
271
}
212
- // M > 256: use default 128x128x128 config with Cooperative (Auto) schedule
213
- using Gemm = typename sm120_blockwise_fp8_config_default<OutType>::Gemm;
214
- return cutlass_gemm_caller_blockwise<Gemm>(
215
- out, a, b, a_scales, b_scales);
272
}
273
274
} // namespace vllm