AI 技术
#GQA#MQA#KV Cache#Attention#LLM 推理#模型架构

Grouped-Query Attention 如何压缩 KV Cache:MHA、MQA 与 GQA 的工程权衡

围绕长上下文在线推理中的 KV Cache 显存压力,解释 Grouped-Query Attention 如何让多组 Query Head 共享更少的 Key/Value Head,并分析它与 MHA、MQA、PagedAttention、FlashAttention 的边界、吞吐收益和质量权衡。

一个长上下文聊天服务很容易遇到这样的现象:模型权重能够完整装进 GPU,单条请求也能正常运行,但并发一上升,显存很快被吃满。继续增加 batch 或并发请求时,系统首先撞上的往往不是矩阵乘法算力,而是不断增长的 KV Cache。每个活跃请求都要保存历史 Token 对应的 Key 和 Value,序列越长、层数越多、KV Head 越多,缓存占用就越大。

Grouped-Query Attention(GQA)解决的正是这个结构性成本。它不改变 Transformer 自回归生成的基本流程,也不等同于一种运行时缓存算法,而是在模型 Attention 结构里减少需要独立保存的 Key/Value Head 数量。对在线推理而言,这会直接改变每个 Token 写入 KV Cache 的数据量,从而影响可容纳的并发请求数、显存带宽压力以及长上下文服务的部署边界。

从 MHA 开始看 KV Cache 为什么会随 Head 数量增长

标准 Multi-Head Attention(MHA)把隐藏表示投影成多组 Query、Key 和 Value。假设一个 Attention 层有 32 个 Query Head,典型 MHA 也会配置 32 个 Key Head 和 32 个 Value Head。每个 Head 在不同子空间里完成注意力计算,再把结果拼接回统一表示。

训练阶段处理一段完整序列时,Q、K、V 都会参与当前前向计算。进入自回归推理后,情况发生变化。模型每次只生成少量新 Token,历史 Token 对应的 Key 和 Value 不需要重复计算,因此推理引擎会把它们保存在 KV Cache 中。新 Token 的 Query 不需要长期缓存,因为它完成当前 Attention 后就不再被未来步骤直接复用;历史 Key 和 Value 则要一直保留到请求结束或上下文被淘汰。

因此,单个请求的 KV Cache 规模与多个因素相关:模型层数、历史序列长度、每个 KV Head 的维度、数据类型,以及最关键的 KV Head 数量。其他条件不变时,KV Head 从 32 减到 8,Key 和 Value 的缓存元素数量也会按相同比例下降。这里减少的不是模型权重,而是每个活跃请求随上下文持续增长的运行时状态。

对于一个在线对话服务,这种区别会被并发放大。假设一张 GPU 同时承载几十个长对话,每个请求都维护自己的历史 KV。模型权重是共享的,但 KV Cache 是请求级状态。单请求节省一部分缓存,看起来只是几百 MB 或更少;乘上层数、上下文长度和并发数之后,就可能决定系统能否继续接纳新请求。

MQA 把共享推到极致,但也压缩了表示自由度

Multi-Query Attention(MQA)采用更激进的结构:多个 Query Head 共享同一组 Key 和 Value。仍以 32 个 Query Head 为例,MHA 可以是 32 个 Query Head 对应 32 个 KV Head,而 MQA 则保留 32 个 Query Head,却只使用 1 个 KV Head。

从推理系统角度看,这非常有吸引力。历史每新增一个 Token,所有 Query Head 不再需要分别写入 32 组 Key/Value,只需保存共享的那一组。decode 阶段读取历史 KV 时,数据量也显著下降。对于受 KV Cache 容量和内存带宽限制的在线推理,MQA 能把结构性开销压到很低。

代价来自模型表达能力。Query Head 仍然可以学习不同的查询方向,但它们看到的 Key/Value 表示被强制共享。可以把它理解为多个查询视角仍然存在,但所有视角检索的是同一套索引表示。共享越彻底,KV 成本越低,同时模型在不同 Head 之间保留独立 Key/Value 表示的自由度也越少。

用户提供的 GQA 论文材料正是在这一背景下讨论折中方案:MQA 具有很好的推理效率,但质量可能相对 MHA 出现退化。这里不应把原因简单归结为“某些 Head 负责语法、某些 Head 负责实体”这样的固定角色,因为实际 Head 行为并没有如此整齐的人工分工。更准确的理解是,不同 Head 可以学习不同投影子空间,强制所有 Query Head 共享同一个 K/V 投影,会减少模型可用的表示容量。

GQA 的核心不是新增 Head,而是改变 Q Head 与 KV Head 的映射

Grouped-Query Attention 位于 MHA 和 MQA 之间。它保留较多 Query Head,但只配置较少的 KV Head,并让一组 Query Head 共享一个 Key Head 和一个 Value Head。

例如一个层有 32 个 Query Head,可以配置 8 个 KV Head。此时每 4 个 Query Head 归入一组,共享对应的 K/V。模型仍然拥有 32 个 Query 投影去形成不同查询,但历史 Token 只需要存储 8 组 Key 和 Value。与 32 个 KV Head 的 MHA 相比,单层 KV Cache 的这部分规模降为四分之一;与只有 1 个 KV Head 的 MQA 相比,GQA 又保留了更多独立的 K/V 表示。

flowchart TD
  A[新 Token 隐藏状态] --> B[生成 32 个 Query Head]
  A --> C[生成 8 个 Key Head]
  A --> D[生成 8 个 Value Head]
  C --> E[写入 8 组历史 K Cache]
  D --> F[写入 8 组历史 V Cache]
  B --> G[按组映射到对应 KV Head]
  E --> G
  F --> G
  G --> H[计算各 Query Head Attention]
  H --> I[合并 Head 输出]
  I --> J[进入后续层]

这个数据流里最重要的边界是:Query Head 的数量和 KV Head 的数量被解耦了。Query 侧仍然可以保持较高并行度和多个表示子空间,KV 侧则通过分组减少需要长期缓存的数据。推理引擎在执行 Attention 时,只要知道某个 Query Head 属于哪个 KV Group,就能读取对应的历史 K/V。

因此,GQA 并不是在运行时“把已经生成的 32 份 KV 合并成 8 份”,而是模型从结构上就只产生指定数量的 KV Head。缓存节省发生在源头。已经训练好的 MHA 模型也不能仅靠修改一个推理配置把 32 个独立 KV Head 无损变成 8 个,因为原模型参数本身定义了不同的 K/V 投影。

KV Head 数量如何直接影响长上下文与并发能力

在不考虑量化、分页和缓存淘汰等额外机制时,可以把 KV Cache 的规模理解为与“层数 × 历史 Token 数 × KV Head 数 × 单 Head 维度 × K/V 两份数据 × 数据类型字节数”相关。GQA 能直接改变其中的 KV Head 数量,因此它的收益会随着上下文长度和活跃请求数量放大。

假设一个模型原本使用 32 个 KV Head,改为 8 个 KV Head,而单 Head 维度和其他条件保持一致,那么理论上的 KV 元素数量会下降到原来的四分之一。这不意味着整张 GPU 的总显存占用也一定下降四分之三,因为模型权重、激活、临时工作区和框架开销仍然存在。但当 KV Cache 已经是主要增量内存时,这种结构性缩减会明显扩大可用空间。

对长上下文聊天服务来说,收益通常体现在两个方向。第一,同一并发下,每个请求能保留更长历史而不那么快触发显存上限。第二,同一上下文长度下,GPU 可以同时维护更多活跃序列,从而给 continuous batching 更大的调度空间。后者是否真的转化为更高吞吐,还取决于计算、内存带宽和调度器能否处理新增并发。

GQA 也会影响 decode 阶段的带宽压力。每生成一个新 Token,Attention 需要读取历史 KV。KV Head 更少意味着需要访问的数据量下降,尤其在长上下文下更明显。但它不是一个保证端到端延迟固定降低多少的公式:模型其他层、权重读取、跨卡通信、采样和调度都可能成为新的瓶颈。

MHA、GQA 与 MQA 的选择本质上是容量和效率的交换

把三种结构放在一起看,GQA 的位置就很清楚。它不是“比 MHA 和 MQA 都更先进”的绝对答案,而是在 Key/Value 表示容量与推理资源成本之间提供更多可调档位。

架构Query HeadKV HeadKV Cache 压力表示自由度典型工程取舍
MHA与 Query Head 通常相同最高每个 Head 拥有独立 K/V 投影更重视模型容量,能承担较高缓存成本
GQA少于 Query Head、通常大于 1中等多个 Query Head 分组共享 K/V在质量、显存和带宽之间折中
MQA1最低所有 Query Head 共享同一组 K/V极度重视推理效率,可接受更强共享约束

用户材料引用的 GQA 工作报告称,GQA 在其实验设置中能够在接近 MQA 的速度特征下获得接近 MHA 的质量表现。这里的“接近”必须保留实验条件,不能扩展成所有模型、任务和 KV Head 配置都没有质量差异。KV Head 应该取多少,本质上仍是架构超参数,需要结合预训练目标、模型规模、训练数据和部署需求确定。

对于模型设计团队,问题可以转化为:当前模型是否主要面向长上下文和高并发服务?KV Cache 是否预计成为部署成本的重要组成?如果答案是肯定的,减少 KV Head 往往具有很强的工程动机。但减少到 8、4 还是 1,需要通过训练和评估验证,而不是仅根据缓存缩减比例决定。

GQA、PagedAttention 与 FlashAttention 处理的是三个不同层次

GQA 经常与 PagedAttention、FlashAttention 同时出现在推理优化讨论中,因此容易被混成同一类技术。它们实际作用在不同位置。

GQA 是模型结构设计。它决定一个新 Token 会产生多少组需要缓存的 Key 和 Value,从源头改变 KV Cache 数据量。

PagedAttention 是缓存管理策略。即使模型已经使用 GQA,多个请求仍然需要分配、增长和释放 KV Cache。分页式管理关注这些缓存如何映射到物理显存块,怎样减少预留和碎片,以及如何支持动态批处理。它不会把模型的 8 个 KV Head 自动变成 1 个。

FlashAttention 主要优化标准 Attention 的执行数据流,重点减少中间结果在 HBM 与片上存储之间的往返。它关注 Attention kernel 怎样计算得更高效,而 GQA 关注需要参与和缓存的 K/V Head 有多少。两者可以在同一个模型中同时存在。

技术作用层级主要优化对象是否改变模型 Attention 结构与其他两者关系
GQA模型架构KV Head 数量与 KV Cache 数据量可与缓存管理和高效 kernel 叠加
PagedAttention推理运行时KV Cache 分配、映射和回收管理 GQA 产生的更小 KV Cache
FlashAttentionGPU kernel / 执行Attention 中间数据流和 IO可加速采用 GQA 的 Attention 计算

对于一个长文档问答服务,可以同时使用 GQA 模型、分页 KV Cache 和优化 Attention kernel。GQA 让每个请求的缓存更小,PagedAttention 让剩余缓存更灵活地装入显存,FlashAttention 类实现则降低某些 Attention 阶段的数据移动成本。三者作用点不同,因此不是三选一。

推理引擎真正需要处理的是 Head 映射与缓存布局

从实现视角看,GQA 的核心关系并不复杂:num_query_heads 大于 num_kv_heads,并且通常要求 Query Head 能按固定分组映射到 KV Head。真正的工程难点出现在高性能 kernel、张量并行和缓存布局上。

在 KV Cache 中,系统只需要为 KV Head 分配存储。执行 Attention 时,多个 Query Head 读取同一 KV Head 对应的历史块。如果 kernel 仍然假设 Query、Key、Value Head 数完全一致,就需要修改索引和广播逻辑。高性能实现还要避免因为重复扩展 K/V 而把原本节省的内存重新复制回来。

张量并行会进一步影响布局。模型可能把 Attention Head 分散到多张 GPU,每张卡需要拥有与本地 Query Head 对应的 KV Head。若分片方式与 GQA 分组不匹配,通信或复制成本可能增加。因此,模型配置支持 GQA 只是第一步,推理框架还需要在本地 Head 映射、KV Cache 分片和 kernel 选择上正确适配。

这也是为什么“主流框架支持 GQA”不能等价于“任何模型配置都能自动获得最佳性能”。生产环境应确认实际 Attention backend 是否支持当前 num_query_heads / num_kv_heads 组合,缓存是否按 KV Head 数量真实分配,以及张量并行后是否出现额外复制。一个逻辑上正确但实现上把 K/V 展开到 Query Head 数量的路径,可能保留模型输出,却损失很大一部分预期缓存收益。

Uptraining 说明 GQA 与模型参数本身绑定

用户提供的 GQA 论文还讨论了把已有 MHA checkpoint 转换到 GQA 的 uptraining 方法,并报告使用相对于原始预训练较少的额外计算进行适配。原稿给出的数字约为原始预训练计算量的 5%,这个数字应限定在该论文方法和实验设置内,不能理解为任意 MHA 模型都能用固定 5% 成本完成无损转换。

从工程上看,这一方法说明了一个重要事实:GQA 是参数化架构的一部分。把多个独立 K/V 投影合并为更少的 KV Head,会改变模型内部表示方式,需要通过重新训练或适配让参数学会在新的共享关系下工作。仅修改配置文件,最多只能造成张量形状不匹配或错误解释已有权重。

因此,对于已经部署的大型 MHA 模型,团队通常需要在两条路线之间选择:继续保留原模型,通过 KV Cache 量化、分页管理、卸载等运行时技术降低成本;或者投入模型适配和重新评估成本,迁移到 GQA/MQA 架构。前者改动小但无法从根源减少 KV Head,后者潜在收益更大,却涉及权重转换、继续训练、质量回归和部署兼容性验证。

哪些情况下减少 KV Head 不一定带来更好的系统结果

第一种情况是 KV Cache 并非当前瓶颈。如果服务主要处理短 Prompt 和短输出,缓存占用很小,而计算时间集中在 MLP、权重读取或跨卡通信,那么减少 KV Head 对端到端吞吐的影响可能有限。此时为了 GQA 重新训练模型,收益未必覆盖迁移成本。

第二种情况是质量退化超过业务容忍范围。不同 KV Head 数量对应不同共享强度。某个模型在通用基准上差异很小,不代表特定领域、长尾任务或长上下文行为完全一致。模型团队需要在目标任务上重新评估,而不是只依赖平均分数。

第三种情况是推理栈没有真正利用 GQA。若 kernel 或缓存实现为了兼容旧路径,把共享 K/V 显式复制到每个 Query Head,显存和带宽优势会被部分抵消。上线前应通过 profiler 和显存统计确认缓存张量真实形状,而不是只查看模型配置中的 num_key_value_heads

第四种情况是并发提高后瓶颈转移。更小的 KV Cache 允许加入更多请求,但更多请求也会增加调度、权重计算和其他层的工作。系统可能从“显存装不下”变成“计算吞吐不足”。这不是 GQA 失效,而是优化释放了一个约束后,新的最长路径开始主导性能。

生产环境应该观察哪些信号

评估 GQA 是否带来真实收益,可以从缓存、请求和硬件三个层面看。

缓存层面需要确认每个请求的 KV Cache 实际字节数、KV Cache 总使用率、缓存块分配失败或抢占频率,以及上下文增长时的斜率。如果模型从 MHA 迁移到 GQA 后,这些指标没有按预期改善,应检查 KV Head 配置和缓存布局是否真正生效。

请求层面关注可维持的并发序列数、等待队列长度、TTFT、每 Token 延迟和整体吞吐。GQA 更可能直接改善 decode 阶段的 KV 读取压力和并发容量,但首 Token 延迟还包含 prefill 计算,因此不能只看一个指标判断效果。

硬件层面则要看 HBM 带宽、GPU 利用率和 Attention kernel 时间。若 KV 读取原本占据明显带宽,减少 KV Head 后可能降低这部分压力;若 GPU 计算单元已经是主瓶颈,缓存缩减则更多体现为“可以放更多请求”,而不是单请求立即变快。

最可靠的验证方式是固定模型规模和服务条件,对比不同 KV Head 架构或官方提供的对应模型版本,在真实 Prompt/输出长度分布下测量。测试应包含单请求长上下文、固定高并发以及混合长度流量。只有同时观察质量回归和服务指标,才能判断共享程度是否落在合适的位置。

GQA 最值得记住的不是“多个 Head 共用一份资料”这个类比,而是它把 Query Head 数量与 KV Head 数量从一一对应关系中解耦出来。模型仍然保留较多 Query 投影去提出不同查询,却只为较少的 K/V 投影维护历史缓存。

对于长上下文在线服务,这种结构直接减少每个请求随序列增长的状态量,也降低了读取历史 KV 的压力。但它的收益取决于 KV Cache 是否真是瓶颈、推理引擎是否正确支持共享布局,以及减少 KV Head 后的模型质量是否满足目标任务。GQA 因此更像一个模型架构层面的资源预算旋钮:从 MHA 向 MQA 移动时,缓存和带宽成本持续下降,共享约束持续增强,而工程上真正需要决定的是哪一个位置最适合当前模型与部署环境。

资料来源

  1. GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints