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 2, 14:45 UTC
[Perf] Optimize moe permute by pre-allocate buffer, 9~14% kernel performance improvement (#43014)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
May 28, 2026 at 09:18 UTC
64e121867382d0d8f34b10fce194492be9b36745
8 files changed
+446
-58
benchmarks/kernels/benchmark_moe_permute_unpermute.py
+21
@@ -10,6 +10,7 @@ from transformers import AutoConfig
10
11
from vllm.model_executor.layers.fused_moe import fused_topk
12
from vllm.model_executor.layers.fused_moe.moe_permute_unpermute import (
13
+ MoEPermuteScratch,
14
moe_permute,
15
moe_unpermute,
16
)
@@ -54,6 +55,15 @@ def benchmark_permute(
55
topk_weights, topk_ids, token_expert_indices = fused_topk(
56
qhidden_states, input_gating, topk, False
57
)
58
+ scratch = MoEPermuteScratch(
59
+ max_num_tokens=num_tokens,
60
+ topk=topk,
61
+ num_experts=num_experts,
62
+ num_local_experts=num_experts,
63
+ device=qhidden_states.device,
64
+ hidden_size=hidden_size,
65
+ hidden_dtype=qhidden_states.dtype,
66
+ )
67
68
def prepare(i: int):
69
input_gating.copy_(gating_output[i])
@@ -65,6 +75,7 @@ def benchmark_permute(
75
topk_ids=topk_ids,
76
n_expert=num_experts,
77
expert_map=None,
78
+ scratch=scratch,
79
)
80
81
# JIT compilation & warmup
@@ -123,6 +134,15 @@ def benchmark_unpermute(
134
topk_weights, topk_ids, token_expert_indices = fused_topk(
135
qhidden_states, input_gating, topk, False
136
)
137
+ scratch = MoEPermuteScratch(
138
+ max_num_tokens=num_tokens,
139
+ topk=topk,
140
+ num_experts=num_experts,
141
+ num_local_experts=num_experts,
142
+ device=qhidden_states.device,
143
+ hidden_size=hidden_size,
144
+ hidden_dtype=qhidden_states.dtype,
145
+ )
146
147
def prepare():
148
(
@@ -137,6 +157,7 @@ def benchmark_unpermute(
157
topk_ids=topk_ids,
158
n_expert=num_experts,
159
expert_map=None,
160
+ scratch=scratch,
161
)
162
# convert to fp16/bf16 as gemm output
163
return (
csrc/moe/moe_ops.h
+3
@@ -62,6 +62,9 @@ std::tuple<torch::Tensor, torch::Tensor> grouped_topk(
62
63
bool moe_permute_unpermute_supported();
64
65
+int64_t moe_permute_sort_workspace_size(int64_t num_expanded_rows,
66
+ int64_t num_experts);
67
+
68
void shuffle_rows(const torch::Tensor& input_tensor,
69
const torch::Tensor& dst2src_map,
70
torch::Tensor& output_tensor);
csrc/moe/moe_permute_unpermute_op.cu
+107
-24
@@ -8,7 +8,35 @@
8
// moe_permute kernels require at least CUDA 12.0
9
#if defined(CUDA_VERSION) && (CUDA_VERSION >= 12000)
10
11
-void moe_permute(
11
+namespace {
12
+
13
+torch::Tensor maybe_allocate_tensor(
14
+ const std::optional<torch::Tensor>& maybe_tensor,
15
+ at::IntArrayRef expected_sizes, torch::ScalarType dtype, c10::Device device,
16
+ char const* name) {
17
+ auto expected_numel = c10::multiply_integers(expected_sizes);
18
+ if (maybe_tensor.has_value()) {
19
+ auto tensor = maybe_tensor.value();
20
+ TORCH_CHECK(tensor.device() == device, name, " must be on the same device");
21
+ TORCH_CHECK(tensor.scalar_type() == dtype, name, " has incorrect dtype");
22
+ TORCH_CHECK(tensor.is_contiguous(), name, " must be contiguous");
23
+ TORCH_CHECK(tensor.numel() >= expected_numel, name,
24
+ " is too small for the requested shape");
25
+ auto flat_tensor = tensor.view({tensor.numel()});
26
+ return flat_tensor.narrow(0, 0, expected_numel).view(expected_sizes);
27
+ }
28
+ return torch::empty(expected_sizes, torch::dtype(dtype).device(device));
29
+}
30
+
31
+} // namespace
32
+
33
+int64_t moe_permute_sort_workspace_size(int64_t num_expanded_rows,
34
+ int64_t n_expert) {
35
+ return static_cast<int64_t>(
36
+ CubKeyValueSorter::getWorkspaceSize(num_expanded_rows, n_expert));
37
+}
38
+
39
+void moe_permute_impl(
40
const torch::Tensor& input, // [n_token, hidden]
41
const torch::Tensor& topk_ids, // [n_token, topk]
42
const torch::Tensor& token_expert_indices, // [n_token, topk]
@@ -17,7 +45,11 @@ void moe_permute(
45
torch::Tensor& permuted_input, // [permuted_size, hidden]
46
torch::Tensor& expert_first_token_offset, // [n_local_expert + 1]
47
torch::Tensor& inv_permuted_idx, // [n_token, topk]
20
- torch::Tensor& permuted_idx) { // [permute_size]
48
+ torch::Tensor& permuted_idx, // [permute_size]
49
+ const std::optional<torch::Tensor>& maybe_sort_workspace,
50
+ const std::optional<torch::Tensor>& maybe_permuted_experts_id,
51
+ const std::optional<torch::Tensor>& maybe_sorted_row_idx,
52
+ const std::optional<torch::Tensor>& maybe_topk_ids_for_sort) {
53
TORCH_CHECK(expert_first_token_offset.scalar_type() == at::ScalarType::Long,
54
"expert_first_token_offset must be int64");
55
TORCH_CHECK(topk_ids.scalar_type() == at::ScalarType::Int,
@@ -27,49 +59,48 @@ void moe_permute(
59
TORCH_CHECK(inv_permuted_idx.scalar_type() == at::ScalarType::Int,
60
"inv_permuted_idx must be int32");
61
TORCH_CHECK(expert_first_token_offset.size(0) == n_local_expert + 1,
30
- "expert_first_token_offset shape != n_local_expert+1")
62
+ "expert_first_token_offset shape != n_local_expert+1");
63
TORCH_CHECK(inv_permuted_idx.sizes() == token_expert_indices.sizes(),
64
"token_expert_indices shape must be same as inv_permuted_idx");
65
+ auto device = input.device();
66
auto n_token = input.sizes()[0];
67
auto n_hidden = input.sizes()[1];
68
+ auto expanded_rows = n_token * topk;
69
auto stream = at::cuda::getCurrentCUDAStream().stream();
36
- const long sorter_size =
37
- CubKeyValueSorter::getWorkspaceSize(n_token * topk, n_expert);
38
- auto sort_workspace = torch::empty(
39
- {sorter_size},
40
- torch::dtype(torch::kInt8).device(torch::kCUDA).requires_grad(false));
41
- torch::Tensor topk_ids_for_sort = topk_ids;
42
- auto permuted_experts_id = torch::empty_like(topk_ids);
43
- auto sorted_row_idx = torch::empty_like(inv_permuted_idx);
70
+
71
+ auto sorter_size = moe_permute_sort_workspace_size(expanded_rows, n_expert);
72
+ auto sort_workspace =
73
+ maybe_allocate_tensor(maybe_sort_workspace, {sorter_size}, torch::kInt8,
74
+ device, "sort_workspace");
75
+ auto permuted_experts_id =
76
+ maybe_allocate_tensor(maybe_permuted_experts_id, topk_ids.sizes(),
77
+ at::ScalarType::Int, device, "permuted_experts_id");
78
+ auto sorted_row_idx =
79
+ maybe_allocate_tensor(maybe_sorted_row_idx, inv_permuted_idx.sizes(),
80
+ at::ScalarType::Int, device, "sorted_row_idx");
81
82
CubKeyValueSorter sorter{};
83
int64_t* valid_num_ptr = nullptr;
47
- // pre-process kernel for expert-parallelism:
48
- // no local expert id plus "n_expert" offset for priority to local expert
49
- // map local expert id [n, .., n+n_local_expert-1] to [0, n_local_expert -1]
50
- // For example, 4 expert with ep_size=2. ep_rank=1 owns global expert id
51
- // [2,3] with expert_map[-1, -1, 0, 1], preprocess_topk_id process topk_ids
52
- // and map global expert id [2, 3] to local_expert id [0, 1] and map global
53
- // expert id [0, 1] ( not in ep rank=1) to [4, 5] by plus n_expert. This map
54
- // operation is to make local expert high priority in following sort topk_ids
55
- // and scan local expert_first_token_offset for each ep rank for next group
56
- // gemm.
84
+ torch::Tensor topk_ids_for_sort = topk_ids;
85
+
86
if (expert_map.has_value()) {
87
const int* expert_map_ptr = get_ptr<int>(expert_map.value());
88
valid_num_ptr =
89
get_ptr<int64_t>(expert_first_token_offset) + n_local_expert;
61
- topk_ids_for_sort = topk_ids.clone();
90
+ topk_ids_for_sort =
91
+ maybe_allocate_tensor(maybe_topk_ids_for_sort, topk_ids.sizes(),
92
+ at::ScalarType::Int, device, "topk_ids_for_sort");
93
+ topk_ids_for_sort.copy_(topk_ids);
94
preprocessTopkIdLauncher(get_ptr<int>(topk_ids_for_sort), n_token * topk,
95
expert_map_ptr, n_expert, stream);
96
}
65
- // expert sort topk expert id and scan expert id get expert_first_token_offset
97
+
98
sortAndScanExpert(
99
get_ptr<const int>(topk_ids_for_sort), get_ptr<int>(token_expert_indices),
100
get_ptr<int>(permuted_experts_id), get_ptr<int>(sorted_row_idx),
101
get_ptr<int64_t>(expert_first_token_offset), n_token, n_expert,
102
n_local_expert, topk, sorter, get_ptr<int>(sort_workspace), stream);
103
72
- // dispatch expandInputRowsKernelLauncher
104
MOE_DISPATCH(input.scalar_type(), [&] {
105
expandInputRowsKernelLauncher<scalar_t>(
106
get_ptr<scalar_t>(input), get_ptr<scalar_t>(permuted_input),
@@ -79,6 +110,38 @@ void moe_permute(
110
});
111
}
112
113
+void moe_permute(
114
+ const torch::Tensor& input, // [n_token, hidden]
115
+ const torch::Tensor& topk_ids, // [n_token, topk]
116
+ const torch::Tensor& token_expert_indices, // [n_token, topk]
117
+ const std::optional<torch::Tensor>& expert_map, // [n_expert]
118
+ int64_t n_expert, int64_t n_local_expert, int64_t topk,
119
+ torch::Tensor& permuted_input, // [permuted_size, hidden]
120
+ torch::Tensor& expert_first_token_offset, // [n_local_expert + 1]
121
+ torch::Tensor& inv_permuted_idx, // [n_token, topk]
122
+ torch::Tensor& permuted_idx) { // [permute_size]
123
+ moe_permute_impl(input, topk_ids, token_expert_indices, expert_map, n_expert,
124
+ n_local_expert, topk, permuted_input,
125
+ expert_first_token_offset, inv_permuted_idx, permuted_idx,
126
+ std::nullopt, std::nullopt, std::nullopt, std::nullopt);
127
+}
128
+
129
+void moe_permute_with_scratch(
130
+ const torch::Tensor& input, const torch::Tensor& topk_ids,
131
+ const torch::Tensor& token_expert_indices,
132
+ const std::optional<torch::Tensor>& expert_map, int64_t n_expert,
133
+ int64_t n_local_expert, int64_t topk, torch::Tensor& permuted_input,
134
+ torch::Tensor& expert_first_token_offset, torch::Tensor& inv_permuted_idx,
135
+ torch::Tensor& permuted_idx, torch::Tensor& sort_workspace,
136
+ torch::Tensor& permuted_experts_id, torch::Tensor& sorted_row_idx,
137
+ torch::Tensor& topk_ids_for_sort) {
138
+ moe_permute_impl(input, topk_ids, token_expert_indices, expert_map, n_expert,
139
+ n_local_expert, topk, permuted_input,
140
+ expert_first_token_offset, inv_permuted_idx, permuted_idx,
141
+ sort_workspace, permuted_experts_id, sorted_row_idx,
142
+ topk_ids_for_sort);
143
+}
144
+
145
void moe_unpermute(
146
const torch::Tensor& permuted_hidden_states, // [n_token * topk, hidden]
147
const torch::Tensor& topk_weights, // [n_token, topk]
@@ -169,6 +232,12 @@ void shuffle_rows(const torch::Tensor& input_tensor,
232
233
#else
234
235
+int64_t moe_permute_sort_workspace_size(int64_t num_expanded_rows,
236
+ int64_t n_expert) {
237
+ TORCH_CHECK(
238
+ false, "moe_permute_sort_workspace_size is not supported on CUDA < 12.0");
239
+}
240
+
241
void moe_permute(const torch::Tensor& input, const torch::Tensor& topk_ids,
242
const torch::Tensor& token_expert_indices,
243
const std::optional<torch::Tensor>& expert_map,
@@ -179,6 +248,19 @@ void moe_permute(const torch::Tensor& input, const torch::Tensor& topk_ids,
248
TORCH_CHECK(false, "moe_permute is not supported on CUDA < 12.0");
249
}
250
251
+void moe_permute_with_scratch(
252
+ const torch::Tensor& input, const torch::Tensor& topk_ids,
253
+ const torch::Tensor& token_expert_indices,
254
+ const std::optional<torch::Tensor>& expert_map, int64_t n_expert,
255
+ int64_t n_local_expert, int64_t topk, torch::Tensor& permuted_input,
256
+ torch::Tensor& expert_first_token_offset, torch::Tensor& inv_permuted_idx,
257
+ torch::Tensor& permuted_idx, torch::Tensor& sort_workspace,
258
+ torch::Tensor& permuted_experts_id, torch::Tensor& sorted_row_idx,
259
+ torch::Tensor& topk_ids_for_sort) {
260
+ TORCH_CHECK(false,
261
+ "moe_permute_with_scratch is not supported on CUDA < 12.0");
262
+}
263
+
264
void moe_unpermute(
265
const torch::Tensor& permuted_hidden_states,
266
const torch::Tensor& topk_weights, const torch::Tensor& inv_permuted_idx,
@@ -199,5 +281,6 @@ bool moe_permute_unpermute_supported() {
281
282
TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) {
283
m.impl("moe_permute", &moe_permute);
284
+ m.impl("moe_permute_with_scratch", &moe_permute_with_scratch);
285
m.impl("moe_unpermute", &moe_unpermute);
286
}
\ No newline at end of file
csrc/moe/torch_bindings.cpp
+13
@@ -100,13 +100,26 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, m) {
100
"expert_first_token_offset, Tensor! inv_permuted_idx, Tensor! "
101
"permuted_idx)->()");
102
103
+ m.def(
104
+ "moe_permute_with_scratch(Tensor input, Tensor topk_ids,"
105
+ "Tensor token_expert_indices, Tensor? expert_map, int n_expert,"
106
+ "int n_local_expert,"
107
+ "int topk, Tensor! permuted_input, Tensor! "
108
+ "expert_first_token_offset, Tensor! inv_permuted_idx, Tensor! "
109
+ "permuted_idx, Tensor! sort_workspace, Tensor! permuted_experts_id, "
110
+ "Tensor! sorted_row_idx, Tensor! topk_ids_for_sort)->()");
111
+
112
m.def(
113
"moe_unpermute(Tensor permuted_hidden_states, Tensor topk_weights,"
114
"Tensor inv_permuted_idx, Tensor? expert_first_token_offset, "
115
"int topk, Tensor! hidden_states)->()");
116
117
m.def("moe_permute_unpermute_supported() -> bool");
118
+ m.def(
119
+ "moe_permute_sort_workspace_size(int num_expanded_rows, int n_expert) -> "
120
+ "int");
121
m.impl("moe_permute_unpermute_supported", &moe_permute_unpermute_supported);
122
+ m.impl("moe_permute_sort_workspace_size", &moe_permute_sort_workspace_size);
123
124
// Row shuffle for MoE
125
m.def(
tests/kernels/moe/test_moe_permute_unpermute.py
+77
@@ -14,6 +14,7 @@ from vllm.model_executor.layers.fused_moe.expert_map_manager import (
14
determine_expert_map,
15
)
16
from vllm.model_executor.layers.fused_moe.moe_permute_unpermute import (
17
+ MoEPermuteScratch,
18
moe_permute,
19
moe_permute_unpermute_supported,
20
moe_unpermute,
@@ -209,3 +210,79 @@ def test_moe_permute_unpermute(
210
)
211
# check unpermuted hidden
212
torch.testing.assert_close(result4, gold4, atol=2e-2, rtol=0)
213
+
214
+
215
+@pytest.mark.parametrize("dtype", [torch.bfloat16])
216
+def test_moe_permute_reuses_scratch_buffers(dtype: torch.dtype):
217
+ if not moe_permute_unpermute_supported():
218
+ pytest.skip("moe_permute_unpermute is not supported on this platform.")
219
+
220
+ n_token = 64
221
+ n_hidden = 2048
222
+ n_expert = 16
223
+ topk = 4
224
+ hidden_states = torch.randn((n_token, n_hidden), device="cuda").to(dtype)
225
+ gating_output = torch.randn((n_token, n_expert), device="cuda").to(dtype)
226
+ _, topk_ids, _ = fused_topk(hidden_states, gating_output, topk, False)
227
+
228
+ scratch = MoEPermuteScratch(
229
+ max_num_tokens=n_token,
230
+ topk=topk,
231
+ num_experts=n_expert,
232
+ num_local_experts=n_expert,
233
+ device=hidden_states.device,
234
+ hidden_size=n_hidden,
235
+ hidden_dtype=hidden_states.dtype,
236
+ )
237
+
238
+ first = moe_permute(
239
+ hidden_states=hidden_states,
240
+ a1q_scale=None,
241
+ topk_ids=topk_ids,
242
+ n_expert=n_expert,
243
+ scratch=scratch,
244
+ )
245
+ second = moe_permute(
246
+ hidden_states=hidden_states,
247
+ a1q_scale=None,
248
+ topk_ids=topk_ids,
249
+ n_expert=n_expert,
250
+ scratch=scratch,
251
+ )
252
+
253
+ (
254
+ permuted_hidden_states_1,
255
+ _,
256
+ expert_first_token_offset_1,
257
+ inv_permuted_idx_1,
258
+ permuted_idx_1,
259
+ ) = first
260
+ (
261
+ permuted_hidden_states_2,
262
+ _,
263
+ expert_first_token_offset_2,
264
+ inv_permuted_idx_2,
265
+ permuted_idx_2,
266
+ ) = second
267
+
268
+ torch.testing.assert_close(permuted_hidden_states_1, permuted_hidden_states_2)
269
+ torch.testing.assert_close(expert_first_token_offset_1, expert_first_token_offset_2)
270
+ torch.testing.assert_close(inv_permuted_idx_1, inv_permuted_idx_2)
271
+ torch.testing.assert_close(permuted_idx_1, permuted_idx_2)
272
+
273
+ assert (
274
+ permuted_hidden_states_1.untyped_storage().data_ptr()
275
+ == permuted_hidden_states_2.untyped_storage().data_ptr()
276
+ )
277
+ assert (
278
+ expert_first_token_offset_1.untyped_storage().data_ptr()
279
+ == expert_first_token_offset_2.untyped_storage().data_ptr()
280
+ )
281
+ assert (
282
+ inv_permuted_idx_1.untyped_storage().data_ptr()
283
+ == scratch.inv_permuted_idx.untyped_storage().data_ptr()
284
+ )
285
+ assert (
286
+ permuted_idx_1.untyped_storage().data_ptr()
287
+ == scratch.permuted_idx.untyped_storage().data_ptr()
288
+ )
vllm/model_executor/layers/fused_moe/experts/cutlass_moe.py
+32
@@ -17,7 +17,9 @@ from vllm.model_executor.layers.fused_moe.config import (
17
FusedMoEQuantConfig,
18
)
19
from vllm.model_executor.layers.fused_moe.moe_permute_unpermute import (
20
+ MoEPermuteScratch,
21
moe_permute,
22
+ moe_permute_unpermute_supported,
23
moe_unpermute,
24
)
25
from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
@@ -73,6 +75,7 @@ def run_cutlass_moe_fp8(
75
per_out_ch: bool,
76
use_batched_format: bool,
77
topk_weights: torch.Tensor | None,
78
+ permute_scratch: MoEPermuteScratch | None,
79
):
80
a1q = hidden_states
81
@@ -198,6 +201,7 @@ def run_cutlass_moe_fp8(
201
local_E,
202
expert_map,
203
permuted_hidden_states=a1q_perm,
204
+ scratch=permute_scratch,
205
)
206
# swap_ab is a CUTLASS grouped-GEMM optimization (M <= 64 reduces padding).
207
swap_ab = a1q.size(0) <= 64
@@ -291,6 +295,7 @@ class CutlassExpertsFp8Base(mk.FusedMoEExpertsModular):
295
self.ab_strides2 = ab_strides2
296
self.c_strides1 = c_strides1
297
self.c_strides2 = ab_strides1_c_strides2
298
+ self._permute_scratch: MoEPermuteScratch | None = None
299
300
@staticmethod
301
def _supports_current_device() -> bool:
@@ -324,6 +329,17 @@ class CutlassExpertsFp8Base(mk.FusedMoEExpertsModular):
329
# Let PrepareAndFinalize::finalize() decide the impl.
330
return TopKWeightAndReduceDelegate()
331
332
+ def _get_permute_scratch(self) -> MoEPermuteScratch | None:
333
+ if self._permute_scratch is None and moe_permute_unpermute_supported():
334
+ self._permute_scratch = MoEPermuteScratch(
335
+ max_num_tokens=self.moe_config.max_num_tokens,
336
+ topk=self.moe_config.experts_per_token,
337
+ num_experts=self.moe_config.num_experts,
338
+ num_local_experts=self.moe_config.num_local_experts,
339
+ device=torch.device(self.moe_config.device),
340
+ )
341
+ return self._permute_scratch
342
+
343
def apply(
344
self,
345
output: torch.Tensor,
@@ -379,6 +395,7 @@ class CutlassExpertsFp8Base(mk.FusedMoEExpertsModular):
395
self.per_out_ch_quant,
396
use_batched_format,
397
topk_weights,
398
+ self._get_permute_scratch(),
399
)
400
401
@@ -1121,6 +1138,7 @@ def run_cutlass_moe_w4a8_fp8(
1138
use_batched_format: bool,
1139
topk_weights: torch.Tensor | None,
1140
group_size: int,
1141
+ permute_scratch: MoEPermuteScratch | None,
1142
):
1143
a1q = hidden_states
1144
M = a1q.size(0)
@@ -1176,6 +1194,7 @@ def run_cutlass_moe_w4a8_fp8(
1194
local_E,
1195
expert_map,
1196
permuted_hidden_states=a1q_perm,
1197
+ scratch=permute_scratch,
1198
)
1199
# for RS gemm SwapAB is always enabled (swap logical M, N in the problem shape).
1200
ops.get_cutlass_moe_mm_problem_sizes_from_expert_offsets(
@@ -1266,6 +1285,7 @@ class CutlassExpertsW4A8Fp8(mk.FusedMoEExpertsModular):
1285
self.s_strides2[:, 0] = k
1286
1287
self.group_size = group_size
1288
+ self._permute_scratch: MoEPermuteScratch | None = None
1289
1290
@staticmethod
1291
def activation_format() -> mk.FusedMoEActivationFormat:
@@ -1330,6 +1350,17 @@ class CutlassExpertsW4A8Fp8(mk.FusedMoEExpertsModular):
1350
def workspace_dtype(self, act_dtype: torch.dtype) -> torch.dtype:
1351
return self.out_dtype if self.out_dtype is not None else act_dtype
1352
1353
+ def _get_permute_scratch(self) -> MoEPermuteScratch | None:
1354
+ if self._permute_scratch is None and moe_permute_unpermute_supported():
1355
+ self._permute_scratch = MoEPermuteScratch(
1356
+ max_num_tokens=self.moe_config.max_num_tokens,
1357
+ topk=self.moe_config.experts_per_token,
1358
+ num_experts=self.moe_config.num_experts,
1359
+ num_local_experts=self.moe_config.num_local_experts,
1360
+ device=torch.device(self.moe_config.device),
1361
+ )
1362
+ return self._permute_scratch
1363
+
1364
def workspace_shapes(
1365
self,
1366
M: int,
@@ -1409,4 +1440,5 @@ class CutlassExpertsW4A8Fp8(mk.FusedMoEExpertsModular):
1440
use_batched_format,
1441
topk_weights,
1442
self.group_size,
1443
+ self._get_permute_scratch(),
1444
)
vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py
+17
@@ -26,7 +26,9 @@ from vllm.model_executor.layers.fused_moe.moe_align_block_size import (
26
)
27
from vllm.model_executor.layers.fused_moe.moe_fused_mul_sum import moe_fused_mul_sum
28
from vllm.model_executor.layers.fused_moe.moe_permute_unpermute import (
29
+ MoEPermuteScratch,
30
moe_permute,
31
+ moe_permute_unpermute_supported,
32
moe_unpermute,
33
)
34
from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
@@ -85,6 +87,7 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular):
87
max_num_tokens=max_num_tokens,
88
num_dispatchers=num_dispatchers,
89
)
90
+ self._permute_scratch: MoEPermuteScratch | None = None
91
92
def init_humming_moe(self):
93
self.compute_config = {
@@ -110,6 +113,19 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular):
113
self.w13_tuning_config_str = json.dumps(self.w13_tuning_config)
114
self.w2_tuning_config_str = json.dumps(self.w2_tuning_config)
115
116
+ def _get_permute_scratch(self) -> MoEPermuteScratch | None:
117
+ if self._permute_scratch is None and moe_permute_unpermute_supported():
118
+ self._permute_scratch = MoEPermuteScratch(
119
+ max_num_tokens=self.moe_config.max_num_tokens,
120
+ topk=self.moe_config.experts_per_token,
121
+ num_experts=self.moe_config.num_experts,
122
+ num_local_experts=self.moe_config.num_local_experts,
123
+ device=torch.device(self.moe_config.device),
124
+ hidden_size=self.moe_config.hidden_dim,
125
+ hidden_dtype=self.moe_config.in_dtype,
126
+ )
127
+ return self._permute_scratch
128
+
129
def get_global_valid_shape_m(self, topk_ids: torch.Tensor):
130
num_tokens = topk_ids.size(0)
131
ctx = get_forward_context()
@@ -598,6 +614,7 @@ class HummingGroupedExperts(HummingExpertsBase):
614
n_expert=self.global_num_experts,
615
n_local_expert=self.num_experts,
616
expert_map=self.layer.expert_map,
617
+ scratch=self._get_permute_scratch(),
618
)
619
620
inputs, input_scale = HummingMethod.may_quant_input(
vllm/model_executor/layers/fused_moe/moe_permute_unpermute.py
+176
-34
@@ -1,9 +1,104 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
+from dataclasses import dataclass, field
5
+
6
import torch
7
8
9
+@dataclass
10
+class MoEPermuteScratch:
11
+ # Reused metadata buffers for repeated grouped-MoE permutes.
12
+ max_num_tokens: int
13
+ topk: int
14
+ num_experts: int
15
+ num_local_experts: int
16
+ device: torch.device
17
+ hidden_size: int | None = None
18
+ hidden_dtype: torch.dtype | None = None
19
+ token_expert_indices: torch.Tensor = field(init=False)
20
+ expert_first_token_offset: torch.Tensor = field(init=False)
21
+ permuted_idx: torch.Tensor = field(init=False)
22
+ inv_permuted_idx: torch.Tensor = field(init=False)
23
+ permuted_hidden_states: torch.Tensor | None = field(init=False, default=None)
24
+ sort_workspace: torch.Tensor = field(init=False)
25
+ permuted_experts_id: torch.Tensor = field(init=False)
26
+ sorted_row_idx: torch.Tensor = field(init=False)
27
+ topk_ids_int32: torch.Tensor = field(init=False)
28
+ topk_ids_for_sort: torch.Tensor = field(init=False)
29
+ max_expanded_rows: int = field(init=False)
30
+
31
+ def __post_init__(self) -> None:
32
+ assert self.max_num_tokens > 0
33
+ assert self.topk > 0
34
+ assert self.num_experts > 0
35
+ assert self.num_local_experts > 0
36
+ if self.hidden_size is None:
37
+ assert self.hidden_dtype is None
38
+ else:
39
+ assert self.hidden_dtype is not None
40
+
41
+ self.max_expanded_rows = self.max_num_tokens * self.topk
42
+ self.token_expert_indices = torch.arange(
43
+ self.max_expanded_rows, dtype=torch.int32, device=self.device
44
+ )
45
+ self.expert_first_token_offset = torch.empty(
46
+ self.num_local_experts + 1, dtype=torch.int64, device=self.device
47
+ )
48
+ self.permuted_idx = torch.empty(
49
+ self.max_expanded_rows, dtype=torch.int32, device=self.device
50
+ )
51
+ self.inv_permuted_idx = torch.empty(
52
+ self.max_expanded_rows, dtype=torch.int32, device=self.device
53
+ )
54
+ if self.hidden_size is not None:
55
+ hidden_numel = self.max_expanded_rows * self.hidden_size
56
+ self.permuted_hidden_states = torch.empty(
57
+ hidden_numel, dtype=self.hidden_dtype, device=self.device
58
+ )
59
+ self.permuted_experts_id = torch.empty(
60
+ self.max_expanded_rows, dtype=torch.int32, device=self.device
61
+ )
62
+ self.sorted_row_idx = torch.empty(
63
+ self.max_expanded_rows, dtype=torch.int32, device=self.device
64
+ )
65
+ self.topk_ids_int32 = torch.empty(
66
+ self.max_expanded_rows, dtype=torch.int32, device=self.device
67
+ )
68
+ self.topk_ids_for_sort = torch.empty(
69
+ self.max_expanded_rows, dtype=torch.int32, device=self.device
70
+ )
71
+ sorter_size = torch.ops._moe_C.moe_permute_sort_workspace_size(
72
+ self.max_expanded_rows, self.num_experts
73
+ )
74
+ self.sort_workspace = torch.empty(
75
+ sorter_size, dtype=torch.int8, device=self.device
76
+ )
77
+
78
+ def validate(self, hidden_states: torch.Tensor, topk_ids: torch.Tensor) -> None:
79
+ n_token, n_hidden = hidden_states.shape
80
+ assert hidden_states.device == self.device
81
+ assert topk_ids.device == self.device
82
+ assert n_token <= self.max_num_tokens
83
+ assert topk_ids.size(1) == self.topk
84
+ assert topk_ids.size(0) == n_token
85
+ if self.hidden_size is not None:
86
+ assert n_hidden == self.hidden_size
87
+ assert hidden_states.dtype == self.hidden_dtype
88
+ assert self.permuted_hidden_states is not None
89
+
90
+ def token_expert_indices_view(self, n_token: int) -> torch.Tensor:
91
+ return self.token_expert_indices[: n_token * self.topk].view(n_token, self.topk)
92
+
93
+ def prepare_topk_ids(self, topk_ids: torch.Tensor) -> torch.Tensor:
94
+ if topk_ids.dtype == torch.int32:
95
+ return topk_ids
96
+ numel = topk_ids.numel()
97
+ topk_ids_int32 = self.topk_ids_int32[:numel].view_as(topk_ids)
98
+ topk_ids_int32.copy_(topk_ids)
99
+ return topk_ids_int32
100
+
101
+
102
def moe_permute(
103
hidden_states: torch.Tensor,
104
a1q_scale: torch.Tensor | None,
@@ -12,6 +107,7 @@ def moe_permute(
107
n_local_expert: int = -1,
108
expert_map: torch.Tensor | None = None,
109
permuted_hidden_states: torch.Tensor | None = None,
110
+ scratch: MoEPermuteScratch | None = None,
111
) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor, torch.Tensor, torch.Tensor]:
112
"""
113
This function expands and permutes activation to gather uncontinuous tokens
@@ -45,46 +141,92 @@ def moe_permute(
141
if n_local_expert == -1:
142
n_local_expert = n_expert
143
if permuted_hidden_states is None:
48
- permuted_hidden_states = torch.empty(
49
- (permuted_row_size, n_hidden),
50
- dtype=hidden_states.dtype,
51
- device=hidden_states.device,
52
- )
144
+ if scratch is None:
145
+ permuted_hidden_states = torch.empty(
146
+ (permuted_row_size, n_hidden),
147
+ dtype=hidden_states.dtype,
148
+ device=hidden_states.device,
149
+ )
150
+ else:
151
+ scratch.validate(hidden_states, topk_ids)
152
+ hidden_numel = permuted_row_size * n_hidden
153
+ scratch_hidden_states = scratch.permuted_hidden_states
154
+ assert scratch_hidden_states is not None
155
+ permuted_hidden_states = scratch_hidden_states[:hidden_numel].view(
156
+ permuted_row_size, n_hidden
157
+ )
158
assert permuted_hidden_states.size() == (permuted_row_size, n_hidden), (
159
f"Expected permuted hidden states to be {(permuted_row_size, n_hidden)}"
160
f" but got {permuted_hidden_states.size()}"
161
)
162
58
- token_expert_indices = torch.arange(
59
- 0, n_token * topk, dtype=torch.int32, device=hidden_states.device
60
- ).reshape((n_token, topk))
163
+ if scratch is None:
164
+ token_expert_indices = torch.arange(
165
+ 0, n_token * topk, dtype=torch.int32, device=hidden_states.device
166
+ ).reshape((n_token, topk))
167
62
- expert_first_token_offset = torch.empty(
63
- n_local_expert + 1, dtype=torch.int64, device=hidden_states.device
64
- )
65
- permuted_idx = torch.full(
66
- (permuted_row_size,),
67
- n_token * topk,
68
- dtype=torch.int32,
69
- device=hidden_states.device,
70
- )
71
- inv_permuted_idx = torch.empty(
72
- (n_token, topk), dtype=torch.int32, device=hidden_states.device
73
- )
74
- topk_ids = topk_ids.to(torch.int32)
75
- torch.ops._moe_C.moe_permute(
76
- hidden_states,
77
- topk_ids,
78
- token_expert_indices,
79
- expert_map,
80
- n_expert,
81
- n_local_expert,
82
- topk,
83
- permuted_hidden_states,
84
- expert_first_token_offset,
85
- inv_permuted_idx,
86
- permuted_idx,
87
- )
168
+ expert_first_token_offset = torch.empty(
169
+ n_local_expert + 1, dtype=torch.int64, device=hidden_states.device
170
+ )
171
+ permuted_idx = torch.full(
172
+ (permuted_row_size,),
173
+ n_token * topk,
174
+ dtype=torch.int32,
175
+ device=hidden_states.device,
176
+ )
177
+ inv_permuted_idx = torch.empty(
178
+ (n_token, topk), dtype=torch.int32, device=hidden_states.device
179
+ )
180
+ topk_ids_int32 = topk_ids.to(torch.int32)
181
+ torch.ops._moe_C.moe_permute(
182
+ hidden_states,
183
+ topk_ids_int32,
184
+ token_expert_indices,
185
+ expert_map,
186
+ n_expert,
187
+ n_local_expert,
188
+ topk,
189
+ permuted_hidden_states,
190
+ expert_first_token_offset,
191
+ inv_permuted_idx,
192
+ permuted_idx,
193
+ )
194
+ else:
195
+ scratch.validate(hidden_states, topk_ids)
196
+ assert n_expert == scratch.num_experts
197
+ assert n_local_expert == scratch.num_local_experts
198
+ token_expert_indices = scratch.token_expert_indices_view(n_token)
199
+ expert_first_token_offset = scratch.expert_first_token_offset
200
+ permuted_idx = scratch.permuted_idx[:permuted_row_size]
201
+ permuted_idx.fill_(permuted_row_size)
202
+ inv_permuted_idx = scratch.inv_permuted_idx[:permuted_row_size].view(
203
+ n_token, topk
204
+ )
205
+ permuted_experts_id = scratch.permuted_experts_id[:permuted_row_size].view(
206
+ n_token, topk
207
+ )
208
+ sorted_row_idx = scratch.sorted_row_idx[:permuted_row_size].view(n_token, topk)
209
+ topk_ids_for_sort = scratch.topk_ids_for_sort[:permuted_row_size].view(
210
+ n_token, topk
211
+ )
212
+ topk_ids_int32 = scratch.prepare_topk_ids(topk_ids)
213
+ torch.ops._moe_C.moe_permute_with_scratch(
214
+ hidden_states,
215
+ topk_ids_int32,
216
+ token_expert_indices,
217
+ expert_map,
218
+ n_expert,
219
+ n_local_expert,
220
+ topk,
221
+ permuted_hidden_states,
222
+ expert_first_token_offset,
223
+ inv_permuted_idx,
224
+ permuted_idx,
225
+ scratch.sort_workspace,
226
+ permuted_experts_id,
227
+ sorted_row_idx,
228
+ topk_ids_for_sort,
229
+ )
230
231
if a1q_scale is not None and a1q_scale.dim() > 1:
232
a1q_scale = a1q_scale[permuted_idx.clamp(max=n_token * topk - 1) // topk]