四 k 上下文的缓存，为什么会比模型权重更占地方

端侧跑不动长上下文，常被算成模型太大的问题。缓存按层数乘宽度乘序列长度长，七 B 模型在四 k 上下文要 2.1 GB，八 k 要 4.3 GB，而 int4 权重只有 1.75 GB。把权重再压一半，缓存一点没少。

每个 token 都要留一份

自回归推理里，每生成一个 token，模型要把这个 token 在每一层的键和值存下来，供后面的 token 做注意力计算。这份存储叫 KV 缓存，它的大小和三个量成正比：层数、键值头的维度、已经看过的 token 数。前两个由模型结构定死，改不动；只有第三个随对话一直涨，端侧内存规划的难点全在它身上。

写成式子是缓存字节数等于 2 乘层数乘键值头数乘每头维度乘序列长度乘每元素字节数。前面的 2 是键和值各一份。

键和值都要留，原因在注意力的算法里。每个新 token 要拿自己的查询向量和所有历史 token 的键做点积，算出权重，再用这些权重对历史 token 的值加权求和。键和值在第几层、第几个头被算出来之后就不会变，但每次生成都要用，所以只能存着。

查询向量不用存。它只服务当前这一步，算完就丢。这也是缓存的字节数里只有键和值、没有查询的原因。

以常见的七 B 模型为例，32 层，隐藏宽度 4096，32 个注意力头，每头 128 维，半精度存两字节。每个 token 的缓存是 2 乘 32 乘 32 乘 128 乘 2，等于 524288 字节，也就是 512 kB。

512 kB 看着不大，乘上序列长度就不一样了。

四 k 就是 2.1 GB

4096 个 token 乘 512 kB，2.1 GB。8192 个 token 乘同样的数，4.3 GB。

分摊到每一层看，这个数字更直观。32 层，每层每 token 16 kB，四 k 序列每层 64 MB，32 层摞起来 2.1 GB。任何一层都不能少存，因为解码时每层都要重新算注意力。

这个数字和权重的量级放在一起才有意义。七 B 模型用 int4 量化，权重 1.75 GB；用 int8，3.5 GB。四 k 上下文的缓存 2.1 GB，已经超过 int4 权重，接近 int8 权重。

缓存还有一个特点，它是动态的。权重加载一次就固定，缓存随对话变长一直涨。用户聊到第 50 轮，缓存可能比刚启动时大十倍。显存规划如果只按权重算，跑到一半就会撞墙。

缓存还会直接改出字速度。解码阶段每生成一个 token，要把权重和全部缓存从内存读一遍，这一步是带宽受限的。手机内存带宽按 50 GB/s 算，1.75 GB 权重加 2.1 GB 缓存是 3.85 GB，读一遍 77 毫秒，每秒出 13 个字。序列拉到 8k，缓存 4.3 GB，读一遍 121 毫秒，掉到每秒 8 个字。上下文长了以后出字变慢，原因是这个，不是算力不够。

八 G 的手机只剩 6 G 可用

一台 8 GB 内存的手机，操作系统和常驻服务占掉 2 GB 左右，留给应用的是 6 GB。int4 权重 1.75 GB 加四 k 缓存 2.1 GB，合计 3.85 GB，还剩 2 GB 给激活值、算子临时缓冲和框架本身。

序列拉到 8k，缓存 4.3 GB，加权重 6.05 GB，正好越过 6 GB 的线。这就是为什么很多端侧应用把上下文卡在 4k，不是模型不支持更长，是内存装不下。

换成 fp16 权重，3.5 GB 加 2.1 GB 是 5.6 GB，四 k 就已经很紧。量化在这里救的是权重那一半，缓存那一半纹丝不动。

并发还要再乘一次。端侧应用常常要同时维护系统提示、历史对话和当前输入三份序列，两路并发意味着缓存翻倍到 4.2 GB。多数手机上的助手其实是单路串行的，一旦加上后台摘要或者多轮预取，内存立刻不够。

图 1　KV 缓存随序列长度的增长与内存边界（依据原文第 2、3 节绘制）

系统提示那一份还特别贵。一份两千 token 的系统提示，缓存 1 GB，每次开新会话都要重新算一遍预填充。把它常驻下来省了计算，可内存又被占住 1 GB。端侧要在这两者之间选：要么每次重算，首字延迟多几百毫秒；要么常驻，可用内存少 1 GB。

分组查询把 KV 头砍到四分之一

键值头不必和查询头一样多。分组查询的做法是让多个查询头共享一组键值头，32 个查询头配 8 组键值头，每组的缓存只有原来的四分之一。

按这个配置重算：2 乘 32 乘 8 乘 128 乘 2，每 token 131072 字节，128 kB。四 k 上下文的缓存降到 0.54 GB，8k 降到 1.08 GB。

这一招把缓存从 2.1 GB 压到 0.54 GB，四倍。手机上的 8k 上下文从“装不下”变成“装得下还有余量”，代价不在内存，在别处。

分组数是连续的旋钮。1 组键值头是最早的多查询注意力，缓存再降 8 倍，精度掉得明显；32 组就是标准多头注意力，缓存最大，精度最好；8 组是中间档，这也是多数开源模型默认选 8 的原因。旋钮的位置决定了缓存和精度怎么分账，两件事没有同时拿满的配置。

图 2　多头与分组查询的缓存占用对比（依据原文第 4 节绘制）

省下的四倍从注意力里扣

查询头从 32 个减到 8 组共享，每个键值头要服务 4 个查询头。注意力的表达力跟着降，模型在需要区分细粒度位置关系时容易混。

具体表现是长文检索类的任务掉点。在一个十几万 token 的文档里找一句具体的话，多头注意力的不同头可以分别盯不同的位置模式，共享之后这些模式被压到同一个键值子空间里，召回率往下走。公开的对比里，这种配置在短任务上几乎无损，在长文精确检索上会有可测的下降。

这条取舍的边界很清楚：对话、摘要、改写这类任务，缓存省四倍几乎不付代价；需要在大段材料里精确定位的任务，省下来的内存要用精度还。

实测里的差别不在平均值，在尾部。同一套评测，短问题的得分几乎一样，长文档里找具体数字的那几题，分组查询的版本会多错一两道。厂商报的均值看不出来，用的人能感觉到，而这类任务恰恰是长上下文最常被用到的场景。

还有一条路不改缓存结构，改看多远。滑动窗口只保留最近几千个 token 的缓存，超出窗口的丢掉；稀疏注意力按块保留，只留和当前查询相关性高的部分。两种做法的缓存都能固定在一个上限，代价是远距离的信息真的看不到，模型不知道窗口外面发生过什么。

量化权重救不了缓存

端侧优化的惯性是继续压权重。int4 压到 int3，1.75 GB 变 1.3 GB，省 0.45 GB；同一时间四 k 缓存还是 2.1 GB，纹丝不动。

权重和缓存的压缩比不是一回事。权重的数值分布集中，量化误差可以靠校准集补；缓存的数值是激活值，分布随输入变化，动态范围大，压到 8 位以下就会明显掉点。这也是为什么权重量化已经做到 4 位、2 位，缓存量化还停在 8 位附近。

把两边的账放在一起看：int4 权重 1.75 GB 加 2.1 GB 缓存，缓存占 55%；换成 int8 权重 3.5 GB，缓存占 38%。权重压得越狠，缓存占比越高。继续压权重，边际收益越来越小。

把缓存换到闪存也不是出路。手机的闪存顺序读按 2 GB/s 算，2.1 GB 的缓存换入一次要 1 秒，而生成一个 token 的预算是 77 毫秒。差一个数量级，换入换出这条路在端侧走不通，这也是稀疏注意力和分组查询被反复提起的原因，它们动的是缓存的增量，不是缓存的存放位置。

一句话：端侧长上下文的门槛不在权重，在按 token 线性增长的缓存，压权重压不到它。

内容来源：arXiv 预印本（开放获取）

原文标题：GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints

作者：Joshua Ainslie, James Lee-Thorp, Michiel de Jong, Yury Zemlyanskiy, Federico Lebrón, Sumit Sanghai

DOI：https://arxiv.org/abs/2305.13245 （arXiv:2305.13245，2023）

配图说明：两张配图均由作者自行绘制，为原理示意，非原文原图复制。

改动声明：本文为对开放获取文献的技术编译与解读，标题、行文与配图为原创。

来源：21ic电子网