用于 KV 缓存量化的 RoPE 感知位分配

现有的低位 KV 缓存量化器通常将每个缓存的键视为平面向量。然而，在 RoPE 下，密钥对未来注意力 logits 的贡献分解为二维频率块上的位置相关总和。这使得密钥缓存量化成为逐块比特分配问题：高能量 RoPE 块对量化误差更敏感，应该接收更多比特。我们引入了 Block-GTQ，这是一种基于 TurboQuant-MSE (TQ-MSE) 构建的用于键缓存量化的 RoPE 感知位分配器。对于每一层和 KV 头，Block-GTQ 计算每个 RoPE 块的无标签能量得分，并通过边际增益贪婪地分配整数位宽度。在匹配的 K/V 位预算下，Block-GTQ 在十模型诊断面板上更好地保留了 RoPE 查询密钥 logits，在 2 和 3 b/dim K 量化时将每层 MAE 降低了 32-80%，并赢得了与统一 TQ-MSE 的所有 367/367 层比较。这些保真度的提高转化为更强大的下游长上下文检索、理解和推理。在 Llama-3.1-8B-Instruct 的 K2V2 上，Block-GTQ 将六任务 NIAH 平均值从 70.6 提高到 97.4，将 LongBench-EN 平均值从 36.87 提高到 53.31。在使用 DeepSeek-R1-Distill-Qwen-7B 的 AIME 2024/2025 上，如果没有 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 实现了 3.24 倍的 KV 缓存压缩，质量与 fp16 相当，在 128K 环境下运行速度比 fp16 FlashAttention2 快 1.34 倍，将峰值内存从 56.31 GB 减少到 19.85 GB，并且在 256K 和 512K（其中 fp16）下仍然可行。 OOM。代码可在 https://github.com/JIA-Lab-research/blockgtq 获取。