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 14, 05:30 UTC
[Model] Apply shared experts overlap optimization to all models with shared experts (#26145)
Signed-off-by: Bill Nell <bnell@redhat.com>
bnellnm committed
Oct 9, 2025 at 11:31 UTC
47e66c24e2770b094ba8e160ddb282d6861737f6
15 files changed
+271
-283
vllm/model_executor/layers/fused_moe/__init__.py
+2
@@ -15,6 +15,7 @@ from vllm.model_executor.layers.fused_moe.modular_kernel import (
15
FusedMoEPermuteExpertsUnpermute,
16
FusedMoEPrepareAndFinalize,
17
)
18
+from vllm.model_executor.layers.fused_moe.shared_fused_moe import SharedFusedMoE
19
from vllm.model_executor.layers.fused_moe.utils import activation_without_mul
20
from vllm.triton_utils import HAS_TRITON
21
@@ -42,6 +43,7 @@ __all__ = [
43
"FusedMoEPermuteExpertsUnpermute",
44
"FusedMoEActivationFormat",
45
"FusedMoEPrepareAndFinalize",
46
+ "SharedFusedMoE",
47
"activation_without_mul",
48
"override_config",
49
"get_config",
vllm/model_executor/layers/fused_moe/shared_fused_moe.py
renamed
+23
-12
@@ -18,13 +18,21 @@ class SharedFusedMoE(FusedMoE):
18
19
def __init__(
20
self,
21
- shared_experts: torch.nn.Module,
21
+ shared_experts: Optional[torch.nn.Module],
22
use_overlapped: bool = True,
23
**kwargs,
24
):
25
super().__init__(**kwargs)
26
self._shared_experts = shared_experts
27
- self.use_overlapped = use_overlapped
27
+ # Disable shared expert overlap if EP is disabled or we are not using
28
+ # flashinfer + DP since there is nothing to be gained in this case.
29
+ # Disabling the overlap optimization also prevents the shared experts
30
+ # from being hidden from torch.compile.
31
+ self.use_overlapped = (
32
+ use_overlapped
33
+ and not (self.use_ep or self.use_flashinfer_cutlass_kernels)
34
+ and self._shared_experts is not None
35
+ )
36
37
@property
38
def shared_experts(self) -> Optional[torch.nn.Module]:
@@ -36,16 +44,19 @@ class SharedFusedMoE(FusedMoE):
44
router_logits: torch.Tensor,
45
) -> tuple[torch.Tensor, torch.Tensor]:
46
if not self.use_overlapped:
39
- shared_out = self._shared_experts(hidden_states)
40
-
41
- # Reduce outputs if necessary, since the MLP should
42
- # have been created with reduce_results=False.
43
- if (
44
- self.reduce_results
45
- and self.tp_size > 1
46
- and self.must_reduce_shared_expert_outputs()
47
- ):
48
- shared_out = tensor_model_parallel_all_reduce(shared_out)
47
+ if self._shared_experts is not None:
48
+ shared_out = self._shared_experts(hidden_states)
49
+
50
+ # Reduce shared expert outputs if necessary, since the MLP
51
+ # should have been created with reduce_results=False.
52
+ if (
53
+ self.reduce_results
54
+ and self.tp_size > 1
55
+ and self.must_reduce_shared_expert_outputs()
56
+ ):
57
+ shared_out = tensor_model_parallel_all_reduce(shared_out)
58
+ else:
59
+ shared_out = None
60
61
fused_out = super().forward(
62
hidden_states=hidden_states,
vllm/model_executor/layers/quantization/fp8.py
+2
@@ -741,6 +741,8 @@ class Fp8MoEMethod(FusedMoEMethodBase):
741
layer.w13_input_scale = None
742
layer.w2_input_scale = None
743
744
+ self.rocm_aiter_moe_enabled = False
745
+
746
def process_weights_after_loading(self, layer: Module) -> None:
747
# Lazy import to avoid importing triton too early.
748
from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (
vllm/model_executor/layers/shared_fused_moe/__init__.py
deleted
-5
@@ -1,5 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-from vllm.model_executor.layers.shared_fused_moe.shared_fused_moe import SharedFusedMoE
4
-
5
-__all__ = ["SharedFusedMoE"]
vllm/model_executor/models/aria.py
+15
-13
@@ -13,7 +13,7 @@ from vllm.config import VllmConfig
13
from vllm.config.multimodal import BaseDummyOptions
14
from vllm.distributed import get_tensor_model_parallel_rank
15
from vllm.model_executor.layers.activation import get_act_fn
16
-from vllm.model_executor.layers.fused_moe import FusedMoE
16
+from vllm.model_executor.layers.fused_moe import SharedFusedMoE
17
from vllm.model_executor.layers.linear import ColumnParallelLinear, RowParallelLinear
18
from vllm.model_executor.layers.logits_processor import LogitsProcessor
19
from vllm.model_executor.layers.quantization import QuantizationConfig
@@ -206,7 +206,7 @@ class AriaProjector(nn.Module):
206
return out
207
208
209
-class AriaFusedMoE(FusedMoE):
209
+class AriaFusedMoE(SharedFusedMoE):
210
def weight_loader(
211
self, param: nn.Parameter, loaded_weight: torch.Tensor, shard_id: str
212
) -> None:
@@ -260,7 +260,16 @@ class AriaTextMoELayer(nn.Module):
260
torch.empty((self.config.moe_num_experts, self.config.hidden_size))
261
)
262
263
+ self.shared_experts = LlamaMLP(
264
+ config.hidden_size,
265
+ config.intermediate_size * config.moe_num_shared_experts,
266
+ "silu",
267
+ quant_config=quant_config,
268
+ bias=config.mlp_bias,
269
+ )
270
+
271
self.experts = AriaFusedMoE(
272
+ shared_experts=self.shared_experts,
273
num_experts=config.moe_num_experts,
274
top_k=config.moe_topk,
275
hidden_size=config.hidden_size,
@@ -269,13 +278,6 @@ class AriaTextMoELayer(nn.Module):
278
reduce_results=True,
279
prefix=f"{prefix}.experts",
280
)
272
- self.shared_experts = LlamaMLP(
273
- config.hidden_size,
274
- config.intermediate_size * config.moe_num_shared_experts,
275
- "silu",
276
- quant_config=quant_config,
277
- bias=config.mlp_bias,
278
- )
281
282
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
283
"""
@@ -291,12 +293,12 @@ class AriaTextMoELayer(nn.Module):
293
294
router_output = torch.nn.functional.linear(hidden_states, self.router_weight)
295
294
- hidden_states_copy = hidden_states.clone()
295
- # NOTE: hidden_states will be modified inplace by `FusedMoE`
296
sparse_expert_output = self.experts(hidden_states, router_output)
297
- shared_expert_output = self.shared_experts(hidden_states_copy)
297
299
- return sparse_expert_output + shared_expert_output
298
+ if self.shared_experts is not None:
299
+ return sparse_expert_output[0] + sparse_expert_output[1]
300
+ else:
301
+ return sparse_expert_output
302
303
304
class AriaTextDecoderLayer(LlamaDecoderLayer):
vllm/model_executor/models/bailing_moe.py
+26
-21
@@ -43,7 +43,7 @@ from vllm.distributed import (
43
tensor_model_parallel_all_reduce,
44
)
45
from vllm.model_executor.layers.activation import SiluAndMul
46
-from vllm.model_executor.layers.fused_moe import FusedMoE
46
+from vllm.model_executor.layers.fused_moe import SharedFusedMoE
47
from vllm.model_executor.layers.layernorm import RMSNorm
48
from vllm.model_executor.layers.linear import (
49
MergedColumnParallelLinear,
@@ -276,22 +276,6 @@ class BailingMoE(nn.Module):
276
# default value for scoring_func
277
self.score_function = "softmax"
278
279
- self.experts = FusedMoE(
280
- num_experts=self.num_experts,
281
- top_k=self.top_k,
282
- hidden_size=self.hidden_size,
283
- intermediate_size=config.moe_intermediate_size,
284
- reduce_results=False,
285
- renormalize=self.norm_expert_prob,
286
- quant_config=quant_config,
287
- prefix=f"{prefix}.experts",
288
- scoring_func=self.score_function,
289
- e_score_correction_bias=self.gate.expert_bias,
290
- num_expert_group=self.n_group,
291
- topk_group=self.topk_group,
292
- use_grouped_topk=self.use_grouped_topk,
293
- )
294
-
279
if self.num_shared_experts > 0:
280
if hasattr(config, "moe_shared_expert_intermediate_size"):
281
intermediate_size = config.moe_shared_expert_intermediate_size
@@ -308,11 +292,27 @@ class BailingMoE(nn.Module):
292
else:
293
self.shared_experts = None
294
295
+ self.experts = SharedFusedMoE(
296
+ shared_experts=self.shared_experts,
297
+ num_experts=self.num_experts,
298
+ top_k=self.top_k,
299
+ hidden_size=self.hidden_size,
300
+ intermediate_size=config.moe_intermediate_size,
301
+ reduce_results=False,
302
+ renormalize=self.norm_expert_prob,
303
+ quant_config=quant_config,
304
+ prefix=f"{prefix}.experts",
305
+ scoring_func=self.score_function,
306
+ e_score_correction_bias=self.gate.expert_bias,
307
+ num_expert_group=self.n_group,
308
+ topk_group=self.topk_group,
309
+ use_grouped_topk=self.use_grouped_topk,
310
+ )
311
+
312
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
313
num_tokens, hidden_size = hidden_states.shape
314
hidden_states = hidden_states.view(-1, hidden_size)
314
- if self.shared_experts:
315
- shared_output = self.shared_experts(hidden_states)
315
+
316
# router_logits: (num_tokens, n_experts)
317
router_logits = self.gate(hidden_states.to(self.router_dtype))
318
router_logits = router_logits.to(hidden_states.dtype)
@@ -321,9 +321,14 @@ class BailingMoE(nn.Module):
321
hidden_states=hidden_states, router_logits=router_logits
322
)
323
324
+ if self.shared_experts is not None:
325
+ shared_output, final_hidden_states = final_hidden_states
326
+ else:
327
+ shared_output = None
328
+
329
final_hidden_states *= self.routed_scaling_factor
330
326
- if self.shared_experts:
331
+ if shared_output is not None:
332
final_hidden_states = final_hidden_states + shared_output
333
334
if self.tp_size > 1:
@@ -475,7 +480,7 @@ class BailingMoeModel(nn.Module):
480
return hidden_states
481
482
def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
478
- return FusedMoE.make_expert_params_mapping(
483
+ return SharedFusedMoE.make_expert_params_mapping(
484
ckpt_gate_proj_name="gate_proj",
485
ckpt_down_proj_name="down_proj",
486
ckpt_up_proj_name="up_proj",
vllm/model_executor/models/deepseek_v2.py
+24
-45
@@ -49,7 +49,7 @@ from vllm.forward_context import get_forward_context
49
from vllm.logger import init_logger
50
from vllm.model_executor.layers.activation import SiluAndMul
51
from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
52
-from vllm.model_executor.layers.fused_moe import FusedMoE
52
+from vllm.model_executor.layers.fused_moe import SharedFusedMoE
53
from vllm.model_executor.layers.layernorm import LayerNorm, RMSNorm
54
from vllm.model_executor.layers.linear import (
55
ColumnParallelLinear,
@@ -64,7 +64,6 @@ from vllm.model_executor.layers.quantization.utils.fp8_utils import (
64
per_token_group_quant_fp8,
65
)
66
from vllm.model_executor.layers.rotary_embedding import get_rope
67
-from vllm.model_executor.layers.shared_fused_moe import SharedFusedMoE
67
from vllm.model_executor.layers.vocab_parallel_embedding import (
68
ParallelLMHead,
69
VocabParallelEmbedding,
@@ -205,26 +204,6 @@ class DeepseekV2MoE(nn.Module):
204
)
205
206
if config.n_shared_experts is None:
208
- self.experts = FusedMoE(
209
- num_experts=config.n_routed_experts,
210
- top_k=config.num_experts_per_tok,
211
- hidden_size=config.hidden_size,
212
- intermediate_size=config.moe_intermediate_size,
213
- reduce_results=False,
214
- renormalize=config.norm_topk_prob,
215
- quant_config=quant_config,
216
- use_grouped_topk=True,
217
- num_expert_group=config.n_group,
218
- topk_group=config.topk_group,
219
- prefix=f"{prefix}.experts",
220
- scoring_func=config.scoring_func,
221
- # we do scaling outside, set factor to 1.0 to avoid double mul
222
- routed_scaling_factor=1.0,
223
- e_score_correction_bias=self.gate.e_score_correction_bias,
224
- enable_eplb=self.enable_eplb,
225
- num_redundant_experts=self.n_redundant_experts,
226
- is_sequence_parallel=self.is_sequence_parallel,
227
- )
207
self.shared_experts = None
208
else:
209
intermediate_size = config.moe_intermediate_size * config.n_shared_experts
@@ -239,27 +218,27 @@ class DeepseekV2MoE(nn.Module):
218
prefix=f"{prefix}.shared_experts",
219
)
220
242
- self.experts = SharedFusedMoE(
243
- shared_experts=self.shared_experts,
244
- num_experts=config.n_routed_experts,
245
- top_k=config.num_experts_per_tok,
246
- hidden_size=config.hidden_size,
247
- intermediate_size=config.moe_intermediate_size,
248
- reduce_results=False,
249
- renormalize=config.norm_topk_prob,
250
- quant_config=quant_config,
251
- use_grouped_topk=True,
252
- num_expert_group=config.n_group,
253
- topk_group=config.topk_group,
254
- prefix=f"{prefix}.experts",
255
- scoring_func=config.scoring_func,
256
- # we do scaling outside, set factor to 1.0 to avoid double mul
257
- routed_scaling_factor=1.0,
258
- e_score_correction_bias=self.gate.e_score_correction_bias,
259
- enable_eplb=self.enable_eplb,
260
- num_redundant_experts=self.n_redundant_experts,
261
- is_sequence_parallel=self.is_sequence_parallel,
262
- )
221
+ self.experts = SharedFusedMoE(
222
+ shared_experts=self.shared_experts,
223
+ num_experts=config.n_routed_experts,
224
+ top_k=config.num_experts_per_tok,
225
+ hidden_size=config.hidden_size,
226
+ intermediate_size=config.moe_intermediate_size,
227
+ reduce_results=False,
228
+ renormalize=config.norm_topk_prob,
229
+ quant_config=quant_config,
230
+ use_grouped_topk=True,
231
+ num_expert_group=config.n_group,
232
+ topk_group=config.topk_group,
233
+ prefix=f"{prefix}.experts",
234
+ scoring_func=config.scoring_func,
235
+ # we do scaling outside, set factor to 1.0 to avoid double mul
236
+ routed_scaling_factor=1.0,
237
+ e_score_correction_bias=self.gate.e_score_correction_bias,
238
+ enable_eplb=self.enable_eplb,
239
+ num_redundant_experts=self.n_redundant_experts,
240
+ is_sequence_parallel=self.is_sequence_parallel,
241
+ )
242
243
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
244
num_tokens, hidden_dim = hidden_states.shape
@@ -1306,7 +1285,7 @@ class DeepseekV2ForCausalLM(nn.Module, SupportsPP, MixtureOfExperts, SupportsLoR
1285
self.num_moe_layers = config.num_hidden_layers - config.first_k_dense_replace
1286
self.num_expert_groups = config.n_group
1287
1309
- self.moe_layers: list[FusedMoE] = []
1288
+ self.moe_layers: list[SharedFusedMoE] = []
1289
example_moe = None
1290
for layer in self.model.layers:
1291
if isinstance(layer, PPMissingLayer):
@@ -1394,7 +1373,7 @@ class DeepseekV2ForCausalLM(nn.Module, SupportsPP, MixtureOfExperts, SupportsLoR
1373
1374
# Params for weights, fp8 weight scales, fp8 activation scales
1375
# (param_name, weight_name, expert_id, shard_id)
1397
- expert_params_mapping = FusedMoE.make_expert_params_mapping(
1376
+ expert_params_mapping = SharedFusedMoE.make_expert_params_mapping(
1377
ckpt_gate_proj_name="gate_proj",
1378
ckpt_down_proj_name="down_proj",
1379
ckpt_up_proj_name="up_proj",
vllm/model_executor/models/dots1.py
+22
-18
@@ -42,7 +42,7 @@ from vllm.distributed import (
42
tensor_model_parallel_all_reduce,
43
)
44
from vllm.model_executor.layers.activation import SiluAndMul
45
-from vllm.model_executor.layers.fused_moe import FusedMoE
45
+from vllm.model_executor.layers.fused_moe import SharedFusedMoE
46
from vllm.model_executor.layers.layernorm import RMSNorm
47
from vllm.model_executor.layers.linear import (
48
MergedColumnParallelLinear,
@@ -145,7 +145,21 @@ class Dots1MoE(nn.Module):
145
else:
146
self.gate.e_score_correction_bias = None
147
148
- self.experts = FusedMoE(
148
+ if config.n_shared_experts is not None:
149
+ intermediate_size = config.moe_intermediate_size * config.n_shared_experts
150
+ self.shared_experts = Dots1MLP(
151
+ hidden_size=config.hidden_size,
152
+ intermediate_size=intermediate_size,
153
+ hidden_act=config.hidden_act,
154
+ quant_config=quant_config,
155
+ reduce_results=False,
156
+ prefix=f"{prefix}.shared_experts",
157
+ )
158
+ else:
159
+ self.shared_experts = None
160
+
161
+ self.experts = SharedFusedMoE(
162
+ shared_experts=self.shared_experts,
163
num_experts=config.n_routed_experts,
164
top_k=config.num_experts_per_tok,
165
hidden_size=config.hidden_size,
@@ -163,29 +177,19 @@ class Dots1MoE(nn.Module):
177
e_score_correction_bias=self.gate.e_score_correction_bias,
178
)
179
166
- if config.n_shared_experts is not None:
167
- intermediate_size = config.moe_intermediate_size * config.n_shared_experts
168
- self.shared_experts = Dots1MLP(
169
- hidden_size=config.hidden_size,
170
- intermediate_size=intermediate_size,
171
- hidden_act=config.hidden_act,
172
- quant_config=quant_config,
173
- reduce_results=False,
174
- prefix=f"{prefix}.shared_experts",
175
- )
176
-
180
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
181
num_tokens, hidden_dim = hidden_states.shape
182
hidden_states = hidden_states.view(-1, hidden_dim)
180
- if self.n_shared_experts is not None:
181
- shared_output = self.shared_experts(hidden_states)
183
+
184
router_logits, _ = self.gate(hidden_states)
185
final_hidden_states = (
186
self.experts(hidden_states=hidden_states, router_logits=router_logits)
187
* self.routed_scaling_factor
188
)
187
- if shared_output is not None:
188
- final_hidden_states = final_hidden_states + shared_output
189
+
190
+ if self.shared_experts is not None:
191
+ final_hidden_states = final_hidden_states[0] + final_hidden_states[1]
192
+
193
if self.tp_size > 1:
194
final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states)
195
return final_hidden_states.view(num_tokens, hidden_dim)
@@ -426,7 +430,7 @@ class Dots1Model(nn.Module):
430
return hidden_states
431
432
def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
429
- return FusedMoE.make_expert_params_mapping(
433
+ return SharedFusedMoE.make_expert_params_mapping(
434
ckpt_gate_proj_name="gate_proj",
435
ckpt_down_proj_name="down_proj",
436
ckpt_up_proj_name="up_proj",
vllm/model_executor/models/ernie45_moe.py
+20
-20
@@ -37,7 +37,7 @@ from vllm.config import CacheConfig, VllmConfig
37
from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
38
from vllm.logger import init_logger
39
from vllm.model_executor.layers.activation import SiluAndMul
40
-from vllm.model_executor.layers.fused_moe import FusedMoE
40
+from vllm.model_executor.layers.fused_moe import SharedFusedMoE
41
from vllm.model_executor.layers.layernorm import RMSNorm
42
from vllm.model_executor.layers.linear import (
43
MergedColumnParallelLinear,
@@ -145,18 +145,6 @@ class Ernie4_5_MoeMoE(nn.Module):
145
torch.empty(config.moe_num_experts, dtype=torch.float32)
146
)
147
148
- self.experts = FusedMoE(
149
- num_experts=config.moe_num_experts,
150
- top_k=config.moe_k,
151
- hidden_size=config.hidden_size,
152
- intermediate_size=config.moe_intermediate_size,
153
- reduce_results=False,
154
- renormalize=True,
155
- quant_config=quant_config,
156
- prefix=f"{prefix}.experts",
157
- e_score_correction_bias=self.gate.e_score_correction_bias,
158
- )
159
-
148
if self.has_shared_experts:
149
intermediate_size = (
150
config.moe_intermediate_size * config.moe_num_shared_experts
@@ -167,16 +155,28 @@ class Ernie4_5_MoeMoE(nn.Module):
155
hidden_act=config.hidden_act,
156
quant_config=quant_config,
157
prefix=f"{prefix}.shared_experts",
170
- reduce_results=self.experts.must_reduce_shared_expert_outputs(),
158
+ reduce_results=False,
159
)
160
+ else:
161
+ self.shared_experts = None
162
+
163
+ self.experts = SharedFusedMoE(
164
+ shared_experts=self.shared_experts,
165
+ num_experts=config.moe_num_experts,
166
+ top_k=config.moe_k,
167
+ hidden_size=config.hidden_size,
168
+ intermediate_size=config.moe_intermediate_size,
169
+ reduce_results=False,
170
+ renormalize=True,
171
+ quant_config=quant_config,
172
+ prefix=f"{prefix}.experts",
173
+ e_score_correction_bias=self.gate.e_score_correction_bias,
174
+ )
175
176
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
177
orig_shape = hidden_states.shape
178
hidden_dim = hidden_states.shape[-1]
179
hidden_states = hidden_states.view(-1, hidden_dim)
177
- shared_output = None
178
- if self.has_shared_experts:
179
- shared_output = self.shared_experts(hidden_states)
180
181
router_logits, _ = self.gate(hidden_states.to(dtype=torch.float32))
182
@@ -184,8 +184,8 @@ class Ernie4_5_MoeMoE(nn.Module):
184
hidden_states=hidden_states, router_logits=router_logits
185
)
186
187
- if self.has_shared_experts and shared_output is not None:
188
- final_hidden_states = final_hidden_states + shared_output
187
+ if self.has_shared_experts:
188
+ final_hidden_states = final_hidden_states[0] + final_hidden_states[1]
189
190
if self.tp_size > 1:
191
final_hidden_states = self.experts.maybe_all_reduce_tensor_model_parallel(
@@ -460,7 +460,7 @@ class Ernie4_5_MoeModel(nn.Module):
460
def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
461
# Params for weights, fp8 weight scales, fp8 activation scales
462
# (param_name, weight_name, expert_id, shard_id)
463
- return FusedMoE.make_expert_params_mapping(
463
+ return SharedFusedMoE.make_expert_params_mapping(
464
ckpt_gate_proj_name="gate_proj",
465
ckpt_down_proj_name="down_proj",
466
ckpt_up_proj_name="up_proj",
vllm/model_executor/models/ernie45_vl_moe.py
+34
-23
@@ -37,7 +37,7 @@ from vllm.attention import Attention
37
from vllm.config import CacheConfig, VllmConfig
38
from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
39
from vllm.logger import init_logger
40
-from vllm.model_executor.layers.fused_moe import FusedMoE
40
+from vllm.model_executor.layers.fused_moe import SharedFusedMoE
41
from vllm.model_executor.layers.layernorm import RMSNorm
42
from vllm.model_executor.layers.linear import (
43
QKVParallelLinear,
@@ -74,7 +74,15 @@ logger = init_logger(__name__)
74
75
76
class Ernie4_5_VLMoeMLP(Ernie4_5_MoeMLP):
77
- pass
77
+ def __init__(self, shared_experts: Optional[torch.nn.Module] = None, **kwargs):
78
+ super().__init__(**kwargs)
79
+ self.shared_experts = shared_experts
80
+
81
+ def forward(self, x):
82
+ if self.shared_experts is not None:
83
+ return self.shared_experts(x) + super().forward(x)
84
+ else:
85
+ return super().forward(x)
86
87
88
class Ernie4_5_VLMoeAttention(nn.Module):
@@ -223,6 +231,21 @@ class Ernie4_5_VLMoeMoE(nn.Module):
231
232
assert text_moe_layer_start_index <= text_moe_layer_end_index
233
234
+ if self.has_shared_experts:
235
+ intermediate_size = (
236
+ config.moe_intermediate_size[0] * config.moe_num_shared_experts
237
+ )
238
+ self.shared_experts = Ernie4_5_VLMoeMLP(
239
+ hidden_size=config.hidden_size,
240
+ intermediate_size=intermediate_size,
241
+ hidden_act=config.hidden_act,
242
+ quant_config=quant_config,
243
+ prefix=f"{prefix}.shared_experts",
244
+ reduce_results=False,
245
+ )
246
+ else:
247
+ self.shared_experts = None
248
+
249
if (
250
layer_idx >= text_moe_layer_start_index
251
and layer_idx <= text_moe_layer_end_index
@@ -236,7 +259,8 @@ class Ernie4_5_VLMoeMoE(nn.Module):
259
prefix=f"{prefix}.text_experts_gate",
260
)
261
239
- self.text_experts = FusedMoE(
262
+ self.text_experts = SharedFusedMoE(
263
+ shared_experts=self.shared_experts,
264
num_experts=config.moe_num_experts[0],
265
top_k=config.moe_k,
266
hidden_size=config.hidden_size,
@@ -249,6 +273,7 @@ class Ernie4_5_VLMoeMoE(nn.Module):
273
)
274
else:
275
self.text_experts = Ernie4_5_VLMoeMLP(
276
+ shared_experts=self.shared_experts,
277
hidden_size=config.hidden_size,
278
intermediate_size=config.intermediate_size,
279
hidden_act=config.hidden_act,
@@ -271,7 +296,8 @@ class Ernie4_5_VLMoeMoE(nn.Module):
296
prefix=f"{prefix}.vision_experts_gate",
297
)
298
274
- self.vision_experts = FusedMoE(
299
+ self.vision_experts = SharedFusedMoE(
300
+ shared_experts=self.shared_experts,
301
num_experts=config.moe_num_experts[1],
302
top_k=config.moe_k,
303
hidden_size=config.hidden_size,
@@ -284,6 +310,7 @@ class Ernie4_5_VLMoeMoE(nn.Module):
310
)
311
else:
312
self.vision_experts = Ernie4_5_VLMoeMLP(
313
+ shared_experts=self.shared_experts,
314
hidden_size=config.hidden_size,
315
intermediate_size=config.intermediate_size,
316
hidden_act=config.hidden_act,
@@ -292,19 +319,6 @@ class Ernie4_5_VLMoeMoE(nn.Module):
319
prefix=f"{prefix}.mlp",
320
)
321
295
- if self.has_shared_experts:
296
- intermediate_size = (
297
- config.moe_intermediate_size[0] * config.moe_num_shared_experts
298
- )
299
- self.shared_experts = Ernie4_5_VLMoeMLP(
300
- hidden_size=config.hidden_size,
301
- intermediate_size=intermediate_size,
302
- hidden_act=config.hidden_act,
303
- quant_config=quant_config,
304
- prefix=f"{prefix}.shared_experts",
305
- reduce_results=self.text_experts.must_reduce_shared_expert_outputs(),
306
- )
307
-
322
def forward(
323
self,
324
hidden_states: torch.Tensor,
@@ -315,9 +329,6 @@ class Ernie4_5_VLMoeMoE(nn.Module):
329
hidden_dim = hidden_states.shape[-1]
330
hidden_states = hidden_states.view(-1, hidden_dim)
331
318
- if self.has_shared_experts:
319
- shared_output = self.shared_experts(hidden_states)
320
-
332
if visual_token_mask is not None and visual_token_mask.all():
333
# only vision modal input
334
router_logits, _ = self.vision_experts_gate(
@@ -362,8 +373,8 @@ class Ernie4_5_VLMoeMoE(nn.Module):
373
hidden_states=hidden_states, router_logits=text_router_logits
374
)
375
365
- if self.has_shared_experts and shared_output is not None:
366
- final_hidden_states = final_hidden_states + shared_output
376
+ if self.has_shared_experts:
377
+ final_hidden_states = final_hidden_states[0] + final_hidden_states[1]
378
379
if self.tp_size > 1:
380
final_hidden_states = (
@@ -649,7 +660,7 @@ class Ernie4_5_VLMoeForCausalLM(nn.Module, SupportsPP):
660
661
# Params for weights, fp8 weight scales, fp8 activation scales
662
# (param_name, weight_name, expert_id, shard_id)
652
- expert_params_mapping = FusedMoE.make_expert_params_mapping(
663
+ expert_params_mapping = SharedFusedMoE.make_expert_params_mapping(
664
ckpt_gate_proj_name="gate_proj",
665
ckpt_down_proj_name="down_proj",
666
ckpt_up_proj_name="up_proj",
vllm/model_executor/models/glm4_moe.py
+25
-43
@@ -42,7 +42,7 @@ from vllm.distributed import (
42
)
43
from vllm.logger import init_logger
44
from vllm.model_executor.layers.activation import SiluAndMul
45
-from vllm.model_executor.layers.fused_moe import FusedMoE
45
+from vllm.model_executor.layers.fused_moe import SharedFusedMoE
46
from vllm.model_executor.layers.layernorm import RMSNorm
47
from vllm.model_executor.layers.linear import (
48
MergedColumnParallelLinear,
@@ -52,7 +52,6 @@ from vllm.model_executor.layers.linear import (
52
from vllm.model_executor.layers.logits_processor import LogitsProcessor
53
from vllm.model_executor.layers.quantization import QuantizationConfig
54
from vllm.model_executor.layers.rotary_embedding import get_rope
55
-from vllm.model_executor.layers.shared_fused_moe import SharedFusedMoE
55
from vllm.model_executor.layers.vocab_parallel_embedding import (
56
ParallelLMHead,
57
VocabParallelEmbedding,
@@ -176,46 +175,29 @@ class Glm4MoE(nn.Module):
175
reduce_results=False,
176
prefix=f"{prefix}.shared_experts",
177
)
179
- self.experts = SharedFusedMoE(
180
- shared_experts=self.shared_experts,
181
- num_experts=config.n_routed_experts,
182
- top_k=config.num_experts_per_tok,
183
- hidden_size=config.hidden_size,
184
- intermediate_size=config.moe_intermediate_size,
185
- reduce_results=False,
186
- renormalize=config.norm_topk_prob,
187
- quant_config=quant_config,
188
- use_grouped_topk=True,
189
- num_expert_group=config.n_group,
190
- topk_group=config.topk_group,
191
- prefix=f"{prefix}.experts",
192
- scoring_func="sigmoid",
193
- # we do scaling outside, set factor to 1.0 to avoid double mul
194
- routed_scaling_factor=1.0,
195
- e_score_correction_bias=self.gate.e_score_correction_bias,
196
- enable_eplb=self.enable_eplb,
197
- num_redundant_experts=self.n_redundant_experts,
198
- )
178
else:
200
- self.experts = FusedMoE(
201
- num_experts=config.n_routed_experts,
202
- top_k=config.num_experts_per_tok,
203
- hidden_size=config.hidden_size,
204
- intermediate_size=config.moe_intermediate_size,
205
- reduce_results=False,
206
- renormalize=config.norm_topk_prob,
207
- quant_config=quant_config,
208
- use_grouped_topk=True,
209
- num_expert_group=config.n_group,
210
- topk_group=config.topk_group,
211
- prefix=f"{prefix}.experts",
212
- scoring_func="sigmoid",
213
- # we do scaling outside, set factor to 1.0 to avoid double mul
214
- routed_scaling_factor=1.0,
215
- e_score_correction_bias=self.gate.e_score_correction_bias,
216
- enable_eplb=self.enable_eplb,
217
- num_redundant_experts=self.n_redundant_experts,
218
- )
179
+ self.shared_experts = None
180
+
181
+ self.experts = SharedFusedMoE(
182
+ shared_experts=self.shared_experts,
183
+ num_experts=config.n_routed_experts,
184
+ top_k=config.num_experts_per_tok,
185
+ hidden_size=config.hidden_size,
186
+ intermediate_size=config.moe_intermediate_size,
187
+ reduce_results=False,
188
+ renormalize=config.norm_topk_prob,
189
+ quant_config=quant_config,
190
+ use_grouped_topk=True,
191
+ num_expert_group=config.n_group,
192
+ topk_group=config.topk_group,
193
+ prefix=f"{prefix}.experts",
194
+ scoring_func="sigmoid",
195
+ # we do scaling outside, set factor to 1.0 to avoid double mul
196
+ routed_scaling_factor=1.0,
197
+ e_score_correction_bias=self.gate.e_score_correction_bias,
198
+ enable_eplb=self.enable_eplb,
199
+ num_redundant_experts=self.n_redundant_experts,
200
+ )
201
202
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
203
num_tokens, hidden_dim = hidden_states.shape
@@ -522,7 +504,7 @@ class Glm4MoeModel(nn.Module):
504
def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
505
# Params for weights, fp8 weight scales, fp8 activation scales
506
# (param_name, weight_name, expert_id, shard_id)
525
- return FusedMoE.make_expert_params_mapping(
507
+ return SharedFusedMoE.make_expert_params_mapping(
508
ckpt_gate_proj_name="gate_proj",
509
ckpt_down_proj_name="down_proj",
510
ckpt_up_proj_name="up_proj",
@@ -677,7 +659,7 @@ class Glm4MoeForCausalLM(nn.Module, SupportsPP, SupportsLoRA):
659
self.num_moe_layers = config.num_hidden_layers - config.first_k_dense_replace
660
self.num_expert_groups = config.n_group
661
680
- self.moe_layers: list[FusedMoE] = []
662
+ self.moe_layers: list[SharedFusedMoE] = []
663
example_moe = None
664
for layer in self.model.layers:
665
if isinstance(layer, PPMissingLayer):
vllm/model_executor/models/hunyuan_v1.py
+20
-21
@@ -44,7 +44,7 @@ from vllm.distributed import (
44
tensor_model_parallel_all_reduce,
45
)
46
from vllm.model_executor.layers.activation import SiluAndMul
47
-from vllm.model_executor.layers.fused_moe import FusedMoE
47
+from vllm.model_executor.layers.fused_moe import SharedFusedMoE
48
from vllm.model_executor.layers.layernorm import RMSNorm
49
from vllm.model_executor.layers.linear import (
50
ColumnParallelLinear,
@@ -415,19 +415,6 @@ class HunYuanSparseMoeBlock(nn.Module):
415
self.physical_expert_start + self.n_local_physical_experts
416
)
417
418
- self.experts = FusedMoE(
419
- num_experts=self.n_routed_experts,
420
- top_k=top_k,
421
- hidden_size=config.hidden_size,
422
- intermediate_size=intermediate_size,
423
- reduce_results=False,
424
- renormalize=top_k > 1,
425
- quant_config=quant_config,
426
- prefix=f"{prefix}.experts",
427
- enable_eplb=self.enable_eplb,
428
- num_redundant_experts=self.n_redundant_experts,
429
- )
430
-
418
self.gate = ReplicatedLinear(
419
config.hidden_size,
420
config.num_experts,
@@ -455,22 +442,34 @@ class HunYuanSparseMoeBlock(nn.Module):
442
else:
443
self.shared_mlp = None
444
445
+ self.experts = SharedFusedMoE(
446
+ shared_experts=self.shared_mlp,
447
+ num_experts=self.n_routed_experts,
448
+ top_k=top_k,
449
+ hidden_size=config.hidden_size,
450
+ intermediate_size=intermediate_size,
451
+ reduce_results=False,
452
+ renormalize=top_k > 1,
453
+ quant_config=quant_config,
454
+ prefix=f"{prefix}.experts",
455
+ enable_eplb=self.enable_eplb,
456
+ num_redundant_experts=self.n_redundant_experts,
457
+ )
458
+
459
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
460
# NOTE: hidden_states can have either 1D or 2D shape.
461
orig_shape = hidden_states.shape
462
hidden_dim = hidden_states.shape[-1]
463
hidden_states = hidden_states.view(-1, hidden_dim)
463
- shared_output = None
464
- if self.shared_mlp is not None:
465
- shared_output = self.shared_mlp(hidden_states)
464
465
# router_logits: (num_tokens, n_experts)
466
router_logits, _ = self.gate(hidden_states)
467
final_hidden_states = self.experts(
468
hidden_states=hidden_states, router_logits=router_logits
469
)
472
- if shared_output is not None:
473
- final_hidden_states = final_hidden_states + shared_output
470
+ if self.shared_mlp is not None:
471
+ final_hidden_states = final_hidden_states[0] + final_hidden_states[1]
472
+
473
if self.tp_size > 1:
474
final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states)
475
@@ -724,7 +723,7 @@ class HunYuanModel(nn.Module):
723
if _is_moe(self.config):
724
# Params for weights, fp8 weight scales, fp8 activation scales
725
# (param_name, weight_name, expert_id, shard_id)
727
- return FusedMoE.make_expert_params_mapping(
726
+ return SharedFusedMoE.make_expert_params_mapping(
727
ckpt_gate_proj_name="gate_proj",
728
ckpt_down_proj_name="down_proj",
729
ckpt_up_proj_name="up_proj",
@@ -1008,7 +1007,7 @@ class HunYuanMoEV1Base(HunyuanV1ModelBase, MixtureOfExperts):
1007
# Set MoE hyperparameters
1008
self.expert_weights = []
1009
self.num_expert_groups = 1
1011
- self.moe_layers: list[FusedMoE] = []
1010
+ self.moe_layers: list[SharedFusedMoE] = []
1011
example_layer = None
1012
for layer in self.model.layers:
1013
if isinstance(layer, PPMissingLayer):
vllm/model_executor/models/llama4.py
+4
-5
@@ -33,7 +33,7 @@ from vllm.distributed import (
33
get_tensor_model_parallel_world_size,
34
tensor_model_parallel_all_gather,
35
)
36
-from vllm.model_executor.layers.fused_moe import FusedMoE
36
+from vllm.model_executor.layers.fused_moe import SharedFusedMoE
37
from vllm.model_executor.layers.layernorm import RMSNorm
38
from vllm.model_executor.layers.linear import (
39
QKVParallelLinear,
@@ -42,7 +42,6 @@ from vllm.model_executor.layers.linear import (
42
)
43
from vllm.model_executor.layers.quantization import QuantizationConfig
44
from vllm.model_executor.layers.rotary_embedding import get_rope
45
-from vllm.model_executor.layers.shared_fused_moe import SharedFusedMoE
45
from vllm.model_executor.model_loader.weight_utils import (
46
default_weight_loader,
47
maybe_remap_kv_scale_name,
@@ -399,7 +398,7 @@ class Llama4Model(LlamaModel):
398
params_dict: The dictionary of module parameters.
399
loaded_params: The set of already loaded parameters.
400
expert_params_mapping: The mapping of expert parameters. Must be
402
- generated by FusedMoE.make_expert_params_mapping().
401
+ generated by SharedFusedMoE.make_expert_params_mapping().
402
fused: Whether the expert weights are fused into a single weight
403
tensor or are separate weight tensors for each expert.
404
When fused is True, loaded_weight should have shape of:
@@ -522,7 +521,7 @@ class Llama4Model(LlamaModel):
521
fused_experts_params = False
522
# Expert parameter mapping for the case where the expert weights are
523
# not fused into a single weight tensor.
525
- expert_params_mapping = FusedMoE.make_expert_params_mapping(
524
+ expert_params_mapping = SharedFusedMoE.make_expert_params_mapping(
525
ckpt_gate_proj_name="gate_proj",
526
ckpt_down_proj_name="down_proj",
527
ckpt_up_proj_name="up_proj",
@@ -530,7 +529,7 @@ class Llama4Model(LlamaModel):
529
)
530
# Expert parameter mapping for the case where the expert weights are
531
# fused into a single weight tensor.
533
- expert_params_mapping_fused = FusedMoE.make_expert_params_mapping(
532
+ expert_params_mapping_fused = SharedFusedMoE.make_expert_params_mapping(
533
ckpt_gate_proj_name="gate_up_proj",
534
ckpt_down_proj_name="down_proj",
535
ckpt_up_proj_name="gate_up_proj",
vllm/model_executor/models/qwen2_moe.py
+30
-27
@@ -40,7 +40,7 @@ from vllm.config import CacheConfig, VllmConfig
40
from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
41
from vllm.logger import init_logger
42
from vllm.model_executor.layers.activation import SiluAndMul
43
-from vllm.model_executor.layers.fused_moe import FusedMoE
43
+from vllm.model_executor.layers.fused_moe import SharedFusedMoE
44
from vllm.model_executor.layers.layernorm import RMSNorm
45
from vllm.model_executor.layers.linear import (
46
MergedColumnParallelLinear,
@@ -79,6 +79,7 @@ class Qwen2MoeMLP(nn.Module):
79
hidden_act: str,
80
quant_config: Optional[QuantizationConfig] = None,
81
reduce_results: bool = True,
82
+ expert_gate: Optional[torch.nn.Linear] = None,
83
prefix: str = "",
84
) -> None:
85
super().__init__()
@@ -102,12 +103,17 @@ class Qwen2MoeMLP(nn.Module):
103
f"Unsupported activation: {hidden_act}. Only silu is supported for now."
104
)
105
self.act_fn = SiluAndMul()
106
+ self.expert_gate = expert_gate
107
108
def forward(self, x):
109
gate_up, _ = self.gate_up_proj(x)
108
- x = self.act_fn(gate_up)
109
- x, _ = self.down_proj(x)
110
- return x
110
+ out = self.act_fn(gate_up)
111
+ out, _ = self.down_proj(out)
112
+
113
+ if self.expert_gate is not None:
114
+ out = F.sigmoid(self.expert_gate(x)) * out
115
+
116
+ return out
117
118
119
class Qwen2MoeSparseMoeBlock(nn.Module):
@@ -126,17 +132,6 @@ class Qwen2MoeSparseMoeBlock(nn.Module):
132
f"the number of experts {config.num_experts}."
133
)
134
129
- self.experts = FusedMoE(
130
- num_experts=config.num_experts,
131
- top_k=config.num_experts_per_tok,
132
- hidden_size=config.hidden_size,
133
- intermediate_size=config.moe_intermediate_size,
134
- reduce_results=False,
135
- renormalize=config.norm_topk_prob,
136
- quant_config=quant_config,
137
- prefix=f"{prefix}.experts",
138
- )
139
-
135
self.gate = ReplicatedLinear(
136
config.hidden_size,
137
config.num_experts,
@@ -144,39 +139,47 @@ class Qwen2MoeSparseMoeBlock(nn.Module):
139
quant_config=None,
140
prefix=f"{prefix}.gate",
141
)
142
+
143
+ self.shared_expert_gate = torch.nn.Linear(config.hidden_size, 1, bias=False)
144
+
145
if config.shared_expert_intermediate_size > 0:
146
self.shared_expert = Qwen2MoeMLP(
147
hidden_size=config.hidden_size,
148
intermediate_size=config.shared_expert_intermediate_size,
149
hidden_act=config.hidden_act,
150
quant_config=quant_config,
153
- reduce_results=self.experts.must_reduce_shared_expert_outputs(),
151
+ reduce_results=False,
152
+ expert_gate=self.shared_expert_gate,
153
prefix=f"{prefix}.shared_expert",
154
)
155
else:
156
self.shared_expert = None
158
- self.shared_expert_gate = torch.nn.Linear(config.hidden_size, 1, bias=False)
157
+
158
+ self.experts = SharedFusedMoE(
159
+ shared_experts=self.shared_expert,
160
+ num_experts=config.num_experts,
161
+ top_k=config.num_experts_per_tok,
162
+ hidden_size=config.hidden_size,
163
+ intermediate_size=config.moe_intermediate_size,
164
+ reduce_results=False,
165
+ renormalize=config.norm_topk_prob,
166
+ quant_config=quant_config,
167
+ prefix=f"{prefix}.experts",
168
+ )
169
170
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
171
# NOTE: hidden_states can have either 1D or 2D shape.
172
orig_shape = hidden_states.shape
173
hidden_dim = hidden_states.shape[-1]
174
hidden_states = hidden_states.view(-1, hidden_dim)
165
- shared_output = None
166
- if self.shared_expert is not None:
167
- shared_output = self.shared_expert(hidden_states)
168
- if self.shared_expert_gate is not None:
169
- shared_output = (
170
- F.sigmoid(self.shared_expert_gate(hidden_states)) * shared_output
171
- )
175
176
# router_logits: (num_tokens, n_experts)
177
router_logits, _ = self.gate(hidden_states)
178
final_hidden_states = self.experts(
179
hidden_states=hidden_states, router_logits=router_logits
180
)
178
- if shared_output is not None:
179
- final_hidden_states = final_hidden_states + shared_output
181
+ if self.shared_expert is not None:
182
+ final_hidden_states = final_hidden_states[0] + final_hidden_states[1]
183
if self.tp_size > 1:
184
final_hidden_states = self.experts.maybe_all_reduce_tensor_model_parallel( # noqa E501
185
final_hidden_states
@@ -418,7 +421,7 @@ class Qwen2MoeModel(nn.Module):
421
def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
422
# Params for weights, fp8 weight scales, fp8 activation scales
423
# (param_name, weight_name, expert_id, shard_id)
421
- return FusedMoE.make_expert_params_mapping(
424
+ return SharedFusedMoE.make_expert_params_mapping(
425
ckpt_gate_proj_name="gate_proj",
426
ckpt_down_proj_name="down_proj",
427
ckpt_up_proj_name="up_proj",
vllm/model_executor/models/qwen3_next.py
+24
-30
@@ -7,7 +7,6 @@ from itertools import islice
7
from typing import Optional
8
9
import torch
10
-import torch.nn.functional as F
10
from einops import rearrange
11
from torch import nn
12
from transformers.activations import ACT2FN
@@ -36,7 +35,7 @@ from vllm.model_executor.layers.fla.ops import (
35
chunk_gated_delta_rule,
36
fused_recurrent_gated_delta_rule,
37
)
39
-from vllm.model_executor.layers.fused_moe import FusedMoE
38
+from vllm.model_executor.layers.fused_moe import SharedFusedMoE
39
from vllm.model_executor.layers.layernorm import GemmaRMSNorm as Qwen3NextRMSNorm
40
from vllm.model_executor.layers.linear import (
41
ColumnParallelLinear,
@@ -136,20 +135,6 @@ class Qwen3NextSparseMoeBlock(nn.Module):
135
self.physical_expert_start + self.n_local_physical_experts
136
)
137
139
- self.experts = FusedMoE(
140
- num_experts=self.n_routed_experts,
141
- top_k=config.num_experts_per_tok,
142
- hidden_size=config.hidden_size,
143
- intermediate_size=config.moe_intermediate_size,
144
- reduce_results=False,
145
- renormalize=config.norm_topk_prob,
146
- quant_config=quant_config,
147
- prefix=f"{prefix}.experts",
148
- enable_eplb=self.enable_eplb,
149
- num_redundant_experts=self.n_redundant_experts,
150
- is_sequence_parallel=self.is_sequence_parallel,
151
- )
152
-
138
self.gate = ReplicatedLinear(
139
config.hidden_size,
140
config.num_experts,
@@ -158,18 +143,35 @@ class Qwen3NextSparseMoeBlock(nn.Module):
143
prefix=f"{prefix}.gate",
144
)
145
146
+ self.shared_expert_gate = torch.nn.Linear(config.hidden_size, 1, bias=False)
147
+
148
if config.shared_expert_intermediate_size > 0:
149
self.shared_expert = Qwen3NextMLP(
150
hidden_size=config.hidden_size,
151
intermediate_size=config.shared_expert_intermediate_size,
152
hidden_act=config.hidden_act,
153
quant_config=quant_config,
167
- reduce_results=self.experts.must_reduce_shared_expert_outputs(),
154
+ reduce_results=False,
155
+ expert_gate=self.shared_expert_gate,
156
prefix=f"{prefix}.shared_expert",
157
)
158
else:
159
self.shared_expert = None
172
- self.shared_expert_gate = torch.nn.Linear(config.hidden_size, 1, bias=False)
160
+
161
+ self.experts = SharedFusedMoE(
162
+ shared_experts=self.shared_expert,
163
+ num_experts=self.n_routed_experts,
164
+ top_k=config.num_experts_per_tok,
165
+ hidden_size=config.hidden_size,
166
+ intermediate_size=config.moe_intermediate_size,
167
+ reduce_results=False,
168
+ renormalize=config.norm_topk_prob,
169
+ quant_config=quant_config,
170
+ prefix=f"{prefix}.experts",
171
+ enable_eplb=self.enable_eplb,
172
+ num_redundant_experts=self.n_redundant_experts,
173
+ is_sequence_parallel=self.is_sequence_parallel,
174
+ )
175
176
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
177
# NOTE: hidden_states can have either 1D or 2D shape.
@@ -180,22 +182,14 @@ class Qwen3NextSparseMoeBlock(nn.Module):
182
if self.is_sequence_parallel:
183
hidden_states = sequence_parallel_chunk(hidden_states)
184
183
- shared_output = None
184
- if self.shared_expert is not None:
185
- shared_output = self.shared_expert(hidden_states)
186
- if self.shared_expert_gate is not None:
187
- shared_output = (
188
- F.sigmoid(self.shared_expert_gate(hidden_states)) * shared_output
189
- )
190
-
185
# router_logits: (num_tokens, n_experts)
186
router_logits, _ = self.gate(hidden_states)
187
final_hidden_states = self.experts(
188
hidden_states=hidden_states, router_logits=router_logits
189
)
190
197
- if shared_output is not None:
198
- final_hidden_states = final_hidden_states + shared_output
191
+ if self.shared_expert is not None:
192
+ final_hidden_states = final_hidden_states[0] + final_hidden_states[1]
193
194
if self.is_sequence_parallel:
195
final_hidden_states = tensor_model_parallel_all_gather(
@@ -1008,7 +1002,7 @@ class Qwen3NextModel(nn.Module):
1002
def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
1003
# Params for weights, fp8 weight scales, fp8 activation scales
1004
# (param_name, weight_name, expert_id, shard_id)
1011
- return FusedMoE.make_expert_params_mapping(
1005
+ return SharedFusedMoE.make_expert_params_mapping(
1006
ckpt_gate_proj_name="gate_proj",
1007
ckpt_down_proj_name="down_proj",
1008
ckpt_up_proj_name="up_proj",
@@ -1150,7 +1144,7 @@ class Qwen3NextForCausalLM(
1144
# Set MoE hyperparameters
1145
self.expert_weights = []
1146
1153
- self.moe_layers: list[FusedMoE] = []
1147
+ self.moe_layers: list[SharedFusedMoE] = []
1148
example_layer = None
1149
for layer in self.model.layers:
1150
if isinstance(layer, PPMissingLayer):