KV-COBRA: KV Cache Compression via Co-Optimized Bit-Rank Allocation

2026-09-21

AI Summary

KV-COBRA通过优化每个注意力头的秩和位宽分配，解决了在极低比特率下KV缓存压缩效率问题，使用标准低秩投影和标量量化方法。

Abstract

What limits KV-cache compression at extreme bit-rates? We argue that it is not the choice of compression scheme, but how its budget is allocated across attention heads. Existing methods apply rank and bit-width uniformly, ignoring that each head has a different optimal mix of rank truncation and quantization. We show that co-optimizing rank and bit-width per head, using only standard low-rank projection and scalar quantization, dominates uniform allocation, with the largest gains at low bit-rates. Our method, KV-COBRA (Co-Optimized Bit-Rank Allocation), formalizes this as a resource-allocation problem: it balances rank-truncation loss against quantization loss within each head, then redistributes budget across heads to minimize total distortion. A fused Hadamard rotation equalizes per-channel variance, and reordering the SVD basis by attention-KL importance makes the solver query-aware. The same allocator extends to joint $K{+}V$ compression. On perplexity, zero-shot, and long-context benchmarks from $0.5$ to $4$ bits per dimension (bpd), KV-COBRA shows the smallest accuracy degradation among evaluated methods at low bpd, with no per-token overhead.

Problem

• KV-cache compression
• bit-rates
• attention heads
• rank truncation
• quantization

Innovation

• Co-Optimized Bit-Rank Allocation
• Resource-Allocation Problem
• Hadamard Rotation
• Attention-KL Importance