Skip to content

Commit e43ac31

Browse files
menogreymenogrey
authored
Revert "[Misc][Quantization] Remove kvcache quant weight related patches." (vllm-project#11615)
Reverts vllm-project#11053 - vLLM version: v0.23.0 - vLLM main: vllm-project/vllm@1f486d9 Signed-off-by: menogrey <1299267905@qq.com>
1 parent 53dd79c commit e43ac31

6 files changed

Lines changed: 253 additions & 149 deletions

File tree

‎tests/ut/quantization/test_modelslim_config.py‎

Lines changed: 0 additions & 111 deletions
Original file line numberDiff line numberDiff line change
@@ -353,117 +353,6 @@ def test_step3p5_mtp_maps_direct_and_step3p7_wrapped_quant_keys(self):
353353
self.assertEqual(prefix, expected)
354354

355355

356-
class TestGetCacheScaleMapper(TestBase):
357-
def test_c8_kv_cache_type_returns_mapper(self):
358-
config = AscendModelSlimConfig({"kv_cache_type": "C8"})
359-
mapper = config.get_cache_scale_mapper()
360-
self.assertIsNotNone(mapper)
361-
# C8 mappings: k_proj → attn
362-
self.assertEqual(
363-
mapper._map_name("model.layers.0.k_proj.kv_cache_scale"),
364-
"model.layers.0.attn.k_cache_scale",
365-
)
366-
self.assertEqual(
367-
mapper._map_name("model.layers.0.k_proj.kv_cache_offset"),
368-
"model.layers.0.attn.k_cache_offset",
369-
)
370-
self.assertEqual(
371-
mapper._map_name("model.layers.0.v_proj.kv_cache_scale"),
372-
"model.layers.0.attn.v_cache_scale",
373-
)
374-
self.assertEqual(
375-
mapper._map_name("model.layers.0.v_proj.kv_cache_offset"),
376-
"model.layers.0.attn.v_cache_offset",
377-
)
378-
379-
def test_no_match_returns_none(self):
380-
config = AscendModelSlimConfig({"kv_cache_type": "FLOAT"})
381-
mapper = config.get_cache_scale_mapper()
382-
self.assertIsNone(mapper)
383-
384-
config = AscendModelSlimConfig({})
385-
mapper = config.get_cache_scale_mapper()
386-
self.assertIsNone(mapper)
387-
388-
def test_fa_quant_returns_mapper(self):
389-
config = AscendModelSlimConfig(
390-
{
391-
"fa_quant_type": "C8",
392-
"layers.1.fa_k.scale": "C8",
393-
}
394-
)
395-
mapper = config.get_cache_scale_mapper()
396-
self.assertIsNotNone(mapper)
397-
self.assertEqual(
398-
mapper._map_name("model.layers.1.fa_k.scale"),
399-
"model.layers.1.mla_attn.mla_attn.fa_k.scale",
400-
)
401-
self.assertEqual(
402-
mapper._map_name("model.layers.1.fa_q.scale"),
403-
"model.layers.1.mla_attn.mla_attn.fa_q.scale",
404-
)
405-
self.assertEqual(
406-
mapper._map_name("model.layers.1.fa_v.offset"),
407-
"model.layers.1.mla_attn.mla_attn.fa_v.offset",
408-
)
409-
410-
def test_indexer_quant_returns_mapper(self):
411-
config = AscendModelSlimConfig(
412-
{
413-
"indexer_quant_type": "INT8",
414-
"layers.1.indexer.quant_type": "INT8",
415-
}
416-
)
417-
mapper = config.get_cache_scale_mapper()
418-
self.assertIsNotNone(mapper)
419-
self.assertEqual(
420-
mapper._map_name("model.layers.1.indexer.q_rot"),
421-
"model.layers.1.mla_attn.mla_attn.indexer.q_rot",
422-
)
423-
self.assertEqual(
424-
mapper._map_name("model.layers.1.indexer.k_rot"),
425-
"model.layers.1.mla_attn.mla_attn.indexer.k_rot",
426-
)
427-
428-
def test_combined_quant_types_returns_mapper_with_all_suffixes(self):
429-
config = AscendModelSlimConfig(
430-
{
431-
"kv_cache_type": "C8",
432-
"fa_quant_type": "C8",
433-
"layers.1.fa_k.scale": "C8",
434-
"indexer_quant_type": "INT8",
435-
"layers.2.indexer.quant_type": "INT8",
436-
}
437-
)
438-
mapper = config.get_cache_scale_mapper()
439-
self.assertIsNotNone(mapper)
440-
# C8 mapping
441-
self.assertEqual(
442-
mapper._map_name("model.layers.0.k_proj.kv_cache_scale"),
443-
"model.layers.0.attn.k_cache_scale",
444-
)
445-
# FA quant mapping
446-
self.assertEqual(
447-
mapper._map_name("model.layers.1.fa_k.scale"),
448-
"model.layers.1.mla_attn.mla_attn.fa_k.scale",
449-
)
450-
# Indexer quant mapping
451-
self.assertEqual(
452-
mapper._map_name("model.layers.2.indexer.k_rot"),
453-
"model.layers.2.mla_attn.mla_attn.indexer.k_rot",
454-
)
455-
456-
def test_unmapped_key_returns_unchanged(self):
457-
config = AscendModelSlimConfig({"kv_cache_type": "C8"})
458-
mapper = config.get_cache_scale_mapper()
459-
self.assertIsNotNone(mapper)
460-
# C8 mapper doesn't cover FA quant keys
461-
self.assertEqual(
462-
mapper._map_name("model.layers.0.other_key"),
463-
"model.layers.0.other_key",
464-
)
465-
466-
467356
class TestGetKvQuantDtype(TestBase):
468357
def test_enable_fa_quant(self):
469358
config = AscendModelSlimConfig(

‎vllm_ascend/patch/__init__.py‎

Lines changed: 67 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -898,7 +898,44 @@
898898
# Future Plan:
899899
# Remove this patch when:
900900
# vLLM itself supports kv transfer for mamba
901-
# ** 21. File: worker/patch_v2/patch_input_batch.py
901+
# ** 21. File: worker/patch_weight_utils.py**
902+
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
903+
# 1. `vllm.model_executor.models.deepseek_v2.DeepseekV2ForCausalLM.load_weights`
904+
# Why:
905+
# The C8 weight quantized by modelslim will modify the model structure,
906+
# and the scale and offset required for kvcache quantization will increase.
907+
# In addition, the names of the quantization parameters are different from
908+
# those in the community.
909+
# How：
910+
# we have enhanced the maybe_remap_kv_scale_name function.
911+
# Future Plan:
912+
# The maybe_remap_kv_scale_name function of the community is reconstructed to support
913+
# multiple backends.
914+
# ** 21b. File: worker/patch_process_weights_after_loading.py**
915+
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
916+
# 1. `vllm.model_executor.model_loader.utils.process_weights_after_loading`
917+
# `vllm.model_executor.model_loader.base_loader.process_weights_after_loading`
918+
# and imported references in vllm-ascend model loaders
919+
# Why:
920+
# DSA attention is implemented in vllm-ascend as the plugin layer
921+
# `DSAAttention`. Upstream vLLM only runs post-load attention weight
922+
# processing for built-in attention classes, so
923+
# `DSAAttention.process_weights_after_loading()` is skipped in the
924+
# original loader flow. DSV4 DSA-CP o-proj TP initialization must run in
925+
# this post-load phase rather than being initialized lazily in forward.
926+
# How:
927+
# Rebind the upstream `process_weights_after_loading` helper, including
928+
# already-imported loader references, so `DSAAttention` participates in
929+
# the same post-load traversal while preserving the original quant-method
930+
# and torchao reload behavior.
931+
# Related PR (if no, explain why):
932+
# https://github.com/vllm-project/vllm-ascend/pull/10694
933+
# https://github.com/vllm-project/vllm/pull/46828
934+
# Future Plan:
935+
# Remove this patch once the supported vLLM version includes PR #46828.
936+
# Then register `DSAAttention` through vLLM's post-load weight-processing
937+
# registry instead of monkey-patching model-loader helpers.
938+
# ** 22. File: worker/patch_v2/patch_input_batch.py**
902939
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
903940
# 1. `vllm.v1.worker.gpu.input_batch.InputBatch`
904941
# Why:
@@ -908,7 +945,7 @@
908945
# replace InputBatch with AscendInputBatch.
909946
# Future Plan:
910947
# remove this patch when vLLM-ascend's make_dummy behavior aligns with vLLM.
911-
# ** 22. File: worker/patch_v2/patch_block_table.py**
948+
# ** 23. File: worker/patch_v2/patch_block_table.py**
912949
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
913950
# 1. `vllm.v1.worker.gpu.block_table.BlockTables`
914951
# Why:
@@ -920,7 +957,7 @@
920957
# Future Plan:
921958
# remove this patch when vLLM-ascend's BlockTables can initialize
922959
# slot mapping as torch.int64 dtype.
923-
# ** 23. File: worker/patch_v2/patch_model_state.py**
960+
# ** 24. File: worker/patch_v2/patch_model_state.py**
924961
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
925962
# 1. `vllm.v1.worker.gpu.model_states.default.init_model_state`
926963
# Why:
@@ -930,7 +967,7 @@
930967
# Define AscendModelState and initialize it in init_model_state.
931968
# Future Plan:
932969
# remove this when vllm-ascend's attention metadata is align with vllm.
933-
# ** 24. File: worker/patch_v2/patch_triton.py**
970+
# ** 25. File: worker/patch_v2/patch_triton.py**
934971
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
935972
# 1. `vllm.v1.worker.gpu.sample.logprob`, `vllm.v1.worker.gpu.sample.penalties.apply_penalties`,
936973
# `vllm.v1.worker.gpu.sample.gumbel.gumbel_sample`
@@ -943,7 +980,30 @@
943980
# Future Plan:
944981
# Remove this patch when vLLM support the dispatch function.
945982
#
946-
# ** 25. File: worker/patch_qwen3vl.py**
983+
# ** 26. File: worker/patch_gqa_c8.py**
984+
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
985+
# 1. `vllm.model_executor.models.qwen3.Qwen3ForCausalLM.load_weights`
986+
# Why:
987+
# The GQA W8A8C8 model stores per-channel KV cache scales and offsets
988+
# (k_cache_scale, k_cache_offset, v_cache_scale, v_cache_offset) under
989+
# weight names that AutoWeightsLoader does not recognise and would
990+
# silently discard. Without these scales the INT8 KV cache cannot be
991+
# dequantised correctly at inference time.
992+
# How:
993+
# Wrap load_weights to intercept the C8 scale/offset tensors before they
994+
# reach the base loader. Each intercepted tensor is routed to the
995+
# corresponding nn.Parameter via its weight_loader, then excluded from
996+
# the remaining weight stream so the base loader never sees it.
997+
# Related PR (if no, explain why):
998+
# This PR (Qwen3-32B and GLM4.7 W8A8C8 support). Upstream vLLM's weight-loading
999+
# pipeline does not yet have a generic hook for hardware-plugin-defined
1000+
# KV cache parameters.
1001+
# Future Plan:
1002+
# Remove this patch when vLLM provides a first-class extension point
1003+
# for loading extra KV cache quantisation parameters in model load_weights,
1004+
# or when the GQA model's weight names are aligned with the parameter
1005+
# names expected by the quantisation backend.
1006+
# ** 27. File: worker/patch_qwen3vl.py**
9471007
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
9481008
# 1. `vllm.model_executor.models.qwen3.Qwen3Attention.forward` and
9491009
# `vllm.model_executor.models.qwen3_moe.Qwen3MoeAttention.forward`
@@ -954,7 +1014,7 @@
9541014
# when using mrope.
9551015
# Future Plan:
9561016
# Remove this patch when vllm-ascend supports pattern matching for this fused kernel.
957-
# ** 26. File: worker/patch_qwen3_dflash.py**
1017+
# ** 28. File: worker/patch_qwen3_dflash.py**
9581018
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
9591019
# 1. `vllm.model_executor.models.qwen3_dflash.DFlashQwen3Model.precompute_and_store_context_kv`
9601020
# Why:
@@ -982,7 +1042,7 @@
9821042
# Remove this patch when upstream vLLM supports MoE communication type abstraction that
9831043
# can be extended by hardware plugins like vllm-ascend.
9841044
#
985-
# ** 27. File: platform/patch_mamba_manager.py**
1045+
# ** 29. File: platform/patch_mamba_manager.py**
9861046
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
9871047
# 1. `vllm.v1.core.single_type_kv_cache_manager.MambaManager`
9881048
# Why:

‎vllm_ascend/patch/worker/__init__.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,7 @@
3838

3939

4040
import vllm_ascend.patch.worker.patch_process_weights_after_loading # noqa
41+
import vllm_ascend.patch.worker.patch_weight_utils # noqa
4142
import vllm_ascend.patch.worker.patch_distributed # noqa
4243
import vllm_ascend.patch.worker.patch_minimax_m2 # noqa
4344
import vllm_ascend.patch.worker.patch_minimax_m2_linear_attn # noqa
@@ -65,6 +66,7 @@
6566
import vllm_ascend.patch.worker.patch_cudagraph # noqa
6667
import vllm_ascend.patch.worker.patch_deepseek_mtp # noqa
6768
import vllm_ascend.patch.worker.patch_deepseek_v2 # noqa
69+
import vllm_ascend.patch.worker.patch_gqa_c8 # noqa
6870

6971
# vLLM's use_v2_model_runner may enable the v2 runner without the
7072
# VLLM_USE_V2_MODEL_RUNNER env var (e.g. based on model architecture).

‎vllm_ascend/patch/worker/patch_gqa_c8.py‎

Lines changed: 78 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,78 @@
1+
#
2+
# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved.
3+
# This file is a part of the vllm-ascend project.
4+
#
5+
# Licensed under the Apache License, Version 2.0 (the "License");
6+
# you may not use this file except in compliance with the License.
7+
# You may obtain a copy of the License at
8+
#
9+
# http://www.apache.org/licenses/LICENSE-2.0
10+
#
11+
# Unless required by applicable law or agreed to in writing, software
12+
# distributed under the License is distributed on an "AS IS" BASIS,
13+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14+
# See the License for the specific language governing permissions and
15+
# limitations under the License.
16+
#
17+
18+
import logging
19+
from collections.abc import Callable, Iterable
20+
21+
import torch
22+
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
23+
from vllm.model_executor.models.glm4_moe import Glm4MoeForCausalLM
24+
from vllm.model_executor.models.minimax_m2 import MiniMaxM2ForCausalLM
25+
from vllm.model_executor.models.qwen3 import Qwen3ForCausalLM
26+
27+
logger = logging.getLogger(__name__)
28+
29+
_orig_qwen3_causal_lm_load_weights = Qwen3ForCausalLM.load_weights
30+
_orig_Glm4_causal_lm_load_weights = Glm4MoeForCausalLM.load_weights
31+
_orig_Minimax_m2_causal_lm_load_weights = MiniMaxM2ForCausalLM.load_weights
32+
33+
34+
def _patched_causal_lm_load_weights(
35+
self, weights: Iterable[tuple[str, torch.Tensor]], original_load_weights: Callable
36+
) -> set[str]:
37+
quant_config = self.quant_config
38+
if quant_config is None or not callable(getattr(quant_config, "get_cache_scale", None)):
39+
return original_load_weights(self, weights)
40+
41+
params_dict = dict(self.named_parameters())
42+
c8_loaded_params: set[str] = set()
43+
44+
def _intercept_c8_scales(
45+
raw_weights: Iterable[tuple[str, torch.Tensor]],
46+
) -> Iterable[tuple[str, torch.Tensor]]:
47+
for name, loaded_weight in raw_weights:
48+
scale_name = quant_config.get_cache_scale(name)
49+
if scale_name is not None:
50+
if scale_name in params_dict:
51+
param = params_dict[scale_name]
52+
weight_loader = getattr(param, "weight_loader", default_weight_loader)
53+
weight_loader(param, loaded_weight.squeeze())
54+
c8_loaded_params.add(scale_name)
55+
else:
56+
logger.warning(
57+
"Cache scale %s found in quant_config for weight %s "
58+
"but not found in model parameters; weight will be skipped.",
59+
scale_name,
60+
name,
61+
)
62+
else:
63+
yield name, loaded_weight
64+
65+
loaded_params = original_load_weights(self, _intercept_c8_scales(weights))
66+
loaded_params.update(c8_loaded_params)
67+
return loaded_params
68+
69+
70+
Qwen3ForCausalLM.load_weights = lambda self, weights: _patched_causal_lm_load_weights(
71+
self, weights, _orig_qwen3_causal_lm_load_weights
72+
)
73+
Glm4MoeForCausalLM.load_weights = lambda self, weights: _patched_causal_lm_load_weights(
74+
self, weights, _orig_Glm4_causal_lm_load_weights
75+
)
76+
MiniMaxM2ForCausalLM.load_weights = lambda self, weights: _patched_causal_lm_load_weights(
77+
self, weights, _orig_Minimax_m2_causal_lm_load_weights
78+
)

0 commit comments

Comments
 (0)