fix: #552 int8 kv cache quantization cannot be enabled

csrc/layers/attention/attention.cpp

@@ -39,10 +39,9 @@ Attention::Attention(std::shared_ptr<infinilm::config::ModelConfig> model_config
     rotary_emb_ = infinilm::layers::rotary_embedding::get_rope(model_config, device);
 
     float scaling = 1.0f / std::sqrt(static_cast<float>(head_dim_));
+    init_kv_cache_quant_params(register_fn, device, kv_cache_k_scale_, kv_cache_v_scale_);
     attn_ = std::make_shared<AttentionLayer>(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::function<void(const std::string &, infinico
         break;
     case infinilm::quantization::KVQuantAlgo::INT8:
         kv_cache_k_scale = infinicore::nn::Parameter({1}, infinicore::DataType::F32, device, 0, 0, 1);
+        kv_cache_k_scale.load(infinicore::Tensor::ones({1}, infinicore::DataType::F32, device));
         register_fn("kv_cache_k_scale", kv_cache_k_scale);
         kv_cache_v_scale = infinicore::nn::Parameter({1}, infinicore::DataType::F32, device, 0, 0, 1);
+        kv_cache_v_scale.load(infinicore::Tensor::ones({1}, infinicore::DataType::F32, device));
         register_fn("kv_cache_v_scale", kv_cache_v_scale);
         break;
     default:

csrc/layers/attention/backends/static_attn.cpp

@@ -30,18 +30,19 @@ infinicore::Tensor StaticAttentionImpl::forward(const AttentionLayer &layer,
 
     auto k_scale = layer.get_k_scale();
     auto v_scale = layer.get_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]
+
     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];

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());

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,

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,
     )

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:

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

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":

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(

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] = []

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()