[ROCm][Bugfix][Perf] enable shared expert fusion for Qwen3.5 (#44434)
Signed-off-by: Nico Holmberg <nico.holmberg@amd.com>
Nico Holmberg committed
Jun 23, 2026 at 16:19 UTC
80e511772f3e83461fc3f73112b0ee1664713cde
2 files changed
+48
-4
vllm/model_executor/models/qwen3_5.py
+19
@@ -30,6 +30,7 @@ from collections.abc import Callable, Iterable
30
import torch
31
from torch import nn
32
33
+from vllm._aiter_ops import rocm_aiter_ops
34
from vllm.compilation.decorators import support_torch_compile
35
from vllm.config import VllmConfig
36
from vllm.distributed import (
@@ -314,6 +315,15 @@ class Qwen3_5Model(Qwen3NextModel):
315
num_experts = (
316
self.config.num_experts if hasattr(self.config, "num_experts") else 0
317
)
318
+ from vllm.config import get_current_vllm_config
319
+
320
+ from .qwen3_next import _is_shared_expert_fse_compatible
321
+
322
+ is_fse = (
323
+ rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()
324
+ and _is_shared_expert_fse_compatible(get_current_vllm_config().quant_config)
325
+ )
326
+
327
for name, loaded_weight in weights:
328
if "rotary_emb.inv_freq" in name:
329
continue
@@ -327,6 +337,15 @@ class Qwen3_5Model(Qwen3NextModel):
337
if name is None:
338
continue
339
340
+ # FSE: remap shared_expert weights to fused expert slot
341
+ if is_fse and "mlp.shared_expert." in name:
342
+ name = name.replace(
343
+ "mlp.shared_expert.",
344
+ f"mlp.experts.{num_experts}.",
345
+ )
346
+ is_fused_expert = False
347
+ expert_params_mapping = self.get_expert_mapping()
348
+
349
for param_name, weight_name, shard_id in stacked_params_mapping:
350
if "experts.gate_up_proj" in name or "experts.down_proj" in name:
351
is_fused_expert = True
vllm/model_executor/models/qwen3_next.py
+29
-4
@@ -88,6 +88,26 @@ logger = init_logger(__name__)
88
KVCache = tuple[torch.Tensor, torch.Tensor]
89
90
91
+def _is_shared_expert_fse_compatible(quant_config) -> bool:
92
+ """Check if shared expert can be fused with routed experts.
93
+
94
+ FSE requires that shared and routed expert weights use the same
95
+ quantization format. Returns False when the shared expert is
96
+ excluded from quantization (e.g. float32 shared in an MXFP4 model)
97
+ or has a different quant spec than routed experts.
98
+ """
99
+ if quant_config is None:
100
+ return True
101
+ # Quark stores its full config dict in quant_config.quant_config
102
+ raw_config = getattr(quant_config, "quant_config", None)
103
+ if not isinstance(raw_config, dict):
104
+ return True
105
+ exclude = raw_config.get("exclude", [])
106
+ if not exclude:
107
+ return True
108
+ return not any("shared_expert." in str(e) for e in exclude)
109
+
110
+
111
class Qwen3NextSparseMoeBlock(nn.Module):
112
def __init__(self, vllm_config: VllmConfig, prefix: str = ""):
113
super().__init__()
@@ -142,10 +162,15 @@ class Qwen3NextSparseMoeBlock(nn.Module):
162
prefix=f"{prefix}.shared_expert_gate",
163
)
164
145
- if (
146
- rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()
147
- or config.shared_expert_intermediate_size <= 0
148
- ):
165
+ _fse_requested = rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()
166
+ _fse_enabled = _fse_requested and _is_shared_expert_fse_compatible(quant_config)
167
+ if _fse_requested and not _fse_enabled:
168
+ logger.warning(
169
+ "VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS is enabled but "
170
+ "shared expert has a different quantization spec than routed "
171
+ "experts. Falling back to non-fused shared expert path."
172
+ )
173
+ if _fse_enabled or config.shared_expert_intermediate_size <= 0:
174
self.shared_expert = None
175
else:
176
self.shared_expert = Qwen3NextMLP(