From c8d1cf077a7878ea30040a33ac61262003a3953a Mon Sep 17 00:00:00 2001 From: Colin Z <59755453+ColinZ22@users.noreply.github.com> Date: Tue, 15 Sep 2026 18:25:29 -0700 Subject: [PATCH] [ROCm] [Bugfix] Enable Load and Inference of GLM-5.3-Flash Quark MXFP4 Checkpoint (#56176) Signed-off-by: Colin Zeng Co-authored-by: Andreas Karatzas --- tests/quantization/test_quark.py | 145 ++++++++++++++++++ vllm/models/glm5next/__init__.py | 4 +- vllm/models/glm5next/common/__init__.py | 2 + .../glm5next/{nvidia => common}/attention.py | 0 .../models/glm5next/{nvidia => common}/kda.py | 0 .../glm5next/{nvidia => common}/model.py | 15 +- .../models/glm5next/{nvidia => common}/mtp.py | 2 +- .../glm5next/{nvidia => common}/multimodal.py | 0 8 files changed, 164 insertions(+), 4 deletions(-) create mode 100644 vllm/models/glm5next/common/__init__.py rename vllm/models/glm5next/{nvidia => common}/attention.py (100%) rename vllm/models/glm5next/{nvidia => common}/kda.py (100%) rename vllm/models/glm5next/{nvidia => common}/model.py (98%) rename vllm/models/glm5next/{nvidia => common}/mtp.py (99%) rename vllm/models/glm5next/{nvidia => common}/multimodal.py (100%) diff --git a/tests/quantization/test_quark.py b/tests/quantization/test_quark.py index ae236d71f8..9f00adf783 100644 --- a/tests/quantization/test_quark.py +++ b/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)) diff --git a/vllm/models/glm5next/__init__.py b/vllm/models/glm5next/__init__.py index a1146e2602..b38806260b 100644 --- a/vllm/models/glm5next/__init__.py +++ b/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", diff --git a/vllm/models/glm5next/common/__init__.py b/vllm/models/glm5next/common/__init__.py new file mode 100644 index 0000000000..208f01a7cb --- /dev/null +++ b/vllm/models/glm5next/common/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm/models/glm5next/nvidia/attention.py b/vllm/models/glm5next/common/attention.py similarity index 100% rename from vllm/models/glm5next/nvidia/attention.py rename to vllm/models/glm5next/common/attention.py diff --git a/vllm/models/glm5next/nvidia/kda.py b/vllm/models/glm5next/common/kda.py similarity index 100% rename from vllm/models/glm5next/nvidia/kda.py rename to vllm/models/glm5next/common/kda.py diff --git a/vllm/models/glm5next/nvidia/model.py b/vllm/models/glm5next/common/model.py similarity index 98% rename from vllm/models/glm5next/nvidia/model.py rename to vllm/models/glm5next/common/model.py index 21918bdf6b..e4f669835e 100644 --- a/vllm/models/glm5next/nvidia/model.py +++ b/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 diff --git a/vllm/models/glm5next/nvidia/mtp.py b/vllm/models/glm5next/common/mtp.py similarity index 99% rename from vllm/models/glm5next/nvidia/mtp.py rename to vllm/models/glm5next/common/mtp.py index 35b452f44f..33a5402b8f 100644 --- a/vllm/models/glm5next/nvidia/mtp.py +++ b/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): diff --git a/vllm/models/glm5next/nvidia/multimodal.py b/vllm/models/glm5next/common/multimodal.py similarity index 100% rename from vllm/models/glm5next/nvidia/multimodal.py rename to vllm/models/glm5next/common/multimodal.py