[ROCm][Perf][GLM-5.3-Flash] Enable AITER fused shared experts on gfx950 (#59221)

Signed-off-by: Jin Tao <jin.tao@amd.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Simon Danielsson <70206058+simondanielsson@users.noreply.github.com>

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)],

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

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,