KV Cache int8

简介

此量化方式将k cache和v cache量化为8 bit，通过减少KV Cache的显存占用，提升吞吐。

KV Cache int8搭配W8A8量化后权重目录结构：

├─ config.json
├─ quant_model_weight_w8a8.safetensors
├─ quant_model_description_w8a8.json
├─ tokenizer_config.json
├─ tokenizer.json
└─ tokenizer.model

• 量化输出包含：权重文件quant_model_weight_w8a8.safetensors和权重描述文件quant_model_description_w8a8.json。
• 目录中的其余文件为推理时所需的配置文件，不同模型略有差异。

以下展示了量化后权重描述文件quant_model_description_w8a8.json中的部分内容：

{
  "model_quant_type": "W8A8",
  "kv_cache_type": "C8",
  "model.embed_tokens.weight": "FLOAT",
  "model.layers.0.self_attn.q_proj.weight": "FLOAT",
  "model.layers.0.self_attn.k_proj.weight": "FLOAT",
  "model.layers.0.self_attn.k_proj.kv_cache_scale": "W8A8",
  "model.layers.0.self_attn.k_proj.kv_cache_offset": "W8A8",
  "model.layers.0.self_attn.v_proj.weight": "FLOAT",
  "model.layers.0.self_attn.v_proj.kv_cache_scale": "W8A8",
  "model.layers.0.self_attn.v_proj.kv_cache_offset": "W8A8",
}

和W8A8量化权重相比，新增kv_cache_type描述字段，新增kv linear激活值量化缩放因子权重文件kv_cache_scale和kv linear激活值量化偏移值权重文件kv_cache_offset。推理时会基于这两个权重，推导出

k_quant_scale，k_dequant_scale，v_quant_scale，v_dequant_scale，k_quant_offset，k_dequant_offset，v_quant_offset和v_dequant_offset。其中quant_scale和quant_offset用于将k和v特征量化为int8类型，dequant_scale和dequant_offset用于将Paged attention的输出反量化为浮点类型。

生成权重

以LLaMa3.1-70B为例，您可以使用以下指令生成W8A8 + KV Cache int8量化权重。

export ASCEND_RT_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
cd ${ATB_SPEED_HOME_PATH}
python examples/models/llama3/convert_quant_weights.py --model_path {模型权重路径} 
--save_directory {量化模型保存路径} --w_bit 8 --a_bit 8 --device_type npu 
--act_method 3 
--anti_method m3 --disable_level L0 
--calib_file { boolq.jsonl路径}
--use_kvcache_quant True 

• 相比于W8A8的量化方式，需要新增use_kvcache_quant参数。
• KV Cache int8量化权重的config.json中应包含quantization_config字段，其中包含"kv_quant_type": "C8"键值对。

执行推理

和W8A8量化权重执行推理的方式相同，请参考W8A8。