EngineX/xc-llm-ascend
3
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

[Feature]Supports DSv3.1 PD separation and C8 quantization (#7222)

Browse Source 
Co-authored-by: kunpengW-code <1289706727@qq.com>
Co-authored-by: linsheng1 <1950916997@qq.com>

### What this PR does / why we need it?
Currently, chunked prefill is forcibly enabled. DeepSeek V3.1 W8A8C8
supports only the PD separation scenario. C8 refers to quantizing the KV
cache to int8, which aims to reduce the GPU memory usage of the KV cache
and improve the inference throughput.
Constraints: 
1. Only the PD separation mode can be used and
MooncakeLayerwiseConnector can be used to run the model.
2. Currently, only the activation value supports dynamic quantization,
and the KV cache supports static quantization. C8 quantization with MTP
is not supported. You can use ModelSlim for quantization. The
quantization procedure is as follows:
pip install transformers==4.48.2
git clone https://gitcode.com/Ascend/msmodelslim.git
cd msmodelslim
bash install.sh
cd example/DeepSeek/
python3 quant_deepseek_w8a8.py --model_path <path/weight> --save_path
<path/quant_weight>
--anti_dataset../common/deepseek_anti_prompt_50_v3_1.json
--calib_dataset../common/deepseek_calib_prompt_50_v3_1.json --rot
--trust_remote_code True --fa_quant --dynamic --anti_method m6

### Does this PR introduce _any_ user-facing change?
no
### How was this patch tested?

- vLLM version: v0.17.0
- vLLM main:
4034c3d32e

---------

Signed-off-by: pichangping <1337510399@qq.com>
Signed-off-by: Wang Kunpeng <1289706727@qq.com>
Co-authored-by: Wang Kunpeng <1289706727@qq.com>
This commit is contained in:
pichangping
2026-03-16 22:49:05 +08:00
committed by GitHub
parent a6f6e919e6
commit 3f39ac9c8d
15 changed files with 1112 additions and 161 deletions
Download Patch File Download Diff File Expand all files Collapse all files

57
vllm_ascend/quantization/modelslim_config.py
View File

@@ -24,6 +24,7 @@ configs generated by the ModelSlim tool, along with model-specific mappings.
import glob
import json
import os
import re
from collections.abc import Mapping
from types import MappingProxyType
from typing import Any, Optional
@@ -31,6 +32,7 @@ from typing import Any, Optional
import torch
from vllm.config import get_current_vllm_config
from vllm.logger import init_logger
from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
from vllm.model_executor.layers.fused_moe import FusedMoE
from vllm.model_executor.layers.linear import LinearBase
from vllm.model_executor.layers.quantization import register_quantization_config
@@ -38,7 +40,7 @@ from vllm.model_executor.layers.quantization.base_config import QuantizationConf
from vllm.model_executor.layers.vocab_parallel_embedding import UnquantizedEmbeddingMethod, VocabParallelEmbedding
from vllm.model_executor.models.utils import WeightsMapper
from vllm_ascend.utils import ASCEND_QUANTIZATION_METHOD
from vllm_ascend.utils import ASCEND_QUANTIZATION_METHOD, calc_split_factor
from .methods import get_scheme_class
@@ -438,6 +440,7 @@ class AscendModelSlimConfig(QuantizationConfig):
new_k = k.replace("weight_packed", "weight")
extra_quant_dict[new_k] = self.quant_description[k]
self.quant_description.update(extra_quant_dict)
self._add_kvcache_quant_metadata()
def __repr__(self) -> str:
return "AscendModelSlimConfig:\n" + super().__repr__()
@@ -509,8 +512,6 @@ class AscendModelSlimConfig(QuantizationConfig):
self.packed_modules_mapping = packed_modules_model_mapping[model_type]
prefix = self.quant_prefix_mapper(model_type, prefix)
from vllm.model_executor.layers.attention import Attention
if model_type != "kimi_k2":
if prefix.startswith("language_model"):
prefix = prefix.split(".", 1)[-1]
@@ -522,11 +523,7 @@ class AscendModelSlimConfig(QuantizationConfig):
return AscendUnquantizedLinearMethod()
scheme = create_scheme_for_layer(self.quant_description, prefix, "linear", self.packed_modules_mapping)
return AscendLinearMethod(scheme)
elif (
isinstance(layer, Attention)
and "fa_quant_type" in self.quant_description
and self.quant_description["fa_quant_type"] is not None
):
elif isinstance(layer, AttentionLayerBase) and self.is_fa_quant_layer(prefix):
scheme = create_scheme_for_layer(self.quant_description, prefix, "attention", self.packed_modules_mapping)
return AscendKVCacheMethod(scheme)
elif isinstance(layer, FusedMoE):
@@ -573,6 +570,39 @@ class AscendModelSlimConfig(QuantizationConfig):
assert is_skipped is not None
return is_skipped
def is_fa_quant_layer(self, prefix):
if self.enable_fa_quant:
layer_id_str = "".join(re.findall(r"\.(\d+)\.", prefix))
if layer_id_str.isdigit() and int(layer_id_str) in self.kvcache_quant_layers:
return True
return False
def enabling_fa_quant(self, vllm_config, layer_name) -> bool:
is_decode_instance = (
vllm_config.kv_transfer_config is not None
and vllm_config.kv_transfer_config.is_kv_consumer
and not vllm_config.kv_transfer_config.is_kv_producer
)
return bool(is_decode_instance and self.is_fa_quant_layer(layer_name))
def get_kv_quant_dtype(self, layer_name, cache_dtype, model_config):
if self.enable_fa_quant and self.is_fa_quant_layer(layer_name):
ori_dtype = model_config.dtype
quant_dtype = torch.int8
# For MLA models like deepseek, we only quantify K cache to ensure accuracy
if model_config.use_mla:
return quant_dtype, ori_dtype
else:
return quant_dtype, quant_dtype
return cache_dtype, cache_dtype
def get_kv_quant_split_factor(self, layer_name, kv_head_dim_list):
if self.enable_fa_quant and self.is_fa_quant_layer(layer_name):
k_quant_head_dim = kv_head_dim_list[0]
v_quant_head_dim = kv_head_dim_list[1] * 2
kv_head_dim_list = [k_quant_head_dim, v_quant_head_dim]
return calc_split_factor(kv_head_dim_list)
def maybe_update_config(self, model_name: str, revision: str | None = None) -> None:
"""Load the ModelSlim quantization config from model directory.
@@ -606,6 +636,7 @@ class AscendModelSlimConfig(QuantizationConfig):
with open(config_path) as f:
self.quant_description = json.load(f)
self._apply_extra_quant_adaptations()
self._add_kvcache_quant_metadata()
return
# Collect diagnostic info for the error message
@@ -678,3 +709,13 @@ class AscendModelSlimConfig(QuantizationConfig):
def get_scaled_act_names(self) -> list[str]:
return []
def _add_kvcache_quant_metadata(self):
fa_quant_type = self.quant_description.get("fa_quant_type", "")
self.enable_fa_quant = fa_quant_type != ""
self.kvcache_quant_layers = []
if self.enable_fa_quant:
for key in self.quant_description:
if "fa_k.scale" in key:
_id = "".join(re.findall(r"\.(\d+)\.", key))
self.kvcache_quant_layers.append(int(_id))