跳转到内容
LLM Primer

KV Cache 与量化：让大模型跑得更快

直觉版：记住已经算过的东西

Section titled “直觉版：记住已经算过的东西”

LLM 生成文本是一个词一个词地”续写”。如果每生成一个新词都要重新看一遍前面所有内容，那就会浪费大量计算。KV Cache 的直觉很简单：把前面 token 计算好的 Key 和 Value 向量存起来，下次直接用，不用再算一遍。

量化则是另一个降低成本的思路：模型权重通常是 32 位浮点数，如果改成 16 位、8 位甚至 4 位，显存占用和计算量都会大幅下降。虽然会损失一点精度，但工程上往往完全可接受。

KV Cache 内存占用示意图：

graph TB
    subgraph "无 KV Cache 的注意力计算"
        A1[Token 1] --> |Q1,K1,V1| B1[计算注意力]
        A2[Token 2] --> |Q2,K2,V2| B1
        A3[Token 3] --> |Q3,K3,V3| B1
        B1 --> C1[输出]
        style B1 fill:#ff9999
    end


    subgraph "使用 KV Cache 的生成过程"
        D1[步骤1: Token 1] --> |存储 K1,V1| E1[Cache: K1,V1]
        E1 --> F1[生成 Token 2]


        F1 --> |存储 K2,V2| E2[Cache: K1,V1,K2,V2]
        E2 --> F2[生成 Token 3]


        F2 --> |复用 K1,V1,K2,V2<br/>+ 新的 K3,V3| E3[Cache: K1,V1,K2,V2,K3,V3]
        E3 --> F3[生成 Token 4]


        style E1 fill:#99ff99
        style E2 fill:#99ff99
        style E3 fill:#99ff99
    end

内存占用计算公式：

KV Cache 大小=2×batch_size×num_layers×seq_len×num_kv_heads×head_dim×bytes_per_param\text{KV Cache 大小} = 2 \times \text{batch\_size} \times \text{num\_layers} \times \text{seq\_len} \times \text{num\_kv\_heads} \times \text{head\_dim} \times \text{bytes\_per\_param}

其中：

  • 2 表示 Key 和 Value 两个张量
  • 对于 7B 模型，batch_size=1, seq_len=4096 时约需 2-4 GB

工程版：内存墙与精度取舍

Section titled “工程版：内存墙与精度取舍”

KV Cache 是推理时的主要内存消耗之一，尤其在长上下文场景中。优化方向包括：

  • PagedAttention（vLLM）：把 KV Cache 分页管理，减少内存碎片，提高批处理效率。
  • KV Cache 压缩：通过量化、剪枝、滑动窗口或 Eviction 策略减少缓存体积。
  • GQA / MQA：让多个注意力头共享 Key/Value，降低缓存量。

量化技术按粒度分类：

  • PTQ（训练后量化）：GPTQ、AWQ、SmoothQuant 等，对已经训练好的模型直接量化，无需重新训练。
  • QAT（量化感知训练）：在训练中模拟低精度，效果通常更好但成本更高。
  • GGML/GGUF：社区常用的 4-bit 量化格式，让大模型能在笔记本上运行。

工程实践中，要评估量化后的困惑度（perplexity）、下游任务精度和端到端延迟，不能只看显存省了百分之几。不同层对精度敏感度不同，混合精度或逐层调优往往效果最好。

示例代码：KV Cache 的基本原理

Section titled “示例代码：KV Cache 的基本原理”
import numpy as np


class KVCacheAttention:
    """带 KV Cache 的简化注意力机制"""


    def __init__(self, d_model):
        self.d_model = d_model
        self.kv_cache = {"keys": [], "values": []}


    def compute_attention_with_cache(self, query, key, value, use_cache=True):
        """
        计算注意力，可选使用 KV Cache
        query: [1, d_model] - 新生成的 token
        key, value: [1, d_model] - 新 token 的 K, V
        """
        if use_cache:
            # 将新的 K, V 添加到缓存
            self.kv_cache["keys"].append(key)
            self.kv_cache["values"].append(value)


            # 使用所有历史 K, V
            all_keys = np.vstack(self.kv_cache["keys"])  # [seq_len, d_model]
            all_values = np.vstack(self.kv_cache["values"])
        else:
            all_keys = key
            all_values = value


        # 计算注意力分数
        scores = np.matmul(query, all_keys.T) / np.sqrt(self.d_model)
        weights = np.exp(scores) / np.sum(np.exp(scores))


        # 加权求和
        output = np.matmul(weights, all_values)
        return output, len(self.kv_cache["keys"])


# 示例：模拟自回归生成
d_model = 64
kv_cache_attn = KVCacheAttention(d_model)


print("模拟生成 5 个 token:")
for i in range(5):
    # 新 token 的 Q, K, V
    q = np.random.randn(1, d_model)
    k = np.random.randn(1, d_model)
    v = np.random.randn(1, d_model)


    output, cache_len = kv_cache_attn.compute_attention_with_cache(q, k, v)
    print(f"Token {i+1}: Cache 长度 = {cache_len}, 输出形状 = {output.shape}")


print(f"\n总缓存大小: {len(kv_cache_attn.kv_cache['keys'])} 个 token")

示例代码：简单的量化

Section titled “示例代码：简单的量化”
import numpy as np


def quantize_int8(tensor, symmetric=True):
    """将浮点张量量化为 int8"""
    if symmetric:
        # 对称量化: [-127, 127]
        scale = np.max(np.abs(tensor)) / 127
        quantized = np.round(tensor / scale).astype(np.int8)
    else:
        # 非对称量化: [0, 255]
        min_val, max_val = tensor.min(), tensor.max()
        scale = (max_val - min_val) / 255
        zero_point = -np.round(min_val / scale)
        quantized = np.round(tensor / scale + zero_point).astype(np.uint8)


    return quantized, scale


def dequantize_int8(quantized, scale):
    """反量化"""
    return quantized.astype(np.float32) * scale


# 示例
weights = np.random.randn(100) * 10  # 模拟权重
quantized, scale = quantize_int8(weights)
dequantized = dequantize_int8(quantized, scale)


print(f"原始权重范围: [{weights.min():.2f}, {weights.max():.2f}]")
print(f"量化后: int8, scale={scale:.4f}")
print(f"反量化误差 (MAE): {np.mean(np.abs(weights - dequantized)):.4f}")
print(f"内存占用: {weights.nbytes} → {quantized.nbytes} bytes")

研究版：精度与效率的边界

Section titled “研究版：精度与效率的边界”

研究上，量化的极限在哪里？1-bit 或 ternary 权重是否仍能保留语言能力？如何通过激活分布分析找到最优的裁剪阈值和缩放因子？

另一个前沿是推测解码（speculative decoding）：用小模型快速生成候选序列，大模型并行验证并修正，从而在不损失质量的前提下加速 2-3 倍。这本质上是计算与内存之间的重新平衡。

🔬 开放研究问题

该领域的关键问题与研究方向：

  1. KV cache 的内存占用能否进一步压缩而不显著损失精度？稀疏化、量化、蒸馏哪种方案更优？
  2. 混合精度推理的最优策略是什么？哪些层、哪些参数更适合低精度？
  3. 量化感知训练是否必要？后训练量化能否达到相同效果？

相关阅读

高效注意力：突破序列长度平方瓶颈

解释 FlashAttention、稀疏注意力与长上下文推理优化。

直觉 工程 研究

Attention：让每个位置选择上下文

解释注意力权重、多头注意力与自注意力。

直觉 工程 研究

预训练与 Scaling Law：模型怎么学

解释预训练目标、数据工程、Scaling Law 与计算最优训练。

直觉 工程 研究

长上下文：让模型读得更远

解释上下文窗口扩展、位置编码外推与长文本评测。

直觉 工程 研究

本文引用论文