大规模语言推理：KV Cache 机制与注意力优化深度解析

引言

大语言模型（LLM）的推理服务面临一个核心矛盾：自回归生成需要反复计算注意力，而每一新生成的 token 都要和所有历史 token 做交互。如果每次都重新计算全部 query-key-value，推理开销会随序列长度呈平方级增长。KV Cache 及其一系列优化技术正是为了解决这个瓶颈而生。

1. KV Cache 基本原理

在 Transformer 的自注意力层中，每个 token 会生成 Query、Key、Value 三组向量。当模型自回归生成第 t 个 token 时，它需要与前面 t-1 个 token 计算注意力权重。这些历史 token 的 Key 和 Value 在生成过程中不会改变，因此可以将它们缓存起来，避免重复计算。

具体而言，KV Cache 为每个注意力层、每个注意力头维护一块连续内存区域：

• K Cache：存储所有已生成 token 的 Key 向量，形状为 [num_layers, num_heads, seq_len, head_dim]
• V Cache：存储所有已生成 token 的 Value 向量，形状与 K Cache 相同

每次新 token 到来时，只需计算其 Q、K、V，将 K 和 V 追加到 Cache 中，然后用新 Q 与全部 K 计算注意力。这样每步的计算复杂度从 O(n²) 降低到 O(n)。

2. 显存瓶颈与 PagedAttention

尽管 KV Cache 减少了计算量，但显存压力随着序列长度线性增长。对于 Llama-2-70B 这类大模型，单个请求的 KV Cache 可达数 GB。更糟糕的是，现有系统通常为每个请求预分配一块连续显存，导致严重的内存碎片和利用率低下——研究表明实际有效利用率仅为 20%~40%。

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

• 将 KV Cache 分割成固定大小的块（Block），通常为 16 个 token
• 每个请求维护一个块表（Block Table），将逻辑块映射到物理块
• 物理块按需分配，不要求连续，消除内存碎片
• 支持 Copy-on-Write，使 Beam Search 等场景下的序列共享效率极高

这一设计将 KV Cache 利用率从不到 40% 提升至接近 100%，在相同硬件上支持 2~4 倍的吞吐。

3. GQA：在质量与效率之间取平衡

标准的多头注意力（MHA）为每个 Query 头配备独立的 Key 和 Value 头。对于 70B 级别的模型，这意味着大量的 KV Cache 占用。Grouped-Query Attention (GQA) 将多个 Query 头共享一组 Key/Value 头：

• MHA：Q 头数 = K 头数 = V 头数（如 64:64:64）
• GQA：Q 头数 > K 头数（如 8 个 Q 头共享 1 个 KV 头，即 GQA-8）
• MQA：所有 Q 头共享单一 KV 头（极端的压缩比）

Llama-2-70B 采用 GQA-8 方案，KV Cache 缩减为 MHA 的 1/8，几乎不损失生成质量。推理时 KV 头通过广播机制与所有 Q 头完成注意力计算，硬件利用率更高。

4. FlashAttention：IO 感知的精确算法

即使有了 KV Cache，每步的注意力计算仍可能成为瓶颈。GPU 的显存层次结构（HBM → Shared Memory → Registers）中，HBM 访问延迟是寄存器访问的百倍以上。传统注意力实现需要将 Q、K、V 矩阵从 HBM 反复读写。

FlashAttention 通过 Tiling（分块）和 Recomputation（重计算）实现 IO 优化的精确注意力：

• 将 Q、K、V 切分成适合 Shared Memory 的小块
• 在块级别计算 Softmax，通过保留 "logsumexp" 在线校正，避免保存中间注意力矩阵
• 将 HBM 读写复杂度从 O(N²) 降至 O(N²/M)，其中 M 为 Shared Memory 大小
• 在不改变数学结果的前提下，实现 2~4× 加速并大幅节省显存

FlashAttention-2 进一步优化并行策略和 warp 分工，逼近 GPU 硬件峰值算力。

5. 推测解码与并行生成策略

在推理服务场景中，除了单 token 的逐字生成，研究者还开发了多种并行化方案来突破自回归瓶颈：

• 推测解码 (Speculative Decoding)：使用小型草稿模型快速生成候选 token，大模型并行验证。草稿模型命中率高时，有效吞吐可提升 2~3 倍，且输出分布与原始大模型一致。
• Medusa 多头并行：为大模型增加多个解码头，每头预测不同位置的 token，通过树状验证筛选最优序列。
• Chunked Prefill：将长 Prompt 的 Prefill 阶段分成多个 Chunk，与 Decode 阶段交错执行，保障 TTFT 和 TBT 延迟 SLO。

6. 实际系统中的端到端优化

工业级推理引擎（vLLM、TGI、TensorRT-LLM、SGLang）将上述技术组合成完整的优化栈：

• Continuous Batching：取消静态 Batch Size 限制，新请求可在任意时刻加入当前 Batch
• Quantized KV Cache：采用 FP8/INT8 量化 KV Cache，显存减半，带宽瓶颈缓解
• Prefix Caching：对 System Prompt 等共享前缀复用已计算的 KV Cache
• Tensor/Pipeline Parallelism：多机场景下将 KV Cache 按头或层分布到不同设备

这些技术协同作用，使得单卡 A100-80GB 上可以服务 70B 参数模型，多机架设下可达到每秒数千 token 的吞吐。

总结

LLM 推理优化是一场在数学精确性、计算效率、显存约束之间的精密权衡。KV Cache 是基石，PagedAttention 解决了显存利用效率，GQA 在压缩与质量间取中，FlashAttention 优化了计算 IO，推测解码则突破了自回归的本质限速。理解这些技术的组合逻辑，是构建高效推理服务的必修课。