diff --git a/tests/model_executor/layers/test_fused_shared_expert.py b/tests/model_executor/layers/test_fused_shared_expert.py index 6721a4c63b..d8d5c4fb03 100644 --- a/tests/model_executor/layers/test_fused_shared_expert.py +++ b/tests/model_executor/layers/test_fused_shared_expert.py @@ -14,7 +14,7 @@ import torch from torch import nn import vllm.config as vllm_config_module -from vllm.config import VllmConfig, set_current_vllm_config +from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config 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.base_config import QuantizationConfig @@ -145,6 +145,12 @@ def _stub_quant_config() -> QuantizationConfig: return cast(QuantizationConfig, object()) +class _StubGlm5NextAttention(nn.Module): + def __init__(self, *args: Any, **kwargs: Any) -> None: + super().__init__() + self.o_proj = SimpleNamespace(reduce_results=True) + + def get_fse_test_model_config( model_type: str, quantization_config: dict[str, Any], @@ -279,6 +285,33 @@ def get_fse_test_model_config( quantization_config=quantization_config, ) return config, Glm4MoeModel + if model_type == "glm5_next": + from transformers import Glm5NextTextConfig + + from vllm.models.glm5next.common.model import Glm5NextModel + + config = Glm5NextTextConfig( + vocab_size=256, + hidden_size=128, + intermediate_size=32, + moe_intermediate_size=32, + num_hidden_layers=1, + 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"], + index_topk=1, + index_kpool=1, + pad_token_id=None, + scoring_func="sigmoid", + topk_method="noaux_tc", + mhc=False, + quantization_config=quantization_config, + ) + return config, Glm5NextModel raise ValueError(f"Unsupported FSE test model: {model_type}") @@ -796,7 +829,7 @@ def test_is_model_fused_shared_expert_compatible() -> None: @pytest.mark.parametrize( "model_type", - ["minimax_m3", "deepseek_v4", "qwen3_5", "glm4_moe", "deepseek_v2"], + ["minimax_m3", "deepseek_v4", "qwen3_5", "glm4_moe", "deepseek_v2", "glm5_next"], ) @pytest.mark.parametrize( ("use_fse", "exclude"), @@ -850,6 +883,16 @@ def test_models_fse_init( ) else: vllm_config.quant_config = QuarkConfig(quantization_config) + if model_type == "glm5_next": + from vllm.models.glm5next.common import model as glm5_next_model + + for attention in ("Glm5NextLinearAttention", "Glm5NextMLAAttention"): + monkeypatch.setattr(glm5_next_model, attention, _StubGlm5NextAttention) + # The gate allows fusion only at TP4/TP8 on gfx950, and this test + # runs at TP1; the gate has its own test. + monkeypatch.setattr( + glm5_next_model, "_fused_shared_experts_tuned", lambda _: True + ) import vllm.envs as envs from vllm._aiter_ops import rocm_aiter_ops @@ -902,6 +945,10 @@ def test_models_fse_init( ) mtp = DeepSeekV4MTP(vllm_config=vllm_config) 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 + ) # The dummy quant config here uses mixed mxfp4/fp8 for experts/shared_expert # so should just raise a warning. @@ -926,6 +973,74 @@ def test_models_fse_init( rocm_aiter_ops.refresh_env_variables() +@pytest.mark.parametrize( + ("setting", "value", "reason"), + [ + ("tensor_parallel_size", 4, None), + ("tensor_parallel_size", 8, None), + ("tensor_parallel_size", 2, "tensor_parallel_size is 2"), + ("data_parallel_size", 2, "data_parallel_size is 2"), + ("prefill_context_parallel_size", 2, "prefill_context_parallel_size is 2"), + ("enable_expert_parallel", True, "expert parallelism is enabled"), + ("on_gfx950", False, "the GPU is not gfx950"), + ], +) +def test_glm5_next_fuses_shared_experts_only_in_tuned_setups( + monkeypatch: pytest.MonkeyPatch, + setting: str, + value: object, + reason: str | None, +) -> None: + from vllm.models.glm5next.common import model as glm5_next_model + + parallel_config = SimpleNamespace( + tensor_parallel_size=4, + data_parallel_size=1, + prefill_context_parallel_size=1, + enable_expert_parallel=False, + ) + on_gfx950 = value if setting == "on_gfx950" else True + if setting != "on_gfx950": + setattr(parallel_config, setting, value) + monkeypatch.setattr("vllm.platforms.rocm.on_gfx950", lambda: on_gfx950) + + with patch.object(glm5_next_model.logger, "warning_once") as warning: + tuned = glm5_next_model._fused_shared_experts_tuned( + cast(ParallelConfig, parallel_config) + ) + + assert tuned is (reason is None) + if reason is None: + warning.assert_not_called() + else: + warning.assert_called_once() + assert warning.call_args.args[1] == reason + + +def test_glm5_next_fused_shared_expert_loads_into_the_shared_slot() -> None: + from vllm.models.glm5next.common import model as glm5_next_model + + assert glm5_next_model._num_fused_shared_experts(1, True) == 1 + assert glm5_next_model._num_fused_shared_experts(1, False) == 0 + assert glm5_next_model._num_fused_shared_experts(None, True) == 0 + with pytest.raises(NotImplementedError, match="only 1 shared expert"): + glm5_next_model._num_fused_shared_experts(2, True) + + for shared, fused in [ + ( + "model.layers.3.mlp.shared_experts.gate_proj.weight_scale_inv", + "model.layers.3.mlp.experts.288.gate_proj.weight_scale_inv", + ), + ( + "model.layers.45.mtp_block.mlp.shared_experts.down_proj.weight", + "model.layers.45.mtp_block.mlp.experts.288.down_proj.weight", + ), + ]: + assert glm5_next_model._fused_shared_expert_name(shared, 288) == fused + routed = "model.layers.3.mlp.experts.7.down_proj.weight" + assert glm5_next_model._fused_shared_expert_name(routed, 288) == routed + + @pytest.mark.parametrize( ("exclude", "expected"), [([], True), ([r"re:.*\.shared_expert\..*"], False)], diff --git a/vllm/models/glm5next/common/model.py b/vllm/models/glm5next/common/model.py index e15c00ecb8..778738c137 100644 --- a/vllm/models/glm5next/common/model.py +++ b/vllm/models/glm5next/common/model.py @@ -23,6 +23,10 @@ from vllm.model_executor.layers.fused_moe import ( GateLinear, 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 from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.linear import ( @@ -258,7 +262,12 @@ class Glm5NextMoE(nn.Module): ) swiglu_limit = config.swiglu_limit - if config.n_shared_experts is None: + self.is_fused_shared_expert_enabled = False + if config.n_shared_experts is not None: + self.is_fused_shared_expert_enabled = resolve_layer_fused_shared_expert( + quant_config, prefix + ) and _fused_shared_experts_tuned(parallel_config) + if config.n_shared_experts is None or self.is_fused_shared_expert_enabled: self.shared_experts = None else: intermediate_size = config.moe_intermediate_size * config.n_shared_experts @@ -294,7 +303,10 @@ class Glm5NextMoE(nn.Module): enable_eplb=self.enable_eplb, num_redundant_experts=self.n_redundant_experts, is_sequence_parallel=self.is_sequence_parallel, - n_shared_experts=None, + n_shared_experts=config.n_shared_experts + if self.is_fused_shared_expert_enabled + else None, + fuse_shared_experts=self.is_fused_shared_expert_enabled, router_logits_dtype=self.gate.out_dtype, swiglu_limit=swiglu_limit, ) @@ -733,6 +745,9 @@ class Glm5NextModel(nn.Module): # 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" + ) if get_pp_group().is_last_rank: self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) @@ -834,12 +849,15 @@ class Glm5NextModel(nn.Module): ), 0, ) + num_fused_shared = _num_fused_shared_experts( + self.config.n_shared_experts, self.is_fused_shared_expert_enabled + ) 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", - num_experts=self.config.n_routed_experts, + num_experts=self.config.n_routed_experts + num_fused_shared, num_redundant_experts=num_redundant_experts, ) else: @@ -868,6 +886,8 @@ class Glm5NextModel(nn.Module): # Models trained using ColossalAI may include these tensors in # the checkpoint. Skip them. continue + if self.is_fused_shared_expert_enabled: + name = _fused_shared_expert_name(name, self.config.n_routed_experts) # Handle FP8 indexer WK: dequantize to BF16 for fusion with # weights_proj into wk_weights_proj. @@ -1228,6 +1248,65 @@ def get_spec_layer_idx_from_weight_name( return None +def _fused_shared_experts_tuned(parallel_config: ParallelConfig) -> bool: + """AITER has fused-MoE configs tuned for the fused shared-expert shape + (one more expert and one more top-k slot than the routed MoE) only on + gfx950, with every expert on each rank and its weights split by TP4 or + TP8. Data, prefill context and expert parallelism change that split, so + any other GPU or parallel layout would run untuned fallback kernels.""" + from vllm.platforms.rocm import on_gfx950 + + reasons: list[str] = [] + if not on_gfx950(): + reasons.append("the GPU is not gfx950") + if parallel_config.tensor_parallel_size not in (4, 8): + reasons.append( + f"tensor_parallel_size is {parallel_config.tensor_parallel_size}" + ) + if parallel_config.data_parallel_size != 1: + reasons.append(f"data_parallel_size is {parallel_config.data_parallel_size}") + if parallel_config.prefill_context_parallel_size != 1: + reasons.append( + "prefill_context_parallel_size is " + f"{parallel_config.prefill_context_parallel_size}" + ) + if parallel_config.enable_expert_parallel: + reasons.append("expert parallelism is enabled") + + if not reasons: + return True + logger.warning_once( + "VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS is ignored for GLM-5.3-Flash: " + "%s. AITER has tuned configs for its fused shared-expert MoE only on " + "gfx950 at TP4 and TP8, without data, prefill context or expert " + "parallelism. Running the shared experts as a separate MLP.", + "; ".join(reasons), + ) + return False + + +def _num_fused_shared_experts(n_shared_experts: int | None, enabled: bool) -> int: + """Expert slots the fused MoE appends for the shared expert; must match the + ``num_fused_shared_experts`` that ``FusedMoE`` allocates.""" + if not enabled or n_shared_experts is None: + return 0 + if n_shared_experts > 1: + raise NotImplementedError( + "Fused shared-expert loading supports only 1 shared expert per " + f"layer, but config.n_shared_experts is {n_shared_experts}. Set " + "VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS=0 to run the shared " + "experts as a separate MLP." + ) + return n_shared_experts + + +def _fused_shared_expert_name(name: str, n_routed_experts: int) -> str: + """Point a checkpoint ``mlp.shared_experts.*`` tensor at the fused MoE's + shared-expert slot, which follows the routed experts; other names are + returned unchanged.""" + return name.replace("mlp.shared_experts.", f"mlp.experts.{n_routed_experts}.", 1) + + def _try_load_fp8_indexer_wk(name, tensor, buf, params_dict, loaded_params): if "indexer.wk." not in name or "wk_weights" in name: return False diff --git a/vllm/models/glm5next/common/mtp.py b/vllm/models/glm5next/common/mtp.py index 816cc46fdc..2bb8b10d07 100644 --- a/vllm/models/glm5next/common/mtp.py +++ b/vllm/models/glm5next/common/mtp.py @@ -10,6 +10,9 @@ 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 ( @@ -30,6 +33,8 @@ from .model import ( Glm5NextDecoderLayer, Glm5NextMLAAttention, Glm5NextMoE, + _fused_shared_expert_name, + _num_fused_shared_experts, _try_load_fp8_attn_proj, _try_load_fp8_indexer_wk, get_spec_layer_idx_from_weight_name, @@ -217,6 +222,9 @@ class Glm5NextMTP(nn.Module, DeepseekV2MixtureOfExperts): 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" + ) def set_moe_parameters(self): self.num_moe_layers = self.config.num_nextn_predict_layers @@ -298,12 +306,15 @@ class Glm5NextMTP(nn.Module, DeepseekV2MixtureOfExperts): ("wk_weights_proj", "wk", 0), ("wk_weights_proj", "weights_proj", 1), ] + num_fused_shared = _num_fused_shared_experts( + self.config.n_shared_experts, self.is_fused_shared_expert_enabled + ) 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", - num_experts=self.config.n_routed_experts, + num_experts=self.config.n_routed_experts + num_fused_shared, num_redundant_experts=self.num_redundant_experts, ) @@ -328,6 +339,8 @@ class Glm5NextMTP(nn.Module, DeepseekV2MixtureOfExperts): if spec_layer is None: continue name = self._rewrite_spec_layer_name(spec_layer, name) + if self.is_fused_shared_expert_enabled: + name = _fused_shared_expert_name(name, self.config.n_routed_experts) if _try_load_fp8_indexer_wk( name,