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 4, 02:30 UTC
[ROCm][Perf] Add Fused Shared Expert (FSE) support for Qwen3-Next (#39280)
Signed-off-by: nholmber <nholmber@users.noreply.github.com> Signed-off-by: Tres Popp <tres.popp@amd.com> Signed-off-by: Doug Lehr <douglehr@amd.com> Co-authored-by: nholmber <nholmber@users.noreply.github.com> Co-authored-by: Tres <tpopp@users.noreply.github.com> Co-authored-by: Tres Popp <tres.popp@amd.com> Co-authored-by: Doug Lehr <douglehr@amd.com> Co-authored-by: Cursor <cursoragent@cursor.com> Co-authored-by: Douglas Lehr <91553416+dllehr-amd@users.noreply.github.com>
Nico Holmberg committed
May 8, 2026 at 12:38 UTC
2c6b59b80771ac3bfa1c789abe3cb27c379bf3a1
8 files changed
+359
-10
vllm/_aiter_ops.py
+69
-2
@@ -261,11 +261,19 @@ def _rocm_aiter_topk_softmax_impl(
261
token_expert_indices: torch.Tensor,
262
gating_output: torch.Tensor,
263
renormalize: bool,
264
+ num_shared_experts: int = 0,
265
+ shared_expert_scoring_func: str = "",
266
) -> None:
267
from aiter import topk_softmax
268
269
topk_softmax(
268
- topk_weights, topk_indices, token_expert_indices, gating_output, renormalize
270
+ topk_weights,
271
+ topk_indices,
272
+ token_expert_indices,
273
+ gating_output,
274
+ renormalize,
275
+ num_shared_experts,
276
+ shared_expert_scoring_func,
277
)
278
279
@@ -275,6 +283,8 @@ def _rocm_aiter_topk_softmax_fake(
283
token_expert_indices: torch.Tensor,
284
gating_output: torch.Tensor,
285
renormalize: bool,
286
+ num_shared_experts: int = 0,
287
+ shared_expert_scoring_func: str = "",
288
) -> None:
289
pass
290
@@ -1206,6 +1216,9 @@ class rocm_aiter_ops:
1216
_MOE_SHARED_EXPERTS_ENABLED = envs.VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS
1217
# TODO: Consolidate under _LINEAR_ENABLED
1218
_TRITON_UNQUANT_GEMM = envs.VLLM_ROCM_USE_AITER_TRITON_GEMM
1219
+ # Lazily probed: whether aiter.topk_softmax supports the
1220
+ # num_shared_experts / shared_expert_scoring_func args (7-arg form).
1221
+ _TOPK_SOFTMAX_FUSED_SIGMOID: bool | None = None
1222
1223
_ALL_REDUCE_MAX_SIZE: int = 8192 * 1024 * 8 * 2
1224
_CUSTOM_ALL_REDUCE: AiterCustomAllreduceProto | None = None
@@ -1322,6 +1335,52 @@ class rocm_aiter_ops:
1335
def is_fusion_moe_shared_experts_enabled(cls) -> bool:
1336
return cls.is_fused_moe_enabled() and cls._MOE_SHARED_EXPERTS_ENABLED
1337
1338
+ @classmethod
1339
+ @if_aiter_supported
1340
+ def topk_softmax_supports_fused_sigmoid(cls) -> bool:
1341
+ """Check if topk_softmax supports fused shared expert activation."""
1342
+ if cls._TOPK_SOFTMAX_FUSED_SIGMOID is None:
1343
+ try:
1344
+ import inspect
1345
+
1346
+ from aiter import topk_softmax
1347
+
1348
+ params = inspect.signature(topk_softmax).parameters
1349
+ if "num_shared_experts" in params:
1350
+ cls._TOPK_SOFTMAX_FUSED_SIGMOID = True
1351
+ else:
1352
+ # @compile_ops wrapper loses the original signature.
1353
+ # Fall back to the torch custom op schema.
1354
+ import torch
1355
+
1356
+ schema = getattr(
1357
+ getattr(torch.ops.aiter, "topk_softmax", None), "default", None
1358
+ )
1359
+ schema_str = str(getattr(schema, "_schema", ""))
1360
+ cls._TOPK_SOFTMAX_FUSED_SIGMOID = "num_shared_experts" in schema_str
1361
+ except (ImportError, ValueError):
1362
+ cls._TOPK_SOFTMAX_FUSED_SIGMOID = False
1363
+ return cls._TOPK_SOFTMAX_FUSED_SIGMOID
1364
+
1365
+ @classmethod
1366
+ @if_aiter_supported
1367
+ def fuse_sigmoid_in_kernel(cls, aiter_topK_meta_data: object) -> bool:
1368
+ """Whether fused shared-expert sigmoid in the topk kernel is usable.
1369
+
1370
+ Combines the cached static capability checks (FSE enabled, fused-moe
1371
+ enabled, topk_softmax supports fused sigmoid) with the runtime
1372
+ readiness check (topK meta-data buffer initialized).
1373
+
1374
+ ``aiter_topK_meta_data`` is accepted as a parameter rather than
1375
+ imported internally so callers cannot hit initialization-order
1376
+ issues where the module-level global has not been set yet.
1377
+ """
1378
+ return (
1379
+ cls.is_fusion_moe_shared_experts_enabled()
1380
+ and cls.topk_softmax_supports_fused_sigmoid()
1381
+ and aiter_topK_meta_data is not None
1382
+ )
1383
+
1384
@classmethod
1385
@if_aiter_supported
1386
def is_mla_enabled(cls) -> bool:
@@ -1795,9 +1854,17 @@ class rocm_aiter_ops:
1854
token_expert_indices: torch.Tensor,
1855
gating_output: torch.Tensor,
1856
renormalize: bool,
1857
+ num_shared_experts: int = 0,
1858
+ shared_expert_scoring_func: str = "",
1859
) -> tuple[torch.Tensor, ...]:
1860
torch.ops.vllm.rocm_aiter_topk_softmax(
1800
- topk_weights, topk_indices, token_expert_indices, gating_output, renormalize
1861
+ topk_weights,
1862
+ topk_indices,
1863
+ token_expert_indices,
1864
+ gating_output,
1865
+ renormalize,
1866
+ num_shared_experts,
1867
+ shared_expert_scoring_func,
1868
)
1869
return topk_weights, topk_indices
1870
vllm/model_executor/layers/fused_moe/layer.py
+4
@@ -285,6 +285,7 @@ class FusedMoE(PluggableLayer):
285
router_logits_dtype: torch.dtype | None = None,
286
gate: torch.nn.Module | None = None,
287
shared_experts: torch.nn.Module | None = None,
288
+ shared_expert_gate: torch.nn.Module | None = None,
289
routed_input_transform: torch.nn.Module | None = None,
290
routed_output_transform: torch.nn.Module | None = None,
291
apply_routed_scale_to_output: bool = False,
@@ -370,6 +371,8 @@ class FusedMoE(PluggableLayer):
371
if n_shared_experts is not None and self.aiter_fmoe_shared_expert_enabled
372
else 0
373
)
374
+ self.shared_expert_gate = shared_expert_gate
375
+
376
if (
377
not self.aiter_fmoe_shared_expert_enabled
378
and self.num_fused_shared_experts != 0
@@ -608,6 +611,7 @@ class FusedMoE(PluggableLayer):
611
router=self.router,
612
gate=gate,
613
shared_experts=shared_experts,
614
+ shared_expert_gate=self.shared_expert_gate,
615
quant_method=self.quant_method,
616
enable_dbo=self.vllm_config.parallel_config.enable_dbo,
617
routed_input_transform=routed_input_transform,
vllm/model_executor/layers/fused_moe/rocm_aiter_fused_moe.py
+49
@@ -109,6 +109,55 @@ def init_aiter_topK_meta_data(
109
aiter_topK_meta_data = (total_topk_weights, total_topk_ids)
110
111
112
+def inject_shared_expert_weights(
113
+ topk_weights: torch.Tensor,
114
+ topk_ids: torch.Tensor,
115
+ topk: int,
116
+ num_fused_shared_experts: int,
117
+ shared_expert_weights: torch.Tensor | None = None,
118
+) -> tuple[torch.Tensor, torch.Tensor]:
119
+ """Merge routed topk results with the shared expert buffer and inject
120
+ dynamic per-token shared expert gate values for AITER fusion.
121
+
122
+ For routers that already return the combined buffer (e.g. GroupedTopKRouter
123
+ via rocm_aiter_grouped_topk), only the dynamic weight injection is needed.
124
+ For routers that return only routed slots (e.g. FusedTopKRouter), this also
125
+ copies the routed results into the pre-allocated combined buffer.
126
+ """
127
+ if num_fused_shared_experts == 0:
128
+ return topk_weights, topk_ids
129
+
130
+ assert aiter_topK_meta_data is not None, (
131
+ "aiter_topK_meta_data is not initialized but "
132
+ "num_fused_shared_experts > 0. Ensure init_aiter_topK_meta_data "
133
+ "is called before routing."
134
+ )
135
+
136
+ total_topk_weights, total_topk_ids = aiter_topK_meta_data
137
+ token = topk_weights.shape[0]
138
+
139
+ assert total_topk_weights.shape[0] >= token, (
140
+ f"AITER topK meta data supports {total_topk_weights.shape[0]} "
141
+ f"tokens, but got {token} tokens."
142
+ )
143
+
144
+ total_topk_weights_slice = total_topk_weights[:token]
145
+ total_topk_ids_slice = total_topk_ids[:token]
146
+
147
+ if topk_weights.shape[1] == topk:
148
+ total_topk_weights_slice[:, :topk] = topk_weights
149
+ total_topk_ids_slice[:, :topk] = topk_ids
150
+ topk_weights = total_topk_weights_slice
151
+ topk_ids = total_topk_ids_slice
152
+
153
+ if shared_expert_weights is not None:
154
+ topk_weights[:, topk : topk + num_fused_shared_experts] = shared_expert_weights[
155
+ :token
156
+ ]
157
+
158
+ return topk_weights, topk_ids
159
+
160
+
161
def rocm_aiter_grouped_topk(
162
hidden_states: torch.Tensor,
163
gating_output: torch.Tensor,
vllm/model_executor/layers/fused_moe/router/aiter_shared_routed_fused_moe_router.py
new
+143
@@ -0,0 +1,143 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from collections.abc import Callable
4
+
5
+import torch
6
+
7
+from vllm._aiter_ops import rocm_aiter_ops
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.router.base_router import BaseRouter
14
+from vllm.model_executor.layers.fused_moe.router.fused_topk_router import (
15
+ dispatch_topk_softmax_func,
16
+)
17
+
18
+
19
+class AiterSharedRoutedFusedMoERouter(BaseRouter):
20
+ """
21
+ ROCm AITER router for models with fused shared experts (e.g. Qwen3-MoE).
22
+
23
+ When the AITER topk_softmax kernel supports sigmoid fusion, the routing
24
+ softmax and shared-expert sigmoid are computed in a single kernel launch.
25
+ Otherwise the shared-expert weights are injected into the pre-allocated
26
+ AITER buffer via a fallback path.
27
+
28
+ Only instantiated when rocm_aiter fused-MoE is active and
29
+ num_fused_shared_experts > 0.
30
+ """
31
+
32
+ def __init__(
33
+ self,
34
+ top_k: int,
35
+ global_num_experts: int,
36
+ eplb_state: EplbLayerState,
37
+ num_fused_shared_experts: int,
38
+ scoring_func: str = "softmax",
39
+ renormalize: bool = True,
40
+ enable_eplb: bool = False,
41
+ indices_type_getter: Callable[[], torch.dtype | None] | None = None,
42
+ ):
43
+ super().__init__(
44
+ top_k=top_k,
45
+ global_num_experts=global_num_experts,
46
+ eplb_state=eplb_state,
47
+ enable_eplb=enable_eplb,
48
+ indices_type_getter=indices_type_getter,
49
+ )
50
+ self.renormalize = renormalize
51
+ self.scoring_func = scoring_func
52
+ self.num_fused_shared_experts = num_fused_shared_experts
53
+
54
+ @property
55
+ def routing_method_type(self) -> RoutingMethodType:
56
+ return get_routing_method_type(
57
+ scoring_func=self.scoring_func,
58
+ top_k=self.top_k,
59
+ renormalize=self.renormalize,
60
+ num_expert_group=None,
61
+ has_e_score_bias=False,
62
+ )
63
+
64
+ def _compute_routing(
65
+ self,
66
+ hidden_states: torch.Tensor,
67
+ router_logits: torch.Tensor,
68
+ indices_type: torch.dtype | None,
69
+ *,
70
+ input_ids: torch.Tensor | None = None,
71
+ ) -> tuple[torch.Tensor, torch.Tensor]:
72
+ assert hidden_states.size(0) == router_logits.size(0), (
73
+ "Number of tokens mismatch"
74
+ )
75
+
76
+ from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (
77
+ aiter_topK_meta_data,
78
+ )
79
+
80
+ M = hidden_states.size(0)
81
+ topk = self.top_k
82
+ num_fse = self.num_fused_shared_experts
83
+
84
+ token_expert_indices = torch.empty(
85
+ M, topk, dtype=torch.int32, device=hidden_states.device
86
+ )
87
+
88
+ if rocm_aiter_ops.fuse_sigmoid_in_kernel(aiter_topK_meta_data):
89
+ total_topk_weights, total_topk_ids = aiter_topK_meta_data # type: ignore[misc]
90
+ total_topk_weights_slice = total_topk_weights[:M]
91
+ topk_ids_slice = total_topk_ids[:M, :topk]
92
+
93
+ topk_func = dispatch_topk_softmax_func(use_rocm_aiter=True)
94
+ topk_func(
95
+ total_topk_weights_slice,
96
+ topk_ids_slice,
97
+ token_expert_indices,
98
+ router_logits,
99
+ self.renormalize,
100
+ num_fse,
101
+ "sigmoid",
102
+ )
103
+ return total_topk_weights_slice, total_topk_ids[:M]
104
+
105
+ routing_logits = router_logits[:, :-num_fse]
106
+ shared_logits = router_logits[:, -num_fse:]
107
+
108
+ topk_weights = torch.empty(
109
+ M, topk, dtype=torch.float32, device=hidden_states.device
110
+ )
111
+ topk_ids = torch.empty(
112
+ M,
113
+ topk,
114
+ dtype=torch.int32 if indices_type is None else indices_type,
115
+ device=hidden_states.device,
116
+ )
117
+
118
+ topk_func = dispatch_topk_softmax_func(
119
+ use_rocm_aiter=rocm_aiter_ops.is_fused_moe_enabled()
120
+ )
121
+ topk_weights, topk_ids = topk_func(
122
+ topk_weights,
123
+ topk_ids,
124
+ token_expert_indices,
125
+ routing_logits,
126
+ self.renormalize,
127
+ )
128
+
129
+ if aiter_topK_meta_data is not None:
130
+ from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (
131
+ inject_shared_expert_weights,
132
+ )
133
+
134
+ shared_weights = torch.sigmoid(shared_logits)
135
+ topk_weights, topk_ids = inject_shared_expert_weights(
136
+ topk_weights,
137
+ topk_ids,
138
+ topk=topk,
139
+ num_fused_shared_experts=num_fse,
140
+ shared_expert_weights=shared_weights,
141
+ )
142
+
143
+ return topk_weights, topk_ids
vllm/model_executor/layers/fused_moe/router/router_factory.py
+25
-2
@@ -5,8 +5,14 @@ from collections.abc import Callable
5
import torch
6
7
import vllm.envs as envs
8
+from vllm._aiter_ops import rocm_aiter_ops
9
from vllm.distributed.eplb.eplb_state import EplbLayerState
9
-from vllm.model_executor.layers.fused_moe.config import RoutingMethodType
10
+from vllm.model_executor.layers.fused_moe.config import (
11
+ RoutingMethodType,
12
+)
13
+from vllm.model_executor.layers.fused_moe.router.aiter_shared_routed_fused_moe_router import ( # noqa: E501
14
+ AiterSharedRoutedFusedMoERouter,
15
+)
16
from vllm.model_executor.layers.fused_moe.router.custom_routing_router import (
17
CustomRoutingRouter,
18
)
@@ -67,7 +73,8 @@ def create_fused_moe_router(
73
3. GroupedTopKRouter - if use_grouped_topk is True
74
4. CustomRoutingRouter - if custom_routing_function is not None
75
5. FusedTopKBiasRouter - if e_score_correction_bias is not None
70
- 6. FusedTopKRouter - default fallback
76
+ 6. AiterSharedRoutedFusedMoERouter - if num_fused_shared_experts > 0
77
+ 7. FusedTopKRouter - default fallback
78
79
Common arguments:
80
top_k: Number of experts to select per token
@@ -199,6 +206,22 @@ def create_fused_moe_router(
206
hash_indices_table=hash_indices_table,
207
)
208
209
+ if (
210
+ num_fused_shared_experts > 0
211
+ and scoring_func == "softmax"
212
+ and rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()
213
+ ):
214
+ return AiterSharedRoutedFusedMoERouter(
215
+ top_k=top_k,
216
+ global_num_experts=global_num_experts,
217
+ eplb_state=eplb_state,
218
+ num_fused_shared_experts=num_fused_shared_experts,
219
+ renormalize=renormalize,
220
+ scoring_func=scoring_func,
221
+ enable_eplb=enable_eplb,
222
+ indices_type_getter=indices_type_getter,
223
+ )
224
+
225
return FusedTopKRouter(
226
top_k=top_k,
227
global_num_experts=global_num_experts,
vllm/model_executor/layers/fused_moe/runner/moe_runner.py
+29
-1
@@ -223,6 +223,7 @@ class MoERunner(MoERunnerInterface):
223
shared_experts: torch.nn.Module | None,
224
quant_method: FusedMoEMethodBase,
225
enable_dbo: bool,
226
+ shared_expert_gate: torch.nn.Module | None = None,
227
routed_output_transform: torch.nn.Module | None = None,
228
routed_scaling_factor: float = 1.0,
229
):
@@ -233,9 +234,18 @@ class MoERunner(MoERunnerInterface):
234
self.routed_output_transform = routed_output_transform
235
self.routed_scaling_factor = routed_scaling_factor
236
self.gate = gate
237
+ self.shared_expert_gate = shared_expert_gate
238
self._quant_method = quant_method
239
self.enable_dbo = enable_dbo
240
241
+ # When both gates are present and FSE is enabled, fuse their
242
+ # weight matrices into [num_experts + num_shared, hidden] so one
243
+ # F.linear produces combined logits. The topk kernel can then
244
+ # apply routing softmax and shared expert activation (sigmoid)
245
+ # in a single launch.
246
+ self._fse_fuse_gate = gate is not None and shared_expert_gate is not None
247
+ self._combined_gate_weight: torch.Tensor | None = None
248
+
249
self._shared_experts: SharedExperts | None = None
250
if shared_experts is not None:
251
self._shared_experts = SharedExperts(
@@ -278,6 +288,20 @@ class MoERunner(MoERunnerInterface):
288
self._shared_experts._quant_method = quant_method
289
self._quant_method = quant_method
290
291
+ def _maybe_fuse_gate_weights(self):
292
+ """Fuse router and shared expert gate weights on first call.
293
+
294
+ Cannot be done at __init__ because gate weights are loaded after
295
+ module construction (via weight_loader). Called once from
296
+ _forward_impl before the first forward pass.
297
+ """
298
+ if self._combined_gate_weight is None:
299
+ assert self.gate is not None and self.shared_expert_gate is not None
300
+ self._combined_gate_weight = torch.cat(
301
+ [self.gate.weight, self.shared_expert_gate.weight],
302
+ dim=0,
303
+ )
304
+
305
def is_internal_router(self) -> bool:
306
return self.gate is not None
307
@@ -752,7 +776,11 @@ class MoERunner(MoERunnerInterface):
776
# so it can run overlapped with the
777
# NOTE: in future PR, MoE runner will always hold the gate.
778
if self.gate is not None:
755
- router_logits, _ = self.gate(hidden_states)
779
+ if self._fse_fuse_gate:
780
+ self._maybe_fuse_gate_weights()
781
+ router_logits = F.linear(hidden_states, self._combined_gate_weight)
782
+ else:
783
+ router_logits, _ = self.gate(hidden_states)
784
785
with self._sequence_parallel_context():
786
# TODO(bnell): parts of the dispatch/combine steps will go away once
vllm/model_executor/models/qwen3_next.py
+26
-4
@@ -8,6 +8,7 @@ from itertools import islice
8
import torch
9
from torch import nn
10
11
+from vllm._aiter_ops import rocm_aiter_ops
12
from vllm.compilation.decorators import support_torch_compile
13
from vllm.config import (
14
CacheConfig,
@@ -135,7 +136,12 @@ class Qwen3NextSparseMoeBlock(nn.Module):
136
prefix=f"{prefix}.shared_expert_gate",
137
)
138
138
- if config.shared_expert_intermediate_size > 0:
139
+ if (
140
+ rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()
141
+ or config.shared_expert_intermediate_size <= 0
142
+ ):
143
+ self.shared_expert = None
144
+ else:
145
self.shared_expert = Qwen3NextMLP(
146
hidden_size=config.hidden_size,
147
intermediate_size=config.shared_expert_intermediate_size,
@@ -146,8 +152,6 @@ class Qwen3NextSparseMoeBlock(nn.Module):
152
is_sequence_parallel=self.is_sequence_parallel,
153
prefix=f"{prefix}.shared_expert",
154
)
149
- else:
150
- self.shared_expert = None
155
156
self.experts = FusedMoE(
157
shared_experts=self.shared_expert,
@@ -162,6 +166,10 @@ class Qwen3NextSparseMoeBlock(nn.Module):
166
enable_eplb=self.enable_eplb,
167
num_redundant_experts=self.n_redundant_experts,
168
is_sequence_parallel=self.is_sequence_parallel,
169
+ n_shared_experts=1 if self.shared_expert is None else None,
170
+ shared_expert_gate=self.shared_expert_gate
171
+ if self.shared_expert is None
172
+ else None,
173
)
174
175
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
@@ -536,12 +544,15 @@ class Qwen3NextModel(nn.Module, EagleModelMixin):
544
def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
545
# Params for weights, fp8 weight scales, fp8 activation scales
546
# (param_name, weight_name, expert_id, shard_id)
547
+ num_experts = getattr(self.config, "num_experts", 0)
548
+ if rocm_aiter_ops.is_fusion_moe_shared_experts_enabled():
549
+ num_experts += 1
550
return fused_moe_make_expert_params_mapping(
551
self,
552
ckpt_gate_proj_name="gate_proj",
553
ckpt_down_proj_name="down_proj",
554
ckpt_up_proj_name="up_proj",
544
- num_experts=getattr(self.config, "num_experts", 0),
555
+ num_experts=num_experts,
556
num_redundant_experts=self.num_redundant_experts,
557
)
558
@@ -558,6 +569,10 @@ class Qwen3NextModel(nn.Module, EagleModelMixin):
569
params_dict = dict(self.named_parameters())
570
loaded_params: set[str] = set()
571
expert_params_mapping = self.get_expert_mapping()
572
+
573
+ is_fse = rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()
574
+ num_routed = getattr(self.config, "num_experts", 0)
575
+
576
for name, loaded_weight in weights:
577
if "rotary_emb.inv_freq" in name:
578
continue
@@ -571,6 +586,13 @@ class Qwen3NextModel(nn.Module, EagleModelMixin):
586
if name is None:
587
continue
588
589
+ # FSE: remap shared_expert weights to the fused expert slot
590
+ if is_fse and "mlp.shared_expert." in name:
591
+ name = name.replace(
592
+ "mlp.shared_expert.",
593
+ f"mlp.experts.{num_routed}.",
594
+ )
595
+
596
for param_name, weight_name, shard_id in stacked_params_mapping:
597
if weight_name not in name:
598
continue
vllm/model_executor/models/qwen3_next_mtp.py
+14
-1
@@ -7,6 +7,7 @@ from collections.abc import Iterable
7
import torch
8
from torch import nn
9
10
+from vllm._aiter_ops import rocm_aiter_ops
11
from vllm.compilation.decorators import support_torch_compile
12
from vllm.config import VllmConfig
13
from vllm.distributed.parallel_state import get_pp_group
@@ -147,13 +148,18 @@ class Qwen3NextMultiTokenPredictor(nn.Module):
148
149
# Params for weights, fp8 weight scales, fp8 activation scales
150
# (param_name, weight_name, expert_id, shard_id)
151
+ is_fse = rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()
152
+ num_experts = self.config.num_experts
153
+ if is_fse:
154
+ num_experts += 1
155
expert_params_mapping = fused_moe_make_expert_params_mapping(
156
self,
157
ckpt_gate_proj_name="gate_proj",
158
ckpt_down_proj_name="down_proj",
159
ckpt_up_proj_name="up_proj",
155
- num_experts=self.config.num_experts,
160
+ num_experts=num_experts,
161
)
162
+ num_routed = self.config.num_experts
163
164
params_dict = dict(self.named_parameters())
165
loaded_params: set[str] = set()
@@ -161,6 +167,13 @@ class Qwen3NextMultiTokenPredictor(nn.Module):
167
if "rotary_emb.inv_freq" in name:
168
continue
169
170
+ # FSE: remap shared_expert weights to the fused expert slot
171
+ if is_fse and "mlp.shared_expert." in name:
172
+ name = name.replace(
173
+ "mlp.shared_expert.",
174
+ f"mlp.experts.{num_routed}.",
175
+ )
176
+
177
for param_name, weight_name, shard_id in stacked_params_mapping:
178
if weight_name not in name:
179
continue