[ROCm] [Bugfix] Enable Load and Inference of GLM-5.3-Flash Quark MXFP4 Checkpoint (#56176)

Signed-off-by: Colin Zeng <Colin.Zeng@amd.com>
Co-authored-by: Andreas Karatzas <akaratza@amd.com>

tests/quantization/test_quark.py

@@ -1668,3 +1668,148 @@ def test_suffix_match_at_module_boundary(prefix, ignored_layer, expected):
         )
         is expected
     )
+
+
+_GLM5_MXFP4_WEIGHT = {
+    "dtype": "fp4",
+    "qscheme": "per_group",
+    "group_size": 32,
+    "scale_format": "e8m0",
+    "is_dynamic": False,
+}
+_GLM5_BLOCK_FP8_WEIGHT = {
+    "dtype": "fp8_e4m3",
+    "qscheme": "per_block",
+    "is_dynamic": False,
+    "block_size": [128, 128],
+    "symmetric": True,
+}
+_GLM5_BLOCK_FP8_INPUT = {
+    "dtype": "fp8_e4m3",
+    "qscheme": "per_group",
+    "is_dynamic": True,
+    "group_size": 128,
+    "symmetric": True,
+}
+_GLM5_GATE_UP = "language_model.model.layers.0.mlp.gate_up_proj"
+
+
+def _glm5_mixed_precision_config() -> QuarkConfig:
+    """Global MXFP4 with the dense-MLP gate/up shards marked block-FP8."""
+    fp8 = {"weight": _GLM5_BLOCK_FP8_WEIGHT, "input_tensors": _GLM5_BLOCK_FP8_INPUT}
+    return QuarkConfig(
+        {
+            "global_quant_config": {
+                "weight": _GLM5_MXFP4_WEIGHT,
+                "input_tensors": None,
+            },
+            "layer_type_quant_config": {},
+            "layer_quant_config": {"*mlp.gate_proj": fp8, "*mlp.up_proj": fp8},
+            "exclude": [],
+        }
+    )
+
+
+class _GLM5RecordingParam:
+    """Stand-in for a BF16 target param that records what gets loaded."""
+
+    def __init__(self):
+        self.loaded_weight: torch.Tensor | None = None
+        self.loaded_shard: object = "unset"
+
+    def weight_loader(self, param, loaded_weight, shard_id=None):
+        self.loaded_weight = loaded_weight
+        self.loaded_shard = shard_id
+
+
+def _glm5_block_fp8(out_dim: int, in_dim: int, block: int = 128):
+    """Return a (fp8 weight, f32 per-block scale) pair."""
+    weight = (torch.randn(out_dim, in_dim) * 0.1).to(torch.float8_e4m3fn)
+    scale = torch.rand(out_dim // block, in_dim // block, dtype=torch.float32) + 0.5
+    return weight, scale
+
+
+def test_glm5next_gate_up_proj_mapping_is_genuine_fusion():
+    from vllm.models.glm5next.common.model import Glm5NextForConditionalGeneration
+
+    mapping = Glm5NextForConditionalGeneration.packed_modules_mapping
+    assert mapping["gate_up_proj"] == ["gate_proj", "up_proj"]
+
+
+def test_glm5next_genuine_fusion_resolves_gate_up_to_block_fp8():
+    # The fused module must expand to its real shard names so each resolves to
+    # the per-layer block-FP8 entry (rather than the global MXFP4 scheme).
+    config = _glm5_mixed_precision_config()
+    config.packed_modules_mapping = {"gate_up_proj": ["gate_proj", "up_proj"]}
+    _, _, scheme_cls = config.get_scheme_cls(LinearBase, _GLM5_GATE_UP)
+    assert scheme_cls is QuarkW8A8Fp8PerBlock
+
+
+def test_glm5next_attn_loader_accepts_quark_weight_scale():
+    # Regression for KeyError on '...kv_a_proj_with_mqa.weight_scale': the fused
+    # q_a/kv_a projection is kept BF16 and dequantized on load; the Quark scale
+    # name must be recognized and routed to the correct fused shard.
+    from vllm.models.glm5next.common.model import (
+        _dequant_fp8_block,
+        _try_load_fp8_attn_proj,
+    )
+
+    prefix = "layers.3.self_attn"
+    target = _GLM5RecordingParam()
+    params_dict = {f"{prefix}.fused_qkv_a_proj.weight": target}
+    buf: dict = {}
+    loaded: set = set()
+
+    weight, scale = _glm5_block_fp8(256, 256)
+    # fp8 weight arrives first -> buffered, nothing loaded yet.
+    assert (
+        _try_load_fp8_attn_proj(
+            f"{prefix}.kv_a_proj_with_mqa.weight", weight, buf, params_dict, loaded, 0
+        )
+        is True
+    )
+    assert target.loaded_weight is None
+    # Quark-named scale completes the pair -> dequantize + load.
+    assert (
+        _try_load_fp8_attn_proj(
+            f"{prefix}.kv_a_proj_with_mqa.weight_scale",
+            scale,
+            buf,
+            params_dict,
+            loaded,
+            0,
+        )
+        is True
+    )
+    assert target.loaded_weight is not None
+    assert target.loaded_weight.dtype == torch.bfloat16
+    assert target.loaded_shard == 1  # kv_a is shard 1 of fused_qkv_a_proj
+    assert torch.equal(target.loaded_weight, _dequant_fp8_block(weight, scale, 128))
+    assert f"{prefix}.fused_qkv_a_proj.weight" in loaded
+
+
+def test_glm5next_attn_loader_accepts_deepseek_weight_scale_inv():
+    # DeepSeek scale name keeps working unchanged.
+    from vllm.models.glm5next.common.model import (
+        _dequant_fp8_block,
+        _try_load_fp8_attn_proj,
+    )
+
+    prefix = "layers.7.self_attn"
+    target = _GLM5RecordingParam()
+    params_dict = {f"{prefix}.o_proj.weight": target}
+    buf: dict = {}
+    loaded: set = set()
+
+    weight, scale = _glm5_block_fp8(128, 256)
+    _try_load_fp8_attn_proj(
+        f"{prefix}.o_proj.weight", weight, buf, params_dict, loaded, 0
+    )
+    assert (
+        _try_load_fp8_attn_proj(
+            f"{prefix}.o_proj.weight_scale_inv", scale, buf, params_dict, loaded, 0
+        )
+        is True
+    )
+    assert target.loaded_shard is None  # o_proj is a direct (non-fused) proj
+    assert torch.equal(target.loaded_weight, _dequant_fp8_block(weight, scale, 128))

vllm/models/glm5next/__init__.py

@@ -6,8 +6,8 @@ from vllm.platforms import current_platform
 if current_platform.is_xpu():
     raise NotImplementedError("GLM-5.3-Flash does not currently support XPU.")
 
-from .nvidia.model import Glm5NextForCausalLM, Glm5NextForConditionalGeneration
-from .nvidia.mtp import Glm5NextMTP
+from .common.model import Glm5NextForCausalLM, Glm5NextForConditionalGeneration
+from .common.mtp import Glm5NextMTP
 
 __all__ = [
     "Glm5NextForCausalLM",

vllm/models/glm5next/common/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

vllm/models/glm5next/common/model.py

@@ -1036,6 +1036,17 @@ class Glm5NextForConditionalGeneration(
     has_inner_state: ClassVar[Literal[True]] = True
     is_hybrid: ClassVar[Literal[True]] = True
 
+    # GLM-5.3-Flash stores the dense-MLP gate/up as separate tensors (like
+    # ``Glm4vMoeForConditionalGeneration``, ``glm4_moe`` and ``deepseek_v2``),
+    # so the fused ``gate_up_proj`` must expand to its real shard names for
+    # per-layer quant-scheme resolution. The identity ``gate_up_proj`` entry
+    # inherited from ``Glm4vForConditionalGeneration`` (pre-fused gate_up_proj)
+    # would otherwise route the module to ``global_quant_config`` and mismatch
+    # at load for mixed-precision Quark checkpoints.
+    packed_modules_mapping = {
+        "gate_up_proj": ["gate_proj", "up_proj"],
+    }
+
     # NOTE: weight-prefix mapping is inherited from Glm4vForConditionalGeneration
     # (``model.visual.`` -> ``visual.``, ``model.language_model.`` ->
     # ``language_model.model.``, ``lm_head.`` -> ``language_model.lm_head.``),
@@ -1259,7 +1270,9 @@ def _try_load_fp8_attn_proj(
         return False
     suffix, (key, target_base, shard_id, is_kva) = matched
     is_weight = name.endswith(".weight") and tensor.dtype == torch.float8_e4m3fn
-    is_scale = "weight_scale_inv" in name
+    # Need to accept both the DeepSeek-native ``weight_scale_inv`` and the Quark
+    # ``weight_scale`` names before feeding the shared block dequant below.
+    is_scale = "weight_scale_inv" in name or name.endswith(".weight_scale")
     if not is_weight and not is_scale:
         return False
 

vllm/models/glm5next/common/mtp.py

@@ -22,6 +22,7 @@ from vllm.model_executor.model_loader.weight_utils import (
 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.models.glm5next.nvidia.ops.fused_eh_norm import fused_eh_norm
 from vllm.platforms import current_platform
 from vllm.sequence import IntermediateTensors
 
@@ -33,7 +34,6 @@ from .model import (
     _try_load_fp8_indexer_wk,
     get_spec_layer_idx_from_weight_name,
 )
-from .ops.fused_eh_norm import fused_eh_norm
 
 
 class Glm5NextMultiTokenPredictorLayer(nn.Module):