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

• Standard serving systems lack an input for AM's per-key logit bias, requiring a separate biased path across engines and cache layouts.Dropping the bias avoids that path but gives up the mass correction AM fits.Source: §1 · pp. 1–2
• Moving retained keys along one fitted direction per head folds the bias into a plain cache with AM's size, head dimension, and unchanged query path.The output error is bounded in terms of |q · u −1|.Source: §1 · p. 2
• 13 of 14 equal-budget settings were within noise of AM's accuracy, and 9 of 10 published per-head-budget settings were within noise.With equal budgets, fold KL to the full cache was lower than dropping the bias in all 12 settings where both ran.Source: §1 · p. 2
• The fold also applies structurally to query-independent per-entry logit terms, with tests on capped log-count merging and two KV-Surgeon compactors.The tested cases ran within noise or with no detected loss, respectively.Source: §1 · p. 2
• All 51 Qwen3-4B answers matched Hugging Face, and greedy output matched on 180 of 200 Qwen3-14B LongBench samples using unmodified vLLM.These serving results use equal budgets per head and a KV connector.Source: §1 · p. 2

• AM fits its compact cache on reference queries produced by the model over the context; grouped-query attention pools queries sharing a key-value head.The benchmark queries come from a second prefill that asks the model to repeat the context.Source: §2 · pp. 2–3
• AM keeps high-attention keys and fits nonnegative weights so retained keys carry the full context's attention mass on reference queries.Each weight makes a retained key count as multiple copies in the softmax.Source: §2 · p. 3
• AM refits the retained values by least squares after selecting and weighting the keys.Source: §2 · p. 3
• Each retained key's weight becomes a logit bias βj = log wj, and the compact cache contains retained keys, biases, and refitted values.The bias is added to the key's query logit.Source: §2 · p. 3

3.1 Goal

• No single shifted key can match a nonzero constant AM bias for every possible query, so the method fits its direction using context queries as a proxy for future ones.Source: §3.1 · p. 3

3.2 One shared direction per head

• Moving each retained key by its bias times a shared direction makes its logit shift βj(q · u), matching AM when q · u = 1.The fitted direction sets the context-query mean projection to one and minimizes its spread, with a small ridge for conditioning.Source: §3.2 · pp. 3–4

3.3 Error bound

• The output error is bounded by a factor proportional to |q · u −1| and the attended-value spread; later-query drift or projection spread can increase it.To first order, the fold's error relative to AM is |q · u −1| times the error from dropping the bias with the same values; the bound also applies to static per-entry logit terms.Source: §3.3 · pp. 4–5

3.4 Value refit

• Refitting values on the folded keys absorbs part of the residual on reference queries; on the median head, the default fold stays within about 5% of AM at every tested distance.The keys are rounded to cache precision before fitting the values.Source: §3.4 · p. 5, Figure 2 · p. 6

3.5 Cost

• Per head, folding adds a covariance, a d × d solve, and a rank-one key update; it returns r keys and values of dimension d, as in a plain cache.The value refit replaces AM's at the same size.Source: §3.5 · p. 6

4 Where the Fold Tracks AM, Head by Head

• The fit sets the mean of q · u to 1 on observed queries, but nothing ensures it remains there thousands of tokens later.Across query sets, fold error grows with the spread of q · u and the drift of its mean from 1.Source: §4 · p. 6
• The study compacts the first 4096 tokens of Pride and Prejudice to 409 keys per head, then measures error across near and far query windows.It compares a base fold fitted on the last 1024 context queries with a default fold fitted on all 4096; errors are relative to attention renormalized over the compacted tokens.Source: §4 · p. 6
• Across far windows, median fold/AM error ratio stays within 2.2% of 1, while the worst tenth carries 90–96% of excess, concentrated mostly in layers 0–6.On those pairs, the mean of q · u has drifted from 1.Source: Table 1 · p. 7, §4 · p. 7
• Near context, the default fold's median error ratio is within 1.3% of AM, retaining at least 97% of AM's error reduction over keep.Fitting u on all context queries rather than a subset removes more than half of the base fold's excess.Source: §4 · p. 7

5 Downstream Accuracy

• Evaluation used 30 QuALITY articles and 20 LongHealth patients; paired bootstrap intervals use contexts, but LongHealth's 20 clusters may make intervals too narrow.Intervals also miss GPU nondeterminism; on Qwen3-8B QuALITY at 10%, the 5% of questions with largest KL carry about half the mean.Source: §5 · pp. 7–8
• Under published per-head budgets, fold is level with AM in accuracy in 9 of 10 cells, but has higher KL in all ten; only one KL gap exceeds noise.The fold loses 4.5 points on Llama-3.1-8B LongHealth at 5%; fold and drop each win once in accuracy across the compared settings.Source: §5 · pp. 8–9
• On QuALITY, every AM−keep accuracy interval contains zero; fold KL beats AM on Qwen3-8B but retains 85% and 88% of AM's KL reduction on Qwen3-1.7B while remaining above AM.Source: §5 · p. 8
• Against drop, equal-budget fold has lower KL in all 12 comparisons (beyond noise in 10), while accuracy differs beyond noise once, by 4.5 points for fold.Drop also has more questions with KL above 1 in 11 of 12 comparisons.Source: §5 · p. 8
• With equal per-head budgets, fold and AM show no detected accuracy difference in 13 of 14 QuALITY and LongHealth settings.In the other 13, intervals span −3.8 to +5.0 points, so the benchmarks resolve differences of about 4–5 points or more.Source: §5 · p. 7, Table 2 · p. 9, Table 3 · p. 10
• Folding log-count merging reproduces the biased cache within noise on both benchmarks, while folding two compactors inside KV-Surgeon's pipeline shows no detected loss.For log-count merging, every accuracy difference is within 0.6 points; in KV-Surgeon, neither the 3.3-point gap above KV-Surgeon-fast nor the 2.8-point gap below KVMerger is beyond noise.Source: §5 · p. 9
• On LongHealth, fold preserves AM's 6.3–11.1-point gains over keep across Qwen3-4B and 8B, but loses the gain on Qwen3-1.7B at 10%.Source: §5 · p. 8, Table 3 · p. 10

Answers in vLLM

• All 51 Qwen3-4B questions received the same answer letter as Hugging Face at both cache budgets, and 128 greedy tokens matched exactly on 42 and 44 questions.The 10% and 5% caches averaged 781 and 462 prompt tokens, respectively, versus 6,533 for the full prompt.Source: §6 · p. 10, Table 4 · p. 11

Capacity and speed

• At 10% and 5% cache budgets, the T4 KV pool held 31 and 44 requests versus 4 full-cache requests, with peak decode rates of 200 and 297 tokens/s versus 29.At equal batches from 4 to 32, folded-cache decode ran at 97–98% of the rate for plain prompts of the same length.Source: §6 · p. 11

Kernel cost of the bias

• Where a biased kernel exists, its per-call cost is 1.02–1.29 times the fastest plain kernel and reaches 1.90 times at batch 8.The paper frames the fold’s benefit as availability in engines whose caches and kernels have no bias input.Source: §6 · p. 11

Answers from a 14B model

• 180 of 200 Qwen3-14B LongBench samples matched Hugging Face’s greedy output, with mean F1 29.2 versus 28.7 in Hugging Face.The same first token matched on 197 samples; the fold’s mean F1 was within noise of AM’s.Source: §6 · p. 11, Table 5 · p. 12

Larger models

• On Qwen3-32B, the A100 held 45 requests at 10% and peak decode rose from 135 to 872 tokens/s; at equal batch sizes, injected caches on both larger models ran at 99–101% of same-length plain prompts.These capacity and speed tests used random weights and did not test answers; the fold retained head dimension 128 and the unchanged query path.Source: §6 · p. 11, Table 6 · p. 12

Ablations

• The lower-KL fit varied by model and dataset: AM-style on QuALITY Qwen3-1.7B, the authors’ fit on QuALITY 4B and 8B, and the authors’ fit on all three LongHealth models.Both fits solve NNLS by projected gradient and differ in key ranking, starting point, and number of steps.Source: §7 · p. 12
• On Qwen3-4B and 8B, value refitting lowered the fold’s KL by 0.36–0.67, while using only the mean direction raised KL in all four columns, beyond noise in two.On 1.7B under the authors’ fit, both effects reversed; under the AM-style fit, refitting lowered KL there too.Source: §7 · p. 12
• The fold handles static per-entry logit terms structurally, but tests beyond AM covered only one capped log-count merge and two KV-Surgeon compactors.Table 7 classifies methods by whether each entry’s scalar logit term is the same for every query.Source: Table 7 · p. 13

Compaction with a per-key bias

• AM successors retain per-key biases: ARC-KV learns key selection and fits a mass bias, TAM splits budgets across reasoning segments, and Latent Compaction trains biased entries at test time [21] [8].KV-Surgeon retains original keys with per-key biases and serves them through an additive attention mask; other methods drop the bias or use explicit PyTorch operations [8].Source: §8 · p. 13

Which methods carry a static logit term

• The fold applies to static per-entry terms such as log-counts, but large unbounded counts can increase its error, and only the static count component folds when a method adds query-dependent corrections.SelKV reports counts near 200 at 75% compression; ResKV and MomentKV add query-dependent corrections [31] [12].Source: §8 · p. 13

Serving an attention bias

• FlashBias computes low-rank biases inside FlashAttention, while an exact plain-cache rank-one construction for per-key bias enlarges keys and changes the query path.For head dimension 128, the described vLLM backends may store the augmented cache at 136 or 256 dimensions; the fold keeps the original dimension and query path but is exact only when q · u = 1.Source: §8 · pp. 13–14

Dropping the bias and merging keys

• QREAM removes per-key biases and fits values directly, while a value-only merge is exact for every future query only when eviction loses nothing.The cited analysis reports accuracy loss for compensated merges that rewrite retained keys; the fold also rewrites keys, with error bounded by a multiple of |q · u −1|.Source: §8 · p. 14

Finite-query compilation

• Finite-Query KV Cache Compilation certifies caches for known queries, whereas the fold fits a direction to transfer the condition to unseen queries.Duplicating keys can represent positive bias with log n but may require about 20 copies and cannot express negative bias; CertKV reserves a tail-summary slot [29].Source: §8 · p. 14

Eviction and quantization of plain caches

• Eviction by attention, observation window, or position, and 2-bit quantization, produce plain caches that standard kernels can run [13] [28].Source: §8 · p. 14

Limitations

• The equal-budget KL advantage over dropping the bias is not established under published per-head budgets, and the fold loses to drop in accuracy in the Llama cell where it loses to AM.QuALITY shows no detectable AM-versus-keep gain; only three main LongHealth settings show a clear AM gain, and none of six per-head-budget LongHealth cells does.Source: §9 · p. 14

Where the fold loses to AM

• On Llama-3.1-8B LongHealth at 5% under per-head budgets, fold KL is higher on 19 of 20 patients; the ±1 bias box removes the loss but was chosen on this failure and is untested elsewhere.The sign of the fold–AM KL difference also depends on the model.Source: §9 · p. 14

Queries far from the context

• Far-query behavior was studied on one 4096-token prefix, and the layer-output error estimate overstates the effect by a distance-dependent factor that was not logged.The QA benchmarks ask questions immediately after the document and therefore do not probe this regime; a multi-turn task would.Source: §9 · pp. 14–15

Serving

• Answer serving was checked only for equal budgets on Qwen3-4B and Qwen3-14B; the 14B and 32B capacity and speed runs used random weights and tested no answers.Per-head-budget accuracy runs pad ragged heads and mask the padding, so they do not run as compact plain caches.Source: §9 · p. 15

Scope

• Beyond AM, evaluation covered one log-count merge and two KV-Surgeon compactors on 10 articles; other methods in the structural survey were not tested.Self-study queries were tested only on Qwen3-1.7B and in the Llama diagnosis, and already compacted-and-folded caches were not tested.Source: §9 · p. 15

10 Conclusion

• 13 of 14 equal-budget settings and 9 of 10 published-budget settings put the fold within noise of AM’s accuracy, at about 5-point resolution.With equal per-head budgets, unmodified vLLM also serves folded caches with its stock kernels.Source: §10 · p. 15
• One fitted direction per head folds each AM bias into a plain cache without changing its size, head dimension, or query path.The output error is bounded by a multiple of |q · u −1|.Source: §10 · p. 15
• The fold applies structurally to static per-entry logit terms, with tests on one capped log-count merge and two KV-Surgeon compactors.The conclusion notes that error grows with the spread of the term.Source: §10 · p. 15

A.1 Protocol details

• For LongHealth, two of its 400 questions whose correct answer appears under two letters are dropped; prompts present the patient documents, question and options, and score option-letter log-probabilities at the reply’s first token.Source: §A.1 · p. 17
• Reference queries come from a second prefill asking the model to repeat the context, and each key-value head retains 10% or 5% of context tokens.The queries cover the whole document, so Qc = Q.Source: §A.1 · p. 17
• Value fits use double-precision least squares, while both NNLS fits take projected-gradient steps and clip each wj to [e−3, e3].The step size is 1/Λ, where Λ estimates ∥A∥2 2 for the NNLS matrix.Source: §A.1 · p. 17
• The fold ranks keys by mean attention, initializes retained-key mass from dropped keys with similar cosine direction, and runs 300 steps.The AM-style fit instead uses root-mean-square ranking, a clamped least-squares start, and two steps.Source: §A.1 · p. 17

A.3 Per-head budgets

• Under published per-head budgets, the fold is level with AM in 9 of 10 cells; the exception is Llama-3.1-8B LongHealth at 5%.The remaining cell has a 4.5-point fold accuracy loss to AM. The AM-style fit is an approximation, not a reproduction of published AM, and the QuALITY protocol and fit queries differ; under the published per-head schedule, the AM-style fit reaches 67.8 and 65.2% on Qwen3-4B QuALITY.Source: §A.3 · p. 18, §A.2 · p. 18, Table 8 · p. 19

A.4 The Llama-3.1-8B LongHealth failure

• On Llama-3.1-8B LongHealth at 5%, boxing β to ±1 removes the fold’s gap to AM in the 20-patient rerun.On the original five-patient diagnosis, query drift was larger than on fit queries, and keeping AM’s original values made the fold level with AM; fitting with self-study queries or boxing β to ±1 also removed the gap, though those initial runs were only diagnostic. In the 20-patient rerun, the ±1 box raised fold accuracy by 10.8 points and lowered its KL by 0.17 versus the ±3 run; query drift persisted, but attended β variance was smaller.Source: §A.4 · pp. 20–21, Table 9 · p. 21

A.5 The fold against dropping the bias

• To first order, the fold’s error is |q · u −1| times the error from dropping β, but this relation depends on attended-bias variation.The bound applies with fixed values, before refitting; the diagnostic run used refitted values and did not log the quantities needed to evaluate the bound. Across heads, the fold-to-drop error ratio rose with query drift, and the median fold error was 0.13 of drop’s under |β| ≤3 and 1.1 times the first-order prediction under |β| ≤1.Source: §A.5 · pp. 21–22

A.6 Folding other biased caches

• On QuALITY and LongHealth, merge-fold reproduces merge-log within noise when folding the log-count bias into kept keys.The test used clusters capped at 16 keys, so β ≤ log 16 = 2.77; few evicted keys merged, and the uncapped regime was not tested at scale.Source: §A.6 · p. 22, Table 10 · p. 23
• In a third-party KV-Surgeon pipeline test, neither folded compactor’s accuracy difference from its base was beyond noise.These were 10-article, 181-question Qwen3-4B QuALITY runs with thinking enabled; the study did not run the full cache or the compactors with β dropped, so it tested folding inside the pipeline rather than comparative cache quality.Source: §A.6 · pp. 22–23, Table 11 · p. 24

A.7 Ablations

• On QuALITY ablations, value refitting and the covariance term in u lower fold KL on Qwen3-4B and 8B, while Qwen3-1.7B reverses both patterns under the authors’ fit.The mean direction raises fold KL in all four Qwen3-4B and 8B columns; restricting u to slow-rotating pairs did not lower error in the layer study and shifted accuracy by at most 1.1 points.Source: Table 12 · p. 24, §A.7 · p. 25

A.8 Other budgets and the AM-style fit

• Across QuALITY settings, only one fold–AM accuracy interval excludes zero, favoring the fold; on LongHealth, the AM-style fit serves AM worse than the authors’ fit.The QuALITY exception is +4.1 points on Qwen3-1.7B at 10%, beyond the stated noise margin, though a repeat run falls short of that margin. On LongHealth, only one fold–AM accuracy interval excludes zero, and it is short of the margin; the fold–AM KL interval excludes zero only for Qwen3-8B at 5%.Source: §A.8 · p. 25, Table 13 · p. 26, Table 14 · p. 26

A.11 Serving details; A.12 Far-query diagnosis

• Under stock bf16 kernels, beyond-noise accuracy and KL differences retain their signs, while fitting u on context queries does not eliminate far-query drift.vLLM’s connector wrote folded caches exactly across 232 injections; Qwen3-4B’s folded-answer log-probability differences from Hugging Face were 0.22 and 0.28 at 10% and 5%. Kernel timing found the biased memory-efficient backend could be 10 to 27 times slower than the fastest plain kernel at 16k–32k entries and batch 8; fitting u on Qc removed 54, 34, and 0% of far-context mean excess on Qwen3-1.7B, 4B, and 8B.Source: §A.11 · pp. 28–29, Table 17 · p. 28, §A.12 · pp. 29–30

A.12 Far-query diagnosis

• The worst tenth of far pairs carries 90 to 96% of default-fold excess, with 81 to 89% of those pairs in layers 0 to 6.Their median |m −1| is 0.18, 0.30 and 0.34, versus 0.09, 0.12 and 0.11 for other pairs; median σ differs by at most a factor of 1.5.Source: §A.12 · p. 30
• At ∆=16k, the fold’s relative error reaches 98.6 on Qwen3-8B and 37.6 on 4B, both in layer 3, versus medians of 0.47 and 0.41.From ∆=4k, mean error is 1.1 to 2.6 times AM’s on 4B and 8B; on 8B, the ten worst heads carry three quarters of summed fold error and all lie in layers 0 to 6.Source: §A.12 · p. 30