Asignación de bits consciente de RoPE para la cuantización del caché KV

Resumen

Los cuantificadores existentes de caché KV de bajo bit a menudo tratan cada clave almacenada como un vector plano. Sin embargo, bajo RoPE, la contribución de una clave a un logit de atención futuro se descompone en una suma dependiente de la posición sobre bloques de frecuencia bidimensionales. Esto convierte la cuantificación de la caché de claves en un problema de asignación de bits por bloques: los bloques RoPE de alta energía son más sensibles al error de cuantificación y deberían recibir más bits. Presentamos Block-GTQ, un asignador de bits consciente de RoPE para la cuantificación de la caché de claves, construido sobre TurboQuant-MSE (TQ-MSE). Para cada capa y cabeza KV, Block-GTQ calcula una puntuación de energía sin etiquetas para cada bloque RoPE y asigna de forma greedy anchos de bits enteros mediante ganancia marginal. Bajo presupuestos de bits K/V igualados, Block-GTQ preserva mejor los logits de consulta-clave RoPE en un panel de diagnóstico de diez modelos, reduciendo el MAE por capa en un 32-80% en cuantificación solo de K a 2 y 3 b/dim, y ganando todas las 367/367 comparaciones de capas frente a TQ-MSE uniforme. Estas ganancias de fidelidad se traducen en un mejor rendimiento en tareas de recuperación, comprensión y razonamiento de contexto largo. En K2V2 con Llama-3.1-8B-Instruct, Block-GTQ eleva el promedio de seis tareas NIAH de 70.6 a 97.4, y el promedio de LongBench-EN de 36.87 a 53.31. En AIME 2024/2025 con DeepSeek-R1-Distill-Qwen-7B, sin un buffer de claves recientes fp16, Block-GTQ en K3V2 obtiene 51.7/37.5, cercano a los 54.2/37.9 de fp16, mientras que TQ-MSE uniforme colapsa a 0.0/0.0. Además, implementamos una ruta de servicio de caché empaquetada. En una sola GPU H800 con Qwen2.5-3B-Instruct, K3V3 empaquetado logra una compresión de caché KV de 3.24x con calidad comparable a fp16, se ejecuta 1.34x más rápido que FlashAttention2 en fp16 con contexto de 128K, reduce la memoria máxima de 56.31 GB a 19.85 GB, y sigue siendo factible en 256K y 512K donde fp16 se queda sin memoria (OOM). El código está disponible en 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.