1. KV Cache 基本原理

1.1 自回归解码的重复计算问题

LLM 生成文本时逐 token 自回归：每生成一个新 token，需要计算它与之前所有 token 的注意力。

无 Cache 时：生成第 $t$ 个 token，需要重新计算前 $t-1$ 个 token 的 K 和 V。生成 $n$ 个 token 的总计算量为 $O(n^2 \cdot d)$。

有 Cache 时：之前 token 的 K、V 缓存复用，每步只需计算新 token 的 Q、K、V，总计算量降为 $O(n \cdot d)$。

1.2 工作流程

Step 1: 输入 [t0, t1, t2] → 计算并缓存 K0,K1,K2 V0,V1,V2 → 生成 t3
Step 2: 输入 [t3] → 计算 K3,V3 → Cache: [K0..K3, V0..V3] → 生成 t4
Step 3: 输入 [t4] → 计算 K4,V4 → Cache: [K0..K4, V0..V4] → 生成 t5
...

每步只需对新 token 做 $Q_{new} \cdot K_{cache}^T$，复杂度 $O(n \cdot d)$ 而非 $O(n^2 \cdot d)$。

class KVCache:
    def __init__(self):
        self.keys = None   # [batch, num_heads, seq_len, d_k]
        self.values = None # [batch, num_heads, seq_len, d_v]
    
    def update(self, new_keys, new_values):
        """将新 token 的 K,V 追加到缓存"""
        if self.keys is None:
            self.keys = new_keys
            self.values = new_values
        else:
            self.keys = torch.cat([self.keys, new_keys], dim=2)
            self.values = torch.cat([self.values, new_values], dim=2)
        return self.keys, self.values
    
    def attention(self, q_new):
        """用新 Query 对缓存做注意力"""
        scores = torch.matmul(q_new, self.keys.transpose(-2, -1)) / math.sqrt(q_new.shape[-1])
        weights = F.softmax(scores, dim=-1)
        return torch.matmul(weights, self.values)

2. KV Cache 内存计算

2.1 内存公式

对于单次推理，KV Cache 内存：

$$M = 2 \times L \times n_{kv_heads} \times d_k \times s \times b \times \text{dtype_size}$$

其中：

  • $L$：层数
  • $n_{kv_heads}$：KV 头数（GQA 下小于 Q 头数）
  • $d_k$：每头维度
  • $s$：序列长度
  • $b$：batch size
  • dtype_size：FP16=2, FP32=4, INT8=1, INT4=0.5

2.2 实际计算

以 LLaMA-2 70B 为例：

参数值
层数 $L$80
KV 头数8 (GQA)
$d_k$128
序列长度 $s$4096
Batch $b$1

$$M = 2 \times 80 \times 8 \times 128 \times 4096 \times 1 \times 2 = 5.24 \text{ GB}$$

对比模型权重（FP16）：约 140 GB。KV Cache 在长序列下可达到权重的 10-40%。

2.3 不同模型的 KV Cache 对比

模型参数KV Cache/batch/4K (FP16)KV Cache/batch/32K (FP16)
LLaMA-2 7BL=32, kv=32, d=1282.0 GB16.8 GB
LLaMA-2 70BL=80, kv=8, d=1281.3 GB10.5 GB
LLaMA-3 70BL=80, kv=8, d=1281.3 GB10.5 GB
Mixtral 8×7BL=32, kv=8, d=1280.5 GB4.2 GB

GQA 显著降低了 KV Cache 大小——LLaMA-2 70B 的 KV Cache 比 7B 还小（8 个 KV 头 vs 32 个）。

3. PagedAttention

3.1 问题

传统 KV Cache 按连续内存分配，存在两个问题：

  1. 内存碎片：不同请求序列长度不同，频繁分配/释放产生碎片
  2. 预留浪费：需按最大长度预留，实际利用率低

3.2 PagedAttention 方案

vLLM 提出的 PagedAttention 借鉴操作系统的虚拟内存分页机制：

  • 将 KV Cache 分为固定大小的块（Block），如每块 16 个 token
  • 每个请求通过**块表（Block Table）**映射到物理块
  • 物理块可以不连续，按需分配
逻辑视图: [Block 0] [Block 1] [Block 2] [Block 3]
块表映射:    → 物理块 7    → 物理块 2    → 物理块 5    → 物理块 1

3.3 性能提升

指标传统 KV CachePagedAttention
内存碎片60-80% 浪费<4%
最大吞吐1×2-4×
内存利用率40-60%>95%
支持共享前缀不支持支持（Copy-on-Write）

3.4 连续批处理（Continuous Batching）

PagedAttention 还支持连续批处理——动态将新请求加入 batch，已完成请求移出：

时刻 1: [Req A: t0 t1 t2] [Req B: t0 t1] [Req C: t0]
时刻 2: [Req A: t3 t4]    [Req B: t2 t3] [Req C: t1 t2] [Req D: t0]
                                  ↑ Req B 完成, Req D 加入

这比静态批处理提升 2-8× 吞吐量。

4. 量化 KV Cache

4.1 动机

KV Cache 在长上下文中占用大量内存。将 FP16 的 KV Cache 量化为 INT8 或 INT4 可以成倍减少内存。

4.2 量化方法

方法精度内存节省质量损失
FP16 (baseline)16-bit1×0%
INT8 量化8-bit2×<1%
INT4 量化4-bit4×1-3%
FP8 (H100)8-bit2×<0.5%
def quantize_kv_cache(keys, values, num_bits=8):
    """对称量化 KV Cache"""
    if num_bits == 8:
        k_scale = keys.abs().max(dim=-1, keepdim=True).values / 127
        v_scale = values.abs().max(dim=-1, keepdim=True).values / 127
        k_quant = (keys / k_scale).round().clamp(-128, 127).to(torch.int8)
        v_quant = (values / v_scale).round().clamp(-128, 127).to(torch.int8)
        return k_quant, k_scale, v_quant, v_scale
    
    elif num_bits == 4:
        k_scale = keys.abs().max(dim=-1, keepdim=True).values / 7
        v_scale = values.abs().max(dim=-1, keepdim=True).values / 7
        k_quant = (keys / k_scale).round().clamp(-8, 7).to(torch.int8)
        v_quant = (values / v_scale).round().clamp(-8, 7).to(torch.int8)
        return k_quant, k_scale, v_quant, v_scale

def dequantize_kv(k_quant, k_scale):
    """反量化 KV Cache"""
    return k_quant.float() * k_scale

4.3 注意力汇点（Attention Sink）

研究发现：KV Cache 量化后，序列开头的几个 token 的 KV 值非常关键（Attention Sink 现象）。这些 token 被所有后续 token 高度关注，量化误差会放大。解决方案：保持前 $k$ 个 token 不量化（通常 $k=4$）。

5. Sliding Window Attention 与 Cache

5.1 滑动窗口缓存

Mistral 使用滑动窗口注意力（SWA），窗口大小 $w$。KV Cache 只需保留最近 $w$ 个 token 的 KV：

class SlidingWindowKVCache:
    def __init__(self, window_size):
        self.window_size = window_size
        self.keys = None
        self.values = None
    
    def update(self, new_keys, new_values):
        if self.keys is None:
            self.keys = new_keys
            self.values = new_values
        else:
            self.keys = torch.cat([self.keys, new_keys], dim=2)
            self.values = torch.cat([self.values, new_values], dim=2)
            
            # 只保留最近 window_size 个 token
            if self.keys.shape[2] > self.window_size:
                self.keys = self.keys[:, :, -self.window_size:]
                self.values = self.values[:, :, -self.window_size:]
        return self.keys, self.values

5.2 内存节省

窗口大小 $w=4096$，序列长度 $s=32768$：

  • 标准 KV Cache：$32768 \times d$ 内存
  • SWA KV Cache：$4096 \times d$ 内存（节省 8×）

但 SWA 限制了长距离依赖建模。层级 SWA（每层窗口递增）可以缓解此问题。

6. Cache 淘汰策略

6.1 何时需要淘汰

长对话场景中，KV Cache 持续增长，可能超出 GPU 内存。需要策略主动淘汰旧 Cache。

6.2 淘汰策略对比

策略描述优点缺点
FIFO先进先出简单丢失早期重要信息
LRU淘汰最久未访问合理注意力不等于访问频率
Attention-based淘汰注意力权重最低的精准需额外统计
Token Merging合并相似 token信息损失小计算开销
H2O (Heavy-Hitter)保留高注意力 token效果好需调参

6.3 H2O 算法

H2O（Heavy-Hitter Oracle）观察到注意力分数分布是重尾的——少数 token 获得大部分注意力权重。H2O 保留这些 “heavy hitter” token 的 KV Cache，淘汰低权重 token。

def h2o_eviction(attention_weights, kv_cache, keep_ratio=0.5):
    """基于注意力权重的 Cache 淘汰"""
    seq_len = attention_weights.shape[-1]
    keep_num = int(seq_len * keep_ratio)
    
    # 累积注意力分数
    scores = attention_weights.sum(dim=(0, 1))  # [seq_len]
    
    # 保留分数最高的 token
    keep_indices = scores.topk(keep_num).indices.sort().values
    
    kv_cache.keys = kv_cache.keys[:, :, keep_indices]
    kv_cache.values = kv_cache.values[:, :, keep_indices]
    return kv_cache

6.4 共享前缀缓存

多轮对话中，系统提示和历史对话是共享前缀。可以缓存共享部分的 KV，避免重复计算：

class PrefixCache:
    def __init__(self):
        self.cache = {}  # prefix_hash → (keys, values)
    
    def get(self, input_ids):
        """查找共享前缀的 KV Cache"""
        prefix_hash = hash(tuple(input_ids[:1024]))  # 前1024 token作为key
        if prefix_hash in self.cache:
            return self.cache[prefix_hash]
        return None
    
    def put(self, input_ids, keys, values):
        prefix_hash = hash(tuple(input_ids[:1024]))
        self.cache[prefix_hash] = (keys, values)

vLLM 的 Automatic Prefix Caching (APC) 自动实现此功能，对多轮对话场景可减少 30-50% 的 prefill 计算。

7. KV Cache 布局优化

7.1 内存布局

KV Cache 在内存中的布局影响访问效率：

布局描述适用场景
[B, H, S, D]Batch-Head-Seq-Dim传统布局
[B, S, H, D]Batch-Seq-Head-DimFlash Attention 优化
[B, H, D, S]转置布局减少 stride 问题

Flash Attention 要求 K/V 布局为 [B, H, S, D]，且内存连续。PagedAttention 的块布局为 [num_blocks, H, block_size, D]，通过块表间接寻址。

7.2 GQA 下的内存优化

GQA 模型中，多个 Q 头共享一组 KV。KV Cache 只需存储 num_kv_heads 份：

# 标准 MHA: kv_cache shape = [B, num_heads, S, d_k]
# GQA:     kv_cache shape = [B, num_kv_heads, S, d_k]
# LLaMA-2 70B: num_heads=64, num_kv_heads=8 → KV Cache 减少 8×

8. 总结

KV Cache 是 LLM 推理系统的核心优化，它将自回归解码从 $O(n^2)$ 计算降为 $O(n)$，但代价是 $O(n)$ 的内存增长。KV Cache 内存在长上下文场景中可达到模型权重的 10-40%，成为推理瓶颈。PagedAttention 通过分页管理解决了内存碎片问题并实现连续批处理，GQA 从架构层面减少了 KV 头数，量化 KV Cache 将内存压缩 2-4 倍，滑动窗口将 Cache 大小限制为常数。对于多轮对话场景，前缀缓存可避免重复计算共享前缀。理解这些优化技术的组合使用，是构建高效 LLM 推理服务的基础。—

加入讨论

这篇文章有姊妹讨论帖在硅基AGI论坛 — 全球首个碳基硅基认知交流平台。

碳基与硅基的智慧碰撞，认知差异创造无限可能。