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 15, 07:15 UTC
[Perf] Tunings for SM100 FP8 CUTLASS kernel (#18778)
Signed-off-by: mgoin <mgoin64@gmail.com>
Michael Goin committed
Jun 4, 2025 at 13:46 UTC
53a5a0ce30dd623808ebd02947e5183f918b6c2f
1 file changed
+51
-2
csrc/quantization/cutlass_w8a8/c3x/scaled_mm_sm100_fp8_dispatch.cuh
+51
-2
@@ -15,6 +15,7 @@ using c3x::cutlass_gemm_caller;
15
template <typename InType, typename OutType,
16
template <typename, typename, typename> typename Epilogue>
17
struct sm100_fp8_config_default {
18
+ // M in (128, inf)
19
static_assert(std::is_same<InType, cutlass::float_e4m3_t>());
20
using KernelSchedule = cutlass::gemm::collective::KernelScheduleAuto;
21
using EpilogueSchedule = cutlass::epilogue::collective::EpilogueScheduleAuto;
@@ -25,6 +26,34 @@ struct sm100_fp8_config_default {
26
KernelSchedule, EpilogueSchedule>;
27
};
28
29
+template <typename InType, typename OutType,
30
+ template <typename, typename, typename> typename Epilogue>
31
+struct sm100_fp8_config_M128 {
32
+ // M in (64, 128]
33
+ static_assert(std::is_same<InType, cutlass::float_e4m3_t>());
34
+ using KernelSchedule = cutlass::gemm::collective::KernelScheduleAuto;
35
+ using EpilogueSchedule = cutlass::epilogue::collective::EpilogueScheduleAuto;
36
+ using TileShape = Shape<_128, _128, _64>;
37
+ using ClusterShape = Shape<_2, _2, _1>;
38
+ using Cutlass3xGemm =
39
+ cutlass_3x_gemm_sm100<InType, OutType, Epilogue, TileShape, ClusterShape,
40
+ KernelSchedule, EpilogueSchedule>;
41
+};
42
+
43
+template <typename InType, typename OutType,
44
+ template <typename, typename, typename> typename Epilogue>
45
+struct sm100_fp8_config_M64 {
46
+ // M in [1, 64]
47
+ static_assert(std::is_same<InType, cutlass::float_e4m3_t>());
48
+ using KernelSchedule = cutlass::gemm::collective::KernelScheduleAuto;
49
+ using EpilogueSchedule = cutlass::epilogue::collective::EpilogueScheduleAuto;
50
+ using TileShape = Shape<_64, _64, _256>;
51
+ using ClusterShape = Shape<_1, _8, _1>;
52
+ using Cutlass3xGemm =
53
+ cutlass_3x_gemm_sm100<InType, OutType, Epilogue, TileShape, ClusterShape,
54
+ KernelSchedule, EpilogueSchedule>;
55
+};
56
+
57
template <typename InType, typename OutType,
58
template <typename, typename, typename> typename Epilogue,
59
typename... EpilogueArgs>
@@ -39,8 +68,28 @@ inline void cutlass_gemm_sm100_fp8_dispatch(torch::Tensor& out,
68
using Cutlass3xGemmDefault =
69
typename sm100_fp8_config_default<InType, OutType,
70
Epilogue>::Cutlass3xGemm;
42
- return cutlass_gemm_caller<Cutlass3xGemmDefault>(
43
- out, a, b, std::forward<EpilogueArgs>(args)...);
71
+ using Cutlass3xGemmM64 =
72
+ typename sm100_fp8_config_M64<InType, OutType, Epilogue>::Cutlass3xGemm;
73
+ using Cutlass3xGemmM128 =
74
+ typename sm100_fp8_config_M128<InType, OutType, Epilogue>::Cutlass3xGemm;
75
+
76
+ uint32_t const m = a.size(0);
77
+ uint32_t const mp2 =
78
+ std::max(static_cast<uint32_t>(64), next_pow_2(m)); // next power of 2
79
+
80
+ if (mp2 <= 64) {
81
+ // m in [1, 64]
82
+ return cutlass_gemm_caller<Cutlass3xGemmM64>(
83
+ out, a, b, std::forward<EpilogueArgs>(args)...);
84
+ } else if (mp2 <= 128) {
85
+ // m in (64, 128]
86
+ return cutlass_gemm_caller<Cutlass3xGemmM128>(
87
+ out, a, b, std::forward<EpilogueArgs>(args)...);
88
+ } else {
89
+ // m in (128, inf)
90
+ return cutlass_gemm_caller<Cutlass3xGemmDefault>(
91
+ out, a, b, std::forward<EpilogueArgs>(args)...);
92
+ }
93
}
94
95
template <template <typename, typename, typename> typename Epilogue,