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 9, 17:06 UTC
[Models] Add `SharedFusedMoE` support to Qwen3MoE (#32082)
Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn>
Isotr0py committed
Jan 24, 2026 at 15:36 UTC
8edaf3857027c75382672ade255f7cb531f96844
1 file changed
+56
-16
vllm/model_executor/models/qwen3_moe.py
+56
-16
@@ -29,6 +29,7 @@ from itertools import islice
29
from typing import Any
30
31
import torch
32
+import torch.nn.functional as F
33
from torch import nn
34
35
from vllm.attention.layer import Attention
@@ -42,7 +43,7 @@ from vllm.distributed import (
43
)
44
from vllm.logger import init_logger
45
from vllm.model_executor.layers.activation import SiluAndMul
45
-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,
@@ -86,6 +87,7 @@ class Qwen3MoeMLP(nn.Module):
87
hidden_act: str,
88
quant_config: QuantizationConfig | None = None,
89
reduce_results: bool = True,
90
+ expert_gate: torch.nn.Linear | None = None,
91
prefix: str = "",
92
) -> None:
93
super().__init__()
@@ -109,12 +111,17 @@ class Qwen3MoeMLP(nn.Module):
111
f"Unsupported activation: {hidden_act}. Only silu is supported for now."
112
)
113
self.act_fn = SiluAndMul()
114
+ self.expert_gate = expert_gate
115
116
def forward(self, x):
117
gate_up, _ = self.gate_up_proj(x)
115
- x = self.act_fn(gate_up)
116
- x, _ = self.down_proj(x)
117
- return x
118
+ out = self.act_fn(gate_up)
119
+ out, _ = self.down_proj(out)
120
+
121
+ if self.expert_gate is not None:
122
+ out = F.sigmoid(self.expert_gate(x)[0]) * out
123
+
124
+ return out
125
126
127
class Qwen3MoeSparseMoeBlock(nn.Module):
@@ -159,12 +166,46 @@ class Qwen3MoeSparseMoeBlock(nn.Module):
166
self.physical_expert_start + self.n_local_physical_experts
167
)
168
162
- self.experts = FusedMoE(
169
+ self.gate = ReplicatedLinear(
170
+ config.hidden_size,
171
+ config.num_experts,
172
+ bias=False,
173
+ quant_config=quant_config,
174
+ prefix=f"{prefix}.gate",
175
+ )
176
+
177
+ shared_expert_intermediate_size = getattr(
178
+ config, "shared_expert_intermediate_size", 0
179
+ )
180
+ if shared_expert_intermediate_size > 0:
181
+ self.shared_expert_gate = ReplicatedLinear(
182
+ config.hidden_size,
183
+ 1,
184
+ bias=False,
185
+ quant_config=None,
186
+ prefix=f"{prefix}.shared_expert_gate",
187
+ )
188
+ self.shared_expert = Qwen3MoeMLP(
189
+ hidden_size=config.hidden_size,
190
+ intermediate_size=shared_expert_intermediate_size,
191
+ hidden_act=config.hidden_act,
192
+ quant_config=quant_config,
193
+ reduce_results=False,
194
+ expert_gate=self.shared_expert_gate,
195
+ prefix=f"{prefix}.shared_expert",
196
+ )
197
+ else:
198
+ self.shared_expert_gate = None
199
+ self.shared_expert = None
200
+
201
+ self.experts = SharedFusedMoE(
202
+ shared_experts=self.shared_expert,
203
+ gate=self.gate,
204
num_experts=self.n_routed_experts,
205
top_k=config.num_experts_per_tok,
206
hidden_size=config.hidden_size,
207
intermediate_size=config.moe_intermediate_size,
167
- reduce_results=True,
208
+ reduce_results=False,
209
renormalize=config.norm_topk_prob,
210
quant_config=quant_config,
211
prefix=f"{prefix}.experts",
@@ -173,14 +214,6 @@ class Qwen3MoeSparseMoeBlock(nn.Module):
214
is_sequence_parallel=self.is_sequence_parallel,
215
)
216
176
- self.gate = ReplicatedLinear(
177
- config.hidden_size,
178
- config.num_experts,
179
- bias=False,
180
- quant_config=quant_config,
181
- prefix=f"{prefix}.gate",
182
- )
183
-
217
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
218
assert hidden_states.dim() <= 2, (
219
"Qwen3MoeSparseMoeBlock only supports 1D or 2D inputs"
@@ -194,15 +227,22 @@ class Qwen3MoeSparseMoeBlock(nn.Module):
227
228
# router_logits: (num_tokens, n_experts)
229
router_logits, _ = self.gate(hidden_states)
197
- final_hidden_states = self.experts(
230
+ shared_out, fused_out = self.experts(
231
hidden_states=hidden_states, router_logits=router_logits
232
)
233
+ final_hidden_states = (
234
+ shared_out + fused_out if shared_out is not None else fused_out
235
+ )
236
237
if self.is_sequence_parallel:
238
final_hidden_states = tensor_model_parallel_all_gather(
239
final_hidden_states, 0
240
)
241
final_hidden_states = final_hidden_states[:num_tokens]
242
+ elif self.tp_size > 1:
243
+ final_hidden_states = self.experts.maybe_all_reduce_tensor_model_parallel( # noqa E501
244
+ final_hidden_states
245
+ )
246
247
# return to 1d if input is 1d
248
return final_hidden_states.squeeze(0) if is_input_1d else final_hidden_states
@@ -467,7 +507,7 @@ class Qwen3MoeModel(nn.Module):
507
def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
508
# Params for weights, fp8 weight scales, fp8 activation scales
509
# (param_name, weight_name, expert_id, shard_id)
470
- return FusedMoE.make_expert_params_mapping(
510
+ return SharedFusedMoE.make_expert_params_mapping(
511
self,
512
ckpt_gate_proj_name="gate_proj",
513
ckpt_down_proj_name="down_proj",