大模型逐词生成时,每出来一个新词元,都要回头看前面的键和值。服务端通常把这些结果留在 KV 缓存里,免得每一步重新计算。上下文越长、并发请求越多,这块缓存越占显存,读取它也会吃掉内存带宽。
我把多查询注意力论文、GQA 原论文和 Llama 2 技术报告对了一遍。几种注意力的差别集中在一处,查询头到底各自保留一份键和值,还是与别的查询头共享。这个选择会同时改变缓存大小、生成速度和模型质量。
从每头一份到几头共用
标准多头注意力会为每个头计算查询、键和值。多个头可以从不同角度比较词元,代价是生成时要反复读取许多组键和值。2019 年的多查询注意力把所有查询头改成共用一组键和值。论文把目标写得很直接,减少增量解码时张量大小和内存带宽,同时接受少量质量损失。
GQA把共享程度放在两端之间。假如模型有三十二个查询头,可以把它们分成八组,每组共用一个键头和一个值头。KV 缓存只需保存八组,查询仍保留三十二个头。组数越少,缓存通常越省;共享越强,能保留的键值差异也越少。
GQA论文还给出一条改造旧模型的路。研究者先把多头注意力的键和值按组聚合,再用原预训练计算量约百分之五继续训练。实验中,转换后的 GQA 质量接近多头注意力,速度接近多查询注意力。这个结论来自论文模型和任务,具体部署仍会受头数、上下文长度、批量大小与推理引擎影响。
Llama 2 的技术报告提供了一个落地例子。70B 模型采用 GQA,7B 和 13B 使用多头注意力。报告把这项选择与大模型的推理可扩展性联系起来。由此也能看出,GQA通常在模型训练或结构设计时确定,并非给任意现成模型改个服务参数就能安全切换。
省下多少要在目标负载上量
验收这类模型,我会先记录单请求在不同上下文长度下的 KV 缓存占用,再逐步增加并发,观察吞吐和每个词元的延迟。随后用同一批任务比较原模型与共享键值后的正确率,尤其留意需要分辨多个相近对象、长距离引用和代码修改的样本。
缓存缩小不保证端到端延迟按同样比例下降。预填充还要处理整段输入,采样、网络和其他算子也占时间。短回答、低并发场景可能看不出明显收益;长上下文和高并发更容易让内存带宽成为限制。
GQA是一种模型结构取舍。它用较少的键值头服务较多查询头,给 KV 缓存减负,又比所有头共用一组留出更多表达空间。最终值不值得,要看目标机器上的显存、速度和任务质量三份记录。