OSCAR:离线谱协方差感知旋转实现 2-bit KV 缓存量化 OSCAR: Offline Spectral Covariance-Aware Rotation for 2-bit KV Cache Quantization
用注意力感知协方差校准的2-bit KV缓存量化,接近BF16精度且吞吐最高7.83倍。
前置知识
KV 缓存 (Key-Value Cache)
自回归解码时每个 token 都需访问历史所有 token 的 Key 与 Value 矩阵,显存随序列长度、批大小、层数线性增长,是长上下文 LLM 推理的主要内存与带宽瓶颈。Qwen3-4B 在 128K 上下文下 KV 缓存可达数十 GB。
OSCAR 的全部动机是压缩 KV 缓存,理解 KV 缓存是显存(而非权重)主导的,才能体会 INT2 量化为何能在几乎不增加算力的情况下获得 8× 内存节省与 7× 吞吐提升。
对称 / 非对称均匀量化与剪裁 (Clipping)
b-bit 量化将 $x$ 映射到 $2^b$ 个离散码,对称形式 $Q(x)=s\cdot\mathrm{clip}(\lfloor x/s\rceil,-\tau,\tau)$;OSCAR 用 per-channel 非对称 INT2 + per-token 剪裁 $\tau$ 控制离群值。
INT2 只有 4 个量化级,极少数离群通道会吞掉整个 scale,导致大部分正常值被压到 1-2 个码,精度直接崩塌。OSCAR 的全部收益都建立在"通过旋转+剪裁为 INT2 创造良好输入分布"这一前提上。
Hadamard 旋转与 QuaRot/Saw-INT4 等旋转量化基线
正交旋转 $R$ 把激活变换到新基底 $x\to xR$,使原本集中在少数通道的离群值被均匀分散到所有通道,便于低比特量化;Hadamard 矩阵是计算极其友好的正交矩阵。QuaRot/Saw-INT4 等用 Hadamard 旋转,但它们的旋转目标都是"最小化原始缓存重建误差"。
OSCAR 论证了这种 data-oblivious 旋转的不足:它平滑了缓存张量,却不识别哪些方向对 attention 重要。理解这一缺口,才能明白 OSCAR 为何要从 query / score 协方差出发重新选择旋转。
注意力机制与 $S = \mathrm{softmax}(QK^\top/\sqrt{d})$
单头注意力 $S=\mathrm{softmax}_{\mathrm{row}}(QK^\top/\sqrt{d})$,输出 $O=SV$;Key 经 logit 被消费、Value 经 score-加权聚合被消费,二者对 $K,V$ 误差敏感方向不同,这是 OSCAR 旋转目标的关键。
论文核心论证:Key 量化误差经 $Q^\top Q$ 加权传递到 logit,Value 误差经 $S^\top S$ 加权传递到输出。因此 Key 应沿 query 协方差 $C_Q=Q^\top Q$ 的特征基旋转,Value 应沿 score-加权协方差 $C_S=V^\top S^\top SV$ 的特征基旋转。
Paged KV-Cache 与 SGLang 调度
PagedAttention 把 KV 切成固定大小页面,支持非连续显存与跨请求共享,vLLM/SGLang 采用;SGLang 进一步加入 radix prefix-cache 复用前缀 KV。OSCAR 必须保持 paged-cache 兼容才能上线。
OSCAR 强调"不是又一个研究原型"。它的工程价值在于把旋转+量化+INT2 内核塞进 SGLang 的 paged serving 流水线,使得 sink+recent 保留 BF16、中间历史走 INT2 路径而不破坏 prefix-cache 共享与 fused kernel。
研究动机
长上下文推理使 KV 缓存成为 LLM 服务的首要内存/带宽瓶颈。INT2 量化理论上能换 8× 内存,但现有方法在 2-bit 档几乎全部崩溃:Table 2 显示 naive INT2 与 QuaRot-INT2 在 Qwen3-4B/8B/32B/GLM-4.7-FP8 四个模型、AIME25/MATH500/HumanEval/LiveCodeBench/GPQA 五个推理编程基准上几乎全部降到 0,平均分相对 BF16 下降 60-75 分;TurboQuant(关 mixed precision) 在 4B 上掉 43.90、32B 上掉 66.29 分;Table 3 的 RULER-NIAH 显示 QuaRot-INT2 在 4k 就跌到 0,16k 后完全无法检索。
本文的目标是OSCAR 旨在用一种可落地的方案让 INT2 KV 缓存同时做到(a)在 4B-358B 多种规模、多种模型家族上接近 BF16 精度,(b)与 SGLang 的 paged KV cache、prefix cache、fused decode kernel 兼容,(c)在长上下文(100k+)和大批大小(BS=32)场景下给出可量化的吞吐与时延收益。具体目标是把 Qwen3-4B/8B 上相对 BF16 的平均精度差从 70+ 分压到 3 分以内,在 32B 和 358B 上做到与 BF16 实质持平或更优,并实现约 8× KV 内存压缩、batch=1 长上下文下 3× 解码加速、batch=32 下 7.83× 吞吐提升。
与已有工作不同的是,已有 INT2 方法共有的盲点是"选错了旋转目标":QuaRot/Saw-INT4/TurboQuant 用数据无关 Hadamard 或原始缓存重建目标 $K^\top K$ / $V^\top V$,而 attention 实际按 logit 与 score-weighted 聚合消费 KV,敏感方向由 $C_Q=Q^\top Q$ 与 $C_S=V^\top S^\top SV$ 决定。OSCAR 的切入角度是:离线校准估计 $C_Q$、$C_S$ 并做特征分解得到 attention-aware 旋转基 $U_Q$、$U_S$,与 Hadamard、bit-reversal 组合成 $R_K=U_Q H_{\mathrm{Had}} P_{\mathrm{br}}$、$R_V=U_S H_{\mathrm{Had}} P_{\mathrm{br}}$,让被量化的方向恰好是 attention 不读的方向,在 2.28 BPE 下同时击败所有 2.x-bit 旋转基线。
核心方法
OSCAR 的核心直觉是"对的方向做量化、不重要的方向承担误差"。离线阶段用 8k MMLU/GPQA token 做一次 teacher-forced 前向,dump 每层每头 Q/K/V 激活;对 Key 估计 query 协方差 $C_Q=\frac{1}{N}\sum_n q_n^\top q_n=U_Q\Lambda_Q U_Q^\top$,取 $U_Q$;对 Value 用 $S=\mathrm{softmax}_{\mathrm{row}}(QK^\top/\sqrt{d})$ 算 score-加权协方差 $C_S=\frac{1}{N}V^\top S^\top SV=U_S\Lambda_S U_S^\top$,取 $U_S$。在线合成 $R_K=U_Q H_{\mathrm{Had}} P_{\mathrm{br}}$、$R_V=U_S H_{\mathrm{Had}} P_{\mathrm{br}}$,per-layer 自适应剪裁 $(c_K,c_V)\approx(0.96,0.92)$ 配 per-channel 非对称 INT2 部署到 SGLang。
OSCAR 与 QuaRot/Saw-INT4/TurboQuant 的本质区别在于"用对角化哪个矩阵作为目标":已有方法最小化 $\|K-Q(K)\|_F^2$ 或 $\|V-Q(V)\|_F^2$,即让量化后缓存接近原缓存;OSCAR 从 attention 下游目标出发,最小化 logit 失真 $\|QK^\top-Q\hat K^\top\|_F^2=\mathrm{tr}((K-\hat K)Q^\top Q(K-\hat K)^\top)$与 attention-block 输出失真 $\|SV-S\hat V\|_F^2=\mathrm{tr}((V-\hat V)^\top S^\top S(V-\hat V))$。在 ambient-residual 对角假设下,这两个目标的最优旋转基正好是 $C_Q$、$C_S$ 的特征向量 $U_Q$、$U_S$,Key 沿 query 投影能量大的方向保持精细、Value 按 score 加权;这一从下游误差反推旋转基的视角是论文最系统的贡献。
方法步骤详情
方法分四步。(1) 离线校准:8k MMLU/GPQA token 在 SGLang 内 teacher-forced 前向一次,dump 每层每头 Q/K/V 激活。(2) 算协方差并构造旋转:对每 head 算 $C_Q=\frac{1}{N}\sum_n q_n^\top q_n$、$S=\mathrm{softmax}_{\mathrm{row}}(QK^\top/\sqrt{d})$、$C_S=\frac{1}{N}V^\top S^\top SV$,特征分解得 $U_Q$、$U_S$,合成 $R_K=U_Q H_{\mathrm{Had}} P_{\mathrm{br}}$、$R_V=U_S H_{\mathrm{Had}} P_{\mathrm{br}}$。(3) PTQ 估 per-layer 剪裁与 per-channel asymmetric INT2 scale/zero-point。(4) SGLang 集成:prefill fused Triton kernel 旋转+量化 K/V 为 2-bit;cache 保持 sink+recent BF16、中间 INT2(S=64,R=256);decode byte-unpack + dequant + 逆旋转;$R_V$ 吸进 $W_O$。
技术新颖性
新颖性体现在三层。(1) 形式化创新:首次把 Key/Value 量化目标矩阵显式拆成 $C_Q$、$C_S$,Theorem 1 证明 frozen-error surrogate 下 $U_Q$、$U_S$ 最优。(2) 组合旋转的实证必要性:Table 4 显示去掉 $U$(只留 Hadamard+Pbr)在 Qwen3-8B 上从 70.01 掉到 32.82,去掉 Hadamard 掉到 51.74,换 raw $K^\top K$/$V^\top V$ PCA 掉到 31.12,说明"对哪个矩阵做 PCA"比"是否做 PCA"重要。(3) 工程完整性:SGLang 兼容 fused Triton kernel(per-step 8.0ms vs BF16 9.9ms)、paged-cache layout、$R_V$ 吸进 $W_O$、mixed-precision sink/recent;Table 8 显示 BS=64/128 INT2 仍 attn 主导,23.4/39.3ms 接近 BF16 在 BS=1/2 时的容量边界,工程可部署性是区别于同期研究的关键。
实验结果
(1) 主精度(Table 2):OSCAR 是唯一在 4B-358B 范围均接近 BF16 的近 2-bit 方法——Qwen3-4B 71.86(差距 3.78)、Qwen3-8B 69.42(1.42)、Qwen3-32B 74.17(-0.02,几乎打平)、GLM-4.7-FP8 78.16(+0.27 反而略好);QuaRot-INT2 1.40/10.14/7.90/75.14,TurboQuant 31.74/23.13/17.60/77.95。(2) 长上下文(Table 3):4B 上 OSCAR 4k→128k 取得 99.7→39.5,QuaRot-INT2 在 4k 即 0.0;8B 全程 99.5→45.0;GLM 128k 仍 97.2 与 BF16 持平。(3) 系统吞吐(图 4):batch=1 4B 在 30k/60k/100k 加速 1.98×/2.52×/3.08×;batch=32 100k GLM-4.7 取得 7.83× 吞吐,大幅领先 Saw-INT4 的 3.44×。(4) 服务级(Table 9):32 并发 8k/1k,Qwen3-4B OSCAR 63.3 tok/s/user、1723.9 tok/s/GPU、精度 71.86。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| Qwen3-4B-Thinking-2507 五基准平均精度 (GPQA/HumanEval/LCB v6/AIME25/MATH500) | 平均分 (5 seeds mean) | 71.86(BPE 2.28,GPQA 64.95,HumanE 92.24,LCB 45.38,AIME 64.00,MATH 92.75) | BF16 75.64;QuaRot-INT2 1.40;TurboQuant 31.74;Saw-INT4 73.11 | 相对 QuaRot-INT2 提升 +70.46 分;相对 TurboQuant 提升 +40.12 分;距 BF16 仅 -3.78 |
| Qwen3-8B 五基准平均精度 | 平均分 (5 seeds mean) | 69.42(BPE 2.28,GPQA 55.05,HumanE 87.88,LCB 46.32,AIME 66.67,MATH 92.22) | BF16 70.84;QuaRot-INT2 10.14;TurboQuant 23.13;Saw-INT4 69.97 | 相对 QuaRot-INT2 提升 +59.28 分;距 BF16 仅 -1.42;略优于 Saw-INT4 4-bit |
| Qwen3-32B 五基准平均精度 | 平均分 (5 seeds mean) | 74.17(BPE 2.28,GPQA 60.40,HumanE 90.12,LCB 53.57,AIME 74.00,MATH 92.75) | BF16 74.19;QuaRot-INT2 7.90;TurboQuant 17.60;Saw-INT4 74.43 | 相对 BF16 -0.02(实质持平);相对 QuaRot-INT2 提升 +66.27 |
| GLM-4.7-FP8 (358B) 五基准平均精度 | 平均分 (3 runs mean) | 78.16(BPE 2.28,GPQA 73.57,HumanE 91.06,LCB 52.63,AIME 78.89,MATH 94.66) | BF16 77.89;QuaRot-INT2 75.14;TurboQuant 77.95;Saw-INT4 78.15 | 相对 BF16 +0.27(反超);相对 QuaRot-INT2 提升 +3.02;BPE 比 Saw-INT4 的 4.25 低 47% |
| RULER-NIAH 128k 长上下文检索 | 检索准确率 (3 seeds mean) | Qwen3-4B 39.5;Qwen3-8B 45.0;GLM-4.7-FP8 97.2 | Qwen3-4B BF16 81.0 / QuaRot-INT2 0.0;Qwen3-8B BF16 78.2 / QuaRot-INT2 0.0;GLM BF16 97.2 / QuaRot 96.3 | 4B/8B 上 OSCAR 在 16-128k 区间显著优于 QuaRot-INT2;GLM 上 OSCAR 与 BF16 持平并略胜 QuaRot |
| Decode 吞吐 100k 输入 BS=1 加速 (vs BF16) | speedup × | Qwen3-4B 3.08×;Qwen3-8B 2.88×;GLM-4.7-FP8 2.83× | Saw-INT4 同条件约 1.5-2.0×;QuaRot-INT2 接近 OSCAR 但精度不可用 | 100k 上下文时比 Saw-INT4 再快 1.5-2×,逼近 QuaRot-INT2 的吞吐上限同时保住精度 |
| 大批量吞吐 100k 输入 BS=32 | speedup × | Qwen3-4B 3.35×;Qwen3-8B 3.08×;GLM-4.7-FP8 7.83× | Qwen3-4B BF16 1.00×;Qwen3-8B BF16 1.00×;GLM BF16 1.00×;Saw-INT4 4B 2.43×/8B 2.36×/GLM 3.63× | GLM 上 OSCAR 较 BF16 提升 7.83×,较 Saw-INT4 再提升 2.16× |
局限与改进
作者明示的局限:(1)calibration-derived clip(0.96,0.92) 0.15 分内接近最优但仍依赖一次 offline pass,无 zero-calibration 路径;(2)只在 2-bit INT 上做实验,未推广到 INT3/INT4 或 mixed-precision 联合搜索;(3)prefix-cache 优势(图 6)只覆盖 Qwen3-4B 与 GLM-4.7,8B/32B 的 prefix-cache sweep 未充分报告。我观察的隐含限制:$R_K,R_V$ 是 fixed per layer,长程 query/attention 分布可能漂移(尤其 chain-of-thought reasoning),分布外鲁棒性无 ablation;sink+recent BF16 隐含假设 SGLang radix cache 边界与 OSCAR 对齐,prefix < S+R 时浪费 BF16 配额;未量化 INT2 内核 attention score 数值精度(RTL floor),只测端到端;5 个基准均为英文,中文与多语种、长 code generation 在 32K 截断下未被覆盖。
独立分析的弱点
(1) 校准数据域依赖:Table 7 显示 MMLU 2k/8k/16k/32k 平均 68.28/69.15/70.88/69.90,GPQA 8k 71.01,WikiText 8k 69.50,同 8k token 域差 1.86 分,改进方向是多域混合或在线 EMA 更新 $C_Q,C_S$。(2) mixed-precision 边界不灵活:(S,R)=(64,256) 是网格最优但未与 context length/batch size/模型深度联合搜索,改进方向是 page 级动态选 BF16/INT2。(3) 小 batch attention fused 优势被高估:Table 8 显示 Qwen3-8B BS=1 attention 段从 3.8ms 降到 1.3ms,total step 仅从 9.9ms 降到 8.0ms,主因 GEMM 5.6ms 未变,改进方向是结合 INT2 权重/GEMV 量化。(4) rotation 矩阵吸收只对 $R_V$ 做了,未对 $R_K$ 吸收,改进方向是 $W_Q$、$W_K$ 两端都吸收并讨论对 RoPE 的兼容性。
未来方向
作者暗示方向:(1)把 covariance-aware 框架推广到 INT3/INT4 乃至 per-layer 比特分配,在 2.5-3.0 BPE 甜点区再换一档精度;(2)把 $R_K$、$R_V$ 视为可学习参数而非 PCA 固定基,允许少量微调进一步精调;(3)与 prefix-cache 共享更深耦合,做"仅共享 INT2 history、每请求独享 sink/recent"的细粒度 sharing。延伸方向:第一,与 outlier-aware 微调(QLoRA/AffineQuant) 结合把精度推到 BF16+0.5 分以内;第二,迁移到 MLA(DeepSeek 系列)验证对压缩 KV latent 量化是否同样有效;第三,在 reasoning 模型上研究长 chain-of-thought 是否让 $C_Q,C_S$ 分布漂移,设计 dynamic rotation 或 sliding-window calibration refresh;第四,扩展到 encoder-decoder / 多模态 KV cache;第五,把 fused INT2 decode kernel 进一步融合进 prefill 实现 prefill+decode 全 INT2 一体化。
复现评估
代码层面,论文链接了 Code + Website + RotationZoo 三个仓库,给出完整 calibration 脚本与 SGLang 集成补丁,作者承诺开源;数据层面,校准仅需 8k MMLU/GPQA token,benchmark 使用 AIME25、GPQA-Diamond、HumanEval、LiveCodeBench v6、MATH500、RULER-NIAH 均为公开基准;算力是主要门槛——Qwen3-4B/8B 单卡 H100 80GB 可复现 INT2 精度,32B/358B 与长上下文(>64k)分别需 2×H100、8×H100 张量并行,完整复现所有表约 300-500 H100-hour。复现难度:(a)修改 SGLang 源码插入 fused INT2 attention kernel 与 mixed-precision paged layout,熟悉 SGLang 调度器与 radix-cache 机制;(b)calibration dump 每层每头 Q/K/V 激活,8k token × 36-80 层体积可观,需预留 ~50GB 磁盘;(c)GLM-4.7-FP8 在 8×H100 上跑一次 32K generation 数小时,加 5 seed × 3 prefix-cache regime,3 runs 共 45 次端到端,小算力下完整表格不可达。
论文图表
四联子图,横轴为 layer 0-35:(a)K/V cache 相对 MSE,naive INT2 与 Clip-only 在 100 量级,Hadamard 在 10^0,OSCAR 在 10^{-1} 附近;(b)FP16 vs 量化 attention 分布的 KL 散度,naive/Hadamard 在 10^{-2}-10^{-1} 区间,OSCAR 普遍低 1-2 个数量级;(c)attention block 输出相对 MSE,OSCAR 维持 10^{-1} 而 baseline 漂到 10^1;(d)跨层 hidden state 累积 MSE,OSCAR 增速明显慢于其它方案。
直接展示'每一步都把 OSCAR vs 其它' 拆开来的误差传播,论证 OSCAR 的精度优势源自 attention-aware 旋转在每一层、每一种误差度上都更小,而不是某一处的局部 trick。这是论文 motivation 的最强可视化证据。