Распределение битов с учетом RoPE для квантования KV-кэша

Аннотация

Существующие низкобитные квантизаторы KV-кэша часто обрабатывают каждый кэшированный ключ как плоский вектор. Однако при использовании RoPE вклад ключа в будущее внимание по логиту раскладывается на сумму, зависящую от позиции, по двумерным частотным блокам. Это превращает квантизацию ключевого кэша в задачу поблочного распределения битов: высокоэнергетические блоки RoPE более чувствительны к ошибке квантизации и должны получать больше битов. Мы представляем Block-GTQ — чувствительный к RoPE распределитель битов для квантизации ключевого кэша, построенный на основе TurboQuant-MSE (TQ-MSE). Для каждого слоя и KV-головы Block-GTQ вычисляет безметочную оценку энергии для каждого блока RoPE и жадно распределяет целочисленные разрядности по предельному выигрышу. При одинаковых K/V битовых бюджетах Block-GTQ лучше сохраняет логиты запрос-ключ RoPE на диагностической панели из десяти моделей, снижая среднюю абсолютную ошибку (MAE) на слой на 32–80% при K-только квантизации с 2 и 3 битами на размерность и выигрывая все 367/367 сравнений слоёв против равномерного TQ-MSE. Эти выигрыши в точности транслируются в более сильное дальнее извлечение контекста, понимание и рассуждение. На K2V2 для Llama-3.1-8B-Instruct Block-GTQ повышает среднее по шести задачам NIAH с 70,6 до 97,4, а среднее по LongBench-EN — с 36,87 до 53,31. На AIME 2024/2025 с DeepSeek-R1-Distill-Qwen-7B, без буфера недавних ключей в fp16, Block-GTQ при K3V2 достигает 51,7/37,5, что близко к fp16 с 54,2/37,9, тогда как равномерный TQ-MSE падает до 0,0/0,0. Мы также реализовали путь обслуживания с упакованным кэшем. На одном графическом процессоре H800 с Qwen2.5-3B-Instruct упакованный K3V3 обеспечивает 3,24-кратное сжатие KV-кэша с качеством, сравнимым с fp16, работает в 1,34 раза быстрее, чем FlashAttention2 в fp16, при контексте 128K, снижает пиковое использование памяти с 56,31 ГБ до 19,85 ГБ и остаётся работоспособным при 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.