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
[Hardware][NVIDIA] FP4 MoE kernel optimization (#19110)
Signed-off-by: Chiyue Wei <chiyuew@nvidia.com> Co-authored-by: Chiyue Wei <chiyuew@nvidia.com>
Chiyue Wei committed
Jun 5, 2025 at 09:48 UTC
61059bee40511b6f6c044053cf921da81cf89985
12 files changed
+165
-38
benchmarks/kernels/benchmark_cutlass_fp4_moe.py
+1
-1
@@ -91,7 +91,7 @@ def bench_run(
91
92
score = torch.randn((m, num_experts), device=device, dtype=dtype)
93
94
- topk_weights, topk_ids = fused_topk(a, score, topk, renormalize=False)
94
+ topk_weights, topk_ids, _ = fused_topk(a, score, topk, renormalize=False)
95
96
quant_blocksize = 16
97
w1_blockscale = torch.empty(
csrc/moe/moe_ops.h
+5
-1
@@ -30,4 +30,8 @@ torch::Tensor moe_wna16_gemm(torch::Tensor input, torch::Tensor output,
30
int64_t BLOCK_SIZE_K, int64_t bit);
31
#endif
32
33
-bool moe_permute_unpermute_supported();
\ No newline at end of file
33
+bool moe_permute_unpermute_supported();
34
+
35
+void shuffle_rows(const torch::Tensor& input_tensor,
36
+ const torch::Tensor& dst2src_map,
37
+ torch::Tensor& output_tensor);
\ No newline at end of file
csrc/moe/moe_permute_unpermute_op.cu
+56
@@ -130,6 +130,62 @@ void moe_unpermute(
130
});
131
}
132
133
+template <typename T>
134
+__global__ void shuffleInputRowsKernel(const T* input,
135
+ const int32_t* dst2src_map, T* output,
136
+ int64_t num_src_rows,
137
+ int64_t num_dst_rows, int64_t num_cols) {
138
+ int64_t dest_row_idx = blockIdx.x;
139
+ int64_t const source_row_idx = dst2src_map[dest_row_idx];
140
+
141
+ if (blockIdx.x < num_dst_rows) {
142
+ // Load 128-bits per thread
143
+ constexpr int64_t ELEM_PER_THREAD = 128 / sizeof(T) / 8;
144
+ using DataElem = cutlass::Array<T, ELEM_PER_THREAD>;
145
+
146
+ // Duplicate and permute rows
147
+ auto const* source_row_ptr =
148
+ reinterpret_cast<DataElem const*>(input + source_row_idx * num_cols);
149
+ auto* dest_row_ptr =
150
+ reinterpret_cast<DataElem*>(output + dest_row_idx * num_cols);
151
+
152
+ int64_t const start_offset = threadIdx.x;
153
+ int64_t const stride = blockDim.x;
154
+ int64_t const num_elems_in_col = num_cols / ELEM_PER_THREAD;
155
+
156
+ for (int elem_index = start_offset; elem_index < num_elems_in_col;
157
+ elem_index += stride) {
158
+ dest_row_ptr[elem_index] = source_row_ptr[elem_index];
159
+ }
160
+ }
161
+}
162
+
163
+void shuffle_rows(const torch::Tensor& input_tensor,
164
+ const torch::Tensor& dst2src_map,
165
+ torch::Tensor& output_tensor) {
166
+ TORCH_CHECK(input_tensor.scalar_type() == output_tensor.scalar_type(),
167
+ "Input and output tensors must have the same data type");
168
+
169
+ auto stream = at::cuda::getCurrentCUDAStream().stream();
170
+ int64_t const blocks = output_tensor.size(0);
171
+ int64_t const threads = 256;
172
+ int64_t const num_dest_rows = output_tensor.size(0);
173
+ int64_t const num_src_rows = input_tensor.size(0);
174
+ int64_t const num_cols = input_tensor.size(1);
175
+
176
+ TORCH_CHECK(!(num_cols % (128 / sizeof(input_tensor.scalar_type()) / 8)),
177
+ "num_cols must be divisible by 128 / "
178
+ "sizeof(input_tensor.scalar_type()) / 8");
179
+
180
+ MOE_DISPATCH(input_tensor.scalar_type(), [&] {
181
+ shuffleInputRowsKernel<scalar_t><<<blocks, threads, 0, stream>>>(
182
+ reinterpret_cast<scalar_t*>(input_tensor.data_ptr()),
183
+ dst2src_map.data_ptr<int32_t>(),
184
+ reinterpret_cast<scalar_t*>(output_tensor.data_ptr()), num_src_rows,
185
+ num_dest_rows, num_cols);
186
+ });
187
+}
188
+
189
#else
190
191
void moe_permute(const torch::Tensor& input, const torch::Tensor& topk_weights,
csrc/moe/permute_unpermute_kernels/dispatch.h
+12
-6
@@ -14,12 +14,13 @@
14
__VA_ARGS__(); \
15
break; \
16
}
17
-#define MOE_DISPATCH_FLOAT_CASE(...) \
18
- MOE_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \
19
- MOE_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__) \
20
- MOE_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__) \
21
- MOE_DISPATCH_CASE(at::ScalarType::Float8_e5m2, __VA_ARGS__) \
22
- MOE_DISPATCH_CASE(at::ScalarType::Float8_e4m3fn, __VA_ARGS__)
17
+#define MOE_DISPATCH_FLOAT_CASE(...) \
18
+ MOE_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \
19
+ MOE_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__) \
20
+ MOE_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__) \
21
+ MOE_DISPATCH_CASE(at::ScalarType::Float8_e5m2, __VA_ARGS__) \
22
+ MOE_DISPATCH_CASE(at::ScalarType::Float8_e4m3fn, __VA_ARGS__) \
23
+ MOE_DISPATCH_CASE(at::ScalarType::Byte, __VA_ARGS__)
24
25
#define MOE_DISPATCH(TYPE, ...) \
26
MOE_SWITCH(TYPE, MOE_DISPATCH_FLOAT_CASE(__VA_ARGS__))
@@ -39,6 +40,11 @@ template <>
40
struct ScalarType2CudaType<at::ScalarType::BFloat16> {
41
using type = __nv_bfloat16;
42
};
43
+// uint8 for packed fp4
44
+template <>
45
+struct ScalarType2CudaType<at::ScalarType::Byte> {
46
+ using type = uint8_t;
47
+};
48
49
// #if __CUDA_ARCH__ >= 890
50
// fp8
csrc/moe/torch_bindings.cpp
+6
@@ -81,6 +81,12 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, m) {
81
m.def("moe_permute_unpermute_supported() -> bool");
82
m.impl("moe_permute_unpermute_supported", &moe_permute_unpermute_supported);
83
84
+ // Row shuffle for MoE
85
+ m.def(
86
+ "shuffle_rows(Tensor input_tensor, Tensor dst2src_map, Tensor! "
87
+ "output_tensor) -> ()");
88
+ m.impl("shuffle_rows", torch::kCUDA, &shuffle_rows);
89
+
90
#endif
91
}
92
csrc/ops.h
+2
-1
@@ -248,7 +248,8 @@ void get_cutlass_moe_mm_data(
248
const torch::Tensor& topk_ids, torch::Tensor& expert_offsets,
249
torch::Tensor& problem_sizes1, torch::Tensor& problem_sizes2,
250
torch::Tensor& input_permutation, torch::Tensor& output_permutation,
251
- const int64_t num_experts, const int64_t n, const int64_t k);
251
+ const int64_t num_experts, const int64_t n, const int64_t k,
252
+ const std::optional<torch::Tensor>& blockscale_offsets);
253
254
void cutlass_scaled_mm_azp(torch::Tensor& out, torch::Tensor const& a,
255
torch::Tensor const& b,
csrc/quantization/cutlass_w8a8/moe/moe_data.cu
+31
-5
@@ -45,6 +45,23 @@ __global__ void compute_expert_offsets(
45
}
46
}
47
48
+__global__ void compute_expert_blockscale_offsets(
49
+ const int32_t* __restrict__ problem_sizes1, int32_t* expert_offsets,
50
+ int32_t* blockscale_offsets, int32_t* atomic_buffer,
51
+ const int num_experts) {
52
+ int32_t tot_offset = 0;
53
+ int32_t tot_offset_round = 0;
54
+ expert_offsets[0] = 0;
55
+ blockscale_offsets[0] = 0;
56
+ for (int i = 0; i < num_experts; ++i) {
57
+ atomic_buffer[i] = tot_offset;
58
+ tot_offset += problem_sizes1[i * 3];
59
+ expert_offsets[i + 1] = tot_offset;
60
+ tot_offset_round += (problem_sizes1[i * 3] + (128 - 1)) / 128 * 128;
61
+ blockscale_offsets[i + 1] = tot_offset_round;
62
+ }
63
+}
64
+
65
__global__ void compute_arg_sorts(const int* __restrict__ topk_ids,
66
const int32_t* __restrict__ expert_offsets,
67
int32_t* input_permutation,
@@ -77,7 +94,8 @@ void get_cutlass_moe_mm_data_caller(
94
const torch::Tensor& topk_ids, torch::Tensor& expert_offsets,
95
torch::Tensor& problem_sizes1, torch::Tensor& problem_sizes2,
96
torch::Tensor& input_permutation, torch::Tensor& output_permutation,
80
- const int64_t num_experts, const int64_t n, const int64_t k) {
97
+ const int64_t num_experts, const int64_t n, const int64_t k,
98
+ const std::optional<torch::Tensor>& blockscale_offsets) {
99
auto stream = at::cuda::getCurrentCUDAStream(topk_ids.device().index());
100
auto options_int32 =
101
torch::TensorOptions().dtype(torch::kInt32).device(topk_ids.device());
@@ -89,10 +107,18 @@ void get_cutlass_moe_mm_data_caller(
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);
92
- compute_expert_offsets<<<1, 1, 0, stream>>>(
93
- static_cast<const int32_t*>(problem_sizes1.data_ptr()),
94
- static_cast<int32_t*>(expert_offsets.data_ptr()),
95
- static_cast<int32_t*>(atomic_buffer.data_ptr()), num_experts);
110
+ if (blockscale_offsets.has_value()) {
111
+ compute_expert_blockscale_offsets<<<1, 1, 0, stream>>>(
112
+ static_cast<const int32_t*>(problem_sizes1.data_ptr()),
113
+ static_cast<int32_t*>(expert_offsets.data_ptr()),
114
+ static_cast<int32_t*>(blockscale_offsets.value().data_ptr()),
115
+ static_cast<int32_t*>(atomic_buffer.data_ptr()), num_experts);
116
+ } else {
117
+ compute_expert_offsets<<<1, 1, 0, stream>>>(
118
+ static_cast<const int32_t*>(problem_sizes1.data_ptr()),
119
+ static_cast<int32_t*>(expert_offsets.data_ptr()),
120
+ static_cast<int32_t*>(atomic_buffer.data_ptr()), num_experts);
121
+ }
122
compute_arg_sorts<<<num_experts, num_threads, 0, stream>>>(
123
static_cast<const int32_t*>(topk_ids.data_ptr()),
124
static_cast<const int32_t*>(expert_offsets.data_ptr()),
csrc/quantization/cutlass_w8a8/scaled_mm_entry.cu
+6
-3
@@ -54,7 +54,8 @@ void get_cutlass_moe_mm_data_caller(
54
const torch::Tensor& topk_ids, torch::Tensor& expert_offsets,
55
torch::Tensor& problem_sizes1, torch::Tensor& problem_sizes2,
56
torch::Tensor& input_permutation, torch::Tensor& output_permutation,
57
- const int64_t num_experts, const int64_t n, const int64_t k);
57
+ const int64_t num_experts, const int64_t n, const int64_t k,
58
+ const std::optional<torch::Tensor>& blockscale_offsets);
59
#endif
60
61
void cutlass_scaled_mm_azp_sm75(torch::Tensor& c, torch::Tensor const& a,
@@ -224,7 +225,8 @@ void get_cutlass_moe_mm_data(
225
const torch::Tensor& topk_ids, torch::Tensor& expert_offsets,
226
torch::Tensor& problem_sizes1, torch::Tensor& problem_sizes2,
227
torch::Tensor& input_permutation, torch::Tensor& output_permutation,
227
- const int64_t num_experts, const int64_t n, const int64_t k) {
228
+ const int64_t num_experts, const int64_t n, const int64_t k,
229
+ const std::optional<torch::Tensor>& blockscale_offsets) {
230
// This function currently gets compiled only if we have a valid cutlass moe
231
// mm to run it for.
232
int32_t version_num = get_sm_version_num();
@@ -232,7 +234,8 @@ void get_cutlass_moe_mm_data(
234
(defined ENABLE_SCALED_MM_SM100 && ENABLE_SCALED_MM_SM90)
235
get_cutlass_moe_mm_data_caller(topk_ids, expert_offsets, problem_sizes1,
236
problem_sizes2, input_permutation,
235
- output_permutation, num_experts, n, k);
237
+ output_permutation, num_experts, n, k,
238
+ blockscale_offsets);
239
return;
240
#endif
241
TORCH_CHECK_NOT_IMPLEMENTED(
csrc/torch_bindings.cpp
+1
-1
@@ -450,7 +450,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
450
" Tensor! problem_sizes1, Tensor! problem_sizes2, "
451
" Tensor! input_permutation, "
452
" Tensor! output_permutation, int num_experts, "
453
- " int n, int k) -> ()",
453
+ " int n, int k, Tensor? blockscale_offsets) -> ()",
454
{stride_tag});
455
ops.impl("get_cutlass_moe_mm_data", torch::kCUDA, &get_cutlass_moe_mm_data);
456
tests/kernels/moe/test_nvfp4_moe.py
+4
-1
@@ -80,7 +80,10 @@ def test_cutlass_fp4_moe_no_graph(m: int, n: int, k: int, e: int, topk: int,
80
w2[expert], w2_gs[expert])
81
82
score = torch.randn((m, e), device="cuda", dtype=dtype)
83
- topk_weights, topk_ids = fused_topk(a, score, topk, renormalize=False)
83
+ topk_weights, topk_ids, _ = fused_topk(a,
84
+ score,
85
+ topk,
86
+ renormalize=False)
87
88
a1_gs = torch.ones((e, ), device="cuda", dtype=torch.float32)
89
a2_gs = torch.ones((e, ), device="cuda", dtype=torch.float32)
vllm/_custom_ops.py
+33
-12
@@ -845,11 +845,16 @@ def cutlass_scaled_sparse_mm(
845
return out
846
847
848
-def get_cutlass_moe_mm_data(
849
- topk_ids: torch.Tensor, expert_offsets: torch.Tensor,
850
- problem_sizes1: torch.Tensor, problem_sizes2: torch.Tensor,
851
- input_permutation: torch.Tensor, output_permutation: torch.Tensor,
852
- num_experts: int, n: int, k: int):
848
+def get_cutlass_moe_mm_data(topk_ids: torch.Tensor,
849
+ expert_offsets: torch.Tensor,
850
+ problem_sizes1: torch.Tensor,
851
+ problem_sizes2: torch.Tensor,
852
+ input_permutation: torch.Tensor,
853
+ output_permutation: torch.Tensor,
854
+ num_experts: int,
855
+ n: int,
856
+ k: int,
857
+ blockscale_offsets: Optional[torch.Tensor] = None):
858
"""
859
Prepare data necessary to perform CUTLASS grouped matrix multiplications
860
used in CUTLASS-based fused MoE.
@@ -867,12 +872,31 @@ def get_cutlass_moe_mm_data(
872
before executing the MMs.
873
- output_permutation: Permutation that must be used to shuffle the output
874
after executing the MMs.
875
+ - blockscale_offsets: Optional argument passed for fp4 moe. Indices that
876
+ mark at which block scale index each expert begins
877
+ its computation. The number of block scale rows
878
+ computed with expert E is blockscale_offsets[E + 1] -
879
+ blockscale_offsets[E]
880
"""
881
return torch.ops._C.get_cutlass_moe_mm_data(topk_ids, expert_offsets,
882
problem_sizes1, problem_sizes2,
883
input_permutation,
884
output_permutation,
875
- num_experts, n, k)
885
+ num_experts, n, k,
886
+ blockscale_offsets)
887
+
888
+
889
+def shuffle_rows(input_tensor: torch.Tensor, dst2src_map: torch.Tensor):
890
+ """
891
+ Shuffle and expand the input tensor according to the dst2src_map and store the result in output_tensor.
892
+ This is used in MoE to permute the input tensor before performing grouped matrix multiplications.
893
+ """
894
+ num_tokens_permuted = dst2src_map.shape[0]
895
+ output_tensor = torch.empty((num_tokens_permuted, input_tensor.shape[1]),
896
+ device=input_tensor.device,
897
+ dtype=input_tensor.dtype)
898
+ torch.ops._moe_C.shuffle_rows(input_tensor, dst2src_map, output_tensor)
899
+ return output_tensor
900
901
902
def cutlass_moe_mm(out_tensors: torch.Tensor, a_tensors: torch.Tensor,
@@ -1124,14 +1148,12 @@ def scaled_fp4_experts_quant(
1148
expert_offsets: torch.Tensor,
1149
blockscale_offsets: torch.Tensor,
1150
topk: int,
1127
- expert_map: Optional[torch.Tensor] = None,
1151
) -> tuple[torch.Tensor, torch.Tensor]:
1152
"""
1153
Quantize input tensor to FP4 and return quantized tensor and scale, for
1154
packed MoE Inputs.
1155
Args:
1133
- input: The input tensor to be quantized to FP4
1134
- expert_map: The expert map tensor
1156
+ input_tensor: The input tensor to be quantized to FP4
1157
input_global_scale: A scalar scaling factor for the entire tensor.
1158
expert_offsets: The expert offsets tensor
1159
blockscale_offsets: The blockscale offsets tensor
@@ -1143,14 +1165,13 @@ def scaled_fp4_experts_quant(
1165
assert input_tensor.ndim == 2, (
1166
f'input.ndim needs to be == 2, but got {input_tensor.ndim}.')
1167
1146
- input_tensor = input_tensor[
1147
- expert_map] if expert_map is not None else input_tensor
1148
- m_numtopk, k = input_tensor.shape
1168
# Control the maximum number of tokens per expert supported by the
1169
# NVFP4 MoE Expert Quantization. This is used to prevent the kernel
1170
# from running out of memory. This value can also be increased to support
1171
# larger models.
1172
MAX_TOKENS_PER_EXPERT = envs.VLLM_MAX_TOKENS_PER_EXPERT_FP4_MOE
1173
+ m_numtopk, k = input_tensor.shape
1174
+
1175
assert (m_numtopk <= MAX_TOKENS_PER_EXPERT * topk), (
1176
f"m_numtopk must be less than MAX_TOKENS_PER_EXPERT("
1177
f"{MAX_TOKENS_PER_EXPERT})"
vllm/model_executor/layers/fused_moe/cutlass_moe.py
+8
-7
@@ -333,6 +333,7 @@ def cutlass_moe_fp4(a: torch.Tensor, a1_gscale: torch.Tensor,
333
num_topk = topk_ids.shape[1]
334
335
expert_offsets = torch.empty((e + 1), dtype=torch.int32, device=device)
336
+ blockscale_offsets = torch.empty((e + 1), dtype=torch.int32, device=device)
337
# Problem size: (num_experts, (m,2n,k))
338
problem_sizes1 = torch.empty((e, 3), dtype=torch.int32, device=device)
339
# Problem size: (num_experts, (m,n,k))
@@ -344,12 +345,10 @@ def cutlass_moe_fp4(a: torch.Tensor, a1_gscale: torch.Tensor,
345
# problem shapes should have [m, n, k]
346
# Note that problem sizes are based on logical number of elements.
347
ops.get_cutlass_moe_mm_data(topk_ids, expert_offsets, problem_sizes1,
347
- problem_sizes2, a_map, c_map, e, n, k)
348
+ problem_sizes2, a_map, c_map, e, n, k,
349
+ blockscale_offsets)
350
349
- tokens_per_expert = problem_sizes1[:, 0]
350
- rounded_tokens_per_expert = (tokens_per_expert + (128 - 1)) // 128 * 128
351
- blockscale_offsets = torch.zeros(e + 1, dtype=torch.int32, device=device)
352
- blockscale_offsets[1:] = torch.cumsum(rounded_tokens_per_expert, dim=0)
351
+ a = ops.shuffle_rows(a, a_map)
352
353
rep_a_fp4, rep_a_blockscale = ops.scaled_fp4_experts_quant(
354
a,
@@ -357,7 +356,7 @@ def cutlass_moe_fp4(a: torch.Tensor, a1_gscale: torch.Tensor,
356
expert_offsets,
357
blockscale_offsets,
358
num_topk,
360
- expert_map=a_map)
359
+ )
360
361
c1 = ops.cutlass_fp4_moe_mm(rep_a_fp4, w1_fp4, rep_a_blockscale,
362
w1_blockscale, w1_alphas, problem_sizes1,
@@ -378,6 +377,8 @@ def cutlass_moe_fp4(a: torch.Tensor, a1_gscale: torch.Tensor,
377
w2_alphas, problem_sizes2, expert_offsets[:-1],
378
blockscale_offsets[:-1], out_dtype, device)
379
del int_fp4, int_blockscale
381
- out = (c2[c_map].view(m, num_topk, k) *
380
+
381
+ c2 = ops.shuffle_rows(c2, c_map)
382
+ out = (c2.view(m, num_topk, k) *
383
topk_weights.view(m, num_topk, 1).half()).sum(dim=1)
384
return out.to(dtype=out_dtype)