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 6, 00:45 UTC
[Kernel] Optimize SM120 CUTLASS blockwise FP8 GEMM (#37970)
Signed-off-by: Necofish <liuxiangyang@mail.ustc.edu.cn> Co-authored-by: Michael Goin <mgoin64@gmail.com>
Necofish committed
Mar 25, 2026 at 23:20 UTC
e7221180e1d67a91e84a93aea08841fef8c67680
1 file changed
+36
-5
csrc/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm120_fp8_dispatch.cuh
+36
-5
@@ -110,6 +110,33 @@ struct cutlass_3x_gemm_fp8_blockwise {
110
struct GemmKernel : public KernelType {};
111
};
112
113
+// Tile configurations for different M ranges
114
+template <typename OutType>
115
+struct sm120_blockwise_fp8_config_default {
116
+ // M > 256: use 128x128x128 tile with Cooperative (Auto) schedule
117
+ using KernelSchedule = cutlass::gemm::collective::KernelScheduleAuto;
118
+ using EpilogueSchedule = cutlass::epilogue::collective::EpilogueScheduleAuto;
119
+ using TileShape = Shape<_128, _128, _128>;
120
+ using ClusterShape = Shape<_1, _1, _1>;
121
+ // ScaleGranularity must match the actual quantization block size (1, 128, 128)
122
+ using Gemm = cutlass_3x_gemm_fp8_blockwise<
123
+ OutType, 1, 128, 128, TileShape, ClusterShape,
124
+ EpilogueSchedule, KernelSchedule>;
125
+};
126
+
127
+template <typename OutType>
128
+struct sm120_blockwise_fp8_config_M64 {
129
+ // M in [1, 256]: use 64x128x128 tile with Pingpong schedule
130
+ using KernelSchedule = cutlass::gemm::KernelTmaWarpSpecializedBlockwisePingpongSm120;
131
+ using EpilogueSchedule = cutlass::epilogue::collective::EpilogueScheduleAuto;
132
+ using TileShape = Shape<_64, _128, _128>;
133
+ using ClusterShape = Shape<_1, _1, _1>;
134
+ // ScaleGranularity stays (1, 128, 128) to match actual quantization data
135
+ using Gemm = cutlass_3x_gemm_fp8_blockwise<
136
+ OutType, 1, 128, 128, TileShape, ClusterShape,
137
+ EpilogueSchedule, KernelSchedule>;
138
+};
139
+
140
template <typename Gemm>
141
void cutlass_gemm_caller_blockwise(torch::Tensor& out, torch::Tensor const& a,
142
torch::Tensor const& b,
@@ -174,11 +201,15 @@ void cutlass_gemm_blockwise_sm120_fp8_dispatch(torch::Tensor& out,
201
torch::Tensor const& b,
202
torch::Tensor const& a_scales,
203
torch::Tensor const& b_scales) {
177
- // TODO: better heuristics
178
- cutlass_gemm_caller_blockwise<cutlass_3x_gemm_fp8_blockwise<
179
- OutType, 1, 128, 128, Shape<_128, _128, _128>,
180
- Shape<_1, _1, _1>, cutlass::epilogue::collective::EpilogueScheduleAuto,
181
- cutlass::gemm::collective::KernelScheduleAuto>>(
204
+ int M = a.size(0);
205
+ if (M <= 256) {
206
+ using Gemm = typename sm120_blockwise_fp8_config_M64<OutType>::Gemm;
207
+ return cutlass_gemm_caller_blockwise<Gemm>(
208
+ out, a, b, a_scales, b_scales);
209
+ }
210
+ // M > 256: use default 128x128x128 config with Cooperative (Auto) schedule
211
+ using Gemm = typename sm120_blockwise_fp8_config_default<OutType>::Gemm;
212
+ return cutlass_gemm_caller_blockwise<Gemm>(
213
out, a, b, a_scales, b_scales);
214
}
215