KV Cache：大模型推理的内存瓶颈与优化全景

title: "KV Cache：大模型推理的内存瓶颈与优化全景" date: "2026-02-18" tags: ["LLM", "Inference", "KV Cache", "Performance"] description: "从 KV Cache 的基本原理出发，深入解析 MQA、GQA、PagedAttention、量化压缩等主流优化方案，理解大模型推理的核心内存瓶颈。"#

如果你部署过大模型推理服务，一定对这个现象不陌生：GPU 显存还剩很多，但吞吐量就是上不去。罪魁祸首往往不是计算，而是 KV Cache —— Transformer 解码阶段的内存大户。

今天我们从原理到工程，把 KV Cache 优化这件事讲透。

为什么需要 KV Cache#

Transformer 的 self-attention 机制要求每个 token 都和序列中所有之前的 token 做注意力计算。在自回归生成（autoregressive generation）时，模型每一步只生成一个新 token，但需要重新访问所有历史 token 的 Key 和 Value 向量。

如果不做缓存，每生成一个 token，就要对整个序列重新跑一遍 attention，复杂度是 O(n²)。KV Cache 的做法很直接：把已经算过的 K、V 向量存起来，下一步直接拼接新 token 的 K、V 就好。

# 伪代码：带 KV Cache 的 attention 前向
def attention_with_cache(q, k, v, kv_cache):
    # q: [batch, 1, d_model]  -- 只有当前 token
    # kv_cache: (cached_k, cached_v)
    
    cached_k, cached_v = kv_cache
    
    # 拼接新的 K、V
    k = torch.cat([cached_k, k], dim=1)  # [batch, seq_len+1, d_model]
    v = torch.cat([cached_v, v], dim=1)
    
    # 标准 attention
    scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(d_k)
    attn = torch.softmax(scores, dim=-1)
    output = torch.matmul(attn, v)
    
    # 更新 cache
    new_cache = (k, v)
    return output, new_cache

这样每步只需要做 1 × n 的 attention 计算，而不是 n × n。

KV Cache 到底占多少内存#

来算一笔账。对于一个典型的 LLaMA-70B 模型：

每个 token 的 KV Cache 大小：

单 token KV = 2 × L × kv_heads × d_head × dtype_bytes
           = 2 × 80 × 8 × 128 × 2 (FP16)
           = 327,680 bytes ≈ 320 KB

一个 4K 长度的序列：320 KB × 4096 ≈ 1.28 GB。

一个 batch 16 个请求？20.5 GB —— 光 KV Cache 就吃掉了一张 A100 的四分之一显存。而如果是 128K 上下文窗口，单个请求的 KV Cache 就要 40 GB。

这就是问题的核心：KV Cache 的内存消耗随序列长度线性增长，随 batch size 线性增长，它直接决定了你能同时服务多少用户。

优化方向一：减少 KV 头数#

Multi-Query Attention (MQA)#

最直觉的优化：既然 KV Cache 太大，那就让多个 Query 头共享同一组 K、V。

MQA 由 Noam Shazeer 在 2019 年提出，核心思想是所有 Query 头共享 一组 K 和 V：

class MultiQueryAttention(nn.Module):
    def __init__(self, d_model, n_heads):
        super().__init__()
        self.n_heads = n_heads
        self.d_head = d_model // n_heads
        
        # Q: 多头
        self.w_q = nn.Linear(d_model, d_model)
        # K, V: 单头（所有 query 头共享）
        self.w_k = nn.Linear(d_model, self.d_head)
        self.w_v = nn.

如果原来有 64 个 KV 头，MQA 只保留 1 个，KV Cache 直接缩小 64 倍。代价是什么？模型质量会有一点下降，尤其在需要精细区分不同注意力模式的任务上。

Grouped-Query Attention (GQA)#

GQA 是 MQA 和 MHA（Multi-Head Attention）之间的折中，由 Google 在 2023 年提出。它把 Query 头分成若干组，每组共享一组 K、V：

MHA:  Q头数 = K头数 = V头数 = 64
MQA:  Q头数 = 64, K头数 = V头数 = 1
GQA-8: Q头数 = 64, K头数 = V头数 = 8  (每 8 个 Q 头共享一组 KV)

LLaMA 2 70B 使用 GQA-8，KV Cache 缩小了 8 倍，模型质量几乎没有损失。这已经成为当前大模型的标配方案。

工程直觉：GQA 的组数是一个 trade-off 旋钮。组数越少，内存省得越多，但模型表达能力越受限。实践中 4-8 组是甜点区。

优化方向二：内存管理#

PagedAttention (vLLM)#

传统 KV Cache 的一个大问题是 内存碎片化。每个请求的序列长度不同，但你必须预分配一个最大长度的连续内存块。短序列浪费空间，长序列可能分配失败。

vLLM 提出的 PagedAttention 借鉴了操作系统的虚拟内存管理：

1. 把 KV Cache 切成固定大小的 块（block），类似内存页
2. 用一个 块表（block table）记录每个序列的 KV Cache 分布在哪些块里
3. 块可以不连续，按需分配和释放

传统方式（连续分配）：
请求A: [████████░░░░░░░░]  (8/16 used, 50% 浪费)
请求B: [██████████████░░]  (14/16 used)
请求C: [██░░░░░░░░░░░░░░]  (2/16 used, 87.5% 浪费)
 
PagedAttention（分页）：
块池: [A][A][B][B][B][C][A][B][空][空]
请求A 块表: [0, 1, 6]
请求B 块表: [2, 3, 4, 7]
请求C 块表: [5]
→ 内存利用率接近 100%

实测 vLLM 通过 PagedAttention 可以把吞吐量提升 2-4 倍，核心原因就是内存利用率从 50-60% 提升到了接近 100%。

Prefix Caching#

很多场景下，多个请求共享相同的前缀（system prompt、few-shot examples）。Prefix Caching 让这些公共前缀的 KV Cache 只算一次，然后在请求间共享：

请求1: [System Prompt] + [用户问题A]
请求2: [System Prompt] + [用户问题B]
请求3: [System Prompt] + [用户问题C]
 
→ System Prompt 的 KV Cache 只需要存一份

vLLM 的 Automatic Prefix Caching 和 SGLang 的 RadixAttention 都实现了这个优化。对于有长 system prompt 的应用（比如 agent 场景），这个优化可以节省大量内存和首 token 延迟。

优化方向三：压缩 KV Cache#

量化（Quantization）#

最直接的压缩：用更低的精度存储 KV Cache。

FP16 KV Cache → INT8 KV Cache: 内存减半
FP16 KV Cache → INT4 KV Cache: 内存减 75%

KIVI（2024）提出的方案：Key 用 2-bit 量化，Value 用 2-bit 量化，配合少量 FP16 的 residual，模型质量几乎无损。

实践要点：

• Key 的量化比 Value 更敏感（因为 Key 参与 softmax 前的点积，误差会被放大）
• Per-channel 量化比 per-tensor 效果好
• 最近的几个 token 保持高精度（sliding window），历史 token 激进压缩

Token 级别的压缩#

另一个思路：不是所有历史 token 的 KV 都同等重要，能不能只保留重要的？

H2O (Heavy-Hitter Oracle) 观察到 attention 分数有明显的幂律分布——少数 token 占了大部分注意力权重。H2O 保留这些 "heavy hitter" token 的 KV Cache，丢弃其余的：

def h2o_evict(kv_cache, attention_scores, budget):
    # 统计每个 token 被关注的累计分数
    importance = attention_scores.sum(dim=-2)  # 跨 query 位置求和
    
    # 保留 top-k 重要的 + 最近的 window
    recent = set(range(len(kv_cache) - window_size, len(kv_cache)))
    topk = set(importance.topk(budget - window_size).indices.tolist())
    
    keep = recent | topk
    return kv_cache[sorted

StreamingLLM 则发现一个有趣的现象：attention 分数中有大量集中在 最开头几个 token 的 "attention sink"。只要保留开头几个 token + 最近的滑动窗口，模型就能稳定地无限长度生成：

保留策略: [sink tokens (4个)] + ... 丢弃 ... + [recent window (最近 1024 个)]

优化方向四：计算优化#

Flash Attention 与 KV Cache#

Flash Attention 本身是优化 attention 计算的 IO 效率，但它对 KV Cache 的读取模式有直接影响。Flash Attention 通过分块（tiling）计算，把 KV Cache 按块从 HBM 搬到 SRAM，减少了显存带宽的瓶颈：

传统 Attention:
  读取完整 K, V → 计算 → 写回
  显存带宽成为瓶颈
 
Flash Attention:
  分块读取 K, V → 在 SRAM 中计算 → 增量更新输出
  减少 HBM 访问次数

在 decode 阶段（单 token 生成），Flash Decoding 进一步优化：把 KV Cache 沿序列维度切分，多个 thread block 并行处理不同的 KV 块，最后 reduce 结果。这对长上下文场景特别有效。

Speculative Decoding 与 KV Cache#

投机解码用小模型快速生成若干候选 token，大模型一次性验证。这改变了 KV Cache 的访问模式：

传统: 生成 1 token → 更新 cache → 生成 1 token → ...
投机: 小模型生成 [t1, t2, t3, t4] → 大模型并行验证 → 接受 [t1, t2, t3] → 回滚 t4 的 cache

需要注意 KV Cache 的回滚机制——被拒绝的 token 对应的 KV 要从 cache 中移除。

工程实践：选型指南#

部署建议：

1. 先用 vLLM/SGLang：它们已经集成了 PagedAttention、GQA、Flash Attention 等基础优化，开箱即用
2. 监控 KV Cache 利用率：vLLM 暴露了 gpu_cache_usage_perc 指标，这是你最重要的容量指标
3. 根据场景叠加：不需要一次全上，先跑基准，找到瓶颈，对症下药

# vLLM 启动示例，启用 KV Cache 量化
python -m vllm.entrypoints.openai.api_server \\
    --model meta-llama/Llama-2-70b-chat-hf \\
    --tensor-parallel-size 4 \\
    --kv-cache-dtype fp8 \\
    --enable-prefix-caching \\
    --max-model-len 32768

总结#

KV Cache 优化是大模型推理工程的核心战场。理解它，需要从三个层面思考：

1. 算法层面：MQA/GQA 从根本上减少 KV 的数量
2. 系统层面：PagedAttention、Prefix Caching 提升内存利用率
3. 数值层面：量化、token 剪枝压缩每个 KV 的大小

这些优化不是互斥的，而是可以叠加的。现代推理框架（vLLM、SGLang、TensorRT-LLM）已经集成了大部分方案，但理解底层原理能帮你做出更好的部署决策。

下次当你的推理服务 OOM 或者吞吐上不去时，先看看 KV Cache —— 大概率，答案就在那里。