针对注意力缓存量化，提出感知位置编码的分块比特分配方法。
RoPE-Aware Bit Allocation for KV-Cache Quantization
- 根据RoPE中频率块的能量分布，动态分配比特数，高能块多分比特。
- 在2-3比特/维度下，将逐层误差降低32%-80%，全层胜过均匀量化。
- 适用于长文本生成与推理，显著提升模型在长上下文任务中的表现。
现有低比特键值缓存量化器通常将每个缓存键视为扁平向量。然而，在RoPE机制下，键对后续注意力日志概率的贡献可分解为二维频率块的位置相关求和。这使得键缓存量化成为分块比特分配问题：高能量的RoPE块对量化误差更敏感，应分配更多比特。本文提出Block-GTQ，一种基于TurboQuant-MSE（TQ-MSE）的RoPE感知比特分配器。对每层每头，计算各RoPE块的无标签能量得分，并按边际增益贪婪分配整数比特宽度。在匹配的K/V比特预算下，Block-GTQ在十模型诊断面板上更好保留了RoPE查询-键日志概率，2和3比特/维度的键仅量化时，逐层平均绝对误差降低32%-80%，且在全部367个层比较中战胜均匀TQ-MSE。这些保真度提升转化为更强的下游长上下文检索、理解与推理能力。在Llama-3.1-8B-Instruct上，K2V2设置下，六任务NIAH平均得分从70.6提升至97.4，LongBench-EN平均得分从36.87升至53.31。在AIME 2024/2025测试中，使用DeepSeek-R1-Distill-Qwen-7B，无fp16近期键缓冲时，Block-GTQ在K3V2下得分为51.7/37.5，接近fp16的54.2/37.9；而均匀TQ-MSE降至0.0/0.0。此外，我们实现了压缩缓存服务路径：单张H800 GPU上，Qwen2.5-3B-Instruct的K3V3实现3.24倍缓存压缩，质量媲美fp16，速度达fp16 FlashAttention2的1.34倍（128K上下文），峰值内存由56.31 GB降至19.85 GB，且在256K与512K上下文仍可用，而fp16在此会内存溢出。代码已开源：https://github.com/JIA-Lab-research/blockgtq。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。