[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>
Colin Z committed
Sep 15, 2026 at 18:25 UTC
c8d1cf077a7878ea30040a33ac61262003a3953a
8 files changed
+164
-4
tests/quantization/test_quark.py
+145
@@ -1668,3 +1668,148 @@ def test_suffix_match_at_module_boundary(prefix, ignored_layer, expected):
1668
)
1669
is expected
1670
)
1671
+
1672
+
1673
+_GLM5_MXFP4_WEIGHT = {
1674
+ "dtype": "fp4",
1675
+ "qscheme": "per_group",
1676
+ "group_size": 32,
1677
+ "scale_format": "e8m0",
1678
+ "is_dynamic": False,
1679
+}
1680
+_GLM5_BLOCK_FP8_WEIGHT = {
1681
+ "dtype": "fp8_e4m3",
1682
+ "qscheme": "per_block",
1683
+ "is_dynamic": False,
1684
+ "block_size": [128, 128],
1685
+ "symmetric": True,
1686
+}
1687
+_GLM5_BLOCK_FP8_INPUT = {
1688
+ "dtype": "fp8_e4m3",
1689
+ "qscheme": "per_group",
1690
+ "is_dynamic": True,
1691
+ "group_size": 128,
1692
+ "symmetric": True,
1693
+}
1694
+_GLM5_GATE_UP = "language_model.model.layers.0.mlp.gate_up_proj"
1695
+
1696
+
1697
+def _glm5_mixed_precision_config() -> QuarkConfig:
1698
+ """Global MXFP4 with the dense-MLP gate/up shards marked block-FP8."""
1699
+ fp8 = {"weight": _GLM5_BLOCK_FP8_WEIGHT, "input_tensors": _GLM5_BLOCK_FP8_INPUT}
1700
+ return QuarkConfig(
1701
+ {
1702
+ "global_quant_config": {
1703
+ "weight": _GLM5_MXFP4_WEIGHT,
1704
+ "input_tensors": None,
1705
+ },
1706
+ "layer_type_quant_config": {},
1707
+ "layer_quant_config": {"*mlp.gate_proj": fp8, "*mlp.up_proj": fp8},
1708
+ "exclude": [],
1709
+ }
1710
+ )
1711
+
1712
+
1713
+class _GLM5RecordingParam:
1714
+ """Stand-in for a BF16 target param that records what gets loaded."""
1715
+
1716
+ def __init__(self):
1717
+ self.loaded_weight: torch.Tensor | None = None
1718
+ self.loaded_shard: object = "unset"
1719
+
1720
+ def weight_loader(self, param, loaded_weight, shard_id=None):
1721
+ self.loaded_weight = loaded_weight
1722
+ self.loaded_shard = shard_id
1723
+
1724
+
1725
+def _glm5_block_fp8(out_dim: int, in_dim: int, block: int = 128):
1726
+ """Return a (fp8 weight, f32 per-block scale) pair."""
1727
+ weight = (torch.randn(out_dim, in_dim) * 0.1).to(torch.float8_e4m3fn)
1728
+ scale = torch.rand(out_dim // block, in_dim // block, dtype=torch.float32) + 0.5
1729
+ return weight, scale
1730
+
1731
+
1732
+def test_glm5next_gate_up_proj_mapping_is_genuine_fusion():
1733
+ from vllm.models.glm5next.common.model import Glm5NextForConditionalGeneration
1734
+
1735
+ mapping = Glm5NextForConditionalGeneration.packed_modules_mapping
1736
+ assert mapping["gate_up_proj"] == ["gate_proj", "up_proj"]
1737
+
1738
+
1739
+def test_glm5next_genuine_fusion_resolves_gate_up_to_block_fp8():
1740
+ # The fused module must expand to its real shard names so each resolves to
1741
+ # the per-layer block-FP8 entry (rather than the global MXFP4 scheme).
1742
+ config = _glm5_mixed_precision_config()
1743
+ config.packed_modules_mapping = {"gate_up_proj": ["gate_proj", "up_proj"]}
1744
+ _, _, scheme_cls = config.get_scheme_cls(LinearBase, _GLM5_GATE_UP)
1745
+ assert scheme_cls is QuarkW8A8Fp8PerBlock
1746
+
1747
+
1748
+def test_glm5next_attn_loader_accepts_quark_weight_scale():
1749
+ # Regression for KeyError on '...kv_a_proj_with_mqa.weight_scale': the fused
1750
+ # q_a/kv_a projection is kept BF16 and dequantized on load; the Quark scale
1751
+ # name must be recognized and routed to the correct fused shard.
1752
+ from vllm.models.glm5next.common.model import (
1753
+ _dequant_fp8_block,
1754
+ _try_load_fp8_attn_proj,
1755
+ )
1756
+
1757
+ prefix = "layers.3.self_attn"
1758
+ target = _GLM5RecordingParam()
1759
+ params_dict = {f"{prefix}.fused_qkv_a_proj.weight": target}
1760
+ buf: dict = {}
1761
+ loaded: set = set()
1762
+
1763
+ weight, scale = _glm5_block_fp8(256, 256)
1764
+ # fp8 weight arrives first -> buffered, nothing loaded yet.
1765
+ assert (
1766
+ _try_load_fp8_attn_proj(
1767
+ f"{prefix}.kv_a_proj_with_mqa.weight", weight, buf, params_dict, loaded, 0
1768
+ )
1769
+ is True
1770
+ )
1771
+ assert target.loaded_weight is None
1772
+ # Quark-named scale completes the pair -> dequantize + load.
1773
+ assert (
1774
+ _try_load_fp8_attn_proj(
1775
+ f"{prefix}.kv_a_proj_with_mqa.weight_scale",
1776
+ scale,
1777
+ buf,
1778
+ params_dict,
1779
+ loaded,
1780
+ 0,
1781
+ )
1782
+ is True
1783
+ )
1784
+ assert target.loaded_weight is not None
1785
+ assert target.loaded_weight.dtype == torch.bfloat16
1786
+ assert target.loaded_shard == 1 # kv_a is shard 1 of fused_qkv_a_proj
1787
+ assert torch.equal(target.loaded_weight, _dequant_fp8_block(weight, scale, 128))
1788
+ assert f"{prefix}.fused_qkv_a_proj.weight" in loaded
1789
+
1790
+
1791
+def test_glm5next_attn_loader_accepts_deepseek_weight_scale_inv():
1792
+ # DeepSeek scale name keeps working unchanged.
1793
+ from vllm.models.glm5next.common.model import (
1794
+ _dequant_fp8_block,
1795
+ _try_load_fp8_attn_proj,
1796
+ )
1797
+
1798
+ prefix = "layers.7.self_attn"
1799
+ target = _GLM5RecordingParam()
1800
+ params_dict = {f"{prefix}.o_proj.weight": target}
1801
+ buf: dict = {}
1802
+ loaded: set = set()
1803
+
1804
+ weight, scale = _glm5_block_fp8(128, 256)
1805
+ _try_load_fp8_attn_proj(
1806
+ f"{prefix}.o_proj.weight", weight, buf, params_dict, loaded, 0
1807
+ )
1808
+ assert (
1809
+ _try_load_fp8_attn_proj(
1810
+ f"{prefix}.o_proj.weight_scale_inv", scale, buf, params_dict, loaded, 0
1811
+ )
1812
+ is True
1813
+ )
1814
+ assert target.loaded_shard is None # o_proj is a direct (non-fused) proj
1815
+ assert torch.equal(target.loaded_weight, _dequant_fp8_block(weight, scale, 128))
vllm/models/glm5next/__init__.py
+2
-2
@@ -6,8 +6,8 @@ from vllm.platforms import current_platform
6
if current_platform.is_xpu():
7
raise NotImplementedError("GLM-5.3-Flash does not currently support XPU.")
8
9
-from .nvidia.model import Glm5NextForCausalLM, Glm5NextForConditionalGeneration
10
-from .nvidia.mtp import Glm5NextMTP
9
+from .common.model import Glm5NextForCausalLM, Glm5NextForConditionalGeneration
10
+from .common.mtp import Glm5NextMTP
11
12
__all__ = [
13
"Glm5NextForCausalLM",
vllm/models/glm5next/common/__init__.py
new
+2
@@ -0,0 +1,2 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
vllm/models/glm5next/common/attention.py
renamed
vllm/models/glm5next/common/kda.py
renamed
vllm/models/glm5next/common/model.py
renamed
+14
-1
@@ -1036,6 +1036,17 @@ class Glm5NextForConditionalGeneration(
1036
has_inner_state: ClassVar[Literal[True]] = True
1037
is_hybrid: ClassVar[Literal[True]] = True
1038
1039
+ # GLM-5.3-Flash stores the dense-MLP gate/up as separate tensors (like
1040
+ # ``Glm4vMoeForConditionalGeneration``, ``glm4_moe`` and ``deepseek_v2``),
1041
+ # so the fused ``gate_up_proj`` must expand to its real shard names for
1042
+ # per-layer quant-scheme resolution. The identity ``gate_up_proj`` entry
1043
+ # inherited from ``Glm4vForConditionalGeneration`` (pre-fused gate_up_proj)
1044
+ # would otherwise route the module to ``global_quant_config`` and mismatch
1045
+ # at load for mixed-precision Quark checkpoints.
1046
+ packed_modules_mapping = {
1047
+ "gate_up_proj": ["gate_proj", "up_proj"],
1048
+ }
1049
+
1050
# NOTE: weight-prefix mapping is inherited from Glm4vForConditionalGeneration
1051
# (``model.visual.`` -> ``visual.``, ``model.language_model.`` ->
1052
# ``language_model.model.``, ``lm_head.`` -> ``language_model.lm_head.``),
@@ -1259,7 +1270,9 @@ def _try_load_fp8_attn_proj(
1270
return False
1271
suffix, (key, target_base, shard_id, is_kva) = matched
1272
is_weight = name.endswith(".weight") and tensor.dtype == torch.float8_e4m3fn
1262
- is_scale = "weight_scale_inv" in name
1273
+ # Need to accept both the DeepSeek-native ``weight_scale_inv`` and the Quark
1274
+ # ``weight_scale`` names before feeding the shared block dequant below.
1275
+ is_scale = "weight_scale_inv" in name or name.endswith(".weight_scale")
1276
if not is_weight and not is_scale:
1277
return False
1278
vllm/models/glm5next/common/mtp.py
renamed
+1
-1
@@ -22,6 +22,7 @@ from vllm.model_executor.model_loader.weight_utils import (
22
from vllm.model_executor.models.deepseek_mtp import SharedHead
23
from vllm.model_executor.models.deepseek_v2 import DeepseekV2MixtureOfExperts
24
from vllm.model_executor.models.utils import maybe_prefix
25
+from vllm.models.glm5next.nvidia.ops.fused_eh_norm import fused_eh_norm
26
from vllm.platforms import current_platform
27
from vllm.sequence import IntermediateTensors
28
@@ -33,7 +34,6 @@ from .model import (
34
_try_load_fp8_indexer_wk,
35
get_spec_layer_idx_from_weight_name,
36
)
36
-from .ops.fused_eh_norm import fused_eh_norm
37
38
39
class Glm5NextMultiTokenPredictorLayer(nn.Module):