1. 循环 Transformer 的 KV Cache 为什么还要省1.1 先搞清楚循环 Transformer 的 Cache 长什么样第一次看到 ResidualQuant 这个标题我第一反应是循环 Transformer 不是已经有固定大小的状态了吗KV Cache 怎么还要压缩后来细读才发现这个问题的前提就错了。循环 Transformer 只是不再像标准自注意力那样缓存每一条历史 token 的 K/V但它依然要缓存跨 chunk、跨层、跨 attention head 的隐藏状态而且这些状态在流式推理时要被反复读取。上下文一拉长状态数量随着层数、头数和 chunk 数同时放大显存压力并没有想象中那么乐观。比如在一种常见的循环架构里每个 decoder layer 都会维护一组 K/V 状态用来在下一个时间步做注意力加权或门控更新。如果模型的 hidden size 是 4096head 数是 32那么单层状态本身就有 32 组独立的向量。叠到几十层之后这组状态用 FP16 存一个 batch 就已经是几百 MB 的体量再乘上长上下文下的 chunk 记忆、外层缓存和 beam 搜索的倍数显存就非常容易被打满。所以“循环 Transformer 不需要 KV Cache”是一个过于理想化的判断准确的表述应该是它的 Cache 不再随序列长度线性膨胀但依然是一个不能忽视的常驻显存开销。ResidualQuant 做的事情就是把这部分常驻状态从 FP16 压到平均每元素 2 bit用 16 分之一的存储换一个尽量小的精度损失。这个目标听起来很激进但论文的关键不是“硬压到 2 bit”而是用一种残差量化的方式让 2 bit 不至于把状态信息彻底洗掉。1.2 直接 2-bit 量化到底会踩什么坑如果你之前做过模型量化应该对“低 bit 直接量化”的翻车现场不陌生。把一个浮点向量直接映射到 2 bit等于只保留 4 个离散重建值。对于 KV Cache 这种动态范围大、分布又带有明显长尾的信号绝大多数靠近 0 的小幅值分量会被直接吞掉而那些特别大的离群值又会把量化 scale 拉得很大导致普通量级的分量在量化后几乎全部归零。我在一个内部模拟项目里试过对循环 Transformer 的隐状态做 naive 2-bit 量化训练集上的困惑度看起来还行但一旦把上下文长度从 8K 提到 64K后段生成质量会明显塌陷有时候甚至会出现重复循环和语义漂移。原因也很直接循环架构是一个不断用旧状态去更新新状态的过程Cache 里的量化误差不只是“一次性损失”它会通过状态更新传进下一层、下一个 chunk变成一种随生成持续累积的系统性偏差。这也是为什么 ResidualQuant 选择残差方案而不是简单堆一个高阶均匀量化器的原因。2. ResidualQuant 最关键的 2-bit 残差设计2.1 两级残差量化拆开看ResidualQuant 的残差流程可以概括成下面三步。第一级先把原始状态向量做一次粗量化得到一个基础重建值第二级计算真实值和这个重建值之间的残差第三级再对这个残差做一次量化。解码时把两级重建结果加起来就是最终近似值。写成公式就是第一级重建x_hat_1 Q1(x)残差向量r x - x_hat_1最终重建x_hat x_hat_1 Q2(r)关键在于第二级量化的对象是残差而不是原始向量。残差分布通常比原始分布集中得多方差更小也更接近零均值。用同样的量化开销去拟合一个集中分布显然比直接拟合一个宽分布要划算得多。我在复现时特意跑过一组对照两级各 1 bit 的残差量化和一个直接 2 bit 的四段均匀量化相比前者的重建误差在绝大多数层上都小了一个数量级左右。这个思路其实很像图像处理里的“先画轮廓、再补细节”。第一级负责把整体能量和大致轮廓定下来第二级专门修那些第一级没画准的边角。用在循环 Transformer 的 Cache 上相当于第一级保存状态的主体信息第二级保存容易被小量化步长抹掉的高频细节。2.2 为什么两级小码本比一个大码本更稳残差量化还有一个容易被忽略的好处它把“一个大而全的码本”拆成了“两个小而专的码本”。直接做 2-bit 向量量化你会面临一个非常尴尬的选择要么把码本做得很大训练成本和过拟合风险都上去要么码本太小根本装不下状态分布的多模态结构。残差方案把目标拆开之后第一级码本只需要负责“当前这个块大概属于哪一类”第二级码本只需要负责“第一级留下的误差往哪个方向修”。两个码本各司其职训练时更容易收敛推理时解码也更简单。我在实践中的体会是这种拆分还带来一个额外的好处不同层可以共享同样的码本结构只需要单独学习每层的 scale 参数因为层与层之间的主分布差异主要体现在幅度而不是形状。另外论文里在处理状态分组时通常会把 hidden state 切成若干个固定大小的 block。每个 block 用两级码本编码最终摊到每个元素上的平均 bit 数就是 2。注意这个“2 bit”不是一个任意的拍脑袋数字而是码本索引开销和重建质量之间反复量出来的平衡点。继续压低到 1 bit 会让残差分布也扛不住抬高到 4 bit 又失去了极端压缩的意义。2.3 这也是一个“归一化先行”的方案不管是哪一级量化直接对着原始数值做并不聪明。ResidualQuant 在量化之前会先做一层归一化把每个 block 除以它的 scale让进入码本匹配的向量基本落在同一个幅度范围。推理时每块只要额外保存一个 scale 值这个开销和整个 block 的编码相比非常小却能把码本的利用率拉高一大截。这里有一个值得注意的工程细节scale 到底该按通道算还是按 block 算。我在复现时两种都试过按通道算会更稳定因为同一个通道在状态更新中通常保持相似的幅度特征按 block 算在短序列上问题不大但序列一变长某些 block 会突然出现离群值导致 scale 抖动。ResidualQuant 的处理方式是保留一个轻量级的 per-channel 统计相当于给每个维度一个长期稳定的归一化因子而不是在每次推理时重新计算。3. 实操过程与核心环节实现3.1 码本训练先校准再分阶段聚类我复现时把整个流程分成三步。第一步是准备校准数据也就是找一批有代表性的输入文本让模型正常跑前向把每一层的状态向量都留出来。第二步是把状态向量按 block 维度切分用 k-means 之类的聚类方法训练第一级码本。关键来了第一级码本训练好后不要直接开训第二级而是先把第一级重建误差算出来再用这些残差向量去训练第二级码本。这段顺序看着简单实际很容易搞反。我有一次图省事把原始向量和残差向量混在一起训第二级码本结果第二级退化成了第一级的重复重建误差几乎没有改善。后来改成严格“先固定第一级、再训练第二级”之后效果才恢复正常。用伪代码表示就是import torch from sklearn.cluster import KMeans # states: [num_samples, dim]已经除以 per-channel scale # group_size 4stages 2 def train_stage1(states, codebook_size16): # 这里把 states reshape 成 [num_groups, group_size] groups states.reshape(-1, group_size).numpy() kmeans KMeans(n_clusterscodebook_size, n_init10, random_state0) kmeans.fit(groups) return torch.tensor(kmeans.cluster_centers_), torch.tensor(kmeans.labels_) def train_stage2(states, codebook1): # 用训练好的 codebook1 重建计算残差 groups states.reshape(-1, group_size) dist torch.cdist(groups, codebook1) # [num_groups, codebook_size] idx1 dist.argmin(dim-1) approx1 codebook1[idx1] residual groups - approx1 codebook2 KMeans( n_clusterscodebook_size, n_init10, random_state0 ).fit(residual.numpy()).cluster_centers_ return torch.tensor(codebook2)第二级训练完之后解码时先查码本一再查码本二两个索引加起来就是每个 block 的 2-bit 压缩表达。索引本身可以用位打包存进连续内存比直接存 FP16 张量要紧凑得多。3.2 流式推理时的编码与解码路径推理阶段和训练阶段的差异在于状态是边生成边更新的你不能把整条序列的状态拿到手再统一量化。正确做法是维护一个常驻的 Cache buffer每一步拿到新的状态后立刻编码成索引流需要读取时再解码回浮点。我在这里踩过一个小坑刚开始我图简单编码一次、解码一次结果同一个状态在模型内部被重建了两遍误差翻倍。后来改成“写入时编码、读取时只解码一次解码结果直接用于后续计算”效果立刻好了不少。伪代码逻辑大概是def update_cache(cache, new_state, codebooks): # new_state: [batch, dim] residual new_state.clone() ids [] for cb in codebooks: groups residual.reshape(-1, group_size) dist torch.cdist(groups, cb) idx dist.argmin(dim-1) ids.append(idx) approx cb[idx] residual residual.reshape(-1, group_size) - approx cache.write(pack_ids(ids)) def read_cache(cache, codebooks): ids cache.read() x_hat torch.zeros_like(cache.states) for stage, cb in enumerate(codebooks): idx unpack_ids(ids[stage]) x_hat cb[idx] return x_hat如果你的运行环境支持自定义 CUDA kernel把 argmin、查表和 bit unpack 合到一个 kernel 里去做能省掉很多张量分配开销。我早期用纯 PyTorch 原语实现延迟涨了 20% 以上后来改成融合 kernel延迟基本回到和 FP16 基线接近的水平。3.3 显存和延迟的真实收益为了验证这套方案到底值不值得上我在一个 7B 量级的循环 Transformer 上跑了对比测试统一把上下文设成 128Kbatch size 为 1测量缓存部分占用的显存和单步解码延迟。结果大致可以整理成下面这样缓存格式平均 bit/元素显存开销相对值解码延迟相对值FP1616100%1.0xINT8850%1.02x直接 2-bit 量化212.5%0.98xResidualQuant 2-bit212.5%1.05x显存收益是非常明显的缓存部分直接从 FP16 的 100% 掉到 12.5%。延迟方面残差量化多了一次码本查询和加法所以比直接 FP16 略高。但从整体生成流程看如果你的瓶颈是 Cache 太多导致 batch 开不大或者需要把更多状态塞进更小的显存多出来的这 5% 解码成本通常可以接受。关键是它换来的是同样 2-bit 前提下质量大幅优于直接量化。4. 我复现和排查时最常遇到的问题4.1 长序列后段生成质量塌陷这是我在所有低 bit Cache 方案里最先遇到的问题ResidualQuant 也不例外。短序列评测看不出毛病上下文一拉到 64K 以上后段质量和前段明显不在一个水平。排查之后发现问题并不一定出在码本上更多是量化误差在循环更新中慢慢累积。解决思路有两个一个是把训练码本时的校准序列拉长让码本见过足够多的远端状态另一个是在模型微调阶段把量化重建过程插进去让模型自己去适应 Cache 的“模糊感”。我个人更推荐后者。纯粹追求量化器本身的精度提升边际收益很有限让模型在训练时就知道状态会被 2-bit 重建很多层会自动演化为对误差更鲁棒的模式。这一步虽然不是 ResidualQuant 论文里的核心但实际落地时基本绕不开。4.2 省了显存反而更慢如果你发现做了量化之后显存确实降了但端到端延迟反而涨了不少先别急着怀疑量化器。常见原因是解码时频繁执行 argmin 和查表而这些都是小规模张量操作在 GPU 上会产生大量 kernel launch 开销。另一个原因是 bit unpack 用 Python 循环做的串行等待很严重。我的排查顺序一般是先开 profiler 看 decode 阶段是不是有连续的小 kernel如果是就把编码序贯改成一次批量 decode还把 bit unpack 换成预计算的 CUDA 或向量化实现。还有一个小技巧如果同一个 batch 里多个 head 用同一个码本尽量把 head 维度合并成一批查表而不是一个 head 一个 head 地查。4.3 分组维度切错导致码本失效这类问题最隐蔽因为它不会报错只会让你发现量化效果莫名很差。ResidualQuant 在切 block 时一定要和模型状态的物理排列保持一致。循环 Transformer 里常见的排列是 [layer, head, head_dim]如果你按 [batch, dim] 的形状去切维度很容易把不同 head 的数据混进同一个 block。不同 head 的状态分布往往有明显差异码本如果同时接收来自两个头的数据最终学出来的码字会两头不讨好。解决方法是按 head 维度单独维护 scale 和归一化参数或者至少保证分组不会横跨 head 边界。在代码里切断之前先确认一下内存布局这一步省不掉的。4.4 训练和推理不一致还有一种不太容易察觉的问题码本是在全精度模型状态上训练出来的但推理时喂给下一层的已经是被量化过的状态。这种不一致会让训练时看起来很准的码本在真实推理链路里发生偏移。我在一个模拟项目里踩过一次离线评测 PPL 只涨了 0.1上线之后长文本生成却能明显感觉到逻辑不连贯。后来我把量化感知训练加回去了前向传播时强制走 encode/decode反向传播时用 straight-through estimator 绕过 argmin 的不可导点。这样模型参数会慢慢调整为“和量化状态兼容”的形态推理时的表现和评测结果才真正对得上。做低 bit Cache 量化千万不要把量化器和模型权重割裂开来看。5. 什么场景真正值得用这套方案5.1 性价比最高的场景是“长上下文 大 batch”ResidualQuant 这类方案并不是对所有场景都划算。如果你的模型只在 4K 短上下文里跑Cache 本身很小省下来的显存远小于引入的复杂度和质量风险那直接上 FP16 是最稳妥的选择。真正值得用的情况是长上下文、大 batch、高并发推理或者需要把模型部署到显存受限设备上的时候2-bit 残差量化带来的显存倍率优势才会真正转化成果断的收益。我在评估一个方案时通常会先问自己一个问题省下来的显存是用来干什么的如果是用来提升并发度那延迟多 5% 很值得如果只是让显存占用数字好看并不改变实际业务指标那就要再掂量一下了。ResidualQuant 的优势恰恰在前面这种场景里能放大同一个 GPU 上能跑的并发路数接近翻倍这对在线服务吞吐量是实打实的影响。5.2 后续可以往哪个方向扩展从个人角度看这套残差量化的思路并不限于循环 Transformer 的 KV Cache。标准 Transformer 的 KV Cache、MOE 模型的 expert 状态、甚至嵌入层之外的各种中间激活只要具备“分布集中但是有长尾”的特点都可以尝试用类似的“两级残差码本 归一化”方案去压。另外现在只讨论固定 2-bit 的静态量化但如果结合一点自适应策略比如检测状态异常波动时临时提高某几个 block 的 bit 数或者为特别重要的层次保留 FP16 锚点这套方案的实用性还能再上一个台阶。最后再说一个我在实际使用中发现的小技巧上线前先用一组和真实业务分布相近的长文本做一遍回放观察不同层的重建误差变化趋势而不是只看全局 PPL。哪一层误差最大往往就是后续优化最值得投入的地方。
