WUSH-KV: KV Cache Quantization with Data-Adaptive Transforms
Jiale Chen, Vage Egiazarian, Eldar Kurtić, Torsten Hoefler, Dan Alistarh
cs.LG
2026-09-30
把 WUSH 变换移植到 KV cache:key/value 各学一个可逆变换,value 侧零开销折进权重;2-bit 下 Qwen3-8B 困惑度 10.51,优于 OSCAR 的 13.74。
KV cache 随序列长度和 batch 大小线性增长,解码每步还要整读一遍,是长上下文推理的内存与带宽瓶颈。量化到 2-bit 能把每条缓存条目的存储压到 BF16 的 1/8,但这个精度的误差足以毁掉模型:Qwen3-8B 上 Hadamard 旋转配 2-bit,WikiText-2 困惑度从 9.72(BF16)飙到 258.29,基本不可用。
量化前先做一次可逆线性变换是标准补救。固定变换(Hadamard、随机正交)不看数据分布;学习型旋转(SpinQuant、FlatQuant、RotateKV)会校准,但并发工作 OSCAR 也只把变换约束在正交矩阵内。这篇的出发点是:缓存里的 key 只进 query-key 内积,value 只进输出投影,好变换得同时照顾缓存统计和这两条通路的误差敏感度;正交变换只能旋转,不能各向异性地缩放,这种平衡它做不到。ISTA 与 ETH Zürich 的团队(含 Dan Alistarh、Torsten Hoefler)把原本面向 weight-activation 量化的 WUSH 变换移植到 KV cache,放宽为一般可逆变换,并给出近最优性证明。
WUSH 是闭式构造:对要量化的张量取 Gram 矩阵(自身二阶统计),再取损失对该张量量化扰动的 Hessian,两者组合出变换 T;量化 T X 而非 X,读出时用 T^{-1} 还原。原构造对浮点格式可证最优、对整数格式渐近最优,且块对角、块宽等于量化组大小。
WUSH-KV 在每个 KV head 上学一对变换:
放置方式决定工程成败。value 变换离线折进 value 投影与输出投影两侧的权重,推理零开销。key 变换折不进去:折叠要求它与 RoPE 和 headwise 归一化可交换,联合约束下只剩符号翻转这种平凡解,起不到重分配坐标幅度的作用;放到 RoPE 之前虽保留全部自由度,但补偿项依赖缓存位置,每步解码都要对整条缓存做位置相关运算。权衡结果是 key 变换留在 RoPE 之后的在线路径上,每个新 key 和每个 query 各花一次 d×d 矩阵乘。
变换与量化器解耦,可搭配逐 token 裁剪量化器。理论针对 QuEST 整数量化器:在加性舍入噪声与尾部条件下,理想 WUSH 变换在全部「敏感度平衡」的可逆变换中近最优,期望损失比不超过 1+O(αb^{-2})。缓存管理沿用惯例:端到端评测中前 64 个位置(sink 窗口)与最近 256 个位置保 BF16,中间按页量化,窗口通常不足最大序列长度的 1%。校准只要前向:128 条 FineWeb-Edu 序列累积统计,Qwen3-8B 一次约 12 分钟(单张 L40S,峰值 39 GiB),变换跨 bitwidth、跨任务复用。
验证分三层。第一层固定 QuEST 量化器、只比变换:Qwen3-8B 全部 36 个 attention 模块,key/value 同时 2-bit,模块输出重建误差取几何平均:
| 变换 | 几何平均误差 |
| Hadamard | 0.325 |
| OSCAR | 0.309 |
| WUSH-KV | 0.208 |
| WUSH-A(attention-aware Hessian) | 0.195 |
第二层,Qwen3-8B 的 WikiText-2 困惑度(BF16 基线 9.72):
| 变换 | 4-bit | 3-bit | 2-bit |
| Hadamard | 9.77 | 10.06 | 258.29 |
| OSCAR | 9.77 | 10.19 | 13.74 |
| WUSH-KV | 9.73 | 9.82 | 10.51 |
bit 数越低差距越大:2-bit 上 WUSH-KV 比 OSCAR 低 3.23 个点,离 BF16 只差 0.79。Hadamard 的 2-bit 崩溃有明确机理:第一个 attention 模块的 key 在变换前就有集中在单一坐标的尖峰,Hadamard 把它摊成几乎等幅的系数,恰好落在相邻 2-bit 重建格点的中点附近。
第三层端到端:集成进推理框架 SGLang,用 OSCAR 的百分位裁剪仿射量化器,连裁剪比例都沿用 OSCAR 的调参、未为 WUSH 重调;模型为 Qwen3-4B-Thinking、8B、32B,任务为 AIME 2025、MATH-500、GPQA Diamond、LiveCodeBench v6,三个随机种子。8B 上 WUSH-KV 四项全部高于 OSCAR;OSCAR 大幅下滑处差距更明显,LiveCodeBench v6 在 8B 与 32B 上都是 WUSH-KV 明显占优;4B 与 32B 整体接近。固定变换撑不起 2-bit:Hadamard 在多个任务上频繁耗尽生成预算,TurboQuant 式实现在部分任务得分归零。长上下文 RULER NIAH 与 MRCR 上,最长测试长度处 WUSH-KV 准确率高于 OSCAR。代价是输出变长,平均至多约为 BF16 的 1.5 倍,仍远小于 8 倍的每条目存储缩减。
这回答的是「2-bit KV cache 能不能用于推理模型」。8 倍存储缩减直接换成更大的 batch 或更长的上下文;配方也务实:12 分钟一次性校准、变换复用、value 侧零开销、SGLang 里能跑。与学习型旋转那条线相比,闭式解只要前向统计,不碰反传。相对 OSCAR 的增量说得清楚:正交放宽到一般可逆,加上乘积感知的 Hessian。诚实地说,这是渐进改进,收益集中在 2-bit 与 8B;4-bit 下各家都接近无损,差异不大。
作者自认:key 侧稠密变换必须在线应用,每个 token 多出矩阵乘;吞吐与系统级评估未做。另有几点:近最优性证明只覆盖 QuEST 量化器,端到端用的却是 OSCAR 式仿射量化器,作者自己说后者不易套用同一证明,理论保证与落地配置脱节;证明针对无阻尼的理想变换,实验带阻尼。裁剪比例沿用 OSCAR 调参,未按 WUSH 变换后的分布重调,数字留有 headroom。OSCAR 官方分数未能完全复现,基线是同管线下的重跑,与并发工作的对比要打折扣。attention-aware Hessian(WUSH-A)只小幅降低重建误差,更精细敏感度建模的边际收益存疑。2-bit 困惑度 10.51 对 9.72,便宜不等于免费。