From d9cf85d5e2ee8052672a671dc1ab19ec927dd920 Mon Sep 17 00:00:00 2001 From: akshatvishu Date: Fri, 3 Jul 2026 09:41:28 +0530 Subject: [PATCH] [BugFix] Use KV cache scale mapper in custom loaders (#4810) Signed-off-by: akshatvishu --- .../models/test_kv_cache_scale_mapper.py | 148 ++++++++++++++++++ .../models/hunyuan_image3/hunyuan_image3.py | 10 +- .../models/mammoth_moda2/mammoth_moda2.py | 12 +- .../models/mimo_audio/mimo_audio_llm.py | 13 +- .../models/qwen2_5_omni/qwen2_old.py | 9 -- 5 files changed, 159 insertions(+), 33 deletions(-) create mode 100644 tests/model_executor/models/test_kv_cache_scale_mapper.py diff --git a/tests/model_executor/models/test_kv_cache_scale_mapper.py b/tests/model_executor/models/test_kv_cache_scale_mapper.py new file mode 100644 index 000000000..98f016bb4 --- /dev/null +++ b/tests/model_executor/models/test_kv_cache_scale_mapper.py @@ -0,0 +1,148 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Regression tests for vLLM's KV-cache scale mapper load path.""" + +from collections.abc import Iterable +from pathlib import Path +from types import SimpleNamespace + +import pytest +import torch + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +_REPO_ROOT = Path(__file__).resolve().parents[3] +_STALE_API_FILES = [ + "vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py", + "vllm_omni/model_executor/models/mammoth_moda2/mammoth_moda2.py", + "vllm_omni/model_executor/models/mimo_audio/mimo_audio_llm.py", + "vllm_omni/model_executor/models/qwen2_5_omni/qwen2_old.py", +] + +_SOURCE_SCALE_NAME = "layers.0.self_attn.k_proj.output_scale" +_MAPPED_SCALE_NAME = "layers.0.self_attn.attn.k_scale" + + +class _RecordingCacheScaleMapper: + def __init__(self) -> None: + self.called = False + + def apply( + self, + weights: Iterable[tuple[str, torch.Tensor]], + ) -> Iterable[tuple[str, torch.Tensor]]: + self.called = True + for name, weight in weights: + yield name.replace(".k_proj.output_scale", ".attn.k_scale"), weight + + +class _QuantConfigWithoutOldCacheScale: + def __init__(self) -> None: + self.mapper = _RecordingCacheScaleMapper() + + def get_cache_scale_mapper(self) -> _RecordingCacheScaleMapper: + return self.mapper + + +class _FakeModel: + def __init__(self, quant_config: _QuantConfigWithoutOldCacheScale) -> None: + self.quant_config = quant_config + self.param = torch.nn.Parameter(torch.zeros(()), requires_grad=False) + + def named_parameters( + self, + remove_duplicate: bool = True, + ) -> Iterable[tuple[str, torch.nn.Parameter]]: + del remove_duplicate + return [(_MAPPED_SCALE_NAME, self.param)] + + +class _FakeHunyuanModel(_FakeModel): + def __init__(self, quant_config: _QuantConfigWithoutOldCacheScale) -> None: + super().__init__(quant_config) + self.config = SimpleNamespace( + tie_word_embeddings=False, + num_attention_heads=1, + num_key_value_heads=1, + use_cla=False, + ) + + def get_expert_mapping(self) -> tuple[list[object], dict[str, object]]: + return [], {} + + def _split_qkv_weight(self, weight: torch.Tensor) -> torch.Tensor: + return weight + + +@pytest.mark.parametrize("rel_path", _STALE_API_FILES) +def test_model_loaders_do_not_call_removed_get_cache_scale(rel_path: str) -> None: + """Affected loaders must not call the removed vLLM get_cache_scale API.""" + source = (_REPO_ROOT / rel_path).read_text(encoding="utf-8") + assert ".get_cache_scale(" not in source + + +@pytest.mark.parametrize( + ("module_name", "class_name", "fake_model_cls"), + [ + ( + "vllm_omni.model_executor.models.hunyuan_image3.hunyuan_image3", + "HunyuanModel", + _FakeHunyuanModel, + ), + ( + "vllm_omni.model_executor.models.qwen2_5_omni.qwen2_old", + "Qwen2Model", + _FakeModel, + ), + ], +) +def test_auto_weights_loader_delegated_loaders_accept_mapped_cache_scale( + module_name: str, + class_name: str, + fake_model_cls: type[_FakeModel], + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Delegated loaders must not expect the removed get_cache_scale method.""" + module = __import__(module_name, fromlist=[class_name]) + monkeypatch.setattr(module, "is_pp_missing_parameter", lambda *_: False) + + fake_model = fake_model_cls(_QuantConfigWithoutOldCacheScale()) + load_weights = getattr(module, class_name).load_weights + + loaded = load_weights(fake_model, [(_MAPPED_SCALE_NAME, torch.ones(()))]) + + assert loaded == {_MAPPED_SCALE_NAME} + torch.testing.assert_close(fake_model.param.data, torch.ones(())) + + +@pytest.mark.parametrize( + ("module_name", "class_name"), + [ + ( + "vllm_omni.model_executor.models.mammoth_moda2.mammoth_moda2", + "MammothModa2Qwen2ForCausalLM", + ), + ( + "vllm_omni.model_executor.models.mimo_audio.mimo_audio_llm", + "MiMoAudioLLMForConditionalGeneration", + ), + ], +) +def test_direct_custom_loaders_apply_cache_scale_mapper( + module_name: str, + class_name: str, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Direct custom loaders must apply the new mapper before fall-through load.""" + module = __import__(module_name, fromlist=[class_name]) + monkeypatch.setattr(module, "is_pp_missing_parameter", lambda *_: False) + + quant_config = _QuantConfigWithoutOldCacheScale() + fake_model = _FakeModel(quant_config) + load_weights = getattr(module, class_name).load_weights + + loaded = load_weights(fake_model, [(_SOURCE_SCALE_NAME, torch.ones(()))]) + + assert quant_config.mapper.called + assert loaded == {_MAPPED_SCALE_NAME} + torch.testing.assert_close(fake_model.param.data, torch.ones(())) diff --git a/vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py b/vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py index d9efccd9a..17bbad001 100755 --- a/vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py +++ b/vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py @@ -222,14 +222,8 @@ class HunyuanModel(HunYuanModel): if self.config.tie_word_embeddings and "lm_head.weight" in name: continue - if self.quant_config is not None and (scale_name := self.quant_config.get_cache_scale(name)): - # Loading kv cache scales for compressed-tensors quantization - param = params_dict[scale_name] - weight_loader = getattr(param, "weight_loader", default_weight_loader) - loaded_weight = loaded_weight[0] - weight_loader(param, loaded_weight) - continue - + # KV-cache scales are renamed to .attn.{k,v,q}_scale by the + # quant_config mapper in the outer AutoWeightsLoader. is_found = False for param_name, weight_name, shard_id in stacked_params_mapping: if weight_name not in name: diff --git a/vllm_omni/model_executor/models/mammoth_moda2/mammoth_moda2.py b/vllm_omni/model_executor/models/mammoth_moda2/mammoth_moda2.py index 6311216dc..2fe616cfd 100644 --- a/vllm_omni/model_executor/models/mammoth_moda2/mammoth_moda2.py +++ b/vllm_omni/model_executor/models/mammoth_moda2/mammoth_moda2.py @@ -483,19 +483,15 @@ class MammothModa2Qwen2ForCausalLM(nn.Module, SupportsPP): ] params_dict = dict(self.named_parameters(remove_duplicate=False)) loaded_params: set[str] = set() + if self.quant_config is not None: + cache_scale_mapper = self.quant_config.get_cache_scale_mapper() + if cache_scale_mapper is not None: + weights = cache_scale_mapper.apply(weights) for name, loaded_weight in weights: if "rotary_emb.inv_freq" in name: continue - if self.quant_config is not None and (scale_name := self.quant_config.get_cache_scale(name)): - param = params_dict[scale_name] - weight_loader = getattr(param, "weight_loader", default_weight_loader) - loaded_weight = loaded_weight if loaded_weight.dim() == 0 else loaded_weight[0] - weight_loader(param, loaded_weight) - loaded_params.add(scale_name) - continue - for param_name, weight_name, shard_id in stacked_params_mapping: if weight_name not in name: continue diff --git a/vllm_omni/model_executor/models/mimo_audio/mimo_audio_llm.py b/vllm_omni/model_executor/models/mimo_audio/mimo_audio_llm.py index a640ade3a..48715fb98 100644 --- a/vllm_omni/model_executor/models/mimo_audio/mimo_audio_llm.py +++ b/vllm_omni/model_executor/models/mimo_audio/mimo_audio_llm.py @@ -1153,17 +1153,14 @@ class MiMoAudioLLMForConditionalGeneration(nn.Module, SupportsMultiModal, Suppor params_dict = dict(self.named_parameters(remove_duplicate=False)) loaded_params: set[str] = set() + if self.quant_config is not None: + cache_scale_mapper = self.quant_config.get_cache_scale_mapper() + if cache_scale_mapper is not None: + weights = cache_scale_mapper.apply(weights) + for name, loaded_weight in weights: if name.startswith("model."): name = "model." + name - if self.quant_config is not None and (scale_name := self.quant_config.get_cache_scale(name)): - # Loading kv cache quantization scales - param = params_dict[scale_name] - weight_loader = getattr(param, "weight_loader", default_weight_loader) - loaded_weight = loaded_weight if loaded_weight.dim() == 0 else loaded_weight[0] - weight_loader(param, loaded_weight) - loaded_params.add(scale_name) - continue for param_name, weight_name, shard_id in stacked_params_mapping: if name.startswith("input_local_transformer.") or name.startswith("local_transformer."): diff --git a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_old.py b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_old.py index 67806f951..91eafd2a8 100644 --- a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_old.py +++ b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_old.py @@ -335,19 +335,10 @@ class Qwen2Model(nn.Module): ("gate_up_proj", "up_proj", 1), ] params_dict = dict(self.named_parameters(remove_duplicate=False)) - get_cache_scale = getattr(self.quant_config, "get_cache_scale", None) loaded_params: set[str] = set() for name, loaded_weight in weights: if "rotary_emb.inv_freq" in name: continue - if get_cache_scale is not None and (scale_name := get_cache_scale(name)): - # Loading kv cache quantization scales - param = params_dict[scale_name] - weight_loader = getattr(param, "weight_loader", default_weight_loader) - loaded_weight = loaded_weight if loaded_weight.dim() == 0 else loaded_weight[0] - weight_loader(param, loaded_weight) - loaded_params.add(scale_name) - continue for param_name, weight_name, shard_id in stacked_params_mapping: if weight_name not in name: continue