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 5, 04:45 UTC
[MoE Refactor] Refactor ZeroExpertFusedMoE into new framework (#35549)
Signed-off-by: Bill Nell <bnell@redhat.com>
bnellnm committed
Apr 14, 2026 at 16:11 UTC
19ec9a0a6276d5af449b3011dac2944e8af37ee0
8 files changed
+461
-204
tests/kernels/moe/test_zero_expert_moe.py
new
+282
@@ -0,0 +1,282 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Tests for FusedMoE with zero experts.
4
+
5
+Verifies that:
6
+- The ZeroExpertRouter is properly created and used as the layer router.
7
+- A forward pass through FusedMoE with zero experts produces correct output.
8
+- The output decomposes correctly into real expert + zero expert contributions.
9
+
10
+Note: tests generated with Claude.
11
+"""
12
+
13
+import pytest
14
+import torch
15
+
16
+from vllm.config import VllmConfig, set_current_vllm_config
17
+from vllm.forward_context import get_forward_context, set_forward_context
18
+from vllm.model_executor.layers.fused_moe.layer import FusedMoE
19
+from vllm.model_executor.layers.fused_moe.router.zero_expert_router import (
20
+ ZeroExpertRouter,
21
+)
22
+from vllm.v1.worker.workspace import init_workspace_manager
23
+
24
+
25
+@pytest.fixture
26
+def zero_expert_moe(dist_init, default_vllm_config):
27
+ """Create a FusedMoE layer with zero experts."""
28
+ num_experts = 4
29
+ top_k = 2
30
+ # hidden_size must be >= 256 for the zero expert identity kernel to
31
+ # produce output (its BLOCK_SIZE=256 causes grid=0 when hidden_dim<256).
32
+ hidden_size = 256
33
+ intermediate_size = 512
34
+ zero_expert_num = 1
35
+
36
+ e_score_correction_bias = torch.zeros(
37
+ num_experts + zero_expert_num,
38
+ dtype=torch.float32,
39
+ device="cuda",
40
+ )
41
+
42
+ vllm_config = VllmConfig()
43
+ vllm_config.compilation_config.static_forward_context = dict()
44
+
45
+ with set_current_vllm_config(vllm_config), set_forward_context(None, vllm_config):
46
+ init_workspace_manager(torch.accelerator.current_device_index())
47
+
48
+ layer = FusedMoE(
49
+ zero_expert_type="identity",
50
+ e_score_correction_bias=e_score_correction_bias,
51
+ num_experts=num_experts,
52
+ top_k=top_k,
53
+ hidden_size=hidden_size,
54
+ intermediate_size=intermediate_size,
55
+ params_dtype=torch.bfloat16,
56
+ prefix="test_zero_expert_moe",
57
+ renormalize=False,
58
+ routed_scaling_factor=1.0,
59
+ scoring_func="softmax",
60
+ ).cuda()
61
+
62
+ layer.quant_method.process_weights_after_loading(layer)
63
+
64
+ yield layer, vllm_config
65
+
66
+
67
+@pytest.mark.parametrize("num_tokens", [1, 32])
68
+def test_zero_expert_moe_router_is_zero_expert_router(zero_expert_moe, num_tokens):
69
+ """Verify that FusedMoE with zero_expert_type creates a ZeroExpertRouter."""
70
+ layer, _ = zero_expert_moe
71
+ assert isinstance(layer.router, ZeroExpertRouter), (
72
+ f"Expected ZeroExpertRouter but got {type(layer.router).__name__}."
73
+ )
74
+
75
+
76
+@pytest.mark.parametrize("num_tokens", [1, 32])
77
+def test_zero_expert_moe_no_custom_routing_fn(zero_expert_moe, num_tokens):
78
+ """Verify that custom_routing_function is not set (routing is handled
79
+ by ZeroExpertRouter, not a memoizing closure)."""
80
+ layer, _ = zero_expert_moe
81
+ assert layer.custom_routing_function is None
82
+
83
+
84
+@pytest.mark.parametrize("num_tokens", [1, 32])
85
+def test_zero_expert_moe_forward(zero_expert_moe, num_tokens):
86
+ """Run a forward pass through FusedMoE with zero experts and verify output shape."""
87
+ layer, vllm_config = zero_expert_moe
88
+
89
+ hidden_size = layer.hidden_size
90
+ num_experts = 4
91
+ zero_expert_num = 1
92
+ total_experts = num_experts + zero_expert_num
93
+
94
+ hidden_states = torch.randn(
95
+ num_tokens, hidden_size, dtype=torch.bfloat16, device="cuda"
96
+ )
97
+ router_logits = torch.randn(
98
+ num_tokens, total_experts, dtype=torch.float32, device="cuda"
99
+ )
100
+
101
+ # Initialize weights to small random values to avoid NaN from
102
+ # uninitialized memory.
103
+ with torch.no_grad():
104
+ for param in layer.parameters():
105
+ if param.dtype.is_floating_point:
106
+ param.normal_(0, 0.01)
107
+
108
+ with set_current_vllm_config(vllm_config), set_forward_context(None, vllm_config):
109
+ get_forward_context().all_moe_layers = None
110
+ output = layer.forward(hidden_states, router_logits)
111
+
112
+ assert output.shape == hidden_states.shape, (
113
+ f"Expected output shape {hidden_states.shape}, got {output.shape}"
114
+ )
115
+ assert output.dtype == hidden_states.dtype
116
+ assert not torch.isnan(output).any(), "Output contains NaN values"
117
+
118
+
119
+@pytest.mark.parametrize("num_tokens", [1, 32])
120
+def test_zero_expert_moe_output_decomposition(zero_expert_moe, num_tokens):
121
+ """Validate that the FusedMoE output equals a plain FusedMoE
122
+ output (real experts only) plus the zero expert contribution.
123
+
124
+ The key invariant is:
125
+ zero_layer.forward(h, r_full) == plain_layer.forward(h, r_real)
126
+ + zero_expert_output
127
+
128
+ We create a plain FusedMoE layer with the same weights and real-expert-only
129
+ router logits, compute the zero expert output via the ZeroExpertRouter, and
130
+ verify the sum matches the FusedMoE output.
131
+ """
132
+ layer, vllm_config = zero_expert_moe
133
+ num_experts = 4
134
+ zero_expert_num = 1
135
+ total_experts = num_experts + zero_expert_num
136
+
137
+ hidden_states = torch.randn(
138
+ num_tokens, layer.hidden_size, dtype=torch.bfloat16, device="cuda"
139
+ )
140
+ router_logits = torch.randn(
141
+ num_tokens, total_experts, dtype=torch.float32, device="cuda"
142
+ )
143
+
144
+ with torch.no_grad():
145
+ for param in layer.parameters():
146
+ if param.dtype.is_floating_point:
147
+ param.normal_(0, 0.01)
148
+
149
+ with set_current_vllm_config(vllm_config), set_forward_context(None, vllm_config):
150
+ get_forward_context().all_moe_layers = None
151
+
152
+ # Create a plain FusedMoE layer with the same config but no zero
153
+ # experts. Use a separate prefix to avoid collision.
154
+ plain_layer = FusedMoE(
155
+ num_experts=num_experts,
156
+ top_k=layer.top_k,
157
+ hidden_size=layer.hidden_size,
158
+ intermediate_size=layer.intermediate_size_per_partition,
159
+ params_dtype=torch.bfloat16,
160
+ prefix="test_zero_expert_moe_plain",
161
+ renormalize=False,
162
+ scoring_func="softmax",
163
+ e_score_correction_bias=layer.e_score_correction_bias,
164
+ ).cuda()
165
+
166
+ # Share weights from the zero expert layer.
167
+ plain_layer.w13_weight.data.copy_(layer.w13_weight.data)
168
+ plain_layer.w2_weight.data.copy_(layer.w2_weight.data)
169
+ plain_layer.quant_method.process_weights_after_loading(plain_layer)
170
+
171
+ # Compute routing via the ZeroExpertRouter. This produces masked
172
+ # topk_weights/topk_ids (zero expert entries have weight=0, id=0)
173
+ # and stores zero_expert_output as a side effect.
174
+ topk_weights, topk_ids = layer.router.select_experts(
175
+ hidden_states, router_logits
176
+ )
177
+ zero_output = layer.router.zero_expert_output
178
+
179
+ # Compute real expert output using the plain layer with the masked
180
+ # routing from the ZeroExpertRouter.
181
+ real_output = plain_layer.quant_method.apply(
182
+ layer=plain_layer,
183
+ x=hidden_states,
184
+ topk_weights=topk_weights,
185
+ topk_ids=topk_ids,
186
+ shared_experts_input=None,
187
+ )
188
+
189
+ # Get the combined output from the zero expert layer.
190
+ full_output = layer.forward(hidden_states, router_logits)
191
+
192
+ assert zero_output is not None, "Zero expert output should not be None"
193
+ assert not torch.isnan(real_output).any(), "Real expert output has NaN"
194
+ assert not torch.isnan(zero_output).any(), "Zero expert output has NaN"
195
+ assert not torch.isnan(full_output).any(), "Full output has NaN"
196
+
197
+ expected = real_output + zero_output
198
+ torch.testing.assert_close(
199
+ full_output,
200
+ expected,
201
+ atol=0,
202
+ rtol=0,
203
+ msg="FusedMoE output should equal plain FusedMoE output "
204
+ "plus zero expert contribution",
205
+ )
206
+
207
+
208
+@pytest.mark.parametrize("num_tokens", [1, 32])
209
+def test_zero_expert_moe_zero_expert_is_identity(zero_expert_moe, num_tokens):
210
+ """Validate zero expert identity behavior.
211
+
212
+ When routing strongly favors the zero expert, its contribution should
213
+ be a scaled version of hidden_states (identity operation). We verify
214
+ this by manually computing the expected zero expert output from the
215
+ routing weights and comparing against what the router produces.
216
+ """
217
+ layer, vllm_config = zero_expert_moe
218
+ num_experts = 4
219
+ zero_expert_num = 1
220
+ total_experts = num_experts + zero_expert_num
221
+
222
+ hidden_states = torch.randn(
223
+ num_tokens, layer.hidden_size, dtype=torch.bfloat16, device="cuda"
224
+ )
225
+ # Strongly bias toward the zero expert (index 4).
226
+ router_logits = torch.full(
227
+ (num_tokens, total_experts), -10.0, dtype=torch.float32, device="cuda"
228
+ )
229
+ router_logits[:, num_experts] = 10.0 # zero expert gets high logit
230
+
231
+ with torch.no_grad():
232
+ for param in layer.parameters():
233
+ if param.dtype.is_floating_point:
234
+ param.normal_(0, 0.01)
235
+
236
+ with set_current_vllm_config(vllm_config), set_forward_context(None, vllm_config):
237
+ get_forward_context().all_moe_layers = None
238
+
239
+ # Run routing to get topk_weights/topk_ids before masking.
240
+ from vllm.model_executor.layers.fused_moe.router.fused_topk_bias_router import (
241
+ fused_topk_bias,
242
+ )
243
+
244
+ topk_weights, topk_ids = fused_topk_bias(
245
+ hidden_states=hidden_states,
246
+ gating_output=router_logits,
247
+ e_score_correction_bias=layer.router.e_score_correction_bias.data,
248
+ topk=layer.top_k,
249
+ renormalize=layer.router.renormalize,
250
+ scoring_func=layer.router.scoring_func,
251
+ )
252
+
253
+ # Manually compute expected zero expert identity output:
254
+ # For each token, sum routing weights assigned to zero expert slots,
255
+ # then multiply by hidden_states.
256
+ zero_mask = topk_ids >= num_experts
257
+ zero_weight_per_token = (topk_weights * zero_mask.float()).sum(
258
+ dim=-1, keepdim=True
259
+ )
260
+ expected_zero_output = (hidden_states.float() * zero_weight_per_token).to(
261
+ hidden_states.dtype
262
+ )
263
+
264
+ # Run routing directly to trigger zero expert computation
265
+ # without going through the runner (which consumes the output).
266
+ layer.router.select_experts(hidden_states, router_logits)
267
+ actual_zero_output = layer.router.zero_expert_output
268
+
269
+ assert actual_zero_output is not None
270
+ assert zero_mask.any(), (
271
+ "With high zero expert logit, at least some slots should route "
272
+ "to the zero expert"
273
+ )
274
+
275
+ torch.testing.assert_close(
276
+ actual_zero_output,
277
+ expected_zero_output,
278
+ atol=1e-3,
279
+ rtol=1e-3,
280
+ msg="Zero expert identity output should equal "
281
+ "hidden_states * sum(zero_expert_weights)",
282
+ )
vllm/model_executor/layers/fused_moe/__init__.py
-4
@@ -33,9 +33,6 @@ from vllm.model_executor.layers.fused_moe.shared_fused_moe import SharedFusedMoE
33
from vllm.model_executor.layers.fused_moe.unquantized_fused_moe_method import (
34
UnquantizedFusedMoEMethod,
35
)
36
-from vllm.model_executor.layers.fused_moe.zero_expert_fused_moe import (
37
- ZeroExpertFusedMoE,
38
-)
36
from vllm.triton_utils import HAS_TRITON
37
38
_config: dict[str, Any] | None = None
@@ -68,7 +65,6 @@ __all__ = [
65
"GateLinear",
66
"RoutingMethodType",
67
"SharedFusedMoE",
71
- "ZeroExpertFusedMoE",
68
"activation_without_mul",
69
"apply_moe_activation",
70
"override_config",
vllm/model_executor/layers/fused_moe/layer.py
+3
@@ -274,6 +274,7 @@ class FusedMoE(PluggableLayer):
274
gate: torch.nn.Module | None = None,
275
shared_experts: torch.nn.Module | None = None,
276
routed_input_transform: torch.nn.Module | None = None,
277
+ zero_expert_type: str | None = None,
278
):
279
super().__init__()
280
@@ -462,6 +463,8 @@ class FusedMoE(PluggableLayer):
463
# TODO(bnell): once we can construct the MK at init time, we
464
# can make this a value.
465
indices_type_getter=lambda: self.quant_method.topk_indices_dtype,
466
+ zero_expert_type=zero_expert_type,
467
+ num_logical_experts=self.logical_num_experts,
468
)
469
self.routing_method_type: RoutingMethodType = self.router.routing_method_type
470
vllm/model_executor/layers/fused_moe/router/router_factory.py
+38
-4
@@ -25,6 +25,9 @@ from vllm.model_executor.layers.fused_moe.router.grouped_topk_router import (
25
from vllm.model_executor.layers.fused_moe.router.routing_simulator_router import (
26
RoutingSimulatorRouter,
27
)
28
+from vllm.model_executor.layers.fused_moe.router.zero_expert_router import (
29
+ ZeroExpertRouter,
30
+)
31
32
EMPTY_EPLB_STATE: EplbLayerState = EplbLayerState()
33
@@ -49,6 +52,9 @@ def create_fused_moe_router(
52
# eplb parameters
53
enable_eplb: bool = False,
54
eplb_state: EplbLayerState = EMPTY_EPLB_STATE,
55
+ # zero expert parameters
56
+ zero_expert_type: str | None = None,
57
+ num_logical_experts: int | None = None,
58
) -> FusedMoERouter:
59
"""
60
Factory function to create the appropriate FusedMoERouter subclass based on
@@ -56,10 +62,11 @@ def create_fused_moe_router(
62
63
The selection logic follows this priority order:
64
1. RoutingSimulatorRouter - if VLLM_MOE_ROUTING_SIMULATION_STRATEGY env var is set
59
- 2. GroupedTopKRouter - if use_grouped_topk is True
60
- 3. CustomRoutingRouter - if custom_routing_function is not None
61
- 4. FusedTopKBiasRouter - if e_score_correction_bias is not None
62
- 5. FusedTopKRouter - default fallback
65
+ 2. ZeroExpertRouter - if zero_expert_type is not None
66
+ 3. GroupedTopKRouter - if use_grouped_topk is True
67
+ 4. CustomRoutingRouter - if custom_routing_function is not None
68
+ 5. FusedTopKBiasRouter - if e_score_correction_bias is not None
69
+ 6. FusedTopKRouter - default fallback
70
71
Common arguments:
72
top_k: Number of experts to select per token
@@ -86,6 +93,12 @@ def create_fused_moe_router(
93
enable_eplb: Whether EPLB is enabled
94
eplb_state: EPLB (Expert Parallelism Load Balancing) state
95
96
+ Zero expert arguments:
97
+ zero_expert_type: Type of zero expert (e.g. identity). If not None,
98
+ creates a ZeroExpertRouter.
99
+ num_logical_experts: Number of real (non-zero) experts. Required when
100
+ zero_expert_type is not None.
101
+
102
Returns:
103
An instance of the appropriate FusedMoERouter subclass
104
"""
@@ -100,6 +113,27 @@ def create_fused_moe_router(
113
indices_type_getter=indices_type_getter,
114
)
115
116
+ if zero_expert_type is not None:
117
+ assert num_logical_experts is not None, (
118
+ "num_logical_experts is required when zero_expert_type is set"
119
+ )
120
+ assert e_score_correction_bias is not None, (
121
+ "e_score_correction_bias is required when zero_expert_type is set"
122
+ )
123
+ return ZeroExpertRouter(
124
+ top_k=top_k,
125
+ global_num_experts=global_num_experts,
126
+ eplb_state=eplb_state,
127
+ e_score_correction_bias=e_score_correction_bias,
128
+ num_logical_experts=num_logical_experts,
129
+ zero_expert_type=zero_expert_type,
130
+ scoring_func=scoring_func,
131
+ renormalize=renormalize,
132
+ routed_scaling_factor=routed_scaling_factor,
133
+ enable_eplb=enable_eplb,
134
+ indices_type_getter=indices_type_getter,
135
+ )
136
+
137
if use_grouped_topk:
138
assert custom_routing_function is None
139
if num_expert_group is None or topk_group is None:
vllm/model_executor/layers/fused_moe/router/zero_expert_router.py
new
+115
@@ -0,0 +1,115 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
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.config import (
10
+ RoutingMethodType,
11
+ get_routing_method_type,
12
+)
13
+from vllm.model_executor.layers.fused_moe.fused_moe import (
14
+ zero_experts_compute_triton,
15
+)
16
+from vllm.model_executor.layers.fused_moe.router.base_router import BaseRouter
17
+from vllm.model_executor.layers.fused_moe.router.fused_topk_bias_router import (
18
+ fused_topk_bias,
19
+)
20
+
21
+
22
+class ZeroExpertRouter(BaseRouter):
23
+ """Router that handles zero expert computation as part of routing.
24
+
25
+ Routes over all experts (real + zero) using full e_score_correction_bias.
26
+ Computes zero expert identity contributions as a side effect during routing.
27
+ Remaps zero expert IDs to real expert ID 0 (with weight 0) so downstream
28
+ MoE computation can ignore them.
29
+ """
30
+
31
+ def __init__(
32
+ self,
33
+ top_k: int,
34
+ global_num_experts: int,
35
+ eplb_state: EplbLayerState,
36
+ e_score_correction_bias: torch.Tensor,
37
+ num_logical_experts: int,
38
+ zero_expert_type: str,
39
+ scoring_func: str = "softmax",
40
+ renormalize: bool = False,
41
+ routed_scaling_factor: float = 1.0,
42
+ enable_eplb: bool = False,
43
+ indices_type_getter: Callable[[], torch.dtype | None] | None = None,
44
+ ):
45
+ super().__init__(
46
+ top_k=top_k,
47
+ global_num_experts=global_num_experts,
48
+ eplb_state=eplb_state,
49
+ enable_eplb=enable_eplb,
50
+ indices_type_getter=indices_type_getter,
51
+ )
52
+ self.e_score_correction_bias = e_score_correction_bias
53
+ self.num_logical_experts = num_logical_experts
54
+ self.zero_expert_type = zero_expert_type
55
+ self.scoring_func = scoring_func
56
+ self.renormalize = renormalize
57
+ self.routed_scaling_factor = routed_scaling_factor
58
+ self._zero_expert_output: torch.Tensor | None = None
59
+
60
+ @property
61
+ def routing_method_type(self) -> RoutingMethodType:
62
+ return get_routing_method_type(
63
+ scoring_func=self.scoring_func,
64
+ top_k=self.top_k,
65
+ renormalize=self.renormalize,
66
+ num_expert_group=None,
67
+ has_e_score_bias=True,
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 with full bias, compute zero expert output,
77
+ mask zero expert IDs."""
78
+ topk_weights, topk_ids = fused_topk_bias(
79
+ hidden_states=hidden_states,
80
+ gating_output=router_logits,
81
+ e_score_correction_bias=self.e_score_correction_bias.data,
82
+ topk=self.top_k,
83
+ renormalize=self.renormalize,
84
+ scoring_func=self.scoring_func,
85
+ indices_type=indices_type,
86
+ )
87
+
88
+ if self.routed_scaling_factor != 1.0:
89
+ topk_weights *= self.routed_scaling_factor
90
+
91
+ # Compute zero expert output using pre-EPLB topk_ids/weights.
92
+ # zero_experts_compute_triton modifies its inputs in-place, so
93
+ # pass clones.
94
+ self._zero_expert_output = zero_experts_compute_triton(
95
+ expert_indices=topk_ids.clone(),
96
+ expert_scales=topk_weights.clone(),
97
+ num_experts=self.num_logical_experts,
98
+ zero_expert_type=self.zero_expert_type,
99
+ hidden_states=hidden_states,
100
+ )
101
+
102
+ # Mask zero expert entries: remap zero expert IDs to 0 with weight 0
103
+ # so downstream MoE computation ignores them.
104
+ zero_mask = topk_ids >= self.num_logical_experts
105
+ topk_ids[zero_mask] = 0
106
+ topk_weights[zero_mask] = 0.0
107
+
108
+ return topk_weights, topk_ids
109
+
110
+ @property
111
+ def zero_expert_output(self) -> torch.Tensor | None:
112
+ """Retrieve and clear the zero expert output."""
113
+ output = self._zero_expert_output
114
+ self._zero_expert_output = None
115
+ return output
vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py
+19
-1
@@ -25,6 +25,9 @@ from vllm.model_executor.layers.fused_moe.fused_moe_method_base import (
25
from vllm.model_executor.layers.fused_moe.router.fused_moe_router import (
26
FusedMoERouter,
27
)
28
+from vllm.model_executor.layers.fused_moe.router.zero_expert_router import (
29
+ ZeroExpertRouter,
30
+)
31
from vllm.model_executor.layers.fused_moe.runner.moe_runner import MoERunner
32
from vllm.model_executor.layers.fused_moe.runner.shared_experts import (
33
SharedExperts,
@@ -443,6 +446,19 @@ class MoERunnerBase(MoERunner):
446
if self._shared_experts is not None:
447
self._shared_experts.maybe_sync_shared_experts_stream(shared_experts_input)
448
449
+ def _maybe_add_zero_expert_output(
450
+ self,
451
+ result: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
452
+ ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
453
+ if isinstance(self.router, ZeroExpertRouter):
454
+ zero_expert_output = self.router.zero_expert_output
455
+ assert zero_expert_output is not None
456
+ if isinstance(result, tuple):
457
+ result = (result[0], result[1] + zero_expert_output)
458
+ else:
459
+ result = result + zero_expert_output
460
+ return result
461
+
462
def forward(
463
self,
464
hidden_states: torch.Tensor,
@@ -494,7 +510,9 @@ class MoERunnerBase(MoERunner):
510
self._encode_layer_name(),
511
)
512
497
- return self._maybe_reduce_output(fused_output, og_hidden_dims)
513
+ result = self._maybe_reduce_output(fused_output, og_hidden_dims)
514
+
515
+ return self._maybe_add_zero_expert_output(result)
516
517
def forward_dispatch(
518
self,
vllm/model_executor/layers/fused_moe/zero_expert_fused_moe.py
deleted
-189
@@ -1,189 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
4
-from contextlib import contextmanager
5
-
6
-import torch
7
-from torch import nn
8
-
9
-from vllm.model_executor.layers.fused_moe.fused_moe import zero_experts_compute_triton
10
-from vllm.model_executor.layers.fused_moe.layer import FusedMoE
11
-
12
-
13
-class ZeroExpertFusedMoE(FusedMoE):
14
- """
15
- A FusedMoE operation that also computes the results of zero experts.
16
- Zero experts perform identity operations (scaled pass-through) instead
17
- of full MLP computations.
18
-
19
- This class uses memoization to avoid redundant routing computation:
20
- routing is computed once and reused for both zero expert computation
21
- and the main FusedMoE forward pass.
22
- """
23
-
24
- def __init__(
25
- self,
26
- zero_expert_num: int,
27
- zero_expert_type: str,
28
- router: nn.Module,
29
- **kwargs,
30
- ):
31
- # ZeroExpertFusedMoE manages its own custom_routing_function for memoization
32
- assert (
33
- "custom_routing_function" not in kwargs
34
- or kwargs.get("custom_routing_function") is None
35
- ), (
36
- "ZeroExpertFusedMoE does not support external custom_routing_function. "
37
- "It manages its own for routing memoization."
38
- )
39
-
40
- # Automatically slice router's e_score_correction_bias to only include
41
- # real experts (not zero_experts) for the base FusedMoE.
42
- # The full bias will be used temporarily in forward() for routing.
43
- if hasattr(router, "e_score_correction_bias") and "num_experts" in kwargs:
44
- num_real_experts = kwargs["num_experts"]
45
- router_bias = router.e_score_correction_bias
46
- user_bias = kwargs.get("e_score_correction_bias")
47
-
48
- # Use router's bias if:
49
- # 1. User didn't provide bias, or
50
- # 2. User provided full bias (same size as router)
51
- if user_bias is None or user_bias.shape[0] == router_bias.shape[0]:
52
- kwargs["e_score_correction_bias"] = router_bias[:num_real_experts]
53
-
54
- # FusedMoE no longer accepts zero_expert_num/zero_expert_type.
55
- # We handle zero experts ourselves in forward().
56
- super().__init__(**kwargs)
57
- # Store the actual zero_expert_num and zero_expert_type for our own use
58
- self._actual_zero_expert_num = zero_expert_num
59
- self._actual_zero_expert_type = zero_expert_type
60
- self._router = router # Full router (includes zero experts)
61
-
62
- # Expose zero_expert_num and zero_expert_type as attributes for
63
- # compatibility with quantization methods that check these attributes
64
- self.zero_expert_num = 0
65
- self.zero_expert_type = None
66
-
67
- # Memoization state for routing results
68
- self._memoized_topk_weights: torch.Tensor | None = None
69
- self._memoized_topk_ids: torch.Tensor | None = None
70
-
71
- # Create custom_routing_function to reuse memoized routing results
72
- def custom_routing_function(hidden_states, gating_output, topk, renormalize):
73
- """Return memoized `topk_weights` and `topk_ids`."""
74
- if self._memoized_topk_weights is None or self._memoized_topk_ids is None:
75
- raise RuntimeError(
76
- "ZeroExpertFusedMoE: routing results not memoized. "
77
- "Call select_experts first to compute routing."
78
- )
79
- return self._memoized_topk_weights, self._memoized_topk_ids
80
-
81
- self.custom_routing_function = custom_routing_function
82
-
83
- @contextmanager
84
- def _temporarily_set_attrs(self, **attrs):
85
- """
86
- Temporarily set attributes using object.__setattr__ and restore them.
87
-
88
- This bypasses nn.Module.__setattr__ to avoid Dynamo tracing issues.
89
- When PyTorch Dynamo traces the forward pass, it cannot handle
90
- nn.Module.__setattr__ calls (which include parameter registration logic),
91
- resulting in "Unsupported" errors. Using object.__setattr__ directly
92
- sets the attribute without triggering nn.Module's custom __setattr__,
93
- allowing Dynamo to trace the code successfully.
94
- """
95
- originals = {key: getattr(self, key) for key in attrs}
96
- try:
97
- for key, value in attrs.items():
98
- object.__setattr__(self, key, value)
99
- yield
100
- finally:
101
- for key, value in originals.items():
102
- object.__setattr__(self, key, value)
103
-
104
- def _compute_zero_expert_result(
105
- self,
106
- hidden_states: torch.Tensor,
107
- topk_weights: torch.Tensor,
108
- topk_ids: torch.Tensor,
109
- ) -> torch.Tensor | None:
110
- """Compute zero expert results using pre-computed routing."""
111
- if (
112
- self._actual_zero_expert_num is None
113
- or self._actual_zero_expert_num <= 0
114
- or self._actual_zero_expert_type is None
115
- ):
116
- return None
117
-
118
- return zero_experts_compute_triton(
119
- expert_indices=topk_ids.clone(),
120
- expert_scales=topk_weights.clone(),
121
- num_experts=self.logical_num_experts,
122
- zero_expert_type=self._actual_zero_expert_type,
123
- hidden_states=hidden_states,
124
- )
125
-
126
- def forward(
127
- self,
128
- hidden_states: torch.Tensor,
129
- router_logits: torch.Tensor, # Full logits including zero experts
130
- ) -> torch.Tensor:
131
- """
132
- Forward pass with zero expert support and routing memoization.
133
-
134
- Args:
135
- hidden_states: Input hidden states
136
- router_logits: Full router logits (including zero experts)
137
-
138
- Returns:
139
- Combined output from real experts and zero experts
140
- """
141
- # Prepare temporary attribute overrides for routing computation
142
- temp_attrs = {
143
- "custom_routing_function": None, # Disable for first routing
144
- }
145
- if self._router is not None:
146
- temp_attrs["e_score_correction_bias"] = self._router.e_score_correction_bias
147
-
148
- # Compute routing with temporary attributes
149
- # Pass full router_logits (including zero experts) so that zero experts
150
- # can be properly identified in topk_ids
151
- with self._temporarily_set_attrs(**temp_attrs):
152
- topk_weights, topk_ids = self.select_experts(
153
- hidden_states=hidden_states,
154
- router_logits=router_logits, # Full logits (includes zero experts)
155
- )
156
-
157
- # Compute zero expert result if needed
158
- zero_expert_result = self._compute_zero_expert_result(
159
- hidden_states=hidden_states,
160
- topk_weights=topk_weights,
161
- topk_ids=topk_ids,
162
- )
163
-
164
- # Memoize routing results for reuse in super().forward()
165
- self._memoized_topk_weights = topk_weights
166
- self._memoized_topk_ids = topk_ids
167
-
168
- # Slice router_logits for real experts only
169
- router_logits_sliced = router_logits[..., : self.logical_num_experts]
170
-
171
- # Compute real expert results (will reuse memoized routing via
172
- # custom_routing_function)
173
- # zero_expert_num is already 0, so FusedMoE won't handle zero experts
174
- fused_out = super().forward(
175
- hidden_states=hidden_states,
176
- router_logits=router_logits_sliced,
177
- )
178
-
179
- # Combine results
180
- # Both zero_expert_result and fused_out are computed from the same
181
- # hidden_states, so they should be on the same device.
182
- if zero_expert_result is not None:
183
- fused_out = fused_out + zero_expert_result
184
-
185
- # Clear memoization after use
186
- self._memoized_topk_weights = None
187
- self._memoized_topk_ids = None
188
-
189
- return fused_out
vllm/model_executor/models/longcat_flash.py
+4
-6
@@ -46,7 +46,7 @@ from vllm.config import CacheConfig, VllmConfig
46
from vllm.distributed import get_pp_group
47
from vllm.logger import init_logger
48
from vllm.model_executor.layers.activation import SiluAndMul
49
-from vllm.model_executor.layers.fused_moe import FusedMoE, ZeroExpertFusedMoE
49
+from vllm.model_executor.layers.fused_moe import FusedMoE
50
from vllm.model_executor.layers.layernorm import RMSNorm
51
from vllm.model_executor.layers.linear import (
52
MergedColumnParallelLinear,
@@ -292,12 +292,10 @@ class LongcatMoe(nn.Module):
292
prefix=f"{prefix}.gate",
293
)
294
295
- assert config.zero_expert_num is not None
295
assert config.zero_expert_type is not None
297
- self.experts = ZeroExpertFusedMoE(
298
- zero_expert_num=config.zero_expert_num,
296
+ self.experts = FusedMoE(
297
zero_expert_type=config.zero_expert_type,
300
- router=self.router,
298
+ e_score_correction_bias=self.router.e_score_correction_bias,
299
num_experts=num_experts,
300
top_k=top_k,
301
hidden_size=hidden_size,
@@ -332,7 +330,7 @@ class LongcatMoe(nn.Module):
330
hidden_states_padded.to(self.router_params_dtype)
331
)
332
335
- # ZeroExpertFusedMoE handles routing memoization and zero expert computation
333
+ # FusedMoE handles routing memoization and zero expert computation
334
# internally. Pass full router_logits (including zero experts) so that
335
# zero experts can be properly identified in routing.
336
final_hidden_states = self.experts(