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 7, 19:00 UTC
[Perf] Optimize FP8 gemm of sm120. (#34424)
[Perf] Optimize FP8 gemm of sm120. (#34424) Signed-off-by: wenshuai <wenshuai@xiaomi.com>
wenshuai committed
Feb 25, 2026 at 14:25 UTC
cd4367366814f1d1404e263e621a2b15c117eaf6
1 file changed
+133
-1
csrc/quantization/w8a8/cutlass/c3x/scaled_mm_sm120_fp8_dispatch.cuh
+133
-1
@@ -12,6 +12,68 @@ namespace vllm {
12
13
using c3x::cutlass_gemm_caller;
14
15
+// Custom wrapper to allow specifying EpilogueTile for small M
16
+template <typename ElementAB_, typename ElementD_,
17
+ template <typename, typename, typename> typename Epilogue_,
18
+ typename TileShape, typename ClusterShape, typename KernelSchedule,
19
+ typename EpilogueSchedule, typename EpilogueTile>
20
+struct cutlass_3x_gemm_sm120_custom {
21
+ using ElementAB = ElementAB_;
22
+ using LayoutA = cutlass::layout::RowMajor;
23
+ static constexpr int AlignmentA =
24
+ 128 / cutlass::sizeof_bits<ElementAB>::value;
25
+
26
+ using LayoutB = cutlass::layout::ColumnMajor;
27
+ static constexpr int AlignmentB =
28
+ 128 / cutlass::sizeof_bits<ElementAB>::value;
29
+
30
+ using ElementC = void;
31
+ using LayoutC = cutlass::layout::RowMajor;
32
+ static constexpr int AlignmentC =
33
+ 128 / cutlass::sizeof_bits<ElementD_>::value;
34
+
35
+ using ElementD = ElementD_;
36
+ using LayoutD = cutlass::layout::RowMajor;
37
+ static constexpr int AlignmentD = AlignmentC;
38
+
39
+ using ElementAcc =
40
+ typename std::conditional<std::is_same_v<ElementAB, int8_t>, int32_t,
41
+ float>::type;
42
+ using Epilogue = Epilogue_<ElementAcc, ElementD, TileShape>;
43
+
44
+ // MMA type
45
+ using ElementAccumulator = float;
46
+
47
+ // Epilogue types
48
+ using ElementBias = cutlass::half_t;
49
+ using ElementCompute = float;
50
+ using ElementAux = ElementD;
51
+ using LayoutAux = LayoutD;
52
+ using ElementAmax = float;
53
+
54
+ using EVTCompute = typename Epilogue::EVTCompute;
55
+
56
+ using CollectiveEpilogue =
57
+ typename cutlass::epilogue::collective::CollectiveBuilder<
58
+ cutlass::arch::Sm120, cutlass::arch::OpClassTensorOp, TileShape,
59
+ ClusterShape, EpilogueTile, // Use custom EpilogueTile
60
+ ElementAccumulator, ElementCompute, ElementC, LayoutC, AlignmentC,
61
+ ElementD, LayoutD, AlignmentD, EpilogueSchedule,
62
+ EVTCompute>::CollectiveOp;
63
+
64
+ using CollectiveMainloop =
65
+ typename cutlass::gemm::collective::CollectiveBuilder<
66
+ cutlass::arch::Sm120, cutlass::arch::OpClassTensorOp, ElementAB,
67
+ LayoutA, AlignmentA, ElementAB, LayoutB, AlignmentB,
68
+ ElementAccumulator, TileShape, ClusterShape,
69
+ cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(
70
+ sizeof(typename CollectiveEpilogue::SharedStorage))>,
71
+ KernelSchedule, void>::CollectiveOp;
72
+
73
+ using GemmKernel = enable_sm120_only<cutlass::gemm::kernel::GemmUniversal<
74
+ Shape<int, int, int, int>, CollectiveMainloop, CollectiveEpilogue, void>>;
75
+};
76
+
77
template <typename InType, typename OutType,
78
template <typename, typename, typename> typename Epilogue>
79
struct sm120_fp8_config_default {
@@ -25,6 +87,54 @@ struct sm120_fp8_config_default {
87
KernelSchedule, EpilogueSchedule>;
88
};
89
90
+template <typename InType, typename OutType,
91
+ template <typename, typename, typename> typename Epilogue>
92
+struct sm120_fp8_config_M64 {
93
+ static_assert(std::is_same<InType, cutlass::float_e4m3_t>());
94
+ // SM120 Cooperative kernel requires Tile M >= 128.
95
+ // For M=64 tile, we use Pingpong schedule which is more flexible with small
96
+ // tiles.
97
+ using KernelSchedule = cutlass::gemm::KernelTmaWarpSpecializedPingpong;
98
+ using EpilogueSchedule = cutlass::epilogue::collective::EpilogueScheduleAuto;
99
+ using TileShape = Shape<_64, _64, _128>;
100
+ // CUTLASS 3.x on SM120 currently restricts programmatic multicast (Cluster >
101
+ // 1) for certain schedules/types. Reverting to 1x1x1 to ensure compilation.
102
+ using ClusterShape = Shape<_1, _1, _1>;
103
+ using Cutlass3xGemm =
104
+ cutlass_3x_gemm_sm120<InType, OutType, Epilogue, TileShape, ClusterShape,
105
+ KernelSchedule, EpilogueSchedule>;
106
+};
107
+
108
+template <typename InType, typename OutType,
109
+ template <typename, typename, typename> typename Epilogue>
110
+struct sm120_fp8_config_M32 {
111
+ static_assert(std::is_same<InType, cutlass::float_e4m3_t>());
112
+ using KernelSchedule = cutlass::gemm::KernelTmaWarpSpecializedPingpong;
113
+ using EpilogueSchedule = cutlass::epilogue::collective::EpilogueScheduleAuto;
114
+ using TileShape = Shape<_32, _64, _128>;
115
+ using ClusterShape = Shape<_1, _1, _1>;
116
+ // Use custom gemm to specify EpilogueTile M=32
117
+ using Cutlass3xGemm =
118
+ cutlass_3x_gemm_sm120_custom<InType, OutType, Epilogue, TileShape,
119
+ ClusterShape, KernelSchedule,
120
+ EpilogueSchedule, Shape<_32, _32>>;
121
+};
122
+
123
+template <typename InType, typename OutType,
124
+ template <typename, typename, typename> typename Epilogue>
125
+struct sm120_fp8_config_M16 {
126
+ static_assert(std::is_same<InType, cutlass::float_e4m3_t>());
127
+ using KernelSchedule = cutlass::gemm::KernelTmaWarpSpecializedPingpong;
128
+ using EpilogueSchedule = cutlass::epilogue::collective::EpilogueScheduleAuto;
129
+ using TileShape = Shape<_16, _64, _128>;
130
+ using ClusterShape = Shape<_1, _1, _1>;
131
+ // Use custom gemm to specify EpilogueTile M=16
132
+ using Cutlass3xGemm =
133
+ cutlass_3x_gemm_sm120_custom<InType, OutType, Epilogue, TileShape,
134
+ ClusterShape, KernelSchedule,
135
+ EpilogueSchedule, Shape<_16, _32>>;
136
+};
137
+
138
template <typename InType, typename OutType,
139
template <typename, typename, typename> typename Epilogue,
140
typename... EpilogueArgs>
@@ -36,6 +146,28 @@ inline void cutlass_gemm_sm120_fp8_dispatch(torch::Tensor& out,
146
TORCH_CHECK(a.dtype() == torch::kFloat8_e4m3fn);
147
TORCH_CHECK(b.dtype() == torch::kFloat8_e4m3fn);
148
149
+ int M = a.size(0);
150
+
151
+ if (M <= 16) {
152
+ using Cutlass3xGemmM16 =
153
+ typename sm120_fp8_config_M16<InType, OutType, Epilogue>::Cutlass3xGemm;
154
+ return cutlass_gemm_caller<Cutlass3xGemmM16>(
155
+ out, a, b, std::forward<EpilogueArgs>(args)...);
156
+ }
157
+ if (M <= 32) {
158
+ using Cutlass3xGemmM32 =
159
+ typename sm120_fp8_config_M32<InType, OutType, Epilogue>::Cutlass3xGemm;
160
+ return cutlass_gemm_caller<Cutlass3xGemmM32>(
161
+ out, a, b, std::forward<EpilogueArgs>(args)...);
162
+ }
163
+
164
+ if (M <= 256) {
165
+ using Cutlass3xGemmM64 =
166
+ typename sm120_fp8_config_M64<InType, OutType, Epilogue>::Cutlass3xGemm;
167
+ return cutlass_gemm_caller<Cutlass3xGemmM64>(
168
+ out, a, b, std::forward<EpilogueArgs>(args)...);
169
+ }
170
+
171
using Cutlass3xGemmDefault =
172
typename sm120_fp8_config_default<InType, OutType,
173
Epilogue>::Cutlass3xGemm;
@@ -64,4 +196,4 @@ void cutlass_scaled_mm_sm120_fp8_epilogue(torch::Tensor& out,
196
}
197
}
198
67
-} // namespace vllm
\ No newline at end of file
199
+} // namespace vllm