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 11, 03:45 UTC
[MoE Refactor][5/N] Isolate zero expert to LongCatFlash (#28891)
Signed-off-by: baonudesifeizhai <85092850+baonudesifeizhai@users.noreply.github.com> Signed-off-by: Dongjie Zou <85092850+baonudesifeizhai@users.noreply.github.com> Signed-off-by: baonudesifeizhai <baonudesifeizhai@gmail.com> Signed-off-by: Robert Shaw <robertgshaw2@gmail.com> Co-authored-by: Robert Shaw <robshaw@redhat.com> Co-authored-by: Robert Shaw <robertgshaw2@gmail.com>
baonudesifeizhai committed
Dec 20, 2025 at 13:22 UTC
54c892438479c0a8aec9a10b8db570034af92443
19 files changed
+263
-108
tests/test_routing_simulator.py
+1
-1
@@ -127,7 +127,7 @@ def test_routing_strategy_integration(monkeypatch, device):
127
envs.environment_variables[env_name] = lambda s=strategy: s
128
129
# Test the select_experts method
130
- topk_weights, topk_ids, _ = fused_moe.select_experts(
130
+ topk_weights, topk_ids = fused_moe.select_experts(
131
hidden_states=hidden_states,
132
router_logits=router_logits,
133
)
vllm/model_executor/layers/fused_moe/__init__.py
+4
@@ -25,6 +25,9 @@ from vllm.model_executor.layers.fused_moe.unquantized_fused_moe_method import (
25
UnquantizedFusedMoEMethod,
26
)
27
from vllm.model_executor.layers.fused_moe.utils import activation_without_mul
28
+from vllm.model_executor.layers.fused_moe.zero_expert_fused_moe import (
29
+ ZeroExpertFusedMoE,
30
+)
31
from vllm.triton_utils import HAS_TRITON
32
33
_config: dict[str, Any] | None = None
@@ -54,6 +57,7 @@ __all__ = [
57
"FusedMoEPrepareAndFinalize",
58
"RoutingMethodType",
59
"SharedFusedMoE",
60
+ "ZeroExpertFusedMoE",
61
"activation_without_mul",
62
"override_config",
63
"get_config",
vllm/model_executor/layers/fused_moe/fused_moe_modular_method.py
+2
-8
@@ -92,7 +92,7 @@ class FusedMoEModularMethod(FusedMoEMethodBase, CustomOp):
92
x: torch.Tensor,
93
router_logits: torch.Tensor,
94
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
95
- topk_weights, topk_ids, zero_expert_result = layer.select_experts(
95
+ topk_weights, topk_ids = layer.select_experts(
96
hidden_states=x,
97
router_logits=router_logits,
98
)
@@ -110,10 +110,4 @@ class FusedMoEModularMethod(FusedMoEMethodBase, CustomOp):
110
expert_map=None if self.disable_expert_map else layer.expert_map,
111
)
112
113
- if layer.zero_expert_num != 0 and layer.zero_expert_type is not None:
114
- assert not isinstance(result, tuple), (
115
- "Shared + zero experts are mutually exclusive not yet supported"
116
- )
117
- return result, zero_expert_result
118
- else:
119
- return result
113
+ return result
vllm/model_executor/layers/fused_moe/layer.py
+6
-47
@@ -32,7 +32,6 @@ from vllm.model_executor.layers.fused_moe.config import (
32
FusedMoEQuantConfig,
33
RoutingMethodType,
34
)
35
-from vllm.model_executor.layers.fused_moe.fused_moe import zero_experts_compute_triton
35
from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (
36
init_aiter_topK_meta_data,
37
)
@@ -350,8 +349,6 @@ class FusedMoE(CustomOp):
349
num_redundant_experts: int = 0,
350
has_bias: bool = False,
351
is_sequence_parallel=False,
353
- zero_expert_num: int | None = 0,
354
- zero_expert_type: str | None = None,
352
expert_mapping: list[tuple[str, str, int, str]] | None = None,
353
n_shared_experts: int | None = None,
354
routing_method_type: int | None = None,
@@ -409,8 +406,6 @@ class FusedMoE(CustomOp):
406
407
self.global_num_experts = num_experts + num_redundant_experts
408
self.logical_num_experts = num_experts
412
- self.zero_expert_num = zero_expert_num
413
- self.zero_expert_type = zero_expert_type
409
410
# Expert mapping used in self.load_weights
411
self.expert_mapping = expert_mapping
@@ -1525,15 +1520,15 @@ class FusedMoE(CustomOp):
1520
self,
1521
hidden_states: torch.Tensor,
1522
router_logits: torch.Tensor,
1528
- ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:
1523
+ ) -> tuple[torch.Tensor, torch.Tensor]:
1524
"""
1525
Route the input hidden states to the top-k experts based on the
1526
router logits.
1527
1528
Returns:
1534
- (topk_weights, topk_ids, zero_expert_result)
1535
- (tuple[torch.Tensor, torch.Tensor, torch.Tensor]):
1536
- The weights, expert ids, and zero expert computation result.
1529
+ (topk_weights, topk_ids)
1530
+ (tuple[torch.Tensor, torch.Tensor]):
1531
+ The weights and expert ids.
1532
1533
**Compatibility**: When EPLB is not enabled, the returned ids are
1534
equivalent to global logical ids, so should be compatible with
@@ -1655,23 +1650,7 @@ class FusedMoE(CustomOp):
1650
1651
assert topk_ids.dtype == indices_type or indices_type is None
1652
1658
- # Compute zero expert result if needed
1659
- if (
1660
- self.zero_expert_num is not None
1661
- and self.zero_expert_num > 0
1662
- and self.zero_expert_type is not None
1663
- and self.global_num_experts is not None
1664
- ):
1665
- zero_expert_result = zero_experts_compute_triton(
1666
- expert_indices=topk_ids,
1667
- expert_scales=topk_weights,
1668
- num_experts=self.global_num_experts,
1669
- zero_expert_type=self.zero_expert_type,
1670
- hidden_states=hidden_states,
1671
- )
1672
- else:
1673
- zero_expert_result = None
1674
- return topk_weights, topk_ids, zero_expert_result
1653
+ return topk_weights, topk_ids
1654
1655
def must_reduce_shared_expert_outputs(self) -> bool:
1656
"""
@@ -1736,14 +1715,7 @@ class FusedMoE(CustomOp):
1715
fused_output = torch.ops.vllm.moe_forward(
1716
hidden_states, router_logits, self.layer_name
1717
)
1739
- if self.zero_expert_num is not None and self.zero_expert_num > 0:
1740
- assert isinstance(fused_output, tuple)
1741
- fused_output, zero_expert_result = fused_output
1742
- return (reduce_output(fused_output) + zero_expert_result)[
1743
- ..., :og_hidden_states
1744
- ]
1745
- else:
1746
- return reduce_output(fused_output)[..., :og_hidden_states]
1718
+ return reduce_output(fused_output)[..., :og_hidden_states]
1719
else:
1720
if current_platform.is_tpu() or current_platform.is_cpu():
1721
# TODO: Once the OOM issue for the TPU backend is resolved, we
@@ -1841,13 +1813,6 @@ class FusedMoE(CustomOp):
1813
final_hidden_states,
1814
)
1815
1844
- if self.zero_expert_num is not None and self.zero_expert_num > 0:
1845
- assert isinstance(final_hidden_states, tuple)
1846
- assert self.shared_experts is None
1847
- final_hidden_states, zero_expert_result = final_hidden_states
1848
- if zero_expert_result is not None:
1849
- final_hidden_states += zero_expert_result
1850
-
1816
if not skip_result_store:
1817
if self.shared_experts is None:
1818
full_fused_final_hidden_states[chunk_start:chunk_end, :].copy_(
@@ -2030,9 +1995,6 @@ class FusedMoE(CustomOp):
1995
shared_output,
1996
final_hidden_states,
1997
)
2033
- elif self.zero_expert_num is not None and self.zero_expert_num > 0:
2034
- assert isinstance(final_hidden_states, tuple)
2035
- final_hidden_states, zero_expert_result = final_hidden_states
1998
1999
def combine_output(states: torch.Tensor) -> torch.Tensor:
2000
if do_naive_dispatch_combine:
@@ -2051,9 +2013,6 @@ class FusedMoE(CustomOp):
2013
final_hidden_states[0],
2014
combine_output(final_hidden_states[1]),
2015
)
2054
- elif self.zero_expert_num is not None and self.zero_expert_num > 0:
2055
- assert isinstance(final_hidden_states, torch.Tensor)
2056
- return (combine_output(final_hidden_states), zero_expert_result)
2016
else:
2017
return combine_output(final_hidden_states)
2018
vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py
+2
-8
@@ -295,7 +295,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
295
x: torch.Tensor,
296
router_logits: torch.Tensor,
297
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
298
- topk_weights, topk_ids, zero_expert_result = layer.select_experts(
298
+ topk_weights, topk_ids = layer.select_experts(
299
hidden_states=x,
300
router_logits=router_logits,
301
)
@@ -336,13 +336,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
336
expert_map=layer.expert_map,
337
)
338
339
- if layer.zero_expert_num != 0 and layer.zero_expert_type is not None:
340
- assert not isinstance(result, tuple), (
341
- "Shared + zero experts are mutually exclusive not yet supported"
342
- )
343
- return result, zero_expert_result
344
- else:
345
- return result
339
+ return result
340
341
def forward_cpu(
342
self,
vllm/model_executor/layers/fused_moe/zero_expert_fused_moe.py
new
+189
@@ -0,0 +1,189 @@
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/layers/quantization/awq_marlin.py
+1
-1
@@ -764,7 +764,7 @@ class AWQMarlinMoEMethod(FusedMoEMethodBase):
764
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
765
assert layer.activation == "silu", "Only SiLU activation is supported."
766
767
- topk_weights, topk_ids, _ = layer.select_experts(
767
+ topk_weights, topk_ids = layer.select_experts(
768
hidden_states=x,
769
router_logits=router_logits,
770
)
vllm/model_executor/layers/quantization/bitsandbytes.py
+1
-1
@@ -500,7 +500,7 @@ class BitsAndBytesMoEMethod(FusedMoEMethodBase):
500
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
501
from vllm.model_executor.layers.fused_moe import fused_experts
502
503
- topk_weights, topk_ids, _ = layer.select_experts(
503
+ topk_weights, topk_ids = layer.select_experts(
504
hidden_states=x,
505
router_logits=router_logits,
506
)
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py
+6
-6
@@ -574,7 +574,7 @@ class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod):
574
e_score_correction_bias=layer.e_score_correction_bias,
575
)
576
577
- topk_weights, topk_ids, _ = layer.select_experts(
577
+ topk_weights, topk_ids = layer.select_experts(
578
hidden_states=x,
579
router_logits=router_logits,
580
)
@@ -1166,7 +1166,7 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):
1166
x: torch.Tensor,
1167
router_logits: torch.Tensor,
1168
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
1169
- topk_weights, topk_ids, _ = layer.select_experts(
1169
+ topk_weights, topk_ids = layer.select_experts(
1170
hidden_states=x,
1171
router_logits=router_logits,
1172
)
@@ -1403,7 +1403,7 @@ class CompressedTensorsW8A8Int8MoEMethod(CompressedTensorsMoEMethod):
1403
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
1404
from vllm.model_executor.layers.fused_moe import fused_experts
1405
1406
- topk_weights, topk_ids, _ = layer.select_experts(
1406
+ topk_weights, topk_ids = layer.select_experts(
1407
hidden_states=x,
1408
router_logits=router_logits,
1409
)
@@ -1765,7 +1765,7 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):
1765
f"{layer.activation} not supported for Marlin MoE."
1766
)
1767
1768
- topk_weights, topk_ids, _ = layer.select_experts(
1768
+ topk_weights, topk_ids = layer.select_experts(
1769
hidden_states=x,
1770
router_logits=router_logits,
1771
)
@@ -1991,7 +1991,7 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod):
1991
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
1992
from vllm.model_executor.layers.fused_moe import fused_experts
1993
1994
- topk_weights, topk_ids, _ = layer.select_experts(
1994
+ topk_weights, topk_ids = layer.select_experts(
1995
hidden_states=x,
1996
router_logits=router_logits,
1997
)
@@ -2607,7 +2607,7 @@ class CompressedTensorsW4A8Fp8MoEMethod(CompressedTensorsMoEMethod):
2607
"EPLB not supported for `CompressedTensorsW4A8Fp8MoEMethod` yet."
2608
)
2609
assert self.moe_quant_config is not None
2610
- topk_weights, topk_ids, _ = layer.select_experts(
2610
+ topk_weights, topk_ids = layer.select_experts(
2611
hidden_states=x,
2612
router_logits=router_logits,
2613
)
vllm/model_executor/layers/quantization/experts_int8.py
+1
-1
@@ -142,7 +142,7 @@ class ExpertsInt8MoEMethod(FusedMoEMethodBase):
142
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
143
from vllm.model_executor.layers.fused_moe import fused_experts
144
145
- topk_weights, topk_ids, _ = layer.select_experts(
145
+ topk_weights, topk_ids = layer.select_experts(
146
hidden_states=x,
147
router_logits=router_logits,
148
)
vllm/model_executor/layers/quantization/fp8.py
+2
-10
@@ -1292,13 +1292,11 @@ class Fp8MoEMethod(FusedMoEMethodBase):
1292
apply_router_weight_on_input=layer.apply_router_weight_on_input,
1293
)
1294
1295
- select_result = layer.select_experts(
1295
+ topk_weights, topk_ids = layer.select_experts(
1296
hidden_states=x,
1297
router_logits=router_logits,
1298
)
1299
1300
- topk_weights, topk_ids, zero_expert_result = select_result
1301
-
1300
if self.rocm_aiter_moe_enabled:
1301
from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import ( # noqa: E501
1302
rocm_aiter_fused_experts,
@@ -1353,13 +1351,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):
1351
apply_router_weight_on_input=layer.apply_router_weight_on_input,
1352
)
1353
1356
- if layer.zero_expert_num != 0 and layer.zero_expert_type is not None:
1357
- assert not isinstance(result, tuple), (
1358
- "Shared + zero experts are mutually exclusive not yet supported"
1359
- )
1360
- return result, zero_expert_result
1361
- else:
1362
- return result
1354
+ return result
1355
1356
1357
class Fp8OnlineMoEMethod(Fp8MoEMethod):
vllm/model_executor/layers/quantization/gguf.py
+1
-1
@@ -639,7 +639,7 @@ class GGUFMoEMethod(FusedMoEMethodBase):
639
"fused GGUF MoE method."
640
)
641
642
- topk_weights, topk_ids, _ = layer.select_experts(
642
+ topk_weights, topk_ids = layer.select_experts(
643
hidden_states=x,
644
router_logits=router_logits,
645
)
vllm/model_executor/layers/quantization/gptq_marlin.py
+1
-1
@@ -900,7 +900,7 @@ class GPTQMarlinMoEMethod(FusedMoEMethodBase):
900
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
901
assert layer.activation == "silu", "Only SiLU activation is supported."
902
903
- topk_weights, topk_ids, _ = layer.select_experts(
903
+ topk_weights, topk_ids = layer.select_experts(
904
hidden_states=x,
905
router_logits=router_logits,
906
)
vllm/model_executor/layers/quantization/modelopt.py
+2
-2
@@ -796,7 +796,7 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):
796
)
797
798
# Expert selection
799
- topk_weights, topk_ids, _ = layer.select_experts(
799
+ topk_weights, topk_ids = layer.select_experts(
800
hidden_states=x,
801
router_logits=router_logits,
802
)
@@ -1599,7 +1599,7 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
1599
x_routing, _ = x
1600
else:
1601
x_routing = x
1602
- topk_weights, topk_ids, _ = layer.select_experts(
1602
+ topk_weights, topk_ids = layer.select_experts(
1603
hidden_states=x_routing,
1604
router_logits=router_logits,
1605
)
vllm/model_executor/layers/quantization/moe_wna16.py
+1
-1
@@ -370,7 +370,7 @@ class MoeWNA16Method(FusedMoEMethodBase):
370
from vllm.model_executor.layers.fused_moe import fused_experts
371
372
assert layer.activation == "silu", "Only SiLU activation is supported."
373
- topk_weights, topk_ids, _ = layer.select_experts(
373
+ topk_weights, topk_ids = layer.select_experts(
374
hidden_states=x,
375
router_logits=router_logits,
376
)
vllm/model_executor/layers/quantization/mxfp4.py
+2
-2
@@ -896,7 +896,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
896
raise NotImplementedError("EPLB is not supported for mxfp4")
897
898
if self.mxfp4_backend == Mxfp4Backend.MARLIN:
899
- topk_weights, topk_ids, _ = layer.select_experts(
899
+ topk_weights, topk_ids = layer.select_experts(
900
hidden_states=x,
901
router_logits=router_logits,
902
)
@@ -990,7 +990,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
990
):
991
from vllm.utils.flashinfer import flashinfer_cutlass_fused_moe
992
993
- topk_weights, topk_ids, _ = layer.select_experts(
993
+ topk_weights, topk_ids = layer.select_experts(
994
hidden_states=x,
995
router_logits=router_logits,
996
)
vllm/model_executor/layers/quantization/quark/quark_moe.py
+3
-3
@@ -338,7 +338,7 @@ class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod):
338
x: torch.Tensor,
339
router_logits: torch.Tensor,
340
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
341
- topk_weights, topk_ids, _ = layer.select_experts(
341
+ topk_weights, topk_ids = layer.select_experts(
342
hidden_states=x,
343
router_logits=router_logits,
344
)
@@ -530,7 +530,7 @@ class QuarkW4A8Fp8MoEMethod(QuarkMoEMethod):
530
x: torch.Tensor,
531
router_logits: torch.Tensor,
532
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
533
- topk_weights, topk_ids, _ = layer.select_experts(
533
+ topk_weights, topk_ids = layer.select_experts(
534
hidden_states=x,
535
router_logits=router_logits,
536
)
@@ -738,7 +738,7 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):
738
x: torch.Tensor,
739
router_logits: torch.Tensor,
740
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
741
- topk_weights, topk_ids, _ = layer.select_experts(
741
+ topk_weights, topk_ids = layer.select_experts(
742
hidden_states=x,
743
router_logits=router_logits,
744
)
vllm/model_executor/layers/quantization/rtn.py
+1
-1
@@ -359,7 +359,7 @@ class RTNMoEMethod(FusedMoEMethodBase):
359
x: torch.Tensor,
360
router_logits: torch.Tensor,
361
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
362
- topk_weights, topk_ids, _ = layer.select_experts(
362
+ topk_weights, topk_ids = layer.select_experts(
363
hidden_states=x,
364
router_logits=router_logits,
365
)
vllm/model_executor/models/longcat_flash.py
+37
-14
@@ -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
49
+from vllm.model_executor.layers.fused_moe import FusedMoE, ZeroExpertFusedMoE
50
from vllm.model_executor.layers.layernorm import RMSNorm
51
from vllm.model_executor.layers.linear import (
52
MergedColumnParallelLinear,
@@ -179,7 +179,7 @@ class FlashConfig(PretrainedConfig):
179
self.intermediate_size = (
180
self.ffn_hidden_size
181
if hasattr(self, "ffn_hidden_size")
182
- else self.intermediate_size
182
+ else intermediate_size
183
)
184
if hasattr(self, "moe_intermediate_size"):
185
self.moe_intermediate_size = self.moe_intermediate_size
@@ -280,10 +280,6 @@ class LongcatMoe(nn.Module):
280
):
281
super().__init__()
282
self.hidden_size = hidden_size
283
- self.zero_expert_num = config.zero_expert_num
284
- self.zero_expert_type = config.zero_expert_type
285
- self.routed_scaling_factor = config.routed_scaling_factor
286
- self.enable_eplb = enable_eplb
283
# Gate always runs at half / full precision for now.
284
self.rounter_params_dtype = params_dtype
285
if config.router_dtype == "float32":
@@ -291,25 +287,27 @@ class LongcatMoe(nn.Module):
287
288
self.router = LongcatRouter(
289
config=config,
294
- zero_expert_num=self.zero_expert_num,
290
+ zero_expert_num=config.zero_expert_num,
291
rounter_params_dtype=self.rounter_params_dtype,
292
prefix=f"{prefix}.gate",
293
)
294
299
- self.experts = FusedMoE(
295
+ assert config.zero_expert_num is not None
296
+ assert config.zero_expert_type is not None
297
+ self.experts = ZeroExpertFusedMoE(
298
+ zero_expert_num=config.zero_expert_num,
299
+ zero_expert_type=config.zero_expert_type,
300
+ router=self.router,
301
num_experts=num_experts,
302
top_k=top_k,
303
hidden_size=hidden_size,
304
intermediate_size=intermediate_size,
305
reduce_results=True,
306
params_dtype=params_dtype,
306
- e_score_correction_bias=self.router.e_score_correction_bias,
307
renormalize=False,
308
quant_config=quant_config,
309
prefix=f"{prefix}.experts",
310
- zero_expert_num=self.zero_expert_num,
311
- zero_expert_type=self.zero_expert_type,
312
- enable_eplb=self.enable_eplb,
310
+ enable_eplb=enable_eplb,
311
routed_scaling_factor=config.routed_scaling_factor,
312
)
313
@@ -317,11 +315,34 @@ class LongcatMoe(nn.Module):
315
num_tokens, hidden_dim = hidden_states.shape
316
hidden_states = hidden_states.view(-1, hidden_dim)
317
320
- router_logits = self.router(hidden_states.to(self.rounter_params_dtype))
318
+ # Align to FusedMoE padded hidden size to avoid dim mismatch
319
+ padded_hidden = self.experts.hidden_size
320
+ if hidden_dim < padded_hidden:
321
+ hidden_states_padded = torch.nn.functional.pad(
322
+ hidden_states,
323
+ (0, padded_hidden - hidden_dim),
324
+ mode="constant",
325
+ value=0.0,
326
+ )
327
+ else:
328
+ hidden_states_padded = hidden_states
329
+
330
+ router_logits_full = self.router(
331
+ hidden_states_padded.to(self.rounter_params_dtype)
332
+ )
333
+
334
+ # ZeroExpertFusedMoE handles routing memoization and zero expert computation
335
+ # internally. Pass full router_logits (including zero experts) so that
336
+ # zero experts can be properly identified in routing.
337
final_hidden_states = self.experts(
322
- hidden_states=hidden_states, router_logits=router_logits
338
+ hidden_states=hidden_states_padded,
339
+ router_logits=router_logits_full, # Full logits (includes zero experts)
340
)
341
342
+ # Crop back to original hidden dimension if padded earlier
343
+ if padded_hidden != hidden_dim:
344
+ final_hidden_states = final_hidden_states[..., :hidden_dim]
345
+
346
return final_hidden_states.view(num_tokens, hidden_dim)
347
348
@@ -419,6 +440,7 @@ class FlashDecoderLayer(nn.Module):
440
hidden_states = self.self_attn[0](
441
positions=positions,
442
hidden_states=hidden_states,
443
+ llama_4_scaling=None,
444
)
445
446
hidden_states, residual = self.post_attention_layernorm[0](
@@ -438,6 +460,7 @@ class FlashDecoderLayer(nn.Module):
460
hidden_states = self.self_attn[1](
461
positions=positions,
462
hidden_states=hidden_states,
463
+ llama_4_scaling=None,
464
)
465
hidden_states, residual = self.post_attention_layernorm[1](
466
hidden_states, residual