Allocation de bits consciente du RoPE pour la quantification du cache KV

Résumé

Les quantificateurs de cache KV à faible nombre de bits existants traitent souvent chaque clé mise en cache comme un vecteur plat. Sous RoPE, cependant, la contribution d'une clé à un logit d'attention futur se décompose en une somme dépendante de la position sur des blocs de fréquence bidimensionnels. Cela fait de la quantification du cache de clés un problème d'allocation de bits par bloc : les blocs RoPE à haute énergie sont plus sensibles aux erreurs de quantification et devraient recevoir davantage de bits. Nous introduisons Block-GTQ, un allocateur de bits sensible à RoPE pour la quantification du cache de clés, construit sur TurboQuant-MSE (TQ-MSE). Pour chaque couche et chaque tête KV, Block-GTQ calcule un score d'énergie sans étiquette pour chaque bloc RoPE et alloue de manière gloutonne des largeurs de bits entières par gain marginal. Sous des budgets de bits K/V appariés, Block-GTQ préserve mieux les logits requête-clé RoPE sur un panel de diagnostic de dix modèles, réduisant l'erreur absolue moyenne (MAE) par couche de 32 à 80 % pour une quantification K uniquement à 2 et 3 b/dim, et remportant toutes les 367/367 comparaisons de couches face à TQ-MSE uniforme. Ces gains de fidélité se traduisent par de meilleures performances en aval pour la récupération, la compréhension et le raisonnement en contexte long. En configuration K2V2 sur Llama-3.1-8B-Instruct, Block-GTQ fait passer la moyenne de NIAH sur six tâches de 70,6 à 97,4, et la moyenne de LongBench-EN de 36,87 à 53,31. Sur AIME 2024/2025 avec DeepSeek-R1-Distill-Qwen-7B, sans tampon de clés récentes en fp16, Block-GTQ en K3V2 obtient des scores de 51,7/37,5, proches des 54,2/37,9 du fp16, tandis que TQ-MSE uniforme s'effondre à 0,0/0,0. Nous implémentons en outre un chemin de service de cache compacté. Sur un seul GPU H800 avec Qwen2.5-3B-Instruct, le cache compacté K3V3 atteint une compression du cache KV de 3,24× avec une qualité comparable au fp16, fonctionne 1,34× plus vite que FlashAttention2 en fp16 pour un contexte de 128K, réduit la mémoire maximale de 56,31 Go à 19,85 Go et reste viable à 256K et 512K là où le fp16 provoque un OOM. Le code est disponible à l'adresse 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.