用於KV快取量化的RoPE感知位元分配

摘要

現有的低位元KV快取量化器常將每個快取鍵視為平面向量。然而，在旋轉位置編碼（RoPE）下，鍵對未來注意力對數機率的貢獻會分解為一個與位置相關的和，該和涵蓋二維頻率區塊。這使得鍵快取量化成為一個逐塊的位元分配問題：高能量的RoPE區塊對量化誤差更敏感，應分配更多位元。我們提出Block-GTQ，這是一種基於TurboQuant-MSE（TQ-MSE）構建、感知RoPE的鍵快取位元分配器。對於每個層和KV頭，Block-GTQ為每個RoPE區塊計算無標籤的能量分數，並透過邊際增益貪婪地分配整數位元寬度。在匹配的K/V位元預算下，Block-GTQ在一個包含十個模型的診斷面板上更好地保留了RoPE查詢-鍵對數機率，在僅對K進行2和3位元/維度的量化時，將每層的平均絕對誤差（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。在DeepSeek-R1-Distill-Qwen-7B的AIME 2024/2025測試中，不使用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倍KV快取壓縮，且品質與fp16相當；在128K上下文下，運行速度比fp16 FlashAttention2快1.34倍，峰值記憶體從56.31 GB降至19.85 GB，並能在256K和512K上下文下穩定運行，而fp16在這些長度下會記憶體溢出。程式碼開源於 https://github.com/JIA-Lab-research/blockgtq。

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.