Folding Attention Biases into Keys Makes Compacted KV Caches Run on Standard Kernels

In long-document QA, KV cache compression can already cut a document's cache to a tenth of its size with answers largely unchanged, yet these methods do not run on vLLM or FlashAttention. The reason is concrete: the strongest one, Attention Matching (AM, a method that compresses the cache by fitting attention outputs), adds a bias to every kept key so that the attention score becomes "q·k plus a constant", while standard kernels and paged caches accept only keys and values and have no slot for that constant. This paper "folds" the bias into the keys: for each attention head it finds one direction and moves every key along it by a distance proportional to its bias, so that for most queries the change in score is about equal to the bias. The result is an ordinary KV cache, the same size as AM's, with the model's query computation untouched. With equal budgets per head, it is statistically indistinguishable from AM's accuracy in 13 of 14 settings on two long-document QA sets; on unmodified vLLM on one T4, the number of requests that fit at once rises from 4 to 31 (keeping 10%), and peak decode rises from 29 to 200 tokens per second.

Teams serving long-context inference and RAG should read it. It does not solve compression quality; it solves whether the compressed result can be handed to an off-the-shelf engine, which is where KV compression most often stalls on its way to production. The core idea in one sentence: the bias is a constant independent of the query, while a key's contribution is proportional to the query, so if you find a direction on which queries almost always project to 1, the constant can be imitated by shifting the key. Note that the result rests on every head keeping the same number of keys; AM's original paper does better by giving heads different budgets, and in that setting the fold's advantage shrinks.

What problem is it solving?

First, how AM compresses. After the model reads the document, the authors append a repeat instruction and the full document, then run a second prefill. They take a query from every second position of this prefill as a reference query. These reference queries score the keys cached during the first prefill. Their attention outputs over that full cache are the targets for fitting the values. Each key-value head ranks keys by their mean attention over the reference queries in the authors' fit, or by root-mean-square attention in the AM-style fit, and keeps the top fraction, say 10%. It then refits their values by least squares so the compressed attention output stays close to the original.

Why the bias? Because after dropping 90% of the keys, the document's total weight in the softmax shrinks, and the question tokens that follow look more important than they should. A mental-math example (illustrative): a document has 1,000 tokens with similar key scores; after compression to 100, the document's total in the softmax denominator is 10 times smaller. Adding a bias of ln 10 ≈ 2.3 to each kept key makes it count as 10 copies, which restores the document's total weight. AM fits such a weight for each kept key by non-negative least squares; its logarithm is the bias, limited to between −3 and 3.

The trouble is deployment:

The authors measure the biased path at 1.02 to 1.29 times the fastest plain path's time per single-sequence decode call, reaching 1.90 times at batch size 8. What blocks deployment is availability: paged caches, prefix caches and KV transfer between engines all move plain key and value blocks.

How does the fold work?

The key is one line of algebra. Change each key to $k_j^{\prime }=k_j+{\beta }_j\sqrt{d}u$, where $u$ is a direction shared by all retained keys of this key-value head and $d$ is the head dimension. To fit $u$, the context queries from all query heads sharing this key-value head are pooled. The new score is:
$$\frac{q\cdot k_j^{\prime }}{\sqrt{d}}=\frac{q\cdot k_j}{\sqrt{d}}+{\beta }_j(q\cdot u)$$
AM wants "original score plus ${\beta }_j$"; the fold gives "original score plus ${\beta }_j$ times $q\cdot u$". Whenever $q\cdot u$ equals 1, the two are identical. So the whole problem reduces to finding a direction on which future queries project as close to 1 as possible.

A hand calculation (illustrative): a kept key has bias 0.8. A query with $q\cdot u=1.1$ sees a bias of 0.88, 0.08 too much; another with $q\cdot u=0.95$ sees 0.76. With biases between −3 and 3, the error is "bias times how far the projection is from 1".

Why can such a direction be found? Future queries are unknown, but while reading the document the model already produced thousands of queries whose distribution resembles the queries asked later. The authors solve for $u$ in closed form using a context-query set $Q_c$. In the downstream QA benchmarks, $Q_c$ is the reference-query set $Q$, which the authors also use to fit the biases and values. They take these queries from every second position of the second, repeat prefill. In the head-by-head study, they fit the default direction using queries from all 4,096 positions of the original prefill. They fit the biases and values using queries from its last 1,024 positions. To compute the direction: compute the queries' mean $\mu$ and covariance $\Sigma$, take $a=(\Sigma +\lambda I{)}^{-1}\mu$, then rescale to $u=a/({\mu }^{\top }a)$. This direction guarantees that the queries' projections average exactly 1. Under that constraint, the authors minimize projection variance plus $\lambda \Vert u{\Vert }^2$. Only without the ridge term does the resulting direction strictly minimize projection variance itself ($\lambda$ is a small ridge term that keeps the inversion stable). Intuitively, it picks the direction along which the queries "line up" best.

How large is the error? The authors prove a bound. Before the fold's value refit, let both caches use identical values. Then the fold's output differs from AM's by at most $\frac{3}{2}{D}_V|q\cdot u-1|$, where $D_V$ is the largest distance between values the query attends to. There is a more intuitive result too: dropping the bias is the case $q\cdot u=0$. Let AM, the fold and the cache with the bias dropped all keep AM's values unchanged. If the attended keys have similar biases, the first-order approximation gives the fold's error as about "the error of dropping the bias times $|q\cdot u-1|$". When the projection is 0.1 away from 1, the error is about a tenth of dropping the bias.

Finally, the authors refit the values by least squares using the folded keys. This replaces AM's value fit and absorbs part of the fold residual on the reference queries. The fold's extra cost per head is one covariance, one d×d linear solve and one rank-one update, far smaller than AM's own fitting.

Experiments: is it as accurate as AM?

The experiments answer two questions: does the fold keep AM's accuracy, and is it better than dropping the bias. They use QuALITY (four-option reading comprehension over long articles) and LongHealth (five-option QA over more than ten thousand tokens of clinical records per patient), with every head keeping 10% or 5% of the tokens under equal budgets. All four caches keep the same keys: keep uses the original values and no bias, drop refits values with no bias, AM adds the bias, and fold is the fold.

Compression matters most on LongHealth. Accuracy (%) at 10%:

On Qwen3-4B and 8B the fold fully keeps AM's gain over keep, and on 8B it beats drop by 4.5 points, This is the only case in the equal-budget main tables where the accuracy difference between fold and drop is beyond noise. This shows the bias does help in long clinical QA, and the fold carries it over.

The one failure is Qwen3-1.7B at 10%, where the fold is 4.5 points below AM and about level with keep. The authors explain that AM's gain over keep in this cell was small to begin with.

Beyond accuracy, the authors measure the KL divergence of the answer distribution from the full cache's. In all 12 settings across both datasets where drop was run, the fold's KL is below drop's, beyond noise in 10. This shows that even when the chosen answer is the same, the folded model's confidence is closer to the original model's.

These two benchmarks only resolve differences of about 4 to 5 points or more, so "indistinguishable" does not mean identical.

Inside vLLM

The authors wrote a KV connector (a plugin interface vLLM loads by module path), which, without changing vLLM's source, writes the folded keys and values into vLLM's paged blocks, so vLLM only prefills the question. On one T4 with Qwen3-4B, the answer options on 51 questions match the Hugging Face implementation exactly. With Qwen3-14B on an A100 over 200 LongBench samples, 180 greedy outputs are word-for-word identical to Hugging Face and 197 share the first token. The remaining differences come mainly from floating-point differences between the two implementations, since even the full prompt does not match exactly across them.

The capacity change is the most visible. On the T4 with 256 generated tokens per request:

The throughput gain comes from compression itself; the fold's role is letting the compressed cache run on a standard engine. At equal batch sizes, folded caches decode at 97% to 98% of the speed of plain prompts of the same length, which shows the shifted keys add no overhead.

Limitations as I see them

First, the equal-budget premise weakens the result. AM's original paper gives each head its own budget and reaches 71.3% on Qwen3-4B QuALITY at 10%, while the authors' equal-budget AM gets 63.7%. With per-head budgets, the fold matches AM in 9 of 10 settings, but its KL is no longer better than drop's; on Llama-3.1-8B LongHealth at 5%, drop is 7.8 points more accurate than the fold. Per-head budgets also mean different cache lengths per head, which vLLM does not support either, a deployment obstacle the fold cannot remove.

Second, the approximation worsens with distance from the document. The authors first renormalize attention weights over the original cache's first 4,096 tokens and use the attention output over that span as the reference. They compute both methods' relative output errors against this reference, then divide the fold's error by AM's for each attention head; the median of this ratio stays at or below 1.05 at every distance, so a typical head is within 5% of AM. The authors select every third layer and measure its key-value heads at every distance. Each window starts a specified distance after the context and contains 512 consecutive query positions. For each key-value head, they pool the queries produced by all query heads that share it. Each key-value head in each window counts as one pair, and define extra error as the part of the fold's error above AM's. For queries more than 4,000 tokens past the document, the worst tenth of pairs carries 90% to 96% of all extra error. These pairs sit mostly in the first 7 layers, because the mean projection of queries on u drifts away from 1. This will show up in long conversations with many follow-up questions.

Third, the value of the bias depends on the task. On four LongBench QA tasks, drop's F1 is 30.1 and the fold's 28.7, and dropping the bias costs AM only 0.7 points. On such tasks, drop is simpler.

How a search or RAG team could try it

Step one: compare keep and drop on your own long documents first. Drop only refits values and runs on any engine. If drop is already close to the full cache, there is no need for a bias.

Step two: if drop is clearly behind AM, try the fold. Its extra code is solving for u and one rank-one update, and the paper's KV connector approach can be reused without changing the engine.

Step three: test the case where queries are far from the document, such as the third or fourth follow-up in a conversation, and check whether accuracy drops more than on the first turn. If it does, consider refitting u mid-conversation.

Verdict

In one sentence: each cache entry's bias can be approximately folded into its key if the bias does not depend on the query. If every head has the same budget and the method has no other special mask or output requirements, the resulting cache can run on standard kernels.