FlagRot: Selective Rotation for Low-Bit KV Caches

The key-value (KV) cache has become a major memory and bandwidth bottleneck in long-context large language model inference, making ultra-low-bit quantization increasingly important. However, existing rotation-based INT2 methods optimize cache statistics or proxy errors before the complete attention readout, even though the model is ultimately affected by the error propagated through attention and the output projection $W_O$. To address this mismatch, we propose OptR, an output-aware rotation method that minimizes post-$W_O$ attention-output error. OptR decomposes the post-$W_O$ attention-output error into key- and value-induced terms and learns per-head orthogonal corrections through the full INT2 quantization and attention path. OptR further applies an attention-equivalent key reparameterization to reduce large channel-wise offsets without changing the softmax distribution. Across three models and five reasoning and coding benchmarks, OptR consistently improves both QuaRot and OSCAR and strengthens long-context retrieval, while preserving the paged KV-cache format with negligible inference overhead.

Larger models and longer contexts make the KV cache a memory and bandwidth bottleneck. We introduce Block-GTQ, a framework that uses the two-dimensional frequency blocks of rotary position embeddings (RoPE) to allocate cached-key precision. A blockwise logit-error bound motivates a label-free query-key energy score, which we combine with codec-dependent rate penalties. Under equal-cost budgets and diminishing marginal gains, largest-marginal-gain allocation exactly minimizes the separable integer proxy. We instantiate Block-GTQ with TurboQuant-MSE (TQ-MSE), grouping equal-rate blocks to support attention over packed caches without materializing a decoded fp16 KV cache. Across ten model configurations, K-only diagnostics reduce layer-averaged logit mean absolute error by about 32-80% relative to uniform TQ-MSE, improving all 367 layers at both 2- and 3-bit average logical key rates. At K2V2 (2 average logical bits per K/V coordinate) on Llama-3.1-8B-Instruct, Block-GTQ raises the six-task needle-in-a-haystack (NIAH) average from 70.6 to 97.4 and LongBench-EN from 33.7 to 44.9; K2V2 gains also hold on the tested 32B and 70B models. On DeepSeek-R1-Distill-Qwen-7B at K3V2, without fp16 sink or recent-token buffers, it achieves 49.6/38.8% avg@8 on AIME 2024/2025 versus fp16's 54.2/37.9%, where uniform TQ-MSE scores zero on both. On Qwen2.5-3B-Instruct, packed K3V3 serving on one H800 reduces resident KV footprint by a factor of 3.24 and retains perplexity close to fp16 at 128K, while uniform TQ-MSE degrades severely.

Large language models cache all previously computed key-value (KV) pairs during generation, and this KV cache grows linearly with sequence length, making it a primary memory bottleneck for serving. Quantizing the KV cache to fewer bits reduces this cost, yet all current quantizers assign the same bit-width to every attention head, ignoring the large variation in head importance. A natural idea is to allocate more bits to important heads and fewer to the rest. We show, however, that such mixed-precision allocation has a hidden pitfall: each quantizer follows a different distortion curve D(b)=alpha*beta^{-b}, and the decay rate beta varies from 3.6 to 5.3 across quantizer designs. Applying one quantizer's distortion model to another inverts the allocation order and makes performance worse than uniform quantization. We call this failure mode distortion model mismatch and propose RateQuant to resolve it. RateQuant fits a per-quantizer distortion model from a small calibration set, then solves the resulting bit-allocation problem in closed form via reverse waterfilling from rate-distortion theory. On Qwen3-8B at 2.5 average bits, calibrated RateQuant reduces KIVI's perplexity from 49.3 to 14.9 (70% reduction) and improves QuaRot by 6.6 PPL. The entire calibration takes 1.6 s on a single GPU and adds zero overhead at inference time.