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, 04:45 UTC
[Perf] SM100 FP8 GEMM Optimizations after cutlass_profiler (#20071)
Signed-off-by: ilmarkov <imarkov@redhat.com> Co-authored-by: ilmarkov <imarkov@redhat.com>
Ilya Markov committed
Jun 27, 2025 at 05:50 UTC
2d7779f888f6443c067e0c36bab808ef6b368221
1 file changed
+20
-20
csrc/quantization/cutlass_w8a8/c3x/scaled_mm_sm100_fp8_dispatch.cuh
+20
-20
@@ -29,12 +29,12 @@ struct sm100_fp8_config_default {
29
template <typename InType, typename OutType,
30
template <typename, typename, typename> typename Epilogue>
31
struct sm100_fp8_config_M256 {
32
- // M in (128, 256]
32
+ // M in (64, 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>;
37
+ using ClusterShape = Shape<_2, _1, _1>;
38
using Cutlass3xGemm =
39
cutlass_3x_gemm_sm100<InType, OutType, Epilogue, TileShape, ClusterShape,
40
KernelSchedule, EpilogueSchedule>;
@@ -42,13 +42,13 @@ struct sm100_fp8_config_M256 {
42
43
template <typename InType, typename OutType,
44
template <typename, typename, typename> typename Epilogue>
45
-struct sm100_fp8_config_M128 {
46
- // M in (64, 128]
45
+struct sm100_fp8_config_M64 {
46
+ // M in (16, 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<_128, _128, _256>;
51
- using ClusterShape = Shape<_2, _4, _1>;
50
+ using TileShape = Shape<_64, _64, _128>;
51
+ using ClusterShape = Shape<_1, _1, _1>;
52
using Cutlass3xGemm =
53
cutlass_3x_gemm_sm100<InType, OutType, Epilogue, TileShape, ClusterShape,
54
KernelSchedule, EpilogueSchedule>;
@@ -56,13 +56,13 @@ struct sm100_fp8_config_M128 {
56
57
template <typename InType, typename OutType,
58
template <typename, typename, typename> typename Epilogue>
59
-struct sm100_fp8_config_M64 {
60
- // M in [1, 64]
59
+struct sm100_fp8_config_M16 {
60
+ // M in [1, 16]
61
static_assert(std::is_same<InType, cutlass::float_e4m3_t>());
62
using KernelSchedule = cutlass::gemm::collective::KernelScheduleAuto;
63
using EpilogueSchedule = cutlass::epilogue::collective::EpilogueScheduleAuto;
64
- using TileShape = Shape<_64, _64, _256>;
65
- using ClusterShape = Shape<_1, _8, _1>;
64
+ using TileShape = Shape<_64, _64, _128>;
65
+ using ClusterShape = Shape<_1, _4, _1>;
66
using Cutlass3xGemm =
67
cutlass_3x_gemm_sm100<InType, OutType, Epilogue, TileShape, ClusterShape,
68
KernelSchedule, EpilogueSchedule>;
@@ -82,27 +82,27 @@ inline void cutlass_gemm_sm100_fp8_dispatch(torch::Tensor& out,
82
using Cutlass3xGemmDefault =
83
typename sm100_fp8_config_default<InType, OutType,
84
Epilogue>::Cutlass3xGemm;
85
+ using Cutlass3xGemmM16 =
86
+ typename sm100_fp8_config_M16<InType, OutType, Epilogue>::Cutlass3xGemm;
87
using Cutlass3xGemmM64 =
88
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 =
94
- std::max(static_cast<uint32_t>(64), next_pow_2(m)); // next power of 2
94
+ std::max(static_cast<uint32_t>(16), next_pow_2(m)); // next power of 2
95
96
- if (mp2 <= 64) {
97
- // m in [1, 64]
98
- return cutlass_gemm_caller<Cutlass3xGemmM64>(
96
+ if (mp2 <= 16) {
97
+ // m in [1, 16]
98
+ return cutlass_gemm_caller<Cutlass3xGemmM16>(
99
out, a, b, std::forward<EpilogueArgs>(args)...);
100
- } else if (mp2 <= 128) {
101
- // m in (64, 128]
102
- return cutlass_gemm_caller<Cutlass3xGemmM128>(
100
+ } else if (mp2 <= 64) {
101
+ // m in (16, 64]
102
+ return cutlass_gemm_caller<Cutlass3xGemmM64>(
103
out, a, b, std::forward<EpilogueArgs>(args)...);
104
} else if (mp2 <= 256) {
105
- // m in (128, 256]
105
+ // m in (64, 256]
106
return cutlass_gemm_caller<Cutlass3xGemmM256>(
107
out, a, b, std::forward<EpilogueArgs>(args)...);
108
} else {