针对KV-Cache量化的RoPE感知位分配
RoPE-Aware Bit Allocation for KV-Cache Quantization
摘要
现有的低比特率 KV缓存量化算法通常将每个缓存键视为一个平面向量。而在RoPE机制下，键对未来注意力logit的贡献则可以分解为对二维频率块的依赖式求和。这就使得键-缓存量化成为一个块级比特分配问题：能量较高的RoPE块对量化误差更为敏感，因此应被分配更多的比特数。我们提出了Block-GTQ这一基于TurboQuant-MSE的、能够考虑RoPE特性的键-缓存量化比特分配方法。对于每一层和KV头，Block-GTQ会为每个RoPE块计算一个无标签能量得分，并根据边际收益来贪婪地分配整数比特宽度。在符合规定的K/V比特预算限制下，Block-GTQ能够在十模型诊断平台上更好地保留RoPE查询-键logit信息，在2和3 b/dim K-only量化模式下，每层的MAE降低了32-80%，并且在所有367/367层比较中均优于统一的TQ-MSE算法。这些性能提升有助于实现更强大的下游长上下文检索、理解与推理功能。在Llama-3.1-8B-Instruct上的K2V2任务中，Block-GTQ使得六项NIAH指标平均值从70.6提升到97.4，LongBench-EN指标平均值则从36.87提升到53.31。在AIME 2024/2025比赛中，使用DeepSeek-R1-Distill-Qwen-7B模型时，即使没有fp16格式的最新键缓冲区，K3V2下的Block-GTQ也获得了51.7/37.5的得分，这一数值与fp16格式的54.2/37.9相当；而统一的TQ-MSE则得分为0.0/0.0。我们还实现了一种打包缓存服务方式。在配备Qwen2.5-3B-Instruct模型的单个H800 GPU上，K3V3的打包缓存压缩效果相当于fp16格式FlashAttention2的3.24倍，运行速度则快1.34倍；同时，峰值内存需求从56.31 GB降低到19.85 GB。在256K和512K上下文规模下，该算法仍然可行，因为fp16格式下会出现OOM问题。相关代码可访问于https://github.com/JIA-Lab-research/blockgtq。
English Abstract
Existing low-bit KV-cache quantizers often treat each cached key as a flat vector. Under RoPE, however, a key's contribution to a future attention logit decomposes into a position-dependent sum over two-dimensional frequency blocks. This makes key-cache quantization a block-wise bit-allocation problem: high-energy RoPE blocks are more sensitive to quantization error and should receive more bits. We introduce Block-GTQ, a RoPE-aware bit allocator for key-cache quantization built on TurboQuant-MSE(TQ-MSE). For each layer and KV head, Block-GTQ computes a label-free energy score for each RoPE block and greedily allocates integer bit widths by marginal gain. Under matched K/V bit budgets, Block-GTQ better preserves RoPE query-key logits on a ten-model diagnostic panel, cutting per-layer MAE by 32-80% at 2 and 3 b/dim K-only quantization and winning all 367/367 layer comparisons against uniform TQ-MSE. These fidelity gains translate to stronger downstream long-context retrieval, understanding, and reasoning. At K2V2 on Llama-3.1-8B-Instruct, Block-GTQ raises the six-task NIAH average from 70.6 to 97.4, and the LongBench-EN average from 36.87 to 53.31. On AIME 2024/2025 with DeepSeek-R1-Distill-Qwen-7B, without an fp16 recent-key buffer, Block-GTQ at K3V2 scores 51.7/37.5, close to fp16's 54.2/37.9, whereas uniform TQ-MSE collapses to 0.0/0.0. We further implement a packed-cache serving path. On a single H800 GPU with Qwen2.5-3B-Instruct, packed K3V3 achieves 3.24x KV-cache compression with fp16-comparable quality, runs 1.34x faster than fp16 FlashAttention2 at 128K context, reduces peak memory from 56.31 GB to 19.85 GB, and remains feasible at 256K and 512K where fp16 OOMs. Code is available at https://github.com/JIA-Lab-research/blockgtq.