论文精选

Block-GTQ: RoPE感知的KV缓存量化位分配方法

RoPE-Aware Bit Allocation for KV-Cache Quantization

精选理由

这篇论文用RoPE感知的位分配方案,在KV缓存量化上显著提升长上下文检索和推理,效果逼平fp16,值得研究量化的朋友细读。

AI 摘要

Block-GTQ是一种针对RoPE注意力机制的KV缓存量化位分配方法,基于TurboQuant-MSE构建。它在每个层和注意力头上计算RoPE块的能量得分,通过贪心分配整数位宽。在2和3比特每维度仅量化键的实验中,Block-GTQ在10个模型上使每层平均绝对误差降低32-80%,并赢得全部367个层比较。在Llama-3.1-8B-Instruct上以K2V2配置,NIAH六任务平均从70.6提升至97.4,LongBench英文平均从36.87提升至53.31。在DeepSeek-R1-Distill-Qwen-7B上以K3V2配置,AIME 2024/2025得分51.7/37.5,接近fp16的54.2/37.9,而均匀量化降为0.0/0.0。在H800上对Qwen2.5-3B-Instruct实现3.24倍压缩,128K上下文比fp16 FlashAttention2快1.34倍,峰值内存从56.31GB降至19.85GB。

原文 · arXiv: DeepSeek

RoPE-Aware Bit Allocation for KV-Cache Quantization

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.