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, 06:00 UTC
[Perf] Further tunings for SM100 FP8 CUTLASS kernel (#19566)
Ilya Markov committed
Jun 15, 2025 at 02:25 UTC
e13945f9dd60260ae8ed42d8d316d48401fe850b
1 file changed
+25
-5
csrc/quantization/cutlass_w8a8/c3x/scaled_mm_sm100_fp8_dispatch.cuh
+25
-5
@@ -15,11 +15,25 @@ 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)
18
+ // M in (256, 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;
22
- using TileShape = Shape<_256, _128, _64>;
22
+ using TileShape = Shape<_256, _128, _128>;
23
+ using ClusterShape = Shape<_2, _2, _1>;
24
+ using Cutlass3xGemm =
25
+ cutlass_3x_gemm_sm100<InType, OutType, Epilogue, TileShape, ClusterShape,
26
+ KernelSchedule, EpilogueSchedule>;
27
+};
28
+
29
+template <typename InType, typename OutType,
30
+ template <typename, typename, typename> typename Epilogue>
31
+struct sm100_fp8_config_M256 {
32
+ // M in (128, 256]
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, _128>;
37
using ClusterShape = Shape<_2, _2, _1>;
38
using Cutlass3xGemm =
39
cutlass_3x_gemm_sm100<InType, OutType, Epilogue, TileShape, ClusterShape,
@@ -33,8 +47,8 @@ struct sm100_fp8_config_M128 {
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;
36
- using TileShape = Shape<_128, _128, _64>;
37
- using ClusterShape = Shape<_2, _2, _1>;
50
+ using TileShape = Shape<_128, _128, _256>;
51
+ using ClusterShape = Shape<_2, _4, _1>;
52
using Cutlass3xGemm =
53
cutlass_3x_gemm_sm100<InType, OutType, Epilogue, TileShape, ClusterShape,
54
KernelSchedule, EpilogueSchedule>;
@@ -72,6 +86,8 @@ inline void cutlass_gemm_sm100_fp8_dispatch(torch::Tensor& out,
86
typename sm100_fp8_config_M64<InType, OutType, Epilogue>::Cutlass3xGemm;
87
using Cutlass3xGemmM128 =
88
typename sm100_fp8_config_M128<InType, OutType, Epilogue>::Cutlass3xGemm;
89
+ using Cutlass3xGemmM256 =
90
+ typename sm100_fp8_config_M256<InType, OutType, Epilogue>::Cutlass3xGemm;
91
92
uint32_t const m = a.size(0);
93
uint32_t const mp2 =
@@ -85,8 +101,12 @@ inline void cutlass_gemm_sm100_fp8_dispatch(torch::Tensor& out,
101
// m in (64, 128]
102
return cutlass_gemm_caller<Cutlass3xGemmM128>(
103
out, a, b, std::forward<EpilogueArgs>(args)...);
104
+ } else if (mp2 <= 256) {
105
+ // m in (128, 256]
106
+ return cutlass_gemm_caller<Cutlass3xGemmM256>(
107
+ out, a, b, std::forward<EpilogueArgs>(args)...);
108
} else {
89
- // m in (128, inf)
109
+ // m in (256, inf)
110
return cutlass_gemm_caller<Cutlass3xGemmDefault>(
111
out, a, b, std::forward<EpilogueArgs>(args)...);
112
}