跳到主要内容

KV 缓存这件事,终于有人从架构自己身上找抓手了

画唠
画唠

· 阅读约 5 分钟

KV cache 量化这摊子,我本来以为没什么可看的了——旋转、分组、给 outlier 留条道,翻来覆去就这几板斧,区别只在谁的手更稳。所以看到 ResidualQuant 这篇的时候我没抱什么期待。

但先撂一句:它没那么玄,一句话能讲完——别把每一圈的 KV 当独立的东西存,存它们跟最后一圈的差。

这句话背后那个观察,比论文里任何一个数字都值钱。画开看看。

它压的根本不是 KV

要看懂它为什么能这么干,得先看循环式 Transformer 长什么样。这东西的路子是:同一组模块,转着圈地用。

   x ──►[ 共享块 ]──►[ 共享块 ]──►[ 共享块 ]──► … ──► 输出
          ▲            │
          └────────────┘
          同一组权重,重复转 N 圈

好处明摆着:参数不涨,计算深度照涨。等于用"多转几圈"换"多堆几层",账面划算。

但账不是这么算的。每一圈都要往 KV cache 里写东西,圈数翻倍,KV 就翻倍。参数那头省下来的,显存这头原封不动还回去——而且不是一次性还,是推理的时候一直占着,直接掐住批大小和吞吐。

一个"用计算换参数"的架构,最后被显存掐住脖子。

这就有意思了。论文盯上的就是这中间的缝隙:不同循环产生的 KV 状态,彼此高度相似。

我盯着这句话看了半天。因为如果这是真的,前面那一堆把 KV 当一坨任意张量硬压的方法就有点冤枉——它们在反复压一个本来就没多少新信息的东西,压了一遍又一遍。

于是:

 朴素做法:每圈各存一份
   loop 1  [========================]
   loop 2  [========================]
   loop 3  [========================]
   loop 4  [========================]   显存 ≈ 圈数 × 单份 KV

 ResidualQuant:
   loop 4  [========================]   最后一圈,当参照,精度给足
   loop 3  [··]                         存的是「跟 loop 4 差多少」
   loop 2  [··]
   loop 1  [···]
                                        显存 ≈ 一份 KV + 几份薄残差

图里我把每圈画成一样宽,这是简化。实际是按圈分配精度的混合精度策略,靠近参照的那几圈残差更薄,越往外留的位越多。大方向就是这张图:从"存 N 份完整状态"变成"存 1 份完整 + N-1 份差值"。

咔哒一下扣上了。✨

这个比喻在这儿漏风

"存差值"这说法太顺了,顺到我差点就信了下一步——既然相似,那差值应该接近零,随便压压就行。

不对。

如果差值真接近零,作者不需要往残差上再套一遍最小二乘缩放、再套一遍旋转变换。这两步存在本身就说明:差值不小,直接砍位数会崩,得先把它掰到一个量化友好的形状上,才敢往 INT2 压。

所以"高度相似"的正确读法是:相似到一份 2-bit 残差还能把信息兜住。不是相似到可以扔。这两个说法差得很远,我第一遍就是按后者理解的,然后发现论文里有一堆机制根本没法解释——那种感觉,脑内只有羊叫。

顺便,2-bit 这个数字得单独拎出来。玩过 KV 量化的人都知道线在哪:4-bit 谁都能做,3-bit 开始挑模型,2-bit 是绝大多数方法开始胡言乱语的地方。敢报 INT2,说明它赌的就是"残差里剩下的信息量本来就少"这件事。

数字得掰开看

摘要里最扎眼的那个 80.7%,我建议先放一放。原文写的是"理论 KV 存储需求"——理论两个字在这儿干了很多活。理论省下来的显存和你真跑起来省下来的,中间隔着分配器、碎片、还有 kernel 到底吃不吃你这套布局。

硬的数有另外两个。

一个是同显存预算下的准确率:比旋转类基线最高高 13.0%。这个数重要,因为"省 80% 显存"本身不难,你把位数砍狠一点谁都能省,难的是砍完还准。13 个点就是砍完还准的那部分。

另一个是 RTX 5090 上的实测:固定批量时解码吞吐最高 2.73 倍;批大小能开到 2 倍、峰值吞吐最高 4.15 倍。这俩别混着读,它们说的是两件不同的事——前者是"同一批活干得更快",后者是"显存省出来了,我能往里塞更多活"。后者才是 KV 量化真正想换的东西。

(这儿想岔一句:解码阶段到底卡带宽还是卡算力,我上次画的时候就搞混过一次,回头单独开一期。先拉回来。)

不过这俩数都是最好情况。"最高"两个字每次都出现在关键位置。

我真正想说的那点

回到开头。它压的不是 KV 本身。

大部分 KV 量化论文的默认姿势是:给你一坨张量,不管它从哪来、代表什么,我就是要用尽量少的位编码回去。纯信号处理的路子,通用,但也因此拿不到架构给的便宜。

ResidualQuant 反过来:先问这坨张量是怎么长出来的,发现它是同一个块转圈转出来的,那圈与圈之间就必然带着结构性的冗余——然后才动手,压的是那个冗余。

这个顺序上的差别,比它具体用了哪种变换重要得多。先看结构,再选工具;不是拎着锤子找钉子。

说到这儿我得泼一盆冷水:looped Transformer 现在有多少人真在生产里跑?我没查到,论文摘要也没给代码或权重的发布信息,arXiv 页面上就 PDF、实验性 HTML、TeX 源码那几样常规入口。一个只在特定架构上成立的方法,它那个 80.7% 就不能直接翻译成"你的推理服务能省 80.7%"。

但这不影响它对我有价值。它给的是一条思路:压缩的抓手,可以从架构自己身上长出来。这条思路能挪到别的"结构性重复"上去——挪到哪,我还没想明白,先撂这儿,想明白了再补一期。

自检

照例留个自检:现在你能不能用自己的话,跟一个没看过这篇的人讲明白——为什么"不同循环的 KV 高度相似"这件事,能直接变成显存?

如果在"存差值"这步就停下了,那你停的地方跟我第一遍一样:差值不是零,得先掰形状再压。这个坑我替你踩过了。

下期想画开哪个?我手上还压着两个:一个是"为什么 KV 量化总在 2-bit 崩",一个是"解码到底卡在带宽还是算力"。你说了算。

本期画开:压的不是数据,是数据怎么长出来的。

画唠
画唠

把被讲玄的概念用图 + 比喻 + 动手实验拆到咔哒扣明白,错的也保留。

查看主页 →