perf: dispatch non-grouped bias-less topk routing methods to fused path (#49618)
Signed-off-by: jdebache <jdebache@nvidia.com>
Julien Debache committed
Jul 28, 2026 at 23:57 UTC
bb3b61f2fd2333ab165ebaba13f133db4210b9f2
2 files changed
+159
-23
tests/kernels/moe/test_routing.py
+116
@@ -8,9 +8,19 @@ import torch
8
9
from vllm._aiter_ops import rocm_aiter_ops
10
from vllm.distributed.eplb.eplb_state import EplbLayerState
11
+from vllm.model_executor.layers.fused_moe.config import RoutingMethodType
12
from vllm.model_executor.layers.fused_moe.router.base_router import (
13
eplb_map_to_physical_and_record,
14
)
15
+from vllm.model_executor.layers.fused_moe.router.fused_topk_bias_router import (
16
+ FusedTopKBiasRouter,
17
+)
18
+from vllm.model_executor.layers.fused_moe.router.fused_topk_router import (
19
+ FusedTopKRouter,
20
+)
21
+from vllm.model_executor.layers.fused_moe.router.grouped_topk_router import (
22
+ GroupedTopKRouter,
23
+)
24
from vllm.model_executor.layers.fused_moe.router.router_factory import (
25
create_fused_moe_router,
26
)
@@ -36,6 +46,112 @@ TOP_KS = [2, 4, 6]
46
NUM_EXPERTS = [8, 16, 64]
47
48
49
+def test_degenerate_grouped_config_uses_standard_topk() -> None:
50
+ router = create_fused_moe_router(
51
+ top_k=4,
52
+ global_num_experts=128,
53
+ use_grouped_topk=True,
54
+ num_expert_group=1,
55
+ topk_group=1,
56
+ scoring_func="softmax",
57
+ renormalize=True,
58
+ )
59
+
60
+ assert isinstance(router, FusedTopKRouter)
61
+ hidden_states, router_logits = make_test_data(32, 256, 128)
62
+
63
+ topk_weights, topk_ids = router.select_experts(hidden_states, router_logits)
64
+ baseline_weights, baseline_ids = baseline_fused_topk(
65
+ router_logits,
66
+ top_k=4,
67
+ renormalize=True,
68
+ )
69
+
70
+ assert_routing_results_close(
71
+ topk_weights,
72
+ topk_ids,
73
+ baseline_weights,
74
+ baseline_ids,
75
+ )
76
+
77
+
78
+def test_multiple_expert_groups_use_grouped_topk() -> None:
79
+ router = create_fused_moe_router(
80
+ top_k=4,
81
+ global_num_experts=128,
82
+ use_grouped_topk=True,
83
+ num_expert_group=8,
84
+ topk_group=4,
85
+ scoring_func="softmax",
86
+ renormalize=True,
87
+ )
88
+
89
+ assert isinstance(router, GroupedTopKRouter)
90
+
91
+
92
+def test_degenerate_grouped_config_with_bias_uses_topk_bias() -> None:
93
+ router = create_fused_moe_router(
94
+ top_k=4,
95
+ global_num_experts=128,
96
+ use_grouped_topk=True,
97
+ num_expert_group=1,
98
+ topk_group=1,
99
+ scoring_func="softmax",
100
+ renormalize=True,
101
+ e_score_correction_bias=torch.empty(128),
102
+ )
103
+
104
+ assert isinstance(router, FusedTopKBiasRouter)
105
+
106
+
107
+def test_degenerate_grouped_config_with_bias_keeps_routed_scale() -> None:
108
+ router = create_fused_moe_router(
109
+ top_k=4,
110
+ global_num_experts=128,
111
+ use_grouped_topk=True,
112
+ num_expert_group=1,
113
+ topk_group=1,
114
+ scoring_func="softmax",
115
+ renormalize=True,
116
+ routed_scaling_factor=1.1,
117
+ e_score_correction_bias=torch.empty(128),
118
+ )
119
+
120
+ assert isinstance(router, FusedTopKBiasRouter)
121
+ assert router.routed_scaling_factor == 1.1
122
+
123
+
124
+def test_degenerate_deepseek_v3_routing_stays_grouped() -> None:
125
+ router = create_fused_moe_router(
126
+ top_k=4,
127
+ global_num_experts=128,
128
+ use_grouped_topk=True,
129
+ num_expert_group=1,
130
+ topk_group=1,
131
+ scoring_func="sigmoid",
132
+ renormalize=True,
133
+ e_score_correction_bias=torch.empty(128),
134
+ )
135
+
136
+ assert isinstance(router, GroupedTopKRouter)
137
+ assert router.routing_method_type == RoutingMethodType.DeepSeekV3
138
+
139
+
140
+def test_single_expert_group_with_non_unit_scale_uses_grouped_topk() -> None:
141
+ router = create_fused_moe_router(
142
+ top_k=4,
143
+ global_num_experts=128,
144
+ use_grouped_topk=True,
145
+ num_expert_group=1,
146
+ topk_group=1,
147
+ scoring_func="softmax",
148
+ renormalize=True,
149
+ routed_scaling_factor=1.1,
150
+ )
151
+
152
+ assert isinstance(router, GroupedTopKRouter)
153
+
154
+
155
def setup_eplb_state(
156
enable_eplb: bool, global_num_experts: int
157
) -> EplbLayerState | None:
vllm/model_executor/layers/fused_moe/router/router_factory.py
+43
-23
@@ -9,6 +9,7 @@ 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 (
11
RoutingMethodType,
12
+ get_routing_method_type,
13
)
14
from vllm.model_executor.layers.fused_moe.router.aiter_shared_routed_fused_moe_router import ( # noqa: E501
15
AiterSharedRoutedFusedMoERouter,
@@ -67,7 +68,8 @@ def create_fused_moe_router(
68
The selection logic follows this priority order:
69
1. RoutingSimulatorRouter - if VLLM_MOE_ROUTING_SIMULATION_STRATEGY env var is set
70
2. ZeroExpertRouter - if zero_expert_type is not None
70
- 3. GroupedTopKRouter - if use_grouped_topk is True
71
+ 3. GroupedTopKRouter - if use_grouped_topk is True and the grouping is not
72
+ degenerate (at most one group, with topk_group <= 1)
73
4. CustomRoutingRouter - if custom_routing_function is not None
74
5. FusedTopKBiasRouter - if e_score_correction_bias is not None
75
6. AiterSharedRoutedFusedMoERouter - if num_fused_shared_experts > 0
@@ -143,30 +145,48 @@ def create_fused_moe_router(
145
"num_expert_group and topk_group must be provided when "
146
"use_grouped_topk is True"
147
)
146
- grouped_topk_router = GroupedTopKRouter(
147
- top_k=top_k,
148
- global_num_experts=global_num_experts,
149
- eplb_state=eplb_state,
150
- num_expert_group=num_expert_group,
151
- topk_group=topk_group,
152
- renormalize=renormalize,
153
- scoring_func=scoring_func,
154
- routed_scaling_factor=routed_scaling_factor,
155
- e_score_correction_bias=e_score_correction_bias,
156
- num_fused_shared_experts=num_fused_shared_experts,
148
+
149
+ # For topk_group <= 1, grouped implementation is pure overhead.
150
+ degenerate_grouping = num_expert_group <= 1 and topk_group <= 1
151
+ # FusedTopKRouter cannot apply routed_scaling_factor, FusedTopKBiasRouter can.
152
+ scaling_handled_downstream = (
153
+ routed_scaling_factor == 1.0 or e_score_correction_bias is not None
154
)
158
- if (
159
- grouped_topk_router.routing_method_type != RoutingMethodType.Unspecified
160
- or num_expert_group > 1
161
- or topk_group > 1
162
- ):
163
- return grouped_topk_router
155
165
- # If routing_method for GroupedTopKRouter is Unspecified and there is only
166
- # one group, fallback to standard top-k routing
167
- use_grouped_topk = False
168
- num_expert_group = None
169
- topk_group = None
156
+ # Degenerating must not change the advertised routing method, which drives
157
+ # kernel selection. num_expert_group only affects it for biased routing.
158
+ def advertised_routing_method(groups: int | None) -> RoutingMethodType:
159
+ return get_routing_method_type(
160
+ scoring_func=scoring_func,
161
+ top_k=top_k,
162
+ renormalize=renormalize,
163
+ num_expert_group=groups,
164
+ has_e_score_bias=e_score_correction_bias is not None,
165
+ routed_scaling_factor=routed_scaling_factor,
166
+ )
167
+
168
+ routing_method_preserved = advertised_routing_method(
169
+ num_expert_group
170
+ ) == advertised_routing_method(None)
171
+
172
+ if not (
173
+ degenerate_grouping
174
+ and scaling_handled_downstream
175
+ and routing_method_preserved
176
+ ):
177
+ return GroupedTopKRouter(
178
+ top_k=top_k,
179
+ global_num_experts=global_num_experts,
180
+ eplb_state=eplb_state,
181
+ num_expert_group=num_expert_group,
182
+ topk_group=topk_group,
183
+ renormalize=renormalize,
184
+ scoring_func=scoring_func,
185
+ routed_scaling_factor=routed_scaling_factor,
186
+ e_score_correction_bias=e_score_correction_bias,
187
+ num_fused_shared_experts=num_fused_shared_experts,
188
+ )
189
+ # Otherwise fall through to the non-grouped chain below.
190
191
if custom_routing_function is not None:
192
return CustomRoutingRouter(