高效注意力与长上下文：KV 缓存、GQA/MQA 与 FlashAttention

一句话定义

长上下文工程围绕注意力的两大瓶颈展开：推理端用 KV 缓存及 GQA/MQA 压缩其显存，训练端用 FlashAttention 改写内存访问模式，使百万词元级窗口在工程上可行。

为什么重要

上下文窗口是 LLM 的「工作内存」，其长度与成本直接决定应用形态（长文档分析、代码库理解、超长对话）。这套技术是 2023 年以来各家窗口竞赛（2K → 128K 及更长）的底层支撑，也是读懂推理成本与部署选型（详见推理部署专题站）的必备知识。

前置知识

• kp-011（注意力与复杂度）
• kp-012（位置编码与外推）

核心概念

• KV 缓存（KV Cache）：自回归生成时缓存历史词元的 K、V，把每步复杂度从重算全史降为增量计算。
• MQA（Multi-Query Attention）：所有头共享一份 K/V，缓存缩到 1/h，但质量略降。
• GQA（Grouped-Query Attention）：折中方案，每若干 Q 头共享一组 K/V，成为当代主流。
• FlashAttention：通过分块计算与在线 softmax 减少显存读写，把注意力显存从 O(T²) 降到线性、速度显著提升（近似精确而非近似计算）。
• 上下文扩展：位置插值、NTK 感知缩放等（与 kp-012 衔接），把短窗模型扩到长窗。

原理与机制

生成第 t 个词只需历史 K/V + 当前 Q：缓存历史避免 O(T²) 重算，但缓存体积 = 层数 × 2 × T × d × batch × 精度字节，长窗下成为显存第一大户——MQA/GQA 通过共享 K/V 头直接把常数缩小 4 至 8 倍（典型配置）。FlashAttention 的洞察：注意力慢在 HBM 显存往返而非浮点运算；把 Q/K/V 分块装入高速 SRAM、用递推方式计算 softmax 归一化，一次遍历完成，兼顾速度与精确性。

公式或模型

KV 缓存体积与 GQA 压缩比：

图示

无缓存: 每生成一词 → 对 1..t 全部重算 K,V,Q  （O(T²) 每步）
有缓存: 只算新词 q_t，与缓存 K,V 做注意力     （O(T) 每步）
MQA/GQA: h 组 Q 头共享 1 或 G 组 K,V → 缓存同比缩小

直观类比

KV 缓存像「会议纪要」：不必每次发言都重放全场录音（重算），翻纪要即可（缓存）；GQA 像把逐人记录合并成「小组纪要」，记录本变薄；FlashAttention 像「改变做菜顺序」——食材一样，但不再频繁在厨房与储藏室之间往返，全部在小台面上完成（显存分层优化）。

实例或案例

主流开源模型普遍采用 GQA（如 LLaMA 2 70B 起用 GQA；7B 级 MQA/GQA 混用），同等显存可服务更长窗口与更大批量；FlashAttention 已内置于主流训练与推理框架，成为不可见但无处不在的默认项。窗口竞赛从 4K 到 128K 及以上，均是本节技术叠加 kp-012 位置方案的成果。

常见误区

• 「KV 缓存是可选优化」：无缓存的朴素生成在长文本下慢到不可用，它是推理引擎的必需品。
• 「GQA/MQA 会让质量大幅下降」：GQA 与全多头质量接近，是显存换质量的常见划算交易。
• 「FlashAttention 是近似注意力」：它计算的是精确注意力，优化的是访存路径，不改变数学结果。

与其他知识点的关系

• kp-023 的「迷失中间」说明长窗口 ≠ 长记忆能力；kp-012 的 RoPE/插值决定窗口能「伸多长」，本节决定「伸长后跑不跑得动」；推理部署专题站在此基础上讨论吞吐与成本。

自测题

1. KV 缓存把每步生成复杂度从多少降到多少？
   答案要点：从对全部历史重算（每步 O(T²) 累计）降为增量 O(T) 每步；代价是 O(T) 显存。
2. GQA 与 MQA 的差别及取舍？
   答案要点：MQA 全部 Q 头共享 1 组 K/V（压缩最大、质量略降）；GQA 分组共享（折中），现为主流。
3. FlashAttention 快在哪一层？
   答案要点：不减少浮点计算量（近似精确），而是分块驻留高速缓存、减少 HBM 读写与中间矩阵落盘。

延伸阅读

• Fast Transformer Decoding: One Write-Head is All You Need（Shazeer，2019）：MQA
• GQA: Training Generalized Multi-Query Transformer（Ainslie 等，2023）
• FlashAttention（Dao 等，2022）