Alocação de Bits Consciente de RoPE para Quantização de Cache KV

Resumo

Os quantizadores de cache KV de baixo bit existentes frequentemente tratam cada chave em cache como um vetor plano. No entanto, sob RoPE, a contribuição de uma chave para um logit de atenção futuro se decompõe em uma soma dependente da posição sobre blocos de frequência bidimensionais. Isso torna a quantização do cache de chaves um problema de alocação de bits por bloco: blocos RoPE de alta energia são mais sensíveis ao erro de quantização e devem receber mais bits. Apresentamos o Block-GTQ, um alocador de bits consciente de RoPE para quantização de cache de chaves, construído sobre o TurboQuant-MSE (TQ-MSE). Para cada camada e cabeça KV, o Block-GTQ calcula uma pontuação de energia sem rótulo para cada bloco RoPE e aloca de forma gulosa larguras de bits inteiras por ganho marginal. Sob orçamentos de bits K/V combinados, o Block-GTQ preserva melhor os logits query-key RoPE em um painel de diagnóstico de dez modelos, reduzindo o MAE por camada em 32-80% com quantização somente K de 2 e 3 b/dim e vencendo todas as 367/367 comparações de camadas contra o TQ-MSE uniforme. Esses ganhos de fidelidade se traduzem em recuperação, compreensão e raciocínio de contexto longo a jusante mais fortes. No K2V2 do Llama-3.1-8B-Instruct, o Block-GTQ eleva a média de seis tarefas NIAH de 70,6 para 97,4, e a média LongBench-EN de 36,87 para 53,31. No AIME 2024/2025 com DeepSeek-R1-Distill-Qwen-7B, sem um buffer de chave recente fp16, o Block-GTQ em K3V2 obtém 51,7/37,5, próximo ao fp16 de 54,2/37,9, enquanto o TQ-MSE uniforme colapsa para 0,0/0,0. Implementamos ainda um caminho de serviço de cache empacotado. Em uma única GPU H800 com Qwen2.5-3B-Instruct, o K3V3 empacotado alcança compressão de cache KV de 3,24x com qualidade comparável ao fp16, executa 1,34x mais rápido que o FlashAttention2 fp16 em contexto de 128K, reduz o pico de memória de 56,31 GB para 19,85 GB, e permanece viável em 256K e 512K onde o fp16 esgota a memória. O código está disponível em 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.