diff --git a/tests/kernels/moe/test_fp8_tp_loading.py b/tests/kernels/moe/test_fp8_tp_loading.py index 680ccb812cca..0ac056919693 100644 --- a/tests/kernels/moe/test_fp8_tp_loading.py +++ b/tests/kernels/moe/test_fp8_tp_loading.py @@ -68,6 +68,7 @@ def _make_fp8_tp_experts( or Fp8Config(is_checkpoint_fp8_serialized=True, weight_block_size=[128, 128]), expert_map_manager=SimpleNamespace( local_num_experts=num_experts, + num_fused_shared_experts=0, placement_strategy="linear", expert_map=None, expert_mask=None, diff --git a/tests/model_executor/layers/test_fused_shared_expert.py b/tests/model_executor/layers/test_fused_shared_expert.py index 30b7e2fbd34a..166ccdd197e7 100644 --- a/tests/model_executor/layers/test_fused_shared_expert.py +++ b/tests/model_executor/layers/test_fused_shared_expert.py @@ -17,6 +17,7 @@ import vllm.config as vllm_config_module from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config from vllm.config.quantization import QuantizationConfigArgs +from vllm.logger import _print_warning_once from vllm.model_executor.layers.fused_moe import utils as fused_moe_utils from vllm.model_executor.layers.fused_moe.layer import determine_expert_counts from vllm.model_executor.layers.quantization.experts_int8 import ExpertsInt8Config @@ -31,7 +32,10 @@ from vllm.model_executor.models.qwen3_next import ( _should_replicate_misaligned_shared_expert, ) -from vllm.model_executor.models.utils import PPMissingLayer +from vllm.model_executor.models.utils import ( + get_fused_shared_expert_layers, + maybe_fuse_shared_experts, +) from vllm.models.deepseek_v4 import quant_config as deepseek_v4_quant_config from vllm.models.minimax_m3.amd import model as minimax_m3_model from vllm.platforms import current_platform @@ -148,13 +152,15 @@ class _StubGlm5NextAttention(nn.Module): def __init__(self, *args: Any, **kwargs: Any) -> None: super().__init__() self.o_proj = SimpleNamespace(reduce_results=True) + # GLM5Next's MTP predictor reads this attribute from MLA layers. + self.mla_attn = SimpleNamespace() def get_fse_test_model_config( model_type: str, quantization_config: dict[str, Any], -) -> tuple[object, type[nn.Module]]: - config: object +) -> tuple[Any, type[nn.Module]]: + config: Any if model_type == "minimax_m3": config = MiniMaxM3TextConfig( hidden_size=128, @@ -166,24 +172,29 @@ def get_fse_test_model_config( head_dim=128, num_local_experts=2, num_experts_per_tok=1, - moe_layer_freq=[0, 1], + moe_layer_freq=[1, 1], sparse_attention_config=None, rotary_dim=64, quantization_config=quantization_config, ) return config, minimax_m3_model.MiniMaxM3Model - if model_type == "deepseek_v4": - from vllm.models.deepseek_v4.amd.model import DeepseekV4Model + if model_type in ("deepseek_v4", "deepseek_v4_cpu"): + if model_type == "deepseek_v4": + from vllm.models.deepseek_v4.amd.model import DeepseekV4Model + else: + from vllm.models.deepseek_v4.cpu.model import DeepseekV4Model config = SimpleNamespace( + model_type="deepseek_v4", vocab_size=256, hidden_size=128, - num_hidden_layers=1, + num_hidden_layers=2, num_attention_heads=1, head_dim=128, max_position_embeddings=128, sliding_window=None, - layer_types=["sliding_attention"], + layer_types=["sliding_attention", "sliding_attention"], + mlp_layer_types=[], compress_rates={}, rope_theta=10000.0, compress_rope_theta=10000.0, @@ -214,13 +225,13 @@ def get_fse_test_model_config( quantization_config=quantization_config, ) return config, DeepseekV4Model - if model_type == "qwen3_5": - from vllm.model_executor.models.qwen3_5 import Qwen3_5Model + if model_type in ("qwen3_next", "qwen3_5", "alice_ai"): + from transformers import Qwen3NextConfig - config = Qwen3_5MoeTextConfig( + common = dict( vocab_size=256, hidden_size=128, - num_hidden_layers=1, + num_hidden_layers=2, num_attention_heads=1, num_key_value_heads=1, head_dim=128, @@ -232,20 +243,113 @@ def get_fse_test_model_config( shared_expert_intermediate_size=32, num_experts_per_tok=1, num_experts=2, - layer_types=["full_attention"], + layer_types=["full_attention", "full_attention"], quantization_config=quantization_config, ) - return config, Qwen3_5Model - if model_type == "deepseek_v2": - from transformers import DeepseekV2Config + if model_type == "qwen3_next": + from vllm.model_executor.models.qwen3_next import Qwen3NextModel + + return Qwen3NextConfig(**common), Qwen3NextModel + if model_type == "qwen3_5": + from vllm.model_executor.models.qwen3_5 import Qwen3_5Model + + return Qwen3_5MoeTextConfig(**common), Qwen3_5Model + from vllm.model_executor.models.alice_ai import AliceAIModel + + config = Qwen3NextConfig(**common, linear_conv_kernel_dim=4) + # AliceAI requires the AttnRes grouping metadata. + config.block_attn_res_block_size = 1 + return config, AliceAIModel + if model_type == "qwen4_exp": + from transformers import Qwen4ExpTextConfig + + config = Qwen4ExpTextConfig( + vocab_size=256, + hidden_size=128, + num_hidden_layers=2, + num_attention_heads=1, + num_key_value_heads=1, + head_dim=128, + linear_key_head_dim=64, + linear_value_head_dim=64, + linear_num_key_heads=1, + linear_num_value_heads=1, + moe_intermediate_size=32, + shared_expert_intermediate_size=32, + num_experts_per_tok=1, + num_experts=2, + layer_types=["full_attention", "full_attention"], + eos_token_id=1, + hc_count=2, + hc_lowrank=4, + ple_layer_ids=[], + quantization_config=quantization_config, + ) + from vllm.models.qwen4_exp.amd.model import Qwen4ExpModel + + return config, Qwen4ExpModel + if model_type == "AXK1": + from transformers import AXK1Config - from vllm.model_executor.models.deepseek_v2 import DeepseekV2Model + from vllm.model_executor.models.AXK1 import AXK1Model - config = DeepseekV2Config( + config = AXK1Config( vocab_size=256, hidden_size=128, intermediate_size=32, - num_hidden_layers=1, + moe_intermediate_size=32, + num_hidden_layers=2, + num_attention_heads=1, + num_key_value_heads=1, + n_shared_experts=1, + n_routed_experts=2, + num_experts_per_tok=1, + n_group=1, + topk_group=1, + first_k_dense_replace=0, + q_lora_rank=0, + kv_lora_rank=0, + qk_nope_head_dim=0, + qk_rope_head_dim=0, + v_head_dim=0, + quantization_config=quantization_config, + ) + return config, AXK1Model + if model_type == "glm4_moe_lite": + from transformers.models.glm4_moe_lite import Glm4MoeLiteConfig + + from vllm.model_executor.models.glm4_moe_lite import Glm4MoeLiteModel + + config = Glm4MoeLiteConfig( + vocab_size=256, + hidden_size=128, + intermediate_size=32, + moe_intermediate_size=32, + num_hidden_layers=2, + num_attention_heads=1, + num_key_value_heads=1, + n_routed_experts=2, + n_shared_experts=1, + num_experts_per_tok=1, + q_lora_rank=0, + kv_lora_rank=32, + qk_nope_head_dim=32, + qk_rope_head_dim=16, + v_head_dim=32, + dtype=torch.bfloat16, + quantization_config=quantization_config, + ) + config.first_k_dense_replace = 0 + config.num_nextn_predict_layers = 1 + return config, Glm4MoeLiteModel + if model_type in ("deepseek_v2", "deepseek_v32"): + from transformers import DeepseekV2Config + + common = dict( + vocab_size=256, + hidden_size=128, + intermediate_size=32, + num_hidden_layers=2, num_attention_heads=1, num_key_value_heads=1, head_dim=128, @@ -263,7 +367,18 @@ def get_fse_test_model_config( v_head_dim=0, quantization_config=quantization_config, ) - return config, DeepseekV2Model + config = DeepseekV2Config(**common) + if model_type == "deepseek_v2": + from vllm.model_executor.models.deepseek_v2 import DeepseekV2Model + + # Needed to build the DeepSeek-V2 MTP draft model in this test. + config.num_nextn_predict_layers = 1 + return config, DeepseekV2Model + from vllm.models.deepseek_v32.amd.model import DeepseekV32Model + + config.index_topk = 1 + config.num_nextn_predict_layers = 1 + return config, DeepseekV32Model if model_type == "glm4_moe": from transformers.models.glm4_moe import Glm4MoeConfig @@ -273,7 +388,7 @@ def get_fse_test_model_config( vocab_size=256, hidden_size=128, intermediate_size=32, - num_hidden_layers=1, + num_hidden_layers=2, num_attention_heads=1, num_key_value_heads=1, head_dim=128, @@ -295,14 +410,14 @@ def get_fse_test_model_config( hidden_size=128, intermediate_size=32, moe_intermediate_size=32, - num_hidden_layers=1, + num_hidden_layers=2, num_attention_heads=1, num_key_value_heads=1, n_routed_experts=2, n_shared_experts=1, num_experts_per_tok=1, - layer_types=["linear_attention"], - mlp_layer_types=["sparse"], + layer_types=["linear_attention", "linear_attention"], + mlp_layer_types=["sparse", "sparse"], index_topk=1, index_kpool=1, pad_token_id=None, @@ -311,6 +426,8 @@ def get_fse_test_model_config( mhc=False, quantization_config=quantization_config, ) + config.mla_use_nope = False + config.num_nextn_predict_layers = 1 return config, Glm5NextModel raise ValueError(f"Unsupported FSE test model: {model_type}") @@ -402,9 +519,18 @@ def check_compatibility( ] -def test_resolve_layer_fused_shared_expert_rejects_incompatible_quantization( +@pytest.mark.parametrize( + ("reason", "enable_ep"), + [ + ("shared experts are excluded", False), + ("expert parallelism is enabled", True), + ], +) +def test_resolve_layer_fused_shared_expert_logs_skip_reason( monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture, + reason: str, + enable_ep: bool, ) -> None: monkeypatch.setattr( fused_moe_utils.rocm_aiter_ops, @@ -414,13 +540,20 @@ def test_resolve_layer_fused_shared_expert_rejects_incompatible_quantization( monkeypatch.setattr( fused_moe_utils, "is_shared_expert_quant_fse_compatible", - lambda *_, **__: (False, "shared experts are excluded"), + lambda *_, **__: (False, reason), + ) + monkeypatch.setattr( + "vllm.config.get_current_vllm_config_or_none", + lambda: SimpleNamespace( + parallel_config=SimpleNamespace(enable_expert_parallel=enable_ep) + ), ) + _print_warning_once.cache_clear() assert not fused_moe_utils.resolve_layer_fused_shared_expert( - ExpertsInt8Config(), "model.layers.0.mlp" + ExpertsInt8Config(), "model.layers.0.mlp", disable_under_ep=enable_ep ) - assert "shared experts are excluded" in caplog.text + assert reason in caplog.text def test_deepseek_v4_shared_expert_fse_uses_mtp_quantization_config_prefix( @@ -884,45 +1017,409 @@ def test_aiter_fused_moe_validates_shared_expert_arguments() -> None: assert validate(shared_tensor, shared_tensor, shared_tensor, shared_tensor, 0) -def test_is_model_fused_shared_expert_compatible() -> None: - class MoE(nn.Module): - def __init__(self, enabled: bool) -> None: - super().__init__() - self.is_fused_shared_expert_enabled = enabled +def test_deepseek_v4_shared_expert_mapper_fallback() -> None: + """Mixed and all-excluded DSV4 checkpoints map shared experts correctly.""" + from vllm.models.deepseek_v4.amd.model import ( + _make_deepseek_v4_weights_mapper, + _redirect_shared_expert_name, + ) + + tensor = torch.empty(0) + fuse_by_layer = frozenset({6}) + # Mixed: any fused layer puts the mapper in fused mode, which leaves the + # checkpoint .w2 names untouched; the loader redirects per layer. + mapper = _make_deepseek_v4_weights_mapper("fp8", fuse_shared_experts=True) + for raw, expected in [ + ( + "layers.5.ffn.shared_experts.w2.weight", + "model.layers.5.ffn.shared_experts.down_proj.weight", + ), + ( + "layers.5.ffn.shared_experts.w2.scale", + "model.layers.5.ffn.shared_experts.down_proj.weight_scale_inv", + ), + ( + "layers.6.ffn.shared_experts.w2.weight", + "model.layers.6.ffn.experts.4.w2.weight", + ), + ( + "layers.6.ffn.shared_experts.w2.scale", + "model.layers.6.ffn.experts.4.w2.weight_scale_inv", + ), + ]: + [(mapped, _)] = list(mapper.apply([(raw, tensor)])) + assert _redirect_shared_expert_name(mapped, fuse_by_layer, 4) == expected + + # All layers excluded (fused_shared_expert_layers == frozenset()) selects + # the unfused mapper, which renames .w2 -> .down_proj itself. + mapper = _make_deepseek_v4_weights_mapper("fp8", fuse_shared_experts=False) + for raw, expected in [ + ( + "layers.5.ffn.shared_experts.w2.weight", + "model.layers.5.ffn.shared_experts.down_proj.weight", + ), + ( + "layers.5.ffn.shared_experts.w2.scale", + "model.layers.5.ffn.shared_experts.down_proj.weight_scale_inv", + ), + ]: + [(mapped, _)] = list(mapper.apply([(raw, tensor)])) + assert _redirect_shared_expert_name(mapped, frozenset(), 4) == expected + - class Layer(nn.Module): - def __init__(self, enabled: bool) -> None: - super().__init__() - self.mlp = MoE(enabled) +class _FakeMoE(nn.Module): + def __init__(self, enabled: bool) -> None: + super().__init__() + self.is_fused_shared_expert_enabled = enabled - enabled_layers = nn.ModuleList([Layer(True)]) - disabled_layers = nn.ModuleList([Layer(False)]) - mixed_layers = nn.ModuleList([Layer(True), Layer(False)]) - empty_layers = nn.ModuleList() - pipeline_layers = nn.ModuleList([Layer(True), PPMissingLayer()]) - assert fused_moe_utils.is_model_fused_shared_expert_compatible( - enabled_layers, MoE, "mlp" +def _make_fake_moe_model( + layer_names: list[str], enabled: list[bool], use_dict: bool +) -> nn.Module: + layers = [] + for flag in enabled: + layer = nn.Module() + layer.mlp = _FakeMoE(flag) + layers.append(layer) + model = nn.Module() + model.layers = ( + nn.ModuleDict(dict(zip(layer_names, layers, strict=True))) + if use_dict + else nn.ModuleList(layers) ) - assert not fused_moe_utils.is_model_fused_shared_expert_compatible( - disabled_layers, MoE, "mlp" + return model + + +def test_fused_shared_expert_layers_matches_container_index_space() -> None: + """The fused-layer set follows the module tree, not a hard-coded offset. + + ``nn.ModuleList`` MTP layers (Qwen3 Next/3.5, Qwen4 Exp) are 0-based, + while ``nn.ModuleDict`` MTP layers (GLM4, DeepSeek) are keyed by absolute + layer index. Both must redirect without a per-loader offset. + """ + list_model = _make_fake_moe_model(["0", "1"], [True, False], use_dict=False) + dict_model = _make_fake_moe_model(["28", "29"], [False, True], use_dict=True) + list_layers = get_fused_shared_expert_layers(list_model, _FakeMoE) + dict_layers = get_fused_shared_expert_layers(dict_model, _FakeMoE) + assert list_layers == frozenset({0}) + assert dict_layers == frozenset({29}) + + weight = torch.arange(8.0).reshape(4, 2) + + def redirected(name: str, fused_shared_expert_layers: frozenset[int]) -> bool: + [(out_name, _)] = list( + maybe_fuse_shared_experts( + [(name, weight)], + n_routed_experts=4, + n_shared_experts=1, + fused_shared_expert_layers=fused_shared_expert_layers, + ckpt_prefix="mlp.shared_expert", + ) + ) + return ".experts." in out_name + + # ModuleList: 0-based names, including the Qwen4 Exp MTP case where the + # quant-config prefix uses an absolute index but the weight names do not. + assert redirected("layers.0.mlp.shared_expert.down_proj.weight", list_layers) + assert not redirected("layers.1.mlp.shared_expert.down_proj.weight", list_layers) + # ModuleDict: absolute names, including the MTP layer offset. + assert redirected("layers.29.mlp.shared_expert.down_proj.weight", dict_layers) + assert not redirected("layers.28.mlp.shared_expert.down_proj.weight", dict_layers) + + +@pytest.mark.parametrize( + ("enable_expert_parallel", "match"), + [ + (True, "expert parallelism is not tested"), + (False, "nixl_ep"), + ], +) +def test_mixed_fused_shared_experts_rejected_with_ep_or_nixl_ep( + enable_expert_parallel: bool, match: str +) -> None: + """Mixed FSE must be rejected with EP or a shared-workspace all2all backend. + + With EP, per-layer geometry differences are not validated. ``nixl_ep`` + additionally keeps a single process-wide buffer initialized from the first + MoE layer's ``num_experts_per_rank`` and ignores later geometries. A mixed + model has fused layers with ``n_routed + n_shared`` experts and non-fused + layers with ``n_routed``, so reusing that buffer would silently map experts + to the wrong slots. Construction must fail cleanly instead of running the + untested mixed path. + """ + model = _make_fake_moe_model(["0", "1"], [True, False], use_dict=False) + + vllm_config = VllmConfig() + vllm_config.parallel_config.enable_expert_parallel = enable_expert_parallel + vllm_config.parallel_config.data_parallel_size = 2 + vllm_config.parallel_config.all2all_backend = "nixl_ep" + + with ( + set_current_vllm_config(vllm_config), + pytest.raises(NotImplementedError, match=match), + ): + get_fused_shared_expert_layers(model, _FakeMoE) + + +def _assert_mtp_fused_shared_expert_redirect( + mtp: nn.Module, ckpt_prefix: str, moe: nn.Module, n_routed_experts: int +) -> None: + """Load a fused MTP shared expert and verify it lands in the appended slot. + + The MTP loaders see different name/index spaces (0-based ``ModuleList`` or + absolute ``ModuleDict``); using the live module map must redirect either. + """ + packed_down = torch.full((128, 16), 3, dtype=torch.uint8) + down_scale = torch.full((128, 1), 127, dtype=torch.uint8) + mtp.load_weights( + [ + (f"{ckpt_prefix}.down_proj.weight", packed_down), + (f"{ckpt_prefix}.down_proj.weight_scale", down_scale), + ] ) - assert not fused_moe_utils.is_model_fused_shared_expert_compatible( - empty_layers, MoE, "mlp" + torch.testing.assert_close( + moe.experts.routed_experts.w2_weight[n_routed_experts][:128, :16].float(), + packed_down.float(), ) - assert fused_moe_utils.is_model_fused_shared_expert_compatible( - pipeline_layers, MoE, "mlp" + + +def _assert_mtp_shared_expert_not_fused( + mtp: nn.Module, ckpt_prefix: str, moe: nn.Module, n_routed_experts: int +) -> None: + """Load an unfused MTP shared expert and verify no fused slot is appended. + + Without AITER the shared expert stays a separate MLP, so the routed tensor must + keep exactly ``n_routed_experts`` rows. + """ + packed_down = torch.full((128, 16), 3, dtype=torch.uint8) + down_scale = torch.full((128, 1), 127, dtype=torch.uint8) + mtp.load_weights( + [ + (f"{ckpt_prefix}.down_proj.weight", packed_down), + (f"{ckpt_prefix}.down_proj.weight_scale", down_scale), + ] ) - with pytest.raises(NotImplementedError, match="1 enabled and 1 disabled layers"): - fused_moe_utils.is_model_fused_shared_expert_compatible( - mixed_layers, MoE, "mlp" + assert moe.experts.routed_experts.w2_weight.shape[0] == n_routed_experts + + +def _make_deepseek_v4_shared_expert_checkpoint( + layer: nn.Module, idx: int, fused: bool, exclude: bool +) -> tuple[list[tuple[str, torch.Tensor]], list[tuple[torch.Tensor, torch.Tensor]]]: + """Build DSV4's FP8 shared-expert / MXFP4 routed-slot checkpoint pairs.""" + ffn = layer.ffn + if not fused: + dtype = torch.bfloat16 if exclude else torch.float8_e4m3fn + hidden = ffn.shared_experts.gate_up_proj.weight.shape[1] + intermediate = ffn.shared_experts.down_proj.weight.shape[1] + gate_up = torch.full((2 * intermediate, hidden), idx + 1, dtype=dtype) + down = torch.full((hidden, intermediate), idx + 1, dtype=dtype) + return ( + [ + (f"layers.{idx}.ffn.shared_experts.gate_up_proj.weight", gate_up), + (f"layers.{idx}.ffn.shared_experts.down_proj.weight", down), + ], + [ + (ffn.shared_experts.gate_up_proj.weight, gate_up), + (ffn.shared_experts.down_proj.weight, down), + ], + ) + + routed = ffn.experts.routed_experts + slot = routed.w13_weight.shape[0] - 1 + w13_rows, hidden_half = routed.w13_weight[slot].shape + intermediate = w13_rows // 2 + hidden = hidden_half * 2 + w1 = torch.full((intermediate, hidden_half), idx + 1, dtype=torch.uint8) + w3 = torch.full((intermediate, hidden_half), idx + 1, dtype=torch.uint8) + w2 = torch.full((hidden, intermediate // 2), idx + 1, dtype=torch.uint8) + return ( + [ + (f"layers.{idx}.ffn.shared_experts.w1.weight", w1), + (f"layers.{idx}.ffn.shared_experts.w3.weight", w3), + (f"layers.{idx}.ffn.shared_experts.w2.weight", w2), + ], + [ + (routed.w13_weight[slot].chunk(2, dim=0)[0], w1), + (routed.w13_weight[slot].chunk(2, dim=0)[1], w3), + (routed.w2_weight[slot], w2), + ], + ) + + +def _build_mtp_model( + mtp_cls: Any, + vllm_config: VllmConfig, + config: object, + *, + with_text_config: bool = False, + with_target_model: bool = False, +) -> nn.Module: + draft_model_config: dict[str, object] = {"hf_config": config} + if with_text_config: + draft_model_config["hf_text_config"] = config + speculative_config: dict[str, object] = { + "draft_model_config": SimpleNamespace(**draft_model_config), + "method": "mtp", + "parallel_drafting": False, + "enable_adaptive_verification": False, + "use_dspark": lambda: False, + } + if with_target_model: + speculative_config["target_model_config"] = vllm_config.model_config + vllm_config.speculative_config = SimpleNamespace(**speculative_config) + return mtp_cls(vllm_config=vllm_config) + + +_MTP_FACTORIES: dict[str, tuple[str, str, dict[str, bool]]] = { + "deepseek_v4": ("vllm.models.deepseek_v4.amd.mtp", "DeepSeekV4MTP", {}), + "glm4_moe": ( + "vllm.model_executor.models.glm4_moe_mtp", + "Glm4MoeMTP", + {}, + ), + "glm4_moe_lite": ( + "vllm.model_executor.models.glm4_moe_lite_mtp", + "Glm4MoeLiteMTP", + {}, + ), + "deepseek_v32": ( + "vllm.models.deepseek_v32.amd.mtp", + "DeepseekV32MTP", + {}, + ), + "glm5_next": ( + "vllm.models.glm5next.common.mtp", + "Glm5NextMTP", + {"with_text_config": True}, + ), + "deepseek_v2": ( + "vllm.model_executor.models.deepseek_mtp", + "DeepSeekMTP", + {}, + ), + "qwen3_next": ( + "vllm.model_executor.models.qwen3_next_mtp", + "Qwen3NextMTP", + {"with_text_config": True, "with_target_model": True}, + ), + "qwen3_5": ( + "vllm.model_executor.models.qwen3_5_mtp", + "Qwen3_5MoeMTP", + {"with_text_config": True, "with_target_model": True}, + ), +} + + +def _build_mtp_for_model( + model_type: str, + vllm_config: VllmConfig, + config: object, + enabled: bool, +) -> nn.Module | None: + factory = _MTP_FACTORIES.get(model_type) + if not enabled or factory is None: + return None + module_name, class_name, kwargs = factory + module = importlib.import_module(module_name) + return _build_mtp_model(getattr(module, class_name), vllm_config, config, **kwargs) + + +# TODO: move to conftest.py. +class _SingleRankGroup: + """No-op single-rank stand-in for a compiled ``GroupCoordinator``.""" + + world_size = 1 + rank_in_group = 0 + local_rank = 0 + is_first_rank = True + is_last_rank = True + + @property + def device_group(self): + return self + + @property + def cpu_group(self): + return self + + @property + def device_communicator(self): + return self + + def rank(self) -> int: + return 0 + + def size(self) -> int: + return 1 + + def all_reduce(self, input_: torch.Tensor, *args: Any, **kwargs: Any): + return input_ + + def all_gather(self, input_: torch.Tensor, dim: int = 0): + return input_ + + def all_gatherv(self, inputs: Any, dim: int = 0, sizes: Any = None): + return inputs + + def reduce_scatter(self, input_: torch.Tensor, dim: int = 0): + return input_ + + def reduce_scatterv(self, input_: Any, dim: int = 0, sizes: Any = None): + return input_ + + def broadcast(self, input_: torch.Tensor, src: int = 0): + return input_ + + def send(self, *args: Any, **kwargs: Any) -> None: + return None + + def recv(self, *args: Any, **kwargs: Any) -> None: + return None + + def barrier(self) -> None: + return None + + def destroy(self) -> None: + return None + + +@pytest.fixture +def single_rank_parallel_state(monkeypatch: pytest.MonkeyPatch): + """Single rank parallel groups without initializing a process group. + + This is faster than the `dist_init` fixture. + """ + for name in ("_WORLD", "_TP", "_ETP", "_DCP", "_PP", "_DP", "_EP", "_PCP"): + monkeypatch.setattr( + f"vllm.distributed.parallel_state.{name}", _SingleRankGroup() ) + monkeypatch.setattr( + "vllm.model_executor.parameter.get_tensor_model_parallel_rank", lambda: 0 + ) + monkeypatch.setattr( + "vllm.model_executor.parameter.get_tensor_model_parallel_world_size", lambda: 1 + ) @pytest.mark.parametrize( "model_type", - ["minimax_m3", "deepseek_v4", "qwen3_5", "glm4_moe", "deepseek_v2", "glm5_next"], + [ + "minimax_m3", + "deepseek_v4", + "deepseek_v4_cpu", + "qwen3_next", + "qwen3_5", + "alice_ai", + "glm4_moe", + "deepseek_v2", + "AXK1", + "glm4_moe_lite", + "deepseek_v32", + "glm5_next", + "qwen4_exp", + ], ) +@pytest.mark.parametrize("shared_expert_precision", ["all_mxfp4", "mixed"]) @pytest.mark.parametrize( ("use_fse", "exclude"), [ @@ -933,17 +1430,82 @@ def __init__(self, enabled: bool) -> None: ) def test_models_fse_init( model_type: str, + shared_expert_precision: str, use_fse: bool, exclude: list[str], - dist_init: None, + single_rank_parallel_state: None, monkeypatch: pytest.MonkeyPatch, ) -> None: - """Model construction resolves FSE consistently with Quark quantization.""" + """Construction and loading preserve each shared expert's precision.""" + _print_warning_once.cache_clear() + per_layer_fse = model_type != "deepseek_v4_cpu" + moe_attr = { + "minimax_m3": "block_sparse_moe", + "deepseek_v4": "ffn", + }.get(model_type, "mlp") + shared_name = ( + "shared_experts" + if model_type + in { + "minimax_m3", + "deepseek_v4", + "glm4_moe", + "deepseek_v2", + "AXK1", + "glm4_moe_lite", + "deepseek_v32", + "glm5_next", + } + else "shared_expert" + ) + if shared_expert_precision == "mixed" and not per_layer_fse: + pytest.skip( + "The mixed Quark two-layer setup does not apply to this architecture" + ) quantization_config: dict[str, Any] = ( - get_deepseek_v4_quark_config(["layers.0.ffn.shared_experts"] if exclude else []) + get_deepseek_v4_quark_config( + [f"layers.{idx}.ffn.shared_experts" for idx in range(2)] if exclude else [] + ) if model_type == "deepseek_v4" - else {**_QUARK_FSE_CONFIG, "exclude": exclude} - ) + else {**deepcopy(_QUARK_FSE_CONFIG), "exclude": exclude} + ) + if shared_expert_precision == "mixed": + mxfp8_config = deepcopy(quantization_config["global_quant_config"]) + mxfp8_config["input_tensors"].update( + dtype="fp8_e4m3", + symmetric=True, + scale_type=None, + scale_format=None, + scale_calculation_mode=None, + observer_cls="PerGroupMinMaxObserver", + zero_point_type="int32", + ) + mxfp8_config["weight"].update( + dtype="fp8_e4m3", + qscheme="per_block", + ch_axis=None, + group_size=None, + block_size=[1, 32], + symmetric=True, + scale_type="float8_e8m0fnu", + scale_format=None, + scale_calculation_mode=None, + observer_cls="PerBlock2DMinMaxObserver", + zero_point_type="int32", + ) + if model_type == "deepseek_v4": + quantization_config.setdefault("layer_quant_config", {})[ + "layers.1.ffn.shared_experts.w1" + ] = deepcopy(mxfp8_config) + else: + quantization_config.setdefault("layer_quant_config", {}).update( + { + f"model.layers.1.{moe_attr}.{shared_name}.{projection}": deepcopy( + mxfp8_config + ) + for projection in ("gate_proj", "up_proj", "down_proj") + } + ) config, model_constructor = get_fse_test_model_config( model_type, quantization_config @@ -961,11 +1523,12 @@ def test_models_fse_init( quantization_config=None, runner_type="generate", is_moe=True, + use_mla=False, logits_processors=None, rswa_window=None, ) vllm_config.parallel_config.enable_expert_parallel = False - if model_type == "deepseek_v4": + if model_type in ("deepseek_v4", "deepseek_v4_cpu"): from vllm.models.deepseek_v4.quant_config import DeepseekV4FP8Config vllm_config.cache_config.cache_dtype = "fp8_ds_mla" @@ -975,6 +1538,10 @@ def test_models_fse_init( ) else: vllm_config.quant_config = QuarkConfig(quantization_config) + if per_layer_fse: + vllm_config.quant_config.packed_modules_mapping = { + "gate_up_proj": ["gate_proj", "up_proj"] + } if model_type == "glm5_next": from vllm.models.glm5next.common import model as glm5_next_model @@ -985,18 +1552,40 @@ def test_models_fse_init( monkeypatch.setattr( glm5_next_model, "_fused_shared_experts_tuned", lambda _: True ) + from vllm.models.glm5next.common import mtp as glm5_next_mtp_model + + # The MTP predictor asserts on the concrete MLA attention class and + # reads ``mla_attn``, so patch its imported name to the same stub. + monkeypatch.setattr( + glm5_next_mtp_model, "Glm5NextMLAAttention", _StubGlm5NextAttention + ) + if model_type == "glm4_moe_lite": + from vllm.model_executor.models import glm4_moe_lite as glm4_moe_lite_model + + # Only the MoE/shared-expert wiring is under test here; the MLA backend + # does not accept this test's default dtype. + for attention in ("Glm4MoeLiteAttention", "Glm4MoeLiteMLAAttention"): + monkeypatch.setattr(glm4_moe_lite_model, attention, _StubGlm5NextAttention) + if model_type == "deepseek_v32": + from vllm.models.deepseek_v32.amd import model as deepseek_v32_model + + # The sparse-MLA attention needs a real MLA backend; only the MoE + # wiring is under test here. + monkeypatch.setattr( + deepseek_v32_model, "DeepseekV32MLAAttention", _StubGlm5NextAttention + ) + if model_type == "qwen4_exp": + from vllm.models.qwen4_exp.amd import model as qwen4_exp_model + + # Only the MoE/shared-expert wiring is under test here. + monkeypatch.setattr( + qwen4_exp_model, "Qwen3NextAttention", _StubGlm5NextAttention + ) import vllm.envs as envs from vllm._aiter_ops import rocm_aiter_ops - if model_type == "deepseek_v4": - from vllm.models.deepseek_v4.amd import model as deepseek_v4_model - - warning_logger = deepseek_v4_model.logger - elif model_type == "minimax_m3": - warning_logger = minimax_m3_model.logger - else: - warning_logger = fused_moe_utils.logger + warning_logger = fused_moe_utils.logger with monkeypatch.context() as mp: mp.setenv("VLLM_ROCM_USE_AITER", str(use_fse)) @@ -1010,33 +1599,196 @@ def test_models_fse_init( importlib.reload(envs) mp.setattr(envs, "VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS", use_fse) rocm_aiter_ops.refresh_env_variables() + from vllm.model_executor.layers.fused_moe.experts import rocm_aiter_moe + + mp.setattr(rocm_aiter_moe, "aiter_topK_meta_data", None) + rocm_aiter_moe.init_aiter_topK_meta_data.cache_clear() aiter_fse_enabled = bool(rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()) - # These AMD-specific models currently use the raw environment flag. - # and do not rely on `rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()`. - fse_enabled = use_fse and ( - aiter_fse_enabled or model_type in {"deepseek_v4", "minimax_m3"} - ) + fse_enabled = use_fse and aiter_fse_enabled with ( patch.object(warning_logger, "warning") as warning, set_current_vllm_config(vllm_config), ): - model = model_constructor(vllm_config=vllm_config) - mtp = None - if model_type == "deepseek_v4" and use_fse and not exclude: - from vllm.models.deepseek_v4.amd.mtp import ( - DeepSeekV4MTP, + model = model_constructor( + vllm_config=vllm_config, + **({"prefix": "model"} if per_layer_fse else {}), + ) + mtps = { + model_type: _build_mtp_for_model( + model_type, vllm_config, config, use_fse and not exclude ) - - vllm_config.speculative_config = SimpleNamespace( - draft_model_config=SimpleNamespace(hf_config=config), - method="mtp", - parallel_drafting=False, - enable_adaptive_verification=False, - use_dspark=lambda: False, + } + mtp = mtps.get("deepseek_v4") + if per_layer_fse: + expected_policy = { + f"layers.0.{moe_attr}": bool(fse_enabled and not exclude), + f"layers.1.{moe_attr}": bool( + fse_enabled + and not exclude + and shared_expert_precision == "all_mxfp4" + ), + } + assert model.fused_shared_expert_layers == frozenset( + idx for idx in range(2) if expected_policy[f"layers.{idx}.{moe_attr}"] + ) + if any(expected_policy.values()) and not all(expected_policy.values()): + with pytest.raises(NotImplementedError, match="EPLB"): + get_fused_shared_expert_layers( + model, + type(getattr(model.layers[0], moe_attr)), + enable_eplb=True, + ) + for idx, layer in enumerate(model.layers): + fused = expected_policy[f"layers.{idx}.{moe_attr}"] + moe = getattr(layer, moe_attr) + assert (getattr(moe, shared_name) is None) is fused + assert moe.experts.expert_map_manager.num_fused_shared_experts == int( + fused + ) + assert ( + moe.experts.moe_config.is_fused_shared_expert_requested + is aiter_fse_enabled + ) + if not fused and not exclude and model_type != "deepseek_v4": + scheme = getattr(moe, shared_name).gate_up_proj.scheme + expected_dtype = ( + "mxfp8_e4m3" + if idx == 1 and shared_expert_precision == "mixed" + else "mxfp4" + ) + assert scheme.weight_dtype == expected_dtype + + checkpoint = [] + expected_weights = [] + for idx, layer in enumerate(model.layers): + fused = expected_policy[f"layers.{idx}.{moe_attr}"] + if model_type == "deepseek_v4": + dsv4_checkpoint, dsv4_expected = ( + _make_deepseek_v4_shared_expert_checkpoint( + layer, idx, fused, exclude + ) + ) + checkpoint.extend(dsv4_checkpoint) + expected_weights.extend(dsv4_expected) + continue + high_precision = idx == 1 and shared_expert_precision == "mixed" + dtype = ( + torch.bfloat16 + if exclude + else torch.float8_e4m3fn + if high_precision + else torch.uint8 ) - mtp = DeepSeekV4MTP(vllm_config=vllm_config) - assert model.is_fused_shared_expert_enabled is (fse_enabled and not exclude) + pack_factor = 1 if exclude or high_precision else 2 + for projection in ("gate_proj", "up_proj", "down_proj"): + shape = ( + (128, 32 // pack_factor) + if projection == "down_proj" + else (32, 128 // pack_factor) + ) + weight = torch.full(shape, idx + 1).to(dtype) + name = f"layers.{idx}.{moe_attr}.{shared_name}.{projection}" + checkpoint.append((f"{name}.weight", weight)) + if not exclude: + scale_shape = (128, 1) if projection == "down_proj" else (32, 4) + checkpoint.append( + ( + f"{name}.weight_scale", + torch.full(scale_shape, 127, dtype=torch.uint8), + ) + ) + if fused: + routed = getattr(layer, moe_attr).experts.routed_experts + target = ( + routed.w2_weight[2] + if projection == "down_proj" + else routed.w13_weight[2].chunk(2, dim=0)[ + int(projection == "up_proj") + ] + ) + else: + shared = getattr(getattr(layer, moe_attr), shared_name) + target = ( + shared.down_proj.weight + if projection == "down_proj" + else shared.gate_up_proj.weight.chunk(2, dim=0)[ + int(projection == "up_proj") + ] + ) + expected_weights.append((target, weight)) + if shared_name == "shared_expert": + checkpoint.append( + ( + f"layers.{idx}.{moe_attr}.shared_expert_gate.weight", + torch.full((1, 128), idx + 1, dtype=torch.bfloat16), + ) + ) + with set_current_vllm_config(vllm_config): + loaded = model.load_weights(checkpoint) + base = config.num_hidden_layers + mtp_cases = [ + ( + mtp_model, + f"model.layers.{base}.mlp.{shared}", + mtp_model.model.layers[str(base)].mtp_block.mlp, + getattr(config, "n_routed_experts", 0), + ) + for mtp_model, shared in ( + (mtps.get("glm4_moe"), "shared_experts"), + (mtps.get("glm4_moe_lite"), "shared_experts"), + (mtps.get("deepseek_v2"), "shared_experts"), + (mtps.get("deepseek_v32"), "shared_experts"), + (mtps.get("glm5_next"), "shared_experts"), + ) + if mtp_model is not None + ] + mtp_cases += [ + ( + mtp_model, + "mtp.layers.0.mlp.shared_expert", + mtp_model.model.layers[0].mlp, + getattr(config, "num_experts", 0), + ) + for mtp_model in ( + mtps.get("qwen3_next"), + mtps.get("qwen3_5"), + ) + if mtp_model is not None + ] + for mtp_model, ckpt_prefix, moe, n_routed in mtp_cases: + if fse_enabled and not exclude: + _assert_mtp_fused_shared_expert_redirect( + mtp_model, ckpt_prefix, moe, n_routed + ) + else: + _assert_mtp_shared_expert_not_fused( + mtp_model, ckpt_prefix, moe, n_routed + ) + for idx in range(2): + prefix = f"layers.{idx}.{moe_attr}" + expected_name = ( + f"{prefix}.experts.routed_experts.w2_weight" + if expected_policy[prefix] + else f"{prefix}.{shared_name}.down_proj.weight" + ) + assert expected_name in loaded + if shared_name == "shared_expert": + assert f"{prefix}.shared_expert_gate.weight" in loaded + for target, weight in expected_weights: + torch.testing.assert_close( + target[: weight.shape[0], : weight.shape[1]].float(), weight.float() + ) + elif model_type == "deepseek_v4_cpu": + # CPU never supports fused shared experts; the separate MLP is + # always built and no routed slot is appended. + for layer in model.layers: + assert layer.ffn.shared_experts is not None + assert ( + layer.ffn.experts.expert_map_manager.num_fused_shared_experts == 0 + ) + else: + assert model.is_fused_shared_expert_enabled is (fse_enabled and not exclude) if model_type == "glm5_next": assert (model.layers[0].mlp.shared_experts is None) is ( fse_enabled and not exclude @@ -1046,21 +1798,53 @@ def test_models_fse_init( # so should just raise a warning. if mtp is not None: assert not mtp.model.layers[ - "1" + str(config.num_hidden_layers) ].mtp_block.ffn.is_fused_shared_expert_enabled + if aiter_fse_enabled: + warning.assert_called_once() + assert ( + "DeepSeek-V4/V4.1 shared experts at " + in (warning.call_args.args[1]) + ) + else: + # Without AITER nothing fuses, so no skip reason is reported. + warning.assert_not_called() + if aiter_fse_enabled and shared_expert_precision == "mixed" and not exclude: warning.assert_called_once() - assert ( - "DeepSeek-V4/V4.1 shared experts at mtp.0.ffn.shared_experts" - in (warning.call_args.args[1]) - ) - if aiter_fse_enabled and exclude: + if aiter_fse_enabled and shared_expert_precision == "mixed" and not exclude: warning.assert_called_once() - assert ( - "VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS is enabled" - in (warning.call_args.args[0]) - ) - assert "excludes shared experts" in warning.call_args.args[1] + if model_type == "deepseek_v4": + assert "are not MXFP4" in warning.call_args.args[1] + else: + assert ( + "different quantization configurations" in warning.call_args.args[1] + ) + for mtp_model in ( + mtps.get("glm4_moe"), + mtps.get("deepseek_v2"), + mtps.get("glm4_moe_lite"), + mtps.get("deepseek_v32"), + mtps.get("glm5_next"), + ): + if mtp_model is not None: + assert mtp_model.fused_shared_expert_layers == ( + frozenset({config.num_hidden_layers}) + if fse_enabled + else frozenset() + ) + if aiter_fse_enabled and exclude: + if model_type == "deepseek_v4_cpu": + # CPU has no FSE path, so the resolver never warns. + warning.assert_not_called() + else: + assert warning.call_count == 1 + assert ( + "VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS is enabled" + in (warning.call_args.args[0]) + ) + assert "excludes shared experts" in warning.call_args.args[1] + rocm_aiter_moe.init_aiter_topK_meta_data.cache_clear() importlib.reload(envs) rocm_aiter_ops.refresh_env_variables() diff --git a/tests/model_executor/model_loader/test_reload.py b/tests/model_executor/model_loader/test_reload.py index eac5245bd988..01763dfb831c 100644 --- a/tests/model_executor/model_loader/test_reload.py +++ b/tests/model_executor/model_loader/test_reload.py @@ -720,6 +720,7 @@ def test_padded_moe_reload_releases_each_layer( has_bias=has_bias, tp_rank=tp_rank, tp_shard_with_padding=False, + num_experts=experts, moe_parallel_config=SimpleNamespace(tp_size=2), ) model = torch.nn.ModuleList() @@ -735,7 +736,9 @@ def test_padded_moe_reload_releases_each_layer( layer.moe_config = config layer.quant_config = None layer.quant_method = method - layer.expert_map_manager = SimpleNamespace(map_global_to_local=lambda i: i) + layer.expert_map_manager = SimpleNamespace( + map_global_to_local=lambda i: i, num_fused_shared_experts=0 + ) layer._loaded_expert_biases = set() layer._fused_shared_expert_quantizer = None method.create_weights( diff --git a/tests/models/test_deepseek_v4_vl_rocm.py b/tests/models/test_deepseek_v4_vl_rocm.py index 762f1cd6c0e0..418a2b47c3a5 100644 --- a/tests/models/test_deepseek_v4_vl_rocm.py +++ b/tests/models/test_deepseek_v4_vl_rocm.py @@ -249,19 +249,71 @@ def fake_factory(**kwargs): def test_rocm_fse_rejects_data_parallel_deployment( monkeypatch: pytest.MonkeyPatch, ) -> None: - from vllm.models.deepseek_v4.amd.model import _fuse_shared_experts_enabled + from vllm.config import set_current_vllm_config + from vllm.model_executor.layers.fused_moe.utils import ( + resolve_layer_fused_shared_expert, + ) + from vllm.models.deepseek_v4.amd import model as rocm_model - config = SimpleNamespace(n_shared_experts=1) + class FakeGate(nn.Module): + def __init__(self, **kwargs) -> None: + super().__init__() - monkeypatch.setenv("VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS", "1") - with pytest.raises(ValueError, match="data parallelism"): - _fuse_shared_experts_enabled(config, ParallelConfig(data_parallel_size=2)) + monkeypatch.setattr(rocm_model, "GateLinear", FakeGate) + monkeypatch.setattr(rocm_model, "FusedMoEFactory", lambda **kwargs: nn.Identity()) + monkeypatch.setattr(rocm_model, "get_tensor_model_parallel_world_size", lambda: 1) + monkeypatch.setattr(rocm_model, "get_tensor_model_parallel_rank", lambda: 0) + monkeypatch.setattr( + rocm_model.rocm_aiter_ops, + "is_fusion_moe_shared_experts_enabled", + lambda: True, + ) - # Expert parallelism already disables FSE: no raise, FSE stays off. - assert not _fuse_shared_experts_enabled( - config, ParallelConfig(data_parallel_size=2, enable_expert_parallel=True) + config = SimpleNamespace( + hidden_size=16, + n_routed_experts=4, + num_experts_per_tok=2, + moe_intermediate_size=8, + swiglu_limit=None, + norm_topk_prob=True, + scoring_func="sqrtsoftplus", + num_hash_layers=0, + vocab_size=32, + topk_method="noaux_tc", + vision_n_layers=0, + n_shared_experts=1, ) - assert _fuse_shared_experts_enabled(config, ParallelConfig()) + + def build_vllm_config(parallel_config: ParallelConfig) -> SimpleNamespace: + return SimpleNamespace( + model_config=SimpleNamespace(hf_config=config), + quant_config=None, + parallel_config=parallel_config, + kernel_config=KernelConfig(), + ) + + # Data parallelism without expert parallelism fails fast at MoE construction. + with pytest.raises(ValueError, match="not supported with data parallelism"): + rocm_model.DeepseekV4MoE( + build_vllm_config(ParallelConfig(data_parallel_size=2)), + prefix="model.layers.0.ffn", + ) + + # Under expert parallelism the resolver disables FSE instead of raising. + with set_current_vllm_config( + build_vllm_config( + ParallelConfig(data_parallel_size=2, enable_expert_parallel=True) + ) + ): + assert not resolve_layer_fused_shared_expert( + None, "model.layers.0.ffn", disable_under_ep=True + ) + + # Nothing blocks FSE once expert parallelism is off. + with set_current_vllm_config(build_vllm_config(ParallelConfig())): + assert resolve_layer_fused_shared_expert( + None, "model.layers.0.ffn", disable_under_ep=True + ) def test_rocm_mtp_forwards_input_ids_for_vision_routing( diff --git a/tests/quantization/test_auto_gptq.py b/tests/quantization/test_auto_gptq.py index 4f60a974caed..f0df2aefe10b 100644 --- a/tests/quantization/test_auto_gptq.py +++ b/tests/quantization/test_auto_gptq.py @@ -166,8 +166,10 @@ class Loader: moe_config = SimpleNamespace( is_act_and_mul=True, tp_rank=0, + num_experts=1, moe_parallel_config=SimpleNamespace(tp_size=1), ) + expert_map_manager = SimpleNamespace(num_fused_shared_experts=0) _get_hidden_dim = staticmethod(RoutedExperts._get_hidden_dim) _narrow_expert_data_for_padding = staticmethod( RoutedExperts._narrow_expert_data_for_padding diff --git a/tests/quantization/test_online_shared_expert_fusion.py b/tests/quantization/test_online_shared_expert_fusion.py index 394ba04a1b30..c10ba5e38046 100644 --- a/tests/quantization/test_online_shared_expert_fusion.py +++ b/tests/quantization/test_online_shared_expert_fusion.py @@ -573,18 +573,19 @@ def record_warning(message: str, *args: object, **_kwargs: object) -> None: "cannot be enabled" in warning for warning in logged_warnings ) - if architecture == "MiniMaxM3SparseForCausalLM": - # TODO: Use resolve_layer_fused_shared_expert in MiniMax M3 and remove - # this branch once its virtual fused projections are registered. - assert logged_messages == ["Quantizing 0 layers of types: "] - else: - assert logged_messages == [ - "Quantizing 2 layers of types: " - f"mlp.{shared_expert_name}.down_proj: 1 " - f"(from targets: {target_pattern}, mxfp4); " - f"mlp.{shared_expert_name}.gate_up_proj: 1 " - f"(from targets: {target_pattern}, mxfp4)" - ] + # MiniMax M3 now resolves fused shared experts per layer through + # `resolve_layer_fused_shared_expert`, so its virtual fused projections are + # registered and quantized like any other architecture. + moe_prefix = ( + "block_sparse_moe" if architecture == "MiniMaxM3SparseForCausalLM" else "mlp" + ) + assert logged_messages == [ + "Quantizing 2 layers of types: " + f"{moe_prefix}.{shared_expert_name}.down_proj: 1 " + f"(from targets: {target_pattern}, mxfp4); " + f"{moe_prefix}.{shared_expert_name}.gate_up_proj: 1 " + f"(from targets: {target_pattern}, mxfp4)" + ] def test_online_gate_up_proj_target_disables_fse( @@ -645,6 +646,6 @@ def record_warning(message: str, *args: object, **_kwargs: object) -> None: ] assert logged_warnings == [ "VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS is enabled but cannot be " - "enabled - skipping for this layer: online quantization targets only " - "part of the shared expert at model.layers.0.mlp.shared_experts." + "enabled - skipping shared-expert fusion: online quantization targets only " + "part of the shared expert at ." ] diff --git a/tests/v1/worker/test_spec_decode_embed_sharing_pp.py b/tests/v1/worker/test_spec_decode_embed_sharing_pp.py index 0fef7bb47bdf..901a4d184c73 100644 --- a/tests/v1/worker/test_spec_decode_embed_sharing_pp.py +++ b/tests/v1/worker/test_spec_decode_embed_sharing_pp.py @@ -7,6 +7,7 @@ import torch import torch.nn as nn +from vllm.config import ParallelConfig from vllm.model_executor.model_loader.mtp_validation import ( disable_mtp_completeness_check, ) @@ -158,6 +159,7 @@ def m3_config(): speculative_config=SimpleNamespace( method="mtp", draft_model_config=SimpleNamespace(hf_config=hf_config) ), + parallel_config=ParallelConfig(), quant_config=None, cache_config=None, kernel_config=SimpleNamespace(monokernel=False), diff --git a/vllm/model_executor/layers/fused_moe/config.py b/vllm/model_executor/layers/fused_moe/config.py index e8bd822becab..83df5d16c16e 100644 --- a/vllm/model_executor/layers/fused_moe/config.py +++ b/vllm/model_executor/layers/fused_moe/config.py @@ -1338,7 +1338,7 @@ class FusedMoEConfig: # Use the allocated width as the checkpoint TP stride, including for scales. tp_shard_with_padding: bool = False rocm_aiter_fmoe_enabled: bool = False - aiter_fmoe_shared_expert_enabled: bool = False + is_fused_shared_expert_requested: bool = False # Whether to force MXFP4 (a4w4) MoE activations for DeepSeek V4.1 on # ROCm/AITER. Opt-in via VLLM_ROCM_USE_AITER_MOE_A4W4_DSV4=1; rejected # for any other model type. Resolved here, not in the forward path, @@ -1371,7 +1371,7 @@ def __post_init__(self): if self.is_act_and_mul: self.rocm_aiter_fmoe_enabled = rocm_aiter_ops.is_fused_moe_enabled() - self.aiter_fmoe_shared_expert_enabled = ( + self.is_fused_shared_expert_requested = bool( rocm_aiter_ops.is_fusion_moe_shared_experts_enabled() ) @@ -1402,7 +1402,7 @@ def __post_init__(self): assert self.rocm_aiter_fmoe_enabled, ( "Mori needs to be used with aiter fused_moe for now." ) - assert not self.aiter_fmoe_shared_expert_enabled, ( + assert not self.is_fused_shared_expert_requested, ( "Mori does not support fusion shared expert now. " "Turn it off by setting VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS=0" ) diff --git a/vllm/model_executor/layers/fused_moe/routed_experts.py b/vllm/model_executor/layers/fused_moe/routed_experts.py index 59aea0599b4e..3c4a1ffc7707 100644 --- a/vllm/model_executor/layers/fused_moe/routed_experts.py +++ b/vllm/model_executor/layers/fused_moe/routed_experts.py @@ -725,6 +725,24 @@ def weight_loader( quant_method_name = base_quant_method.__class__.__name__ global_expert_id = expert_id + num_global_experts = ( + self.moe_config.num_experts + + self.expert_map_manager.num_fused_shared_experts + ) + # `num_experts` counts logical and redundant experts; fused shared + # slots are appended after them. + # Expert substitution checkpoints name retained experts by their + # logical ID, which is not bounded by the number of allocated slots. + expert_substitution = getattr(self, "expert_substitution", None) + if expert_substitution is not None: + num_global_experts = ( + expert_substitution.num_logical_experts + + self.expert_map_manager.num_fused_shared_experts + ) + assert global_expert_id < num_global_experts, ( + f"Expert id {global_expert_id} is out of range for " + f"{self.layer_name}; expected < {num_global_experts}" + ) expert_id = self._map_global_expert_id_to_local_expert_id(global_expert_id) use_global_sf = ( diff --git a/vllm/model_executor/layers/fused_moe/utils.py b/vllm/model_executor/layers/fused_moe/utils.py index 27df12668584..6a4911e8387f 100644 --- a/vllm/model_executor/layers/fused_moe/utils.py +++ b/vllm/model_executor/layers/fused_moe/utils.py @@ -1,13 +1,12 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import functools -from collections.abc import Iterable from math import prod from typing import TYPE_CHECKING +import regex as re import torch import torch.nn.functional as F -from torch import nn import vllm.envs as envs from vllm import _custom_ops as ops @@ -40,7 +39,6 @@ from vllm.model_executor.layers.quantization.utils.w8a8_utils import ( per_tensor_dequantize, ) -from vllm.model_executor.models.utils import PPMissingLayer from vllm.platforms import current_platform from vllm.triton_utils import tl, triton from vllm.utils.math_utils import cdiv @@ -52,11 +50,23 @@ logger = init_logger(__name__) +# Matches any layer-qualified module path (e.g. +# ``model.layers.5.mlp.shared_expert`` or +# ``model.language_model.model.layers.5.mlp.shared_expert``) so skip reasons +# are reported once per cause instead of once per layer. +_LAYER_MODULE_PATH_RE = re.compile(r"(?:[\w]+\.)*(?:layers|mtp)\.\d+\.[\w.]*") + + +def _normalize_skip_reason(reason: str) -> str: + return _LAYER_MODULE_PATH_RE.sub("", reason) + + def resolve_layer_fused_shared_expert( quant_config: "QuantizationConfig | None", prefix: str, shared_expert_name: str = "shared_experts", moe_backend: str | None = None, + disable_under_ep: bool = False, ) -> bool: """Resolve whether AITER fused shared-expert execution is enabled. @@ -65,6 +75,9 @@ def resolve_layer_fused_shared_expert( prefix: MoE module prefix. shared_expert_name: Shared-expert module name under ``prefix``. moe_backend: MoE backend that runs the fused shared expert. + disable_under_ep: Whether this model family does not support FSE under + expert parallelism. Kept as an explicit call-site policy instead of + branching on model type here. Returns: Whether AITER fused shared experts are enabled. @@ -76,16 +89,31 @@ def resolve_layer_fused_shared_expert( # NOTE: is_fusion_moe_shared_experts_enabled is decorated with @if_aiter_supported # that returns None if AITER is not available. fse_requested = bool(rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()) - fse_compatible, fse_reason = ( - is_shared_expert_quant_fse_compatible( + + # TODO: should it be disabled for all models? + skip_reason: str | None = None + if fse_requested and disable_under_ep: + from vllm.config import get_current_vllm_config_or_none + + vllm_config = get_current_vllm_config_or_none() + parallel_config = ( + vllm_config.parallel_config if vllm_config is not None else None + ) + if parallel_config is not None and parallel_config.enable_expert_parallel: + fse_requested = False + skip_reason = "expert parallelism is enabled" + + if fse_requested: + fse_compatible, fse_reason = is_shared_expert_quant_fse_compatible( quant_config, f"{prefix}.experts", f"{prefix}.{shared_expert_name}", moe_backend=moe_backend, ) - if fse_requested - else (True, None) - ) + if not fse_compatible: + skip_reason = fse_reason or "unknown reason" + else: + fse_compatible = True is_fused_shared_expert_enabled = fse_requested and fse_compatible # In online shared_expert quantization is used, register it to @@ -121,52 +149,18 @@ def resolve_layer_fused_shared_expert( quant_key_str, target_pattern, ) - if fse_requested and not is_fused_shared_expert_enabled: - logger.warning( + if skip_reason is not None: + # Keep the full per-layer reason at debug level; the warning below is + # deduplicated by cause so a long model logs it once. + logger.debug("Skipping shared-expert fusion at %s: %s", prefix, skip_reason) + logger.warning_once( "VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS is enabled but " - "cannot be enabled - skipping for this layer: %s.", - fse_reason, + "cannot be enabled - skipping shared-expert fusion: %s.", + _normalize_skip_reason(skip_reason), ) return is_fused_shared_expert_enabled -def is_model_fused_shared_expert_compatible( - layers: nn.ModuleList | Iterable[nn.Module], - moe_cls: type[nn.Module], - moe_name: str, -) -> bool: - """Resolve one fused-shared-expert state for a model's MoE layers.""" - - def get_moe_layer(layer: nn.Module) -> nn.Module | None: - for name in moe_name.split("."): - layer = getattr(layer, name, None) - if layer is None: - return None - return layer - - moe_layers = ( - moe_layer - for layer in layers - if not isinstance(layer, PPMissingLayer) - and (moe_layer := get_moe_layer(layer)) is not None - and isinstance(moe_layer, moe_cls) - ) - - enabled = [ - getattr(layer, "is_fused_shared_expert_enabled", False) for layer in moe_layers - ] - enabled_count = sum(enabled) - disabled_count = len(enabled) - enabled_count - if enabled_count > 0 and disabled_count > 0: - raise NotImplementedError( - "Fused shared experts must be enabled for all MoE layers; found " - f"{enabled_count} enabled and {disabled_count} disabled layers. " - "Per-layer fused shared experts is not yet supported. Please open " - "an issue." - ) - return enabled_count > 0 and disabled_count == 0 - - @triton.jit def _count_expert_num_tokens( topk_ids_ptr, diff --git a/vllm/model_executor/models/AXK1.py b/vllm/model_executor/models/AXK1.py index 6a12fc88b255..cf0828e9951d 100644 --- a/vllm/model_executor/models/AXK1.py +++ b/vllm/model_executor/models/AXK1.py @@ -49,7 +49,6 @@ fused_moe_make_expert_params_mapping, ) from vllm.model_executor.layers.fused_moe.utils import ( - is_model_fused_shared_expert_compatible, resolve_layer_fused_shared_expert, ) from vllm.model_executor.layers.layernorm import RMSNorm @@ -84,6 +83,8 @@ from .utils import ( AutoWeightsLoader, PPMissingLayer, + extract_layer_index, + get_fused_shared_expert_layers, get_spec_layer_idx_from_weight_name, is_pp_missing_parameter, make_empty_intermediate_tensors_factory, @@ -742,10 +743,10 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): prefix=f"{prefix}.layers", ) - self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible( - self.layers, + self.fused_shared_expert_layers = get_fused_shared_expert_layers( + self, AXK1MoE, - "mlp", + enable_eplb=vllm_config.parallel_config.enable_eplb, ) if get_pp_group().is_last_rank: @@ -839,10 +840,11 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: ckpt_gate_proj_name="gate_proj", ckpt_down_proj_name="down_proj", ckpt_up_proj_name="up_proj", + # NOTE: global despite per-layer FSE; non-fused layers skip the extras. num_experts=self.config.n_routed_experts + ( (self.config.n_shared_experts or 0) - if self.is_fused_shared_expert_enabled + if self.fused_shared_expert_layers else 0 ), num_redundant_experts=self.num_redundant_experts, @@ -858,8 +860,11 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: if spec_layer is not None: continue # skip spec decode layers for main model + # Only parse the layer index when FSE is enabled. is_fusion_moe_shared_experts_layer = ( - self.is_fused_shared_expert_enabled and ("mlp.shared_experts" in name) + "mlp.shared_experts" in name + and bool(rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()) + and extract_layer_index(name) in self.fused_shared_expert_layers ) for param_name, weight_name, shard_id in stacked_params_mapping: diff --git a/vllm/model_executor/models/alice_ai.py b/vllm/model_executor/models/alice_ai.py index a0c5f53ecf92..9dfdab545916 100644 --- a/vllm/model_executor/models/alice_ai.py +++ b/vllm/model_executor/models/alice_ai.py @@ -20,7 +20,6 @@ ) from vllm.model_executor.layers.fused_moe import FusedMoEFactory, GateLinear from vllm.model_executor.layers.fused_moe.utils import ( - is_model_fused_shared_expert_compatible, resolve_layer_fused_shared_expert, ) from vllm.model_executor.layers.layernorm import GemmaRMSNorm @@ -57,6 +56,7 @@ AutoWeightsLoader, WeightsMapper, extract_layer_index, + get_fused_shared_expert_layers, make_layers, maybe_fuse_shared_experts, maybe_prefix, @@ -500,8 +500,10 @@ def get_layer(prefix: str) -> AliceAIDecoderBlock: layer_fn=get_layer, prefix=f"{prefix}.layers", ) - self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible( - layers=self.layers, moe_cls=AliceAISparseMoeBlock, moe_name="mlp" + self.fused_shared_expert_layers = get_fused_shared_expert_layers( + self, + AliceAISparseMoeBlock, + enable_eplb=vllm_config.parallel_config.enable_eplb, ) self.norm = GemmaRMSNorm( hidden_size=config.hidden_size, eps=config.rms_norm_eps @@ -567,7 +569,7 @@ def forward( def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: weights = maybe_fuse_shared_experts( weights=weights, - enabled=self.is_fused_shared_expert_enabled, + fused_shared_expert_layers=self.fused_shared_expert_layers, n_routed_experts=self.config.num_experts, n_shared_experts=1, ckpt_prefix="mlp.shared_expert", diff --git a/vllm/model_executor/models/deepseek_mtp.py b/vllm/model_executor/models/deepseek_mtp.py index 4f6125a2279b..c870660935bb 100644 --- a/vllm/model_executor/models/deepseek_mtp.py +++ b/vllm/model_executor/models/deepseek_mtp.py @@ -7,15 +7,13 @@ import torch.nn as nn from transformers import PreTrainedConfig +from vllm._aiter_ops import rocm_aiter_ops from vllm.compilation.decorators import support_torch_compile from vllm.config import VllmConfig from vllm.distributed import tensor_model_parallel_all_gather from vllm.model_executor.layers.fused_moe import ( fused_moe_make_expert_params_mapping, ) -from vllm.model_executor.layers.fused_moe.utils import ( - is_model_fused_shared_expert_compatible, -) from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.logits_processor import LogitsProcessor from vllm.model_executor.layers.quantization import QuantizationConfig @@ -41,6 +39,8 @@ ) from .interfaces import SupportsPP from .utils import ( + extract_layer_index, + get_fused_shared_expert_layers, get_pp_missing_layer_names, get_spec_layer_idx_from_weight_name, make_empty_intermediate_tensors_factory, @@ -242,6 +242,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__() self.config = vllm_config.model_config.hf_config self.quant_config = vllm_config.quant_config + self.vllm_config = vllm_config self.model = DeepSeekMultiTokenPredictor( vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") ) @@ -267,10 +268,10 @@ def set_moe_parameters(self): self.moe_mlp_layers.append(layer.mlp) self.moe_layers.append(layer.mlp.experts) self.extract_moe_parameters(example_moe) - self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible( - self.model.layers.values(), + self.fused_shared_expert_layers = get_fused_shared_expert_layers( + self, DeepseekV2MoE, - "mtp_block.mlp", + enable_eplb=self.vllm_config.parallel_config.enable_eplb, ) def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: @@ -316,17 +317,15 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: ] stacked_params_mapping.extend(indexer_fused_mapping) + fused_shared_expert_layers = self.fused_shared_expert_layers expert_params_mapping = fused_moe_make_expert_params_mapping( self, ckpt_gate_proj_name="gate_proj", ckpt_down_proj_name="down_proj", ckpt_up_proj_name="up_proj", + # NOTE: global despite per-layer FSE; non-fused layers skip the extras. num_experts=self.config.n_routed_experts - + ( - self.config.n_shared_experts - if self.is_fused_shared_expert_enabled - else 0 - ), + + (self.config.n_shared_experts if fused_shared_expert_layers else 0), ) pp_missing_layer_names = get_pp_missing_layer_names(self) @@ -353,8 +352,11 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: weight_loader(param, loaded_weight) loaded_params.add(name) continue + # Only parse the layer index when FSE is enabled. is_fusion_moe_shared_experts_layer = ( - self.is_fused_shared_expert_enabled and ("mlp.shared_experts" in name) + "mlp.shared_experts" in name + and bool(rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()) + and extract_layer_index(name) in fused_shared_expert_layers ) name = self._rewrite_spec_layer_name(spec_layer, name) diff --git a/vllm/model_executor/models/deepseek_v2.py b/vllm/model_executor/models/deepseek_v2.py index a1a919fae644..c3befbc1a3a5 100644 --- a/vllm/model_executor/models/deepseek_v2.py +++ b/vllm/model_executor/models/deepseek_v2.py @@ -59,7 +59,6 @@ fused_moe_make_expert_params_mapping, ) from vllm.model_executor.layers.fused_moe.utils import ( - is_model_fused_shared_expert_compatible, resolve_layer_fused_shared_expert, ) from vllm.model_executor.layers.layernorm import LayerNorm, RMSNorm @@ -97,6 +96,7 @@ from vllm.model_executor.models.utils import ( AutoWeightsLoader, extract_layer_index, + get_fused_shared_expert_layers, sequence_parallel_chunk, ) from vllm.platforms import current_platform @@ -1476,10 +1476,10 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): prefix=f"{prefix}.layers", ) - self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible( - self.layers, + self.fused_shared_expert_layers = get_fused_shared_expert_layers( + self, DeepseekV2MoE, - "mlp", + enable_eplb=vllm_config.parallel_config.enable_eplb, ) if get_pp_group().is_last_rank: @@ -1629,12 +1629,9 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: ckpt_gate_proj_name="gate_proj", ckpt_down_proj_name="down_proj", ckpt_up_proj_name="up_proj", + # NOTE: global despite per-layer FSE; non-fused layers skip the extras. num_experts=self.config.n_routed_experts - + ( - self.config.n_shared_experts - if self.is_fused_shared_expert_enabled - else 0 - ), + + (self.config.n_shared_experts if self.fused_shared_expert_layers else 0), num_redundant_experts=self.num_redundant_experts, ) @@ -1659,8 +1656,11 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: ): continue # this layer has no indexer; drop its checkpoint weights + # Only parse the layer index when FSE is enabled. is_fusion_moe_shared_experts_layer = ( - self.is_fused_shared_expert_enabled and ("mlp.shared_experts" in name) + "mlp.shared_experts" in name + and bool(rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()) + and extract_layer_index(name) in self.fused_shared_expert_layers ) if _try_load_fp8_indexer_wk( diff --git a/vllm/model_executor/models/glm4_moe.py b/vllm/model_executor/models/glm4_moe.py index 4e541a49413e..c43f12e41fad 100644 --- a/vllm/model_executor/models/glm4_moe.py +++ b/vllm/model_executor/models/glm4_moe.py @@ -46,7 +46,6 @@ GateLinear, ) from vllm.model_executor.layers.fused_moe.utils import ( - is_model_fused_shared_expert_compatible, resolve_layer_fused_shared_expert, ) from vllm.model_executor.layers.layernorm import RMSNorm @@ -69,6 +68,7 @@ AutoWeightsLoader, PPMissingLayer, WeightsMapper, + get_fused_shared_expert_layers, make_empty_intermediate_tensors_factory, make_layers, maybe_fuse_shared_experts, @@ -444,10 +444,10 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): prefix=f"{prefix}.layers", ) - self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible( - self.layers, + self.fused_shared_expert_layers = get_fused_shared_expert_layers( + self, Glm4MoE, - "mlp", + enable_eplb=enable_eplb, ) if get_pp_group().is_last_rank: @@ -495,7 +495,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: skip_spec_layers(weights, self.config), n_routed_experts=self.config.n_routed_experts, n_shared_experts=self.config.n_shared_experts or 1, - enabled=self.is_fused_shared_expert_enabled, + fused_shared_expert_layers=self.fused_shared_expert_layers, ) loader = AutoWeightsLoader(self) return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) diff --git a/vllm/model_executor/models/glm4_moe_lite.py b/vllm/model_executor/models/glm4_moe_lite.py index 019380540d08..367fbbbcd96a 100644 --- a/vllm/model_executor/models/glm4_moe_lite.py +++ b/vllm/model_executor/models/glm4_moe_lite.py @@ -34,6 +34,7 @@ if TYPE_CHECKING: from transformers.models.glm4_moe_lite import Glm4MoeLiteConfig +from vllm._aiter_ops import rocm_aiter_ops from vllm.compilation.decorators import support_torch_compile from vllm.config import VllmConfig from vllm.distributed import ( @@ -43,9 +44,6 @@ from vllm.model_executor.layers.fused_moe import ( fused_moe_make_expert_params_mapping, ) -from vllm.model_executor.layers.fused_moe.utils import ( - is_model_fused_shared_expert_compatible, -) from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.logits_processor import LogitsProcessor from vllm.model_executor.layers.vocab_parallel_embedding import ( @@ -76,6 +74,8 @@ from .utils import ( AutoWeightsLoader, PPMissingLayer, + extract_layer_index, + get_fused_shared_expert_layers, get_spec_layer_idx_from_weight_name, is_pp_missing_parameter, make_layers, @@ -276,10 +276,10 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): prefix=f"{prefix}.layers", ) - self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible( - self.layers, + self.fused_shared_expert_layers = get_fused_shared_expert_layers( + self, Glm4MoeLite, - "mlp", + enable_eplb=vllm_config.parallel_config.enable_eplb, ) if get_pp_group().is_last_rank: @@ -364,12 +364,9 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: ckpt_gate_proj_name="gate_proj", ckpt_down_proj_name="down_proj", ckpt_up_proj_name="up_proj", + # NOTE: global despite per-layer FSE; non-fused layers skip the extras. num_experts=self.config.n_routed_experts - + ( - self.config.n_shared_experts - if self.is_fused_shared_expert_enabled - else 0 - ), + + (self.config.n_shared_experts if self.fused_shared_expert_layers else 0), ) params_dict = dict(self.named_parameters()) @@ -382,8 +379,11 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: if spec_layer is not None: continue # skip spec decode layers for main model + # Only parse the layer index when FSE is enabled. is_fusion_moe_shared_experts_layer = ( - self.is_fused_shared_expert_enabled and ("mlp.shared_experts" in name) + "mlp.shared_experts" in name + and bool(rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()) + and extract_layer_index(name) in self.fused_shared_expert_layers ) for param_name, weight_name, shard_id in stacked_params_mapping: diff --git a/vllm/model_executor/models/glm4_moe_lite_mtp.py b/vllm/model_executor/models/glm4_moe_lite_mtp.py index f5a82b72314b..a59eeb8f60fb 100644 --- a/vllm/model_executor/models/glm4_moe_lite_mtp.py +++ b/vllm/model_executor/models/glm4_moe_lite_mtp.py @@ -30,14 +30,12 @@ import torch.nn as nn from transformers import PreTrainedConfig +from vllm._aiter_ops import rocm_aiter_ops from vllm.config import VllmConfig from vllm.model_executor.layers.fused_moe import ( MoERunner, fused_moe_make_expert_params_mapping, ) -from vllm.model_executor.layers.fused_moe.utils import ( - is_model_fused_shared_expert_compatible, -) from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.logits_processor import LogitsProcessor from vllm.model_executor.layers.quantization import QuantizationConfig @@ -62,7 +60,12 @@ Glm4MoeLiteDecoderLayer, ) from .interfaces import SupportsPP -from .utils import get_spec_layer_idx_from_weight_name, maybe_prefix +from .utils import ( + extract_layer_index, + get_fused_shared_expert_layers, + get_spec_layer_idx_from_weight_name, + maybe_prefix, +) class SharedHead(nn.Module): @@ -241,10 +244,10 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self.moe_mlp_layers.append(layer.mlp) self.moe_layers.append(layer.mlp.experts) self.extract_moe_parameters(example_moe) - self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible( - self.model.layers.values(), + self.fused_shared_expert_layers = get_fused_shared_expert_layers( + self, Glm4MoeLite, - "mtp_block.mlp", + enable_eplb=vllm_config.parallel_config.enable_eplb, ) def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: @@ -279,17 +282,15 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: ("fused_qkv_a_proj", "kv_a_proj_with_mqa", 1), ] + fused_shared_expert_layers = self.fused_shared_expert_layers expert_params_mapping = fused_moe_make_expert_params_mapping( self, ckpt_gate_proj_name="gate_proj", ckpt_down_proj_name="down_proj", ckpt_up_proj_name="up_proj", + # NOTE: global despite per-layer FSE; non-fused layers skip the extras. num_experts=self.config.n_routed_experts - + ( - self.config.n_shared_experts - if self.is_fused_shared_expert_enabled - else 0 - ), + + (self.config.n_shared_experts if fused_shared_expert_layers else 0), ) params_dict = dict(self.named_parameters()) @@ -300,8 +301,11 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: spec_layer = get_spec_layer_idx_from_weight_name(self.config, name) if spec_layer is None: continue + # Only parse the layer index when FSE is enabled. is_fusion_moe_shared_experts_layer = ( - self.is_fused_shared_expert_enabled and ("mlp.shared_experts" in name) + "mlp.shared_experts" in name + and bool(rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()) + and extract_layer_index(name) in fused_shared_expert_layers ) name = self._rewrite_spec_layer_name(spec_layer, name) for param_name, weight_name, shard_id in stacked_params_mapping: diff --git a/vllm/model_executor/models/glm4_moe_mtp.py b/vllm/model_executor/models/glm4_moe_mtp.py index 330f891d1cf3..0464322cae02 100644 --- a/vllm/model_executor/models/glm4_moe_mtp.py +++ b/vllm/model_executor/models/glm4_moe_mtp.py @@ -31,14 +31,12 @@ import torch.nn as nn from transformers import PreTrainedConfig +from vllm._aiter_ops import rocm_aiter_ops from vllm.config import CacheConfig, ParallelConfig, VllmConfig from vllm.model_executor.layers.fused_moe import ( MoERunner, fused_moe_make_expert_params_mapping, ) -from vllm.model_executor.layers.fused_moe.utils import ( - is_model_fused_shared_expert_compatible, -) from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.logits_processor import LogitsProcessor from vllm.model_executor.layers.quantization import QuantizationConfig @@ -54,7 +52,12 @@ Glm4MoE, Glm4MoeDecoderLayer, ) -from .utils import get_spec_layer_idx_from_weight_name, maybe_prefix +from .utils import ( + extract_layer_index, + get_fused_shared_expert_layers, + get_spec_layer_idx_from_weight_name, + maybe_prefix, +) class SharedHead(nn.Module): @@ -215,10 +218,10 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self.moe_mlp_layers.append(layer.mlp) self.moe_layers.append(layer.mlp.experts) self.extract_moe_parameters(example_moe) - self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible( - self.model.layers.values(), + self.fused_shared_expert_layers = get_fused_shared_expert_layers( + self, Glm4MoE, - "mtp_block.mlp", + enable_eplb=vllm_config.parallel_config.enable_eplb, ) def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: @@ -257,8 +260,10 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: # Params for weights, fp8 weight scales, fp8 activation scales # (param_name, weight_name, expert_id, shard_id) + fused_shared_expert_layers = self.fused_shared_expert_layers + # NOTE: global despite per-layer FSE; non-fused layers skip the extras. num_experts = self.config.n_routed_experts - if self.is_fused_shared_expert_enabled and self.config.n_shared_experts: + if fused_shared_expert_layers and self.config.n_shared_experts: num_experts += self.config.n_shared_experts expert_params_mapping = fused_moe_make_expert_params_mapping( self, @@ -282,8 +287,11 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: continue name = self._rewrite_spec_layer_name(spec_layer, name) + # Only parse the layer index when FSE is enabled. is_fusion_moe_shared_experts_layer = ( - self.is_fused_shared_expert_enabled and ("mlp.shared_experts" in name) + "mlp.shared_experts" in name + and bool(rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()) + and extract_layer_index(name) in fused_shared_expert_layers ) for param_name, weight_name, shard_id in stacked_params_mapping: diff --git a/vllm/model_executor/models/qwen3_5.py b/vllm/model_executor/models/qwen3_5.py index f105948f9283..6e370f180da5 100644 --- a/vllm/model_executor/models/qwen3_5.py +++ b/vllm/model_executor/models/qwen3_5.py @@ -42,9 +42,6 @@ get_pp_group, ) from vllm.logger import init_logger -from vllm.model_executor.layers.fused_moe.utils import ( - is_model_fused_shared_expert_compatible, -) from vllm.model_executor.layers.fusion.mm_input_norm import build_mm_input_norm from vllm.model_executor.layers.layernorm import GemmaRMSNorm as Qwen3_5RMSNorm from vllm.model_executor.layers.logits_processor import LogitsProcessor @@ -99,6 +96,7 @@ WeightsMapper, _merge_multimodal_embeddings, extract_layer_index, + get_fused_shared_expert_layers, make_empty_intermediate_tensors_factory, make_layers, maybe_fuse_shared_experts, @@ -263,10 +261,10 @@ def get_layer(prefix: str): self.start_layer, self.end_layer, self.layers = make_layers( config.num_hidden_layers, get_layer, prefix=f"{prefix}.layers" ) - self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible( - self.layers, + self.fused_shared_expert_layers = get_fused_shared_expert_layers( + self, Qwen3NextSparseMoeBlock, - "mlp", + enable_eplb=vllm_config.parallel_config.enable_eplb, ) self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory( ["hidden_states", "residual"], config.hidden_size @@ -286,7 +284,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: if "moe" in self.config.model_type: weights = maybe_fuse_shared_experts( weights, - enabled=self.is_fused_shared_expert_enabled, + fused_shared_expert_layers=self.fused_shared_expert_layers, n_routed_experts=self.config.num_experts, n_shared_experts=1, ckpt_prefix="mlp.shared_expert", diff --git a/vllm/model_executor/models/qwen3_5_mtp.py b/vllm/model_executor/models/qwen3_5_mtp.py index 2acac6b5b7c5..aa7e6c06e338 100644 --- a/vllm/model_executor/models/qwen3_5_mtp.py +++ b/vllm/model_executor/models/qwen3_5_mtp.py @@ -12,9 +12,6 @@ from vllm.config import VllmConfig from vllm.distributed import get_pp_group, tensor_model_parallel_all_gather from vllm.logger import init_logger -from vllm.model_executor.layers.fused_moe.utils import ( - is_model_fused_shared_expert_compatible, -) from vllm.model_executor.layers.linear import ColumnParallelLinear from vllm.model_executor.layers.logits_processor import LogitsProcessor from vllm.model_executor.layers.vocab_parallel_embedding import ( @@ -44,6 +41,7 @@ AutoWeightsLoader, PPMissingLayer, _merge_multimodal_embeddings, + get_fused_shared_expert_layers, make_empty_intermediate_tensors_factory, maybe_fuse_shared_experts, maybe_prefix, @@ -124,10 +122,10 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): for idx in range(self.num_mtp_layers) ) vllm_config.quant_config = original_quant - self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible( - self.layers, + self.fused_shared_expert_layers = get_fused_shared_expert_layers( + self, Qwen3NextSparseMoeBlock, - "mlp", + enable_eplb=vllm_config.parallel_config.enable_eplb, ) self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory( ["hidden_states", "residual"], config.hidden_size @@ -193,7 +191,7 @@ def forward( def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: weights = maybe_fuse_shared_experts( weights, - enabled=self.is_fused_shared_expert_enabled, + fused_shared_expert_layers=self.fused_shared_expert_layers, n_routed_experts=getattr(self.config, "num_experts", 0), n_shared_experts=1, ckpt_prefix="mlp.shared_expert", diff --git a/vllm/model_executor/models/qwen3_next.py b/vllm/model_executor/models/qwen3_next.py index ee5c05d0a4ae..f6df90594a69 100644 --- a/vllm/model_executor/models/qwen3_next.py +++ b/vllm/model_executor/models/qwen3_next.py @@ -24,7 +24,6 @@ GateLinear, ) from vllm.model_executor.layers.fused_moe.utils import ( - is_model_fused_shared_expert_compatible, resolve_layer_fused_shared_expert, ) from vllm.model_executor.layers.fused_qk_norm_rope import fused_qk_rmsnorm_rope_gate @@ -75,6 +74,7 @@ PPMissingLayer, WeightsMapper, extract_layer_index, + get_fused_shared_expert_layers, make_empty_intermediate_tensors_factory, make_layers, maybe_fuse_shared_experts, @@ -672,10 +672,10 @@ def get_layer(prefix: str): self.start_layer, self.end_layer, self.layers = make_layers( config.num_hidden_layers, get_layer, prefix=f"{prefix}.layers" ) - self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible( - self.layers, + self.fused_shared_expert_layers = get_fused_shared_expert_layers( + self, Qwen3NextSparseMoeBlock, - "mlp", + enable_eplb=parallel_config.enable_eplb, ) self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory( ["hidden_states", "residual"], config.hidden_size @@ -756,7 +756,7 @@ def forward( def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: weights = maybe_fuse_shared_experts( weights, - enabled=self.is_fused_shared_expert_enabled, + fused_shared_expert_layers=self.fused_shared_expert_layers, n_routed_experts=getattr(self.config, "num_experts", 0), n_shared_experts=1, ckpt_prefix="mlp.shared_expert", diff --git a/vllm/model_executor/models/qwen3_next_mtp.py b/vllm/model_executor/models/qwen3_next_mtp.py index 4ae8f21840a1..47d038b8e287 100644 --- a/vllm/model_executor/models/qwen3_next_mtp.py +++ b/vllm/model_executor/models/qwen3_next_mtp.py @@ -12,9 +12,6 @@ from vllm.config import VllmConfig from vllm.distributed import get_pp_group, tensor_model_parallel_all_gather from vllm.logger import init_logger -from vllm.model_executor.layers.fused_moe.utils import ( - is_model_fused_shared_expert_compatible, -) from vllm.model_executor.layers.linear import ColumnParallelLinear from vllm.model_executor.layers.logits_processor import LogitsProcessor from vllm.model_executor.layers.vocab_parallel_embedding import ( @@ -33,6 +30,7 @@ from .utils import ( AutoWeightsLoader, + get_fused_shared_expert_layers, make_empty_intermediate_tensors_factory, maybe_fuse_shared_experts, maybe_prefix, @@ -95,10 +93,10 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): for idx in range(self.num_mtp_layers) ) - self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible( - self.layers, + self.fused_shared_expert_layers = get_fused_shared_expert_layers( + self, Qwen3NextSparseMoeBlock, - "mlp", + enable_eplb=vllm_config.parallel_config.enable_eplb, ) self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory( ["hidden_states", "residual"], config.hidden_size @@ -164,7 +162,7 @@ def forward( def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: weights = maybe_fuse_shared_experts( weights, - enabled=self.is_fused_shared_expert_enabled, + fused_shared_expert_layers=self.fused_shared_expert_layers, n_routed_experts=self.config.num_experts, n_shared_experts=1, ckpt_prefix="mlp.shared_expert", diff --git a/vllm/model_executor/models/utils.py b/vllm/model_executor/models/utils.py index a86496b48a9b..3386a76a0ea6 100644 --- a/vllm/model_executor/models/utils.py +++ b/vllm/model_executor/models/utils.py @@ -12,7 +12,7 @@ import torch.nn as nn from torch.nn.modules.module import register_module_module_registration_hook -from vllm.config import VllmConfig +from vllm.config import VllmConfig, get_current_vllm_config_or_none from vllm.distributed import ( get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size, @@ -502,13 +502,88 @@ def _check_skipped_aliases(self, autoloaded_weights: set[str]) -> None: ) +def get_fused_shared_expert_layers( + module: nn.Module, + moe_cls: type[nn.Module], + enable_eplb: bool = False, +) -> frozenset[int]: + """Resolve fused-shared-expert layer indices from the live module tree. + + The index is taken from each MoE module's qualified name, so it follows the + same index space as the caller's weight names: 0-based for modules stored in + an ``nn.ModuleList``, and absolute for ``nn.ModuleDict`` MTP layers keyed by + layer index. No loader needs to hard-code an offset. + + Also rejects mixed fusion when EPLB is enabled or the ``nixl_ep`` all2all + backend is selected, and logs the resolved policy. + """ + fused_layers: set[int] = set() + moe_layers: set[int] = set() + for name, mod in module.named_modules(): + if not isinstance(mod, moe_cls): + continue + idx = extract_layer_index(name) + moe_layers.add(idx) + if getattr(mod, "is_fused_shared_expert_enabled", False): + fused_layers.add(idx) + + # EPLB bookkeeping is per rank/model: `EplbState` builds each model's + # `expert_weights`/`num_moe_layers` from the local `MixtureOfExperts` (one + # pipeline stage here), while main and draft models have separate + # `EplbModelState`s keyed by model config. The uniformity requirement is + # therefore stage/model-local by design: mixed within a stage is rejected, + # mixed across PP stages (or main vs draft) is allowed. + if 0 < len(fused_layers) < len(moe_layers): + if enable_eplb: + raise NotImplementedError( + "Mixed per-layer shared-expert fusion within a pipeline " + "stage/model is not supported with EPLB" + ) + vllm_config = get_current_vllm_config_or_none() + parallel_config = ( + vllm_config.parallel_config if vllm_config is not None else None + ) + if parallel_config is not None and parallel_config.enable_expert_parallel: + raise NotImplementedError( + "Mixed per-layer shared-expert fusion with expert parallelism is " + "not tested. Please open an issue if you need this combination." + ) + # These backends size one shared buffer from the first MoE layer's + # geometry (all2all.py NixlEPAll2AllManager._init_buffer/get_handle). + all2all_backend = ( + parallel_config.all2all_backend if parallel_config is not None else None + ) + if all2all_backend == "nixl_ep": + raise NotImplementedError( + "Mixed per-layer shared-expert fusion is not supported with " + f"all2all_backend={all2all_backend!r}" + ) + + if fused_layers: + logger.info( + "AITER fused shared experts: enabled on MoE layer(s) %s; " + "disabled on MoE layer(s) %s.", + sorted(fused_layers), + sorted(moe_layers - fused_layers) or "none", + ) + elif moe_layers: + from vllm._aiter_ops import rocm_aiter_ops + + if rocm_aiter_ops.is_fusion_moe_shared_experts_enabled(): + logger.info( + "AITER fused shared experts: disabled on all %d MoE layer(s).", + len(moe_layers), + ) + return frozenset(fused_layers) + + def maybe_fuse_shared_experts( weights: Iterable[tuple[str, torch.Tensor]], *, n_routed_experts: int, n_shared_experts: int, + fused_shared_expert_layers: frozenset[int], ckpt_prefix: str = "mlp.shared_experts", - enabled: bool | None = None, ) -> Iterable[tuple[str, torch.Tensor]]: """Route AITER fused-shared-expert checkpoint weights into fused slots. @@ -524,21 +599,16 @@ def maybe_fuse_shared_experts( n_routed_experts: Number of routed experts; offsets the fused slots. n_shared_experts: Number of shared experts packed into the tensor. ckpt_prefix: Checkpoint module name of the shared experts. - enabled: Whether AITER fused-shared-experts is active. Defaults to - `rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()`; pass an - explicit value only when the model gates on something more (e.g. - quant-spec compatibility) and it must match its construction-time - decision. + fused_shared_expert_layers: Layer indices whose shared expert was fused + at construction time, as returned by + `get_fused_shared_expert_layers`. Only weights belonging to these + layers are routed to fused slots. Yields: `(name, tensor)` pairs with shared experts routed to fused slots. """ - if enabled is None: - from vllm._aiter_ops import rocm_aiter_ops - - enabled = rocm_aiter_ops.is_fusion_moe_shared_experts_enabled() - if not enabled: + if not fused_shared_expert_layers: yield from weights return @@ -549,6 +619,9 @@ def maybe_fuse_shared_experts( if prefix not in name: yield name, loaded_weight continue + if extract_layer_index(name) not in fused_shared_expert_layers: + yield name, loaded_weight + continue # gate/up split on the output dim; down_proj on the input dim. split_dim = 1 if ("down_proj.weight" in name and loaded_weight.ndim > 1) else 0 total = loaded_weight.shape[split_dim] diff --git a/vllm/models/deepseek_v32/amd/model.py b/vllm/models/deepseek_v32/amd/model.py index 225e9b0ffe10..2338b0d030ef 100644 --- a/vllm/models/deepseek_v32/amd/model.py +++ b/vllm/models/deepseek_v32/amd/model.py @@ -7,14 +7,12 @@ import torch +from vllm._aiter_ops import rocm_aiter_ops from vllm.config import VllmConfig from vllm.distributed import get_pp_group, get_tensor_model_parallel_rank from vllm.model_executor.layers.fused_moe import ( fused_moe_make_expert_params_mapping, ) -from vllm.model_executor.layers.fused_moe.utils import ( - is_model_fused_shared_expert_compatible, -) from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.vocab_parallel_embedding import ( VocabParallelEmbedding, @@ -32,6 +30,8 @@ ) from vllm.model_executor.models.utils import ( PPMissingLayer, + extract_layer_index, + get_fused_shared_expert_layers, get_pp_missing_layer_names, is_pp_missing_parameter, make_empty_intermediate_tensors_factory, @@ -172,10 +172,10 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): prefix=f"{prefix}.layers", ) - self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible( - self.layers, + self.fused_shared_expert_layers = get_fused_shared_expert_layers( + self, DeepseekV2MoE, - "mlp", + enable_eplb=vllm_config.parallel_config.enable_eplb, ) if get_pp_group().is_last_rank: @@ -246,12 +246,9 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: ckpt_gate_proj_name="gate_proj", ckpt_down_proj_name="down_proj", ckpt_up_proj_name="up_proj", + # NOTE: global despite per-layer FSE; non-fused layers skip the extras. num_experts=self.config.n_routed_experts - + ( - self.config.n_shared_experts - if self.is_fused_shared_expert_enabled - else 0 - ), + + (self.config.n_shared_experts if self.fused_shared_expert_layers else 0), num_redundant_experts=self.num_redundant_experts, ) @@ -269,8 +266,11 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: continue if get_spec_layer_idx_from_weight_name(self.config, name) is not None: continue + # Only parse the layer index when FSE is enabled. is_fusion_moe_shared_experts_layer = ( - self.is_fused_shared_expert_enabled and ("mlp.shared_experts" in name) + "mlp.shared_experts" in name + and bool(rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()) + and extract_layer_index(name) in self.fused_shared_expert_layers ) if _try_load_fp8_indexer_wk( name, diff --git a/vllm/models/deepseek_v32/amd/mtp.py b/vllm/models/deepseek_v32/amd/mtp.py index d5066b39de8d..b773035e9199 100644 --- a/vllm/models/deepseek_v32/amd/mtp.py +++ b/vllm/models/deepseek_v32/amd/mtp.py @@ -8,14 +8,12 @@ import torch import torch.nn as nn +from vllm._aiter_ops import rocm_aiter_ops from vllm.config import VllmConfig from vllm.distributed import tensor_model_parallel_all_reduce from vllm.model_executor.layers.fused_moe import ( fused_moe_make_expert_params_mapping, ) -from vllm.model_executor.layers.fused_moe.utils import ( - is_model_fused_shared_expert_compatible, -) from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.logits_processor import LogitsProcessor from vllm.model_executor.layers.vocab_parallel_embedding import ( @@ -34,6 +32,8 @@ ) from vllm.model_executor.models.interfaces import SupportsPP from vllm.model_executor.models.utils import ( + extract_layer_index, + get_fused_shared_expert_layers, get_pp_missing_layer_names, make_empty_intermediate_tensors_factory, maybe_prefix, @@ -183,10 +183,10 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory( ["hidden_states", "residual"], self.config.hidden_size ) - self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible( - self.model.layers.values(), + self.fused_shared_expert_layers = get_fused_shared_expert_layers( + self, DeepseekV2MoE, - "mtp_block.mlp", + enable_eplb=vllm_config.parallel_config.enable_eplb, ) def set_moe_parameters(self): @@ -266,17 +266,15 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: ("wk_weights_proj", "wk", 0), ("wk_weights_proj", "weights_proj", 1), ] + fused_shared_expert_layers = self.fused_shared_expert_layers expert_params_mapping = fused_moe_make_expert_params_mapping( self, ckpt_gate_proj_name="gate_proj", ckpt_down_proj_name="down_proj", ckpt_up_proj_name="up_proj", + # NOTE: global despite per-layer FSE; non-fused layers skip the extras. num_experts=self.config.n_routed_experts - + ( - self.config.n_shared_experts - if self.is_fused_shared_expert_enabled - else 0 - ), + + (self.config.n_shared_experts if fused_shared_expert_layers else 0), ) pp_missing_layer_names = get_pp_missing_layer_names(self) @@ -303,8 +301,11 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: weight_loader(param, loaded_weight) loaded_params.add(name) continue + # Only parse the layer index when FSE is enabled. is_fusion_moe_shared_experts_layer = ( - self.is_fused_shared_expert_enabled and ("mlp.shared_experts" in name) + "mlp.shared_experts" in name + and bool(rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()) + and extract_layer_index(name) in fused_shared_expert_layers ) name = self._rewrite_spec_layer_name(spec_layer, name) diff --git a/vllm/models/deepseek_v4/amd/model.py b/vllm/models/deepseek_v4/amd/model.py index b0d798033f11..b60189b0b57d 100644 --- a/vllm/models/deepseek_v4/amd/model.py +++ b/vllm/models/deepseek_v4/amd/model.py @@ -13,7 +13,7 @@ import vllm.envs as envs from vllm._aiter_ops import rocm_aiter_ops -from vllm.config import ParallelConfig, VllmConfig +from vllm.config import VllmConfig from vllm.distributed import ( get_pp_group, get_tensor_model_parallel_rank, @@ -35,7 +35,6 @@ fused_topk_bias, ) from vllm.model_executor.layers.fused_moe.utils import ( - is_model_fused_shared_expert_compatible, resolve_layer_fused_shared_expert, ) from vllm.model_executor.layers.layernorm import RMSNorm @@ -54,9 +53,6 @@ ) from vllm.model_executor.layers.quantization import QuantizationConfig from vllm.model_executor.layers.quantization.mxfp4 import Mxfp4MoEMethod -from vllm.model_executor.layers.quantization.utils.config_utils import ( - is_shared_expert_quant_fse_compatible, -) from vllm.model_executor.layers.vocab_parallel_embedding import ( ParallelLMHead, VocabParallelEmbedding, @@ -72,6 +68,7 @@ PPMissingLayer, WeightsMapper, extract_layer_index, + get_fused_shared_expert_layers, is_pp_missing_parameter, make_layers, maybe_prefix, @@ -525,31 +522,6 @@ def forward_modular( return routed + shared_out -def _fuse_shared_experts_enabled(config, parallel_config: ParallelConfig) -> bool: - if ( - getattr(config, "n_shared_experts", None) - and envs.VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS - and not parallel_config.enable_expert_parallel - and parallel_config.data_parallel_size > 1 - ): - # Fused shared experts are not supported under data parallelism: - # the fused path cannot load the FP8 shared expert into the MXFP4 - # routed slot, which fails obscurely during weight loading. Fail - # fast with an actionable message instead. - raise ValueError( - "DeepSeek-V4 fused shared experts " - "(VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS=1) are not supported " - "with data parallelism (data_parallel_size > 1). Set " - "VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS=0 to run the shared " - "expert as an unfused FP8 MLP." - ) - return bool( - getattr(config, "n_shared_experts", None) - and envs.VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS - and not parallel_config.enable_expert_parallel - ) - - def _validate_aiter_mega_moe_config(vllm_config: VllmConfig) -> None: parallel_config = vllm_config.parallel_config if not rocm_aiter_ops.is_fused_moe_enabled(): @@ -658,28 +630,31 @@ def __init__( self.n_shared_experts = config.n_shared_experts - # TODO: Historically, only `VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS=1` - # is checked to enable FSE for DeepSeek-v4, despite AITER not being used. - # This should be cleaned up and use `resolve_layer_fused_shared_expert`. self.fuse_heterogeneous_shared_expert = fuse_heterogeneous_shared_expert - fse_requested = ( - _fuse_shared_experts_enabled(config, vllm_config.parallel_config) + self.is_fused_shared_expert_enabled = False + if ( + getattr(config, "n_shared_experts", None) and not self.fuse_heterogeneous_shared_expert - ) - fse_compatible = False - if fse_requested: - fse_compatible, fse_reason = is_shared_expert_quant_fse_compatible( - quant_config, - f"{prefix}.experts", - f"{prefix}.shared_experts", - ) - if not fse_compatible: - logger.warning( - "VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS is enabled but " - "cannot be enabled: %s.", - fse_reason, + ): + # Fused shared experts are not supported under data parallelism: + # the fused path cannot load the FP8 shared expert into the MXFP4 + # routed slot. Fail fast with an actionable message. + if ( + not vllm_config.parallel_config.enable_expert_parallel + and vllm_config.parallel_config.data_parallel_size > 1 + and rocm_aiter_ops.is_fusion_moe_shared_experts_enabled() + ): + raise ValueError( + "DeepSeek-V4 fused shared experts " + "(VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS=1) are not " + "supported with data parallelism " + "(data_parallel_size > 1). Set " + "VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS=0 to run the " + "shared expert as an unfused FP8 MLP." ) - self.is_fused_shared_expert_enabled = fse_requested and fse_compatible + self.is_fused_shared_expert_enabled = resolve_layer_fused_shared_expert( + quant_config, prefix, disable_under_ep=True + ) self.use_mega_moe = vllm_config.kernel_config.moe_backend == "aiter_mega_moe" fuse_mega_moe_shared_expert = ( @@ -1160,10 +1135,10 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): ), prefix=f"{prefix}.layers", ) - self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible( - self.layers, + self.fused_shared_expert_layers = get_fused_shared_expert_layers( + self, DeepseekV4MoE, - "ffn", + enable_eplb=vllm_config.parallel_config.enable_eplb, ) if get_pp_group().is_last_rank: @@ -1344,16 +1319,12 @@ def _resolve_param_name(name: str) -> str: # at init) as the single source of truth, so the redirect below cannot # diverge from how the module was built if per-layer quantization ever # mixes fused and non-fused layers. - fuse_by_layer = { - extract_layer_index(mod_name): mod.is_fused_shared_expert_enabled - for mod_name, mod in self.named_modules() - if isinstance(mod, DeepseekV4MoE) - } + fuse_by_layer = self.fused_shared_expert_layers n_routed = self.config.n_routed_experts # The redirect below maps the single shared-expert tensor group to one # appended slot; multiple shared experts would need per-expert slicing # (see deepseek_v2.py). DeepSeek-V4 has n_shared_experts == 1. - if any(fuse_by_layer.values()) and self.config.n_shared_experts != 1: + if fuse_by_layer and self.config.n_shared_experts != 1: raise NotImplementedError( "deepseek-v4 fused shared-expert loading supports only " f"n_shared_experts == 1, got {self.config.n_shared_experts}" @@ -1364,13 +1335,7 @@ def _resolve_param_name(name: str) -> str: # into appended routed-expert slot ``.ffn.experts.{n_routed}`` # so the MXFP4-quantized shared expert loads through the routed # expert loader (grouped GEMM). Single shared expert only. - if ".ffn.shared_experts.w" in name and fuse_by_layer.get( - extract_layer_index(name), False - ): - name = name.replace( - ".ffn.shared_experts.w", - f".ffn.experts.{n_routed}.w", - ) + name = _redirect_shared_expert_name(name, fuse_by_layer, n_routed) for param_name, weight_name, shard_id in stacked_params_mapping: # Skip non-stacked layers and experts (experts handled below). @@ -1448,7 +1413,7 @@ def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: # redirected shared-expert weights route through the expert loader. n_shared = getattr(self.config, "n_shared_experts", 0) or 0 num_experts = self.config.n_routed_experts + ( - n_shared if self.is_fused_shared_expert_enabled else 0 + n_shared if self.fused_shared_expert_layers else 0 ) return fused_moe_make_expert_params_mapping( self, @@ -1460,6 +1425,24 @@ def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: ) +def _redirect_shared_expert_name( + name: str, fuse_by_layer: frozenset[int], n_routed_experts: int +) -> str: + """Route a DSV4 shared-expert checkpoint name to its final parameter name. + + Fused layers go to the appended routed-expert slot. When any layer is fused + the weight mapper leaves the checkpoint ``.w2`` names untouched, so + non-fused layers still need the native ``down_proj`` parameter name. + """ + if ".ffn.shared_experts.w" not in name: + return name + if extract_layer_index(name) in fuse_by_layer: + return name.replace( + ".ffn.shared_experts.w", f".ffn.experts.{n_routed_experts}.w" + ) + return name.replace(".ffn.shared_experts.w2", ".ffn.shared_experts.down_proj") + + def _make_deepseek_v4_weights_mapper( expert_dtype: str, fuse_shared_experts: bool = False ) -> WeightsMapper: @@ -1533,10 +1516,10 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self.model = self.model_cls( vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") ) - if expert_dtype != "fp4" or self.model.is_fused_shared_expert_enabled: + if expert_dtype != "fp4" or self.model.fused_shared_expert_layers: self.hf_to_vllm_mapper = _make_deepseek_v4_weights_mapper( expert_dtype, - fuse_shared_experts=self.model.is_fused_shared_expert_enabled, + fuse_shared_experts=bool(self.model.fused_shared_expert_layers), ) if get_pp_group().is_last_rank: self.lm_head = ParallelLMHead( diff --git a/vllm/models/deepseek_v4/amd/mtp.py b/vllm/models/deepseek_v4/amd/mtp.py index d49769c843d9..1144bbde9604 100644 --- a/vllm/models/deepseek_v4/amd/mtp.py +++ b/vllm/models/deepseek_v4/amd/mtp.py @@ -302,6 +302,18 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self.model = DeepSeekV4MultiTokenPredictor( vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") ) + if any( + getattr( + layer.mtp_block.ffn, + "is_fused_shared_expert_enabled", + False, + ) + for layer in self.model.layers.values() + ): + raise NotImplementedError( + "DeepSeek-V4 MTP fused shared experts are not supported yet; " + "please open an issue." + ) def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: return self.model.embed_input_ids(input_ids) diff --git a/vllm/models/deepseek_v4/cpu/model.py b/vllm/models/deepseek_v4/cpu/model.py index 63ec337fa48b..2e6f76c09e07 100644 --- a/vllm/models/deepseek_v4/cpu/model.py +++ b/vllm/models/deepseek_v4/cpu/model.py @@ -21,9 +21,6 @@ GateLinear, fused_moe_make_expert_params_mapping, ) -from vllm.model_executor.layers.fused_moe.utils import ( - is_model_fused_shared_expert_compatible, -) from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.linear import ( MergedColumnParallelLinear, @@ -493,11 +490,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): ), prefix=f"{prefix}.layers", ) - self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible( - self.layers, - DeepseekV4MoE, - "ffn", - ) + self.is_fused_shared_expert_enabled = False if get_pp_group().is_last_rank: self.norm = RMSNorm(config.hidden_size, self.rms_norm_eps) diff --git a/vllm/models/deepseek_v41/amd/model.py b/vllm/models/deepseek_v41/amd/model.py index b9094c697f50..68c2f47ec3ef 100644 --- a/vllm/models/deepseek_v41/amd/model.py +++ b/vllm/models/deepseek_v41/amd/model.py @@ -1344,6 +1344,18 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self.model = self.model_cls( vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") ) + if any( + getattr( + getattr(layer, "ffn", None), + "is_fused_shared_expert_enabled", + False, + ) + for layer in self.model.layers + ): + raise NotImplementedError( + "DeepSeek-V4.1 AMD fused shared experts are not supported yet; " + "please open an issue." + ) if get_pp_group().is_last_rank: self.lm_head = ParallelLMHead( config.vocab_size, diff --git a/vllm/models/glm5next/common/model.py b/vllm/models/glm5next/common/model.py index 6b0cbe4ad61b..e944ad38f151 100644 --- a/vllm/models/glm5next/common/model.py +++ b/vllm/models/glm5next/common/model.py @@ -8,6 +8,7 @@ from torch import nn from transformers import Glm5NextTextConfig +from vllm._aiter_ops import rocm_aiter_ops from vllm.config import ParallelConfig, VllmConfig from vllm.distributed import ( get_ep_group, @@ -24,7 +25,6 @@ fused_moe_make_expert_params_mapping, ) from vllm.model_executor.layers.fused_moe.utils import ( - is_model_fused_shared_expert_compatible, resolve_layer_fused_shared_expert, ) from vllm.model_executor.layers.fusion.mm_input_norm import build_mm_input_norm @@ -76,6 +76,8 @@ from vllm.model_executor.models.utils import ( AutoWeightsLoader, PPMissingLayer, + extract_layer_index, + get_fused_shared_expert_layers, init_vllm_registered_model, is_pp_missing_parameter, make_layers, @@ -749,8 +751,10 @@ def get_layer(prefix: str): # The active slice is fixed after construction; cache it so forward # doesn't rebuild the slice (a fresh list) every step. self._active_layers = self.layers[self.start_layer : self.end_layer] - self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible( - self.layers, Glm5NextMoE, "mlp" + self.fused_shared_expert_layers = get_fused_shared_expert_layers( + self, + Glm5NextMoE, + enable_eplb=vllm_config.parallel_config.enable_eplb, ) if get_pp_group().is_last_rank: @@ -894,13 +898,15 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: 0, ) num_fused_shared = _num_fused_shared_experts( - self.config.n_shared_experts, self.is_fused_shared_expert_enabled + self.config.n_shared_experts, + bool(self.fused_shared_expert_layers), ) expert_params_mapping = fused_moe_make_expert_params_mapping( self, ckpt_gate_proj_name="gate_proj", ckpt_down_proj_name="down_proj", ckpt_up_proj_name="up_proj", + # NOTE: global despite per-layer FSE; non-fused layers skip the extras. num_experts=self.config.n_routed_experts + num_fused_shared, num_redundant_experts=num_redundant_experts, ) @@ -930,7 +936,12 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: # Models trained using ColossalAI may include these tensors in # the checkpoint. Skip them. continue - if self.is_fused_shared_expert_enabled: + # Only parse the layer index when FSE is enabled. + if ( + "mlp.shared_experts." in name + and bool(rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()) + and extract_layer_index(name) in self.fused_shared_expert_layers + ): name = _fused_shared_expert_name(name, self.config.n_routed_experts) # Handle FP8 indexer WK: dequantize to BF16 for fusion with diff --git a/vllm/models/glm5next/common/mtp.py b/vllm/models/glm5next/common/mtp.py index 6e6950ca700f..deb3c61a1dfb 100644 --- a/vllm/models/glm5next/common/mtp.py +++ b/vllm/models/glm5next/common/mtp.py @@ -6,13 +6,11 @@ import torch import torch.nn as nn +from vllm._aiter_ops import rocm_aiter_ops from vllm.config import VllmConfig from vllm.model_executor.layers.fused_moe import ( fused_moe_make_expert_params_mapping, ) -from vllm.model_executor.layers.fused_moe.utils import ( - is_model_fused_shared_expert_compatible, -) from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.logits_processor import LogitsProcessor from vllm.model_executor.layers.vocab_parallel_embedding import ( @@ -24,7 +22,11 @@ ) from vllm.model_executor.models.deepseek_mtp import SharedHead from vllm.model_executor.models.deepseek_v2 import DeepseekV2MixtureOfExperts -from vllm.model_executor.models.utils import maybe_prefix +from vllm.model_executor.models.utils import ( + extract_layer_index, + get_fused_shared_expert_layers, + maybe_prefix, +) from vllm.models.glm5next.nvidia.ops.fused_eh_norm import fused_eh_norm from vllm.platforms import current_platform from vllm.sequence import IntermediateTensors @@ -212,8 +214,10 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") ) self.set_moe_parameters() - self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible( - self.model.layers.values(), Glm5NextMoE, "mtp_block.mlp" + self.fused_shared_expert_layers = get_fused_shared_expert_layers( + self, + Glm5NextMoE, + enable_eplb=vllm_config.parallel_config.enable_eplb, ) def set_moe_parameters(self): @@ -296,14 +300,17 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: ("wk_weights_proj", "wk", 0), ("wk_weights_proj", "weights_proj", 1), ] + fused_shared_expert_layers = self.fused_shared_expert_layers num_fused_shared = _num_fused_shared_experts( - self.config.n_shared_experts, self.is_fused_shared_expert_enabled + self.config.n_shared_experts, + bool(fused_shared_expert_layers), ) expert_params_mapping = fused_moe_make_expert_params_mapping( self, ckpt_gate_proj_name="gate_proj", ckpt_down_proj_name="down_proj", ckpt_up_proj_name="up_proj", + # NOTE: global despite per-layer FSE; non-fused layers skip the extras. num_experts=self.config.n_routed_experts + num_fused_shared, num_redundant_experts=self.num_redundant_experts, ) @@ -329,7 +336,12 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: if spec_layer is None: continue name = self._rewrite_spec_layer_name(spec_layer, name) - if self.is_fused_shared_expert_enabled: + # Only parse the layer index when FSE is enabled. + if ( + "mlp.shared_experts." in name + and bool(rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()) + and extract_layer_index(name) in fused_shared_expert_layers + ): name = _fused_shared_expert_name(name, self.config.n_routed_experts) if _try_load_fp8_indexer_wk( diff --git a/vllm/models/minimax_m3/amd/model.py b/vllm/models/minimax_m3/amd/model.py index 7b7459770d87..424b24e51a87 100644 --- a/vllm/models/minimax_m3/amd/model.py +++ b/vllm/models/minimax_m3/amd/model.py @@ -51,7 +51,7 @@ fused_moe_make_expert_params_mapping, ) from vllm.model_executor.layers.fused_moe.utils import ( - is_model_fused_shared_expert_compatible, + resolve_layer_fused_shared_expert, ) from vllm.model_executor.layers.fusion.quant_activation import ( QuantizedActivation, @@ -65,9 +65,6 @@ ) from vllm.model_executor.layers.logits_processor import LogitsProcessor from vllm.model_executor.layers.quantization import QuantizationConfig -from vllm.model_executor.layers.quantization.utils.config_utils import ( - is_shared_expert_quant_fse_compatible, -) from vllm.model_executor.layers.quantization.utils.quant_utils import kMxfp8Dynamic from vllm.model_executor.layers.rotary_embedding import get_rope from vllm.model_executor.layers.vocab_parallel_embedding import ( @@ -89,6 +86,8 @@ AutoWeightsLoader, PPMissingLayer, WeightsMapper, + extract_layer_index, + get_fused_shared_expert_layers, init_vllm_registered_model, is_pp_missing_parameter, make_empty_intermediate_tensors_factory, @@ -405,32 +404,11 @@ def __init__( # MoE call as the last expert slot, so we don't build a separate module. # On gfx950 with aiter MoE the append is fused inside aiter's grouped # top-k kernel; otherwise it goes through the vLLM top-k bias router. - # It is disabled under expert parallelism (the shared slot is appended to - # the routed top-k, which the EP expert-mapping path does not handle). - # TODO: Historically, only `VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS=1` - # is checked to enable FSE for MiniMax-M3, despite AITER not being used. - # This should be cleaned up and use `resolve_layer_fused_shared_expert`. - fse_requested = envs.VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS - self.is_fused_shared_expert_enabled = False - if ( - fse_requested - and bool(getattr(config, "n_shared_experts", None)) - and not get_current_vllm_config().parallel_config.enable_expert_parallel - ): - fse_compatible, fse_reason = is_shared_expert_quant_fse_compatible( - quant_config, - f"{prefix}.experts", - f"{prefix}.shared_experts", + if bool(getattr(config, "n_shared_experts", None)): + self.is_fused_shared_expert_enabled = resolve_layer_fused_shared_expert( + quant_config, prefix, disable_under_ep=True ) - if not fse_compatible: - logger.warning( - "VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS is enabled but " - "cannot be enabled: %s.", - fse_reason, - ) - else: - self.is_fused_shared_expert_enabled = True # When additionally on gfx950 with an active aiter MoE backend, the shared # expert is appended inside aiter's an active aiter MoE backend, the shared @@ -1538,10 +1516,10 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): ), prefix=f"{prefix}.layers", ) - self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible( - self.layers, + self.fused_shared_expert_layers = get_fused_shared_expert_layers( + self, MiniMaxM3MoE, - "block_sparse_moe", + enable_eplb=vllm_config.parallel_config.enable_eplb, ) if get_pp_group().is_last_rank: @@ -1622,8 +1600,9 @@ def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: # Checkpoint experts use w1=gate, w2=down, w3=up. When fusing the shared # expert, include the appended slot (id == num_local_experts). n_shared = getattr(self.config, "n_shared_experts", 0) or 0 + # NOTE: global despite per-layer FSE; non-fused layers skip the extras. num_experts = self.config.num_local_experts + ( - n_shared if self.is_fused_shared_expert_enabled else 0 + n_shared if self.fused_shared_expert_layers else 0 ) return fused_moe_make_expert_params_mapping( self, @@ -1672,7 +1651,12 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: # down->w2) so it loads via the routed expert loader. Runs before the # stacked/dense mappings so shared_experts.gate_proj/up_proj are not # captured by the dense gate_up_proj mapping. - if self.is_fused_shared_expert_enabled and ".shared_experts." in name: + # Only parse the layer index when FSE is enabled. + if ( + ".shared_experts." in name + and bool(rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()) + and extract_layer_index(name) in self.fused_shared_expert_layers + ): sid = self.config.num_local_experts name = name.replace(".shared_experts.gate_proj.", f".experts.{sid}.w1.") name = name.replace(".shared_experts.up_proj.", f".experts.{sid}.w3.") diff --git a/vllm/models/qwen4_exp/amd/model.py b/vllm/models/qwen4_exp/amd/model.py index 5715d038dd1a..675b9a9a6cff 100644 --- a/vllm/models/qwen4_exp/amd/model.py +++ b/vllm/models/qwen4_exp/amd/model.py @@ -12,9 +12,6 @@ from vllm.compilation.decorators import support_torch_compile from vllm.config import VllmConfig from vllm.distributed import get_pp_group -from vllm.model_executor.layers.fused_moe.utils import ( - is_model_fused_shared_expert_compatible, -) from vllm.model_executor.layers.fusion.mm_input_norm import build_mm_input_norm from vllm.model_executor.layers.logits_processor import LogitsProcessor from vllm.model_executor.layers.mamba.gdn.qwen_gdn_linear_attn import ( @@ -64,6 +61,7 @@ WeightsMapper, _merge_multimodal_embeddings, extract_layer_index, + get_fused_shared_expert_layers, make_empty_intermediate_tensors_factory, make_layers, maybe_fuse_shared_experts, @@ -446,10 +444,10 @@ def get_layer(prefix: str) -> Qwen4ExpDecoderLayer: self.start_layer, self.end_layer, self.layers = make_layers( config.num_hidden_layers, get_layer, prefix=f"{prefix}.layers" ) - self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible( - self.layers, + self.fused_shared_expert_layers = get_fused_shared_expert_layers( + self, Qwen4ExpSparseMoeBlock, - "mlp", + enable_eplb=vllm_config.parallel_config.enable_eplb, ) intermediate_size = config.hidden_size * config.hc_count self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory( @@ -620,7 +618,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: ) weights = maybe_fuse_shared_experts( weights, - enabled=self.is_fused_shared_expert_enabled, + fused_shared_expert_layers=self.fused_shared_expert_layers, n_routed_experts=getattr(self.config, "num_experts", 0) or 0, n_shared_experts=1, ckpt_prefix="mlp.shared_expert", diff --git a/vllm/models/qwen4_exp/amd/mtp.py b/vllm/models/qwen4_exp/amd/mtp.py index 0805d102b0c9..7a73ac8f33ed 100644 --- a/vllm/models/qwen4_exp/amd/mtp.py +++ b/vllm/models/qwen4_exp/amd/mtp.py @@ -23,9 +23,6 @@ from vllm.compilation.decorators import support_torch_compile from vllm.config import VllmConfig, replace, set_current_vllm_config from vllm.distributed import get_pp_group -from vllm.model_executor.layers.fused_moe.utils import ( - is_model_fused_shared_expert_compatible, -) from vllm.model_executor.layers.layernorm import GemmaRMSNorm from vllm.model_executor.layers.linear import ColumnParallelLinear from vllm.model_executor.layers.logits_processor import LogitsProcessor @@ -41,6 +38,7 @@ PPMissingLayer, WeightsMapper, get_draft_quant_config, + get_fused_shared_expert_layers, make_empty_intermediate_tensors_factory, maybe_fuse_shared_experts, maybe_prefix, @@ -229,10 +227,10 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: ) for idx in range(self.num_mtp_layers) ) - self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible( - self.layers, + self.fused_shared_expert_layers = get_fused_shared_expert_layers( + self, Qwen4ExpSparseMoeBlock, - "mlp", + enable_eplb=vllm_config.parallel_config.enable_eplb, ) self.pre_fc_norm_embedding = GemmaRMSNorm( @@ -373,7 +371,7 @@ def forward( def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: weights = maybe_fuse_shared_experts( weights, - enabled=self.is_fused_shared_expert_enabled, + fused_shared_expert_layers=self.fused_shared_expert_layers, n_routed_experts=getattr(self.config, "num_experts", 0) or 0, n_shared_experts=1, ckpt_prefix="mlp.shared_expert", diff --git a/vllm/models/qwen4_exp/nvidia/model.py b/vllm/models/qwen4_exp/nvidia/model.py index ed9aaa927733..08ad38361e14 100644 --- a/vllm/models/qwen4_exp/nvidia/model.py +++ b/vllm/models/qwen4_exp/nvidia/model.py @@ -13,9 +13,6 @@ from vllm.config import VllmConfig from vllm.distributed import get_pp_group from vllm.forward_context import get_forward_context, is_forward_context_available -from vllm.model_executor.layers.fused_moe.utils import ( - is_model_fused_shared_expert_compatible, -) from vllm.model_executor.layers.fusion.mm_input_norm import build_mm_input_norm from vllm.model_executor.layers.logits_processor import LogitsProcessor from vllm.model_executor.layers.mamba.gdn.qwen_gdn_linear_attn import ( @@ -65,6 +62,7 @@ WeightsMapper, _merge_multimodal_embeddings, extract_layer_index, + get_fused_shared_expert_layers, make_empty_intermediate_tensors_factory, make_layers, maybe_fuse_shared_experts, @@ -446,10 +444,10 @@ def get_layer(prefix: str) -> Qwen4ExpDecoderLayer: self.start_layer, self.end_layer, self.layers = make_layers( config.num_hidden_layers, get_layer, prefix=f"{prefix}.layers" ) - self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible( - self.layers, + self.fused_shared_expert_layers = get_fused_shared_expert_layers( + self, Qwen4ExpSparseMoeBlock, - "mlp", + enable_eplb=vllm_config.parallel_config.enable_eplb, ) intermediate_size = config.hidden_size * config.hc_count self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory( @@ -656,7 +654,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: ) weights = maybe_fuse_shared_experts( weights, - enabled=self.is_fused_shared_expert_enabled, + fused_shared_expert_layers=self.fused_shared_expert_layers, n_routed_experts=getattr(self.config, "num_experts", 0) or 0, n_shared_experts=1, ckpt_prefix="mlp.shared_expert", diff --git a/vllm/models/qwen4_exp/nvidia/mtp.py b/vllm/models/qwen4_exp/nvidia/mtp.py index fb15d076467a..30401ca56dcb 100644 --- a/vllm/models/qwen4_exp/nvidia/mtp.py +++ b/vllm/models/qwen4_exp/nvidia/mtp.py @@ -24,9 +24,6 @@ from vllm.config import VllmConfig, replace, set_current_vllm_config from vllm.distributed import get_pp_group from vllm.forward_context import get_forward_context, is_forward_context_available -from vllm.model_executor.layers.fused_moe.utils import ( - is_model_fused_shared_expert_compatible, -) from vllm.model_executor.layers.layernorm import GemmaRMSNorm from vllm.model_executor.layers.linear import ColumnParallelLinear from vllm.model_executor.layers.logits_processor import LogitsProcessor @@ -42,6 +39,7 @@ PPMissingLayer, WeightsMapper, get_draft_quant_config, + get_fused_shared_expert_layers, make_empty_intermediate_tensors_factory, maybe_fuse_shared_experts, maybe_prefix, @@ -228,10 +226,10 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: ) for idx in range(self.num_mtp_layers) ) - self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible( - self.layers, + self.fused_shared_expert_layers = get_fused_shared_expert_layers( + self, Qwen4ExpSparseMoeBlock, - "mlp", + enable_eplb=vllm_config.parallel_config.enable_eplb, ) self.pre_fc_norm_embedding = GemmaRMSNorm( @@ -384,7 +382,7 @@ def forward( def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: weights = maybe_fuse_shared_experts( weights, - enabled=self.is_fused_shared_expert_enabled, + fused_shared_expert_layers=self.fused_shared_expert_layers, n_routed_experts=getattr(self.config, "num_experts", 0) or 0, n_shared_experts=1, ckpt_prefix="mlp.shared_expert", diff --git a/vllm/triton_utils/dispatcher.py b/vllm/triton_utils/dispatcher.py index c7e2f9e46986..7d1c2e79aa0e 100644 --- a/vllm/triton_utils/dispatcher.py +++ b/vllm/triton_utils/dispatcher.py @@ -198,8 +198,22 @@ def lookup(value: Any) -> Callable[..., Any] | None: # kernel as an instance attribute; read the instance dict directly # so module-level __getattr__ hooks are never triggered, and so # this scan needs no dependency on the warmup machinery. - if hasattr(value, "__dict__"): - replacement = lookup(vars(value).get("kernel")) + # + # `hasattr(value, "__dict__")` is not a sufficient guard: proxy + # objects such as llvmlite's `binding.ffi.lib` synthesise every + # attribute through `__getattr__`, so `__dict__` resolves to a + # `_lib_fn_wrapper` and `.get()` would raise AttributeError. Only + # treat a real dict as an instance namespace. + try: + instance_dict = vars(value) + except TypeError: + # Slots-only objects have no instance dict. + instance_dict = None + replacement = ( + lookup(instance_dict.get("kernel")) + if isinstance(instance_dict, dict) + else None + ) if replacement is None: continue value.kernel = replacement