RoPE-bewuste bitallocatie voor KV-cache kwantisering

Samenvatting

Bestaande laag-bit KV-cache quantizers behandelen elke gecachete sleutel vaak als een platte vector. Onder RoPE decomposeert de bijdrage van een sleutel aan een toekomstige attention logit echter in een positie-afhankelijke som over tweedimensionale frequentieblokken. Dit maakt key-cache kwantisatie een bloksgewijs bit-toewijzingsprobleem: hoogenergetische RoPE-blokken zijn gevoeliger voor kwantiseringsfouten en zouden meer bits moeten krijgen. We introduceren Block-GTQ, een RoPE-bewuste bit-toewijzer voor key-cache kwantisatie gebouwd op TurboQuant-MSE (TQ-MSE). Voor elke laag en KV-kop berekent Block-GTQ een labelvrije energiescore voor elk RoPE-blok en kent op hebberige wijze gehele bitbreedtes toe op basis van marginale winst. Onder gelijke K/V-bitbudgetten behoudt Block-GTQ de RoPE query-key logits beter op een tien-model diagnostisch paneel, met een verlaging van de per-laag MAE met 32-80% bij 2 en 3 b/dim K-only kwantisatie, en wint het alle 367/367 laagvergelijkingen tegen uniforme TQ-MSE. Deze getrouwheidswinsten vertalen zich naar sterkere downstream lange-context retrieval, begrip en redeneren. Bij K2V2 op Llama-3.1-8B-Instruct verhoogt Block-GTQ het zes-taak NIAH-gemiddelde van 70,6 naar 97,4, en het LongBench-EN-gemiddelde van 36,87 naar 53,31. Op AIME 2024/2025 met DeepSeek-R1-Distill-Qwen-7B, zonder een fp16 recent-key buffer, scoort Block-GTQ bij K3V2 51,7/37,5, dicht bij fp16's 54,2/37,9, terwijl uniforme TQ-MSE instort tot 0,0/0,0. We implementeren verder een packed-cache serving pad. Op een enkele H800 GPU met Qwen2.5-3B-Instruct bereikt packed K3V3 3,24x KV-cache compressie met fp16-vergelijkbare kwaliteit, draait 1,34x sneller dan fp16 FlashAttention2 bij 128K context, vermindert piekgeheugen van 56,31 GB naar 19,85 GB, en blijft haalbaar bij 256K en 512K waar fp16 OOM geeft. Code is beschikbaar op 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.