@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 88b2bff2c6

[MOE] Standardize and abstract fused shared expert optimization selection (#51695)

Signed-off-by: Felix Marty <Felix.Marty@amd.com>

fxmarty-amd committed Aug 18, 2026 at 18:56 UTC 88b2bff2c63d0f28396451f1199d09ee0f3e2d88
20 files changed +1288 -281
tests/model_executor/layers/test_fused_shared_expert.py new
+700
@@ -0,0 +1,700 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import importlib
5 +from copy import deepcopy
6 +from types import SimpleNamespace
7 +from typing import Any, cast
8 +from unittest.mock import patch
9 +
10 +import pytest
11 +import torch
12 +from torch import nn
13 +
14 +import vllm.config as vllm_config_module
15 +from vllm.config import VllmConfig, set_current_vllm_config
16 +from vllm.model_executor.layers.fused_moe import utils as fused_moe_utils
17 +from vllm.model_executor.layers.fused_moe.layer import determine_expert_counts
18 +from vllm.model_executor.layers.quantization.quark.quark import QuarkConfig
19 +from vllm.model_executor.layers.quantization.utils.config_utils import (
20 + is_shared_expert_quant_fse_compatible,
21 +)
22 +from vllm.model_executor.models.utils import PPMissingLayer
23 +from vllm.models.deepseek_v4 import quant_config as deepseek_v4_quant_config
24 +from vllm.models.minimax_m3.amd import model as minimax_m3_model
25 +from vllm.transformers_utils.configs.minimax_m3 import MiniMaxM3TextConfig
26 +from vllm.transformers_utils.configs.qwen3_5_moe import Qwen3_5MoeTextConfig
27 +
28 +_QUARK_FSE_CONFIG: dict[str, Any] = {
29 + "global_quant_config": {
30 + "input_tensors": {
31 + "dtype": "fp4",
32 + "is_dynamic": True,
33 + "qscheme": "per_group",
34 + "ch_axis": -1,
35 + "group_size": 32,
36 + "block_size": None,
37 + "symmetric": None,
38 + "round_method": "half_even",
39 + "scale_type": "float",
40 + "scale_format": "e8m0",
41 + "scale_calculation_mode": "even",
42 + "mx_element_dtype": None,
43 + "observer_cls": "PerBlockMXObserver",
44 + "is_scale_quant": False,
45 + "enable_buffer_reuse": False,
46 + "max_input_numel": 4194304,
47 + },
48 + "output_tensors": None,
49 + "weight": {
50 + "dtype": "fp4",
51 + "is_dynamic": False,
52 + "qscheme": "per_group",
53 + "ch_axis": -1,
54 + "group_size": 32,
55 + "block_size": None,
56 + "symmetric": None,
57 + "round_method": "half_even",
58 + "scale_type": "float",
59 + "scale_format": "e8m0",
60 + "scale_calculation_mode": "even",
61 + "mx_element_dtype": None,
62 + "observer_cls": "PerBlockMXObserver",
63 + "is_scale_quant": False,
64 + "enable_buffer_reuse": False,
65 + "max_input_numel": 4194304,
66 + },
67 + "bias": None,
68 + "target_device": None,
69 + },
70 + "algo_config": None,
71 + "softmax_quant_spec": None,
72 + "quant_method": "quark",
73 + "layer_type_quant_config": {},
74 + "layer_quant_config": {},
75 + "kv_cache_quant_config": {},
76 + "kv_cache_post_rope": False,
77 + "quant_mode": "eager_mode",
78 + "version": "0.12+9d3d471cdf1",
79 + "export": {
80 + "kv_cache_group": [],
81 + "min_kv_scale": 0.0,
82 + "pack_method": "reorder",
83 + "weight_format": "real_quantized",
84 + "weight_merge_groups": None,
85 + },
86 +}
87 +
88 +
89 +def get_deepseek_v4_quark_config(exclude: list[str]) -> dict[str, Any]:
90 + """Return the DeepSeek-V4-Pro-MXFP4 FSE quantization layout."""
91 + # Mimics https://huggingface.co/amd/DeepSeek-V4-Pro-MXFP4.
92 + quantization_config: dict[str, Any] = deepcopy(_QUARK_FSE_CONFIG)
93 + quantization_config["exclude"] = exclude
94 + mxfp4_config = cast(
95 + dict[str, Any], deepcopy(quantization_config["global_quant_config"])
96 + )
97 + fp8_config = deepcopy(mxfp4_config)
98 + fp8_config["input_tensors"].update(
99 + {
100 + "dtype": "fp8_e4m3",
101 + "group_size": 128,
102 + "symmetric": True,
103 + "scale_type": None,
104 + "scale_format": None,
105 + "scale_calculation_mode": None,
106 + "observer_cls": "PerGroupMinMaxObserver",
107 + }
108 + )
109 + fp8_config["weight"].update(
110 + {
111 + "dtype": "fp8_e4m3",
112 + "qscheme": "per_block",
113 + "ch_axis": None,
114 + "group_size": None,
115 + "block_size": [128, 128],
116 + "symmetric": True,
117 + "scale_type": "float8_e8m0fnu",
118 + "scale_format": None,
119 + "scale_calculation_mode": None,
120 + "observer_cls": "PerBlock2DMinMaxObserver",
121 + }
122 + )
123 + quantization_config["layer_quant_config"] = {
124 + r"re:mtp\.0\.ffn\.experts\.\d+\.w1": mxfp4_config,
125 + r"re:mtp\.0\.ffn\.shared_experts\.w1": fp8_config,
126 + }
127 + return quantization_config
128 +
129 +
130 +def get_fse_test_model_config(
131 + model_type: str,
132 + quantization_config: dict[str, Any],
133 +) -> tuple[object, type[nn.Module]]:
134 + if model_type == "minimax_m3":
135 + config = MiniMaxM3TextConfig(
136 + hidden_size=128,
137 + intermediate_size=32,
138 + dense_intermediate_size=32,
139 + num_hidden_layers=2,
140 + num_attention_heads=1,
141 + num_key_value_heads=1,
142 + head_dim=128,
143 + num_local_experts=2,
144 + num_experts_per_tok=1,
145 + moe_layer_freq=[0, 1],
146 + sparse_attention_config=None,
147 + rotary_dim=64,
148 + quantization_config=quantization_config,
149 + )
150 + return config, minimax_m3_model.MiniMaxM3Model
151 + if model_type == "deepseek_v4":
152 + from vllm.models.deepseek_v4.amd.model import DeepseekV4Model
153 +
154 + config = SimpleNamespace(
155 + vocab_size=256,
156 + hidden_size=128,
157 + num_hidden_layers=1,
158 + num_attention_heads=1,
159 + head_dim=128,
160 + max_position_embeddings=128,
161 + sliding_window=None,
162 + compress_ratios=[1],
163 + rope_theta=10000.0,
164 + compress_rope_theta=10000.0,
165 + rope_parameters={"rope_type": "default"},
166 + rms_norm_eps=1e-6,
167 + hidden_act="silu",
168 + q_lora_rank=0,
169 + o_lora_rank=0,
170 + o_groups=1,
171 + qk_rope_head_dim=64,
172 + index_head_dim=64,
173 + index_n_heads=1,
174 + index_topk=1,
175 + n_routed_experts=2,
176 + n_shared_experts=1,
177 + num_experts_per_tok=1,
178 + n_group=1,
179 + topk_group=1,
180 + moe_intermediate_size=32,
181 + norm_topk_prob=False,
182 + num_hash_layers=0,
183 + swiglu_limit=0.0,
184 + hc_eps=1e-6,
185 + hc_mult=1,
186 + hc_sinkhorn_iters=1,
187 + expert_dtype="fp4",
188 + num_nextn_predict_layers=1,
189 + quantization_config=quantization_config,
190 + )
191 + return config, DeepseekV4Model
192 + if model_type == "qwen3_5":
193 + from vllm.model_executor.models.qwen3_5 import Qwen3_5Model
194 +
195 + config = Qwen3_5MoeTextConfig(
196 + vocab_size=256,
197 + hidden_size=128,
198 + num_hidden_layers=1,
199 + num_attention_heads=1,
200 + num_key_value_heads=1,
201 + head_dim=128,
202 + linear_key_head_dim=64,
203 + linear_value_head_dim=64,
204 + linear_num_key_heads=1,
205 + linear_num_value_heads=1,
206 + moe_intermediate_size=32,
207 + shared_expert_intermediate_size=32,
208 + num_experts_per_tok=1,
209 + num_experts=2,
210 + layer_types=["full_attention"],
211 + quantization_config=quantization_config,
212 + )
213 + return config, Qwen3_5Model
214 + if model_type == "deepseek_v2":
215 + from transformers import DeepseekV2Config
216 +
217 + from vllm.model_executor.models.deepseek_v2 import DeepseekV2Model
218 +
219 + config = DeepseekV2Config(
220 + vocab_size=256,
221 + hidden_size=128,
222 + intermediate_size=32,
223 + num_hidden_layers=1,
224 + num_attention_heads=1,
225 + num_key_value_heads=1,
226 + head_dim=128,
227 + first_k_dense_replace=0,
228 + n_routed_experts=2,
229 + n_shared_experts=1,
230 + num_experts_per_tok=1,
231 + n_group=1,
232 + topk_group=1,
233 + moe_intermediate_size=32,
234 + q_lora_rank=None,
235 + kv_lora_rank=0,
236 + qk_nope_head_dim=0,
237 + qk_rope_head_dim=0,
238 + v_head_dim=0,
239 + quantization_config=quantization_config,
240 + )
241 + return config, DeepseekV2Model
242 + if model_type == "glm4_moe":
243 + from transformers.models.glm4_moe import Glm4MoeConfig
244 +
245 + from vllm.model_executor.models.glm4_moe import Glm4MoeModel
246 +
247 + config = Glm4MoeConfig(
248 + vocab_size=256,
249 + hidden_size=128,
250 + intermediate_size=32,
251 + num_hidden_layers=1,
252 + num_attention_heads=1,
253 + num_key_value_heads=1,
254 + head_dim=128,
255 + moe_intermediate_size=32,
256 + n_routed_experts=2,
257 + n_shared_experts=1,
258 + num_experts_per_tok=1,
259 + first_k_dense_replace=0,
260 + quantization_config=quantization_config,
261 + )
262 + return config, Glm4MoeModel
263 + raise ValueError(f"Unsupported FSE test model: {model_type}")
264 +
265 +
266 +def test_determine_expert_counts_fuse_shared_experts_override(
267 + monkeypatch: pytest.MonkeyPatch,
268 +) -> None:
269 + common_args = (8, 0, 2)
270 + assert determine_expert_counts(*common_args, True)[2] == 2
271 + assert determine_expert_counts(*common_args, False)[2] == 0
272 +
273 +
274 +def test_resolve_layer_fused_shared_expert_skips_compatibility_when_disabled(
275 + monkeypatch: pytest.MonkeyPatch,
276 +) -> None:
277 + monkeypatch.setattr(
278 + fused_moe_utils.rocm_aiter_ops,
279 + "is_fusion_moe_shared_experts_enabled",
280 + lambda: False,
281 + )
282 + monkeypatch.setattr(
283 + fused_moe_utils,
284 + "is_shared_expert_quant_fse_compatible",
285 + lambda *_: pytest.fail(
286 + "compatibility must not be checked when FSE is disabled"
287 + ),
288 + )
289 +
290 + assert not fused_moe_utils.resolve_layer_fused_shared_expert(
291 + object(), "model.layers.0.mlp"
292 + )
293 +
294 +
295 +def test_resolve_layer_fused_shared_expert_normalizes_unavailable_aiter(
296 + monkeypatch: pytest.MonkeyPatch,
297 +) -> None:
298 + monkeypatch.setattr(
299 + fused_moe_utils.rocm_aiter_ops,
300 + "is_fusion_moe_shared_experts_enabled",
301 + lambda: None,
302 + )
303 +
304 + assert (
305 + fused_moe_utils.resolve_layer_fused_shared_expert(
306 + object(), "model.layers.0.mlp"
307 + )
308 + is False
309 + )
310 +
311 +
312 +def test_resolve_layer_fused_shared_expert_passes_module_prefixes(
313 + monkeypatch: pytest.MonkeyPatch,
314 +) -> None:
315 + quant_config = object()
316 + monkeypatch.setattr(
317 + fused_moe_utils.rocm_aiter_ops,
318 + "is_fusion_moe_shared_experts_enabled",
319 + lambda: True,
320 + )
321 +
322 + calls: list[tuple[object, str, str]] = []
323 +
324 + def check_compatibility(
325 + config: object, expert_prefix: str, shared_expert_prefix: str
326 + ) -> tuple[bool, str | None]:
327 + calls.append((config, expert_prefix, shared_expert_prefix))
328 + return True, None
329 +
330 + monkeypatch.setattr(
331 + fused_moe_utils, "is_shared_expert_quant_fse_compatible", check_compatibility
332 + )
333 +
334 + assert fused_moe_utils.resolve_layer_fused_shared_expert(
335 + quant_config,
336 + "model.layers.0.mlp",
337 + shared_expert_name="shared_expert",
338 + )
339 + assert calls == [
340 + (
341 + quant_config,
342 + "model.layers.0.mlp.experts",
343 + "model.layers.0.mlp.shared_expert",
344 + )
345 + ]
346 +
347 +
348 +def test_resolve_layer_fused_shared_expert_rejects_incompatible_quantization(
349 + monkeypatch: pytest.MonkeyPatch,
350 + caplog: pytest.LogCaptureFixture,
351 +) -> None:
352 + monkeypatch.setattr(
353 + fused_moe_utils.rocm_aiter_ops,
354 + "is_fusion_moe_shared_experts_enabled",
355 + lambda: True,
356 + )
357 + monkeypatch.setattr(
358 + fused_moe_utils,
359 + "is_shared_expert_quant_fse_compatible",
360 + lambda *_: (False, "shared experts are excluded"),
361 + )
362 +
363 + assert not fused_moe_utils.resolve_layer_fused_shared_expert(
364 + object(), "model.layers.0.mlp"
365 + )
366 + assert "shared experts are excluded" in caplog.text
367 +
368 +
369 +def test_deepseek_v4_shared_expert_fse_uses_mtp_quantization_config_prefix(
370 + monkeypatch: pytest.MonkeyPatch,
371 +) -> None:
372 + class DeepseekV4Config:
373 + expert_dtype = "fp4"
374 +
375 + def _is_quark_mxfp4_ocp(self, hf_config: object) -> bool:
376 + return True
377 +
378 + hf_config = SimpleNamespace(
379 + num_hidden_layers=2,
380 + quantization_config={
381 + "layer_quant_config": {
382 + r"re:mtp\.0\.ffn\.shared_experts\.w1": {"weight": {"dtype": "fp4"}}
383 + },
384 + "global_quant_config": {"weight": {"dtype": "fp8"}},
385 + },
386 + )
387 + monkeypatch.setattr(
388 + deepseek_v4_quant_config, "DeepseekV4FP8Config", DeepseekV4Config
389 + )
390 + monkeypatch.setattr(
391 + vllm_config_module,
392 + "get_current_vllm_config",
393 + lambda: SimpleNamespace(model_config=SimpleNamespace(hf_config=hf_config)),
394 + )
395 +
396 + compatible, reason = is_shared_expert_quant_fse_compatible(
397 + DeepseekV4Config(),
398 + "model.layers.2.ffn.experts",
399 + "model.layers.2.ffn.shared_experts",
400 + )
401 +
402 + assert compatible
403 + assert reason is None
404 +
405 +
406 +def test_is_model_fused_shared_expert_compatible() -> None:
407 + class MoE(nn.Module):
408 + def __init__(self, enabled: bool) -> None:
409 + super().__init__()
410 + self.is_fused_shared_expert_enabled = enabled
411 +
412 + class Layer(nn.Module):
413 + def __init__(self, enabled: bool) -> None:
414 + super().__init__()
415 + self.mlp = MoE(enabled)
416 +
417 + enabled_layers = nn.ModuleList([Layer(True)])
418 + disabled_layers = nn.ModuleList([Layer(False)])
419 + mixed_layers = nn.ModuleList([Layer(True), Layer(False)])
420 + empty_layers = nn.ModuleList()
421 + pipeline_layers = nn.ModuleList([Layer(True), PPMissingLayer()])
422 +
423 + assert fused_moe_utils.is_model_fused_shared_expert_compatible(
424 + enabled_layers, MoE, "mlp"
425 + )
426 + assert not fused_moe_utils.is_model_fused_shared_expert_compatible(
427 + disabled_layers, MoE, "mlp"
428 + )
429 + assert not fused_moe_utils.is_model_fused_shared_expert_compatible(
430 + empty_layers, MoE, "mlp"
431 + )
432 + assert fused_moe_utils.is_model_fused_shared_expert_compatible(
433 + pipeline_layers, MoE, "mlp"
434 + )
435 + with pytest.raises(NotImplementedError, match="1 enabled and 1 disabled layers"):
436 + fused_moe_utils.is_model_fused_shared_expert_compatible(
437 + mixed_layers, MoE, "mlp"
438 + )
439 +
440 +
441 +@pytest.mark.parametrize(
442 + "model_type",
443 + ["minimax_m3", "deepseek_v4", "qwen3_5", "glm4_moe", "deepseek_v2"],
444 +)
445 +@pytest.mark.parametrize(
446 + ("use_fse", "exclude"),
447 + [
448 + (False, []),
449 + (True, []),
450 + (True, ["*.shared_experts.*"]),
451 + ],
452 +)
453 +def test_models_fse_init(
454 + model_type: str,
455 + use_fse: bool,
456 + exclude: list[str],
457 + dist_init: None,
458 + monkeypatch: pytest.MonkeyPatch,
459 +) -> None:
460 + """Model construction resolves FSE consistently with Quark quantization."""
461 +
462 + quantization_config: dict[str, Any] = (
463 + get_deepseek_v4_quark_config(["layers.0.ffn.shared_experts"] if exclude else [])
464 + if model_type == "deepseek_v4"
465 + else {**_QUARK_FSE_CONFIG, "exclude": exclude}
466 + )
467 +
468 + config, model_constructor = get_fse_test_model_config(
469 + model_type, quantization_config
470 + )
471 + vllm_config = VllmConfig()
472 + vllm_config.model_config = SimpleNamespace(
473 + hf_config=config,
474 + hf_text_config=config,
475 + dtype=torch.bfloat16,
476 + max_model_len=128,
477 + is_diffusion=False,
478 + head_dtype=None,
479 + is_mm_prefix_lm=False,
480 + multimodal_config=None,
481 + quantization_config=None,
482 + )
483 + vllm_config.parallel_config.enable_expert_parallel = False
484 + if model_type == "deepseek_v4":
485 + from vllm.models.deepseek_v4.quant_config import DeepseekV4FP8Config
486 +
487 + vllm_config.cache_config.cache_dtype = "fp8_ds_mla"
488 + vllm_config.quant_config = DeepseekV4FP8Config(
489 + is_checkpoint_fp8_serialized=True,
490 + weight_block_size=[128, 128],
491 + )
492 + else:
493 + vllm_config.quant_config = QuarkConfig(quantization_config)
494 +
495 + import vllm.envs as envs
496 + from vllm._aiter_ops import rocm_aiter_ops
497 +
498 + if model_type == "deepseek_v4":
499 + from vllm.models.deepseek_v4.amd import model as deepseek_v4_model
500 +
501 + warning_logger = deepseek_v4_model.logger
502 + elif model_type == "minimax_m3":
503 + warning_logger = minimax_m3_model.logger
504 + else:
505 + warning_logger = fused_moe_utils.logger
506 +
507 + with monkeypatch.context() as mp:
508 + mp.setenv("VLLM_ROCM_USE_AITER", str(use_fse))
509 + mp.setenv("VLLM_ROCM_USE_AITER_MOE", str(use_fse))
510 + mp.setenv("VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS", str(use_fse))
511 + mp.setattr(
512 + "vllm.model_executor.layers.fused_moe.experts."
513 + "ocp_mx_emulation_moe.has_quark",
514 + lambda: True,
515 + )
516 + importlib.reload(envs)
517 + mp.setattr(envs, "VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS", use_fse)
518 + rocm_aiter_ops.refresh_env_variables()
519 + aiter_fse_enabled = bool(rocm_aiter_ops.is_fusion_moe_shared_experts_enabled())
520 + # These AMD-specific models currently use the raw environment flag.
521 + # and do not rely on `rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()`.
522 + fse_enabled = use_fse and (
523 + aiter_fse_enabled or model_type in {"deepseek_v4", "minimax_m3"}
524 + )
525 +
526 + with (
527 + patch.object(warning_logger, "warning") as warning,
528 + set_current_vllm_config(vllm_config),
529 + ):
530 + model = model_constructor(vllm_config=vllm_config)
531 + mtp = None
532 + if model_type == "deepseek_v4" and use_fse and not exclude:
533 + from vllm.models.deepseek_v4.amd.mtp import (
534 + DeepSeekV4MTP,
535 + )
536 +
537 + vllm_config.speculative_config = SimpleNamespace(
538 + draft_model_config=SimpleNamespace(hf_config=config)
539 + )
540 + mtp = DeepSeekV4MTP(vllm_config=vllm_config)
541 + assert model.is_fused_shared_expert_enabled is (fse_enabled and not exclude)
542 +
543 + # The dummy quant config here uses mixed mxfp4/fp8 for experts/shared_expert
544 + # so should just raise a warning.
545 + if mtp is not None:
546 + assert not mtp.model.layers[
547 + "1"
548 + ].mtp_block.ffn.is_fused_shared_expert_enabled
549 + warning.assert_called_once()
550 + assert (
551 + "DeepSeek-V4 shared experts at mtp.0.ffn.shared_experts"
552 + in (warning.call_args.args[1])
553 + )
554 + if aiter_fse_enabled and exclude:
555 + warning.assert_called_once()
556 + assert (
557 + "VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS is enabled"
558 + in (warning.call_args.args[0])
559 + )
560 + assert "excludes shared experts" in warning.call_args.args[1]
561 +
562 + importlib.reload(envs)
563 + rocm_aiter_ops.refresh_env_variables()
564 +
565 +
566 +@pytest.mark.parametrize(
567 + ("exclude", "expected"),
568 + [([], True), (["*.shared_expert.*"], False)],
569 +)
570 +def test_quark_shared_expert_fse_compatibility(
571 + exclude: list[str], expected: bool
572 +) -> None:
573 + compatible, reason = is_shared_expert_quant_fse_compatible(
574 + QuarkConfig(
575 + {
576 + "exclude": exclude,
577 + "global_quant_config": {},
578 + "layer_quant_config": {},
579 + }
580 + ),
581 + "model.layers.0.mlp.experts",
582 + "model.layers.0.mlp.shared_expert",
583 + )
584 +
585 + assert compatible is expected
586 + if expected:
587 + assert reason is None
588 + else:
589 + assert (
590 + reason
591 + == "Quark excludes shared experts at model.layers.0.mlp.shared_expert"
592 + )
593 +
594 +
595 +def test_quark_shared_expert_fse_requires_matching_layer_quant_configs() -> None:
596 + global_quant_config = {"weight": {"dtype": "fp4"}}
597 + quant_config = QuarkConfig(
598 + {
599 + "exclude": [],
600 + "global_quant_config": global_quant_config,
601 + "layer_quant_config": {
602 + "model.layers.0.mlp.shared_expert.gate_up_proj": {
603 + "weight": {"dtype": "fp8"}
604 + },
605 + "model.layers.0.mlp.shared_expert.down_proj": {
606 + "weight": {"dtype": "fp8"}
607 + },
608 + },
609 + }
610 + )
611 +
612 + compatible, reason = is_shared_expert_quant_fse_compatible(
613 + quant_config,
614 + "model.layers.0.mlp.experts",
615 + "model.layers.0.mlp.shared_expert",
616 + )
617 +
618 + assert not compatible
619 + assert reason == (
620 + "Quark uses different quantization configurations for routed and "
621 + "shared experts at model.layers.0.mlp.shared_expert"
622 + )
623 +
624 +
625 +def test_quark_shared_expert_fse_rejects_partial_packed_projection_override() -> None:
626 + fp4_config = {"weight": {"dtype": "fp4"}}
627 + fp8_config = {"weight": {"dtype": "fp8"}}
628 + quant_config = QuarkConfig(
629 + {
630 + "exclude": [],
631 + "global_quant_config": fp4_config,
632 + "layer_quant_config": {
633 + "model.layers.0.mlp.experts": fp8_config,
634 + "model.layers.0.mlp.shared_expert.gate_proj": fp8_config,
635 + "model.layers.0.mlp.shared_expert.down_proj": fp8_config,
636 + },
637 + }
638 + )
639 + quant_config.packed_modules_mapping = {"gate_up_proj": ["gate_proj", "up_proj"]}
640 +
641 + compatible, reason = is_shared_expert_quant_fse_compatible(
642 + quant_config,
643 + "model.layers.0.mlp.experts",
644 + "model.layers.0.mlp.shared_expert",
645 + )
646 +
647 + assert not compatible
648 + assert reason == (
649 + "Quark uses different quantization configurations for routed and "
650 + "shared experts at model.layers.0.mlp.shared_expert"
651 + )
652 +
653 +
654 +def test_quark_layer_config_from_name_checks_packed_projections() -> None:
655 + quant_config = QuarkConfig(
656 + {
657 + "global_quant_config": {"weight": {"dtype": "fp4"}},
658 + "layer_quant_config": {
659 + "model.layers.0.mlp.shared_expert.w1": {"weight": {"dtype": "fp4"}},
660 + "model.layers.0.mlp.shared_expert.w3": {"weight": {"dtype": "fp8"}},
661 + },
662 + }
663 + )
664 + quant_config.packed_modules_mapping = {"gate_up_proj": ["w1", "w3"]}
665 +
666 + with pytest.raises(ValueError, match="requires all to use the same scheme"):
667 + quant_config.get_layer_quant_config_from_name(
668 + "model.layers.0.mlp.shared_expert.gate_up_proj"
669 + )
670 +
671 +
672 +def test_quark_packed_layer_config_must_match_global_config() -> None:
673 + quant_config = QuarkConfig(
674 + {
675 + "global_quant_config": {"weight": {"dtype": "fp4"}},
676 + "layer_type_quant_config": {},
677 + "layer_quant_config": {
678 + "model.layers.0.mlp.shared_expert.w1": {"weight": {"dtype": "fp8"}},
679 + },
680 + }
681 + )
682 + quant_config.packed_modules_mapping = {"gate_up_proj": ["w1", "w3"]}
683 +
684 + with pytest.raises(ValueError, match="requires all to use the same scheme"):
685 + quant_config._find_matched_config(
686 + "model.layers.0.mlp.shared_expert.gate_up_proj", nn.Module()
687 + )
688 +
689 +
690 +def test_non_quark_shared_expert_fse_is_incompatible() -> None:
691 + compatible, reason = is_shared_expert_quant_fse_compatible(
692 + object(),
693 + "model.layers.0.mlp.experts",
694 + "model.layers.0.mlp.shared_experts",
695 + )
696 +
697 + assert not compatible
698 + assert reason == (
699 + "shared-expert FSE quantization compatibility is not implemented for object"
700 + )
vllm/model_executor/layers/fused_moe/layer.py
+5 -14
@@ -6,7 +6,6 @@ from typing import Any
6
7 import torch
8
9 -import vllm.envs as envs
9 from vllm._aiter_ops import rocm_aiter_ops
10 from vllm.config import ParallelConfig, get_current_vllm_config
11 from vllm.distributed import (
@@ -74,23 +73,13 @@ def determine_expert_counts(
73 num_experts: int,
74 num_redundant_experts: int,
75 n_shared_experts: int | None,
77 - is_act_and_mul: bool,
76 + fuse_shared_experts: bool,
77 ) -> tuple[int, int, int]:
78 global_num_experts = num_experts + num_redundant_experts
79 logical_num_experts = num_experts
81 - # Shared-expert fusion: append the shared expert(s) as routed-expert slots
82 - # so they run in the same grouped GEMM. Gated by
83 - # VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS: either the native aiter fused-MoE
84 - # path (env + master switch, via is_fusion_moe_shared_experts_enabled) or the
85 - # backend-neutral router-append path (env alone, independent of the master
86 - # switch; e.g. the MM3 triton/flydsl mxfp8 MoE). Gated activations only.
87 - fuse_shared_enabled = (
88 - rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()
89 - or envs.VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS
90 - ) and is_act_and_mul
80
81 num_fused_shared_experts = (
93 - n_shared_experts if n_shared_experts is not None and fuse_shared_enabled else 0
82 + n_shared_experts if n_shared_experts is not None and fuse_shared_experts else 0
83 )
84
85 return global_num_experts, logical_num_experts, num_fused_shared_experts
@@ -132,6 +121,7 @@ def FusedMoEFactory(
121 ckpt_names: tuple[str, str, str] = ("gate_proj", "down_proj", "up_proj"),
122 is_fused_checkpoint_transposed: bool = False,
123 n_shared_experts: int | None = None,
124 + fuse_shared_experts: bool = False,
125 router_logits_dtype: torch.dtype | None = None,
126 gate: torch.nn.Module | None = None,
127 shared_experts: torch.nn.Module | None = None,
@@ -198,6 +188,7 @@ def FusedMoEFactory(
188 block scales use transposed storage.
189 n_shared_experts: Number of shared experts to fuse into the routed
190 grouped GEMM (ROCm; requires aiter FSE or the router-append path)
191 + fuse_shared_experts: Whether to enable shared-expert fusion.
192 router_logits_dtype: Data type for router logits buffers
193 gate: Pre-configured gate module
194 shared_experts: Pre-configured shared experts module
@@ -244,7 +235,7 @@ def FusedMoEFactory(
235 num_experts,
236 num_redundant_experts,
237 n_shared_experts,
247 - is_act_and_mul,
238 + fuse_shared_experts,
239 )
240 )
241
vllm/model_executor/layers/fused_moe/utils.py
+85
@@ -1,15 +1,21 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 import functools
4 +from collections.abc import Iterable
5 from math import prod
6 from typing import TYPE_CHECKING
7
8 import torch
9 import torch.nn.functional as F
10 +from torch import nn
11
12 import vllm.envs as envs
13 from vllm import _custom_ops as ops
14 +from vllm._aiter_ops import rocm_aiter_ops
15 from vllm.logger import init_logger
16 +from vllm.model_executor.layers.quantization.utils.config_utils import (
17 + is_shared_expert_quant_fse_compatible,
18 +)
19 from vllm.model_executor.layers.quantization.utils.fp8_utils import (
20 per_token_group_quant_fp8,
21 )
@@ -34,16 +40,95 @@ from vllm.model_executor.layers.quantization.utils.nvfp4_emulation_utils import
40 from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
41 per_tensor_dequantize,
42 )
43 +from vllm.model_executor.models.utils import PPMissingLayer
44 from vllm.platforms import current_platform
45 from vllm.triton_utils import tl, triton
46 from vllm.utils.math_utils import cdiv
47
48 if TYPE_CHECKING:
49 from vllm.model_executor.layers.fused_moe.config import FusedMoEConfig
50 + from vllm.model_executor.layers.quantization import QuantizationConfig
51
52 logger = init_logger(__name__)
53
54
55 +def resolve_layer_fused_shared_expert(
56 + quant_config: "QuantizationConfig | None",
57 + prefix: str,
58 + shared_expert_name: str = "shared_experts",
59 +) -> bool:
60 + """Resolve whether AITER fused shared-expert execution is enabled.
61 +
62 + Args:
63 + quant_config: Model quantization configuration.
64 + prefix: MoE module prefix.
65 + shared_expert_name: Shared-expert module name under ``prefix``.
66 +
67 + Returns:
68 + Whether AITER fused shared experts are enabled.
69 +
70 + Raises:
71 + ValueError: If requested shared-expert fusion is quantization-incompatible.
72 + """
73 + # NOTE: is_fusion_moe_shared_experts_enabled is decorated with @if_aiter_supported
74 + # that returns None if AITER is not available.
75 + fse_requested = bool(rocm_aiter_ops.is_fusion_moe_shared_experts_enabled())
76 + fse_compatible, fse_reason = (
77 + is_shared_expert_quant_fse_compatible(
78 + quant_config,
79 + f"{prefix}.experts",
80 + f"{prefix}.{shared_expert_name}",
81 + )
82 + if fse_requested
83 + else (True, None)
84 + )
85 + is_fused_shared_expert_enabled = fse_requested and fse_compatible
86 + if fse_requested and not is_fused_shared_expert_enabled:
87 + logger.warning(
88 + "VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS is enabled but "
89 + "cannot be enabled: %s.",
90 + fse_reason,
91 + )
92 + return is_fused_shared_expert_enabled
93 +
94 +
95 +def is_model_fused_shared_expert_compatible(
96 + layers: nn.ModuleList | Iterable[nn.Module],
97 + moe_cls: type[nn.Module],
98 + moe_name: str,
99 +) -> bool:
100 + """Resolve one fused-shared-expert state for a model's MoE layers."""
101 +
102 + def get_moe_layer(layer: nn.Module) -> nn.Module | None:
103 + for name in moe_name.split("."):
104 + layer = getattr(layer, name, None)
105 + if layer is None:
106 + return None
107 + return layer
108 +
109 + moe_layers = (
110 + moe_layer
111 + for layer in layers
112 + if not isinstance(layer, PPMissingLayer)
113 + and (moe_layer := get_moe_layer(layer)) is not None
114 + and isinstance(moe_layer, moe_cls)
115 + )
116 +
117 + enabled = [
118 + getattr(layer, "is_fused_shared_expert_enabled", False) for layer in moe_layers
119 + ]
120 + enabled_count = sum(enabled)
121 + disabled_count = len(enabled) - enabled_count
122 + if enabled_count > 0 and disabled_count > 0:
123 + raise NotImplementedError(
124 + "Fused shared experts must be enabled for all MoE layers; found "
125 + f"{enabled_count} enabled and {disabled_count} disabled layers. "
126 + "Per-layer fused shared experts is not yet supported. Please open "
127 + "an issue."
128 + )
129 + return enabled_count > 0 and disabled_count == 0
130 +
131 +
132 @triton.jit
133 def _count_expert_num_tokens(
134 topk_ids_ptr,
vllm/model_executor/layers/quantization/quark/quark.py
+65 -37
@@ -576,63 +576,91 @@ class QuarkConfig(QuantizationConfig):
576
577 return True
578
579 - def _find_matched_config(
580 - self, layer_name: str, module: torch.nn.Module
581 - ) -> dict[str, Any]:
579 + def get_layer_quant_config_from_name(
580 + self, layer_name: str
581 + ) -> dict[str, Any] | None:
582 proj_name = layer_name.split(".")[-1]
583 if proj_name in self.packed_modules_mapping:
584 shard_proj_names = self.packed_modules_mapping[proj_name]
585 -
586 - # Convert fused_name --> [shard_names]
587 - shard_names = [
588 - layer_name.replace(proj_name, shard_proj_name)
589 - for shard_proj_name in shard_proj_names
590 - ]
591 -
585 shard_configs = []
593 - for shard_name in shard_names:
594 - if shard_name == layer_name:
595 - config = cast(
596 - dict[str, Any], self.quant_config.get("global_quant_config")
597 - )
586 + for shard_proj_name in shard_proj_names:
587 + shard_name = layer_name.replace(proj_name, shard_proj_name)
588 + if shard_name != layer_name:
589 + config = self.get_layer_quant_config_from_name(shard_name)
590 else:
599 - config = self._find_matched_config(shard_name, module)
591 + config = None
592 shard_configs.append(config)
593
602 - if not all(
603 - deep_compare(q_config, shard_configs[0]) for q_config in shard_configs
594 + matched_configs = [config for config in shard_configs if config is not None]
595 + if matched_configs and not all(
596 + deep_compare(config, matched_configs[0]) for config in matched_configs
597 ):
598 raise ValueError(
599 f"Found a different quantization configuration for "
600 f"{shard_proj_names} in {layer_name}. vLLM "
601 "requires all to use the same scheme."
602 )
610 - return shard_configs[0]
603 + if matched_configs:
604 + return matched_configs[0]
605 + return None
606 else:
607 layer_quant_config = cast(
613 - dict[str, Any], self.quant_config.get("layer_quant_config")
608 + dict[str, Any], self.quant_config.get("layer_quant_config") or {}
609 )
615 -
616 - def _matches_pattern(layer_name, pattern):
617 - if "*" not in pattern:
618 - return layer_name in pattern
619 - return fnmatch.fnmatch(layer_name, pattern)
620 -
610 for name_pattern, config in layer_quant_config.items():
622 - if _matches_pattern(layer_name, name_pattern):
611 + if "*" not in name_pattern:
612 + matches = layer_name in name_pattern
613 + else:
614 + matches = fnmatch.fnmatch(layer_name, name_pattern)
615 + if matches:
616 return config
617 + return None
618
625 - layer_type = cast(str, type(module))
626 - layer_type_quant_config = cast(
627 - dict[str, Any], self.quant_config.get("layer_type_quant_config")
628 - )
629 - if layer_type in layer_type_quant_config:
630 - return layer_type_quant_config[layer_type]
619 + def _find_matched_config(
620 + self, layer_name: str, module: torch.nn.Module
621 + ) -> dict[str, Any]:
622 + # Priority order:
623 + # 1. layer_quant_config,
624 + # 2. layer_type_quant_config,
625 + # 3. global_quant_config.
626 +
627 + layer_type = cast(str, type(module))
628 + layer_type_quant_config = cast(
629 + dict[str, Any], self.quant_config.get("layer_type_quant_config")
630 + )
631 + global_quant_config = cast(
632 + dict[str, Any], self.quant_config.get("global_quant_config")
633 + )
634 + fallback_config = layer_type_quant_config.get(layer_type, global_quant_config)
635
632 - global_quant_config = cast(
633 - dict[str, Any], self.quant_config.get("global_quant_config")
634 - )
635 - return global_quant_config
636 + proj_name = layer_name.split(".")[-1]
637 + if proj_name in self.packed_modules_mapping:
638 + shard_proj_names = self.packed_modules_mapping[proj_name]
639 + shard_configs = []
640 + for shard_proj_name in shard_proj_names:
641 + shard_name = layer_name.replace(proj_name, shard_proj_name)
642 + if shard_name == layer_name:
643 + config = fallback_config
644 + else:
645 + config = self.get_layer_quant_config_from_name(shard_name)
646 + if config is None:
647 + config = fallback_config
648 + shard_configs.append(config)
649 +
650 + if not all(
651 + deep_compare(config, shard_configs[0]) for config in shard_configs
652 + ):
653 + raise ValueError(
654 + f"Found a different quantization configuration for "
655 + f"{shard_proj_names} in {layer_name}. vLLM requires all "
656 + "to use the same scheme."
657 + )
658 + return shard_configs[0]
659 + else:
660 + layer_quant_config = self.get_layer_quant_config_from_name(layer_name)
661 + if layer_quant_config is not None:
662 + return layer_quant_config
663 + return fallback_config
664
665 def _get_scheme_from_config(
666 self, config: dict[str, Any], dynamic_mxfp4_quant: bool = False
vllm/model_executor/layers/quantization/utils/config_utils.py new
+155
@@ -0,0 +1,155 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from typing import TYPE_CHECKING
5 +
6 +import regex as re
7 +
8 +if TYPE_CHECKING:
9 + from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
10 +
11 +
12 +def is_shared_expert_quant_fse_compatible(
13 + quant_config: "QuantizationConfig | None",
14 + expert_prefix: str,
15 + shared_expert_prefix: str,
16 + projection_names: list[str] | None = None,
17 +) -> tuple[bool, str | None]:
18 + """Check whether quantization permits fused shared-expert execution.
19 +
20 + Args:
21 + quant_config: Model quantization configuration.
22 + expert_prefix: Routed-expert module prefix.
23 + shared_expert_prefix: Shared-expert module prefix.
24 + projection_names: Shared-expert projection names.
25 +
26 + Returns:
27 + A compatibility flag and, when incompatible, the reason.
28 + """
29 + if projection_names is None:
30 + projection_names = ["gate_up_proj", "down_proj"]
31 +
32 + if quant_config is None:
33 + return True, None
34 +
35 + from vllm.model_executor.layers.quantization.quark.quark import QuarkConfig
36 + from vllm.models.deepseek_v4.quant_config import DeepseekV4FP8Config
37 +
38 + if isinstance(quant_config, DeepseekV4FP8Config):
39 + from vllm.config import get_current_vllm_config
40 + from vllm.model_executor.models.utils import extract_layer_index
41 +
42 + if quant_config.expert_dtype != "fp4":
43 + return False, "DeepSeek-V4 routed experts are not MXFP4"
44 +
45 + hf_config = get_current_vllm_config().model_config.hf_config
46 +
47 + # TODO: This is adapted from former `_shared_experts_are_fp4`, and
48 + # needs to be cleaned up this . There should not be Quark-specific
49 + # logic in DeepseekV4FP8Config.
50 + quantization_config = getattr(hf_config, "quantization_config", None)
51 + if quantization_config is None:
52 + return False, "DeepSeek-V4 has no quantization configuration"
53 +
54 + if not quant_config._is_quark_mxfp4_ocp(quantization_config):
55 + return False, "DeepSeek-v4 FSE is only implemented/tested with Quark MXFP4"
56 +
57 + layer_idx = extract_layer_index(shared_expert_prefix)
58 + if layer_idx >= hf_config.num_hidden_layers:
59 + shared_expert_prefix = (
60 + f"mtp.{layer_idx - hf_config.num_hidden_layers}.ffn.shared_experts"
61 + )
62 + else:
63 + shared_expert_prefix = f"layers.{layer_idx}.ffn.shared_experts"
64 +
65 + if any(
66 + entry.startswith(shared_expert_prefix)
67 + for entry in quantization_config.get("exclude") or []
68 + if isinstance(entry, str)
69 + ):
70 + return (
71 + False,
72 + f"DeepSeek-V4 excludes shared experts at {shared_expert_prefix}",
73 + )
74 +
75 + shared_expert_weight_name = f"{shared_expert_prefix}.w1"
76 + layer_quant_config = quantization_config.get("layer_quant_config") or {}
77 + layer_config = layer_quant_config.get(shared_expert_weight_name)
78 + if layer_config is None:
79 + layer_config = next(
80 + (
81 + config
82 + for pattern, config in layer_quant_config.items()
83 + if isinstance(pattern, str)
84 + and pattern.startswith("re:")
85 + and re.fullmatch(
86 + pattern.removeprefix("re:"), shared_expert_weight_name
87 + )
88 + ),
89 + None,
90 + )
91 + shared_weight_config = (
92 + layer_config or quantization_config.get("global_quant_config") or {}
93 + ).get("weight") or {}
94 + if shared_weight_config.get("dtype") == "fp4":
95 + return True, None
96 + return (
97 + False,
98 + f"DeepSeek-V4 shared experts at {shared_expert_prefix} are not MXFP4",
99 + )
100 +
101 + if isinstance(quant_config, QuarkConfig):
102 + # TODO: layer_type_quant_config is not taken into account here.
103 + assert "exclude" in quant_config.quant_config
104 + assert "global_quant_config" in quant_config.quant_config
105 +
106 + is_compatible = not any(
107 + "shared_expert" in str(entry)
108 + for entry in quant_config.quant_config["exclude"]
109 + )
110 + if not is_compatible:
111 + return False, f"Quark excludes shared experts at {shared_expert_prefix}"
112 +
113 + global_quant_config = quant_config.quant_config["global_quant_config"]
114 +
115 + def get_projection_quant_configs(layer_name: str) -> list[object]:
116 + module_prefix, _, projection_name = layer_name.rpartition(".")
117 + packed_projection_names = quant_config.packed_modules_mapping.get(
118 + projection_name, [projection_name]
119 + )
120 + return [
121 + quant_config.get_layer_quant_config_from_name(
122 + f"{module_prefix}.{packed_projection_name}"
123 + )
124 + or global_quant_config
125 + for packed_projection_name in packed_projection_names
126 + ]
127 +
128 + expert_quant_config = (
129 + quant_config.get_layer_quant_config_from_name(expert_prefix)
130 + or global_quant_config
131 + )
132 + shared_expert_quant_configs = [
133 + config
134 + for projection_name in projection_names
135 + for config in get_projection_quant_configs(
136 + f"{shared_expert_prefix}.{projection_name}"
137 + )
138 + ]
139 + if all(config == expert_quant_config for config in shared_expert_quant_configs):
140 + return True, None
141 + return (
142 + False,
143 + "Quark uses different quantization configurations for routed and "
144 + f"shared experts at {shared_expert_prefix}",
145 + )
146 +
147 + # TODO: Extend FSE support detection to other quantization methods. Typically,
148 + # one would check that the experts and shared_experts use the same
149 + # quantization config. This may be refactored as part of QuantizationConfig later.
150 +
151 + return (
152 + False,
153 + "shared-expert FSE quantization compatibility is not implemented for "
154 + f"{type(quant_config).__name__}",
155 + )
vllm/model_executor/models/AXK1.py
+22 -10
@@ -46,6 +46,10 @@ from vllm.model_executor.layers.fused_moe import (
46 FusedMoEFactory,
47 fused_moe_make_expert_params_mapping,
48 )
49 +from vllm.model_executor.layers.fused_moe.utils import (
50 + is_model_fused_shared_expert_compatible,
51 + resolve_layer_fused_shared_expert,
52 +)
53 from vllm.model_executor.layers.layernorm import RMSNorm
54 from vllm.model_executor.layers.linear import (
55 ColumnParallelLinear,
@@ -148,10 +152,14 @@ class AXK1MoE(nn.Module):
152 self.n_local_physical_experts = self.n_physical_experts // self.ep_size
153
154 self.is_rocm_aiter_moe_enabled = rocm_aiter_ops.is_fused_moe_enabled()
151 - self.is_fusion_moe_shared_experts_enabled = (
152 - rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()
153 - )
154 - if config.n_shared_experts is None or self.is_fusion_moe_shared_experts_enabled:
155 +
156 + self.is_fused_shared_expert_enabled = False
157 + if config.n_shared_experts is not None:
158 + self.is_fused_shared_expert_enabled = resolve_layer_fused_shared_expert(
159 + quant_config, prefix
160 + )
161 +
162 + if config.n_shared_experts is None or self.is_fused_shared_expert_enabled:
163 self.shared_experts = None
164 else:
165 intermediate_size = config.moe_intermediate_size * config.n_shared_experts
@@ -189,8 +197,9 @@ class AXK1MoE(nn.Module):
197 num_redundant_experts=self.n_redundant_experts,
198 is_sequence_parallel=self.is_sequence_parallel,
199 n_shared_experts=config.n_shared_experts
192 - if self.is_fusion_moe_shared_experts_enabled
200 + if self.is_fused_shared_expert_enabled
201 else None,
202 + fuse_shared_experts=self.is_fused_shared_expert_enabled,
203 )
204
205 def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
@@ -725,6 +734,12 @@ class AXK1Model(nn.Module):
734 prefix=f"{prefix}.layers",
735 )
736
737 + self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible(
738 + self.layers,
739 + AXK1MoE,
740 + "mlp",
741 + )
742 +
743 if get_pp_group().is_last_rank:
744 self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
745 else:
@@ -789,9 +804,6 @@ class AXK1Model(nn.Module):
804 return hidden_states
805
806 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
792 - rocm_aiter_moe_shared_expert_enabled = (
793 - rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()
794 - )
807 stacked_params_mapping: list[tuple[str, str, int | str]] = [
808 # (param_name, shard_name, shard_id)
809 ("gate_up_proj", "gate_proj", 0),
@@ -822,7 +834,7 @@ class AXK1Model(nn.Module):
834 num_experts=self.config.n_routed_experts
835 + (
836 (self.config.n_shared_experts or 0)
825 - if rocm_aiter_moe_shared_expert_enabled
837 + if self.is_fused_shared_expert_enabled
838 else 0
839 ),
840 num_redundant_experts=self.num_redundant_experts,
@@ -839,7 +851,7 @@ class AXK1Model(nn.Module):
851 continue # skip spec decode layers for main model
852
853 is_fusion_moe_shared_experts_layer = (
842 - rocm_aiter_moe_shared_expert_enabled and ("mlp.shared_experts" in name)
854 + self.is_fused_shared_expert_enabled and ("mlp.shared_experts" in name)
855 )
856
857 for param_name, weight_name, shard_id in stacked_params_mapping:
vllm/model_executor/models/deepseek_mtp.py
+10 -6
@@ -7,13 +7,15 @@ import torch
7 import torch.nn as nn
8 from transformers import PretrainedConfig
9
10 -from vllm._aiter_ops import rocm_aiter_ops
10 from vllm.compilation.decorators import support_torch_compile
11 from vllm.config import VllmConfig
12 from vllm.distributed import tensor_model_parallel_all_gather
13 from vllm.model_executor.layers.fused_moe import (
14 fused_moe_make_expert_params_mapping,
15 )
16 +from vllm.model_executor.layers.fused_moe.utils import (
17 + is_model_fused_shared_expert_compatible,
18 +)
19 from vllm.model_executor.layers.layernorm import RMSNorm
20 from vllm.model_executor.layers.logits_processor import LogitsProcessor
21 from vllm.model_executor.layers.quantization import QuantizationConfig
@@ -255,6 +257,11 @@ class DeepSeekMTP(nn.Module, DeepseekV2MixtureOfExperts):
257 self.moe_mlp_layers.append(layer.mlp)
258 self.moe_layers.append(layer.mlp.experts)
259 self.extract_moe_parameters(example_moe)
260 + self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible(
261 + self.model.layers.values(),
262 + DeepseekV2MoE,
263 + "mtp_block.mlp",
264 + )
265
266 def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
267 return self.model.embed_input_ids(input_ids)
@@ -285,9 +292,6 @@ class DeepSeekMTP(nn.Module, DeepseekV2MixtureOfExperts):
292 return self.model.compute_logits(hidden_states, spec_step_idx)
293
294 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
288 - rocm_aiter_moe_shared_expert_enabled = (
289 - rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()
290 - )
295 stacked_params_mapping = [
296 ("gate_up_proj", "gate_proj", 0),
297 ("gate_up_proj", "up_proj", 1),
@@ -310,7 +314,7 @@ class DeepSeekMTP(nn.Module, DeepseekV2MixtureOfExperts):
314 num_experts=self.config.n_routed_experts
315 + (
316 self.config.n_shared_experts
313 - if rocm_aiter_moe_shared_expert_enabled
317 + if self.is_fused_shared_expert_enabled
318 else 0
319 ),
320 )
@@ -326,7 +330,7 @@ class DeepSeekMTP(nn.Module, DeepseekV2MixtureOfExperts):
330 if spec_layer is None:
331 continue
332 is_fusion_moe_shared_experts_layer = (
329 - rocm_aiter_moe_shared_expert_enabled and ("mlp.shared_experts" in name)
333 + self.is_fused_shared_expert_enabled and ("mlp.shared_experts" in name)
334 )
335 name = self._rewrite_spec_layer_name(spec_layer, name)
336
vllm/model_executor/models/deepseek_v2.py
+22 -10
@@ -54,6 +54,10 @@ from vllm.model_executor.layers.fused_moe import (
54 GateLinear,
55 fused_moe_make_expert_params_mapping,
56 )
57 +from vllm.model_executor.layers.fused_moe.utils import (
58 + is_model_fused_shared_expert_compatible,
59 + resolve_layer_fused_shared_expert,
60 +)
61 from vllm.model_executor.layers.layernorm import LayerNorm, RMSNorm
62 from vllm.model_executor.layers.linear import (
63 ColumnParallelLinear,
@@ -329,9 +333,13 @@ class DeepseekV2MoE(nn.Module):
333 self.n_local_physical_experts = self.n_physical_experts // self.ep_size
334
335 self.is_rocm_aiter_moe_enabled = rocm_aiter_ops.is_fused_moe_enabled()
332 - self.is_fusion_moe_shared_experts_enabled = (
333 - rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()
334 - )
336 +
337 + self.is_fused_shared_expert_enabled = False
338 + if config.n_shared_experts is not None:
339 + self.is_fused_shared_expert_enabled = resolve_layer_fused_shared_expert(
340 + quant_config, prefix
341 + )
342 +
343 if (
344 self.is_rocm_aiter_moe_enabled
345 and self.gate.e_score_correction_bias is not None
@@ -340,7 +348,7 @@ class DeepseekV2MoE(nn.Module):
348 # Accumulates in fp32; avoids bf16->fp32 cast.
349 self.gate.set_out_dtype(self.gate.weight.dtype)
350
343 - if config.n_shared_experts is None or self.is_fusion_moe_shared_experts_enabled:
351 + if config.n_shared_experts is None or self.is_fused_shared_expert_enabled:
352 self.shared_experts = None
353 else:
354 intermediate_size = config.moe_intermediate_size * config.n_shared_experts
@@ -377,8 +385,9 @@ class DeepseekV2MoE(nn.Module):
385 is_sequence_parallel=self.is_sequence_parallel,
386 reduce_results=reduce_results,
387 n_shared_experts=config.n_shared_experts
380 - if self.is_fusion_moe_shared_experts_enabled
388 + if self.is_fused_shared_expert_enabled
389 else None,
390 + fuse_shared_experts=self.is_fused_shared_expert_enabled,
391 router_logits_dtype=self.gate.out_dtype,
392 )
393
@@ -1401,6 +1410,12 @@ class DeepseekV2Model(nn.Module):
1410 prefix=f"{prefix}.layers",
1411 )
1412
1413 + self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible(
1414 + self.layers,
1415 + DeepseekV2MoE,
1416 + "mlp",
1417 + )
1418 +
1419 if get_pp_group().is_last_rank:
1420 self.norm = RMSNorm(self.hidden_size, eps=config.rms_norm_eps)
1421 else:
@@ -1511,9 +1526,6 @@ class DeepseekV2Model(nn.Module):
1526 return hidden_states
1527
1528 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
1514 - rocm_aiter_moe_shared_expert_enabled = (
1515 - rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()
1516 - )
1529 stacked_params_mapping: list[tuple[str, str, int | str]] = [
1530 # (param_name, shard_name, shard_id)
1531 ("gate_up_proj", "gate_proj", 0),
@@ -1554,7 +1566,7 @@ class DeepseekV2Model(nn.Module):
1566 num_experts=self.config.n_routed_experts
1567 + (
1568 self.config.n_shared_experts
1557 - if rocm_aiter_moe_shared_expert_enabled
1569 + if self.is_fused_shared_expert_enabled
1570 else 0
1571 ),
1572 num_redundant_experts=self.num_redundant_experts,
@@ -1582,7 +1594,7 @@ class DeepseekV2Model(nn.Module):
1594 continue # this layer has no indexer; drop its checkpoint weights
1595
1596 is_fusion_moe_shared_experts_layer = (
1585 - rocm_aiter_moe_shared_expert_enabled and ("mlp.shared_experts" in name)
1597 + self.is_fused_shared_expert_enabled and ("mlp.shared_experts" in name)
1598 )
1599
1600 if _try_load_fp8_indexer_wk(
vllm/model_executor/models/glm4_moe.py
+20 -12
@@ -39,12 +39,15 @@ from vllm.distributed import (
39 get_pp_group,
40 get_tensor_model_parallel_world_size,
41 )
42 -from vllm.logger import init_logger
42 from vllm.model_executor.layers.activation import SiluAndMul
43 from vllm.model_executor.layers.attention import Attention
44 from vllm.model_executor.layers.fused_moe import (
45 FusedMoEFactory,
46 )
47 +from vllm.model_executor.layers.fused_moe.utils import (
48 + is_model_fused_shared_expert_compatible,
49 + resolve_layer_fused_shared_expert,
50 +)
51 from vllm.model_executor.layers.layernorm import RMSNorm
52 from vllm.model_executor.layers.linear import (
53 MergedColumnParallelLinear,
@@ -72,8 +75,6 @@ from .utils import (
75 skip_spec_layers,
76 )
77
75 -logger = init_logger(__name__)
76 -
78
79 class Glm4MoeMLP(nn.Module):
80 def __init__(
@@ -159,14 +160,15 @@ class Glm4MoE(nn.Module):
160 self.n_physical_experts = self.n_logical_experts + self.n_redundant_experts
161 self.n_local_physical_experts = self.n_physical_experts // self.ep_size
162
162 - # AITER fused shared-expert (FSE) gate; mirrors the deepseek_v2.py
163 - # pattern (see Glm4MoE / MoERunner wiring there).
163 self.is_rocm_aiter_moe_enabled = rocm_aiter_ops.is_fused_moe_enabled()
165 - self.is_fusion_moe_shared_experts_enabled = (
166 - rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()
167 - )
164
169 - if config.n_shared_experts is None or self.is_fusion_moe_shared_experts_enabled:
165 + self.is_fused_shared_expert_enabled = False
166 + if config.n_shared_experts is not None:
167 + self.is_fused_shared_expert_enabled = resolve_layer_fused_shared_expert(
168 + quant_config, prefix
169 + )
170 +
171 + if config.n_shared_experts is None or self.is_fused_shared_expert_enabled:
172 self.shared_experts = None
173 else:
174 intermediate_size = config.moe_intermediate_size * config.n_shared_experts
@@ -200,10 +202,9 @@ class Glm4MoE(nn.Module):
202 num_redundant_experts=self.n_redundant_experts,
203 router_logits_dtype=torch.float32,
204 n_shared_experts=(
203 - config.n_shared_experts
204 - if self.is_fusion_moe_shared_experts_enabled
205 - else None
205 + config.n_shared_experts if self.is_fused_shared_expert_enabled else None
206 ),
207 + fuse_shared_experts=self.is_fused_shared_expert_enabled,
208 )
209
210 def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
@@ -443,6 +444,12 @@ class Glm4MoeModel(nn.Module):
444 prefix=f"{prefix}.layers",
445 )
446
447 + self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible(
448 + self.layers,
449 + Glm4MoE,
450 + "mlp",
451 + )
452 +
453 if get_pp_group().is_last_rank:
454 self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
455 else:
@@ -488,6 +495,7 @@ class Glm4MoeModel(nn.Module):
495 skip_spec_layers(weights, self.config),
496 n_routed_experts=self.config.n_routed_experts,
497 n_shared_experts=self.config.n_shared_experts or 1,
498 + enabled=self.is_fused_shared_expert_enabled,
499 )
500 loader = AutoWeightsLoader(self)
501 return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
vllm/model_executor/models/glm4_moe_lite.py
+11 -6
@@ -34,7 +34,6 @@ from torch import nn
34 if TYPE_CHECKING:
35 from transformers.models.glm4_moe_lite import Glm4MoeLiteConfig
36
37 -from vllm._aiter_ops import rocm_aiter_ops
37 from vllm.compilation.decorators import support_torch_compile
38 from vllm.config import VllmConfig
39 from vllm.distributed import (
@@ -44,6 +43,9 @@ from vllm.logger import init_logger
43 from vllm.model_executor.layers.fused_moe import (
44 fused_moe_make_expert_params_mapping,
45 )
46 +from vllm.model_executor.layers.fused_moe.utils import (
47 + is_model_fused_shared_expert_compatible,
48 +)
49 from vllm.model_executor.layers.layernorm import RMSNorm
50 from vllm.model_executor.layers.logits_processor import LogitsProcessor
51 from vllm.model_executor.layers.vocab_parallel_embedding import (
@@ -254,6 +256,12 @@ class Glm4MoeLiteModel(nn.Module):
256 prefix=f"{prefix}.layers",
257 )
258
259 + self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible(
260 + self.layers,
261 + Glm4MoeLite,
262 + "mlp",
263 + )
264 +
265 if get_pp_group().is_last_rank:
266 self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
267 else:
@@ -320,9 +328,6 @@ class Glm4MoeLiteModel(nn.Module):
328 )
329
330 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
323 - rocm_aiter_moe_shared_expert_enabled = (
324 - rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()
325 - )
331 stacked_params_mapping = [
332 # (param_name, shard_name, shard_id)
333 ("gate_up_proj", "gate_proj", 0),
@@ -345,7 +350,7 @@ class Glm4MoeLiteModel(nn.Module):
350 num_experts=self.config.n_routed_experts
351 + (
352 self.config.n_shared_experts
348 - if rocm_aiter_moe_shared_expert_enabled
353 + if self.is_fused_shared_expert_enabled
354 else 0
355 ),
356 )
@@ -361,7 +366,7 @@ class Glm4MoeLiteModel(nn.Module):
366 continue # skip spec decode layers for main model
367
368 is_fusion_moe_shared_experts_layer = (
364 - rocm_aiter_moe_shared_expert_enabled and ("mlp.shared_experts" in name)
369 + self.is_fused_shared_expert_enabled and ("mlp.shared_experts" in name)
370 )
371
372 for param_name, weight_name, shard_id in stacked_params_mapping:
vllm/model_executor/models/glm4_moe_lite_mtp.py
+10 -6
@@ -30,12 +30,14 @@ import torch
30 import torch.nn as nn
31 from transformers import PretrainedConfig
32
33 -from vllm._aiter_ops import rocm_aiter_ops
33 from vllm.config import VllmConfig
34 from vllm.model_executor.layers.fused_moe import (
35 MoERunner,
36 fused_moe_make_expert_params_mapping,
37 )
38 +from vllm.model_executor.layers.fused_moe.utils import (
39 + is_model_fused_shared_expert_compatible,
40 +)
41 from vllm.model_executor.layers.layernorm import RMSNorm
42 from vllm.model_executor.layers.logits_processor import LogitsProcessor
43 from vllm.model_executor.layers.quantization import QuantizationConfig
@@ -224,6 +226,11 @@ class Glm4MoeLiteMTP(nn.Module, SupportsPP, Glm4MixtureOfExperts):
226 self.moe_mlp_layers.append(layer.mlp)
227 self.moe_layers.append(layer.mlp.experts)
228 self.extract_moe_parameters(example_moe)
229 + self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible(
230 + self.model.layers.values(),
231 + Glm4MoeLite,
232 + "mtp_block.mlp",
233 + )
234
235 def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
236 return self.model.embed_input_ids(input_ids)
@@ -250,9 +257,6 @@ class Glm4MoeLiteMTP(nn.Module, SupportsPP, Glm4MixtureOfExperts):
257 return self.model.compute_logits(hidden_states, spec_step_idx)
258
259 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
253 - rocm_aiter_moe_shared_expert_enabled = (
254 - rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()
255 - )
260 stacked_params_mapping = [
261 ("gate_up_proj", "gate_proj", 0),
262 ("gate_up_proj", "up_proj", 1),
@@ -268,7 +272,7 @@ class Glm4MoeLiteMTP(nn.Module, SupportsPP, Glm4MixtureOfExperts):
272 num_experts=self.config.n_routed_experts
273 + (
274 self.config.n_shared_experts
271 - if rocm_aiter_moe_shared_expert_enabled
275 + if self.is_fused_shared_expert_enabled
276 else 0
277 ),
278 )
@@ -282,7 +286,7 @@ class Glm4MoeLiteMTP(nn.Module, SupportsPP, Glm4MixtureOfExperts):
286 if spec_layer is None:
287 continue
288 is_fusion_moe_shared_experts_layer = (
285 - rocm_aiter_moe_shared_expert_enabled and ("mlp.shared_experts" in name)
289 + self.is_fused_shared_expert_enabled and ("mlp.shared_experts" in name)
290 )
291 name = self._rewrite_spec_layer_name(spec_layer, name)
292 for param_name, weight_name, shard_id in stacked_params_mapping:
vllm/model_executor/models/glm4_moe_mtp.py
+10 -7
@@ -31,12 +31,14 @@ import torch
31 import torch.nn as nn
32 from transformers import PretrainedConfig
33
34 -from vllm._aiter_ops import rocm_aiter_ops
34 from vllm.config import CacheConfig, ParallelConfig, VllmConfig
35 from vllm.model_executor.layers.fused_moe import (
36 MoERunner,
37 fused_moe_make_expert_params_mapping,
38 )
39 +from vllm.model_executor.layers.fused_moe.utils import (
40 + is_model_fused_shared_expert_compatible,
41 +)
42 from vllm.model_executor.layers.layernorm import RMSNorm
43 from vllm.model_executor.layers.logits_processor import LogitsProcessor
44 from vllm.model_executor.layers.quantization import QuantizationConfig
@@ -212,6 +214,11 @@ class Glm4MoeMTP(nn.Module, Glm4MixtureOfExperts):
214 self.moe_mlp_layers.append(layer.mlp)
215 self.moe_layers.append(layer.mlp.experts)
216 self.extract_moe_parameters(example_moe)
217 + self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible(
218 + self.model.layers.values(),
219 + Glm4MoE,
220 + "mtp_block.mlp",
221 + )
222
223 def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
224 return self.model.embed_input_ids(input_ids)
@@ -238,10 +245,6 @@ class Glm4MoeMTP(nn.Module, Glm4MixtureOfExperts):
245 return self.model.compute_logits(hidden_states, spec_step_idx)
246
247 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
241 - # FSE weight loading mirrors glm4_moe.py / deepseek_mtp.py.
242 - rocm_aiter_moe_shared_expert_enabled = (
243 - rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()
244 - )
248 stacked_params_mapping = [
249 # (param_name, shard_name, shard_id)
250 ("qkv_proj", "q_proj", "q"),
@@ -254,7 +257,7 @@ class Glm4MoeMTP(nn.Module, Glm4MixtureOfExperts):
257 # Params for weights, fp8 weight scales, fp8 activation scales
258 # (param_name, weight_name, expert_id, shard_id)
259 num_experts = self.config.n_routed_experts
257 - if rocm_aiter_moe_shared_expert_enabled and self.config.n_shared_experts:
260 + if self.is_fused_shared_expert_enabled and self.config.n_shared_experts:
261 num_experts += self.config.n_shared_experts
262 expert_params_mapping = fused_moe_make_expert_params_mapping(
263 self,
@@ -279,7 +282,7 @@ class Glm4MoeMTP(nn.Module, Glm4MixtureOfExperts):
282 name = self._rewrite_spec_layer_name(spec_layer, name)
283
284 is_fusion_moe_shared_experts_layer = (
282 - rocm_aiter_moe_shared_expert_enabled and ("mlp.shared_experts" in name)
285 + self.is_fused_shared_expert_enabled and ("mlp.shared_experts" in name)
286 )
287
288 for param_name, weight_name, shard_id in stacked_params_mapping:
vllm/model_executor/models/qwen3_5.py
+9 -4
@@ -29,13 +29,15 @@ from collections.abc import Iterable
29 import torch
30 from torch import nn
31
32 -from vllm._aiter_ops import rocm_aiter_ops
32 from vllm.compilation.decorators import support_torch_compile
33 from vllm.config import VllmConfig
34 from vllm.distributed import (
35 get_pp_group,
36 )
37 from vllm.logger import init_logger
38 +from vllm.model_executor.layers.fused_moe.utils import (
39 + is_model_fused_shared_expert_compatible,
40 +)
41 from vllm.model_executor.layers.layernorm import GemmaRMSNorm as Qwen3_5RMSNorm
42 from vllm.model_executor.layers.logits_processor import LogitsProcessor
43 from vllm.model_executor.layers.mamba.gdn.qwen_gdn_linear_attn import (
@@ -78,7 +80,6 @@ from .qwen3_next import (
80 Qwen3NextModel,
81 Qwen3NextSparseMoeBlock,
82 QwenNextMixtureOfExperts,
81 - _is_shared_expert_fse_compatible,
83 )
84 from .qwen3_vl import (
85 Qwen3_VisionTransformer,
@@ -256,6 +257,11 @@ class Qwen3_5Model(Qwen3NextModel):
257 self.start_layer, self.end_layer, self.layers = make_layers(
258 config.num_hidden_layers, get_layer, prefix=f"{prefix}.layers"
259 )
260 + self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible(
261 + self.layers,
262 + Qwen3NextSparseMoeBlock,
263 + "mlp",
264 + )
265 self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
266 ["hidden_states", "residual"], config.hidden_size
267 )
@@ -274,8 +280,7 @@ class Qwen3_5Model(Qwen3NextModel):
280 if "moe" in self.config.model_type:
281 weights = maybe_fuse_shared_experts(
282 weights,
277 - enabled=rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()
278 - and _is_shared_expert_fse_compatible(self.quant_config),
283 + enabled=self.is_fused_shared_expert_enabled,
284 n_routed_experts=self.config.num_experts,
285 n_shared_experts=1,
286 ckpt_prefix="mlp.shared_expert",
vllm/model_executor/models/qwen3_5_mtp.py
+11 -7
@@ -7,11 +7,13 @@ from collections.abc import Iterable
7 import torch
8 from torch import nn
9
10 -from vllm._aiter_ops import rocm_aiter_ops
10 from vllm.compilation.decorators import support_torch_compile
12 -from vllm.config import VllmConfig, get_current_vllm_config
11 +from vllm.config import VllmConfig
12 from vllm.distributed import get_pp_group, tensor_model_parallel_all_gather
13 from vllm.logger import init_logger
14 +from vllm.model_executor.layers.fused_moe.utils import (
15 + is_model_fused_shared_expert_compatible,
16 +)
17 from vllm.model_executor.layers.linear import ColumnParallelLinear
18 from vllm.model_executor.layers.logits_processor import LogitsProcessor
19 from vllm.model_executor.layers.vocab_parallel_embedding import (
@@ -25,8 +27,8 @@ from vllm.model_executor.models.qwen3_5 import (
27 Qwen3_5RMSNorm,
28 )
29 from vllm.model_executor.models.qwen3_next import (
30 + Qwen3NextSparseMoeBlock,
31 QwenNextMixtureOfExperts,
29 - _is_shared_expert_fse_compatible,
32 )
33 from vllm.model_executor.models.utils import sequence_parallel_chunk
34 from vllm.sequence import IntermediateTensors
@@ -122,6 +124,11 @@ class Qwen3_5MultiTokenPredictor(nn.Module):
124 for idx in range(self.num_mtp_layers)
125 )
126 vllm_config.quant_config = original_quant
127 + self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible(
128 + self.layers,
129 + Qwen3NextSparseMoeBlock,
130 + "mlp",
131 + )
132 self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
133 ["hidden_states", "residual"], config.hidden_size
134 )
@@ -185,10 +192,7 @@ class Qwen3_5MultiTokenPredictor(nn.Module):
192 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
193 weights = maybe_fuse_shared_experts(
194 weights,
188 - enabled=rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()
189 - and _is_shared_expert_fse_compatible(
190 - get_current_vllm_config().quant_config
191 - ),
195 + enabled=self.is_fused_shared_expert_enabled,
196 n_routed_experts=getattr(self.config, "num_experts", 0),
197 n_shared_experts=1,
198 ckpt_prefix="mlp.shared_expert",
vllm/model_executor/models/qwen3_next.py
+22 -32
@@ -8,7 +8,6 @@ from itertools import islice
8 import torch
9 from torch import nn
10
11 -from vllm._aiter_ops import rocm_aiter_ops
11 from vllm.compilation.decorators import support_torch_compile
12 from vllm.config import CacheConfig, ModelConfig, VllmConfig
13 from vllm.distributed import (
@@ -18,9 +17,12 @@ from vllm.distributed import (
17 tensor_model_parallel_all_gather,
18 tensor_model_parallel_reduce_scatter,
19 )
21 -from vllm.logger import init_logger
20 from vllm.model_executor.layers.attention import Attention
21 from vllm.model_executor.layers.fused_moe import FusedMoEFactory
22 +from vllm.model_executor.layers.fused_moe.utils import (
23 + is_model_fused_shared_expert_compatible,
24 + resolve_layer_fused_shared_expert,
25 +)
26 from vllm.model_executor.layers.fused_qk_norm_rope import fused_qk_rmsnorm_rope_gate
27 from vllm.model_executor.layers.layernorm import (
28 GemmaRMSNorm as Qwen3NextRMSNorm,
@@ -73,8 +75,6 @@ from .utils import (
75 maybe_prefix,
76 )
77
76 -logger = init_logger(__name__)
77 -
78 KVCache = tuple[torch.Tensor, torch.Tensor]
79
80
@@ -90,26 +90,6 @@ def _should_use_sequence_parallel(vllm_config: VllmConfig) -> bool:
90 )
91
92
93 -def _is_shared_expert_fse_compatible(quant_config) -> bool:
94 - """Check if shared expert can be fused with routed experts.
95 -
96 - FSE requires that shared and routed expert weights use the same
97 - quantization format. Returns False when the shared expert is
98 - excluded from quantization (e.g. float32 shared in an MXFP4 model)
99 - or has a different quant spec than routed experts.
100 - """
101 - if quant_config is None:
102 - return True
103 - # Quark stores its full config dict in quant_config.quant_config
104 - raw_config = getattr(quant_config, "quant_config", None)
105 - if not isinstance(raw_config, dict):
106 - return True
107 - exclude = raw_config.get("exclude", [])
108 - if not exclude:
109 - return True
110 - return not any("shared_expert." in str(e) for e in exclude)
111 -
112 -
93 class Qwen3NextSparseMoeBlock(nn.Module):
94 def __init__(self, vllm_config: VllmConfig, prefix: str = ""):
95 super().__init__()
@@ -157,15 +137,18 @@ class Qwen3NextSparseMoeBlock(nn.Module):
137 prefix=f"{prefix}.shared_expert_gate",
138 )
139
160 - _fse_requested = rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()
161 - _fse_enabled = _fse_requested and _is_shared_expert_fse_compatible(quant_config)
162 - if _fse_requested and not _fse_enabled:
163 - logger.warning(
164 - "VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS is enabled but "
165 - "shared expert has a different quantization spec than routed "
166 - "experts. Falling back to non-fused shared expert path."
140 + self.is_fused_shared_expert_enabled = False
141 + if config.shared_expert_intermediate_size > 0:
142 + self.is_fused_shared_expert_enabled = resolve_layer_fused_shared_expert(
143 + quant_config,
144 + prefix,
145 + shared_expert_name="shared_expert",
146 )
168 - if _fse_enabled or config.shared_expert_intermediate_size <= 0:
147 +
148 + if (
149 + self.is_fused_shared_expert_enabled
150 + or config.shared_expert_intermediate_size <= 0
151 + ):
152 self.shared_expert = None
153 else:
154 self.shared_expert = Qwen3NextMLP(
@@ -193,6 +176,7 @@ class Qwen3NextSparseMoeBlock(nn.Module):
176 num_redundant_experts=self.n_redundant_experts,
177 is_sequence_parallel=self.is_sequence_parallel,
178 n_shared_experts=1 if self.shared_expert is None else None,
179 + fuse_shared_experts=self.is_fused_shared_expert_enabled,
180 shared_expert_gate=self.shared_expert_gate
181 if self.shared_expert is None
182 else None,
@@ -599,6 +583,11 @@ class Qwen3NextModel(nn.Module, EagleModelMixin):
583 self.start_layer, self.end_layer, self.layers = make_layers(
584 config.num_hidden_layers, get_layer, prefix=f"{prefix}.layers"
585 )
586 + self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible(
587 + self.layers,
588 + Qwen3NextSparseMoeBlock,
589 + "mlp",
590 + )
591 self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
592 ["hidden_states", "residual"], config.hidden_size
593 )
@@ -678,6 +667,7 @@ class Qwen3NextModel(nn.Module, EagleModelMixin):
667 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
668 weights = maybe_fuse_shared_experts(
669 weights,
670 + enabled=self.is_fused_shared_expert_enabled,
671 n_routed_experts=getattr(self.config, "num_experts", 0),
672 n_shared_experts=1,
673 ckpt_prefix="mlp.shared_expert",
vllm/model_executor/models/qwen3_next_mtp.py
+10
@@ -11,6 +11,9 @@ from vllm.compilation.decorators import support_torch_compile
11 from vllm.config import VllmConfig
12 from vllm.distributed import get_pp_group, tensor_model_parallel_all_gather
13 from vllm.logger import init_logger
14 +from vllm.model_executor.layers.fused_moe.utils import (
15 + is_model_fused_shared_expert_compatible,
16 +)
17 from vllm.model_executor.layers.linear import ColumnParallelLinear
18 from vllm.model_executor.layers.logits_processor import LogitsProcessor
19 from vllm.model_executor.layers.vocab_parallel_embedding import (
@@ -21,6 +24,7 @@ from vllm.model_executor.models.qwen3_next import (
24 Qwen3NextDecoderLayer,
25 Qwen3NextModel,
26 Qwen3NextRMSNorm,
27 + Qwen3NextSparseMoeBlock,
28 QwenNextMixtureOfExperts,
29 )
30 from vllm.model_executor.models.utils import sequence_parallel_chunk
@@ -91,6 +95,11 @@ class Qwen3NextMultiTokenPredictor(nn.Module):
95 for idx in range(self.num_mtp_layers)
96 )
97
98 + self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible(
99 + self.layers,
100 + Qwen3NextSparseMoeBlock,
101 + "mlp",
102 + )
103 self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
104 ["hidden_states", "residual"], config.hidden_size
105 )
@@ -155,6 +164,7 @@ class Qwen3NextMultiTokenPredictor(nn.Module):
164 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
165 weights = maybe_fuse_shared_experts(
166 weights,
167 + enabled=self.is_fused_shared_expert_enabled,
168 n_routed_experts=self.config.num_experts,
169 n_shared_experts=1,
170 ckpt_prefix="mlp.shared_expert",
vllm/models/deepseek_v32/amd/mtp.py
+10 -6
@@ -8,12 +8,14 @@ from collections.abc import Callable, Iterable
8 import torch
9 import torch.nn as nn
10
11 -from vllm._aiter_ops import rocm_aiter_ops
11 from vllm.config import VllmConfig
12 from vllm.distributed import tensor_model_parallel_all_reduce
13 from vllm.model_executor.layers.fused_moe import (
14 fused_moe_make_expert_params_mapping,
15 )
16 +from vllm.model_executor.layers.fused_moe.utils import (
17 + is_model_fused_shared_expert_compatible,
18 +)
19 from vllm.model_executor.layers.layernorm import RMSNorm
20 from vllm.model_executor.layers.logits_processor import LogitsProcessor
21 from vllm.model_executor.layers.vocab_parallel_embedding import (
@@ -167,6 +169,11 @@ class DeepseekV32MTP(nn.Module, DeepseekV2MixtureOfExperts):
169 vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
170 )
171 self.set_moe_parameters()
172 + self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible(
173 + self.model.layers.values(),
174 + DeepseekV2MoE,
175 + "mtp_block.mlp",
176 + )
177
178 def set_moe_parameters(self):
179 self.num_moe_layers = self.config.num_nextn_predict_layers
@@ -231,9 +238,6 @@ class DeepseekV32MTP(nn.Module, DeepseekV2MixtureOfExperts):
238 return name
239
240 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
234 - rocm_aiter_moe_shared_expert_enabled = (
235 - rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()
236 - )
241 stacked_params_mapping = [
242 ("gate_up_proj", "gate_proj", 0),
243 ("gate_up_proj", "up_proj", 1),
@@ -250,7 +254,7 @@ class DeepseekV32MTP(nn.Module, DeepseekV2MixtureOfExperts):
254 num_experts=self.config.n_routed_experts
255 + (
256 self.config.n_shared_experts
253 - if rocm_aiter_moe_shared_expert_enabled
257 + if self.is_fused_shared_expert_enabled
258 else 0
259 ),
260 )
@@ -266,7 +270,7 @@ class DeepseekV32MTP(nn.Module, DeepseekV2MixtureOfExperts):
270 if spec_layer is None:
271 continue
272 is_fusion_moe_shared_experts_layer = (
269 - rocm_aiter_moe_shared_expert_enabled and ("mlp.shared_experts" in name)
273 + self.is_fused_shared_expert_enabled and ("mlp.shared_experts" in name)
274 )
275 name = self._rewrite_spec_layer_name(spec_layer, name)
276
vllm/models/deepseek_v32/nvidia/mtp.py
+3 -43
@@ -7,7 +7,6 @@ import torch
7 import torch.nn as nn
8
9 import vllm.envs as envs
10 -from vllm._aiter_ops import rocm_aiter_ops
10 from vllm.config import VllmConfig
11 from vllm.forward_context import get_forward_context, is_forward_context_available
12 from vllm.model_executor.layers.fused_embed_norm import (
@@ -344,9 +343,6 @@ class DeepseekV32MTP(nn.Module, DeepseekV2MixtureOfExperts):
343 return name
344
345 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
347 - rocm_aiter_moe_shared_expert_enabled = (
348 - rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()
349 - )
346 stacked_params_mapping = [
347 ("gate_up_proj", "gate_proj", 0),
348 ("gate_up_proj", "up_proj", 1),
@@ -360,12 +356,7 @@ class DeepseekV32MTP(nn.Module, DeepseekV2MixtureOfExperts):
356 ckpt_gate_proj_name="gate_proj",
357 ckpt_down_proj_name="down_proj",
358 ckpt_up_proj_name="up_proj",
363 - num_experts=self.config.n_routed_experts
364 - + (
365 - self.config.n_shared_experts
366 - if rocm_aiter_moe_shared_expert_enabled
367 - else 0
368 - ),
359 + num_experts=self.config.n_routed_experts,
360 )
361
362 pp_missing_layer_names = get_pp_missing_layer_names(self)
@@ -378,9 +369,6 @@ class DeepseekV32MTP(nn.Module, DeepseekV2MixtureOfExperts):
369 spec_layer = get_spec_layer_idx_from_weight_name(self.config, name)
370 if spec_layer is None:
371 continue
381 - is_fusion_moe_shared_experts_layer = (
382 - rocm_aiter_moe_shared_expert_enabled and ("mlp.shared_experts" in name)
383 - )
372 name = self._rewrite_spec_layer_name(spec_layer, name)
373
374 if _try_load_fp8_indexer_wk(
@@ -398,8 +386,6 @@ class DeepseekV32MTP(nn.Module, DeepseekV2MixtureOfExperts):
386 continue
387 if ("mlp.experts." in name) and name not in params_dict:
388 continue
401 - if is_fusion_moe_shared_experts_layer:
402 - continue
389 name_mapped = name.replace(weight_name, param_name)
390 if (
391 param_name == "fused_qkv_a_proj"
@@ -415,32 +401,10 @@ class DeepseekV32MTP(nn.Module, DeepseekV2MixtureOfExperts):
401 break
402 else:
403 num_chunks = 1
418 - if is_fusion_moe_shared_experts_layer:
419 - num_chunks = getattr(self.config, "n_shared_experts", 1) or 1
420 - split_dim = (
421 - 1
422 - if ("down_proj.weight" in name and loaded_weight.ndim > 1)
423 - else 0
424 - )
425 - total = loaded_weight.shape[split_dim]
426 - assert total % num_chunks == 0
427 - chunk_size = total // num_chunks
404
405 for j in range(num_chunks):
406 chunk_name = name
407 weight_to_load = loaded_weight
432 - if is_fusion_moe_shared_experts_layer:
433 - chunk_slice = slice(j * chunk_size, (j + 1) * chunk_size)
434 - if loaded_weight.ndim == 1:
435 - weight_to_load = loaded_weight[chunk_slice]
436 - elif split_dim == 0:
437 - weight_to_load = loaded_weight[chunk_slice, :]
438 - else:
439 - weight_to_load = loaded_weight[:, chunk_slice]
440 - chunk_name = name.replace(
441 - "mlp.shared_experts",
442 - f"mlp.experts.{self.config.n_routed_experts + j}",
443 - )
408
409 is_expert_weight = False
410 for mapping in expert_params_mapping:
@@ -462,10 +426,7 @@ class DeepseekV32MTP(nn.Module, DeepseekV2MixtureOfExperts):
426 return_success=True,
427 )
428 if success:
465 - if not is_fusion_moe_shared_experts_layer:
466 - name = name_mapped
467 - else:
468 - loaded_params.add(name_mapped)
429 + name = name_mapped
430 break
431 else:
432 if is_expert_weight:
@@ -485,8 +446,7 @@ class DeepseekV32MTP(nn.Module, DeepseekV2MixtureOfExperts):
446 param, "weight_loader", default_weight_loader
447 )
448 weight_loader(param, loaded_weight)
488 - if not is_fusion_moe_shared_experts_layer:
489 - loaded_params.add(name)
449 + loaded_params.add(name)
450
451 loaded_layers: set[int] = set()
452 for param_name in loaded_params:
vllm/models/deepseek_v4/amd/model.py
+44 -51
@@ -16,12 +16,16 @@ from vllm.distributed import (
16 get_tensor_model_parallel_rank,
17 get_tensor_model_parallel_world_size,
18 )
19 +from vllm.logger import init_logger
20 from vllm.model_executor.layers.activation import SiluAndMul, SiluAndMulWithClamp
21 from vllm.model_executor.layers.fused_moe import (
22 FusedMoEFactory,
23 GateLinear,
24 fused_moe_make_expert_params_mapping,
25 )
26 +from vllm.model_executor.layers.fused_moe.utils import (
27 + is_model_fused_shared_expert_compatible,
28 +)
29 from vllm.model_executor.layers.layernorm import RMSNorm
30 from vllm.model_executor.layers.linear import (
31 MergedColumnParallelLinear,
@@ -37,6 +41,9 @@ from vllm.model_executor.layers.mhc import (
41 MHCPreOp,
42 )
43 from vllm.model_executor.layers.quantization import QuantizationConfig
44 +from vllm.model_executor.layers.quantization.utils.config_utils import (
45 + is_shared_expert_quant_fse_compatible,
46 +)
47 from vllm.model_executor.layers.vocab_parallel_embedding import (
48 ParallelLMHead,
49 VocabParallelEmbedding,
@@ -60,6 +67,8 @@ from vllm.models.deepseek_v4.amd.rocm import DeepseekV4ROCMAiterMLAAttention
67 from vllm.platforms import current_platform
68 from vllm.sequence import IntermediateTensors
69
70 +logger = init_logger(__name__)
71 +
72
73 class DeepseekV4MLP(nn.Module):
74 def __init__(
@@ -154,48 +163,11 @@ class DeepseekV4MLP(nn.Module):
163 return x
164
165
157 -def _shared_experts_are_fp4(config, layer_idx: int | None = None) -> bool:
158 - """Whether the shared experts are MXFP4 and thus fusable.
159 -
160 - ``layer_idx=None`` resolves the model-wide default (global scheme), used by
161 - the main-model weight loader / mapper callers that operate per-model.
162 - """
163 - quant_cfg = getattr(config, "quantization_config", None)
164 - if quant_cfg is None:
165 - return False
166 - if layer_idx is None:
167 - base = None
168 - elif layer_idx >= config.num_hidden_layers:
169 - base = f"mtp.{layer_idx - config.num_hidden_layers}.ffn.shared_experts"
170 - else:
171 - base = f"layers.{layer_idx}.ffn.shared_experts"
172 - if base and any(e.startswith(base) for e in (quant_cfg.get("exclude") or [])):
173 - return False
174 - entry = (
175 - (quant_cfg.get("layer_quant_config") or {}).get(f"{base}.w1") if base else None
176 - )
177 - if entry is None:
178 - entry = quant_cfg.get("global_quant_config")
179 - return ((entry or {}).get("weight") or {}).get("dtype") == "fp4"
180 -
181 -
182 -def _fuse_shared_experts_enabled(config, prefix: str = "") -> bool:
183 - """Whether to fuse the shared expert into the routed MXFP4 grouped GEMM.
184 -
185 - Fusion fuses the shared expert into the routed experts' MXFP4 grouped GEMM,
186 - so it only applies where the shared expert is the same precision as the
187 - routed experts. Some layers may carry a shared expert in a different quantization
188 - than the routed experts; when so, it runs as its own linear and must not be fused.
189 - """
190 - if not (
191 - current_platform.is_rocm()
192 - and getattr(config, "n_shared_experts", None)
166 +def _fuse_shared_experts_enabled(config) -> bool:
167 + return bool(
168 + getattr(config, "n_shared_experts", None)
169 and envs.VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS
170 and not get_current_vllm_config().parallel_config.enable_expert_parallel
195 - ):
196 - return False
197 - return _shared_experts_are_fp4(
198 - config, extract_layer_index(prefix) if prefix else None
171 )
172
173
@@ -255,9 +227,25 @@ class DeepseekV4MoE(nn.Module):
227
228 self.n_shared_experts = config.n_shared_experts
229
258 - self.fuse_shared_experts = _fuse_shared_experts_enabled(config, prefix)
230 + # TODO: Historically, only `VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS=1`
231 + # is checked to enable FSE for DeepSeek-v4, despite AITER not being used.
232 + # This should be cleaned up and use `resolve_layer_fused_shared_expert`.
233 + fse_requested = _fuse_shared_experts_enabled(config)
234 + if fse_requested:
235 + fse_compatible, fse_reason = is_shared_expert_quant_fse_compatible(
236 + quant_config,
237 + f"{prefix}.experts",
238 + f"{prefix}.shared_experts",
239 + )
240 + if not fse_compatible:
241 + logger.warning(
242 + "VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS is enabled but "
243 + "cannot be enabled: %s.",
244 + fse_reason,
245 + )
246 + self.is_fused_shared_expert_enabled = fse_requested and fse_compatible
247
260 - if config.n_shared_experts is None or self.fuse_shared_experts:
248 + if config.n_shared_experts is None or self.is_fused_shared_expert_enabled:
249 self.shared_experts = None
250 else:
251 intermediate_size = config.moe_intermediate_size * config.n_shared_experts
@@ -282,8 +270,9 @@ class DeepseekV4MoE(nn.Module):
270 self.experts = FusedMoEFactory(
271 shared_experts=self.shared_experts,
272 n_shared_experts=(
285 - config.n_shared_experts if self.fuse_shared_experts else None
273 + config.n_shared_experts if self.is_fused_shared_expert_enabled else None
274 ),
275 + fuse_shared_experts=self.is_fused_shared_expert_enabled,
276 gate=self.gate,
277 num_experts=config.n_routed_experts,
278 top_k=config.num_experts_per_tok,
@@ -581,6 +570,11 @@ class DeepseekV4Model(nn.Module, EagleModelMixin):
570 ),
571 prefix=f"{prefix}.layers",
572 )
573 + self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible(
574 + self.layers,
575 + DeepseekV4MoE,
576 + "ffn",
577 + )
578
579 if get_pp_group().is_last_rank:
580 self.norm = RMSNorm(config.hidden_size, self.rms_norm_eps)
@@ -763,7 +757,7 @@ class DeepseekV4Model(nn.Module, EagleModelMixin):
757 # diverge from how the module was built if per-layer quantization ever
758 # mixes fused and non-fused layers.
759 fuse_by_layer = {
766 - extract_layer_index(mod_name): mod.fuse_shared_experts
760 + extract_layer_index(mod_name): mod.is_fused_shared_expert_enabled
761 for mod_name, mod in self.named_modules()
762 if isinstance(mod, DeepseekV4MoE)
763 }
@@ -874,7 +868,7 @@ class DeepseekV4Model(nn.Module, EagleModelMixin):
868 # redirected shared-expert weights route through the expert loader.
869 n_shared = getattr(self.config, "n_shared_experts", 0) or 0
870 num_experts = self.config.n_routed_experts + (
877 - n_shared if _fuse_shared_experts_enabled(self.config) else 0
871 + n_shared if self.is_fused_shared_expert_enabled else 0
872 )
873 return fused_moe_make_expert_params_mapping(
874 self,
@@ -956,15 +950,14 @@ class DeepseekV4ForCausalLM(nn.Module, SupportsPP, SupportsEagle3):
950 config = vllm_config.model_config.hf_config
951 self.config = config
952 expert_dtype = getattr(config, "expert_dtype", "fp4")
959 - fuse_shared_experts = _fuse_shared_experts_enabled(config)
960 - if expert_dtype != "fp4" or fuse_shared_experts:
961 - self.hf_to_vllm_mapper = _make_deepseek_v4_weights_mapper(
962 - expert_dtype, fuse_shared_experts=fuse_shared_experts
963 - )
964 -
953 self.model = self.model_cls(
954 vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
955 )
956 + if expert_dtype != "fp4" or self.model.is_fused_shared_expert_enabled:
957 + self.hf_to_vllm_mapper = _make_deepseek_v4_weights_mapper(
958 + expert_dtype,
959 + fuse_shared_experts=self.model.is_fused_shared_expert_enabled,
960 + )
961 if get_pp_group().is_last_rank:
962 self.lm_head = ParallelLMHead(
963 config.vocab_size,
vllm/models/minimax_m3/amd/model.py
+64 -30
@@ -34,6 +34,7 @@ from vllm.config import (
34 )
35 from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
36 from vllm.forward_context import get_forward_context
37 +from vllm.logger import init_logger
38 from vllm.model_executor.layers.attention import Attention
39 from vllm.model_executor.layers.attention.attention import set_default_quant_scales
40 from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
@@ -45,6 +46,9 @@ from vllm.model_executor.layers.fused_moe import (
46 GateLinear,
47 fused_moe_make_expert_params_mapping,
48 )
49 +from vllm.model_executor.layers.fused_moe.utils import (
50 + is_model_fused_shared_expert_compatible,
51 +)
52 from vllm.model_executor.layers.linear import (
53 MergedColumnParallelLinear,
54 MinimaxM3QKVParallelLinearWithIndexer,
@@ -53,6 +57,9 @@ from vllm.model_executor.layers.linear import (
57 )
58 from vllm.model_executor.layers.logits_processor import LogitsProcessor
59 from vllm.model_executor.layers.quantization import QuantizationConfig
60 +from vllm.model_executor.layers.quantization.utils.config_utils import (
61 + is_shared_expert_quant_fse_compatible,
62 +)
63 from vllm.model_executor.layers.rotary_embedding import get_rope
64 from vllm.model_executor.layers.vocab_parallel_embedding import (
65 ParallelLMHead,
@@ -108,23 +115,7 @@ from vllm.v1.kv_cache_interface import (
115 is_quantized_kv_cache,
116 )
117
111 -
112 -def _fuse_shared_experts_enabled(config: PretrainedConfig) -> bool:
113 - """Whether to fuse the shared expert with routed experts.
114 -
115 - ROCm only. Opt-in via ``VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS`` (the
116 - router-append fusion runs on both aiter and non-aiter MoE);
117 - it is disabled under expert parallelism (the shared slot is appended to
118 - the routed top-k, which the EP expert-mapping path does not handle).
119 - """
120 - from vllm.platforms import current_platform
121 -
122 - return bool(
123 - current_platform.is_rocm()
124 - and getattr(config, "n_shared_experts", None)
125 - and envs.VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS
126 - and not get_current_vllm_config().parallel_config.enable_expert_parallel
127 - )
118 +logger = init_logger(__name__)
119
120
121 def _sparse_attention_layer_ids(config: PretrainedConfig) -> set[int]:
@@ -303,10 +294,12 @@ class MiniMaxM3MLP(nn.Module):
294 return x
295
296
306 -def _aiter_moe_fused_shared_experts_enabled(config: PretrainedConfig) -> bool:
297 +def _aiter_moe_fused_shared_experts_enabled(
298 + is_fused_shared_expert_enabled: bool,
299 +) -> bool:
300 """Whether the fused shared expert routes through aiter's grouped top-k MoE.
301
309 - A strict sub-case of :func:`_fuse_shared_experts_enabled`: shared-expert
302 + A strict sub-case of `is_fused_shared_expert_enabled`: shared-expert
303 fusion must already be opted in (``VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS``)
304 and allowed (not under expert parallelism). When additionally on gfx950 with
305 an active aiter MoE backend, the shared expert is appended inside aiter's
@@ -314,11 +307,13 @@ def _aiter_moe_fused_shared_experts_enabled(config: PretrainedConfig) -> bool:
307 vLLM router's torch concat. Otherwise FSE still runs via the vLLM top-k bias
308 router.
309 """
317 - if not _fuse_shared_experts_enabled(config):
318 - return False
310 from vllm.platforms.rocm import on_gfx950
311
321 - return on_gfx950() and rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()
312 + return (
313 + on_gfx950()
314 + and is_fused_shared_expert_enabled
315 + and rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()
316 + )
317
318
319 class MiniMaxM3MoE(nn.Module):
@@ -371,11 +366,46 @@ class MiniMaxM3MoE(nn.Module):
366 # MoE call as the last expert slot, so we don't build a separate module.
367 # On gfx950 with aiter MoE the append is fused inside aiter's grouped
368 # top-k kernel; otherwise it goes through the vLLM top-k bias router.
374 - self.fuse_shared_experts = _fuse_shared_experts_enabled(config)
375 - self.use_aiter_moe_fse = _aiter_moe_fused_shared_experts_enabled(config)
369 + # It is disabled under expert parallelism (the shared slot is appended to
370 + # the routed top-k, which the EP expert-mapping path does not handle).
371 + # TODO: Historically, only `VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS=1`
372 + # is checked to enable FSE for MiniMax-M3, despite AITER not being used.
373 + # This should be cleaned up and use `resolve_layer_fused_shared_expert`.
374 + fse_requested = envs.VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS
375 +
376 + self.is_fused_shared_expert_enabled = False
377 + if (
378 + fse_requested
379 + and bool(getattr(config, "n_shared_experts", None))
380 + and not get_current_vllm_config().parallel_config.enable_expert_parallel
381 + ):
382 + fse_compatible, fse_reason = is_shared_expert_quant_fse_compatible(
383 + quant_config,
384 + f"{prefix}.experts",
385 + f"{prefix}.shared_experts",
386 + )
387 + if not fse_compatible:
388 + logger.warning(
389 + "VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS is enabled but "
390 + "cannot be enabled: %s.",
391 + fse_reason,
392 + )
393 + else:
394 + self.is_fused_shared_expert_enabled = True
395 +
396 + # When additionally on gfx950 with an active aiter MoE backend, the shared
397 + # expert is appended inside aiter's an active aiter MoE backend, the shared
398 + # expert is appended inside aiter's biased grouped top-k kernel
399 + # (``num_fused_shared_experts``) instead of the vLLM router's torch concat.
400 + # Otherwise FSE still runs via the vLLM top-k bias router.
401 + # TODO: `on_gfx950()` check here should not be MiniMax-M3 specific, and
402 + # the check should be done on resolved MOE backend directly.
403 + self.use_aiter_moe_fse = _aiter_moe_fused_shared_experts_enabled(
404 + self.is_fused_shared_expert_enabled
405 + )
406
407 self.shared_experts: MiniMaxM3MLP | None = None
378 - if self.n_shared_experts and not self.fuse_shared_experts:
408 + if self.n_shared_experts and not self.is_fused_shared_expert_enabled:
409 self.shared_experts = MiniMaxM3MLP(
410 config=config,
411 intermediate_size=config.intermediate_size * self.n_shared_experts,
@@ -412,8 +442,9 @@ class MiniMaxM3MoE(nn.Module):
442 router_logits_dtype=self.gate.out_dtype,
443 shared_experts=self.shared_experts,
444 n_shared_experts=(
415 - self.n_shared_experts if self.fuse_shared_experts else None
445 + self.n_shared_experts if self.is_fused_shared_expert_enabled else None
446 ),
447 + fuse_shared_experts=self.is_fused_shared_expert_enabled,
448 quant_config=quant_config,
449 prefix=f"{prefix}.experts",
450 )
@@ -1101,6 +1132,11 @@ class MiniMaxM3Model(nn.Module, EagleModelMixin):
1132 ),
1133 prefix=f"{prefix}.layers",
1134 )
1135 + self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible(
1136 + self.layers,
1137 + MiniMaxM3MoE,
1138 + "block_sparse_moe",
1139 + )
1140
1141 if get_pp_group().is_last_rank:
1142 self.norm = MiniMAXGemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
@@ -1155,7 +1191,7 @@ class MiniMaxM3Model(nn.Module, EagleModelMixin):
1191 # expert, include the appended slot (id == num_local_experts).
1192 n_shared = getattr(self.config, "n_shared_experts", 0) or 0
1193 num_experts = self.config.num_local_experts + (
1158 - n_shared if _fuse_shared_experts_enabled(self.config) else 0
1194 + n_shared if self.is_fused_shared_expert_enabled else 0
1195 )
1196 return fused_moe_make_expert_params_mapping(
1197 self,
@@ -1187,8 +1223,6 @@ class MiniMaxM3Model(nn.Module, EagleModelMixin):
1223 # (param_name, weight_name, expert_id, shard_id)
1224 expert_params_mapping = self.get_expert_mapping()
1225
1190 - _fuse_shared = _fuse_shared_experts_enabled(self.config)
1191 -
1226 params_dict = dict(self.named_parameters())
1227 loaded_params: set[str] = set()
1228 for name, loaded_weight in weights:
@@ -1206,7 +1240,7 @@ class MiniMaxM3Model(nn.Module, EagleModelMixin):
1240 # down->w2) so it loads via the routed expert loader. Runs before the
1241 # stacked/dense mappings so shared_experts.gate_proj/up_proj are not
1242 # captured by the dense gate_up_proj mapping.
1209 - if _fuse_shared and ".shared_experts." in name:
1243 + if self.is_fused_shared_expert_enabled and ".shared_experts." in name:
1244 sid = self.config.num_local_experts
1245 name = name.replace(".shared_experts.gate_proj.", f".experts.{sid}.w1.")
1246 name = name.replace(".shared_experts.up_proj.", f".experts.{sid}.w3.")