针对 KV 缓存量化的 RoPE 感知比特分配

基于 RoPE 块能量感知的自适应 KV 缓存量化方法

前置知识

RoPE（旋转位置编码）

RoPE（Rotary Position Embedding）是一种将位置信息注入注意力机制的方法，通过旋转查询和键向量来编码相对位置。对于查询 $q$ 和键 $k$，RoPE 通过旋转矩阵 $R_{\Delta}$ 将相对位置 $\Delta$ 编码到注意力计算中：$K_{\Delta}(q, k) = q^T R_{\Delta} k$。这个旋转矩阵是块对角结构，每个 2×2 块对应一个频率 $\theta_i$，因此注意力 logit 可以分解为二维频率块的求和。

本文的核心创新在于利用 RoPE 的块结构特性来指导 KV 缓存量化，不理解 RoPE 的数学结构就无法理解 Block-GTQ 的分配策略。

KV Cache（键值缓存）

KV Cache 是大语言模型在自回归推理中为避免重复计算而缓存的历史 token 的键和值向量。对于每个 transformer 层，每个过去的 token 存储一个键向量和一个值向量。在解码步骤中，当前查询与所有缓存的键进行注意力计算，然后使用对应的值向量生成下一个 token。KV Cache 的大小与序列长度成正比，对于长上下文场景（如 128K tokens），KV Cache 可能占用数十 GB 内存。

KV Cache 是长上下文推理的主要内存瓶颈，本文专门解决 KV Cache 的量化压缩问题。

TQ-MSE（TurboQuant-MSE）

TQ-MSE 是一种本地向量量化方法，通过对非零向量 $x$ 进行归一化、应用共享的正交旋转、标量量化旋转后的坐标，然后恢复半径来实现量化。其核心特性是满足误差率律：在 $b$ 比特/坐标下，解码向量满足 $\mathbb{E}\|x - \hat{x}\|_2^2 \leq \|x\|_2^2 C_{TQ} 4^{-b}$，其中 $C_{TQ} = \sqrt{3\pi/2}$。这意味着每增加一个比特，均方误差减少为原来的四分之一。

Block-GTQ 复用 TQ-MSE 作为本地量化器，并利用其 $4^{-b}$ 误差率律进行比特分配。

贪心分配算法

贪心算法是一种在每一步选择当前最优解的策略。在 Block-GTQ 中，算法初始化所有块为最小比特宽度，然后反复将下一个比特分配给边际增益最大的块（即 $\Delta_i(b_i) = s_i 4^{-b_i} - s_i 4^{-(b_i+1)} = \frac{3}{4}s_i 4^{-b_i}$ 最大的块），直到比特预算用完。作者证明了对于目标函数 $J(b) = \sum_i s_i 4^{-b_i}$，这种贪心策略在整数比特分配约束下是最优的。

贪心分配是 Block-GTQ 的核心算法，需要理解其正确性证明（Theorem 1）和计算复杂度。

研究动机

长上下文推理中，KV 缓存成为主要的内存和带宽瓶颈。对于 GQA 风格的 70B 类模型（80 层、8 个 KV 头、128 维头），fp16 KV 缓存每个 token 约需 320 KiB，在 128K token 上下文下约需 40 GiB。现有 KV 缓存量化的方法（如 KIVI、KVQuant、ZipCache 等）将问题视为向量压缩，在头、通道、组或 token 粒度上选择量化粒度，使去量化向量尽可能接近原始向量。这种观点对于存储是自然的，但没有捕获缓存键如何被使用。值误差影响 softmax 后的加权求和，而键误差会扰动未来查询看到的 softmax 前 logit，可能改变注意力分布。对于 RoPE 注意力，键-logit 计算具有块结构，但现有的 RoPE 不可知均匀分配在具有非常不同 logit 敏感度的块上花费相同精度，可能过度保护低影响块而保护不足高影响块。

本文的目标是本文的目标是开发一种轻量级的 RoPE 感知分配器，将键缓存精度分配给未来 logit 最敏感的地方。具体来说，对于每个层和 KV 头，Block-GTQ 计算无标签的 RoPE 块能量分数，结合 TQ-MSE 的 $4^{-b}$ 误差率律，在固定平均比特预算下贪心分配整数比特宽度。被分配相同比特宽度的块分组后由原始 TQ-MSE 本地量化器编码。由于值不进入 RoPE 键-logit 计算，V 使用均匀分配的 TQ-MSE 编码。所有 RoPE 块仍然被存储，Block-GTQ 只改变它们的比特宽度。

与已有工作不同的是，本文的独特切入角度是将键缓存压缩重新表述为二维频率块上的 logit 保留问题，而不是平坦向量重构问题。现有方法大多忽略 RoPE 的块结构或使用不同的机制（如 KVQuant 和 RotateKV 在 RoPE 之前操作、CommVQ 学习与 RoPE 交换的码本、EliteKV 结合 RoPE 频率选择与低秩投影）。Block-GTQ 则从 RoPE logit 误差界导出块能量分数，通过贪心分配为每个 RoPE 频率块分配精度，且不丢弃任何块。与最近的非均匀精度分配方法（如 PM-KVQ 的每层粒度、MixKVQ 和 Kitty 的键通道粒度、Ada-KV 的头级驱逐预算）不同，Block-GTQ 的单位是头内的 RoPE 频率块，分数源自 RoPE logit 误差界。

核心方法

Block-GTQ 的整体思路是利用 RoPE 注意力的块结构特性，将比特分配从传统的均匀分配或基于向量重构误差的分配，转变为基于 RoPE 块能量分数的自适应分配。方法首先从 Q/K 激活计算每个 RoPE 块的无标签能量分数，该分数基于算术-几何平均（AM-GM）不等式，只需要边缘 Q/K 二阶矩，避免了配对查询-键乘积的噪声。然后结合 TQ-MSE 的 $4^{-b}$ 误差率律，使用贪心算法在固定平均比特预算下为每个块分配整数比特宽度。被分配相同比特宽度的块分组后，由原始 TQ-MSE 本地量化器编码。整个方法对 K 和 V 分别处理：K 使用 RoPE 感知的 Block-GTQ，V 使用均匀分配的 TQ-MSE。

Block-GTQ 的核心创新在于发现了 RoPE 块能量分布的非均匀性，并据此设计了自适应比特分配策略。RoPE 注意力的 logit 可以精确分解为二维频率块的求和：$K_{\Delta}(q, k) = \sum_{i=1}^{L} q(i)^T R(\Delta \theta_i) k(i)$，其中 $L = d_h/2$ 是块数量。键缓存量化引起的 logit 误差上界为 $|K_{\Delta}(q, k) - K_{\Delta}(q, \hat{k})| \leq \sum_{i=1}^{L} \|q(i)\|_2 \|e^{(i)}_k\|_2$，每个块独立贡献。实验表明，不同 RoPE 块的能量分数差异巨大（跨越数量级），少数频率块携带大部分查询-键信号。Block-GTQ 使用基于 AM-GM 的无标签能量分数 $s_i = \frac{1}{2}\mathbb{E}[\|q(i)\|_2^2 + \|k(i)\|_2^2]$，结合 TQ-MSE 的 $4^{-b}$ 误差率律，通过贪心算法最小化目标函数 $J(b) = \sum_i s_i 4^{-b_i}$，在整数比特分配约束下实现最优分配。

方法步骤详情

Block-GTQ 的方法步骤如下：（1）块能量分数计算：对于层 $\ell$ 和 KV 头 $h$，从短的未标记校准前缀计算经验能量分数 $s_{\ell,h,i} = \frac{1}{2} \mathbb{E}_{t,g \in G(h)}[\|q(i)^{\ell,g,t}\|_2^2] + \frac{1}{2} \mathbb{E}_t[\|k(i)^{\ell,h,t}\|_2^2]$，其中 $G(h)$ 是读取 KV 头 $h$ 的查询头集合。（2）贪心比特分配：对于每个层和 KV 头，初始化所有块为 $b_{\min}$，然后重复将下一个比特分配给边际增益 $\Delta_i(b_i) = s_i 4^{-b_i} - s_i 4^{-(b_i+1)} = \frac{3}{4}s_i 4^{-b_i}$ 最大的块，直到比特预算 $B$ 用完或达到 $b_{\max}$。作者证明贪心算法对于目标函数 $J(b) = \sum_i s_i 4^{-b_i}$ 在整数比特分配约束下是最优的（Theorem 1）。（3）分组编码：对于每个层和 KV 头，将分配相同比特宽度 $b$ 的 RoPE 块分组 $G(\ell,h)_b = \{i : b_{\ell,h,i} = b\}$，拼接相应的后 RoPE 键块，并使用 TQ-MSE 编码器在 $b$ 比特/维度下编码结果子向量。这保持了分配决策在 RoPE 块粒度，同时避免了为每个二维块使用单独的微小量化器。（4）值编码：由于值不进入 RoPE 键-logit 计算，V 使用均匀分配的 TQ-MSE 编码。

技术新颖性

Block-GTQ 的技术新颖性体现在多个方面：（1）问题重新表述：将键缓存压缩重新表述为二维频率块上的 logit 保留问题，而不是平坦向量重构问题。这是第一个利用 RoPE 块结构指导比特分配的方法。（2）无标签能量分数：提出基于 AM-GM 的无标签能量分数 $s_i = \frac{1}{2}\mathbb{E}[\|q(i)\|_2^2 + \|k(i)\|_2^2]$，只需要边缘 Q/K 二阶矩，避免了配对查询-键乘积的噪声，并且满足 $s^{\star}_i \leq s_i$（不会低估真实块权重）。（3）最优贪心分配：证明贪心算法对于目标函数 $J(b) = \sum_i s_i 4^{-b_i}$ 在整数比特分配约束下是最优的，并给出理论证明（Theorem 1）。（4）打包缓存服务路径：实现了一种压缩 K/V 代码的打包缓存服务路径，避免在 HBM 中实现 fp16 KV 缓存，融合注意力内核只在内核局部临时中解码当前时间块，并在 QK 和 PV 产品中直接消费它们。这在 128K 上下文下实现 3.24× KV 缓存压缩，减少峰值内存从 56.31 GB 到 19.85 GB，并将单请求解码延迟从 70.96 ms 降低到 52.95 ms。

**Figure 1**: RoPE-block allocation. (a) The RoPE attention logit q⊤R∆k decomposes into a sum over two-dimensional frequency blocks. (b) Per-block energy scores for one Qwen3-8B KV head (layer 10, head 4); scores are median-normalized for display and span orders of magnitude. (c) Under the same average bit width ¯b = 3 (dashed line), Block-GTQ reallocates bits from low-energy blocks to high-energy blocks instead of using a uniform 3-bit width.

**Figure 2**: Packed-cache serving path. Persistent HBM stores packed K/V code streams plus norms and metadata. The fused attention kernel decodes only the current tile into kernel-local temporaries and consumes them directly in QK and PV products, avoiding a resident decoded fp16 KV cache.

实验结果

Block-GTQ 在多个实验设置下展现了卓越的性能。在十模型诊断面板（Qwen、Llama、DeepSeek、Mistral 和 GLM 系列模型）上，在 2 和 3 b/dim K-only 下，Block-GTQ 将每层 RoPE-logit MAE 降低 32-80% 跨越模型，并在每个预算下赢得所有 367/367 层比较对抗均匀 TQ-MSE。在长上下文任务上，Block-GTQ 在 Llama-3.1-8B-Instruct 的 K2V2 预算下，将六任务 NIAH 平均从 70.6 提升到 97.4，八任务 LongBench-EN 平均从 36.87 提升到 53.31，相对均匀分配 TQ-MSE 基线。在推理任务上，Block-GTQ 在 DeepSeek-R1-Distill-Qwen-7B 上，不依赖 fp16 最近键缓冲区时，在 K3V2 下得分 51.7/37.5（AIME 2024/2025），接近 fp16 的 54.2/37.9，而均匀分配 TQ-MSE 崩溃到 0.0/0.0。在服务性能上，Block-GTQ 在单个 H800 GPU 上，K3V3 预算下，在 128K 上下文实现 3.24× KV 缓存压缩，峰值内存从 56.31 GB 减少到 19.85 GB，单请求解码延迟从 70.96 ms 降低到 52.95 ms（比 fp16 FlashAttention2 快 1.34×），并在 256K/512K 下仍可行而 fp16 OOM。

**Table 1**: Per-layer RoPE-logit error at the 3 b/dim budget, K-only. Values are mean RoPE-logit MAE across model layers; lower is better. ∆ is the relative reduction versus TQ-MSE; 'Wins' counts layers where Block-GTQ beats uniform TQ-MSE. Definition and protocol are in Appendix B.3.

**Table 2**: Calibration ablations on Llama-3.1-8B-Instruct along two axes: (a) calibration length and (b) calibration corpus. NIAH Overall (%) is the unweighted mean of six NIAH subtasks; higher is better.

**Table 3**: Calibration length sensitivity under prefix noise. Cells are PPL mean ± std across three WT2-train calibration prefixes (offsets 0/10k/20k). Δ is the change in mean PPL from Ncal=128 to Ncal=2048.

**Table 4**: Multi-task NIAH pass-rate (%) on (a) Llama-3.1-8B-Instruct and (b) Qwen2.5-7B-Instruct. Each entry is averaged over context lengths 4K–128K, needle depths (0%–100%), three trials per cell. Block-GTQ uses a 2048-token WikiText-2 calibration; KIVI-ScaleOnly is defined in Appendix B.4.

**Table 5**: LongBench-EN per-subtask scores on Llama-3.1-8B-Instruct. Subtask abbreviations and metrics are listed in Appendix D.1; Avg is the unweighted mean. KIVI denotes KIVI-ScaleOnly (Appendix B.4). Higher is better; bold marks the best quantized method within each budget.

**Table 6**: AIME 2024/2025 pass@1 (%) at K3V2. Protected: 4 sink + 128 recent fp16; No buffer: both 0 (Section 6.3.2). Under no-buffer, KIVI is run as KIVI-ScaleOnly (Appendix B.4).

**Figure 3**: Allocator fingerprint at the 3 b/dim budget. One subplot per model; each vertical slice is one layer, with stacked color bands giving its bit-width distribution (1b red →8b green).

**Figure 4**: Mean softmax KL and Top-10 attended-token overlap across models. K-only quantization (V stays fp16). (a) Mean softmax KL versus fp16 per method, averaged over the ten-model panel at 2, 3, and 4 b/dim budgets. (b) Top-10 attended-token overlap versus softmax KL, one marker per model (color: method, shape: bit-rate); the upper-left corner is best.

**Figure 5**: NIAH single-needle retrieval on Llama-3.1-8B-Instruct. Pass rate is shown over context length (4K–128K) and needle depth (0%–100%), averaged over three trials per cell. The two rows use the same method layout: fp16, KIVI-ScaleOnly (Appendix B.4), TQ-MSE, and Block-GTQ.

**Figure 6**: Kernel-optimization gains, decode latency, and peak memory for Block-GTQ on Qwen2.5-3B-Instruct (single H800). Panel (a) decomposes the speedup contribution of each kernel-optimization stage—each stage targets a specific bottleneck of the packed-cache decode path. As context length grows (panels (b), (c)), Block-GTQ's optimized decode kernel overtakes fp16 FA-2 at T = 128K and continues to run cleanly at T ≥256K where fp16 OOMs. PPL is annotated on in panel (b).

查看结构化数据

任务	指标	本文	基线	提升
RoPE-logit MAE	Mean absolute error across model layers	Block-GTQ achieves 32-80% reduction across 10 models at 2 and 3 b/dim K-only, winning all 367/367 layer comparisons	Uniform TQ-MSE	32-80% reduction, 100% layer win rate
NIAH (Needle in a Haystack)	Pass rate (%) averaged over 6 subtasks	Block-GTQ: 98.4/96.8/97.4 at K3V3/K3V2/K2V2 on Llama-3.1-8B-Instruct	fp16: 99.6, TQ-MSE: 97.7/95.7/70.6, KIVI-ScaleOnly: 35.4/34.4/24.2	TQ-MSE collapses at K2V2 (70.6), Block-GTQ maintains 97.4; KIVI-ScaleOnly never exceeds 35.4
LongBench-EN	Average score (%) across 8 subtasks	Block-GTQ: 59.08/58.84/53.31 at K3V3/K3V2/K2V2 on Llama-3.1-8B-Instruct	fp16: 59.83, TQ-MSE: 57.40/56.65/36.87, KIVI: 47.58/45.02/38.46	At K2V2, TQ-MSE drops to 36.87, Block-GTQ achieves 53.31 (44.6% relative improvement)
AIME 2024/2025 (reasoning)	Pass@1 (%) averaged over 8 samples per problem	Block-GTQ: 51.7/37.5 at K3V2 on DeepSeek-R1-Distill-Qwen-7B (no buffer)	fp16: 54.2/37.9, TQ-MSE: 0.0/0.0, KIVI: 28.8/19.0, PM-KVQ: 40.8/27.5	TQ-MSE completely collapses (0.0/0.0), Block-GTQ stays close to fp16 (95.4%/98.9% of fp16 performance)
Decode latency	Single-request decode latency (ms) at 128K context	Block-GTQ: 52.95 ms on Qwen2.5-3B-Instruct (K3V3)	fp16 FlashAttention2: 70.96 ms, Uniform TQ-MSE: 46.45 ms	Block-GTQ is 1.34× faster than fp16 (25.4% latency reduction) while maintaining comparable quality
Peak GPU memory	Peak memory (GB) at 128K context	Block-GTQ: 19.85 GB on Qwen2.5-3B-Instruct (K3V3)	fp16 FlashAttention2: 56.31 GB	64.7% memory reduction (3.24× KV-cache compression)

局限与改进

作者承认的局限性包括：Block-GTQ 只在 K 上分配比特，保持 V 均匀。添加 V 端分配器、联合 K+V 优化和更密集的打包可以进一步减少内存。融合解码路径是初始单 GPU 实现；多 GPU 和批量服务是开放方向。我观察到的额外局限性包括：贪心分配算法虽然对于目标函数是最优的，但需要校准数据来准确估计块能量分数，如果校准数据与实际推理数据分布不匹配，性能可能下降。Block-GTQ 在 DeepSeek-R1-Distill-Llama-8B 的无缓冲区情况下性能下降较多（AIME 2024/2025 从 fp16 的 43.3/28.8 降到 32.5/23.3），这表明仅 K 侧分配可能在 V 敏感骨干上不够。打包缓存服务路径虽然在长上下文下更快，但在短上下文（T ≤ 64K）下由于内核内解包开销而比 fp16 FlashAttention2 慢。

独立分析的弱点

Block-GTQ 的独立分析弱点包括：（1）仅 K 侧分配：Block-GTQ 只在 K 上分配比特，V 保持均匀分配。这在 V 敏感骨干上可能不够，如在 DeepSeek-R1-Distill-Llama-8B 的无缓冲区情况下性能下降较多。改进方向是添加 V 端分配器或联合 K+V 优化。（2）校准敏感性：块能量分数需要校准数据估计，如果校准数据与实际推理数据分布不匹配，性能可能下降。虽然在实验中 K3V3 表现出鲁棒性（NIAH 在不同校准种子下变化 ≤ 1.07 pp，PPL 在 ±1σ 内），但 K2V2 仍较敏感（NIAH 变化 1.57-4.09 pp，PPL 变化几个 σ）。改进方向是开发更鲁棒的能量分数估计方法或自适应校准。（3）短上下文性能：Block-GTQ 在短上下文（T ≤ 64K）下比 fp16 FlashAttention2 慢，因为打包缓存路径需要内核内解包压缩 K/V 代码的开销。改进方向是优化短上下文路径或开发混合策略。（4）单 GPU 实现：当前的融合解码路径只支持单 GPU，多 GPU 和批量服务是开放方向。改进方向是开发多 GPU 通信和批量服务优化。（5）元数据开销：Block-GTQ 需要每段元数据来支持混合率 K 存储，这比均匀 TQ-MSE（3.88× 压缩）略差（3.24× 压缩）。改进方向是开发更紧凑的元数据编码。

未来方向

作者提出的未来工作方向包括：添加 V 端分配器、联合 K+V 优化和更密集的打包以进一步减少内存；开发多 GPU 和批量服务支持。基于本文成果可以延伸的未来工作方向包括：（1）V 侧分配器：将 Block-GTQ 的 RoPE 感知分配扩展到 V 缓存，可能需要分析 V 在注意力后如何影响输出。（2）联合 K+V 优化：开发联合 K+V 比特分配策略，考虑 K 和 V 在注意力机制中的相互作用，如使用损失梯度敏感性（类似 PM-KVQ）或 V 感知能量分数。（3）自适应校准：开发自适应校准方法，根据实际推理数据动态调整块能量分数，或在运行时更新分配策略。（4）跨层优化：当前 Block-GTQ 在每个层和 KV 头独立分配，可以探索跨层和跨头的联合优化，利用层间和头间的相关性。（5）混合精度策略：结合 Block-GTQ 与其他量化技术（如异常值处理、低秩压缩、稀疏恢复），实现混合精度和混合粒度的量化策略。（6）其他位置编码：将 RoPE 感知分配思想扩展到其他位置编码（如 ALiBi、xPOS），分析它们的块或结构特性。（7）理论分析：进一步分析 RoPE 块能量分数的理论性质，如其与模型架构、训练数据、任务类型的关系，以及不同预算下贪心分配的最优性界。

复现评估

Block-GTQ 的复现评估：开源情况方面，作者在论文中声明代码可在 https://github.com/JIA-Lab-research/blockgtq 获取，这表明方法的开源性。数据方面，实验使用了多个公开数据集：WikiText-2、PG19、C4 和代码语料库用于校准；NIAH（Needle in a Haystack）用于检索任务；LongBench-EN 用于长上下文理解任务；AIME 2024/2025 用于推理任务。这些都是公开可用的数据集。算力方面，论文报告了在单个 H800 GPU（80GB）上的实验，包括不同上下文长度（4K-512K）的评估。这是一个相对昂贵的硬件设置，但不是不可及的。难度方面，实现 Block-GTQ 需要理解 RoPE 注意力的数学结构、TQ-MSE 量化方法、贪心分配算法以及 CUDA 内核优化。论文提供了详细的算法描述和实验设置，但融合注意力内核的实现可能需要深入的 CUDA 编程知识。整体来说，Block-GTQ 的复现难度中等偏高，主要挑战在于 CUDA 内核实现和大规模长上下文评估的算力需求。

针对 KV 缓存量化的 RoPE 感知比特分配 RoPE-Aware Bit Allocation for KV-Cache Quantization