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, 01:00 UTC
[Perf] Add swap_ab to SM90 FP8 non-block CUTLASS moe grouped gemm (#20911)
Signed-off-by: Shixian Cui <shixian@amazon.com> Co-authored-by: Shixian Cui <shixian@amazon.com>
shixianc committed
Jul 17, 2025 at 21:34 UTC
5780121c95bcba6c3abf45d1a7f3a33b2cb70c23
4 files changed
+135
-50
csrc/quantization/cutlass_w8a8/moe/grouped_mm_c3x.cu
+37
-12
@@ -29,19 +29,36 @@ struct sm90_fp8_config_default {
29
30
template <typename InType, typename OutType,
31
template <typename, typename, typename> typename Epilogue>
32
-struct sm90_fp8_config_M16 {
33
- // M in [1, 16]
32
+struct sm90_fp8_config_M4 {
33
+ // M in [1, 4]
34
static_assert(std::is_same<InType, cutlass::float_e4m3_t>());
35
using KernelSchedule =
36
cutlass::gemm::KernelPtrArrayTmaWarpSpecializedPingpongFP8FastAccum;
37
using EpilogueSchedule =
38
cutlass::epilogue::PtrArrayTmaWarpSpecializedPingpong;
39
- using TileShape = cute::Shape<cute::_64, cute::_64, cute::_128>;
40
- using ClusterShape = cute::Shape<cute::_1, cute::_4, cute::_1>;
39
+ using TileShape = cute::Shape<cute::_128, cute::_16, cute::_128>;
40
+ using ClusterShape = cute::Shape<cute::_1, cute::_1, cute::_1>;
41
42
using Cutlass3xGemm =
43
cutlass_3x_group_gemm<InType, OutType, Epilogue, TileShape, ClusterShape,
44
- KernelSchedule, EpilogueSchedule>;
44
+ KernelSchedule, EpilogueSchedule, true>;
45
+};
46
+
47
+template <typename InType, typename OutType,
48
+ template <typename, typename, typename> typename Epilogue>
49
+struct sm90_fp8_config_M64 {
50
+ // M in (4, 64]
51
+ static_assert(std::is_same<InType, cutlass::float_e4m3_t>());
52
+ using KernelSchedule =
53
+ cutlass::gemm::KernelPtrArrayTmaWarpSpecializedPingpongFP8FastAccum;
54
+ using EpilogueSchedule =
55
+ cutlass::epilogue::PtrArrayTmaWarpSpecializedPingpong;
56
+ using TileShape = cute::Shape<cute::_128, cute::_16, cute::_256>;
57
+ using ClusterShape = cute::Shape<cute::_2, cute::_1, cute::_1>;
58
+
59
+ using Cutlass3xGemm =
60
+ cutlass_3x_group_gemm<InType, OutType, Epilogue, TileShape, ClusterShape,
61
+ KernelSchedule, EpilogueSchedule, true>;
62
};
63
64
template <typename InType, typename OutType,
@@ -102,7 +119,9 @@ void run_cutlass_moe_mm_sm90(
119
InType, OutType, vllm::c3x::ScaledEpilogueArray>::Cutlass3xGemm;
120
using Cutlass3xGemmK8192 = typename sm90_fp8_config_K8192<
121
InType, OutType, vllm::c3x::ScaledEpilogueArray>::Cutlass3xGemm;
105
- using Cutlass3xGemmM16 = typename sm90_fp8_config_M16<
122
+ using Cutlass3xGemmM4 = typename sm90_fp8_config_M4<
123
+ InType, OutType, vllm::c3x::ScaledEpilogueArray>::Cutlass3xGemm;
124
+ using Cutlass3xGemmM64 = typename sm90_fp8_config_M64<
125
InType, OutType, vllm::c3x::ScaledEpilogueArray>::Cutlass3xGemm;
126
using Cutlass3xGemmDefault = typename sm90_fp8_config_default<
127
InType, OutType, vllm::c3x::ScaledEpilogueArray>::Cutlass3xGemm;
@@ -111,18 +130,24 @@ void run_cutlass_moe_mm_sm90(
130
uint32_t const n = out_tensors.size(1);
131
uint32_t const k = a_tensors.size(1);
132
114
- if (n >= 8192) {
115
- cutlass_group_gemm_caller<Cutlass3xGemmN8192>(
133
+ // Use swap_ab for M <= 64 by default to reduce padding
134
+ if (m <= 4) {
135
+ cutlass_group_gemm_caller<Cutlass3xGemmM4>(
136
out_tensors, a_tensors, b_tensors, a_scales, b_scales, expert_offsets,
137
problem_sizes, a_strides, b_strides, c_strides, per_act_token,
138
per_out_ch);
119
- } else if (k >= 8192) {
120
- cutlass_group_gemm_caller<Cutlass3xGemmK8192>(
139
+ } else if (m <= 64) {
140
+ cutlass_group_gemm_caller<Cutlass3xGemmM64>(
141
out_tensors, a_tensors, b_tensors, a_scales, b_scales, expert_offsets,
142
problem_sizes, a_strides, b_strides, c_strides, per_act_token,
143
per_out_ch);
124
- } else if (m <= 16) {
125
- cutlass_group_gemm_caller<Cutlass3xGemmM16>(
144
+ } else if (n >= 8192) {
145
+ cutlass_group_gemm_caller<Cutlass3xGemmN8192>(
146
+ out_tensors, a_tensors, b_tensors, a_scales, b_scales, expert_offsets,
147
+ problem_sizes, a_strides, b_strides, c_strides, per_act_token,
148
+ per_out_ch);
149
+ } else if (k >= 8192) {
150
+ cutlass_group_gemm_caller<Cutlass3xGemmK8192>(
151
out_tensors, a_tensors, b_tensors, a_scales, b_scales, expert_offsets,
152
problem_sizes, a_strides, b_strides, c_strides, per_act_token,
153
per_out_ch);
csrc/quantization/cutlass_w8a8/moe/grouped_mm_c3x.cuh
+48
-19
@@ -22,14 +22,23 @@ using ArchTag = cutlass::arch::Sm90;
22
using OperatorClass = cutlass::arch::OpClassTensorOp;
23
24
using LayoutA = cutlass::layout::RowMajor;
25
+using LayoutA_Transpose =
26
+ typename cutlass::layout::LayoutTranspose<LayoutA>::type;
27
using LayoutB = cutlass::layout::ColumnMajor;
26
-using LayoutC = cutlass::layout::RowMajor;
28
+using LayoutB_Transpose =
29
+ typename cutlass::layout::LayoutTranspose<LayoutB>::type;
30
+using LayoutD = cutlass::layout::RowMajor;
31
+using LayoutD_Transpose =
32
+ typename cutlass::layout::LayoutTranspose<LayoutD>::type;
33
+using LayoutC = LayoutD;
34
+using LayoutC_Transpose = LayoutD_Transpose;
35
36
template <typename ElementAB_, typename ElementC_,
37
template <typename, typename, typename> typename Epilogue_,
38
typename TileShape, typename ClusterShape, typename KernelSchedule,
31
- typename EpilogueSchedule>
39
+ typename EpilogueSchedule, bool swap_ab_ = false>
40
struct cutlass_3x_group_gemm {
41
+ static constexpr bool swap_ab = swap_ab_;
42
using ElementAB = ElementAB_;
43
using ElementC = void;
44
using ElementD = ElementC_;
@@ -37,9 +46,6 @@ struct cutlass_3x_group_gemm {
46
47
using Epilogue = Epilogue_<ElementAccumulator, ElementD, TileShape>;
48
40
- using StrideC =
41
- cute::remove_pointer_t<cute::Stride<int64_t, cute::Int<1>, cute::Int<0>>>;
42
-
49
static constexpr int AlignmentAB =
50
128 / cutlass::sizeof_bits<ElementAB>::value;
51
static constexpr int AlignmentC = 128 / cutlass::sizeof_bits<ElementD>::value;
@@ -50,19 +56,26 @@ struct cutlass_3x_group_gemm {
56
typename cutlass::epilogue::collective::CollectiveBuilder<
57
ArchTag, OperatorClass, TileShape, ClusterShape,
58
cutlass::epilogue::collective::EpilogueTileAuto, ElementAccumulator,
53
- ElementAccumulator, ElementC, LayoutC*, AlignmentC, ElementD,
54
- LayoutC*, AlignmentC, EpilogueSchedule, EVTCompute>::CollectiveOp;
59
+ ElementAccumulator, ElementC,
60
+ conditional_t<swap_ab, LayoutC_Transpose*, LayoutC*>, AlignmentC,
61
+ ElementD, conditional_t<swap_ab, LayoutD_Transpose*, LayoutD*>,
62
+ AlignmentC, EpilogueSchedule, EVTCompute>::CollectiveOp;
63
64
static constexpr size_t CEStorageSize =
65
sizeof(typename CollectiveEpilogue::SharedStorage);
66
using Stages = typename cutlass::gemm::collective::StageCountAutoCarveout<
67
static_cast<int>(CEStorageSize)>;
68
61
- using CollectiveMainloop =
69
+ using CollectiveMainloop = conditional_t<
70
+ swap_ab,
71
+ typename cutlass::gemm::collective::CollectiveBuilder<
72
+ ArchTag, OperatorClass, ElementAB, LayoutB_Transpose*, AlignmentAB,
73
+ ElementAB, LayoutA_Transpose*, AlignmentAB, ElementAccumulator,
74
+ TileShape, ClusterShape, Stages, KernelSchedule>::CollectiveOp,
75
typename cutlass::gemm::collective::CollectiveBuilder<
76
ArchTag, OperatorClass, ElementAB, LayoutA*, AlignmentAB, ElementAB,
77
LayoutB*, AlignmentAB, ElementAccumulator, TileShape, ClusterShape,
65
- Stages, KernelSchedule>::CollectiveOp;
78
+ Stages, KernelSchedule>::CollectiveOp>;
79
80
using KernelType = enable_sm90_only<cutlass::gemm::kernel::GemmUniversal<
81
ProblemShape, CollectiveMainloop, CollectiveEpilogue>>;
@@ -78,12 +91,12 @@ void cutlass_group_gemm_caller(
91
torch::Tensor const& problem_sizes, torch::Tensor const& a_strides,
92
torch::Tensor const& b_strides, torch::Tensor const& c_strides,
93
bool per_act_token, bool per_out_ch) {
94
+ static constexpr bool swap_ab = Gemm::swap_ab;
95
+
96
using ElementAB = typename Gemm::ElementAB;
97
using ElementD = typename Gemm::ElementD;
98
99
int num_experts = static_cast<int>(expert_offsets.size(0));
85
- int k_size = a_tensors.size(1);
86
- int n_size = out_tensors.size(1);
100
101
auto stream = at::cuda::getCurrentCUDAStream(a_tensors.device().index());
102
@@ -110,19 +123,35 @@ void cutlass_group_gemm_caller(
123
problem_sizes.data_ptr());
124
ProblemShape prob_shape{num_experts, problem_sizes_as_shapes, nullptr};
125
113
- typename GemmKernel::MainloopArguments mainloop_args{
114
- static_cast<const ElementAB**>(a_ptrs.data_ptr()),
115
- static_cast<StrideA*>(a_strides.data_ptr()),
116
- static_cast<const ElementAB**>(b_ptrs.data_ptr()),
117
- static_cast<StrideB*>(b_strides.data_ptr())};
126
+ typename GemmKernel::MainloopArguments mainloop_args;
127
+ if constexpr (swap_ab) {
128
+ mainloop_args = typename GemmKernel::MainloopArguments{
129
+ static_cast<const ElementAB**>(b_ptrs.data_ptr()),
130
+ static_cast<StrideB*>(b_strides.data_ptr()),
131
+ static_cast<const ElementAB**>(a_ptrs.data_ptr()),
132
+ static_cast<StrideA*>(a_strides.data_ptr())};
133
+ } else {
134
+ mainloop_args = typename GemmKernel::MainloopArguments{
135
+ static_cast<const ElementAB**>(a_ptrs.data_ptr()),
136
+ static_cast<StrideA*>(a_strides.data_ptr()),
137
+ static_cast<const ElementAB**>(b_ptrs.data_ptr()),
138
+ static_cast<StrideB*>(b_strides.data_ptr())};
139
+ }
140
141
// Currently, we are only able to do broadcast on either all or none a_scales
142
// and on either all or none b_scales
143
typename GemmKernel::EpilogueArguments epilogue_args{
144
Gemm::Epilogue::prepare_args(
123
- static_cast<const ElementAccumulator**>(a_scales_ptrs.data_ptr()),
124
- static_cast<const ElementAccumulator**>(b_scales_ptrs.data_ptr()),
125
- per_act_token, per_out_ch),
145
+ swap_ab ? static_cast<const ElementAccumulator**>(
146
+ b_scales_ptrs.data_ptr())
147
+ : static_cast<const ElementAccumulator**>(
148
+ a_scales_ptrs.data_ptr()),
149
+ swap_ab ? static_cast<const ElementAccumulator**>(
150
+ a_scales_ptrs.data_ptr())
151
+ : static_cast<const ElementAccumulator**>(
152
+ b_scales_ptrs.data_ptr()),
153
+ swap_ab ? per_out_ch : per_act_token,
154
+ swap_ab ? per_act_token : per_out_ch),
155
nullptr, static_cast<StrideC*>(c_strides.data_ptr()),
156
static_cast<ElementD**>(out_ptrs.data_ptr()),
157
static_cast<StrideC*>(c_strides.data_ptr())};
csrc/quantization/cutlass_w8a8/moe/moe_data.cu
+49
-19
@@ -6,7 +6,10 @@
6
#include <iostream>
7
8
constexpr uint64_t THREADS_PER_EXPERT = 512;
9
+// threshold must match the dispatch logic in run_cutlass_moe_mm_sm90()
10
+constexpr int SWAP_AB_THRESHOLD = 64;
11
12
+template <bool SWAP_AB>
13
__global__ void compute_problem_sizes(const int32_t* __restrict__ topk_ids,
14
int32_t* problem_sizes1,
15
int32_t* problem_sizes2,
@@ -24,40 +27,53 @@ __global__ void compute_problem_sizes(const int32_t* __restrict__ topk_ids,
27
28
if (threadIdx.x == 0) {
29
int final_occurrences = atomic_buffer[expert_id];
27
- problem_sizes1[expert_id * 3] = final_occurrences;
28
- problem_sizes1[expert_id * 3 + 1] = 2 * n;
29
- problem_sizes1[expert_id * 3 + 2] = k;
30
- problem_sizes2[expert_id * 3] = final_occurrences;
31
- problem_sizes2[expert_id * 3 + 1] = k;
32
- problem_sizes2[expert_id * 3 + 2] = n;
30
+ if constexpr (!SWAP_AB) {
31
+ problem_sizes1[expert_id * 3] = final_occurrences;
32
+ problem_sizes1[expert_id * 3 + 1] = 2 * n;
33
+ problem_sizes1[expert_id * 3 + 2] = k;
34
+ problem_sizes2[expert_id * 3] = final_occurrences;
35
+ problem_sizes2[expert_id * 3 + 1] = k;
36
+ problem_sizes2[expert_id * 3 + 2] = n;
37
+ } else {
38
+ problem_sizes1[expert_id * 3] = 2 * n;
39
+ problem_sizes1[expert_id * 3 + 1] = final_occurrences;
40
+ problem_sizes1[expert_id * 3 + 2] = k;
41
+ problem_sizes2[expert_id * 3] = k;
42
+ problem_sizes2[expert_id * 3 + 1] = final_occurrences;
43
+ problem_sizes2[expert_id * 3 + 2] = n;
44
+ }
45
}
46
}
47
48
__global__ void compute_expert_offsets(
49
const int32_t* __restrict__ problem_sizes1, int32_t* expert_offsets,
38
- int32_t* atomic_buffer, const int num_experts) {
50
+ int32_t* atomic_buffer, const int num_experts, const int topk_length) {
51
int32_t tot_offset = 0;
52
expert_offsets[0] = 0;
53
for (int i = 0; i < num_experts; ++i) {
54
atomic_buffer[i] = tot_offset;
43
- tot_offset += problem_sizes1[i * 3];
55
+ tot_offset += topk_length > SWAP_AB_THRESHOLD ? problem_sizes1[i * 3]
56
+ : problem_sizes1[i * 3 + 1];
57
expert_offsets[i + 1] = tot_offset;
58
}
59
}
60
61
__global__ void compute_expert_blockscale_offsets(
62
const int32_t* __restrict__ problem_sizes1, int32_t* expert_offsets,
50
- int32_t* blockscale_offsets, int32_t* atomic_buffer,
51
- const int num_experts) {
63
+ int32_t* blockscale_offsets, int32_t* atomic_buffer, const int num_experts,
64
+ const int topk_length) {
65
int32_t tot_offset = 0;
66
int32_t tot_offset_round = 0;
67
expert_offsets[0] = 0;
68
blockscale_offsets[0] = 0;
69
for (int i = 0; i < num_experts; ++i) {
70
+ int32_t cur_offset = topk_length > SWAP_AB_THRESHOLD
71
+ ? problem_sizes1[i * 3]
72
+ : problem_sizes1[i * 3 + 1];
73
atomic_buffer[i] = tot_offset;
58
- tot_offset += problem_sizes1[i * 3];
74
+ tot_offset += cur_offset;
75
expert_offsets[i + 1] = tot_offset;
60
- tot_offset_round += (problem_sizes1[i * 3] + (128 - 1)) / 128 * 128;
76
+ tot_offset_round += (cur_offset + (128 - 1)) / 128 * 128;
77
blockscale_offsets[i + 1] = tot_offset_round;
78
}
79
}
@@ -102,22 +118,36 @@ void get_cutlass_moe_mm_data_caller(
118
torch::Tensor atomic_buffer = torch::zeros(num_experts, options_int32);
119
120
int num_threads = min(THREADS_PER_EXPERT, topk_ids.numel());
105
- compute_problem_sizes<<<num_experts, num_threads, 0, stream>>>(
106
- static_cast<const int32_t*>(topk_ids.data_ptr()),
107
- static_cast<int32_t*>(problem_sizes1.data_ptr()),
108
- static_cast<int32_t*>(problem_sizes2.data_ptr()),
109
- static_cast<int32_t*>(atomic_buffer.data_ptr()), topk_ids.numel(), n, k);
121
+
122
+ if (topk_ids.numel() > SWAP_AB_THRESHOLD) {
123
+ compute_problem_sizes<false><<<num_experts, num_threads, 0, stream>>>(
124
+ static_cast<const int32_t*>(topk_ids.data_ptr()),
125
+ static_cast<int32_t*>(problem_sizes1.data_ptr()),
126
+ static_cast<int32_t*>(problem_sizes2.data_ptr()),
127
+ static_cast<int32_t*>(atomic_buffer.data_ptr()), topk_ids.numel(), n,
128
+ k);
129
+ } else {
130
+ compute_problem_sizes<true><<<num_experts, num_threads, 0, stream>>>(
131
+ static_cast<const int32_t*>(topk_ids.data_ptr()),
132
+ static_cast<int32_t*>(problem_sizes1.data_ptr()),
133
+ static_cast<int32_t*>(problem_sizes2.data_ptr()),
134
+ static_cast<int32_t*>(atomic_buffer.data_ptr()), topk_ids.numel(), n,
135
+ k);
136
+ }
137
+
138
if (blockscale_offsets.has_value()) {
139
compute_expert_blockscale_offsets<<<1, 1, 0, stream>>>(
140
static_cast<const int32_t*>(problem_sizes1.data_ptr()),
141
static_cast<int32_t*>(expert_offsets.data_ptr()),
142
static_cast<int32_t*>(blockscale_offsets.value().data_ptr()),
115
- static_cast<int32_t*>(atomic_buffer.data_ptr()), num_experts);
143
+ static_cast<int32_t*>(atomic_buffer.data_ptr()), num_experts,
144
+ topk_ids.numel());
145
} else {
146
compute_expert_offsets<<<1, 1, 0, stream>>>(
147
static_cast<const int32_t*>(problem_sizes1.data_ptr()),
148
static_cast<int32_t*>(expert_offsets.data_ptr()),
120
- static_cast<int32_t*>(atomic_buffer.data_ptr()), num_experts);
149
+ static_cast<int32_t*>(atomic_buffer.data_ptr()), num_experts,
150
+ topk_ids.numel());
151
}
152
compute_arg_sorts<<<num_experts, num_threads, 0, stream>>>(
153
static_cast<const int32_t*>(topk_ids.data_ptr()),
tests/kernels/moe/test_cutlass_moe.py
+1
@@ -25,6 +25,7 @@ MNK_FACTORS = [
25
(2, 1024, 1536),
26
(2, 3072, 1024),
27
(2, 3072, 1536),
28
+ (7, 3072, 1536),
29
(64, 1024, 1024),
30
(64, 1024, 1536),
31
(64, 3072, 1024),