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 10, 01:30 UTC
[MoE Refactor] Separate Router into OO Classes (#30623)
Signed-off-by: Bill Nell <bnell@redhat.com>
bnellnm committed
Jan 18, 2026 at 11:40 UTC
327a02d8db86e57f2488779b0c1b133da5f03fb5
45 files changed
+1754
-688
tests/kernels/moe/modular_kernel_tools/common.py
+1
-1
@@ -21,12 +21,12 @@ from vllm.distributed import (
21
get_tensor_model_parallel_world_size,
22
)
23
from vllm.forward_context import set_forward_context
24
+from vllm.model_executor.layers.fused_moe import fused_topk
25
from vllm.model_executor.layers.fused_moe.config import (
26
FusedMoEConfig,
27
FusedMoEParallelConfig,
28
FusedMoEQuantConfig,
29
)
29
-from vllm.model_executor.layers.fused_moe.fused_moe import fused_topk
30
from vllm.utils.import_utils import has_deep_ep, has_deep_gemm, has_pplx
31
32
from .mk_objects import (
tests/kernels/moe/test_batched_moe.py
+1
-1
@@ -15,10 +15,10 @@ from tests.kernels.moe.utils import (
15
from tests.kernels.quant_utils import native_batched_masked_quant_matmul
16
from tests.kernels.utils import torch_experts
17
from vllm.config import VllmConfig, set_current_vllm_config
18
+from vllm.model_executor.layers.fused_moe import fused_topk
19
from vllm.model_executor.layers.fused_moe.fused_batched_moe import (
20
invoke_moe_batched_triton_kernel,
21
)
21
-from vllm.model_executor.layers.fused_moe.fused_moe import fused_topk
22
from vllm.platforms import current_platform
23
from vllm.triton_utils import tl
24
from vllm.utils.torch_utils import set_random_seed
tests/kernels/moe/test_block_fp8.py
+4
-2
@@ -11,13 +11,15 @@ from tests.kernels.quant_utils import (
11
)
12
from vllm.config import VllmConfig, set_current_vllm_config
13
from vllm.model_executor.layers.activation import SiluAndMul
14
-from vllm.model_executor.layers.fused_moe import fused_experts
14
+from vllm.model_executor.layers.fused_moe import (
15
+ fused_experts,
16
+ fused_topk,
17
+)
18
from vllm.model_executor.layers.fused_moe.deep_gemm_moe import (
19
_valid_deep_gemm_shape,
20
deep_gemm_moe_fp8,
21
)
22
from vllm.model_executor.layers.fused_moe.fused_moe import (
20
- fused_topk,
23
modular_triton_fused_moe,
24
)
25
from vllm.platforms import current_platform
tests/kernels/moe/test_cutlass_moe.py
+1
-1
@@ -10,6 +10,7 @@ import torch
10
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
11
from vllm import _custom_ops as ops
12
from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config
13
+from vllm.model_executor.layers.fused_moe import fused_experts, fused_topk
14
from vllm.model_executor.layers.fused_moe.config import (
15
FUSED_MOE_UNQUANTIZED_CONFIG,
16
FusedMoEQuantConfig,
@@ -19,7 +20,6 @@ from vllm.model_executor.layers.fused_moe.cutlass_moe import (
20
CutlassExpertsFp8,
21
run_cutlass_moe_fp8,
22
)
22
-from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts, fused_topk
23
from vllm.model_executor.layers.fused_moe.prepare_finalize import (
24
MoEPrepareAndFinalizeNoEP,
25
)
tests/kernels/moe/test_flashinfer_moe.py
+1
-1
@@ -12,6 +12,7 @@ from tests.kernels.quantization.nvfp4_utils import (
12
from tests.kernels.utils import torch_moe
13
from vllm import _custom_ops as ops
14
from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config
15
+from vllm.model_executor.layers.fused_moe import fused_topk
16
from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe import (
17
FlashInferExperts,
18
is_valid_flashinfer_cutlass_fused_moe,
@@ -19,7 +20,6 @@ from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe import (
20
from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_prepare_finalize import (
21
create_flashinfer_prepare_finalize,
22
)
22
-from vllm.model_executor.layers.fused_moe.fused_moe import fused_topk
23
from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEModularKernel
24
from vllm.platforms import current_platform
25
from vllm.utils.flashinfer import has_flashinfer_cutlass_fused_moe
tests/kernels/moe/test_grouped_topk.py
+1
-1
@@ -14,7 +14,7 @@ from vllm.config import (
14
get_cached_compilation_config,
15
set_current_vllm_config,
16
)
17
-from vllm.model_executor.layers.fused_moe.fused_moe import (
17
+from vllm.model_executor.layers.fused_moe.router.grouped_topk_router import (
18
GroupedTopk,
19
fused_grouped_topk,
20
)
tests/kernels/moe/test_moe.py
+3
-1
@@ -24,6 +24,9 @@ from vllm._aiter_ops import rocm_aiter_ops
24
from vllm.config import VllmConfig, set_current_vllm_config
25
from vllm.distributed.parallel_state import init_distributed_environment
26
from vllm.forward_context import set_forward_context
27
+from vllm.model_executor.layers.fused_moe import (
28
+ fused_topk,
29
+)
30
from vllm.model_executor.layers.fused_moe.config import (
31
FUSED_MOE_UNQUANTIZED_CONFIG,
32
int4_w4a16_moe_quant_config,
@@ -34,7 +37,6 @@ from vllm.model_executor.layers.fused_moe.fused_marlin_moe import (
37
fused_marlin_moe,
38
)
39
from vllm.model_executor.layers.fused_moe.fused_moe import (
37
- fused_topk,
40
modular_triton_fused_moe,
41
)
42
from vllm.model_executor.layers.quantization.utils.marlin_utils import (
tests/kernels/moe/test_moe_permute_unpermute.py
+1
-1
@@ -9,7 +9,7 @@ import numpy as np
9
import pytest
10
import torch
11
12
-from vllm.model_executor.layers.fused_moe.fused_moe import fused_topk
12
+from vllm.model_executor.layers.fused_moe import fused_topk
13
from vllm.model_executor.layers.fused_moe.layer import determine_expert_map
14
from vllm.model_executor.layers.fused_moe.moe_permute_unpermute import (
15
moe_permute,
tests/kernels/moe/test_nvfp4_moe.py
+1
-1
@@ -13,11 +13,11 @@ from tests.kernels.quantization.nvfp4_utils import (
13
from tests.kernels.utils import torch_moe
14
from vllm import _custom_ops as ops
15
from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config
16
+from vllm.model_executor.layers.fused_moe import fused_topk
17
from vllm.model_executor.layers.fused_moe.config import nvfp4_moe_quant_config
18
from vllm.model_executor.layers.fused_moe.cutlass_moe import (
19
CutlassExpertsFp4,
20
)
20
-from vllm.model_executor.layers.fused_moe.fused_moe import fused_topk
21
from vllm.model_executor.layers.fused_moe.prepare_finalize import (
22
MoEPrepareAndFinalizeNoEP,
23
)
tests/kernels/moe/test_pplx_cutlass_moe.py
+1
-1
@@ -8,9 +8,9 @@ import torch
8
from tests.kernels.utils import torch_experts
9
from vllm import _custom_ops as ops
10
from vllm.config import VllmConfig, set_current_vllm_config
11
+from vllm.model_executor.layers.fused_moe import fused_topk
12
from vllm.model_executor.layers.fused_moe.config import fp8_w8a8_moe_quant_config
13
from vllm.model_executor.layers.fused_moe.cutlass_moe import CutlassBatchedExpertsFp8
13
-from vllm.model_executor.layers.fused_moe.fused_moe import fused_topk
14
from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEModularKernel
15
from vllm.platforms import current_platform
16
from vllm.utils.math_utils import cdiv
tests/kernels/moe/test_routing.py
new
+499
@@ -0,0 +1,499 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from collections.abc import Callable
4
+
5
+import pytest
6
+import torch
7
+
8
+from vllm.distributed.eplb.eplb_state import EplbLayerState
9
+from vllm.model_executor.layers.fused_moe.config import RoutingMethodType
10
+from vllm.model_executor.layers.fused_moe.router.router_factory import (
11
+ create_fused_moe_router,
12
+)
13
+from vllm.model_executor.models.llama4 import Llama4MoE
14
+
15
+# Test parameters
16
+MK_S = [(32, 256), (64, 512)]
17
+TOP_KS = [2, 4, 6]
18
+NUM_EXPERTS = [8, 16, 64]
19
+
20
+
21
+def setup_eplb_state(enable_eplb: bool, global_num_experts: int) -> EplbLayerState:
22
+ if not enable_eplb:
23
+ return EplbLayerState()
24
+
25
+ # Initialize EPLB state with proper tensors for testing
26
+ # For testing purposes, we use a simple 1:1 mapping (no redundant experts)
27
+ # expert_load_view: tracks load on each expert (shape: num_experts)
28
+ expert_load_view = torch.zeros(global_num_experts, dtype=torch.int32, device="cuda")
29
+
30
+ # logical_to_physical_map: maps logical experts to physical experts
31
+ # Shape: (num_logical_experts, max_slots)
32
+ # For testing, use simple 1:1 mapping with single slot per expert
33
+ logical_to_physical_map = torch.arange(
34
+ global_num_experts, dtype=torch.int64, device="cuda"
35
+ ).unsqueeze(-1)
36
+
37
+ # logical_replica_count: number of replicas per logical expert
38
+ # Shape: (num_logical_experts,)
39
+ # For testing, each logical expert has exactly 1 replica
40
+ logical_replica_count = torch.ones(
41
+ global_num_experts, dtype=torch.int64, device="cuda"
42
+ )
43
+
44
+ return EplbLayerState(
45
+ expert_load_view=expert_load_view,
46
+ logical_to_physical_map=logical_to_physical_map,
47
+ logical_replica_count=logical_replica_count,
48
+ )
49
+
50
+
51
+def make_test_data(
52
+ m: int, k: int, num_experts: int
53
+) -> tuple[torch.Tensor, torch.Tensor]:
54
+ hidden_states = torch.randn((m, k), device="cuda") / 10
55
+ logits = torch.randn((m, num_experts), device="cuda")
56
+ return hidden_states, logits
57
+
58
+
59
+def make_e_score_correction_bias(
60
+ e_score_correction_bias_val: float,
61
+ num_experts: int,
62
+) -> torch.Tensor:
63
+ # return torch.randn(num_experts, device="cuda") * e_score_correction_bias_val
64
+ return torch.full(
65
+ (num_experts,), e_score_correction_bias_val, device="cuda", dtype=torch.float32
66
+ )
67
+
68
+
69
+def assert_routing_results_close(
70
+ topk_weights: torch.Tensor,
71
+ topk_ids: torch.Tensor,
72
+ baseline_weights: torch.Tensor,
73
+ baseline_ids: torch.Tensor,
74
+ rtol: float = 1e-3,
75
+ atol: float = 1e-3,
76
+):
77
+ """
78
+ Compare routing results, sorting by expert ID first to handle non-deterministic
79
+ ordering from sorted=False in topk.
80
+ """
81
+ # Sort both results by expert IDs for consistent comparison
82
+ sorted_indices_actual = torch.argsort(topk_ids, dim=-1)
83
+ sorted_indices_baseline = torch.argsort(baseline_ids.to(topk_ids.dtype), dim=-1)
84
+
85
+ # Gather the sorted values
86
+ topk_ids_sorted = torch.gather(topk_ids, 1, sorted_indices_actual)
87
+ topk_weights_sorted = torch.gather(topk_weights, 1, sorted_indices_actual)
88
+ baseline_ids_sorted = torch.gather(
89
+ baseline_ids.to(topk_ids.dtype), 1, sorted_indices_baseline
90
+ )
91
+ baseline_weights_sorted = torch.gather(baseline_weights, 1, sorted_indices_baseline)
92
+
93
+ # Compare
94
+ torch.testing.assert_close(topk_ids_sorted, baseline_ids_sorted)
95
+ torch.testing.assert_close(
96
+ topk_weights_sorted, baseline_weights_sorted, rtol=rtol, atol=atol
97
+ )
98
+
99
+
100
+def baseline_fused_topk(
101
+ router_logits: torch.Tensor, top_k: int, renormalize: bool
102
+) -> tuple[torch.Tensor, torch.Tensor]:
103
+ """
104
+ Baseline for standard fused top-k routing.
105
+
106
+ Algorithm:
107
+ 1. Apply softmax to router logits
108
+ 2. Select top-k experts
109
+ 3. Optionally renormalize the weights
110
+ """
111
+ scores = torch.softmax(router_logits, dim=-1, dtype=torch.float32)
112
+ # Use sorted=False to match vllm implementation (vllm_is_batch_invariant
113
+ # defaults to False)
114
+ topk_weights, topk_ids = torch.topk(scores, top_k, dim=-1, sorted=False)
115
+
116
+ if renormalize:
117
+ topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
118
+
119
+ return topk_weights.to(torch.float32), topk_ids.to(torch.int32)
120
+
121
+
122
+def baseline_fused_topk_bias(
123
+ router_logits: torch.Tensor,
124
+ top_k: int,
125
+ renormalize: bool,
126
+ e_score_correction_bias: torch.Tensor,
127
+ routed_scaling_factor: float,
128
+) -> tuple[torch.Tensor, torch.Tensor]:
129
+ """
130
+ Baseline for fused top-k with bias correction.
131
+
132
+ Algorithm:
133
+ 1. Apply softmax to router logits
134
+ 2. Add bias to scores for expert selection
135
+ 3. Select top-k experts using biased scores
136
+ 4. Get weights from original (unbiased) scores
137
+ 5. Apply routed scaling factor
138
+ 6. Optionally renormalize the weights
139
+ """
140
+ # Apply softmax to get scores
141
+ scores = torch.softmax(router_logits, dim=-1, dtype=torch.float32)
142
+
143
+ # Add bias for expert selection
144
+ scores_for_choice = scores + e_score_correction_bias.unsqueeze(0)
145
+
146
+ # Select top-k using biased scores (sorted=False to match implementation)
147
+ topk_ids = torch.topk(scores_for_choice, k=top_k, dim=-1, sorted=False)[1]
148
+
149
+ # Get weights from original scores (not biased)
150
+ topk_weights = scores.gather(1, topk_ids)
151
+
152
+ # Renormalize if needed (BEFORE applying scaling factor)
153
+ if renormalize:
154
+ topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
155
+
156
+ # Apply scaling factor (AFTER renormalization, if applicable)
157
+ if routed_scaling_factor != 1.0:
158
+ topk_weights *= routed_scaling_factor
159
+
160
+ return topk_weights.to(torch.float32), topk_ids.to(torch.int32)
161
+
162
+
163
+def baseline_grouped_topk(
164
+ router_logits: torch.Tensor,
165
+ top_k: int,
166
+ num_expert_group: int,
167
+ topk_group: int,
168
+ scoring_func: str,
169
+ renormalize: bool,
170
+ e_score_correction_bias: torch.Tensor | None,
171
+ routed_scaling_factor: float,
172
+) -> tuple[torch.Tensor, torch.Tensor]:
173
+ """
174
+ Baseline for grouped top-k routing (e.g., DeepSeek).
175
+
176
+ Algorithm:
177
+ 1. Apply scoring function (softmax or sigmoid)
178
+ 2. Optionally add bias
179
+ 3. Select top-k groups based on max scores within each group
180
+ 4. Mask scores to only include selected groups
181
+ 5. Select top-k experts from masked scores
182
+ 6. Apply scaling factor
183
+ 7. Optionally renormalize
184
+ """
185
+ num_token = router_logits.shape[0]
186
+
187
+ # Apply scoring function
188
+ if scoring_func == "softmax":
189
+ scores = torch.softmax(router_logits, dim=-1, dtype=torch.float32)
190
+ elif scoring_func == "sigmoid":
191
+ scores = torch.sigmoid(router_logits.float())
192
+ else:
193
+ raise ValueError(f"Unsupported scoring function: {scoring_func}")
194
+
195
+ # Handle bias correction
196
+ if e_score_correction_bias is not None:
197
+ original_scores = scores
198
+ scores = scores + e_score_correction_bias.unsqueeze(0)
199
+ # For bias case, use sum of top-2 scores in each group
200
+ group_scores = (
201
+ scores.view(num_token, num_expert_group, -1).topk(2, dim=-1)[0].sum(dim=-1)
202
+ )
203
+ else:
204
+ # Use max score in each group
205
+ group_scores = scores.view(num_token, num_expert_group, -1).max(dim=-1).values
206
+
207
+ # Select top-k groups
208
+ group_idx = torch.topk(group_scores, k=topk_group, dim=-1, sorted=False)[1]
209
+
210
+ # Create mask for selected groups
211
+ group_mask = torch.zeros_like(group_scores)
212
+ group_mask.scatter_(1, group_idx, 1)
213
+
214
+ # Expand mask to all experts
215
+ score_mask = (
216
+ group_mask.unsqueeze(-1)
217
+ .expand(num_token, num_expert_group, scores.shape[-1] // num_expert_group)
218
+ .reshape(num_token, -1)
219
+ )
220
+
221
+ # Mask scores (set non-selected to -inf)
222
+ tmp_scores = scores.masked_fill(~score_mask.bool(), float("-inf"))
223
+
224
+ # Select top-k experts
225
+ if e_score_correction_bias is not None:
226
+ topk_ids = torch.topk(tmp_scores, k=top_k, dim=-1, sorted=False)[1]
227
+ topk_weights = original_scores.gather(1, topk_ids)
228
+ else:
229
+ topk_weights, topk_ids = torch.topk(tmp_scores, k=top_k, dim=-1, sorted=False)
230
+
231
+ # Renormalize if needed
232
+ if renormalize:
233
+ topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
234
+
235
+ # Apply scaling factor
236
+ if routed_scaling_factor != 1.0:
237
+ topk_weights *= routed_scaling_factor
238
+
239
+ return topk_weights.to(torch.float32), topk_ids.to(torch.int32)
240
+
241
+
242
+def baseline_custom_llama4(
243
+ router_logits: torch.Tensor, top_k: int
244
+) -> tuple[torch.Tensor, torch.Tensor]:
245
+ """
246
+ Baseline for Llama4 custom routing.
247
+
248
+ Algorithm:
249
+ 1. Select top-k expert indices (without softmax)
250
+ 2. Apply sigmoid to the selected scores
251
+ """
252
+ router_scores, router_indices = torch.topk(router_logits, top_k, dim=-1)
253
+ router_scores = torch.sigmoid(router_scores.float())
254
+ return router_scores.to(torch.float32), router_indices.to(torch.int32)
255
+
256
+
257
+@pytest.mark.parametrize("m,k", MK_S)
258
+@pytest.mark.parametrize("top_k", TOP_KS)
259
+@pytest.mark.parametrize("global_num_experts", NUM_EXPERTS)
260
+@pytest.mark.parametrize("renormalize", [False, True])
261
+@pytest.mark.parametrize("enable_eplb", [False, True])
262
+def test_fused_topk(
263
+ m: int,
264
+ k: int,
265
+ top_k: int,
266
+ global_num_experts: int,
267
+ renormalize: bool,
268
+ enable_eplb: bool,
269
+):
270
+ if top_k > global_num_experts:
271
+ pytest.skip(f"top_k ({top_k}) > global_num_experts ({global_num_experts})")
272
+
273
+ eplb_state = setup_eplb_state(enable_eplb, global_num_experts)
274
+ router = create_fused_moe_router(
275
+ top_k=top_k,
276
+ global_num_experts=global_num_experts,
277
+ renormalize=renormalize,
278
+ enable_eplb=enable_eplb,
279
+ eplb_state=eplb_state,
280
+ )
281
+
282
+ hidden_states, router_logits = make_test_data(m, k, global_num_experts)
283
+
284
+ # Get router output
285
+ topk_weights, topk_ids = router.select_experts(hidden_states, router_logits)
286
+
287
+ # Compute baseline
288
+ baseline_weights, baseline_ids = baseline_fused_topk(
289
+ router_logits, top_k, renormalize
290
+ )
291
+
292
+ # Compare results
293
+ assert_routing_results_close(topk_weights, topk_ids, baseline_weights, baseline_ids)
294
+
295
+
296
+@pytest.mark.parametrize("m,k", MK_S)
297
+@pytest.mark.parametrize("top_k", TOP_KS)
298
+@pytest.mark.parametrize("global_num_experts", NUM_EXPERTS)
299
+@pytest.mark.parametrize("renormalize", [False, True])
300
+@pytest.mark.parametrize("enable_eplb", [False, True])
301
+@pytest.mark.parametrize("e_score_correction_bias_val", [0.9])
302
+@pytest.mark.parametrize("routed_scaling_factor", [1.0, 1.1])
303
+def test_fused_topk_bias(
304
+ m: int,
305
+ k: int,
306
+ top_k: int,
307
+ global_num_experts: int,
308
+ renormalize: bool,
309
+ enable_eplb: bool,
310
+ e_score_correction_bias_val: float,
311
+ routed_scaling_factor: float,
312
+):
313
+ if top_k > global_num_experts:
314
+ pytest.skip(f"top_k ({top_k}) > global_num_experts ({global_num_experts})")
315
+
316
+ eplb_state = setup_eplb_state(enable_eplb, global_num_experts)
317
+
318
+ e_score_correction_bias = make_e_score_correction_bias(
319
+ e_score_correction_bias_val,
320
+ global_num_experts,
321
+ )
322
+
323
+ router = create_fused_moe_router(
324
+ e_score_correction_bias=e_score_correction_bias,
325
+ routed_scaling_factor=routed_scaling_factor,
326
+ top_k=top_k,
327
+ global_num_experts=global_num_experts,
328
+ renormalize=renormalize,
329
+ enable_eplb=enable_eplb,
330
+ eplb_state=eplb_state,
331
+ )
332
+
333
+ hidden_states, router_logits = make_test_data(m, k, global_num_experts)
334
+
335
+ # Get router output
336
+ topk_weights, topk_ids = router.select_experts(hidden_states, router_logits)
337
+
338
+ # Compute baseline
339
+ baseline_weights, baseline_ids = baseline_fused_topk_bias(
340
+ router_logits,
341
+ top_k,
342
+ renormalize,
343
+ e_score_correction_bias,
344
+ routed_scaling_factor,
345
+ )
346
+
347
+ # Compare results
348
+ assert_routing_results_close(topk_weights, topk_ids, baseline_weights, baseline_ids)
349
+
350
+
351
+@pytest.mark.parametrize("m,k", MK_S)
352
+@pytest.mark.parametrize("top_k", TOP_KS)
353
+@pytest.mark.parametrize(
354
+ "global_num_experts,num_expert_group,topk_group",
355
+ [
356
+ (64, 8, 4), # 8 groups of 8 experts, select 4 groups
357
+ (32, 4, 2), # 4 groups of 8 experts, select 2 groups
358
+ ],
359
+)
360
+@pytest.mark.parametrize("renormalize", [False, True])
361
+@pytest.mark.parametrize("enable_eplb", [False, True])
362
+@pytest.mark.parametrize("e_score_correction_bias_val", [0.9])
363
+@pytest.mark.parametrize("routed_scaling_factor", [1.0, 1.1])
364
+@pytest.mark.parametrize("scoring_func", ["sigmoid", "softmax"])
365
+def test_grouped_topk(
366
+ m: int,
367
+ k: int,
368
+ top_k: int,
369
+ global_num_experts: int,
370
+ renormalize: bool,
371
+ enable_eplb: bool,
372
+ num_expert_group: int,
373
+ topk_group: int,
374
+ scoring_func: str,
375
+ e_score_correction_bias_val: float,
376
+ routed_scaling_factor: float,
377
+):
378
+ if top_k > global_num_experts:
379
+ pytest.skip(f"top_k ({top_k}) > global_num_experts ({global_num_experts})")
380
+
381
+ eplb_state = setup_eplb_state(enable_eplb, global_num_experts)
382
+
383
+ e_score_correction_bias = make_e_score_correction_bias(
384
+ e_score_correction_bias_val,
385
+ global_num_experts,
386
+ )
387
+
388
+ routing_method_type = None
389
+ if scoring_func == "llama4":
390
+ routing_method_type = RoutingMethodType.Llama4
391
+ scoring_func = "sigmoid"
392
+
393
+ router = create_fused_moe_router(
394
+ use_grouped_topk=True,
395
+ num_expert_group=num_expert_group,
396
+ topk_group=topk_group,
397
+ scoring_func=scoring_func,
398
+ routing_method_type=routing_method_type,
399
+ e_score_correction_bias=e_score_correction_bias,
400
+ routed_scaling_factor=routed_scaling_factor,
401
+ top_k=top_k,
402
+ global_num_experts=global_num_experts,
403
+ renormalize=renormalize,
404
+ enable_eplb=enable_eplb,
405
+ eplb_state=eplb_state,
406
+ )
407
+
408
+ hidden_states, router_logits = make_test_data(m, k, global_num_experts)
409
+
410
+ # Get router output
411
+ topk_weights, topk_ids = router.select_experts(hidden_states, router_logits)
412
+
413
+ # Compute baseline
414
+ baseline_weights, baseline_ids = baseline_grouped_topk(
415
+ router_logits,
416
+ top_k,
417
+ num_expert_group,
418
+ topk_group,
419
+ scoring_func,
420
+ renormalize,
421
+ e_score_correction_bias,
422
+ routed_scaling_factor,
423
+ )
424
+
425
+ # Compare results
426
+ assert_routing_results_close(topk_weights, topk_ids, baseline_weights, baseline_ids)
427
+
428
+
429
+@pytest.mark.parametrize("m,k", MK_S)
430
+@pytest.mark.parametrize("top_k", TOP_KS)
431
+@pytest.mark.parametrize("global_num_experts", NUM_EXPERTS)
432
+@pytest.mark.parametrize("renormalize", [False, True])
433
+@pytest.mark.parametrize("enable_eplb", [False, True])
434
+@pytest.mark.parametrize("custom_routing_function", [Llama4MoE.custom_routing_function])
435
+def test_custom(
436
+ m: int,
437
+ k: int,
438
+ top_k: int,
439
+ global_num_experts: int,
440
+ renormalize: bool,
441
+ enable_eplb: bool,
442
+ custom_routing_function: Callable,
443
+):
444
+ if top_k > global_num_experts:
445
+ pytest.skip(f"top_k ({top_k}) > global_num_experts ({global_num_experts})")
446
+
447
+ eplb_state = setup_eplb_state(enable_eplb, global_num_experts)
448
+
449
+ router = create_fused_moe_router(
450
+ top_k=top_k,
451
+ global_num_experts=global_num_experts,
452
+ custom_routing_function=custom_routing_function,
453
+ renormalize=renormalize,
454
+ enable_eplb=enable_eplb,
455
+ eplb_state=eplb_state,
456
+ )
457
+
458
+ hidden_states, router_logits = make_test_data(m, k, global_num_experts)
459
+
460
+ # Get router output
461
+ topk_weights, topk_ids = router.select_experts(hidden_states, router_logits)
462
+
463
+ # Compute baseline (Llama4 uses sigmoid)
464
+ baseline_weights, baseline_ids = baseline_custom_llama4(router_logits, top_k)
465
+
466
+ # Compare results
467
+ assert_routing_results_close(topk_weights, topk_ids, baseline_weights, baseline_ids)
468
+
469
+
470
+# TODO: is other test sufficient?
471
+# # See tests/test_routing_simulatator.py
472
+# @pytest.mark.parametrize("m,k", MK_S)
473
+# @pytest.mark.parametrize("top_k", TOP_KS)
474
+# @pytest.mark.parametrize("global_num_experts", NUM_EXPERTS)
475
+# @pytest.mark.parametrize("renormalize", [False, True])
476
+# @pytest.mark.parametrize("enable_eplb", [False, True])
477
+# @pytest.mark.parameterize("strategy", ["uniform_random", "normal_routing"])
478
+# def test_simulated(
479
+# m: int,
480
+# k: int,
481
+# top_k: int,
482
+# global_num_experts: int,
483
+# renormalize: bool,
484
+# enable_eplb: bool,
485
+# strategy: str,
486
+# monkeypatch,
487
+# ):
488
+# eplb_state = setup_eplb_state(enable_eplb)
489
+
490
+# monkeypatch.setenv("VLLM_MOE_ROUTING_SIMULATION_STRATEGY", strategy)
491
+# router = create_fused_moe_router(
492
+# top_k=top_k,
493
+# global_num_experts=global_num_experts,
494
+# enable_eplb=enable_eplb,
495
+# eplb_state=eplb_state,
496
+# )
497
+
498
+# hidden_states, router_logits = make_test_data(m, k, global_num_experts)
499
+# topk_weights, topk_ids = router.select_experts(hidden_states, router_logits)
tests/kernels/moe/test_routing_simulator.py
renamed
+38
-34
@@ -19,7 +19,7 @@ from vllm.distributed import (
19
init_distributed_environment,
20
initialize_model_parallel,
21
)
22
-from vllm.model_executor.layers.fused_moe.routing_simulator import (
22
+from vllm.model_executor.layers.fused_moe.router.routing_simulator_router import (
23
DistributionBasedRouting,
24
RoutingSimulator,
25
)
@@ -109,40 +109,44 @@ def test_routing_strategy_integration(monkeypatch, device):
109
tensor_model_parallel_size=1,
110
pipeline_model_parallel_size=1,
111
)
112
- fused_moe = FusedMoE(
113
- num_experts=num_experts,
114
- top_k=top_k,
115
- hidden_size=hidden_size,
116
- intermediate_size=0,
117
- use_grouped_topk=False,
118
- renormalize=True,
119
- )
120
-
121
- for strategy in strategies:
122
- # Set environment variable
123
- env_name = "VLLM_MOE_ROUTING_SIMULATION_STRATEGY"
124
- monkeypatch.setenv(env_name, strategy)
112
126
- # Force reload of environment variable
127
- envs.environment_variables[env_name] = lambda s=strategy: s
128
-
129
- # Test the select_experts method
130
- topk_weights, topk_ids = fused_moe.router.select_experts(
131
- hidden_states=hidden_states,
132
- router_logits=router_logits,
133
- )
134
-
135
- # Verify output shapes
136
- assert topk_weights.shape == (num_tokens, top_k), (
137
- f"Wrong weights shape for {strategy}"
138
- )
139
- assert topk_ids.shape == (num_tokens, top_k), f"Wrong ids shape for {strategy}"
140
-
141
- # Verify expert IDs are valid
142
- assert topk_ids.min() >= 0, f"Invalid expert ID (negative) for {strategy}"
143
- assert topk_ids.max() < num_experts, (
144
- f"Invalid expert ID (too large) for {strategy}"
145
- )
113
+ for strategy in strategies:
114
+ fused_moe = FusedMoE(
115
+ num_experts=num_experts,
116
+ top_k=top_k,
117
+ hidden_size=hidden_size,
118
+ intermediate_size=0,
119
+ use_grouped_topk=False,
120
+ renormalize=True,
121
+ prefix=strategy,
122
+ )
123
+
124
+ # Set environment variable
125
+ env_name = "VLLM_MOE_ROUTING_SIMULATION_STRATEGY"
126
+ monkeypatch.setenv(env_name, strategy)
127
+
128
+ # Force reload of environment variable
129
+ envs.environment_variables[env_name] = lambda s=strategy: s
130
+
131
+ # Test the select_experts method
132
+ topk_weights, topk_ids = fused_moe.router.select_experts(
133
+ hidden_states=hidden_states,
134
+ router_logits=router_logits,
135
+ )
136
+
137
+ # Verify output shapes
138
+ assert topk_weights.shape == (num_tokens, top_k), (
139
+ f"Wrong weights shape for {strategy}"
140
+ )
141
+ assert topk_ids.shape == (num_tokens, top_k), (
142
+ f"Wrong ids shape for {strategy}"
143
+ )
144
+
145
+ # Verify expert IDs are valid
146
+ assert topk_ids.min() >= 0, f"Invalid expert ID (negative) for {strategy}"
147
+ assert topk_ids.max() < num_experts, (
148
+ f"Invalid expert ID (too large) for {strategy}"
149
+ )
150
151
152
def test_distribution_based_routing_with_custom_strategy():
tests/model_executor/test_enabled_custom_ops.py
+1
-1
@@ -17,7 +17,7 @@ from vllm.model_executor.layers.activation import (
17
ReLUSquaredActivation,
18
SiluAndMul,
19
)
20
-from vllm.model_executor.layers.fused_moe.fused_moe import (
20
+from vllm.model_executor.layers.fused_moe.router.fused_topk_router import (
21
dispatch_topk_func,
22
vllm_topk_softmax,
23
)
vllm/distributed/eplb/eplb_state.py
+9
@@ -1158,6 +1158,15 @@ class EplbState:
1158
return self._allreduce_list(load_pass_list)
1159
1160
1161
+@dataclass
1162
+class EplbLayerState:
1163
+ """Runtime EPLB data stored in the MoE layer."""
1164
+
1165
+ expert_load_view: torch.Tensor | None = None
1166
+ logical_to_physical_map: torch.Tensor | None = None
1167
+ logical_replica_count: torch.Tensor | None = None
1168
+
1169
+
1170
def _node_count_with_rank_mapping(
1171
pg: ProcessGroup | StatelessProcessGroup,
1172
rank_mapping: dict[int, int],
vllm/model_executor/layers/fused_moe/__init__.py
+9
-5
@@ -11,9 +11,6 @@ from vllm.model_executor.layers.fused_moe.config import (
11
from vllm.model_executor.layers.fused_moe.fused_moe_method_base import (
12
FusedMoEMethodBase,
13
)
14
-from vllm.model_executor.layers.fused_moe.fused_moe_router import (
15
- FusedMoERouter,
16
-)
14
from vllm.model_executor.layers.fused_moe.layer import (
15
FusedMoE,
16
FusedMoeWeightScaleSupported,
@@ -23,6 +20,9 @@ from vllm.model_executor.layers.fused_moe.modular_kernel import (
20
FusedMoEPermuteExpertsUnpermute,
21
FusedMoEPrepareAndFinalize,
22
)
23
+from vllm.model_executor.layers.fused_moe.router.fused_moe_router import (
24
+ FusedMoERouter,
25
+)
26
from vllm.model_executor.layers.fused_moe.shared_fused_moe import SharedFusedMoE
27
from vllm.model_executor.layers.fused_moe.unquantized_fused_moe_method import (
28
UnquantizedFusedMoEMethod,
@@ -83,13 +83,17 @@ if HAS_TRITON:
83
BatchedTritonExperts,
84
)
85
from vllm.model_executor.layers.fused_moe.fused_moe import (
86
- GroupedTopk,
86
TritonExperts,
87
TritonWNA16Experts,
88
fused_experts,
90
- fused_topk,
89
get_config_file_name,
90
)
91
+ from vllm.model_executor.layers.fused_moe.router.fused_topk_router import (
92
+ fused_topk,
93
+ )
94
+ from vllm.model_executor.layers.fused_moe.router.grouped_topk_router import (
95
+ GroupedTopk,
96
+ )
97
from vllm.model_executor.layers.fused_moe.triton_deep_gemm_moe import (
98
TritonOrDeepGemmExperts,
99
)
vllm/model_executor/layers/fused_moe/config.py
+5
-1
@@ -117,8 +117,12 @@ class RoutingMethodType(IntEnum):
117
RenormalizeNaive = (4,)
118
# TopK: TopK (no softmax)
119
TopK = (5,)
120
+ # Custom
121
+ Custom = (6,)
122
+ # Simulated
123
+ Simulated = (7,)
124
# Unspecified
121
- Unspecified = 6.0
125
+ Unspecified = 8.0
126
127
128
@dataclass
vllm/model_executor/layers/fused_moe/fused_moe.py
-375
@@ -13,9 +13,7 @@ import torch
13
import vllm.envs as envs
14
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
15
from vllm import _custom_ops as ops
16
-from vllm._aiter_ops import rocm_aiter_ops
16
from vllm.logger import init_logger
18
-from vllm.model_executor.custom_op import CustomOp
17
from vllm.model_executor.layers.batch_invariant import (
18
vllm_is_batch_invariant,
19
)
@@ -34,9 +32,6 @@ from vllm.model_executor.layers.fused_moe.moe_align_block_size import (
32
from vllm.model_executor.layers.fused_moe.prepare_finalize import (
33
MoEPrepareAndFinalizeNoEP,
34
)
37
-from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import ( # noqa: E501
38
- rocm_aiter_grouped_topk,
39
-)
35
from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
36
TopKWeightAndReduceNoOP,
37
)
@@ -49,7 +44,6 @@ from vllm.model_executor.layers.fused_moe.utils import (
44
from vllm.model_executor.layers.quantization.utils.mxfp4_utils import dequant_mxfp4
45
from vllm.model_executor.layers.quantization.utils.mxfp6_utils import dequant_mxfp6
46
from vllm.model_executor.layers.quantization.utils.ocp_mx_utils import OCP_MX_Scheme
52
-from vllm.model_executor.utils import maybe_disable_graph_partition
47
from vllm.platforms import current_platform
48
from vllm.triton_utils import tl, triton
49
from vllm.utils.deep_gemm import is_deep_gemm_e8m0_used
@@ -1318,375 +1312,6 @@ def try_get_optimal_moe_config(
1312
return config
1313
1314
1321
-def vllm_topk_softmax(
1322
- topk_weights: torch.Tensor,
1323
- topk_indices: torch.Tensor,
1324
- token_expert_indices: torch.Tensor,
1325
- gating_output: torch.Tensor,
1326
- renormalize: bool,
1327
-) -> tuple[torch.Tensor, ...]:
1328
- ops.topk_softmax(
1329
- topk_weights,
1330
- topk_indices,
1331
- token_expert_indices,
1332
- gating_output,
1333
- renormalize,
1334
- )
1335
-
1336
- return topk_weights, topk_indices
1337
-
1338
-
1339
-def dispatch_topk_func(
1340
- use_rocm_aiter: bool = False,
1341
-) -> Callable[..., tuple[torch.Tensor, ...]]:
1342
- if use_rocm_aiter:
1343
- return rocm_aiter_ops.topk_softmax
1344
- return vllm_topk_softmax
1345
-
1346
-
1347
-def fused_topk(
1348
- hidden_states: torch.Tensor,
1349
- gating_output: torch.Tensor,
1350
- topk: int,
1351
- renormalize: bool,
1352
- indices_type: torch.dtype | None = None,
1353
-) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
1354
- assert hidden_states.size(0) == gating_output.size(0), "Number of tokens mismatch"
1355
-
1356
- M, _ = hidden_states.size()
1357
-
1358
- topk_weights = torch.empty(
1359
- M, topk, dtype=torch.float32, device=hidden_states.device
1360
- )
1361
- topk_ids = torch.empty(
1362
- M,
1363
- topk,
1364
- dtype=torch.int32 if indices_type is None else indices_type,
1365
- device=hidden_states.device,
1366
- )
1367
- token_expert_indices = torch.empty(
1368
- M, topk, dtype=torch.int32, device=hidden_states.device
1369
- )
1370
-
1371
- topk_func = dispatch_topk_func(use_rocm_aiter=rocm_aiter_ops.is_fused_moe_enabled())
1372
- topk_weights, topk_ids = topk_func(
1373
- topk_weights, topk_ids, token_expert_indices, gating_output, renormalize
1374
- )
1375
-
1376
- return topk_weights, topk_ids, token_expert_indices
1377
-
1378
-
1379
-def fused_topk_bias(
1380
- hidden_states: torch.Tensor,
1381
- gating_output: torch.Tensor,
1382
- e_score_correction_bias: torch.Tensor,
1383
- topk: int,
1384
- renormalize: bool,
1385
-):
1386
- n_routed_experts = gating_output.shape[-1]
1387
- scores = gating_output.softmax(dim=-1)
1388
- scores_for_choice = scores.view(
1389
- -1, n_routed_experts
1390
- ) + e_score_correction_bias.unsqueeze(0)
1391
-
1392
- # For batch invariance, use sorted=True to ensure deterministic expert selection
1393
- use_sorted = vllm_is_batch_invariant()
1394
- topk_indices = torch.topk(scores_for_choice, k=topk, dim=-1, sorted=use_sorted)[1]
1395
- topk_weights = scores.gather(1, topk_indices)
1396
- if renormalize:
1397
- topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
1398
- return topk_weights.to(torch.float32), topk_indices.to(torch.int32)
1399
-
1400
-
1401
-# This is used by the Deepseek-V2 and Deepseek-V3 model
1402
-@torch.compile(
1403
- dynamic=True,
1404
- backend=current_platform.simple_compile_backend,
1405
- options=maybe_disable_graph_partition(current_platform.simple_compile_backend),
1406
-)
1407
-def grouped_topk(
1408
- hidden_states: torch.Tensor,
1409
- gating_output: torch.Tensor,
1410
- topk: int,
1411
- renormalize: bool,
1412
- num_expert_group: int = 0,
1413
- topk_group: int = 0,
1414
- scoring_func: str = "softmax",
1415
- routed_scaling_factor: float = 1.0,
1416
- e_score_correction_bias: torch.Tensor | None = None,
1417
-) -> tuple[torch.Tensor, torch.Tensor]:
1418
- if (
1419
- envs.VLLM_USE_FUSED_MOE_GROUPED_TOPK
1420
- and current_platform.is_cuda()
1421
- and num_expert_group <= 32
1422
- and topk <= 32
1423
- and e_score_correction_bias is not None
1424
- ):
1425
- return fused_grouped_topk(
1426
- hidden_states=hidden_states,
1427
- gating_output=gating_output,
1428
- topk=topk,
1429
- renormalize=renormalize,
1430
- e_score_correction_bias=e_score_correction_bias,
1431
- num_expert_group=num_expert_group,
1432
- topk_group=topk_group,
1433
- scoring_func=scoring_func,
1434
- routed_scaling_factor=routed_scaling_factor,
1435
- )
1436
-
1437
- assert hidden_states.size(0) == gating_output.size(0), "Number of tokens mismatch"
1438
-
1439
- if scoring_func == "softmax":
1440
- scores = torch.softmax(gating_output, dim=-1)
1441
- elif scoring_func == "sigmoid":
1442
- scores = gating_output.sigmoid()
1443
- else:
1444
- raise ValueError(f"Unsupported scoring function: {scoring_func}")
1445
-
1446
- num_token = scores.size(0)
1447
- if e_score_correction_bias is not None:
1448
- # Store original scores before applying correction bias. We use biased
1449
- # scores for expert selection but original scores for routing weights
1450
- original_scores = scores
1451
- scores = scores + e_score_correction_bias.unsqueeze(0)
1452
- group_scores = (
1453
- scores.view(num_token, num_expert_group, -1).topk(2, dim=-1)[0].sum(dim=-1)
1454
- )
1455
- else:
1456
- group_scores = (
1457
- scores.view(num_token, num_expert_group, -1).max(dim=-1).values
1458
- ) # [n, n_group]
1459
-
1460
- # For batch invariance, use sorted=True to ensure deterministic expert selection
1461
- use_sorted = vllm_is_batch_invariant()
1462
- group_idx = torch.topk(group_scores, k=topk_group, dim=-1, sorted=use_sorted)[
1463
- 1
1464
- ] # [n, top_k_group]
1465
- group_mask = torch.zeros_like(group_scores) # [n, n_group]
1466
- group_mask.scatter_(1, group_idx, 1) # [n, n_group]
1467
- score_mask = (
1468
- group_mask.unsqueeze(-1)
1469
- .expand(num_token, num_expert_group, scores.size(-1) // num_expert_group)
1470
- .reshape(num_token, -1)
1471
- ) # [n, e]
1472
- tmp_scores = scores.masked_fill(~score_mask.bool(), float("-inf")) # [n, e]
1473
-
1474
- if e_score_correction_bias is not None:
1475
- topk_ids = torch.topk(tmp_scores, k=topk, dim=-1, sorted=use_sorted)[1]
1476
- # Use original unbiased scores for the routing weights
1477
- topk_weights = original_scores.gather(1, topk_ids)
1478
- else:
1479
- topk_weights, topk_ids = torch.topk(
1480
- tmp_scores, k=topk, dim=-1, sorted=use_sorted
1481
- )
1482
-
1483
- if renormalize:
1484
- topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
1485
-
1486
- if routed_scaling_factor != 1.0:
1487
- topk_weights = topk_weights * routed_scaling_factor
1488
- return topk_weights.to(torch.float32), topk_ids.to(torch.int32)
1489
-
1490
-
1491
-# --8<-- [start:grouped_topk]
1492
-@CustomOp.register("grouped_topk")
1493
-class GroupedTopk(CustomOp):
1494
- """GroupedTopk used by the Deepseek-V2 and Deepseek-V3 model."""
1495
-
1496
- # --8<-- [end:grouped_topk]
1497
-
1498
- def __init__(
1499
- self,
1500
- topk: int,
1501
- renormalize: bool,
1502
- num_expert_group: int = 0,
1503
- topk_group: int = 0,
1504
- scoring_func: str = "softmax",
1505
- routed_scaling_factor: float = 1.0,
1506
- num_fused_shared_experts: int = 0,
1507
- ) -> None:
1508
- super().__init__()
1509
- self.native_impl = grouped_topk
1510
- self.topk = topk
1511
- self.renormalize = renormalize
1512
- self.num_expert_group = num_expert_group
1513
- self.topk_group = topk_group
1514
- self.scoring_func = scoring_func
1515
- self.routed_scaling_factor = routed_scaling_factor
1516
- self.num_fused_shared_experts = num_fused_shared_experts
1517
-
1518
- def forward_native(
1519
- self,
1520
- hidden_states: torch.Tensor,
1521
- gating_output: torch.Tensor,
1522
- e_score_correction_bias: torch.Tensor | None = None,
1523
- ) -> tuple[torch.Tensor, torch.Tensor]:
1524
- return self.native_impl(
1525
- hidden_states,
1526
- gating_output,
1527
- self.topk,
1528
- self.renormalize,
1529
- self.num_expert_group,
1530
- self.topk_group,
1531
- self.scoring_func,
1532
- self.routed_scaling_factor,
1533
- e_score_correction_bias,
1534
- )
1535
-
1536
- def forward_cuda(
1537
- self,
1538
- hidden_states: torch.Tensor,
1539
- gating_output: torch.Tensor,
1540
- e_score_correction_bias: torch.Tensor | None = None,
1541
- ) -> tuple[torch.Tensor, torch.Tensor]:
1542
- return self.forward_native(
1543
- hidden_states, gating_output, e_score_correction_bias
1544
- )
1545
-
1546
- def forward_hip(
1547
- self,
1548
- hidden_states: torch.Tensor,
1549
- gating_output: torch.Tensor,
1550
- e_score_correction_bias: torch.Tensor | None = None,
1551
- ) -> tuple[torch.Tensor, torch.Tensor]:
1552
- if rocm_aiter_ops.is_fused_moe_enabled():
1553
- if not rocm_aiter_ops.is_fusion_moe_shared_experts_enabled():
1554
- assert self.num_fused_shared_experts == 0
1555
- return rocm_aiter_grouped_topk(
1556
- hidden_states,
1557
- gating_output,
1558
- self.topk,
1559
- self.renormalize,
1560
- self.num_expert_group,
1561
- self.topk_group,
1562
- self.scoring_func,
1563
- self.routed_scaling_factor,
1564
- e_score_correction_bias,
1565
- self.num_fused_shared_experts,
1566
- )
1567
- else:
1568
- return self.forward_native(
1569
- hidden_states, gating_output, e_score_correction_bias
1570
- )
1571
-
1572
-
1573
-@torch.compile(dynamic=True, backend=current_platform.simple_compile_backend)
1574
-def eplb_map_to_physical_and_record(
1575
- topk_ids: torch.Tensor,
1576
- expert_load_view: torch.Tensor,
1577
- logical_to_physical_map: torch.Tensor,
1578
- logical_replica_count: torch.Tensor,
1579
-) -> torch.Tensor:
1580
- """
1581
- Map the logical expert ids to physical expert ids
1582
- and record the expert load metrics.
1583
-
1584
- This will select a pseudo-random replica for each logical expert.
1585
- Only used for EPLB.
1586
-
1587
- Args:
1588
- topk_ids: The logical expert ids.
1589
- expert_load_view: The expert load view.
1590
- logical_to_physical_map: The logical to physical map.
1591
- logical_replica_count: The logical replica count.
1592
-
1593
- Returns:
1594
- The physical expert ids.
1595
- """
1596
-
1597
- # 1. Convert the logical expert ids to physical expert ids
1598
- # Directly select a random replica for each logical expert
1599
-
1600
- # In case `indices_type` is not `torch.long` or `torch.int`,
1601
- # e.g. `torch.uint32` as required by dispatch/combine kernels
1602
- topk_ids_long = topk_ids.long()
1603
- # Use (token position) modulo (replica count)
1604
- # to deterministically choose a replica
1605
- replica_count = logical_replica_count[topk_ids_long]
1606
- # Flatten-position based index, reshaped back to `topk_ids` shape
1607
- pos_indices = torch.arange(
1608
- topk_ids.numel(), device=topk_ids.device, dtype=torch.long
1609
- ).reshape_as(topk_ids)
1610
- # Compute pseudo-random indices by modulo
1611
- replica_indices = (pos_indices % replica_count).unsqueeze(-1)
1612
- physical_ids = (
1613
- logical_to_physical_map[topk_ids_long].gather(-1, replica_indices).squeeze(-1)
1614
- )
1615
-
1616
- topk_ids = physical_ids
1617
-
1618
- # 2. Record expert load metrics.
1619
-
1620
- # TODO(bowen): When using `FusedMoEModularKernel`, this
1621
- # can be done in a more unified way, since
1622
- # `FusedMoEPrepareAndFinalize` will return the expert
1623
- # token count, in some cases directly from the kernel.
1624
- # However, now there are many code paths not using
1625
- # the modular kernel, e.g. calling `fused_experts`,
1626
- # so we decide to keep the logic here.
1627
- #
1628
- # If later refactor moved all the MoE kernel calls
1629
- # to the modular kernel, we can move this logic there
1630
- # to achieve better efficiency.
1631
-
1632
- # `expert_load_view`: (num_physical_experts,)
1633
-
1634
- # `torch.bincount` is not compilable, so use `scatter_add_` instead.
1635
- topk_ids_flatten = topk_ids.flatten()
1636
- expert_load_view.scatter_add_(
1637
- dim=0,
1638
- index=topk_ids_flatten.long(),
1639
- src=torch.ones_like(topk_ids_flatten).to(expert_load_view),
1640
- )
1641
- return topk_ids
1642
-
1643
-
1644
-def fused_grouped_topk(
1645
- hidden_states: torch.Tensor,
1646
- gating_output: torch.Tensor,
1647
- topk: int,
1648
- renormalize: bool,
1649
- e_score_correction_bias: torch.Tensor,
1650
- num_expert_group: int = 0,
1651
- topk_group: int = 0,
1652
- scoring_func: str = "softmax",
1653
- routed_scaling_factor: float = 1.0,
1654
-) -> tuple[torch.Tensor, torch.Tensor]:
1655
- assert hidden_states.size(0) == gating_output.size(0), "Number of tokens mismatch"
1656
-
1657
- if scoring_func == "sigmoid":
1658
- # Fully fused kernel path for sigmoid
1659
- topk_values, topk_indices = ops.grouped_topk(
1660
- gating_output, # raw logits
1661
- num_expert_group,
1662
- topk_group,
1663
- topk,
1664
- renormalize,
1665
- routed_scaling_factor,
1666
- e_score_correction_bias,
1667
- 1, # scoring_func=1 for sigmoid
1668
- )
1669
- elif scoring_func == "softmax":
1670
- # Apply softmax in Python, then use fused kernel
1671
- # TODO: Add support for softmax in kernel
1672
- scores = torch.softmax(gating_output, dim=-1)
1673
- topk_values, topk_indices = ops.grouped_topk(
1674
- scores, # pre-computed scores
1675
- num_expert_group,
1676
- topk_group,
1677
- topk,
1678
- renormalize,
1679
- routed_scaling_factor,
1680
- e_score_correction_bias,
1681
- 0, # scoring_func=0 (no activation, scores already computed)
1682
- )
1683
- else:
1684
- raise ValueError(f"Unsupported scoring function: {scoring_func}")
1685
-
1686
- # Fused kernel outputs float32 values and int32 indices directly
1687
- return topk_values, topk_indices
1688
-
1689
-
1315
def inplace_fused_experts(
1316
hidden_states: torch.Tensor,
1317
w1: torch.Tensor,
vllm/model_executor/layers/fused_moe/fused_moe_method_base.py
+3
-3
@@ -10,13 +10,13 @@ from vllm.model_executor.layers.fused_moe.config import (
10
FusedMoEConfig,
11
FusedMoEQuantConfig,
12
)
13
-from vllm.model_executor.layers.fused_moe.fused_moe_router import (
14
- FusedMoERouter,
15
-)
13
from vllm.model_executor.layers.fused_moe.modular_kernel import (
14
FusedMoEPermuteExpertsUnpermute,
15
FusedMoEPrepareAndFinalize,
16
)
17
+from vllm.model_executor.layers.fused_moe.router.fused_moe_router import (
18
+ FusedMoERouter,
19
+)
20
from vllm.model_executor.layers.quantization.base_config import (
21
QuantizeMethodBase,
22
)
vllm/model_executor/layers/fused_moe/fused_moe_modular_method.py
+3
-1
@@ -12,11 +12,13 @@ from vllm.model_executor.layers.fused_moe.config import (
12
from vllm.model_executor.layers.fused_moe.fused_moe_method_base import (
13
FusedMoEMethodBase,
14
)
15
-from vllm.model_executor.layers.fused_moe.fused_moe_router import FusedMoERouter
15
from vllm.model_executor.layers.fused_moe.modular_kernel import (
16
FusedMoEModularKernel,
17
FusedMoEPrepareAndFinalize,
18
)
19
+from vllm.model_executor.layers.fused_moe.router.fused_moe_router import (
20
+ FusedMoERouter,
21
+)
22
23
logger = init_logger(__name__)
24
vllm/model_executor/layers/fused_moe/layer.py
+60
-224
@@ -21,7 +21,7 @@ from vllm.distributed import (
21
get_tensor_model_parallel_world_size,
22
tensor_model_parallel_all_reduce,
23
)
24
-from vllm.distributed.eplb.eplb_state import EplbState
24
+from vllm.distributed.eplb.eplb_state import EplbLayerState, EplbState
25
from vllm.forward_context import ForwardContext, get_forward_context
26
from vllm.logger import init_logger
27
from vllm.model_executor.custom_op import CustomOp
@@ -31,14 +31,24 @@ from vllm.model_executor.layers.fused_moe.config import (
31
FusedMoEQuantConfig,
32
RoutingMethodType,
33
)
34
-from vllm.model_executor.layers.fused_moe.fused_moe_router import FusedMoERouter
34
+from vllm.model_executor.layers.fused_moe.fused_moe_method_base import (
35
+ FusedMoEMethodBase,
36
+)
37
+from vllm.model_executor.layers.fused_moe.fused_moe_modular_method import (
38
+ FusedMoEModularMethod,
39
+)
40
from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (
41
init_aiter_topK_meta_data,
42
)
43
from vllm.model_executor.layers.fused_moe.routed_experts_capturer import (
44
RoutedExpertsCapturer,
45
)
41
-from vllm.model_executor.layers.fused_moe.routing_simulator import RoutingSimulator
46
+from vllm.model_executor.layers.fused_moe.router.router_factory import (
47
+ create_fused_moe_router,
48
+)
49
+from vllm.model_executor.layers.fused_moe.unquantized_fused_moe_method import (
50
+ UnquantizedFusedMoEMethod,
51
+)
52
from vllm.model_executor.layers.quantization.base_config import (
53
QuantizationConfig,
54
)
@@ -52,31 +62,6 @@ from vllm.utils.torch_utils import (
62
)
63
from vllm.v1.worker.ubatching import dbo_current_ubatch_id
64
55
-if current_platform.is_cuda_alike():
56
- from .fused_moe import eplb_map_to_physical_and_record
57
-else:
58
-
59
- def _eplb_map_to_physical_and_record(
60
- topk_ids: torch.Tensor,
61
- expert_load_view: torch.Tensor,
62
- logical_to_physical_map: torch.Tensor,
63
- logical_replica_count: torch.Tensor,
64
- ) -> torch.Tensor:
65
- # CPU fallback: no EPLB so just return as is
66
- return topk_ids
67
-
68
- eplb_map_to_physical_and_record = _eplb_map_to_physical_and_record
69
-from vllm.model_executor.layers.fused_moe.fused_moe import GroupedTopk
70
-from vllm.model_executor.layers.fused_moe.fused_moe_method_base import (
71
- FusedMoEMethodBase,
72
-)
73
-from vllm.model_executor.layers.fused_moe.fused_moe_modular_method import (
74
- FusedMoEModularMethod,
75
-)
76
-from vllm.model_executor.layers.fused_moe.unquantized_fused_moe_method import (
77
- UnquantizedFusedMoEMethod,
78
-)
79
-
65
logger = init_logger(__name__)
66
67
@@ -288,23 +273,6 @@ def maybe_roundup_hidden_size(
273
return hidden_size
274
275
291
-class FusedMoERouterImpl(FusedMoERouter):
292
- def __init__(self, layer: "FusedMoE"):
293
- super().__init__()
294
- self.layer = layer
295
-
296
- @property
297
- def routing_method_type(self) -> RoutingMethodType:
298
- return self.layer.routing_method_type
299
-
300
- def select_experts(
301
- self,
302
- hidden_states: torch.Tensor,
303
- router_logits: torch.Tensor,
304
- ) -> tuple[torch.Tensor, torch.Tensor]:
305
- return self.layer._select_experts(hidden_states, router_logits)
306
-
307
-
276
# --8<-- [start:fused_moe]
277
@CustomOp.register("fused_moe")
278
class FusedMoE(CustomOp):
@@ -440,9 +408,7 @@ class FusedMoE(CustomOp):
408
self.layer_name = prefix
409
410
self.enable_eplb = enable_eplb
443
- self.expert_load_view: torch.Tensor | None = None
444
- self.logical_to_physical_map: torch.Tensor | None = None
445
- self.logical_replica_count: torch.Tensor | None = None
411
+ self.eplb_state = EplbLayerState()
412
self.expert_placement_strategy: ExpertPlacementStrategy = (
413
vllm_config.parallel_config.expert_placement_strategy
414
)
@@ -538,6 +504,8 @@ class FusedMoE(CustomOp):
504
self.intermediate_size_per_partition = intermediate_size // self.tp_size
505
self.reduce_results = reduce_results
506
self.renormalize = renormalize
507
+
508
+ # TODO(bnell): these attributes are only used by cpu/xpu/mxfp4
509
self.use_grouped_topk = use_grouped_topk
510
if self.use_grouped_topk:
511
assert num_expert_group is not None and topk_group is not None
@@ -547,46 +515,11 @@ class FusedMoE(CustomOp):
515
self.scoring_func = scoring_func
516
self.routed_scaling_factor = routed_scaling_factor
517
self.e_score_correction_bias = e_score_correction_bias
518
+ # TODO(bnell): end attributes
519
+
520
self.apply_router_weight_on_input = apply_router_weight_on_input
521
self.activation = activation
522
553
- self._grouped_topk_impl: GroupedTopk | None = None
554
- if self.use_grouped_topk:
555
- assert self.num_expert_group is not None
556
- assert self.topk_group is not None
557
- self._grouped_topk_impl = GroupedTopk(
558
- topk=self.top_k,
559
- renormalize=self.renormalize,
560
- num_expert_group=self.num_expert_group,
561
- topk_group=self.topk_group,
562
- scoring_func=self.scoring_func,
563
- routed_scaling_factor=self.routed_scaling_factor,
564
- num_fused_shared_experts=self.num_fused_shared_experts,
565
- )
566
-
567
- if self.scoring_func != "softmax" and not self.use_grouped_topk:
568
- raise ValueError(
569
- "Only softmax scoring function is supported for non-grouped topk."
570
- )
571
-
572
- # ToDo: Better logic to determine the routing method type
573
- if routing_method_type is not None:
574
- self.routing_method_type: RoutingMethodType = routing_method_type
575
- else:
576
- if scoring_func == "sigmoid":
577
- if self.use_grouped_topk:
578
- self.routing_method_type = RoutingMethodType.DeepSeekV3
579
- elif self.top_k == 1:
580
- self.routing_method_type = RoutingMethodType.Llama4
581
- elif self.scoring_func == "softmax":
582
- self.routing_method_type = (
583
- RoutingMethodType.Renormalize
584
- if not self.renormalize
585
- else RoutingMethodType.RenormalizeNaive
586
- )
587
- else:
588
- self.routing_method_type = RoutingMethodType.TopK
589
-
523
self.moe_config: FusedMoEConfig = FusedMoEConfig(
524
num_experts=self.global_num_experts,
525
experts_per_token=top_k,
@@ -637,8 +570,7 @@ class FusedMoE(CustomOp):
570
# If you plan to add support for more quantization methods,
571
# please refer to the implementation in `Fp8MoEMethod`.
572
raise NotImplementedError(
640
- f"EPLB is not supported {self.quant_method.__class__.__name__}. "
641
- "EPLB is only supported for FP8 quantization for now."
573
+ f"EPLB is not supported {self.quant_method.__class__.__name__}."
574
)
575
576
moe_quant_params = {
@@ -663,7 +595,38 @@ class FusedMoE(CustomOp):
595
self.batched_hidden_states: torch.Tensor | None = None
596
self.batched_router_logits: torch.Tensor | None = None
597
666
- self.router = FusedMoERouterImpl(self)
598
+ # TODO(bnell): in next PR move capture back to layer
599
+ capture: Callable[[torch.Tensor], None] | None = None
600
+ if (
601
+ self.vllm_config.model_config is not None
602
+ and self.vllm_config.model_config.enable_return_routed_experts
603
+ ):
604
+ # In dummy runs, the capturer is not initialized.
605
+ capturer = RoutedExpertsCapturer.get_instance()
606
+ if capturer is not None:
607
+ capture = lambda topk_ids: capturer.capture(self.layer_id, topk_ids)
608
+
609
+ self.router = create_fused_moe_router(
610
+ top_k=top_k,
611
+ global_num_experts=self.global_num_experts,
612
+ eplb_state=self.eplb_state,
613
+ renormalize=renormalize,
614
+ use_grouped_topk=use_grouped_topk,
615
+ num_expert_group=num_expert_group,
616
+ topk_group=topk_group,
617
+ custom_routing_function=custom_routing_function,
618
+ scoring_func=scoring_func,
619
+ routed_scaling_factor=routed_scaling_factor,
620
+ e_score_correction_bias=e_score_correction_bias,
621
+ num_fused_shared_experts=self.num_fused_shared_experts,
622
+ enable_eplb=enable_eplb,
623
+ # TODO(bnell): once we can construct the MK at init time, we
624
+ # can make this a value.
625
+ indices_type_getter=lambda: self.quant_method.topk_indices_dtype,
626
+ routing_method_type=routing_method_type,
627
+ capture=capture,
628
+ )
629
+ self.routing_method_type: RoutingMethodType = self.router.routing_method_type
630
631
# Note: maybe_init_modular_kernel should only be called by
632
# prepare_communication_buffer_for_model.
@@ -1492,9 +1455,9 @@ class FusedMoE(CustomOp):
1455
This is used later in forward pass, where we get the expert mapping
1456
and record the load metrics in `expert_load_view`.
1457
"""
1495
- self.expert_load_view = expert_load_view[moe_layer_idx]
1496
- self.logical_to_physical_map = logical_to_physical_map[moe_layer_idx]
1497
- self.logical_replica_count = logical_replica_count[moe_layer_idx]
1458
+ self.eplb_state.expert_load_view = expert_load_view[moe_layer_idx]
1459
+ self.eplb_state.logical_to_physical_map = logical_to_physical_map[moe_layer_idx]
1460
+ self.eplb_state.logical_replica_count = logical_replica_count[moe_layer_idx]
1461
1462
def ensure_moe_quant_config_init(self):
1463
if self.quant_method.moe_quant_config is None:
@@ -1535,130 +1498,6 @@ class FusedMoE(CustomOp):
1498
device=torch.cuda.current_device(),
1499
)
1500
1538
- def _select_experts(
1539
- self,
1540
- hidden_states: torch.Tensor,
1541
- router_logits: torch.Tensor,
1542
- ) -> tuple[torch.Tensor, torch.Tensor]:
1543
- """
1544
- Route the input hidden states to the top-k experts based on the
1545
- router logits.
1546
-
1547
- Returns:
1548
- (topk_weights, topk_ids)
1549
- (tuple[torch.Tensor, torch.Tensor]):
1550
- The weights and expert ids.
1551
-
1552
- **Compatibility**: When EPLB is not enabled, the returned ids are
1553
- equivalent to global logical ids, so should be compatible with
1554
- plain MoE implementations without redundant experts.
1555
- """
1556
- from vllm.model_executor.layers.fused_moe.fused_moe import (
1557
- fused_topk,
1558
- fused_topk_bias,
1559
- )
1560
-
1561
- if self.enable_eplb:
1562
- if self.quant_method.supports_eplb:
1563
- if self.expert_load_view is None:
1564
- raise ValueError(
1565
- "enable_eplb=True requiere expert_load_view != None"
1566
- )
1567
- if self.logical_to_physical_map is None:
1568
- raise ValueError(
1569
- "enable_eplb=True requiere logical_to_physical_map != None"
1570
- )
1571
- if self.logical_replica_count is None:
1572
- raise ValueError(
1573
- "enable_eplb=True requiere logical_replica_count != None"
1574
- )
1575
- else:
1576
- raise NotImplementedError(
1577
- f"EPLB is not supported for {self.quant_method.method_name}."
1578
- )
1579
-
1580
- def valid_grouping() -> bool:
1581
- # Check if num_experts is greater than num_expert_group
1582
- # and is divisible by num_expert_group
1583
- num_experts = router_logits.shape[-1]
1584
- if num_experts <= self.num_expert_group:
1585
- return False
1586
- return num_experts % self.num_expert_group == 0
1587
-
1588
- indices_type = self.quant_method.topk_indices_dtype
1589
-
1590
- # Check if we should use a routing simulation strategy
1591
- routing_strategy = envs.VLLM_MOE_ROUTING_SIMULATION_STRATEGY
1592
- if routing_strategy != "":
1593
- topk_weights, topk_ids = RoutingSimulator.simulate_routing(
1594
- hidden_states=hidden_states,
1595
- router_logits=router_logits,
1596
- strategy_name=routing_strategy,
1597
- top_k=self.top_k,
1598
- indices_type=indices_type,
1599
- )
1600
-
1601
- # DeepSeekv2 uses grouped_top_k
1602
- elif self.use_grouped_topk and valid_grouping():
1603
- assert self._grouped_topk_impl is not None
1604
- topk_weights, topk_ids = self._grouped_topk_impl(
1605
- hidden_states=hidden_states,
1606
- gating_output=router_logits,
1607
- e_score_correction_bias=self.e_score_correction_bias,
1608
- )
1609
- elif self.e_score_correction_bias is not None:
1610
- topk_weights, topk_ids = fused_topk_bias(
1611
- hidden_states=hidden_states,
1612
- gating_output=router_logits,
1613
- e_score_correction_bias=self.e_score_correction_bias.data,
1614
- topk=self.top_k,
1615
- renormalize=self.renormalize,
1616
- )
1617
- if self.routed_scaling_factor != 1.0:
1618
- topk_weights *= self.routed_scaling_factor
1619
- elif self.custom_routing_function is None:
1620
- topk_weights, topk_ids, token_expert_indices = fused_topk(
1621
- hidden_states=hidden_states,
1622
- gating_output=router_logits,
1623
- topk=self.top_k,
1624
- renormalize=self.renormalize,
1625
- indices_type=indices_type,
1626
- )
1627
- else:
1628
- topk_weights, topk_ids = self.custom_routing_function(
1629
- hidden_states=hidden_states,
1630
- gating_output=router_logits,
1631
- topk=self.top_k,
1632
- renormalize=self.renormalize,
1633
- )
1634
-
1635
- if self.enable_eplb:
1636
- topk_ids = eplb_map_to_physical_and_record(
1637
- topk_ids=topk_ids,
1638
- expert_load_view=self.expert_load_view,
1639
- logical_to_physical_map=self.logical_to_physical_map,
1640
- logical_replica_count=self.logical_replica_count,
1641
- )
1642
-
1643
- if (indices_type is not None) and topk_ids.dtype != indices_type:
1644
- topk_ids = topk_ids.to(dtype=indices_type)
1645
-
1646
- assert topk_ids.dtype == indices_type or indices_type is None
1647
-
1648
- if (
1649
- self.vllm_config.model_config is not None
1650
- and self.vllm_config.model_config.enable_return_routed_experts
1651
- ):
1652
- # In dummy runs, the capturer is not initialized.
1653
- capturer = RoutedExpertsCapturer.get_instance()
1654
- if capturer is not None: # in dummmy_run may be None
1655
- capturer.capture( # noqa
1656
- layer_id=self.layer_id,
1657
- topk_ids=topk_ids,
1658
- )
1659
-
1660
- return topk_weights, topk_ids
1661
-
1501
def must_reduce_shared_expert_outputs(self) -> bool:
1502
"""
1503
The shared_experts are typically computed using the RowParallelLinear
@@ -1761,8 +1600,12 @@ class FusedMoE(CustomOp):
1600
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
1601
assert self.batched_hidden_states is not None
1602
assert self.batched_router_logits is not None
1764
- assert self.batched_hidden_states.dtype == full_hidden_states.dtype
1765
- assert self.batched_router_logits.dtype == full_router_logits.dtype
1603
+ assert self.batched_hidden_states.dtype == full_hidden_states.dtype, (
1604
+ f"{self.batched_hidden_states.dtype} == {full_hidden_states.dtype}"
1605
+ )
1606
+ assert self.batched_router_logits.dtype == full_router_logits.dtype, (
1607
+ f"{self.batched_router_logits.dtype} == {full_router_logits.dtype}"
1608
+ )
1609
# Check size compatibility.
1610
assert self.batched_hidden_states.size(-1) == full_hidden_states.size(-1)
1611
assert self.batched_router_logits.size(-1) == full_router_logits.size(-1)
@@ -2080,15 +1923,8 @@ class FusedMoE(CustomOp):
1923
f"tp_size={self.tp_size},\n"
1924
f"ep_size={self.ep_size}, "
1925
f"reduce_results={self.reduce_results}, "
2083
- f"renormalize={self.renormalize}, "
2084
- f"use_grouped_topk={self.use_grouped_topk}"
1926
)
1927
2087
- if self.use_grouped_topk:
2088
- s += f", num_expert_group={self.num_expert_group}, topk_group={self.topk_group}" # noqa: E501
2089
-
2090
- s += f", scoring_func='{self.scoring_func}', activation='{self.activation}'" # noqa: E501
2091
-
1928
return s
1929
1930
vllm/model_executor/layers/fused_moe/router/__init__.py
new
+2
@@ -0,0 +1,2 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
vllm/model_executor/layers/fused_moe/router/base_router.py
new
+245
@@ -0,0 +1,245 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from abc import abstractmethod
4
+from collections.abc import Callable
5
+
6
+import torch
7
+
8
+from vllm.distributed.eplb.eplb_state import EplbLayerState
9
+from vllm.model_executor.layers.fused_moe.router.fused_moe_router import (
10
+ FusedMoERouter,
11
+)
12
+from vllm.platforms import current_platform
13
+
14
+if current_platform.is_cuda_alike():
15
+
16
+ @torch.compile(dynamic=True, backend=current_platform.simple_compile_backend)
17
+ def eplb_map_to_physical_and_record(
18
+ topk_ids: torch.Tensor,
19
+ expert_load_view: torch.Tensor,
20
+ logical_to_physical_map: torch.Tensor,
21
+ logical_replica_count: torch.Tensor,
22
+ ) -> torch.Tensor:
23
+ """
24
+ Map the logical expert ids to physical expert ids
25
+ and record the expert load metrics.
26
+
27
+ This will select a pseudo-random replica for each logical expert.
28
+ Only used for EPLB.
29
+
30
+ Args:
31
+ topk_ids: The logical expert ids.
32
+ expert_load_view: The expert load view.
33
+ logical_to_physical_map: The logical to physical map.
34
+ logical_replica_count: The logical replica count.
35
+
36
+ Returns:
37
+ The physical expert ids.
38
+ """
39
+
40
+ # 1. Convert the logical expert ids to physical expert ids
41
+ # Directly select a random replica for each logical expert
42
+
43
+ # In case `indices_type` is not `torch.long` or `torch.int`,
44
+ # e.g. `torch.uint32` as required by dispatch/combine kernels
45
+ topk_ids_long = topk_ids.long()
46
+ # Use (token position) modulo (replica count)
47
+ # to deterministically choose a replica
48
+ replica_count = logical_replica_count[topk_ids_long]
49
+ # Flatten-position based index, reshaped back to `topk_ids` shape
50
+ pos_indices = torch.arange(
51
+ topk_ids.numel(), device=topk_ids.device, dtype=torch.long
52
+ ).reshape_as(topk_ids)
53
+ # Compute pseudo-random indices by modulo
54
+ replica_indices = (pos_indices % replica_count).unsqueeze(-1)
55
+ physical_ids = (
56
+ logical_to_physical_map[topk_ids_long]
57
+ .gather(-1, replica_indices)
58
+ .squeeze(-1)
59
+ )
60
+
61
+ topk_ids = physical_ids
62
+
63
+ # 2. Record expert load metrics.
64
+
65
+ # TODO(bowen): When using `FusedMoEModularKernel`, this
66
+ # can be done in a more unified way, since
67
+ # `FusedMoEPrepareAndFinalize` will return the expert
68
+ # token count, in some cases directly from the kernel.
69
+ # However, now there are many code paths not using
70
+ # the modular kernel, e.g. calling `fused_experts`,
71
+ # so we decide to keep the logic here.
72
+ #
73
+ # If later refactor moved all the MoE kernel calls
74
+ # to the modular kernel, we can move this logic there
75
+ # to achieve better efficiency.
76
+
77
+ # `expert_load_view`: (num_physical_experts,)
78
+
79
+ # `torch.bincount` is not compilable, so use `scatter_add_` instead.
80
+ topk_ids_flatten = topk_ids.flatten()
81
+ expert_load_view.scatter_add_(
82
+ dim=0,
83
+ index=topk_ids_flatten.long(),
84
+ src=torch.ones_like(topk_ids_flatten).to(expert_load_view),
85
+ )
86
+ return topk_ids
87
+else:
88
+
89
+ def eplb_map_to_physical_and_record(
90
+ topk_ids: torch.Tensor,
91
+ expert_load_view: torch.Tensor,
92
+ logical_to_physical_map: torch.Tensor,
93
+ logical_replica_count: torch.Tensor,
94
+ ) -> torch.Tensor:
95
+ # CPU fallback: no EPLB so just return as is
96
+ return topk_ids
97
+
98
+
99
+class BaseRouter(FusedMoERouter):
100
+ """
101
+ Base router class that provides common functionality for all router implementations.
102
+
103
+ This class implements the template method pattern where select_experts() handles
104
+ common pre-processing and post-processing, delegating the actual routing logic
105
+ to the abstract _compute_routing() method.
106
+ """
107
+
108
+ def __init__(
109
+ self,
110
+ top_k: int,
111
+ global_num_experts: int,
112
+ eplb_state: EplbLayerState,
113
+ enable_eplb: bool = False,
114
+ # TODO(bnell): Once the MK is constructed at layer init time, we
115
+ # can make this a plain value instead of a callback.
116
+ indices_type_getter: Callable[[], torch.dtype | None] | None = None,
117
+ ):
118
+ """
119
+ Note: the indices dtype might not be available at router construction
120
+ time, so we need to supply a callback to get it at runtime. This is
121
+ because the indices type is supplied by modular kernels which are
122
+ created after MoE layer/router construction.
123
+ """
124
+ super().__init__()
125
+ self.top_k = top_k
126
+ self.global_num_experts = global_num_experts
127
+ self.eplb_state = eplb_state
128
+ self.enable_eplb = enable_eplb
129
+ self.indices_type_getter = indices_type_getter
130
+ self.capture: Callable[[torch.tensor], None] | None = None
131
+
132
+ def _validate_eplb_state(self) -> None:
133
+ """Validate that EPLB state is properly initialized if EPLB is enabled."""
134
+ if self.enable_eplb:
135
+ if self.eplb_state.expert_load_view is None:
136
+ raise ValueError("enable_eplb=True requires expert_load_view != None")
137
+ if self.eplb_state.logical_to_physical_map is None:
138
+ raise ValueError(
139
+ "enable_eplb=True requires logical_to_physical_map != None"
140
+ )
141
+ if self.eplb_state.logical_replica_count is None:
142
+ raise ValueError(
143
+ "enable_eplb=True requires logical_replica_count != None"
144
+ )
145
+
146
+ def _get_indices_type(self) -> torch.dtype | None:
147
+ """Get the desired indices dtype from the getter function."""
148
+ return (
149
+ self.indices_type_getter() if self.indices_type_getter is not None else None
150
+ )
151
+
152
+ def _apply_eplb_mapping(self, topk_ids: torch.Tensor) -> torch.Tensor:
153
+ """Apply EPLB mapping to convert logical expert IDs to physical expert IDs."""
154
+ if self.enable_eplb:
155
+ assert self.eplb_state.expert_load_view is not None
156
+ assert self.eplb_state.logical_to_physical_map is not None
157
+ assert self.eplb_state.logical_replica_count is not None
158
+ return eplb_map_to_physical_and_record(
159
+ topk_ids=topk_ids,
160
+ expert_load_view=self.eplb_state.expert_load_view,
161
+ logical_to_physical_map=self.eplb_state.logical_to_physical_map,
162
+ logical_replica_count=self.eplb_state.logical_replica_count,
163
+ )
164
+ return topk_ids
165
+
166
+ def _convert_indices_dtype(
167
+ self, topk_ids: torch.Tensor, indices_type: torch.dtype | None
168
+ ) -> torch.Tensor:
169
+ """Convert topk_ids to the desired dtype if needed."""
170
+ if (indices_type is not None) and topk_ids.dtype != indices_type:
171
+ topk_ids = topk_ids.to(dtype=indices_type)
172
+
173
+ assert topk_ids.dtype == indices_type or indices_type is None
174
+ return topk_ids
175
+
176
+ @abstractmethod
177
+ def _compute_routing(
178
+ self,
179
+ hidden_states: torch.Tensor,
180
+ router_logits: torch.Tensor,
181
+ indices_type: torch.dtype | None,
182
+ ) -> tuple[torch.Tensor, torch.Tensor]:
183
+ """
184
+ Compute the actual routing logic.
185
+
186
+ This method must be implemented by subclasses to provide the specific
187
+ routing algorithm (e.g., grouped_topk, fused_topk, custom routing, etc.).
188
+
189
+ Args:
190
+ hidden_states: Input hidden states
191
+ router_logits: Router logits for expert selection
192
+ indices_type: Desired dtype for expert indices (may be None)
193
+
194
+ Returns:
195
+ tuple of (topk_weights, topk_ids)
196
+ """
197
+ raise NotImplementedError
198
+
199
+ def select_experts(
200
+ self,
201
+ hidden_states: torch.Tensor,
202
+ router_logits: torch.Tensor,
203
+ ) -> tuple[torch.Tensor, torch.Tensor]:
204
+ """
205
+ Route the input hidden states to the top-k experts based on the
206
+ router logits.
207
+
208
+ This method implements the template method pattern:
209
+ 1. Validates EPLB state
210
+ 2. Gets indices type
211
+ 3. Calls _compute_routing() to get topk_weights and topk_ids
212
+ 4. Applies EPLB mapping if enabled
213
+ 5. Converts indices dtype if needed
214
+
215
+ Returns:
216
+ (topk_weights, topk_ids)
217
+ (tuple[torch.Tensor, torch.Tensor]):
218
+ The weights and expert ids computation result.
219
+
220
+ **Compatibility**: When EPLB is not enabled, the returned ids are
221
+ equivalent to global logical ids, so should be compatible with
222
+ plain MoE implementations without redundant experts.
223
+ """
224
+ # Step 1: Validate EPLB state
225
+ self._validate_eplb_state()
226
+
227
+ # Step 2: Get indices type.
228
+ indices_type = self._get_indices_type()
229
+
230
+ # Step 3: Compute routing (delegated to subclass)
231
+ topk_weights, topk_ids = self._compute_routing(
232
+ hidden_states, router_logits, indices_type
233
+ )
234
+
235
+ # Step 4: Apply EPLB mapping
236
+ topk_ids = self._apply_eplb_mapping(topk_ids)
237
+
238
+ # Step 5: Convert indices dtype
239
+ topk_ids = self._convert_indices_dtype(topk_ids, indices_type)
240
+
241
+ # TODO(bnell): temporary hack until select_experts is moved into FusedMoE
242
+ if self.capture is not None:
243
+ self.capture(topk_ids)
244
+
245
+ return topk_weights, topk_ids
vllm/model_executor/layers/fused_moe/router/custom_routing_router.py
new
+53
@@ -0,0 +1,53 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from collections.abc import Callable
4
+
5
+import torch
6
+
7
+from vllm.distributed.eplb.eplb_state import EplbLayerState
8
+from vllm.model_executor.layers.fused_moe.config import RoutingMethodType
9
+from vllm.model_executor.layers.fused_moe.router.base_router import BaseRouter
10
+
11
+
12
+class CustomRoutingRouter(BaseRouter):
13
+ """Router using a custom user-provided routing function."""
14
+
15
+ def __init__(
16
+ self,
17
+ top_k: int,
18
+ global_num_experts: int,
19
+ eplb_state: EplbLayerState,
20
+ custom_routing_function: Callable,
21
+ renormalize: bool = True,
22
+ enable_eplb: bool = False,
23
+ indices_type_getter: Callable[[], torch.dtype | None] | None = None,
24
+ ):
25
+ super().__init__(
26
+ top_k=top_k,
27
+ global_num_experts=global_num_experts,
28
+ eplb_state=eplb_state,
29
+ enable_eplb=enable_eplb,
30
+ indices_type_getter=indices_type_getter,
31
+ )
32
+ self.custom_routing_function = custom_routing_function
33
+ self.renormalize = renormalize
34
+
35
+ @property
36
+ def routing_method_type(self) -> RoutingMethodType:
37
+ return RoutingMethodType.Custom
38
+
39
+ def _compute_routing(
40
+ self,
41
+ hidden_states: torch.Tensor,
42
+ router_logits: torch.Tensor,
43
+ indices_type: torch.dtype | None,
44
+ ) -> tuple[torch.Tensor, torch.Tensor]:
45
+ """Compute routing using the custom routing function."""
46
+ topk_weights, topk_ids = self.custom_routing_function(
47
+ hidden_states=hidden_states,
48
+ gating_output=router_logits,
49
+ topk=self.top_k,
50
+ renormalize=self.renormalize,
51
+ )
52
+
53
+ return topk_weights.to(torch.float32), topk_ids
vllm/model_executor/layers/fused_moe/router/fused_moe_router.py
renamed
vllm/model_executor/layers/fused_moe/router/fused_topk_bias_router.py
new
+88
@@ -0,0 +1,88 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from collections.abc import Callable
4
+
5
+import torch
6
+
7
+from vllm.distributed.eplb.eplb_state import EplbLayerState
8
+from vllm.model_executor.layers.batch_invariant import (
9
+ vllm_is_batch_invariant,
10
+)
11
+from vllm.model_executor.layers.fused_moe.config import RoutingMethodType
12
+from vllm.model_executor.layers.fused_moe.router.base_router import BaseRouter
13
+
14
+
15
+def fused_topk_bias(
16
+ hidden_states: torch.Tensor,
17
+ gating_output: torch.Tensor,
18
+ e_score_correction_bias: torch.Tensor,
19
+ topk: int,
20
+ renormalize: bool,
21
+):
22
+ n_routed_experts = gating_output.shape[-1]
23
+ scores = gating_output.softmax(dim=-1)
24
+ scores_for_choice = scores.view(
25
+ -1, n_routed_experts
26
+ ) + e_score_correction_bias.unsqueeze(0)
27
+
28
+ # For batch invariance, use sorted=True to ensure deterministic expert selection
29
+ use_sorted = vllm_is_batch_invariant()
30
+ topk_indices = torch.topk(scores_for_choice, k=topk, dim=-1, sorted=use_sorted)[1]
31
+ topk_weights = scores.gather(1, topk_indices)
32
+ if renormalize:
33
+ topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
34
+ return topk_weights.to(torch.float32), topk_indices.to(torch.int32)
35
+
36
+
37
+class FusedTopKBiasRouter(BaseRouter):
38
+ """Router using fused top-k with e_score_correction_bias."""
39
+
40
+ def __init__(
41
+ self,
42
+ top_k: int,
43
+ global_num_experts: int,
44
+ eplb_state: EplbLayerState,
45
+ e_score_correction_bias: torch.Tensor,
46
+ renormalize: bool = True,
47
+ routed_scaling_factor: float = 1.0,
48
+ enable_eplb: bool = False,
49
+ indices_type_getter: Callable[[], torch.dtype | None] | None = None,
50
+ ):
51
+ super().__init__(
52
+ top_k=top_k,
53
+ global_num_experts=global_num_experts,
54
+ eplb_state=eplb_state,
55
+ enable_eplb=enable_eplb,
56
+ indices_type_getter=indices_type_getter,
57
+ )
58
+ self.e_score_correction_bias = e_score_correction_bias
59
+ self.renormalize = renormalize
60
+ self.routed_scaling_factor = routed_scaling_factor
61
+
62
+ @property
63
+ def routing_method_type(self) -> RoutingMethodType:
64
+ return (
65
+ RoutingMethodType.Renormalize
66
+ if not self.renormalize
67
+ else RoutingMethodType.RenormalizeNaive
68
+ )
69
+
70
+ def _compute_routing(
71
+ self,
72
+ hidden_states: torch.Tensor,
73
+ router_logits: torch.Tensor,
74
+ indices_type: torch.dtype | None,
75
+ ) -> tuple[torch.Tensor, torch.Tensor]:
76
+ """Compute routing using fused top-k with bias."""
77
+ topk_weights, topk_ids = fused_topk_bias(
78
+ hidden_states=hidden_states,
79
+ gating_output=router_logits,
80
+ e_score_correction_bias=self.e_score_correction_bias.data,
81
+ topk=self.top_k,
82
+ renormalize=self.renormalize,
83
+ )
84
+
85
+ if self.routed_scaling_factor != 1.0:
86
+ topk_weights *= self.routed_scaling_factor
87
+
88
+ return topk_weights, topk_ids
vllm/model_executor/layers/fused_moe/router/fused_topk_router.py
new
+118
@@ -0,0 +1,118 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from collections.abc import Callable
4
+
5
+import torch
6
+
7
+import vllm._custom_ops as ops
8
+from vllm._aiter_ops import rocm_aiter_ops
9
+from vllm.distributed.eplb.eplb_state import EplbLayerState
10
+from vllm.model_executor.layers.fused_moe.config import RoutingMethodType
11
+from vllm.model_executor.layers.fused_moe.router.base_router import BaseRouter
12
+
13
+
14
+def vllm_topk_softmax(
15
+ topk_weights: torch.Tensor,
16
+ topk_indices: torch.Tensor,
17
+ token_expert_indices: torch.Tensor,
18
+ gating_output: torch.Tensor,
19
+ renormalize: bool,
20
+) -> tuple[torch.Tensor, ...]:
21
+ ops.topk_softmax(
22
+ topk_weights,
23
+ topk_indices,
24
+ token_expert_indices,
25
+ gating_output,
26
+ renormalize,
27
+ )
28
+
29
+ return topk_weights, topk_indices
30
+
31
+
32
+def dispatch_topk_func(
33
+ use_rocm_aiter: bool = False,
34
+) -> Callable[..., tuple[torch.Tensor, ...]]:
35
+ if use_rocm_aiter:
36
+ return rocm_aiter_ops.topk_softmax
37
+ return vllm_topk_softmax
38
+
39
+
40
+def fused_topk(
41
+ hidden_states: torch.Tensor,
42
+ gating_output: torch.Tensor,
43
+ topk: int,
44
+ renormalize: bool,
45
+ indices_type: torch.dtype | None = None,
46
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
47
+ assert hidden_states.size(0) == gating_output.size(0), "Number of tokens mismatch"
48
+
49
+ M, _ = hidden_states.size()
50
+
51
+ topk_weights = torch.empty(
52
+ M, topk, dtype=torch.float32, device=hidden_states.device
53
+ )
54
+ topk_ids = torch.empty(
55
+ M,
56
+ topk,
57
+ dtype=torch.int32 if indices_type is None else indices_type,
58
+ device=hidden_states.device,
59
+ )
60
+ token_expert_indices = torch.empty(
61
+ M, topk, dtype=torch.int32, device=hidden_states.device
62
+ )
63
+
64
+ topk_func = dispatch_topk_func(use_rocm_aiter=rocm_aiter_ops.is_fused_moe_enabled())
65
+ topk_weights, topk_ids = topk_func(
66
+ topk_weights, topk_ids, token_expert_indices, gating_output, renormalize
67
+ )
68
+
69
+ return topk_weights, topk_ids, token_expert_indices
70
+
71
+
72
+class FusedTopKRouter(BaseRouter):
73
+ """Default router using standard fused top-k routing."""
74
+
75
+ def __init__(
76
+ self,
77
+ top_k: int,
78
+ global_num_experts: int,
79
+ eplb_state: EplbLayerState,
80
+ scoring_func: str = "softmax",
81
+ renormalize: bool = True,
82
+ enable_eplb: bool = False,
83
+ indices_type_getter: Callable[[], torch.dtype | None] | None = None,
84
+ ):
85
+ assert scoring_func == "softmax", "FusedTopKRouter only supports softmax."
86
+ super().__init__(
87
+ top_k=top_k,
88
+ global_num_experts=global_num_experts,
89
+ eplb_state=eplb_state,
90
+ enable_eplb=enable_eplb,
91
+ indices_type_getter=indices_type_getter,
92
+ )
93
+ self.renormalize = renormalize
94
+
95
+ @property
96
+ def routing_method_type(self) -> RoutingMethodType:
97
+ return (
98
+ RoutingMethodType.Renormalize
99
+ if not self.renormalize
100
+ else RoutingMethodType.RenormalizeNaive
101
+ )
102
+
103
+ def _compute_routing(
104
+ self,
105
+ hidden_states: torch.Tensor,
106
+ router_logits: torch.Tensor,
107
+ indices_type: torch.dtype | None,
108
+ ) -> tuple[torch.Tensor, torch.Tensor]:
109
+ """Compute routing using standard fused top-k."""
110
+ topk_weights, topk_ids, token_expert_indices = fused_topk(
111
+ hidden_states=hidden_states,
112
+ gating_output=router_logits,
113
+ topk=self.top_k,
114
+ renormalize=self.renormalize,
115
+ indices_type=indices_type,
116
+ )
117
+
118
+ return topk_weights, topk_ids
vllm/model_executor/layers/fused_moe/router/grouped_topk_router.py
new
+353
@@ -0,0 +1,353 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from collections.abc import Callable
4
+from functools import partial
5
+
6
+import torch
7
+
8
+from vllm import _custom_ops as ops
9
+from vllm import envs as envs
10
+from vllm._aiter_ops import rocm_aiter_ops
11
+from vllm.distributed.eplb.eplb_state import EplbLayerState
12
+from vllm.model_executor.custom_op import CustomOp
13
+from vllm.model_executor.layers.batch_invariant import (
14
+ vllm_is_batch_invariant,
15
+)
16
+from vllm.model_executor.layers.fused_moe.config import RoutingMethodType
17
+from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (
18
+ rocm_aiter_grouped_topk,
19
+)
20
+from vllm.model_executor.layers.fused_moe.router.base_router import BaseRouter
21
+from vllm.model_executor.layers.fused_moe.router.fused_topk_bias_router import (
22
+ fused_topk_bias,
23
+)
24
+from vllm.model_executor.layers.fused_moe.router.fused_topk_router import fused_topk
25
+from vllm.model_executor.utils import maybe_disable_graph_partition
26
+from vllm.platforms import current_platform
27
+
28
+
29
+def fused_grouped_topk(
30
+ hidden_states: torch.Tensor,
31
+ gating_output: torch.Tensor,
32
+ topk: int,
33
+ renormalize: bool,
34
+ e_score_correction_bias: torch.Tensor,
35
+ num_expert_group: int = 0,
36
+ topk_group: int = 0,
37
+ scoring_func: str = "softmax",
38
+ routed_scaling_factor: float = 1.0,
39
+) -> tuple[torch.Tensor, torch.Tensor]:
40
+ assert hidden_states.size(0) == gating_output.size(0), "Number of tokens mismatch"
41
+
42
+ if scoring_func == "sigmoid":
43
+ # Fully fused kernel path for sigmoid
44
+ topk_values, topk_indices = ops.grouped_topk(
45
+ gating_output, # raw logits
46
+ num_expert_group,
47
+ topk_group,
48
+ topk,
49
+ renormalize,
50
+ routed_scaling_factor,
51
+ e_score_correction_bias,
52
+ 1, # scoring_func=1 for sigmoid
53
+ )
54
+ elif scoring_func == "softmax":
55
+ # Apply softmax in Python, then use fused kernel
56
+ # TODO: Add support for softmax in kernel
57
+ scores = torch.softmax(gating_output, dim=-1)
58
+ topk_values, topk_indices = ops.grouped_topk(
59
+ scores, # pre-computed scores
60
+ num_expert_group,
61
+ topk_group,
62
+ topk,
63
+ renormalize,
64
+ routed_scaling_factor,
65
+ e_score_correction_bias,
66
+ 0, # scoring_func=0 (no activation, scores already computed)
67
+ )
68
+ else:
69
+ raise ValueError(f"Unsupported scoring function: {scoring_func}")
70
+
71
+ # Fused kernel outputs float32 values and int32 indices directly
72
+ return topk_values, topk_indices
73
+
74
+
75
+# This is used by the Deepseek-V2 and Deepseek-V3 model
76
+@torch.compile(
77
+ dynamic=True,
78
+ backend=current_platform.simple_compile_backend,
79
+ options=maybe_disable_graph_partition(current_platform.simple_compile_backend),
80
+)
81
+def grouped_topk(
82
+ hidden_states: torch.Tensor,
83
+ gating_output: torch.Tensor,
84
+ topk: int,
85
+ renormalize: bool,
86
+ num_expert_group: int = 0,
87
+ topk_group: int = 0,
88
+ scoring_func: str = "softmax",
89
+ routed_scaling_factor: float = 1.0,
90
+ e_score_correction_bias: torch.Tensor | None = None,
91
+) -> tuple[torch.Tensor, torch.Tensor]:
92
+ if (
93
+ envs.VLLM_USE_FUSED_MOE_GROUPED_TOPK
94
+ and current_platform.is_cuda()
95
+ and num_expert_group <= 32
96
+ and topk <= 32
97
+ and e_score_correction_bias is not None
98
+ ):
99
+ return fused_grouped_topk(
100
+ hidden_states=hidden_states,
101
+ gating_output=gating_output,
102
+ topk=topk,
103
+ renormalize=renormalize,
104
+ e_score_correction_bias=e_score_correction_bias,
105
+ num_expert_group=num_expert_group,
106
+ topk_group=topk_group,
107
+ scoring_func=scoring_func,
108
+ routed_scaling_factor=routed_scaling_factor,
109
+ )
110
+
111
+ assert hidden_states.size(0) == gating_output.size(0), "Number of tokens mismatch"
112
+
113
+ if scoring_func == "softmax":
114
+ scores = torch.softmax(gating_output, dim=-1)
115
+ elif scoring_func == "sigmoid":
116
+ scores = gating_output.sigmoid()
117
+ else:
118
+ raise ValueError(f"Unsupported scoring function: {scoring_func}")
119
+
120
+ num_token = scores.size(0)
121
+ if e_score_correction_bias is not None:
122
+ # Store original scores before applying correction bias. We use biased
123
+ # scores for expert selection but original scores for routing weights
124
+ original_scores = scores
125
+ scores = scores + e_score_correction_bias.unsqueeze(0)
126
+ group_scores = (
127
+ scores.view(num_token, num_expert_group, -1).topk(2, dim=-1)[0].sum(dim=-1)
128
+ )
129
+ else:
130
+ group_scores = (
131
+ scores.view(num_token, num_expert_group, -1).max(dim=-1).values
132
+ ) # [n, n_group]
133
+
134
+ # For batch invariance, use sorted=True to ensure deterministic expert selection
135
+ use_sorted = vllm_is_batch_invariant()
136
+ group_idx = torch.topk(group_scores, k=topk_group, dim=-1, sorted=use_sorted)[
137
+ 1
138
+ ] # [n, top_k_group]
139
+ group_mask = torch.zeros_like(group_scores) # [n, n_group]
140
+ group_mask.scatter_(1, group_idx, 1) # [n, n_group]
141
+ score_mask = (
142
+ group_mask.unsqueeze(-1)
143
+ .expand(num_token, num_expert_group, scores.size(-1) // num_expert_group)
144
+ .reshape(num_token, -1)
145
+ ) # [n, e]
146
+ tmp_scores = scores.masked_fill(~score_mask.bool(), float("-inf")) # [n, e]
147
+
148
+ if e_score_correction_bias is not None:
149
+ topk_ids = torch.topk(tmp_scores, k=topk, dim=-1, sorted=use_sorted)[1]
150
+ # Use original unbiased scores for the routing weights
151
+ topk_weights = original_scores.gather(1, topk_ids)
152
+ else:
153
+ topk_weights, topk_ids = torch.topk(
154
+ tmp_scores, k=topk, dim=-1, sorted=use_sorted
155
+ )
156
+
157
+ if renormalize:
158
+ topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
159
+
160
+ if routed_scaling_factor != 1.0:
161
+ topk_weights = topk_weights * routed_scaling_factor
162
+ return topk_weights.to(torch.float32), topk_ids.to(torch.int32)
163
+
164
+
165
+# --8<-- [start:grouped_topk]
166
+@CustomOp.register("grouped_topk")
167
+class GroupedTopk(CustomOp):
168
+ """GroupedTopk used by the Deepseek-V2 and Deepseek-V3 model."""
169
+
170
+ # --8<-- [end:grouped_topk]
171
+
172
+ def __init__(
173
+ self,
174
+ topk: int,
175
+ renormalize: bool,
176
+ num_expert_group: int = 0,
177
+ topk_group: int = 0,
178
+ scoring_func: str = "softmax",
179
+ routed_scaling_factor: float = 1.0,
180
+ num_fused_shared_experts: int = 0,
181
+ ) -> None:
182
+ super().__init__()
183
+ self.native_impl = grouped_topk
184
+ self.topk = topk
185
+ self.renormalize = renormalize
186
+ self.num_expert_group = num_expert_group
187
+ self.topk_group = topk_group
188
+ self.scoring_func = scoring_func
189
+ self.routed_scaling_factor = routed_scaling_factor
190
+ self.num_fused_shared_experts = num_fused_shared_experts
191
+
192
+ def forward_native(
193
+ self,
194
+ hidden_states: torch.Tensor,
195
+ gating_output: torch.Tensor,
196
+ e_score_correction_bias: torch.Tensor | None = None,
197
+ ) -> tuple[torch.Tensor, torch.Tensor]:
198
+ return self.native_impl(
199
+ hidden_states,
200
+ gating_output,
201
+ self.topk,
202
+ self.renormalize,
203
+ self.num_expert_group,
204
+ self.topk_group,
205
+ self.scoring_func,
206
+ self.routed_scaling_factor,
207
+ e_score_correction_bias,
208
+ )
209
+
210
+ def forward_cuda(
211
+ self,
212
+ hidden_states: torch.Tensor,
213
+ gating_output: torch.Tensor,
214
+ e_score_correction_bias: torch.Tensor | None = None,
215
+ ) -> tuple[torch.Tensor, torch.Tensor]:
216
+ return self.forward_native(
217
+ hidden_states, gating_output, e_score_correction_bias
218
+ )
219
+
220
+ def forward_hip(
221
+ self,
222
+ hidden_states: torch.Tensor,
223
+ gating_output: torch.Tensor,
224
+ e_score_correction_bias: torch.Tensor | None = None,
225
+ ) -> tuple[torch.Tensor, torch.Tensor]:
226
+ if rocm_aiter_ops.is_fused_moe_enabled():
227
+ if not rocm_aiter_ops.is_fusion_moe_shared_experts_enabled():
228
+ assert self.num_fused_shared_experts == 0
229
+ return rocm_aiter_grouped_topk(
230
+ hidden_states,
231
+ gating_output,
232
+ self.topk,
233
+ self.renormalize,
234
+ self.num_expert_group,
235
+ self.topk_group,
236
+ self.scoring_func,
237
+ self.routed_scaling_factor,
238
+ e_score_correction_bias,
239
+ self.num_fused_shared_experts,
240
+ )
241
+ else:
242
+ return self.forward_native(
243
+ hidden_states, gating_output, e_score_correction_bias
244
+ )
245
+
246
+
247
+class GroupedTopKRouter(BaseRouter):
248
+ """Router using grouped top-k routing (e.g., DeepSeekV2/V3)."""
249
+
250
+ def __init__(
251
+ self,
252
+ top_k: int,
253
+ global_num_experts: int,
254
+ eplb_state: EplbLayerState,
255
+ num_expert_group: int,
256
+ topk_group: int,
257
+ renormalize: bool = True,
258
+ scoring_func: str = "softmax",
259
+ routed_scaling_factor: float = 1.0,
260
+ e_score_correction_bias: torch.Tensor | None = None,
261
+ num_fused_shared_experts: int = 0,
262
+ enable_eplb: bool = False,
263
+ indices_type_getter: Callable[[], torch.dtype | None] | None = None,
264
+ routing_method_type: RoutingMethodType | None = None,
265
+ ):
266
+ super().__init__(
267
+ top_k=top_k,
268
+ global_num_experts=global_num_experts,
269
+ eplb_state=eplb_state,
270
+ enable_eplb=enable_eplb,
271
+ indices_type_getter=indices_type_getter,
272
+ )
273
+ self.num_expert_group = num_expert_group
274
+ self.topk_group = topk_group
275
+ self.renormalize = renormalize
276
+ self.scoring_func = scoring_func
277
+ self.routed_scaling_factor = routed_scaling_factor
278
+ self.e_score_correction_bias = e_score_correction_bias
279
+ self.num_fused_shared_experts = num_fused_shared_experts
280
+
281
+ # Determine routing method type
282
+ if routing_method_type is not None:
283
+ self._routing_method_type = routing_method_type
284
+ elif scoring_func == "sigmoid":
285
+ self._routing_method_type = RoutingMethodType.DeepSeekV3
286
+ else:
287
+ self._routing_method_type = RoutingMethodType.TopK
288
+
289
+ @property
290
+ def routing_method_type(self) -> RoutingMethodType:
291
+ return self._routing_method_type
292
+
293
+ def _compute_routing(
294
+ self,
295
+ hidden_states: torch.Tensor,
296
+ router_logits: torch.Tensor,
297
+ indices_type: torch.dtype | None,
298
+ ) -> tuple[torch.Tensor, torch.Tensor]:
299
+ """Compute routing using grouped top-k."""
300
+
301
+ def valid_grouping() -> bool:
302
+ # Check if num_experts is greater than num_expert_group
303
+ # and is divisible by num_expert_group
304
+ num_experts = router_logits.shape[-1]
305
+ if num_experts <= self.num_expert_group:
306
+ return False
307
+ return num_experts % self.num_expert_group == 0
308
+
309
+ if not valid_grouping():
310
+ if self.e_score_correction_bias is not None:
311
+ topk_weights, topk_ids = fused_topk_bias(
312
+ hidden_states=hidden_states,
313
+ gating_output=router_logits,
314
+ e_score_correction_bias=self.e_score_correction_bias.data,
315
+ topk=self.top_k,
316
+ renormalize=self.renormalize,
317
+ )
318
+ if self.routed_scaling_factor != 1.0:
319
+ topk_weights *= self.routed_scaling_factor
320
+ else:
321
+ topk_weights, topk_ids, token_expert_indices = fused_topk(
322
+ hidden_states=hidden_states,
323
+ gating_output=router_logits,
324
+ topk=self.top_k,
325
+ renormalize=self.renormalize,
326
+ indices_type=indices_type,
327
+ )
328
+ return topk_weights, topk_ids
329
+
330
+ # Select grouped_topk implementation
331
+ if rocm_aiter_ops.is_fused_moe_enabled():
332
+ if not rocm_aiter_ops.is_fusion_moe_shared_experts_enabled():
333
+ assert self.num_fused_shared_experts == 0
334
+ grouped_topk_impl = partial(
335
+ rocm_aiter_grouped_topk,
336
+ num_fused_shared_experts=self.num_fused_shared_experts,
337
+ )
338
+ else:
339
+ grouped_topk_impl = grouped_topk
340
+
341
+ topk_weights, topk_ids = grouped_topk_impl(
342
+ hidden_states=hidden_states,
343
+ gating_output=router_logits,
344
+ topk=self.top_k,
345
+ renormalize=self.renormalize,
346
+ num_expert_group=self.num_expert_group,
347
+ topk_group=self.topk_group,
348
+ scoring_func=self.scoring_func,
349
+ routed_scaling_factor=self.routed_scaling_factor,
350
+ e_score_correction_bias=self.e_score_correction_bias,
351
+ )
352
+
353
+ return topk_weights, topk_ids
vllm/model_executor/layers/fused_moe/router/router_factory.py
new
+178
@@ -0,0 +1,178 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from collections.abc import Callable
4
+
5
+import torch
6
+
7
+import vllm.envs as envs
8
+from vllm.distributed.eplb.eplb_state import EplbLayerState
9
+from vllm.model_executor.layers.fused_moe.config import RoutingMethodType
10
+from vllm.model_executor.layers.fused_moe.router.base_router import BaseRouter
11
+from vllm.model_executor.layers.fused_moe.router.custom_routing_router import (
12
+ CustomRoutingRouter,
13
+)
14
+from vllm.model_executor.layers.fused_moe.router.fused_moe_router import (
15
+ FusedMoERouter,
16
+)
17
+from vllm.model_executor.layers.fused_moe.router.fused_topk_bias_router import (
18
+ FusedTopKBiasRouter,
19
+)
20
+from vllm.model_executor.layers.fused_moe.router.fused_topk_router import (
21
+ FusedTopKRouter,
22
+)
23
+from vllm.model_executor.layers.fused_moe.router.grouped_topk_router import (
24
+ GroupedTopKRouter,
25
+)
26
+from vllm.model_executor.layers.fused_moe.router.routing_simulator_router import (
27
+ RoutingSimulatorRouter,
28
+)
29
+
30
+EMPTY_EPLB_STATE: EplbLayerState = EplbLayerState()
31
+
32
+
33
+def create_fused_moe_router(
34
+ # common parameters
35
+ top_k: int,
36
+ global_num_experts: int,
37
+ renormalize: bool = True,
38
+ indices_type_getter: Callable[[], torch.dtype | None] | None = None,
39
+ routing_method_type: RoutingMethodType | None = None,
40
+ # grouped topk parameters
41
+ use_grouped_topk: bool = False,
42
+ num_expert_group: int | None = None,
43
+ topk_group: int | None = None,
44
+ scoring_func: str = "softmax",
45
+ num_fused_shared_experts: int = 0,
46
+ # grouped topk + fused topk bias parameters
47
+ routed_scaling_factor: float = 1.0,
48
+ e_score_correction_bias: torch.Tensor | None = None,
49
+ # custom routing paramaters
50
+ custom_routing_function: Callable | None = None,
51
+ # eplb parameters
52
+ enable_eplb: bool = False,
53
+ eplb_state: EplbLayerState = EMPTY_EPLB_STATE,
54
+ capture: Callable[[torch.tensor], None] | None = None,
55
+) -> FusedMoERouter:
56
+ """
57
+ Factory function to create the appropriate FusedMoERouter subclass based on
58
+ the provided parameters.
59
+
60
+ The selection logic follows this priority order:
61
+ 1. RoutingSimulatorRouter - if VLLM_MOE_ROUTING_SIMULATION_STRATEGY env var is set
62
+ 2. GroupedTopKRouter - if use_grouped_topk is True
63
+ 3. CustomRoutingRouter - if custom_routing_function is not None
64
+ 4. FusedTopKBiasRouter - if e_score_correction_bias is not None
65
+ 5. FusedTopKRouter - default fallback
66
+
67
+ Common arguments:
68
+ top_k: Number of experts to select per token
69
+ global_num_experts: Total number of experts in the model
70
+ renormalize: Whether to renormalize the routing weights
71
+ indices_type_getter: Function to get the desired indices dtype
72
+ routing_method_type: Optional explicit routing method type
73
+
74
+ Grouped topk arguments:
75
+ use_grouped_topk: Whether to use grouped top-k routing
76
+ num_expert_group: Number of expert groups (for grouped routing)
77
+ topk_group: Top-k within each group (for grouped routing)
78
+ scoring_func: Scoring function to use ("softmax" or "sigmoid")
79
+ num_fused_shared_experts: Number of fused shared experts (for ROCm AITER)
80
+
81
+ Grouped topk and fused topk bias arguments:
82
+ routed_scaling_factor: Scaling factor for routed weights
83
+ e_score_correction_bias: Optional bias correction for expert scores
84
+
85
+ Custom routing arguments:
86
+ custom_routing_function: Optional custom routing function
87
+
88
+ EPLB arguments:
89
+ enable_eplb: Whether EPLB is enabled
90
+ eplb_state: EPLB (Expert Parallelism Load Balancing) state
91
+
92
+ Returns:
93
+ An instance of the appropriate FusedMoERouter subclass
94
+ """
95
+ router: BaseRouter
96
+
97
+ routing_strategy = envs.VLLM_MOE_ROUTING_SIMULATION_STRATEGY
98
+ if routing_strategy != "":
99
+ router = RoutingSimulatorRouter(
100
+ top_k=top_k,
101
+ global_num_experts=global_num_experts,
102
+ eplb_state=eplb_state,
103
+ enable_eplb=enable_eplb,
104
+ indices_type_getter=indices_type_getter,
105
+ )
106
+ # TODO(bnell): this is temporary until select_experts is
107
+ # separated from apply.
108
+ router.capture = capture
109
+ return router
110
+
111
+ if use_grouped_topk:
112
+ assert custom_routing_function is None
113
+ if num_expert_group is None or topk_group is None:
114
+ raise ValueError(
115
+ "num_expert_group and topk_group must be provided when "
116
+ "use_grouped_topk is True"
117
+ )
118
+ router = GroupedTopKRouter(
119
+ top_k=top_k,
120
+ global_num_experts=global_num_experts,
121
+ eplb_state=eplb_state,
122
+ num_expert_group=num_expert_group,
123
+ topk_group=topk_group,
124
+ renormalize=renormalize,
125
+ scoring_func=scoring_func,
126
+ routed_scaling_factor=routed_scaling_factor,
127
+ e_score_correction_bias=e_score_correction_bias,
128
+ num_fused_shared_experts=num_fused_shared_experts,
129
+ enable_eplb=enable_eplb,
130
+ indices_type_getter=indices_type_getter,
131
+ routing_method_type=routing_method_type,
132
+ )
133
+ router.capture = capture
134
+ return router
135
+
136
+ if custom_routing_function is not None:
137
+ router = CustomRoutingRouter(
138
+ top_k=top_k,
139
+ global_num_experts=global_num_experts,
140
+ eplb_state=eplb_state,
141
+ custom_routing_function=custom_routing_function,
142
+ renormalize=renormalize,
143
+ enable_eplb=enable_eplb,
144
+ indices_type_getter=indices_type_getter,
145
+ )
146
+ router.capture = capture
147
+ return router
148
+
149
+ if scoring_func != "softmax":
150
+ raise ValueError(
151
+ "Only softmax scoring function is supported for non-grouped topk."
152
+ )
153
+
154
+ if e_score_correction_bias is not None:
155
+ router = FusedTopKBiasRouter(
156
+ top_k=top_k,
157
+ global_num_experts=global_num_experts,
158
+ eplb_state=eplb_state,
159
+ e_score_correction_bias=e_score_correction_bias,
160
+ renormalize=renormalize,
161
+ routed_scaling_factor=routed_scaling_factor,
162
+ enable_eplb=enable_eplb,
163
+ indices_type_getter=indices_type_getter,
164
+ )
165
+ router.capture = capture
166
+ return router
167
+
168
+ router = FusedTopKRouter(
169
+ top_k=top_k,
170
+ global_num_experts=global_num_experts,
171
+ eplb_state=eplb_state,
172
+ renormalize=renormalize,
173
+ scoring_func=scoring_func,
174
+ enable_eplb=enable_eplb,
175
+ indices_type_getter=indices_type_getter,
176
+ )
177
+ router.capture = capture
178
+ return router
vllm/model_executor/layers/fused_moe/router/routing_simulator_router.py
renamed
+46
-9
@@ -1,20 +1,16 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-"""
4
-Token-to-Expert Routing Simulator
5
-
6
-This module provides a framework for simulating and testing different
7
-token-to-expert routing strategies for Mixture of Experts (MoE) models.
8
-It supports routing logic customization and includes example implementations
9
-like uniform random routing.
10
-"""
11
-
3
from abc import ABC, abstractmethod
4
+from collections.abc import Callable
5
from typing import Any
6
7
import torch
8
9
+import vllm.envs as envs
10
+from vllm.distributed.eplb.eplb_state import EplbLayerState
11
from vllm.logger import init_logger
12
+from vllm.model_executor.layers.fused_moe.config import RoutingMethodType
13
+from vllm.model_executor.layers.fused_moe.router.base_router import BaseRouter
14
15
logger = init_logger(__name__)
16
@@ -308,3 +304,44 @@ class RoutingSimulator:
304
top_k=top_k,
305
indices_type=indices_type,
306
)
307
+
308
+
309
+class RoutingSimulatorRouter(BaseRouter):
310
+ """Router that uses routing simulation strategies for testing/debugging."""
311
+
312
+ def __init__(
313
+ self,
314
+ top_k: int,
315
+ global_num_experts: int,
316
+ eplb_state: EplbLayerState,
317
+ enable_eplb: bool = False,
318
+ indices_type_getter: Callable[[], torch.dtype | None] | None = None,
319
+ ):
320
+ super().__init__(
321
+ top_k=top_k,
322
+ global_num_experts=global_num_experts,
323
+ eplb_state=eplb_state,
324
+ enable_eplb=enable_eplb,
325
+ indices_type_getter=indices_type_getter,
326
+ )
327
+
328
+ @property
329
+ def routing_method_type(self) -> RoutingMethodType:
330
+ return RoutingMethodType.Simulated
331
+
332
+ def _compute_routing(
333
+ self,
334
+ hidden_states: torch.Tensor,
335
+ router_logits: torch.Tensor,
336
+ indices_type: torch.dtype | None,
337
+ ) -> tuple[torch.Tensor, torch.Tensor]:
338
+ """Use routing simulator to compute routing."""
339
+ routing_strategy = envs.VLLM_MOE_ROUTING_SIMULATION_STRATEGY
340
+ topk_weights, topk_ids = RoutingSimulator.simulate_routing(
341
+ hidden_states=hidden_states,
342
+ router_logits=router_logits,
343
+ strategy_name=routing_strategy,
344
+ top_k=self.top_k,
345
+ indices_type=indices_type,
346
+ )
347
+ return topk_weights, topk_ids
vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py
+9
-7
@@ -20,7 +20,6 @@ from vllm.model_executor.layers.fused_moe.config import (
20
from vllm.model_executor.layers.fused_moe.fused_moe_method_base import (
21
FusedMoEMethodBase,
22
)
23
-from vllm.model_executor.layers.fused_moe.fused_moe_router import FusedMoERouter
23
from vllm.model_executor.layers.fused_moe.modular_kernel import (
24
FusedMoEActivationFormat,
25
FusedMoEPermuteExpertsUnpermute,
@@ -32,6 +31,9 @@ from vllm.model_executor.layers.fused_moe.oracle.unquantized import (
31
make_unquantized_moe_kernel,
32
select_unquantized_moe_backend,
33
)
34
+from vllm.model_executor.layers.fused_moe.router.fused_moe_router import (
35
+ FusedMoERouter,
36
+)
37
from vllm.model_executor.utils import replace_parameter, set_weight_attrs
38
from vllm.platforms import current_platform
39
from vllm.platforms.interface import CpuArchEnum
@@ -312,9 +314,9 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
314
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
315
if (
316
layer.enable_eplb is not False
315
- or layer.expert_load_view is not None
316
- or layer.logical_to_physical_map is not None
317
- or layer.logical_replica_count is not None
317
+ or layer.eplb_state.expert_load_view is not None
318
+ or layer.eplb_state.logical_to_physical_map is not None
319
+ or layer.eplb_state.logical_replica_count is not None
320
):
321
raise NotImplementedError("Expert load balancing is not supported for CPU.")
322
@@ -346,9 +348,9 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
348
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
349
if (
350
layer.enable_eplb is not False
349
- or layer.expert_load_view is not None
350
- or layer.logical_to_physical_map is not None
351
- or layer.logical_replica_count is not None
351
+ or layer.eplb_state.expert_load_view is not None
352
+ or layer.eplb_state.logical_to_physical_map is not None
353
+ or layer.eplb_state.logical_replica_count is not None
354
):
355
raise NotImplementedError("Expert load balancing is not supported for XPU.")
356
return layer.ipex_fusion(
vllm/model_executor/layers/quantization/awq_marlin.py
+1
-1
@@ -10,12 +10,12 @@ from torch.nn import Parameter
10
import vllm.model_executor.layers.fused_moe # noqa
11
from vllm import _custom_ops as ops
12
from vllm.logger import init_logger
13
+from vllm.model_executor.layers.fused_moe import FusedMoERouter
14
from vllm.model_executor.layers.fused_moe.config import (
15
FusedMoEConfig,
16
FusedMoEQuantConfig,
17
)
18
from vllm.model_executor.layers.fused_moe.fused_marlin_moe import fused_marlin_moe
18
-from vllm.model_executor.layers.fused_moe.fused_moe_router import FusedMoERouter
19
from vllm.model_executor.layers.fused_moe.layer import (
20
FusedMoE,
21
FusedMoEMethodBase,
vllm/model_executor/layers/quantization/bitsandbytes.py
+1
-1
@@ -6,11 +6,11 @@ from typing import Any, Union
6
import torch
7
from packaging import version
8
9
+from vllm.model_executor.layers.fused_moe import FusedMoERouter
10
from vllm.model_executor.layers.fused_moe.config import (
11
FusedMoEConfig,
12
FusedMoEQuantConfig,
13
)
13
-from vllm.model_executor.layers.fused_moe.fused_moe_router import FusedMoERouter
14
from vllm.model_executor.layers.fused_moe.layer import (
15
FusedMoE,
16
FusedMoEMethodBase,
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py
+1
-1
@@ -22,6 +22,7 @@ from vllm.model_executor.layers.fused_moe import (
22
FusedMoEConfig,
23
FusedMoEMethodBase,
24
FusedMoEPermuteExpertsUnpermute,
25
+ FusedMoERouter,
26
FusedMoeWeightScaleSupported,
27
UnquantizedFusedMoEMethod,
28
)
@@ -40,7 +41,6 @@ from vllm.model_executor.layers.fused_moe.fused_marlin_moe import (
41
MarlinExperts,
42
fused_marlin_moe,
43
)
43
-from vllm.model_executor.layers.fused_moe.fused_moe_router import FusedMoERouter
44
from vllm.model_executor.layers.fused_moe.oracle.fp8 import (
45
Fp8MoeBackend,
46
convert_to_fp8_moe_kernel_format,
vllm/model_executor/layers/quantization/experts_int8.py
+1
-1
@@ -10,12 +10,12 @@ from vllm.model_executor.layers.fused_moe import (
10
FusedMoE,
11
FusedMoEConfig,
12
FusedMoEMethodBase,
13
+ FusedMoERouter,
14
)
15
from vllm.model_executor.layers.fused_moe.config import (
16
FusedMoEQuantConfig,
17
int8_w8a16_moe_quant_config,
18
)
18
-from vllm.model_executor.layers.fused_moe.fused_moe_router import FusedMoERouter
19
from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod
20
from vllm.model_executor.layers.quantization import QuantizationMethods
21
from vllm.model_executor.layers.quantization.base_config import (
vllm/model_executor/layers/quantization/fp8.py
+1
-1
@@ -23,13 +23,13 @@ from vllm.model_executor.layers.fused_moe import (
23
FusedMoEMethodBase,
24
FusedMoEPermuteExpertsUnpermute,
25
FusedMoEPrepareAndFinalize,
26
+ FusedMoERouter,
27
FusedMoeWeightScaleSupported,
28
)
29
from vllm.model_executor.layers.fused_moe.config import (
30
FusedMoEQuantConfig,
31
RoutingMethodType,
32
)
32
-from vllm.model_executor.layers.fused_moe.fused_moe_router import FusedMoERouter
33
from vllm.model_executor.layers.fused_moe.layer import UnquantizedFusedMoEMethod
34
from vllm.model_executor.layers.fused_moe.oracle.fp8 import (
35
Fp8MoeBackend,
vllm/model_executor/layers/quantization/gguf.py
+1
-1
@@ -12,11 +12,11 @@ from torch.nn.parameter import Parameter, UninitializedParameter
12
13
from vllm import _custom_ops as ops
14
from vllm.logger import init_logger
15
+from vllm.model_executor.layers.fused_moe import FusedMoERouter
16
from vllm.model_executor.layers.fused_moe.config import (
17
FusedMoEConfig,
18
FusedMoEQuantConfig,
19
)
19
-from vllm.model_executor.layers.fused_moe.fused_moe_router import FusedMoERouter
20
from vllm.model_executor.layers.fused_moe.layer import (
21
FusedMoE,
22
FusedMoEMethodBase,
vllm/model_executor/layers/quantization/gptq_marlin.py
+1
-1
@@ -10,12 +10,12 @@ from safetensors.torch import _TYPES as _SAFETENSORS_TO_TORCH_DTYPE
10
import vllm.model_executor.layers.fused_moe # noqa
11
from vllm import _custom_ops as ops
12
from vllm.logger import init_logger
13
+from vllm.model_executor.layers.fused_moe import FusedMoERouter
14
from vllm.model_executor.layers.fused_moe.config import (
15
FusedMoEConfig,
16
FusedMoEQuantConfig,
17
)
18
from vllm.model_executor.layers.fused_moe.fused_marlin_moe import fused_marlin_moe
18
-from vllm.model_executor.layers.fused_moe.fused_moe_router import FusedMoERouter
19
from vllm.model_executor.layers.fused_moe.layer import (
20
FusedMoE,
21
FusedMoEMethodBase,
vllm/model_executor/layers/quantization/ipex_quant.py
+2
-2
@@ -8,10 +8,10 @@ from packaging import version
8
from torch.nn import Module
9
10
from vllm._ipex_ops import ipex_ops as ops
11
-from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
12
-from vllm.model_executor.layers.fused_moe.fused_moe_router import (
11
+from vllm.model_executor.layers.fused_moe import (
12
FusedMoERouter,
13
)
14
+from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
15
from vllm.model_executor.layers.linear import (
16
LinearBase,
17
LinearMethodBase,
vllm/model_executor/layers/quantization/modelopt.py
+1
-1
@@ -13,11 +13,11 @@ import vllm.model_executor.layers.fused_moe.modular_kernel as mk
13
from vllm._custom_ops import cutlass_scaled_fp4_mm, scaled_fp4_quant
14
from vllm.attention.layer import Attention
15
from vllm.logger import init_logger
16
+from vllm.model_executor.layers.fused_moe import FusedMoERouter
17
from vllm.model_executor.layers.fused_moe.config import (
18
FusedMoEConfig,
19
FusedMoEQuantConfig,
20
)
20
-from vllm.model_executor.layers.fused_moe.fused_moe_router import FusedMoERouter
21
from vllm.model_executor.layers.fused_moe.layer import (
22
FusedMoE,
23
FusedMoEMethodBase,
vllm/model_executor/layers/quantization/moe_wna16.py
+1
-1
@@ -6,12 +6,12 @@ from typing import Any, Optional
6
import torch
7
8
from vllm.distributed import get_tensor_model_parallel_rank, get_tp_group
9
+from vllm.model_executor.layers.fused_moe import FusedMoERouter
10
from vllm.model_executor.layers.fused_moe.config import (
11
FusedMoEQuantConfig,
12
int4_w4a16_moe_quant_config,
13
int8_w8a16_moe_quant_config,
14
)
14
-from vllm.model_executor.layers.fused_moe.fused_moe_router import FusedMoERouter
15
from vllm.model_executor.layers.fused_moe.layer import (
16
FusedMoE,
17
FusedMoEConfig,
vllm/model_executor/layers/quantization/mxfp4.py
+4
-4
@@ -14,6 +14,7 @@ from vllm.model_executor.layers.fused_moe import (
14
FusedMoE,
15
FusedMoEConfig,
16
FusedMoEMethodBase,
17
+ FusedMoERouter,
18
)
19
from vllm.model_executor.layers.fused_moe import modular_kernel as mk
20
from vllm.model_executor.layers.fused_moe.config import (
@@ -27,7 +28,6 @@ from vllm.model_executor.layers.fused_moe.fused_marlin_moe import (
28
MarlinExperts,
29
fused_marlin_moe,
30
)
30
-from vllm.model_executor.layers.fused_moe.fused_moe_router import FusedMoERouter
31
from vllm.model_executor.layers.fused_moe.gpt_oss_triton_kernels_moe import (
32
OAITritonExperts,
33
UnfusedOAITritonExperts,
@@ -936,9 +936,9 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
936
layer.apply_router_weight_on_input,
937
layer.scoring_func,
938
layer.activation,
939
- layer.expert_load_view,
940
- layer.logical_to_physical_map,
941
- layer.logical_replica_count,
939
+ layer.eplb_state.expert_load_view,
940
+ layer.eplb_state.logical_to_physical_map,
941
+ layer.eplb_state.logical_replica_count,
942
), "MXFP4 are not supported with this configuration."
943
944
if (
vllm/model_executor/layers/quantization/quark/quark_moe.py
+1
-1
@@ -548,7 +548,7 @@ class QuarkW4A8Fp8MoEMethod(QuarkMoEMethod):
548
x: torch.Tensor,
549
router_logits: torch.Tensor,
550
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
551
- topk_weights, topk_ids = layer.select_experts(
551
+ topk_weights, topk_ids = router.select_experts(
552
hidden_states=x,
553
router_logits=router_logits,
554
)
vllm/model_executor/layers/quantization/rtn.py
+1
-1
@@ -10,12 +10,12 @@ import torch
10
from torch.nn.parameter import Parameter
11
12
from vllm.logger import init_logger
13
+from vllm.model_executor.layers.fused_moe import FusedMoERouter
14
from vllm.model_executor.layers.fused_moe.config import (
15
FusedMoEConfig,
16
FusedMoEQuantConfig,
17
)
18
from vllm.model_executor.layers.fused_moe.fused_marlin_moe import fused_marlin_moe
18
-from vllm.model_executor.layers.fused_moe.fused_moe_router import FusedMoERouter
19
from vllm.model_executor.layers.fused_moe.layer import (
20
FusedMoE,
21
FusedMoEMethodBase,
vllm/model_executor/models/ernie45_moe.py
+1
@@ -201,6 +201,7 @@ class Ernie4_5_MoeMoE(nn.Module):
201
e_score_correction_bias=self.gate.e_score_correction_bias,
202
enable_eplb=self.enable_eplb,
203
num_redundant_experts=self.n_redundant_experts,
204
+ router_logits_dtype=torch.float32,
205
)
206
207
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
vllm/model_executor/models/ernie45_vl_moe.py
+2
@@ -269,6 +269,7 @@ class Ernie4_5_VLMoeMoE(nn.Module):
269
quant_config=quant_config,
270
e_score_correction_bias=self.e_score_correction_bias[0],
271
prefix=f"{prefix}.text_experts",
272
+ router_logits_dtype=torch.float32,
273
)
274
else:
275
self.text_experts = Ernie4_5_VLMoeMLP(
@@ -306,6 +307,7 @@ class Ernie4_5_VLMoeMoE(nn.Module):
307
quant_config=quant_config,
308
e_score_correction_bias=self.e_score_correction_bias[1],
309
prefix=f"{prefix}.vision_experts",
310
+ router_logits_dtype=torch.float32,
311
)
312
else:
313
self.vision_experts = Ernie4_5_VLMoeMLP(