From e0b03b17a96dde713b504b8e112e37bf90426b22 Mon Sep 17 00:00:00 2001 From: bobodai Date: Fri, 28 Aug 2026 01:06:46 +0800 Subject: [PATCH] fix: correct INT8 kv cache quantization path --- csrc/layers/attention/attention.cpp | 5 +- .../layers/attention/backends/static_attn.cpp | 11 ++- csrc/layers/quantization/kv_quant.cpp | 10 +- examples/bench.py | 93 ++++++++++++++----- examples/test_infer.py | 5 + python/infinilm/config/engine_config.py | 2 + python/infinilm/llm/llm.py | 4 + .../infinilm/llm/model_runner/model_runner.py | 1 + .../llm/model_runner/speculative_runner.py | 1 + python/infinilm/modeling_utils.py | 14 ++- python/infinilm/server/inference_server.py | 5 + python/infinilm/server/pipeline_worker.py | 1 + 12 files changed, 118 insertions(+), 34 deletions(-) diff --git a/csrc/layers/attention/attention.cpp b/csrc/layers/attention/attention.cpp index 16506ef02..0f6859efd 100644 --- a/csrc/layers/attention/attention.cpp +++ b/csrc/layers/attention/attention.cpp @@ -39,10 +39,9 @@ Attention::Attention(std::shared_ptr model_config rotary_emb_ = infinilm::layers::rotary_embedding::get_rope(model_config, device); float scaling = 1.0f / std::sqrt(static_cast(head_dim_)); + init_kv_cache_quant_params(register_fn, device, kv_cache_k_scale_, kv_cache_v_scale_); attn_ = std::make_shared(num_attention_heads_, head_dim_, scaling, num_key_value_heads_, layer_idx_, kv_cache_k_scale_, kv_cache_v_scale_, attention_backend_); - - init_kv_cache_quant_params(register_fn, device, kv_cache_k_scale_, kv_cache_v_scale_); } infinicore::Tensor Attention::forward(const infinicore::Tensor &positions, @@ -146,8 +145,10 @@ void init_kv_cache_quant_params(std::functionpermute({0, 2, 1, 3}); // [bs, n_q_head, seq_len, head_dim] + auto k_permuted = k_reshaped->permute({0, 2, 1, 3}); // [bs, n_kv_head, seq_len, head_dim] + auto v_permuted = v_reshaped->permute({0, 2, 1, 3}); // [bs, n_kv_head, seq_len, head_dim] + if (infinilm::quantization::KVQuantAlgo::NONE != this->kv_quant_scheme_) { infinilm::KVQuantUtils::quantize( - k_reshaped, v_reshaped, + k_permuted, v_permuted, this->kv_quant_scheme_, k_scale, v_scale); } - auto q_reshaped = q_rope->permute({0, 2, 1, 3}); // [bs, n_q_head, seq_len, head_dim] - auto k_permuted = k_reshaped->permute({0, 2, 1, 3}); // [bs, n_kv_head, seq_len, head_dim] - auto v_permuted = v_reshaped->permute({0, 2, 1, 3}); // [bs, n_kv_head, seq_len, head_dim] - // Prepare Attn auto shape = q_reshaped->shape(); size_t batch_size = shape[0]; diff --git a/csrc/layers/quantization/kv_quant.cpp b/csrc/layers/quantization/kv_quant.cpp index 458e568b0..3d1d3e972 100644 --- a/csrc/layers/quantization/kv_quant.cpp +++ b/csrc/layers/quantization/kv_quant.cpp @@ -16,8 +16,11 @@ void KVQuantUtils::quantize( } auto device = k->device(); - auto dtype = k->dtype(); - auto zero_point = infinicore::Tensor::zeros({1}, dtype, device); + + // INT8 symmetric quantization: zero_point must be F32 (the infiniop int8 + // kernels read it as float*); a bf16 zero tensor would be misread and can + // fault or corrupt the output. + auto zero_point = infinicore::Tensor::zeros({1}, infinicore::DataType::F32, device); k = infinicore::op::per_tensor_quant_i8(k, k_scale, zero_point, true); v = infinicore::op::per_tensor_quant_i8(v, v_scale, zero_point, true); @@ -35,7 +38,8 @@ void KVQuantUtils::dequantize( return; // 无需反量化 } - auto zero_point = infinicore::Tensor::zeros({1}, reference->dtype(), reference->device()); + // zero_point must be F32 (int8 dequant kernel reads it as float*) + auto zero_point = infinicore::Tensor::zeros({1}, infinicore::DataType::F32, reference->device()); auto k_dequant = infinicore::Tensor::strided_empty( k->shape(), k->strides(), reference->dtype(), reference->device()); diff --git a/examples/bench.py b/examples/bench.py index 17bfe1a6d..05d22a8da 100644 --- a/examples/bench.py +++ b/examples/bench.py @@ -4,6 +4,7 @@ import sys import time from collections import OrderedDict +from typing import Optional import infinicore import numpy as np @@ -25,6 +26,7 @@ "bfloat16": 2, "float16": 2, "float32": 4, + "int8": 1, } _PAGED_KV_BLOCK_SIZE = 256 @@ -97,12 +99,23 @@ def read_json_file(file_path): return json.load(file) +def _format_bytes(num_bytes: int) -> str: + """Format a byte count with an auto-selected unit (B / KB / MB / GB).""" + value = float(num_bytes) + for unit in ("B", "KB", "MB", "GB"): + if value < 1024 or unit == "GB": + return f"{int(value)} B" if unit == "B" else f"{value:.2f} {unit}" + value /= 1024 + return f"{value:.2f} GB" # unreachable, keeps the type checker happy + + def get_test_cases( model_path: str, batch_size_list: list[int], input_len_list: list[int], output_len_list: list[int], use_mla: bool = False, + kv_cache_dtype: Optional[str] = None, ): model_path = os.path.expanduser(model_path) @@ -127,38 +140,45 @@ def get_test_cases( num_key_value_heads = config.get("num_key_value_heads") num_hidden_layers = config.get("num_hidden_layers") + # KV cache dtype determines the per-element size. Quantized KV cache + # (--kv-cache-dtype int8/fp8) halves the storage vs bf16; the case line + # reports the actual dtype so the memory estimate matches the allocation. + if kv_cache_dtype in DATA_TYPE_BYTES: + data_type = kv_cache_dtype + else: + data_type = "bfloat16" + data_type_bytes = DATA_TYPE_BYTES[data_type] + # Enumerate all batch/input/output combinations and compute KV cache size case_list = [] for batch_size in batch_size_list: for input_len in input_len_list: for output_len in output_len_list: - for data_type in ["bfloat16"]: - data_type_bytes = DATA_TYPE_BYTES[data_type] - - total_seq_len = input_len + output_len - kvcache_memory_bytes = ( - data_type_bytes - * (batch_size * total_seq_len * num_key_value_heads * head_dim) - * num_hidden_layers - ) - kvcache_memory_gb = kvcache_memory_bytes / (1024 * 1024 * 1024) - - case_list.append( - { - "idx": len(case_list), - "batch_size": batch_size, - "input_len": input_len, - "output_len": output_len, - "data_type": data_type, - "kvcache_memory": round(kvcache_memory_gb, 3), - } - ) + total_seq_len = input_len + output_len + # Each KV cache layer stores both K and V (leading dim = 2). + kvcache_memory_bytes = ( + 2 + * data_type_bytes + * (batch_size * total_seq_len * num_key_value_heads * head_dim) + * num_hidden_layers + ) + case_list.append( + { + "idx": len(case_list), + "batch_size": batch_size, + "input_len": input_len, + "output_len": output_len, + "data_type": data_type, + "kvcache_memory": _format_bytes(kvcache_memory_bytes), + "kvcache_memory_bytes": kvcache_memory_bytes, + } + ) # Sort by KV cache size and wrap in OrderedDict with index keys case_dict = OrderedDict( (idx, case) for idx, case in enumerate( - sorted(case_list, key=lambda case: case["kvcache_memory"]) + sorted(case_list, key=lambda case: case["kvcache_memory_bytes"]) ) ) @@ -641,6 +661,20 @@ def get_input_tokens(self): def uses_pipeline_parallel(self) -> bool: return self.pp > 1 + def measure_kv_cache_memory(self) -> int: + """Measure the real allocated KV cache size in bytes (not an estimate).""" + total = 0 + try: + for rank_caches in self.model.get_kv_cache(): + for cache in rank_caches: + dtype_name = str(cache.dtype).split(".")[-1] + total += cache.numel() * DATA_TYPE_BYTES.get(dtype_name, 2) + except Exception: + # Pipeline-parallel / paged setups may not expose the cache list; + # fall back to the estimate already shown in the case line. + pass + return total + def close(self) -> None: if self.uses_pipeline_parallel: self.model.close() @@ -848,7 +882,12 @@ def run( input_len = [natural_input_len] cases_dict = get_test_cases( - model_path, batch_size, input_len, output_len, use_mla=cfg.use_mla + model_path, + batch_size, + input_len, + output_len, + use_mla=cfg.use_mla, + kv_cache_dtype=cfg.kv_cache_dtype, ) max_benchmark_batch_size = max(case["batch_size"] for case in cases_dict.values()) max_benchmark_tokens = max(case["output_len"] for case in cases_dict.values()) @@ -1053,6 +1092,14 @@ def run( ) ) + # Real allocated KV cache size, measured from the live tensors + # (the case-line value above is a config-based estimate). + measured_kv = test.measure_kv_cache_memory() + if measured_kv > 0: + tqdm.write( + f"[bench] measured KV cache memory: {_format_bytes(measured_kv)}" + ) + # run test one case test.run( batch_size=batch_size, diff --git a/examples/test_infer.py b/examples/test_infer.py index f17a8baa0..4c4ee9a48 100644 --- a/examples/test_infer.py +++ b/examples/test_infer.py @@ -41,6 +41,7 @@ def test( use_legacy_moe=False, enable_prefix_caching=True, pre_transpose=False, + kv_cache_dtype=None, ): model_path = os.path.expanduser(model_path) # ---------------------------------------------------------------------------- # @@ -77,6 +78,7 @@ def test( use_legacy_moe=use_legacy_moe, enable_prefix_caching=enable_prefix_caching, pre_transpose=pre_transpose, + kv_cache_dtype=kv_cache_dtype, ) conversations = [ @@ -154,6 +156,8 @@ def test( cfg.tp, cfg.dp, cfg.ep, cfg.moe_ep_backend, cfg.model ) + kv_cache_dtype = cfg.kv_cache_dtype + test( prompts, model_path, @@ -185,4 +189,5 @@ def test( use_legacy_moe=cfg.use_legacy_moe, enable_prefix_caching=cfg.enable_prefix_caching, pre_transpose=cfg.pre_transpose, + kv_cache_dtype=kv_cache_dtype, ) diff --git a/python/infinilm/config/engine_config.py b/python/infinilm/config/engine_config.py index bccb7e758..67fe35065 100644 --- a/python/infinilm/config/engine_config.py +++ b/python/infinilm/config/engine_config.py @@ -37,6 +37,7 @@ class EngineConfig: weight_load_mode: Weight loading mode across tensor-parallel workers. skip_load: Whether to skip loading model weights (for testing). use_legacy_moe: Whether to use the legacy Qwen3 MoE implementation. + kv_cache_dtype: Data type for the KV cache to trade off memory and precision (e.g."fp8", "int8", "bf16"). """ model_path: str @@ -69,6 +70,7 @@ class EngineConfig: use_legacy_moe: bool = False kv_transfer_config: Optional[KVTransferConfig] = None enable_prefix_caching: bool = True + kv_cache_dtype: Optional["str"] = None def __post_init__(self) -> None: if self.num_draft_tokens < 1: diff --git a/python/infinilm/llm/llm.py b/python/infinilm/llm/llm.py index 59d5a1eca..fdece586f 100644 --- a/python/infinilm/llm/llm.py +++ b/python/infinilm/llm/llm.py @@ -366,6 +366,7 @@ def __init__( skip_load: bool = False, use_legacy_moe: bool = False, enable_prefix_caching: bool = True, + kv_cache_dtype: Optional["str"] = None, ): """Initialize LLM. @@ -418,6 +419,7 @@ def __init__( skip_load=skip_load, use_legacy_moe=use_legacy_moe, enable_prefix_caching=enable_prefix_caching, + kv_cache_dtype=kv_cache_dtype, ) self.engine = LLMEngine(config) self.config = config @@ -594,6 +596,7 @@ def __init__( weight_load_mode: str = "async", use_legacy_moe: bool = False, enable_prefix_caching: bool = True, + kv_cache_dtype: Optional["str"] = None, ): """Initialize AsyncLLMEngine. @@ -651,6 +654,7 @@ def __init__( weight_load_mode=weight_load_mode, use_legacy_moe=use_legacy_moe, enable_prefix_caching=enable_prefix_caching, + kv_cache_dtype=kv_cache_dtype, ) self.engine = LLMEngine(config) self.config = config diff --git a/python/infinilm/llm/model_runner/model_runner.py b/python/infinilm/llm/model_runner/model_runner.py index a1696f848..760560586 100644 --- a/python/infinilm/llm/model_runner/model_runner.py +++ b/python/infinilm/llm/model_runner/model_runner.py @@ -91,6 +91,7 @@ def __init__(self, config: EngineConfig, initialize_processor: bool = True): weight_load_mode=config.weight_load_mode, use_legacy_moe=config.use_legacy_moe, pre_transpose=config.pre_transpose, + kv_cache_dtype=config.kv_cache_dtype, ) if self.model_engine.model_type == "minicpm_eagle": diff --git a/python/infinilm/llm/model_runner/speculative_runner.py b/python/infinilm/llm/model_runner/speculative_runner.py index d3c5211d4..359589a33 100644 --- a/python/infinilm/llm/model_runner/speculative_runner.py +++ b/python/infinilm/llm/model_runner/speculative_runner.py @@ -26,6 +26,7 @@ def __init__(self, config, target_model_engine, device): attention_backend="default", use_mla=False, weight_load_mode=config.weight_load_mode, + kv_cache_dtype=config.kv_cache_dtype, ) if self.draft_model_engine.model_type != "minicpm_eagle": raise RuntimeError( diff --git a/python/infinilm/modeling_utils.py b/python/infinilm/modeling_utils.py index b210b0c56..735a36897 100644 --- a/python/infinilm/modeling_utils.py +++ b/python/infinilm/modeling_utils.py @@ -68,6 +68,18 @@ def _is_internal_moe_packed_weight(key: str) -> bool: ) +def _is_internal_kv_cache_scale(key: str) -> bool: + """Per-tensor KV-cache quantization scales are registered by the attention + layer only when --kv-cache-dtype int8 is enabled, and HF checkpoints + never contain them. They are initialized to 1.0 in C++ and are + expected missing keys during checkpoint loading. + """ + + return key.endswith(".self_attn.kv_cache_k_scale") or key.endswith( + ".self_attn.kv_cache_v_scale" + ) + + def check_parameters(model_keys: list, already_loaded_keys: list): model_keys = set(model_keys) already_loaded_keys = set(already_loaded_keys) @@ -76,7 +88,7 @@ def check_parameters(model_keys: list, already_loaded_keys: list): missing_keys = { key for key in model_keys - intersection - if not _is_internal_moe_packed_weight(key) + if not (_is_internal_moe_packed_weight(key) or _is_internal_kv_cache_scale(key)) } unexpected_keys = already_loaded_keys - intersection error_msgs: list[str] = [] diff --git a/python/infinilm/server/inference_server.py b/python/infinilm/server/inference_server.py index 462c31084..0604525a3 100644 --- a/python/infinilm/server/inference_server.py +++ b/python/infinilm/server/inference_server.py @@ -125,6 +125,7 @@ def __init__( kv_transfer_config: Optional[KVTransferConfig] = None, enable_prefix_caching: bool = True, pre_transpose: bool = False, + kv_cache_dtype: Optional["str"] = None, ): """Initialize inference server. @@ -154,6 +155,7 @@ def __init__( weight_load_mode: Weight loading mode across tensor-parallel workers. ignore_eos: Whether to ignore EOS tokens during generation. kv_transfer_config: Optional configuration for the KV transfer mechanism. + kv_cache_dtype: Data type for the KV cache to trade off memory and precision (e.g."fp8", "int8", "bf16"). """ self.model_path = model_path # vLLM-like served model id: directory name of model_path @@ -188,6 +190,7 @@ def __init__( self.kv_transfer_config = kv_transfer_config self.enable_prefix_caching = enable_prefix_caching self.pre_transpose = pre_transpose + self.kv_cache_dtype = kv_cache_dtype self.engine: AsyncLLMEngine = None @@ -232,6 +235,7 @@ async def lifespan(app: FastAPI): kv_transfer_config=self.kv_transfer_config, enable_prefix_caching=self.enable_prefix_caching, pre_transpose=self.pre_transpose, + kv_cache_dtype=self.kv_cache_dtype, ) self.engine.start() logger.info(f"Engine initialized with model at {self.model_path}") @@ -667,6 +671,7 @@ def main(): kv_transfer_config=kv_transfer_config, enable_prefix_caching=cfg.enable_prefix_caching, pre_transpose=cfg.pre_transpose, + kv_cache_dtype=cfg.kv_cache_dtype, ) server.start() diff --git a/python/infinilm/server/pipeline_worker.py b/python/infinilm/server/pipeline_worker.py index 9e4bb8e6b..e414ad6be 100644 --- a/python/infinilm/server/pipeline_worker.py +++ b/python/infinilm/server/pipeline_worker.py @@ -38,6 +38,7 @@ def run_worker(cfg: BaseConfig) -> None: weight_load_mode=cfg.weight_load_mode, skip_load=cfg.skip_load, use_legacy_moe=cfg.use_legacy_moe, + kv_cache_dtype=cfg.kv_cache_dtype, ) runner = ModelRunner(config, initialize_processor=False)