RoPE-bewusste Bit-Allokation für KV-Cache-Quantisierung

Zusammenfassung

Bestehende Low-Bit-KV-Cache-Quantisierer behandeln oft jeden gecachten Schlüssel als flachen Vektor. Unter RoPE zerfällt jedoch der Beitrag eines Schlüssels zu einem zukünftigen Aufmerksamkeitslogit in eine positionsabhängige Summe über zweidimensionale Frequenzblöcke. Dies macht die Schlüssel-Cache-Quantisierung zu einem blockweisen Bit-Zuweisungsproblem: Hochenergie-RoPE-Blöcke sind empfindlicher gegenüber Quantisierungsfehlern und sollten mehr Bits erhalten. Wir stellen Block-GTQ vor, einen RoPE-bewussten Bit-Zuordner für die Schlüssel-Cache-Quantisierung, der auf TurboQuant-MSE (TQ-MSE) aufbaut. Für jede Schicht und jeden KV-Kopf berechnet Block-GTQ einen label-freien Energiewert für jeden RoPE-Block und weist ganzzahlige Bitbreiten gierig nach marginalem Gewinn zu. Bei gleichen K/V-Bitbudgets erhält Block-GTQ die RoPE-Query-Key-Logits auf einem Zehn-Modell-Diagnosepanel besser, reduziert den MAE pro Schicht um 32–80 % bei 2 und 3 b/dim reiner K-Quantisierung und gewinnt alle 367/367 Schichtvergleiche gegen gleichmäßiges TQ-MSE. Diese Genauigkeitsgewinne übertragen sich auf stärkeren downstream-Kontextabruf, -verständnis und -schlussfolgern bei langen Kontexten. Bei K2V2 auf Llama-3.1-8B-Instruct steigert Block-GTQ den Sechs-Aufgaben-NIAH-Durchschnitt von 70,6 auf 97,4 und den LongBench-EN-Durchschnitt von 36,87 auf 53,31. Auf AIME 2024/2025 mit DeepSeek-R1-Distill-Qwen-7B erreicht Block-GTQ bei K3V2 ohne einen fp16-Puffer für aktuelle Schlüssel 51,7/37,5, nahe an fp16 mit 54,2/37,9, während gleichmäßiges TQ-MSE auf 0,0/0,0 einbricht. Wir implementieren zudem einen gepackten Cache-Bereitstellungspfad. Auf einer einzelnen H800-GPU mit Qwen2.5-3B-Instruct erreicht gepacktes K3V3 eine 3,24-fache KV-Cache-Kompression mit fp16-vergleichbarer Qualität, läuft 1,34× schneller als fp16-FlashAttention2 bei 128K Kontext, reduziert den Spitzenarbeitsspeicher von 56,31 GB auf 19,85 GB und bleibt bei 256K und 512K realisierbar, wo fp16 nicht mehr ausreicht. Code ist verfügbar unter 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.