OScaR:极端 KV 缓存量化的奥卡姆剃刀

笔记日期: 2026-06-17 笔记作者: Zhongzhu Zhou 论文标题: OScaR: The Occam’s Razor for Extreme KV Cache Quantization in LLMs and Beyond 论文作者: Zunhai Su, Rui Yang, Chao Zhang 等 arXiv: https://arxiv.org/abs/2605.19660 状态/期刊: 预印本,arXiv 2026年5月

一句话总结

OScaR 发现 KV 缓存 INT2 量化失败的真正元凶——令牌范数不平衡(Token Norm Imbalance,TNI)——并以两个相互不可缺少的外科式干预(Hadamard 通道旋转 + ℓ2 令牌归一化)在无需训练的条件下把 INT2 精度推到逼近 16bit 基线,同时实现 5.3× 内存压缩与 4.1× 吞吐提升。

前置知识

在正式进入论文之前,先把几个必要的背景概念梳理清楚,这些概念将贯穿全文。

KV 缓存的工作原理

Transformer 的自注意力机制在推理时面临一个重复计算问题。对于已经处理过的每个历史 token,每次生成新 token 都需要重新计算其对应的 Key 向量和 Value 向量。KV 缓存(KV Cache)的核心思想很直接:把算过的结果缓存下来,下次直接取用

具体来说,对于序列中第 tt 个位置的 token,其 Key 向量和 Value 向量在第一次计算后被存入缓存:

KV Cache={(K1,V1),(K2,V2),,(Kt,Vt)}\text{KV Cache} = \{(K_1, V_1), (K_2, V_2), \ldots, (K_t, V_t)\}

在解码阶段每生成一个新 token,注意力计算变为:

Attention(Qnew,Kcache,Vcache)=softmax(QnewKcacheTdk)Vcache\text{Attention}(Q_{\text{new}}, K_{\text{cache}}, V_{\text{cache}}) = \text{softmax}\left(\frac{Q_{\text{new}} K_{\text{cache}}^T}{\sqrt{d_k}}\right) V_{\text{cache}}

这里 QnewQ_{\text{new}} 是新 token 的 Query,KcacheK_{\text{cache}}VcacheV_{\text{cache}} 是历史所有 token 的缓存。问题在于缓存大小随序列长度线性增长:序列长度翻倍,缓存也翻倍。对于长上下文场景(比如 128K token),KV 缓存会成为内存瓶颈。

graph LR
    subgraph "解码步骤 t"
        Q_new["新 Query\nQ_new ∈ R^{1×d}"]
        KCache["K Cache\nK ∈ R^{t×d}"]
        VCache["V Cache\nV ∈ R^{t×d}"]
        Attn["注意力计算\nsoftmax(QK^T/√d)V"]
        Out["输出向量\n∈ R^{1×d}"]
        Q_new --> Attn
        KCache --> Attn
        VCache --> Attn
        Attn --> Out
    end
    subgraph "新 token 入缓存"
        NewK["新 K_{t+1}"]
        NewV["新 V_{t+1}"]
        Out --> NewK
        Out --> NewV
    end
    style KCache fill:#ffe4b5
    style VCache fill:#ffe4b5

整数量化基础

整数量化把浮点数映射到低比特整数(比如 INT2,即 2 bit,只有 4 个离散值:0, 1, 2, 3),以减少内存占用和带宽。

均匀量化的基本公式是:

Q(x)=clamp ⁣(xΔ+z,0,2b1)Q(x) = \text{clamp}\!\left(\left\lfloor \frac{x}{\Delta} + z \right\rceil, \, 0, \, 2^b - 1\right)

其中 Δ\Delta 是量化步长(scale),zz 是零点(zero-point),bb 是比特数,\lfloor \cdot \rceil 表示取整。

反量化(从整数恢复近似浮点值):

x^=Δ(Q(x)z)\hat{x} = \Delta \cdot (Q(x) - z)

量化误差的大小直接由步长 Δ\Delta 决定——步长越大,相邻整数代表的浮点值间隔越大,每个浮点值被”近似”时的误差上界也就越大:

xx^Δ2|x - \hat{x}| \leq \frac{\Delta}{2}

步长的确定方式: 对于一组数 {x1,x2,,xn}\{x_1, x_2, \ldots, x_n\},最简单的做法是用极值法:

Δ=maxiximinixi2b1\Delta = \frac{\max_i x_i - \min_i x_i}{2^b - 1}

因此,如果这组数的动态范围(max−min)很大,步长就会大,量化精度就差。这是后面 TNI 问题的关键所在。

Hadamard 变换简介

Hadamard 矩阵 HdRd×dH_d \in \mathbb{R}^{d \times d}dd 是 2 的幂次)的每个元素都是 ±1/d\pm 1 / \sqrt{d},是一个正交矩阵HHT=IH H^T = I)。Fast Hadamard Transform(FHT)可以在 O(dlogd)O(d \log d) 时间内计算矩阵-向量乘法 HvH \mathbf{v},比通用矩阵乘法快得多。

Hadamard 变换的两个关键性质:

  1. 能量均匀分散:如果原向量某几个维度有特别大的值(“离群值”),经过 FHT 后这些能量被均匀分散到所有维度,单个维度的极端值消失。

  2. 正交性保持内积:对于任意两个向量 q\mathbf{q}k\mathbf{k}

(Hq)T(Hk)=qTHTHk=qTk(H\mathbf{q})^T (H\mathbf{k}) = \mathbf{q}^T H^T H \mathbf{k} = \mathbf{q}^T \mathbf{k}

正交变换不改变向量间的内积(即注意力得分),这是 OScaR 应用到 Key-Query 对时能保持注意力计算正确的根本原因。

Attention Sink(注意力汇集令牌)

在 LLM 的注意力机制中,人们观察到一个有趣的现象:少数特定位置(通常是序列开头的几个 token,比如 BOS 或第一个实际 token)会被几乎所有其他 token 的注意力分配到较大权重,即使这些位置在语义上并不特别重要。这类 token 被称为 Attention Sink(注意力汇集)。

Attention Sink 的存在与 Softmax 的归一化性质有关:Softmax 要求所有权重之和为 1,当模型认为当前 token 不需要关注序列中任何特定位置时,注意力权重必须”流向”某处——Attention Sink 扮演了这个”垃圾桶”角色。

Attention Sink 的 Key 向量有一个显著特征:其 ℓ2 范数比普通 token 低得多。这个范数不均衡现象正是 TNI 问题的来源。

问题:为什么 KV 缓存内存很重要

内存瓶颈的规模

以 Llama-3.1-70B 模型为例,在 BF16 精度下服务一个 128K token 的请求:

  • 模型权重约 140 GB
  • KV 缓存 ≈ 2(K+V)× 128000(序列长度)× 8192(隐藏维度)× 80(层数)× 2(BF16 字节数) ≈ 337 GB

KV 缓存已经超过了模型本身的大小!即便在短上下文(如 4K token)场景下,当批大小较大时,KV 缓存也会占用大量 HBM 内存,限制可并发处理的请求数量(即吞吐量)。

内存、延迟、吞吐三角关系

graph TD
    KVSize["KV 缓存大小\n(随序列×批大小增长)"]
    
    KVSize -->|内存占用高| BatchSize["批大小受限\n→ 吞吐量低"]
    KVSize -->|带宽占用大| Latency["解码延迟高\n(每步需加载全部 KV)"]
    KVSize -->|内存不足| OOM["OOM / 无法处理长上下文"]
    
    Quant["INT2 量化\n8× 内存减少"]
    Quant -->|释放内存| BatchSize2["更大批大小\n→ 更高吞吐"]
    Quant -->|减少带宽| Latency2["更低延迟"]
    Quant -->|内存充裕| LongCtx["可处理更长上下文"]
    
    style KVSize fill:#ffcccc
    style Quant fill:#ccffcc

INT2 量化把每个浮点数从 16bit 压缩到 2bit,内存减少 。但精度代价是致命的——直接对 KV 缓存做 INT2 量化会导致严重的精度崩溃。OScaR 的使命就是解决这个”精度崩溃”问题。

背景:KIVI 与逐通道量化范式

KIVI 的核心设计

KIVI 是目前 KV 缓存量化的主流基础框架,其核心设计决策包括:

Key 缓存逐通道量化,Value 缓存逐 token 量化。 这个不对称设计来自一个经验观察:Key 张量的离群值(outlier)主要集中在特定通道(feature dimension),而 Value 张量的分布则在各 token 之间相对均匀。

设 Key 缓存 KRS×dK \in \mathbb{R}^{S \times d},其中 SS 是序列长度,dd 是每头的隐藏维度。KIVI 按通道方向分块,块大小为 GG(通常 G=32G=32G=64G=64)。对于第 jj 个通道的第 gg 个块(包含该块中的所有 token 位置),量化参数计算如下:

Δj,g=maxiblockKi,jminiblockKi,j2b1\Delta_{j,g} = \frac{\max_{i \in \text{block}} K_{i,j} - \min_{i \in \text{block}} K_{i,j}}{2^b - 1} zj,g=miniblockKi,jΔj,gz_{j,g} = -\frac{\min_{i \in \text{block}} K_{i,j}}{\Delta_{j,g}}

每个元素的量化和反量化:

Q(Ki,j)=clamp ⁣(Ki,jΔj,g+zj,g,0,2b1)Q(K_{i,j}) = \text{clamp}\!\left(\left\lfloor \frac{K_{i,j}}{\Delta_{j,g}} + z_{j,g} \right\rceil, \, 0, \, 2^b - 1\right) K^i,j=Δj,g(Q(Ki,j)zj,g)\hat{K}_{i,j} = \Delta_{j,g} \cdot \left(Q(K_{i,j}) - z_{j,g}\right)

Value 缓存 VRS×dV \in \mathbb{R}^{S \times d} 则沿 token 维度分块做逐 token 量化(类似但方向相反)。

高精度残差缓冲

KIVI 还保留一个大小为 R=128R=128 个 token 的高精度残差缓冲(residual window):最新的 RR 个 token 的 KV 保持 BF16 精度,满了再批量量化入低精度缓存。这个设计缓解了新 token 量化误差导致生成错误的问题,但也意味着无论如何都需要一块 BF16 内存。

逐通道量化在 INT2 下的困境

在 INT4 甚至 INT3 精度下,KIVI 已经能取得不错的结果。但在 INT2 时,精度急剧下降,甚至完全崩溃。为什么?

直觉上:INT2 只有 4 个离散值(0, 1, 2, 3),步长 Δ\Delta 必须覆盖整个动态范围。如果块内存在范数差异极大的 token——有些 token 的值很大,有些很小——那么 maxmin\max - \min 会被最大值拉得很宽,步长 Δ\Delta 就会很大,低范数 token 的精度就会极差,甚至被量化到同一个整数值(精度彻底丢失)。

但 KIVI 的作者们只是把这归结为”INT2 太极端”,并没有深入分析根本原因。OScaR 的作者们则做了一步关键的诊断——他们发现,这一切背后有一个统一的机制:TNI。

TNI 诊断:令牌范数不平衡

定义与测量

对于注意力层的 Key 张量,设第 tt 个 token 在第 hh 个注意力头的 Key 向量为 kt,hRdk\mathbf{k}_{t,h} \in \mathbb{R}^{d_k},其 ℓ2 范数为:

kt,h2=j=1dkkt,h,j2\|\mathbf{k}_{t,h}\|_2 = \sqrt{\sum_{j=1}^{d_k} k_{t,h,j}^2}

Token Norm Imbalance(TNI) 就是 token 与 token 之间范数差异巨大的现象。具体来说,定义 token tt 在所有头上的范数集合:

Nt(K)={kt,h2  |  h=1,,H}\mathcal{N}_t^{(K)} = \left\{\|\mathbf{k}_{t,h}\|_2 \;\middle|\; h = 1, \ldots, H\right\}

如果不同 token 的中位数范数差异超过一个数量级,就存在严重的 TNI 问题。

实证观测

作者在 Llama-3.1-8B 和 Qwen3-8B 上系统测量了 Key 向量的 token 范数分布,发现:

  1. 少数 token(Attention Sink)的范数极低,比普通 token 低 10×~50×
  2. 这种范数不平衡在所有注意力头中普遍存在,不是个别头的特殊行为
  3. Value 向量不存在这种问题——Value 的 token 范数分布相对均匀,这解释了为什么 Value 用逐 token 量化效果好
xychart-beta
    title "Key Token 范数分布示意(各 token 的范数中位数)"
    x-axis ["Sink-1", "Sink-2", "Token-10", "Token-50", "Token-100", "Token-200", "Token-500"]
    y-axis "ℓ2 范数(相对单位)" 0 --> 10
    bar [0.3, 0.4, 7.2, 8.1, 7.8, 8.3, 8.0]

上图示意了典型 LLM 中各 token Key 向量范数的分布:Attention Sink token(Sink-1, Sink-2)的范数极低(~0.3-0.4),而普通 token 的范数聚集在较高值(~7-8)。这种极端的范数差异正是 INT2 量化崩溃的根源。

理论分析:TNI 如何破坏逐通道量化

设某个量化块包含 GG 个 token,第 jj 个通道的值为 {K1,j,K2,j,,KG,j}\{K_{1,j}, K_{2,j}, \ldots, K_{G,j}\}。该块的量化步长为:

Δj,g=maxiGKi,jminiGKi,j2b1\Delta_{j,g} = \frac{\max_{i \leq G} K_{i,j} - \min_{i \leq G} K_{i,j}}{2^b - 1}

考虑块内存在一个高范数 token thight_{\text{high}}kthigh28\|\mathbf{k}_{t_{\text{high}}}\|_2 \approx 8)和一个低范数 Attention Sink tsinkt_{\text{sink}}ktsink20.3\|\mathbf{k}_{t_{\text{sink}}}\|_2 \approx 0.3)。

在通道 jj 上,高范数 token 的值可能达到 Kthigh,j+3.0K_{t_{\text{high}}, j} \approx +3.0,低范数 token 的值接近 0。这导致:

Δj,g3.0(3.0)221=6.03=2.0\Delta_{j,g} \approx \frac{3.0 - (-3.0)}{2^2 - 1} = \frac{6.0}{3} = 2.0

步长 Δ=2.0\Delta = 2.0,量化误差上界为 Δ/2=1.0\Delta/2 = 1.0。对于那些值本身只有 0.10.1 量级的低范数 token 而言,1.0 的误差是灾难性的——信噪比约为 1:10,信息几乎全部丢失。

论文附录 G 的定理(公式 11) 给出了严格的上界:逐通道量化块的重建均方误差(MSE)与块内 token 范数极差(max-min)的平方成正比:

E[KblockK^blockF2](maxtkt2mintkt2)2\mathbb{E}\left[\|\mathbf{K}_{\text{block}} - \hat{\mathbf{K}}_{\text{block}}\|_F^2\right] \propto \left(\max_t \|\mathbf{k}_t\|_2 - \min_t \|\mathbf{k}_t\|_2\right)^2

TNI 系统性地放大了这个极差,从而系统性地放大了量化误差。这是 INT2 精度崩溃的理论根源。

那么,TNI 是偶然的吗?

不是。Attention Sink 的存在有其内在原因:Softmax 注意力要求权重归一化,当模型认为无需聚焦任何特定位置时,权重需要流向某个”垃圾桶”。模型在训练中学会让某些 token(通常是 BOS 或位置 0)承担这个角色,这些 token 的 Key 向量范数就会系统性地偏低——这是 LLM 注意力机制的一种内在特性,不是训练数据的偶然噪声。

为什么直接逐 Token 缩放会失败

直觉方案及其问题

看到”token 范数不平衡”,最直觉的修复方案就是:在量化前把每个 token 的 Key 向量归一化到相同范数,量化后再乘回原始范数。这个方案叫Per-Token Scaling(逐 token 缩放)

形式上:

k~t=ktkt2,量化k~t,反量化后乘回kt2\tilde{\mathbf{k}}_t = \frac{\mathbf{k}_t}{\|\mathbf{k}_t\|_2}, \quad \text{量化} \tilde{\mathbf{k}}_t, \quad \text{反量化后乘回} \|\mathbf{k}_t\|_2

这样块内所有 token 的范数都是 1,TNI 消失,量化步长理应减小。

但实验显示:直接逐 token 缩放反而使精度更差! 这个违直觉的结论需要解释。

缩放导致的离群值伪影(Scaling-Induced Outlier Artifact)

问题的关键在于:归一化操作本身会制造新的离群值

考虑一个 Attention Sink token,其 Key 向量范数极低(比如 ksink2=0.3\|\mathbf{k}_{\text{sink}}\|_2 = 0.3),但在通道空间的分布是不均匀的——某些通道上值接近 0,某些通道上有中等幅度的值。

经过缩放后,所有通道都被放大了 1/0.33.3×1/0.3 \approx 3.3×

  • 原来中等幅度的通道值:0.51.670.5 \to 1.67
  • 原来接近 0 的通道值:0.010.030.01 \to 0.03(几乎没变)

现在看普通 token 在同一通道上的值:

  • 普通 token 的范数 8\approx 8,该通道值 0.02\approx 0.02(因为能量集中在其他通道)
  • 缩放后:0.02/80.00250.02 / 8 \approx 0.0025(几乎 0)

对比:

  • Attention Sink 缩放后的该通道值:0.01/0.30.0330.01 / 0.3 \approx 0.033
  • 普通 token 缩放后的该通道值:0.02/80.00250.02 / 8 \approx 0.0025

Attention Sink 缩放后,在普通 token 几乎为零的通道上产生了相对较大的值——这就是”缩放导致的离群值伪影”。这些人工伪影拉大了通道方向的动态范围,使得逐通道量化步长反而更大了!

graph LR
    subgraph "直接缩放的悖论"
        direction TB
        A["Attention Sink\n范数极低(0.3)\n某通道值 0.01"]
        B["归一化 ÷ 0.3\n放大 3.3×\n→ 该通道值 0.033"]
        C["普通 token\n范数正常(8.0)\n同通道值 0.02"]
        D["归一化 ÷ 8.0\n缩小 8×\n→ 该通道值 0.0025"]
        E["逐通道量化块内\n动态范围:0.033 - 0.0025\n比未缩放时更大!"]
        A --> B
        C --> D
        B --> E
        D --> E
    end
    style E fill:#ffcccc

核心矛盾: 逐 token 缩放虽然消除了 token 间范数不均,却在通道方向引入了新的不均匀性(低范数 token 被过度放大,制造了通道方向的离群值)。量化精度仍然被离群值主导,只是换了一种形式。

正确的修复顺序

这个失败案例揭示了一个重要原则:必须先消除通道方向的离群值,再做 token 方向的缩放。如果通道方向已经均匀(没有某个通道特别大),那么缩放低范数 token 也不会在通道方向制造新的伪影。

这正是 OScaR 的设计逻辑:

  1. 第一步(通道旋转):用 Hadamard 变换消除通道方向离群值
  2. 第二步(token 缩放):在通道已均匀的基础上做 ℓ2 归一化消除 TNI

两个步骤缺一不可,顺序不可颠倒。

OScaR 框架:通道旋转 + 全方位令牌缩放

整体架构概览

flowchart TD
    subgraph "OScaR 推理流程"
        Input["输入 token embedding\nx ∈ R^d"]
        
        WK["K 投影\n(离线融合了 Hadamard)\nW_K·x → k̃_raw"]
        WQ["Q 投影\n(离线融合了 Hadamard)\nW_Q·x → q̃"]
        WV["V 投影\n(离线融合了 Hadamard)\nW_V·x → ṽ"]
        
        FHT_K["在线 FHT\nH·k̃_raw → H·k"]
        FHT_Q["在线 FHT\nH·q̃ → H·q"]
        
        Scale["OTS:计算范数缩放因子\ns_t = ‖H·k_t‖₂\nk̃_t = H·k_t / s_t"]
        
        Quant_K["INT2 逐通道量化\nQ(k̃_t) + 存储 s_t"]
        Quant_V["INT2 逐 token 量化\nQ(ṽ_t)"]
        
        Cache["KV 缓存\n(低精度整数 + 标量 s_t)"]
        
        Attn["注意力计算\nQ̃·(s_t·反量化(K̃))^T → softmax → V̂"]
        
        Input --> WK & WQ & WV
        WK --> FHT_K --> Scale --> Quant_K --> Cache
        WQ --> FHT_Q
        WV --> Quant_V --> Cache
        FHT_Q --> Attn
        Cache --> Attn
    end
    style Cache fill:#ffe4b5
    style Scale fill:#ccffcc

步骤一:Canalized Rotation(通道旋转,CR)

目标: 消除 Key/Value 向量中通道方向的离群值。

操作: 对每个 Key 向量应用 Fast Hadamard Transform(FHT):

k~=Hk\tilde{\mathbf{k}} = H \mathbf{k}

其中 HRd×dH \in \mathbb{R}^{d \times d} 是归一化 Hadamard 矩阵,Hij{+1/d,1/d}H_{ij} \in \{+1/\sqrt{d}, -1/\sqrt{d}\},满足 HHT=IH H^T = I

为什么有效:

Hadamard 变换把原向量的能量均匀地”混合”到所有 dd 个维度。原来集中在某几个通道上的大值被分散到所有通道,每个通道的值变为:

k~j=1di=1dHjikik2d\tilde{k}_j = \frac{1}{\sqrt{d}} \sum_{i=1}^{d} H_{ji}^{\prime} k_i \approx \frac{\|\mathbf{k}\|_2}{\sqrt{d}}

(直觉上,能量均匀分配后每个维度的值约为总范数除以 d\sqrt{d}。)

旋转后,通道方向的动态范围从”一个通道很大,其他很小”变为”所有通道差不多大”,逐通道量化的效率大幅提升。

注意力计算的正确性保持:

由于 Hadamard 变换的正交性,对 Query 也同样应用 FHT 后,注意力得分保持不变:

(Hq)T(Hk)=qTHTHk=qTIk=qTk(H\mathbf{q})^T (H\mathbf{k}) = \mathbf{q}^T H^T H \mathbf{k} = \mathbf{q}^T I \mathbf{k} = \mathbf{q}^T \mathbf{k}

实现上的融合优化:

在线 FHT 对延迟有一定影响。OScaR 采用了一个工程优化:

  • Query 和 Key 的 FHT: 在线应用(必须,因为每个 token 实时生成)
  • Value 的 FHT: 离线,通过修改 Value 投影矩阵 WVW_V 来融入(即用 HWVH W_V 替换 WVW_V,推理时正常乘法自动包含了 Hadamard 变换)

Value 的 FHT 融合进权重完全不增加推理开销。

边界条件: FHT 要求维度 dd 是 2 的幂次。对于不满足此条件的模型,可以 padding 到最近的 2 的幂次,但论文未详细讨论这个 edge case。实际上,主流 LLM(Llama 系列、Qwen 系列)的注意力头维度通常是 64、128、256,满足条件。

步骤二:Omni-Token Scaling(全方位令牌缩放,OTS)

目标: 消除 token 间范数不平衡(TNI)。

前提: 通道方向的离群值已被 CR 消除,现在可以安全地做 token 方向的归一化。

操作: 对 Hadamard 变换后的 Key 向量计算 ℓ2 范数并归一化:

st=Hkt2s_t = \|H \mathbf{k}_t\|_2 k~t=Hktst\tilde{\mathbf{k}}_t = \frac{H \mathbf{k}_t}{s_t}

归一化后,所有 token 的范数统一为 1,TNI 消失。量化时只对 k~t\tilde{\mathbf{k}}_t 做逐通道 INT2 量化,同时将缩放因子 sts_t(一个标量)与量化后的整数一起存储在缓存中。

反量化时的恢复:

k^t=stk~^t=stΔj,g(Q(k~t,j)zj,g)\hat{\mathbf{k}}_t = s_t \cdot \hat{\tilde{\mathbf{k}}}_t = s_t \cdot \Delta_{j,g} \cdot \left(Q(\tilde{k}_{t,j}) - z_{j,g}\right)

存储开销分析:

每个 token 额外存储一个 BF16 标量(缩放因子 sts_t),占用 2 字节。对比:

  • 未使用 OTS 的 INT2 Key:d×2 bit=d/4d \times 2\text{ bit} = d/4 字节(d=128d=128 时为 32 字节)
  • 使用 OTS 的 INT2 Key:32+2=3432 + 2 = 34 字节

开销约 6.25%,远小于恢复带来的精度收益。

为什么在 CR 之后 OTS 能正确工作:

旋转后,每个通道的值量级大约是 kt2/d\|\mathbf{k}_t\|_2 / \sqrt{d}(均匀分散)。归一化后每个通道的值量级变为 1/d1/\sqrt{d}。由于各通道量级一致,Attention Sink(低范数)被放大后也不会在某些通道上产生特别大的值——因为没有”某个通道特别小使得放大后很大”的现象了。旋转消除了这种通道不均匀性,是 OTS 能正常工作的前提。

完整 OScaR 算法(伪代码)

Algorithm 1: OScaR KV 缓存量化

离线准备阶段:
  1. 计算 Hadamard 矩阵 H ∈ R^{d×d}
  2. 将 H 融合进 Value 投影矩阵:W_V ← H · W_V
  3. (可选)将 H 融合进 W_K, W_Q 中(减少在线开销)

推理阶段(每个注意力层,每个新 token t):

--- Key 处理 ---
  4. 计算 Key 向量:k_t = W_K · x_t   [原始 Key,R^d]
  5. 在线 FHT:k̃_t = H · k_t          [旋转后 Key,通道均匀]
  6. 计算缩放因子:s_t = ‖k̃_t‖₂      [标量]
  7. 归一化:k̃_t ← k̃_t / s_t         [范数=1,消除 TNI]
  8. 将 k̃_t 加入残差缓冲(BF16,大小 R=128)
  9. 若残差缓冲满:
     a. 对缓冲中的 k̃ 向量做逐通道 INT2 分块量化
     b. 存储量化结果 Q(k̃) + 量化参数 Δ, z + 缩放因子 {s_i}
     c. 清空残差缓冲

--- Query 处理 ---
  10. 计算 Query 向量:q_t = W_Q · x_t  [原始 Query,R^d]
  11. 在线 FHT:q̃_t = H · q_t           [旋转后 Query]
      (注意:不做范数归一化,Query 不量化)

--- Value 处理 ---
  12. 计算 Value 向量:ṽ_t = W_V' · x_t  [V' 已融合 H,等效 ṽ = H·v]
  13. 对 ṽ_t 做逐 token INT2 量化,存入缓存

--- 注意力计算 ---
  14. 对缓存中的量化 Key 反量化:
      k̂_t = s_t · Δ · (Q(k̃_t) - z)
  15. 注意力得分:a_{t'} = q̃_t · k̂_{t'}^T / √d  [对所有历史 t']
  16. 注意力权重:w = softmax({a_{t'}})
  17. 对缓存中的量化 Value 反量化后加权求和:
      output = Σ_{t'} w_{t'} · (Δ_V · (Q(ṽ_{t'}) - z_V))

返回:output ∈ R^d

逐步骤直觉说明:

  • 步骤 5(FHT):将通道离群值打散,是后续一切操作的基础
  • 步骤 6-7(OTS):消除 token 范数差异,使量化块内的值域收窄
  • 步骤 8-9(残差缓冲):新 token 暂时以高精度存储,避免刚生成的 token 量化误差影响当前步输出
  • 步骤 10-11(Query FHT):与 Key 的旋转对消,保证 q~Tk~=qTk\tilde{q}^T \tilde{k} = q^T k(注意力得分不变)
  • 步骤 14(反量化时乘 sts_t:恢复原始尺度,保证注意力计算的数值正确

系统设计与 CUDA 实现

融合内核设计

在 GPU 上,每步解码的主要瓶颈是内存带宽(需要从 HBM 加载全部 KV 缓存)。OScaR 的 CUDA 实现把多个操作融合为单一内核以减少内存访问次数:

Key 量化融合内核:

FHT(k) → 计算范数 s → 归一化 → INT2 量化 → 写入 KV 缓存

整个流程在片上(寄存器 + shared memory)完成,只读一次输入向量,只写一次量化结果。

注意力计算融合内核(基于 FlashDecoding-v2):

OScaR 在 FlashDecoding-v2 的基础上添加了 INT2 反量化路径:

  • 原始 FlashDecoding-v2 的分块注意力计算保留
  • 每个 tile 从 KV 缓存加载 INT2 数据时,立即在寄存器中反量化(先乘量化参数 Δ,z\Delta, z,再乘缩放因子 sts_t
  • 避免了先反量化到 BF16 再做注意力的双倍内存占用
graph LR
    subgraph "OScaR CUDA 内核融合"
        direction TB
        LoadINT2["从 HBM 加载 INT2 KV\n(2 bit/element)"]
        Dequant["在寄存器中反量化\n× Δ, - z, × s_t"]
        FlashAttn["FlashDecoding-v2\n分块注意力计算"]
        WriteOut["写出注意力输出\n到 HBM"]
        
        LoadINT2 --> Dequant --> FlashAttn --> WriteOut
        
        Note1["避免中间 BF16 KV 缓冲\n节省 8× 带宽"]
        Dequant -.-> Note1
    end

内存布局

OScaR 的 KV 缓存内存布局设计需要同时高效存储量化整数、量化参数和缩放因子:

每个注意力头的 Key 缓存布局(序列长度 S,头维度 d,块大小 G):
┌─────────────────────────────────────────────┐
│  量化整数:S × d × 2 bit                    │
│  量化步长 Δ:(S/G) × (d/G) × BF16          │
│  量化零点 z:(S/G) × (d/G) × BF16          │
│  OTS 缩放因子 s:S × BF16                  │
└─────────────────────────────────────────────┘

OTS 缩放因子每 token 一个标量,额外开销为 S×2S \times 2 字节,相比主体整数缓存(S×d/4S \times d / 4 字节)微乎其微。

与 BF16 基线的内存比较

以 Qwen3-8B,序列长度 4K,批大小 48 为例:

方案单 token KV 内存(per head, d=128)总 KV 缓存
BF16128 × 2 × 2 = 512 字节100%
INT2 (OScaR)128/4 + 2 + 参数开销 ≈ 36 字节~18.9%

压缩比约 5.3×(与论文报告数字吻合)。

实验设置

评估基准

语言模型任务:

  • LongBench-E:长上下文理解基准,涵盖文档摘要、问答、代码补全等多个子任务,是目前最全面的长文本评估套件之一
  • NIAH(Needle in a Haystack,大海捞针):在超长文档中检索特定信息片段,测试模型在极长上下文下的精确检索能力
  • 序列长度:LongBench-E 评估 up to 64K;NIAH 评估 up to 128K

多模态任务:

  • OCRBench:图像文字识别基准,测试多模态 LLM 对图像中文字的识别和理解能力

全模态任务:

  • MMAU-Pro:音频理解基准,测试同时处理音频和文本的全模态 LLM

评估模型

  • 纯文本: Llama-3.1-8B、Qwen3-8B
  • 多模态: Qwen3-VL-4B、Qwen3-VL-8B
  • 全模态: Qwen3-Omni-30B

对比基线

基线方法描述
BF16不量化,全精度基线
KIVI逐通道 Key 量化,分块大小 G=32
OTTOnline Token Transformation,另一种旋转方案
TurboQuant最近的 INT2 KV 缓存量化方法
KVQuant混合精度 KV 量化

硬件与软件环境

  • GPU:NVIDIA H20(96GB HBM3)
  • 精度:INT2,group_size=32,残差缓冲大小 R=128
  • 软件:基于 FlashDecoding-v2 构建自定义 CUDA 内核

实验结果

纯文本:LongBench-E

Llama-3.1-8B(INT2,group_size=32):

方法平均分
BF16 基线41.70%
OScaR41.75%
OTT40.74%
KIVI38.21%
KVQuant36.55%

OScaR 的 41.75% 超过了 BF16 基线(41.70%),这在 INT2 量化下是前所未有的结果。领先第二名 OTT 1.01 个百分点。

Qwen3-8B(INT2):

方法平均分
BF16 基线49.6%
OScaR48.7%
OTT47.1%
TurboQuant46.8%

仅损失 1.7%,明显优于其他 INT2 方法。

纯文本:NIAH 大海捞针

大海捞针是对 INT2 量化尤其严苛的测试——需要在极长文档中精确找到一个特定信息点,任何量化误差都可能导致检索失败。

方法精确匹配率
OScaR96.5%
BF16 基线96.0%
OTT92.7%
TurboQuant89.3%
KIVI84.1%

OScaR 以 96.5% 的成绩超过了 BF16 基线(96.0%)!这一结果说明 OScaR 不仅仅是”损失最小的量化方法”,而是在某些任务上真正取得了质的突破。

多模态:OCRBench

Qwen3-VL-8B(INT2):

方法分数
BF16 基线67.4%
OScaR66.6%
KIVI66.2%
TurboQuant65.8%

OScaR 在多模态模型上同样取得最优,损失仅 0.8%。

Qwen3-VL-4B(INT2):

OScaR 领先第二名 2.5 个百分点,显示出在较小模型上优势更明显的趋势。

全模态:MMAU-Pro

Qwen3-Omni-30B(INT2):

方法分数
BF16 基线85.8%
OScaR85.6%
第二名83.1%

仅 0.2% 的损失,说明 OScaR 的方法可以无缝扩展到包含音频模态的全模态 LLM。

效率基准(H20 GPU)

解码延迟(Qwen3-8B,128K 上下文):

方法每步延迟(相对)
BF16 FlashDecoding-v21.0×
OScaR INT20.33×(即 3.0× 加速)

3× 的延迟降低主要来自 INT2 内存带宽节省(每次 KV 缓存访问的字节数减少 8×)。

内存占用(batch=48, ctx=4K):

方法KV 缓存内存
BF16100%
OScaR INT218.9%(压缩 5.3×)

吞吐量(batch=48, ctx=4K):

方法吞吐量(tokens/s)
BF16 FlashDecoding-v2331
OScaR INT21354(提升 4.1×)
xychart-beta
    title "OScaR vs BF16 效率对比(相对值,越高越好)"
    x-axis ["延迟(↓ 越小越好)", "内存(↓ 越小越好)", "吞吐(↑ 越大越好)"]
    y-axis "相对比值" 0 --> 5
    bar [1.0, 1.0, 1.0]
    bar [0.33, 0.19, 4.1]

(蓝色为 BF16 基线,橙色为 OScaR INT2;延迟和内存越小越好,吞吐越大越好。)

局限性与边界条件

已知局限

1. KV 缓存仍随序列线性增长

OScaR 只压缩了每个元素的精度(从 BF16 的 16bit 到 INT2 的 2bit),但缓存中的 token 数量仍然随序列长度 SS 线性增长。在 S=1MS = 1M(百万 token 上下文)的场景下,即使是 INT2 也可能不够。这需要与 token 淘汰(如 StreamingLLM)或稀疏注意力结合才能彻底解决。

2. 残差缓冲不可消除

OScaR 仍然依赖 R=128R=128 个 token 的 BF16 残差缓冲。这一部分内存不受 INT2 压缩的益处,且对于每步解码都需要高精度计算。在极短序列(S<RS < R)时,OScaR 实际上等价于全精度推理。

3. 特定任务仍存在显著下降

论文中的实验数据显示,在 Qasper 问答任务上,OScaR 的 INT2 精度仍然比 BF16 低约 5-8 个百分点。这表明某些需要精确数字理解或细粒度语言推理的任务对量化更敏感,OScaR 并未完全解决这类问题。

4. 架构覆盖局限

当前 CUDA 实现主要针对 Llama 和 Qwen 架构(标准多头注意力或分组查询注意力,GQA)。Multi-head Latent Attention(MLA,多头潜在注意力)——DeepSeek 系列采用的架构——没有在论文中讨论。MLA 对 Key-Value 做了低秩投影,其 KV 缓存结构根本不同,OScaR 的直接适用性存疑。

5. FHT 对头维度的要求

Fast Hadamard Transform 要求头维度 dd 是 2 的幂次。主流模型通常满足(64、128、256),但存在少数不满足的模型需要 padding 处理,可能引入额外开销。

批判性分析:不足与可改进之处

方法论弱点

1. 理论界的适用性边界未充分讨论

论文附录 G 给出了逐通道量化误差与 token 范数极差的理论关系,但这个界是否在实际 INT2 量化场景下是紧的(tight)?OScaR 之所以能在 NIAH 上超过 BF16,是偶然(某些任务对量化不敏感)还是有系统原因(Hadamard 旋转带来了额外的正则化效果)?论文没有深入分析。如果旋转本身有数值稳定化效果(类似随机投影的降维),那么 OScaR 的优势可能部分来自与量化无关的机制,这需要消融实验来排除。

2. Hadamard 基的最优性未证明

通道旋转使用 Hadamard 矩阵,原因是 FHT 的计算效率。但 Hadamard 矩阵是”最优”旋转吗?随机正交矩阵(比如用 QR 分解从随机矩阵生成)或数据驱动的旋转(离线对每层的 Key 分布学习最优旋转)是否能进一步提升精度?论文直接采用了 Hadamard 而未与其他旋转对比。QuaRot 等工作曾讨论随机正交旋转,值得在这里做消融。

3. 残差缓冲大小 R=128 的选取缺乏系统分析

R=128R=128 这个超参数沿用自 KIVI,论文未提供不同 RR 值对 OScaR 的影响分析。理论上,OScaR 的通道旋转和 token 缩放应该减少了”量化误差在时序上传播”的问题,或许更小的 RR(比如 R=32R=32)也能达到相近精度。这一分析的缺失使得”残差缓冲是否仍然必要”的问题悬而未决。

4. OTS 缩放因子的精度选择

论文中缩放因子 sts_t 用 BF16 存储。为什么不是 FP32 或 INT8?对于 Attention Sink(范数极低,st0.3s_t \approx 0.3),BF16 的表示精度(约 3 位有效数字)是否足够?论文没有讨论缩放因子本身的量化误差对整体误差的影响。

实验设计的局限

5. 测试基准集中在中等长度,未充分测试极长上下文

LongBench-E 的多数任务序列长度在 32K 以内,NIAH 测到了 128K。然而在真实长上下文场景(如 Gemini 风格的 1M token 文档理解),OScaR 的性能曲线是否保持?随着序列增长,Attention Sink 的比例会降低(固定数量的 Sink token,但序列越来越长),TNI 的影响理论上会减弱——这意味着 OScaR 相对 KIVI 的优势可能在超长序列下缩小。

6. 批大小的效率测试覆盖不全

效率评估使用固定的 batch=48,ctx=4K。这是一个特定工作点,吞吐量对批大小非常敏感。在在线推理(小批量,延迟优先)vs. 离线推理(大批量,吞吐优先)两种场景下,OScaR 的相对优势可能截然不同。论文缺少 batch=1(在线推理延迟)到 batch=256(最大吞吐量)的完整扫描。

7. 未与 KV 淘汰方法对比

KV 缓存压缩领域有两大方向:精度压缩(量化)和 token 数量压缩(淘汰,如 H2O、StreamingLLM)。OScaR 专注前者,但未与后者对比。一个 2× token 淘汰 + INT4 量化的方案能达到类似的内存目标,且可能在某些任务上精度更好。论文缺失这一维度的比较,给读者留下了不完整的图景。

论文的选择性呈现

8. 失败案例的展示不充分

论文在 LongBench-E 的整体平均分上展示了优异结果,但仅在附录中简略提到了 Qasper 等子任务的下降。更诚实的展示方式应该是列出所有子任务的逐项对比,而不是只展示对 OScaR 有利的平均分。部分子任务的 5-8% 下降在实际应用中可能是不可接受的。

9. 可重复性信息不完整

论文承诺开源,但撰写时尚未公开代码。CUDA 内核的具体实现细节(分块大小、寄存器使用策略、数值稳定性技巧)对复现至关重要,论文中仅有概要描述,这对独立验证是个障碍。

具体改进建议

建议 A:自适应 OTS 缩放精度

对于范数极低的 token(st<ϵs_t < \epsilon,即 Attention Sink),用 FP32 存储缩放因子;对于范数正常的 token,用 INT8 存储(减少开销)。这种混合精度策略可以在不增加平均存储开销的前提下提升 Sink token 的反量化精度。

建议 B:离线优化旋转矩阵

在不增加推理开销的前提下,离线对每个模型/层学习一个正交矩阵(而不是固定用 Hadamard),使得旋转后 Key 向量的通道方差尽可能均匀。目标函数可以是:

minRO(d)tVarj[(Rkt)j2]\min_{R \in O(d)} \sum_{t} \text{Var}_j\left[(R \mathbf{k}_t)_j^2\right]

这需要在少量校准数据上运行一次,但不修改模型权重,仍属于”无训练”范式,而精度收益可能比 Hadamard 更好。

建议 C:动态残差缓冲大小

根据当前块内 TNI 程度动态调整残差缓冲大小:TNI 大时用更大的 RR,TNI 小时用更小的 RR。这可以在安全性和内存效率之间自适应平衡。

建议 D:与 MLA 的适配路径

DeepSeek 的 MLA 把 Key-Value 投影到低秩潜在空间,缓存的是低维潜在向量而非完整的 KV 向量。OScaR 的 Hadamard 旋转需要在潜在空间维度上操作(可能是 64 维而非 128 维),OTS 的范数计算也需要调整为潜在向量的范数。给出这一适配方案(哪怕是理论层面)会显著扩大 OScaR 的实际覆盖范围。

总结

OScaR 的贡献可以用三句话概括:

  1. 诊断清楚了病因:Token Norm Imbalance(TNI)是 INT2 逐通道 Key 量化失败的根本原因,而不是模糊的”INT2 太极端”。

  2. 找到了最小干预集:Hadamard 通道旋转(消除通道离群值)+ ℓ2 令牌归一化(消除 TNI)是充分且最小必要的两个步骤,缺少任何一个都不能奏效,添加任何额外步骤也没有必要——这正是”奥卡姆剃刀”的精髓。

  3. 工程上可落地:无训练、无数据、开销极小(每 token 一个额外标量),在 H20 上实现了 3× 延迟降低、5.3× 内存压缩、4.1× 吞吐提升,且多模态和全模态模型同样适用。

对于 LLM 推理工程师而言,OScaR 提供了一个直接可用的 INT2 KV 缓存量化方案,特别适合长上下文和大批量推理场景。其设计哲学——先诊断根本原因再设计最小干预——也是值得借鉴的方法论范式。

mindmap
  root((OScaR))
    问题诊断
      KV 缓存 INT2 精度崩溃
      KIVI 逐通道量化的局限
      Token Norm Imbalance (TNI)
      Attention Sink 低范数
    方法创新
      Canalized Rotation
        Hadamard FHT
        消除通道离群值
        正交性保持内积
      Omni-Token Scaling
        ℓ2 归一化
        消除 TNI
        依赖 CR 前置
    工程实现
      FHT 融合进权重
      INT2 反量化内核
      FlashDecoding-v2 扩展
    实验成果
      超越 BF16 (NIAH)
      5.3× 内存压缩
      4.1× 吞吐提升
      多模态全模态通用
    局限与未来
      线性 KV 增长问题
      MLA 架构适配
      超长上下文验证

笔记撰写完成于 2026-06-17。论文 arXiv 链接:https://arxiv.org/abs/2605.19660