RoPE感知的KV缓存量化位分配方法Block-GTQ
阅读原文· arxiv.org现有低比特KV量化器将每个key视为平面向量,而RoPE下key的注意力贡献分解为二维频率块的和。Block-GTQ基于TurboQuant-MSE,为每层和KV头计算每个RoPE块的无标签能量分数,通过边际增益贪心分配整数位宽。在10模型诊断集上,2/3 b/dim单K量化下每层MAE降低32–80%,367/367层全部优于均匀TQ-MSE。在Llama-3.1-8B-Instruct的K2V2设置下,六任务NIAH平均分从70.6提升至97.4,LongBench-EN平均分从36.87升至53.31。在DeepSeek-R1-Distill-Qwen-7B的AIME 2024/2025上,K3V2无fp16近期key缓冲区时得分51.7/37.5,接近fp16的54.2/37.9。基于打包缓存的推理路径在单H800、Qwen2.5-3B-Instruct上,K3V3实现3.24倍压缩,128K下比fp16 FlashAttention2快1.34倍,峰值内存从56.31GB降至19.85GB,并在256K和512K下仍可运行。
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.