[KV Cache优化]MQA/GQA/YOCO/CLA笔记: 层内和层间KV Cache共享

主要观点总结

本文是对近期微软和MIT-IBM Watson AI Lab关于KV Cache共享策略的新论文的阅读笔记。文章介绍了层内KV Cache共享（MQA和GQA）和跨层KV Cache共享（YOCO和CLA）的相关内容，包括它们的原理、计算公式以及实验结果。文章还提供了关于LLM推理部署的推荐资源。

关键观点总结

关键观点1: 层内KV Cache共享：介绍MQA和GQA

MQA是极端情况下只保留一个KV Head，多个Query Heads共享相同的KV Head。GQA则采取折中的做法，将Query Heads进行分组，每组对应一个KV Head。这两种方法都可以在减少计算量和KV Cache的同时确保模型效果不受太大影响。

关键观点2: 跨层KV Cache共享：介绍YOCO和CLA

YOCO提出一种全新的KV Cache共享方式，即层间共享。CLA则是另一种跨层KV Cache共享的新方式，通过将产生KV Cache的层交替分布在模型不同深度的层来实现。这两种方法都可以降低KV Cache的显存开销。

关键观点3: KV Cache的计算公式

在MHA下，KV Cache的计算公式为：KV Cache = 4 x B x L x H x N bytes。对于GQA，如果Q的组数为G，则需要的KV Cache为：KV Cache = 4 x B x L x H x N / G bytes。YOCO和CLA的公式则在GQA的基础上进一步考虑了层间共享的因素。

关键观点4: 实验结果和性能优势

YOCO和CLA在实验中表现出了优异的性能，能够在节省显存和提高推理速度的同时保持模型的精度。不过，对于更大规模的模型，如72B参数规模以上的模型实验尚未展示。

关键观点5: 总结与推荐资源

文章总结了KV Cache共享的几种算法，并推荐了关于LLM推理部署的资源。同时，提供了一个公众号二维码供读者关注交流。

文章预览

作者丨DefTruth 来源丨https://zhuanlan.zhihu.com/p/697311739 编辑丨GiantPandaCV 0x00 前言 微软最近新发了一篇论文，提出了YOCO（You Only Cache Once，和RetNet似乎是相同的作者），这是一个KV Cache层间共享的新思路。同期MIT-IBM Watson AI Lab也发了一篇类似的论文，提出了CLA（ Cross-Layer Attention），即KV Cache跨层推理。简直和YOCO不谋而合，因此本文把这两篇论文的阅读笔记放到一起记录了。 由于最近一两年都是在做推理相关的工作（刚毕业的时候做了几年CV算法，偶尔客串一下NLP），对于非推理方向的论文关注已经变少了一些。但是最近在arxiv刷到这篇paper，看完摘要，第一反应就是，“我去，这样做好像还真可以，为啥我没想到（太菜了）”。YOCO的全称是 《You Only Cache Once: Decoder-Decoder Architectures for Language Models》 ，YOCO的名称借鉴了单阶段目标检测始祖YOLO的风格。整篇论 ………………………………