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 9, 14:45 UTC
[Refactor] Remove unused `_moe_permute` function (#33108)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Jan 26, 2026 at 16:06 UTC
8f987883cbc0a562f253fda79783b0f22179e5c6
2 files changed
+38
-169
benchmarks/kernels/benchmark_moe_permute_unpermute.py
+38
-105
@@ -10,8 +10,6 @@ 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
- _moe_permute,
14
- _moe_unpermute_and_reduce,
13
moe_permute,
14
moe_unpermute,
15
)
@@ -41,7 +39,6 @@ def benchmark_permute(
39
use_fp8_w8a8: bool,
40
use_int8_w8a16: bool,
41
num_iters: int = 100,
44
- use_customized_permute: bool = False,
42
) -> float:
43
# init_dtype = torch.float16 if use_fp8_w8a8 else dtype
44
hidden_states = torch.randn(num_tokens, hidden_size, dtype=dtype)
@@ -64,29 +61,14 @@ def benchmark_permute(
61
input_gating.copy_(gating_output[i])
62
63
def run():
67
- if use_customized_permute:
68
- (
69
- permuted_hidden_states,
70
- a1q_scale,
71
- first_token_off,
72
- inv_perm_idx,
73
- m_indices,
74
- ) = moe_permute(
75
- qhidden_states,
76
- a1q_scale=None,
77
- topk_ids=topk_ids,
78
- n_expert=num_experts,
79
- expert_map=None,
80
- align_block_size=align_block_size,
81
- )
82
- else:
83
- (
84
- permuted_hidden_states,
85
- a1q_scale,
86
- sorted_token_ids,
87
- expert_ids,
88
- inv_perm,
89
- ) = _moe_permute(qhidden_states, None, topk_ids, num_experts, None, 16)
64
+ moe_permute(
65
+ qhidden_states,
66
+ a1q_scale=None,
67
+ topk_ids=topk_ids,
68
+ n_expert=num_experts,
69
+ expert_map=None,
70
+ align_block_size=align_block_size,
71
+ )
72
73
# JIT compilation & warmup
74
run()
@@ -131,11 +113,9 @@ def benchmark_unpermute(
113
use_fp8_w8a8: bool,
114
use_int8_w8a16: bool,
115
num_iters: int = 100,
134
- use_customized_permute: bool = False,
116
) -> float:
117
# init_dtype = torch.float16 if use_fp8_w8a8 else dtype
118
hidden_states = torch.randn(num_tokens, hidden_size, dtype=dtype)
138
- output_hidden_states = torch.empty_like(hidden_states)
119
if use_fp8_w8a8:
120
align_block_size = 128 # deepgemm needs 128 m aligned block
121
qhidden_states, scale = _fp8_quantize(hidden_states, None, None)
@@ -150,78 +130,37 @@ def benchmark_unpermute(
130
)
131
132
def prepare():
153
- if use_customized_permute:
154
- (
155
- permuted_hidden_states,
156
- a1q_scale,
157
- first_token_off,
158
- inv_perm_idx,
159
- m_indices,
160
- ) = moe_permute(
161
- qhidden_states,
162
- a1q_scale=None,
163
- topk_ids=topk_ids,
164
- n_expert=num_experts,
165
- expert_map=None,
166
- align_block_size=align_block_size,
167
- )
168
- # convert to fp16/bf16 as gemm output
169
- return (
170
- permuted_hidden_states.to(dtype),
171
- first_token_off,
172
- inv_perm_idx,
173
- m_indices,
174
- )
175
- else:
176
- (
177
- permuted_qhidden_states,
178
- a1q_scale,
179
- sorted_token_ids,
180
- expert_ids,
181
- inv_perm,
182
- ) = _moe_permute(
183
- qhidden_states, None, topk_ids, num_experts, None, block_m=16
184
- )
185
- # convert to fp16/bf16 as gemm output
186
- return (
187
- permuted_qhidden_states.to(dtype),
188
- a1q_scale,
189
- sorted_token_ids,
190
- expert_ids,
191
- inv_perm,
192
- )
133
+ (
134
+ permuted_hidden_states,
135
+ _,
136
+ first_token_off,
137
+ inv_perm_idx,
138
+ _,
139
+ ) = moe_permute(
140
+ qhidden_states,
141
+ a1q_scale=None,
142
+ topk_ids=topk_ids,
143
+ n_expert=num_experts,
144
+ expert_map=None,
145
+ align_block_size=align_block_size,
146
+ )
147
+ # convert to fp16/bf16 as gemm output
148
+ return (
149
+ permuted_hidden_states.to(dtype),
150
+ first_token_off,
151
+ inv_perm_idx,
152
+ )
153
154
def run(input: tuple):
195
- if use_customized_permute:
196
- (
197
- permuted_hidden_states,
198
- first_token_off,
199
- inv_perm_idx,
200
- m_indices,
201
- ) = input
202
- output = torch.empty_like(hidden_states)
203
- moe_unpermute(
204
- output,
205
- permuted_hidden_states,
206
- topk_weights,
207
- inv_perm_idx,
208
- first_token_off,
209
- )
210
- else:
211
- (
212
- permuted_hidden_states,
213
- a1q_scale,
214
- sorted_token_ids,
215
- expert_ids,
216
- inv_perm,
217
- ) = input
218
- _moe_unpermute_and_reduce(
219
- output_hidden_states,
220
- permuted_hidden_states,
221
- inv_perm,
222
- topk_weights,
223
- True,
224
- )
155
+ (permuted_hidden_states, first_token_off, inv_perm_idx) = input
156
+ output = torch.empty_like(hidden_states)
157
+ moe_unpermute(
158
+ output,
159
+ permuted_hidden_states,
160
+ topk_weights,
161
+ inv_perm_idx,
162
+ first_token_off,
163
+ )
164
165
# JIT compilation & warmup
166
input = prepare()
@@ -276,8 +215,7 @@ class BenchmarkWorker:
215
dtype: torch.dtype,
216
use_fp8_w8a8: bool,
217
use_int8_w8a16: bool,
279
- use_customized_permute: bool = False,
280
- ) -> tuple[dict[str, int], float]:
218
+ ) -> tuple[float, float]:
219
set_random_seed(self.seed)
220
221
permute_time = benchmark_permute(
@@ -289,7 +227,6 @@ class BenchmarkWorker:
227
use_fp8_w8a8,
228
use_int8_w8a16,
229
num_iters=100,
292
- use_customized_permute=use_customized_permute,
230
)
231
unpermute_time = benchmark_unpermute(
232
num_tokens,
@@ -300,7 +237,6 @@ class BenchmarkWorker:
237
use_fp8_w8a8,
238
use_int8_w8a16,
239
num_iters=100,
303
- use_customized_permute=use_customized_permute,
240
)
241
return permute_time, unpermute_time
242
@@ -347,7 +283,6 @@ def main(args: argparse.Namespace):
283
dtype = torch.float16 if current_platform.is_rocm() else config.dtype
284
use_fp8_w8a8 = args.dtype == "fp8_w8a8"
285
use_int8_w8a16 = args.dtype == "int8_w8a16"
350
- use_customized_permute = args.use_customized_permute
286
287
if args.batch_size is None:
288
batch_sizes = [
@@ -399,7 +334,6 @@ def main(args: argparse.Namespace):
334
dtype,
335
use_fp8_w8a8,
336
use_int8_w8a16,
402
- use_customized_permute,
337
)
338
for batch_size in batch_sizes
339
],
@@ -419,7 +353,6 @@ if __name__ == "__main__":
353
parser.add_argument(
354
"--dtype", type=str, choices=["auto", "fp8_w8a8", "int8_w8a16"], default="auto"
355
)
422
- parser.add_argument("--use-customized-permute", action="store_true")
356
parser.add_argument("--seed", type=int, default=0)
357
parser.add_argument("--batch-size", type=int, required=False)
358
parser.add_argument("--trust-remote-code", action="store_true")
vllm/model_executor/layers/fused_moe/moe_permute_unpermute.py
-64
@@ -3,70 +3,6 @@
3
4
import torch
5
6
-from vllm import _custom_ops as ops
7
-from vllm.model_executor.layers.fused_moe.moe_align_block_size import (
8
- moe_align_block_size,
9
-)
10
-from vllm.model_executor.layers.fused_moe.utils import _fp8_perm
11
-
12
-
13
-def _moe_permute(
14
- curr_hidden_states: torch.Tensor,
15
- a1q_scale: torch.Tensor | None,
16
- curr_topk_ids: torch.Tensor,
17
- global_num_experts: int,
18
- expert_map: torch.Tensor | None,
19
- block_m: int,
20
-) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor, torch.Tensor, torch.Tensor]:
21
- """
22
- Determine the sorted_token_ids, expert_ids for the given problem size.
23
- Permute the hidden states and scales according to `sorted_token_ids`.
24
- """
25
- top_k_num = curr_topk_ids.size(1)
26
-
27
- tokens_in_chunk = curr_hidden_states.size(0)
28
-
29
- sorted_token_ids, expert_ids, num_tokens_post_padded = moe_align_block_size(
30
- curr_topk_ids, block_m, global_num_experts, expert_map, pad_sorted_ids=True
31
- )
32
-
33
- inv_perm: torch.Tensor | None = None
34
-
35
- num_tokens = top_k_num * tokens_in_chunk
36
- expert_ids = torch.repeat_interleave(expert_ids, block_m, dim=0)
37
- inv_perm = torch.argsort(sorted_token_ids)[:num_tokens]
38
-
39
- # Permute according to sorted token ids.
40
- sorted_token_ids = sorted_token_ids.clamp(max=num_tokens - 1)
41
-
42
- curr_hidden_states = _fp8_perm(curr_hidden_states, sorted_token_ids // top_k_num)
43
-
44
- if a1q_scale is not None:
45
- a1q_scale = a1q_scale[sorted_token_ids // top_k_num]
46
-
47
- return (curr_hidden_states, a1q_scale, sorted_token_ids, expert_ids, inv_perm)
48
-
49
-
50
-def _moe_unpermute_and_reduce(
51
- out: torch.Tensor,
52
- curr_hidden: torch.Tensor,
53
- inv_perm: torch.Tensor | None,
54
- topk_weight: torch.Tensor,
55
- apply_router_weight_on_input: bool,
56
-) -> None:
57
- """
58
- Unpermute the final result and apply topk_weights, then perform the final
59
- reduction on the hidden states.
60
- """
61
- M, topk = topk_weight.size()
62
- K = curr_hidden.size(-1)
63
- if inv_perm is not None:
64
- curr_hidden = curr_hidden[inv_perm, ...]
65
- curr_hidden = curr_hidden.view(-1, topk, K)
66
- if not apply_router_weight_on_input:
67
- curr_hidden.mul_(topk_weight.view(M, -1, 1))
68
- ops.moe_sum(curr_hidden, out)
69
-
6
7
def moe_permute(
8
hidden_states: torch.Tensor,