FlagRot: Selective Rotation for Low-Bit KV Caches

Long-context inference is bound by the KV cache, and squeezing keys to very few bits would save a lot of memory if the error stays small. The usual fix rotates keys before quantizing, assuming rotation always helps. The authors find it does not: rotation flattens a spiky piece of a key but can give a flat piece a new spike, and both kinds often sit in the same key. FlagRot therefore cuts each key into small pieces when it is written to the cache, tries each piece rotated and as is, and stores whichever gives the smaller error in the attention scores. The weights used to measure that error are computed beforehand on calibration data. At read time it does not rotate keys back; it moves the query into the same coordinates instead. At 2.5 bits per value on average, adding FlagRot to KIVI (a common 2-bit KV cache quantizer) lowers the KL divergence from the original model from 18.7 to 12.4, a geometric mean over 12 models. Teams serving long contexts should read it: it shows that rotation does not always help, and the read-side change is small, one extra query rotation per step.

What problem is it actually solving?

Every token generated in long-context inference reads the whole KV cache. At a 32k-token context, the BF16 cache of one Qwen3-8B sequence takes 4.8 GB. Compressing each value to 2 bits, plus the scales the quantizer needs, brings the average down to 2.5 to 3 bits per value, about a sixfold memory saving. The price is accuracy. In the authors' own serving setup, plain per-token 2-bit round-to-nearest quantization drops the LongBench average of Qwen3-8B by 11.0 points. KIVI, which quantizes along channels, already brings that loss down to about 2 points, and FlagRot works on the remaining gap.

A major source of that error is range. During quantization, each group of 32 consecutive numbers shares one scale. Two bits give only 4 levels, so one very large number in the group stretches those levels to cover it, and the small numbers are squeezed into nearly the same level.

The usual remedy is to rotate before quantizing. The rotation here is a Hadamard transform: it spreads the value of each coordinate over all coordinates and keeps the vector's length. As long as the query gets the same rotation, the dot product stays the same, so keys can be stored in the new coordinates. The small 4-dimensional example below (illustrative; the paper uses 32-dimensional pieces) shows what it does. Spikiness is measured by the peak ratio, the largest absolute entry divided by the root mean square, which is close to 1 for a flat piece. A higher peak ratio usually means more low-bit error. This is only a rule of thumb, so when FlagRot actually chooses, it quantizes both versions and measures the error directly.

Rotation pulls every piece toward the same peak ratio: spiky pieces become flat, and flat ones gain a spike. On real keys, the authors measure that rotated pieces settle at a median peak ratio near 2.2. Under 4-bit quantization, among the flattest tenth of the pieces in Llama-3.1-8B and Qwen3-8B, rotation raises the error for about two thirds; among the spikiest tenth, only about one in six gets worse.

What makes it harder is that spiky and flat pieces often sit in the same key. On a 4-bit per-token cache of Llama-3.1-8B, FlagRot's final choices keep 22% of the pieces as they are, rotate 37% on their own, and rotate 41% together with the whole head. The same piece of the same head can also prefer a different choice on a different token. So one rotation for every key, even if it helps on average, leaves extra error on part of the pieces.

The big picture: choose at write time, rotate the query at read time

FlagRot works in three steps:

1. When a key is written to the cache: cut it into pieces of 32 channels. Each piece first picks between "as is" and "rotated on its own", using the weighted error described in the next section. The chosen pieces are put back together into one key, which is then compared once with the version where the whole head is rotated; the one with smaller error wins and is quantized and stored.
2. Record the choice: each piece keeps 1 flag bit saying whether it was rotated, and the head keeps 1 more flag bit saying whether it was rotated as a whole. With a head dimension of 128 that makes 5 flag bits. The flags sit in the sign bits of the scales. Scales are always positive, so the sign bit was unused, and in the packed 2.5-bit layout the cache page does not grow by a single byte.
3. When the cache is read: keys are not rotated back. Instead the query is turned into two copies, one rotated piece by piece and one rotated as a whole head. This happens once per generated token and is shared by all cached keys. A key whose head flag is set takes its dot product with the whole-head copy. For every other key, each piece takes the dot product with the plain or the rotated query piece, according to its flag, and the parts are added up. Rotation does not change dot products, so this score equals that of the restored key rotated back to the original coordinates. The quantization error remains; what is saved is the inverse transform for each key.

The key division of labour: spend a little effort choosing at write time, and rotate the query just once per step at read time. The sections below explain why each step is done this way.

Why compare by the error that attention sees?

The obvious yardstick is reconstruction error: after quantizing and restoring, how far is the key from the original? The authors argue this is the wrong yardstick. A key's error affects attention through its dot product with the query. The queries of one head tend to put their energy into a few channels. A rotation that shrinks the total error of a piece also spreads what remains evenly over all channels, possibly right into the channels the query cares about most.

So FlagRot compares a weighted error. Each candidate is quantized, restored, rotated back to the original coordinates and subtracted from the original key to give its error. What counts is how large the dot product of that error with the queries is: error that lands on channels the queries use counts more. The queries come from calibration data. The model runs once over 4 windows of 2,048 tokens from the WikiText-2 training text, and the queries of each head are recorded. The per-token version uses the average outer product of these queries as the weight. The KIVI version also keeps these calibration queries, together with each query's softmax normalizer on the calibration text. When new keys are written, these give the attention probability each calibration query would assign to each key: the softmax exponential term from the key-query dot product is divided directly by the stored normalizer, and anything above 1 is set to 1. To score, the dot product of each calibration query with the error is squared, multiplied by that attention probability, and all of these are summed into the key's error score. The more attention a key gets, the more its error counts. These statistics are frozen before evaluation and do not change during inference. When new keys are written during inference, the decision uses only these fixed statistics and the new keys themselves.

The ablation shows how much this yardstick matters. It uses the 3.0-bit KIVI layout, so the KL with no choice is 18.4. Choosing by reconstruction error gives 19.1, worse than not choosing. Weighting by channels alone brings it to 13.8; also weighting each key by its attention probability brings it to 12.3. Choosing alone is not enough; choosing by reconstruction error is worse than not choosing.

How often should it decide?

FlagRot has two modes. The fixed mode decides once during calibration for each layer, head and piece, and every token then follows that choice, so writing a key takes a single quantization. The per-key mode decides afresh for every key it writes, quantizing it three times (as is, rotated per piece, rotated as a whole head) and keeping the best.

Which one to use, the authors find, depends on the quantization axis. KIVI is a common 2-bit KV quantizer. It groups the keys of 32 consecutive tokens in one head, and each channel shares one scale across those 32 tokens. A group is quantized only after its last token is written; until then its keys stay in BF16. So on KIVI, "per key" really means that the 32 keys of a group share one set of flags. To compare candidates, the weighted errors of the 32 keys in the group are added up, and the candidate with the smallest total wins. Groups differ little from one another, so the best choice rarely changes, and deciding once at calibration is nearly enough. On KIVI the fixed mode already gets 78% of the KL reduction of per-group decisions, and per-group decisions lower the error by only another 10% over the fixed mode. Another kind of cache quantizes per token: each key is cut into 4 groups of 32 channels and shares no scale with other tokens. Keys then differ much more in shape, and per-key decisions lower the error by another 35% over the fixed mode. The data agrees: on per-token caches, per-key decisions match the fixed one for only 39% to 44% of keys, against 62% to 69% of KIVI groups. Hoot says▾

Hoot saysWhoever shares a scale also shares the rotation choiceClose ✕

The choice follows the quantization unit. On KIVI 32 keys choose together; with per-token quantization each key chooses alone.

KIVI (per channel)

Shares a scale

32 consecutive tokens in one head, one per channel

Who chooses

A group of 32 keys shares one set of flags

How

Add the 32 weighted errors; the smaller total wins

Per group vs fixed

Error only 10% lower

Same as fixed

62% to 69% of groups

The fixed mode already gets 78% of the KL reduction of per-group decisions.

Per-token quantization

Shares a scale

Each key alone, in 4 groups of 32 channels, shared with no other token

Who chooses

Each key chooses alone

How

Quantize three times at write time, keep the best

Per key vs fixed

Error another 35% lower

Same as fixed

Only 39% to 44% of keys

Keys differ a lot, so one fixed choice misses many of them.

Left is KIVI, right is per-token quantization. KIVI groups look alike, so deciding once at calibration is nearly enough. Per-token keys differ a lot, so deciding per key pays off more. Source: paper §4 comparison of the two decision modes (numbers from this section)

The authors build two versions on this basis. FlagRot-KIVI is the main system, adding the choice to KIVI. FlagRot-auto quantizes per token and first subtracts a calibrated key mean, which the authors call centering. The BF16 keys of the first and most recent tokens get the same shift. This is allowed because subtracting the same vector from every key shifts all the attention scores of one query by the same amount, so the softmax weights do not change. Some models squash attention scores under a ceiling, which is called logit soft-capping. If the scores are shifted and then squashed, the weights do change, so this holds only for models that do not squash them. If subtracting the mean makes a model worse on the calibration data, it is switched off automatically, which is what "auto" in the name refers to.

How well does it work?

First, fidelity. On 12 models from 3B to 27B, the authors measure the KL divergence between the next-token distributions of the quantized model and the original BF16 model (reported times 1,000). The table compares 2-bit caches of the same 2.5 bits per value:

RotateKV is the representative approach that rotates all keys the same way. FlagRot-KIVI is significantly below KIVI on 11 of the 12 models and significantly below RotateKV on 9. The RotateKV in the table is the lower-KL of the two settings the authors ran. The authors tuned the method on only 4 models; the other 8 were run once with the final settings and reach 12.1. This shows the gain was not tuned into a few models. The per-piece choice works at 3 and 4 bits as well, cutting the KL of KIVI by 31% and 30%.

Next, real tasks. The authors ship it as a vLLM plugin and run long-context benchmarks on H200:

Where KIVI loses points, FlagRot wins part of them back. Qwen3-32B recovers the most and ends only 0.23 points below BF16. On RULER it recovers 2.9 of the 11.4 points KIVI loses, about a quarter, and is still well below BF16. The +0.42 on Qwen3-8B is within noise; on that model the fixed mode does slightly better, at +0.72. Qwen3.8-27B is a hybrid model with a KV cache in only 16 of its 64 layers; KIVI loses nothing there, and adding the choice lowers the score slightly.

Finally, cost. On Qwen3-8B the compressed cache holds 6.1 times as many tokens as BF16. At a 32k context, with each generated token allowed at most 60 ms, FlagRot-auto generates 3.0 times the tokens per second of BF16 and 1.4 times that of an FP8 cache, because the freed memory fits more sequences and allows a larger batch. These throughput numbers come from the per-token FlagRot-auto. For the main system, FlagRot-KIVI, the price is slower steps at large batch: on the 27B model at batch 64, each output token takes 33% longer than with BF16. Per-group decisions also cost 5% of prefill throughput, which the fixed mode avoids.

The limitations I see

First, the task evaluation is narrow. It covers only 4 models from 8B to 32B, with no mixture-of-experts model, none above 32B, only H20 and H200 GPUs, and only vLLM as the engine. Comparisons with other low-bit caches rest mainly on simulated KL; served task comparisons cover only KIVI and the authors' own per-token version.

Second, it brings no extra gain on models where KIVI already loses nothing. On Qwen3.8-27B, adding the choice lowers LongBench by 0.57 points. So first check that your model actually loses points with a 2-bit cache.

Third, the task gains are much smaller than the KL gains. KL falls by about a third, but the two 8B models gain only 0.42 and 0.86 LongBench points. Lower KL does not guarantee higher task scores, so check your own task metrics before shipping.

How an inference team could try it

1. First confirm the problem exists. Run 2-bit KIVI on your own model and long-context tasks and see how many points it loses against BF16. If the loss is small, there is no need to add the choice.
2. Start with the fixed mode. Calibrate on a few windows of your own text and decide, for each layer, head and piece, whether to rotate. This adds no work at write time and gets most of the gain on KIVI.
3. If your cache quantizes per token, then try the per-key mode and mean subtraction. The per-key choice pays off more on this kind of cache, at a cost of about 5% prefill throughput.
4. Finally, measure throughput under real traffic. Look at how large the batch can grow when memory is the limit, and how much step latency rises at large batch, to judge whether the saved memory turns into higher throughput.

Conclusion

Rotation should not be a fixed preprocessing step but a choice the encoder makes for each piece, judged by the error attention actually reads.