为什么KV Cache如此重要？

在自回归生成中，每生成一个新token都需要计算它与之前所有token的注意力。如果不缓存，生成第N个token需要重新计算前N-1个token的K和V，复杂度为O(N²)。KV Cache将已计算的K/V存储起来，使每步生成只需O(N)的计算。

但缓存本身带来了巨大的显存压力。以Llama-3-70B为例，FP16精度下单个token的KV Cache占用约320KB。在32K上下文、batch_size=32的场景下，KV Cache总量达到约320GB——远超模型本身的140GB。

KV Cache的显存构成

单个token的KV Cache大小为：

cache_per_token = 2 (K和V) × n_layers × n_kv_heads × head_dim × 2 (FP16字节)

以Llama-3-70B为例：

  • n_layers = 80
  • n_kv_heads = 8 (GQA)
  • head_dim = 128
  • FP16 = 2字节
cache_per_token = 2 × 80 × 8 × 128 × 2 = 327,680 bytes ≈ 320KB

理解这个公式是所有优化方案的出发点。

PagedAttention：分页式KV管理

vLLM团队提出的PagedAttention借鉴了操作系统的虚拟内存管理。核心思想：

块式管理

将KV Cache分为固定大小的块（Block），每个块存储若干token的KV。块不需要在物理上连续，通过块表（Block Table）映射逻辑位置到物理位置。

class PagedKVCache:
    def __init__(self, n_blocks, block_size, n_layers, n_kv_heads, head_dim):
        self.block_size = block_size  # 每块的token数
        self.n_blocks = n_blocks
        # 预分配所有块的存储
        self.k_blocks = torch.zeros(n_blocks, n_layers, block_size, n_kv_heads, head_dim)
        self.v_blocks = torch.zeros(n_blocks, n_layers, block_size, n_kv_heads, head_dim)
        self.free_blocks = list(range(n_blocks))
        self.block_tables = {}  # seq_id -> [block_indices]
    
    def allocate(self, seq_id, n_tokens):
        n_blocks_needed = (n_tokens + self.block_size - 1) // self.block_size
        blocks = []
        for _ in range(n_blocks_needed):
            if not self.free_blocks:
                raise OOMError("No free KV cache blocks")
            blocks.append(self.free_blocks.pop())
        self.block_tables[seq_id] = blocks
    
    def append(self, seq_id, token_idx, k, v):
        block_idx = token_idx // self.block_size
        offset = token_idx % self.block_size
        physical_block = self.block_tables[seq_id][block_idx]
        self.k_blocks[physical_block, :, offset] = k
        self.v_blocks[physical_block, :, offset] = v

优势

  • 消除碎片：无需为最大可能长度预分配连续空间
  • 灵活共享：不同序列可以共享相同的块（如系统提示词部分）
  • 按需分配：序列增长时动态分配新块

Prefix Sharing

多个请求共享相同的前缀（如系统提示词）时，PagedAttention可以让这些请求共享同一组物理块。这在对话AI场景中可节省30-50%的显存。

KV Cache量化

将KV Cache从FP16量化到更低精度是另一个重要优化方向。

INT8量化

将KV Cache从FP16量化到INT8可以节省一半显存。由于K和V的分布不同，通常需要分别处理：

def quantize_kv_cache(k_cache, v_cache, scale_k, scale_v):
    """对称量化KV Cache到INT8"""
    k_int8 = (k_cache / scale_k).round().clamp(-128, 127).to(torch.int8)
    v_int8 = (v_cache / scale_v).round().clamp(-128, 127).to(torch.int8)
    return k_int8, v_int8

def dequantize_kv_cache(k_int8, v_int8, scale_k, scale_v):
    return k_int8.float() * scale_k, v_int8.float() * scale_v

关键在于缩放因子的选择。实践中，逐通道量化（per-channel）比逐张量量化（per-tensor）精度损失更小。

INT4量化

更激进的量化方案。INT4 KV Cache将显存节省4倍，但精度损失更为显著。研究表明，K Cache的INT4量化会导致较大的性能退化，而V Cache对INT4更为鲁棒。因此，许多系统采用混合精度——K用INT8、V用INT4。

FP8量化

NVIDIA H100支持的原生FP8格式提供了精度和效率的更好平衡。FP8有两种变体：E4M3和E5M2，前者精度更高，后者动态范围更大。KV Cache量化通常使用E5M2（因为KV值范围较大）。

KV Cache的压缩与驱逐

滑动窗口

只保留最近W个token的KV Cache，超出的直接丢弃。Mistral-7B使用W=4096的滑动窗口。这种方法虽然简单，但会丢失长距离信息。

结构化剪枝

基于注意力分数分析，识别出"不重要"的token并驱逐其KV。H2O（Heavy-Hitter Oracle）算法保持了约5%的token就能维持大部分性能：

def h2o_eviction(attention_scores, kv_cache, keep_ratio=0.05):
    """基于注意力分数的KV Cache驱逐"""
    n_tokens = attention_scores.shape[-1]
    n_keep = int(n_tokens * keep_ratio)
    
    # 计算每个token的累计注意力分数
    cumulative_scores = attention_scores.sum(dim=-2)  # [n_tokens]
    
    # 保留分数最高的token
    keep_indices = torch.topk(cumulative_scores, n_keep).indices
    keep_indices, _ = keep_indices.sort()
    
    return kv_cache[keep_indices], keep_indices

语义聚类压缩

将语义相似的token的KV进行聚类合并，用聚类中心代替原始KV。这种方法比简单的截断保留了更多信息。

跨层KV共享

观察到相邻层的注意力模式高度相似，一些研究提出跨层共享KV Cache：

  • YOCO（You Only Cache Once）：只在模型底部计算一次KV，所有层共享
  • CLA（Cross-Layer Attention）：将层分为若干组，组内共享KV

这种方法可以将KV Cache减少50-80%，但需要特殊的预训练或微调来适应。

2026年的KV Cache前沿

多模态KV分离

在视觉-语言模型中，图像token的KV Cache占用远大于文本。最新的做法是对图像token和文本token使用不同的压缩策略——图像token可以用更激进的压缩（因为视觉信息有冗余）。

推测性KV预取

在推测解码中，draft模型预测的多个候选token的KV可以预先计算和缓存。如果验证通过则直接使用，否则丢弃。这减少了大模型的KV计算次数。

持久化KV存储

对于高频使用的系统提示和文档，将其KV Cache持久化到磁盘或SSD。新请求到来时直接加载，避免重新计算。一些推理引擎已经支持这种"KV Cache仓库"模式。

工程实践建议

  1. 基准测量：先测量你的模型在目标场景下的KV Cache大小，再决定优化策略
  2. GQA/MQA：使用Grouped-Query Attention或Multi-Query Attention可以从架构层面减少KV Cache
  3. 量化优先：INT8量化对性能几乎无损，应作为第一步优化
  4. 分页管理：PagedAttention或类似方案对高并发场景是刚需
  5. 监控碎片率：跟踪KV Cache的碎片率和利用率，动态调整块大小

结语

KV Cache是LLM推理系统中最为关键的工程组件之一。从PagedAttention的分页管理到量化压缩、从注意力驱动的驱逐到跨层共享，每一种技术都在为"在有限显存中服务更长上下文"这一核心目标服务。随着上下文长度从32K向1M迈进，KV Cache优化将变得更加重要。

加入讨论

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

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