KV Cache 与推理优化¶

KV Cache 是 LLM 自回归推理的核心加速机制——用显存换计算，避免重复运算。本文档聚焦它的存储特征和各种优化手段的设计权衡。

相关文档：Transformer 架构 | Mamba 与状态空间模型

1. KV Cache：用显存换速度¶

自回归生成每步只产出一个 token，但注意力计算需要当前 token 的 Query 与所有历史 token 的 Key/Value 做点积（注意力公式详见 Transformer 架构 § 自注意力机制）。

如果每步都从头算，生成第 100 个 token 就要重算前 99 个的 K/V——重复劳动随序列增长越来越严重。KV Cache 把每层算过的 K/V 向量存下来，下一步直接查表复用¹。

为什么只缓存 K/V，不缓存 Q？ 因为 Q 是"提问方"，K/V 是"被查询方"。每一步只需要当前新 token 的 Q 去匹配所有历史 token 的 K，再从对应的 V 里取信息——之前 token 的 Q 算完就没用了（它们的输出在之前的步骤里已经算好了）。而 K/V 每一步都要被新 token 查询，所以值得缓存：

步骤 4（生成"苹果"）:
  需要:   q₄                   ← 当前 token 的 Query（只用一次）
  需要:   K = [k₁,k₂,k₃,k₄]  ← 所有 token 的 Key（每步都要被匹配）
  需要:   V = [v₁,v₂,v₃,v₄]  ← 所有 token 的 Value（匹配后提供信息）
  不需要: q₁,q₂,q₃            ← 已经在之前的步骤里用完了

完整流程：

步骤 t:   缓存 K = [k₁, ..., k_{t-1}]   V = [v₁, ..., v_{t-1}]
          只计算新 token 的 q_t, k_t, v_t
          K ← append(K, k_t)     V ← append(V, v_t)
          output_t = Attention(q_t, K, V)

用「我 喜欢 吃 苹果」来对比每步的计算量：

无 KV Cache（每步重算所有 K/V）：
  步骤 1: 算 [我] 的 K/V                           → 1 个 token
  步骤 2: 重算 [我, 喜欢] 的 K/V                    → 2 个 token
  步骤 3: 重算 [我, 喜欢, 吃] 的 K/V                → 3 个 token
  步骤 4: 重算 [我, 喜欢, 吃, 苹果] 的 K/V          → 4 个 token
  总计: 1+2+3+4 = 10 次                              O(n²)

有 KV Cache（只算新 token，旧的读缓存）：
  步骤 1: 算 [我] → 存入缓存                        → 1 个 token
  步骤 2: 算 [喜欢] → 追加缓存                      → 1 个 token
  步骤 3: 算 [吃] → 追加缓存                        → 1 个 token
  步骤 4: 算 [苹果] → 追加缓存                      → 1 个 token
  总计: 1+1+1+1 = 4 次                               O(n)

核心权衡：计算换显存。 避免了 O(n²) 的重复计算，但要在 GPU 显存里维护一块随序列长度线性增长的缓存。序列越长，这块缓存越大——这就是后续所有优化手段的起因。

2. 存储开销：为什么 KV Cache 是瓶颈¶

KV Cache 必须存在 GPU 显存（HBM） 里——每一步 Decode 都要读取它，放在 CPU 内存或 SSD 上延迟太高。

显存占用和几个因素成正比：

以 Llama-3-70B（80 层, GQA-8, 每头 128 维, FP16）为例：

单请求   4K 上下文:  ≈ 1 GB    → 一张 H100 (80GB) 轻松跑几十个并发
单请求 128K 上下文:  ≈ 34 GB   → 一张 H100 几乎只能服务一个请求

KV Cache 直接决定了能同时服务多少用户、支持多长的上下文。 后面的所有技术，本质上都在想办法让这块缓存变小，或者让它被更高效地使用。

3. 推理的两个阶段与 KV 头数的影响¶

Prefill 与 Decode¶

LLM 推理实际上分两步走，硬件瓶颈完全不同：

graph LR
    Prompt["用户输入\n(prompt)"] --> Prefill

    subgraph Prefill ["① Prefill"]
        P["并行处理全部 prompt\n生成完整 KV Cache"]
    end

    Prefill --> Decode

    subgraph Decode ["② Decode"]
        D1["token₁"] --> D2["token₂"] --> D3["token₃ ..."]
    end

    Decode --> Output["完整回复"]

两个阶段需要的硬件资源正好相反——Prefill 吃算力，Decode 吃带宽。很多推理系统的优化（如 vLLM 的 Chunked Prefill、Sarathi 的 Prefill-Decode 分离）本质上都是在调和这两种需求。

GQA/MQA：从架构层面缩小 Cache¶

Decode 的带宽瓶颈取决于 KV Cache 有多大，而 KV Cache 大小直接由 KV 头数决定。标准 MHA 中每个 Q 头独占一组 KV，GQA/MQA 通过让多个 Q 头共用 KV 来减少头数，从而按比例缩小缓存（架构原理详见 Transformer 架构 § GQA/MQA）。

以 Llama-3-70B（64 个 Q 头, 128K 上下文, FP16）为例：

标准 MHA:  64 组 KV → 需缓存 ~270 GB  ← 不现实
GQA-8:      8 组 KV → 需缓存  ~34 GB  ← 勉强一张 H100
MQA:        1 组 KV → 需缓存   ~4 GB  ← 但质量损失大

为什么 GQA 成为主流？ Llama 2 的实验表明 GQA-8 在质量上几乎无损于完整 MHA，推理吞吐提升 ~1.5x²。这是一个 Pareto 最优点——再减少 KV 头，质量下降加速，而显存收益递减。

4. 缩小 KV Cache 的其他手段¶

除了第 3 节的 KV 头共享，还有多种正交的手段：

H₂O 的洞察⁴：注意力分数的分布高度不均匀——少数 token（如标点、语法词）始终获得高注意力。保留这些"Heavy Hitter"+ 最近的局部窗口，就能在很小的 cache 预算下维持大部分生成质量。

TurboQuant 的双重压缩策略⁷⁸⁹：传统量化方案将数据从 32bit 压到 4bit，理论上省 8 倍空间，但压缩过程本身需要存储量化常量、缩放因子等辅助参数（每个数据块额外占 1-2bit），实际压缩率被这些"附加成本"吃掉一大块。TurboQuant 用两套算法解决这个问题：

1. PolarQuant（极坐标量化）：将向量从直角坐标系转到极坐标系——"往东走 3 个路口再往北走 4 个路口"变成"朝 37 度方向走 5 个路口"。转换后角度的统计分布变得高度集中、可预测，不再需要为每个数据块单独存储归一化参数，附加成本归零。

2. QJL（Johnson-Lindenstrauss 随机投影）：处理 PolarQuant 的压缩残差。用随机投影只花 1 个 bit（记一个正负号），把残差误差的系统性偏差给抹平——类似秤上的"去皮"按钮，保证误差不会朝一个方向累积。

两步组合可在无需重新训练或微调的情况下实现 3bit 量化，在 H100 上注意力计算比未量化快 8 倍。经济价值显著：假设一张 H100 有 80GB 显存，模型权重占 40GB，剩余 40GB 给 KV Cache。没有 TurboQuant 时每个用户 KV Cache 占 4GB 只能同时服务 10 个用户，压缩后每个用户只占 670MB 左右，同一张卡可服务近 60 个并发，推理费用降到原来的 1/6。

工程改动：需要重写 attention kernel。TurboQuant 不是简单的存储格式变化，而是需要在压缩域直接计算注意力（compressed-domain attention）。如果采用"先解量化再计算"的路径（3bit → FP16 → 标准注意力），解量化开销会抵消压缩带来的带宽优势。因此必须重写 attention kernel 的核心计算逻辑，包括点积计算、位操作读取 packed 3bit 数据、内存访问模式等。论文用 JAX 框架测试，生产环境常用 vLLM/TensorRT-LLM/SGLang，需要为每个框架分别适配——这是 TurboQuant 落地的主要工程难点之一⁷。

5. Attention 计算优化：FlashAttention¶

标准 Attention 的实际瓶颈不在算力，而在显存读写。计算 \(QK^T\) 产生 \(n \times n\) 的注意力矩阵，写入 HBM 再读回做 softmax——这些 IO 操作才是时间杀手¹⁰。

FlashAttention 的核心洞察：不把完整注意力矩阵写入 HBM。两个关键技巧：

1. 分块计算（Tiling）：把 Q/K/V 切成小块塞进 GPU 片上 SRAM（比 HBM 快 10-20 倍），在 SRAM 内完成一个块的 \(QK^T → \text{softmax} → \times V\) 后直接输出，不存中间矩阵
2. 在线 Softmax：逐块累积归一化因子（运行最大值和指数和），数学上等价于全局 softmax——计算精确，不是近似

FlashAttention 主要加速 Prefill（多 Query 并行）。Decode 阶段每步只有 1 个 Query，并行度不足——FlashDecoding¹³ 沿序列长度维度切分 KV Cache，每块独立计算局部 Attention 后合并，长序列（>8K tokens）加速 3-8x。

6. 内存管理：PagedAttention¶

传统推理引擎为每个请求预分配连续显存。序列长度不可预知——预分配大了浪费，小了要扩容搬移。多请求交替完成后碎片严重，KV Cache 显存利用率仅 20-40%。

PagedAttention¹⁴ 借鉴操作系统虚拟内存：固定大小的逻辑块（典型 16 tokens）通过 Block Table 映射到物理块，无需连续分配。

效果：KV Cache 利用率 ~30% → ~85%，同等显存下吞吐 2.5-3x¹⁴。已成为 vLLM、SGLang 等主流推理引擎的标准内存管理方案。

7. 分层存储：用廉价介质扩展容量¶

当 KV Cache 超出单卡 HBM 时，可利用低层存储介质，以延迟换容量：

Mooncake¹⁵ 的核心设计哲学是以 KV Cache 为中心：热数据（当前 batch）留在 HBM，温数据（Prefix Cache）放 CPU DRAM，冷数据下沉到 SSD / 远端——用廉价存储换昂贵 GPU 算力。配合 GPUDirect RDMA 零拷贝传输（拓扑感知路径选择 + 多 NIC 带宽聚合），128K tokens 的 KV Cache（~2GB）在 400Gbps InfiniBand 上仅需 ~40ms，远快于重新 Prefill。架构详见 Mooncake 开源库。

8. 长上下文时代的挑战¶

随着上下文窗口从 4K → 128K → 1M+ 扩展，KV Cache 问题被急剧放大：

应对方向：

• 架构层：GQA 减少 KV 头数（§3）；Mamba/RWKV 等 SSM 架构完全绕过 KV Cache → Mamba 与状态空间模型
• 算法层：FlashAttention 降低 IO 开销（§5）；稀疏注意力、Token 驱逐、量化（§4）
• 系统层：PagedAttention 消除碎片（§6）；分层存储扩展容量（§7）；Prefix Caching
• 硬件层：更大显存（H200: 141GB, B200: 192GB）、更高带宽

参考资料¶