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

What it claims

Attention Matching (AM) compaction leaves each kept key with a scalar logit bias β in [−3, 3]. FlashAttention, vLLM paged caches, PyTorch's flash backend, prefix caches and KV transfer have no input for that bias. The paper, by independent researcher Utkarsh Kanwat (arXiv, October 2026), moves the bias into the keys along one closed-form direction u per key-value head.

The result keeps AM's size, the model's own head dimension (128 for Qwen3-4B) and the unchanged query path. The two existing plain-cache options each lose something:

• Rank-one FlashBias. This exact option adds a coordinate to every query and key, so vLLM would store head dimension 136 with FlashAttention or 256 with FlashInfer.
• Dropping β. This is what Still and QREAM do, and it gives up AM's attention-mass correction.

The author frames the obstacle as availability, not speed. A bias-aware kernel costs little: ARC-KV measures 1.06 to 1.09 times FlashAttention-2. But every engine would need that kernel written and tuned per GPU and cache layout.

How it works

Per head, after AM's key selection and NNLS bias fit:

1. Take the mean μ and covariance Σ of the context queries Q_c, which are pooled across query heads under GQA.
2. Set u = (Σ+λI)⁻¹μ / (μᵀ(Σ+λI)⁻¹μ), with λ = 10⁻² tr(Σ)/d. This is the minimum-variance direction with mean q·u = 1 on Q_c. Plain least squares would apply only a fraction s/(1+s) of each bias.
3. Update the keys with k′ = k + β√d·u, a rank-one update, and round to bf16.
4. Refit the values by least squares on the folded keys without a bias, so the values match what the kernel reads.

The output error versus AM is bounded by (3/2)·D_V·|q·u − 1|, where D_V is the largest distance between attended values. Dropping β is the special case q·u = 0. To first order, the fold's error is |q·u − 1| times the error of dropping β. The fold applies to any query-independent per-entry logit term, including log-count merging. In that general case the constant becomes (b_max − b_min)/4.

The added cost is one covariance, one d×d solve and a rank-one update per head. All of these are small next to AM's own fits.

Evidence

Layer study. On Pride and Prejudice, the first 4,096 tokens are compacted to 409 keys per head (10%) for Qwen3-1.7B, 4B and 8B. The median fold/AM error ratio is within 1.3% of 1 near the context and within 2.2% far out (Δ ≥ 4k tokens). The worst tenth of far (head, window) pairs carries 90 to 96% of the excess error, mostly in layers 0 to 6, where the mean of q·u drifts from 1.

Downstream accuracy. QuALITY uses 30 articles and 534 questions with Qwen3-1.7B, 4B, 8B and Llama-3.1-8B-Instruct. LongHealth uses 20 patients and 398 questions with the three Qwen3 models. Budgets are 10% and 5% per head. Results with equal budgets:

• The fold is within noise of AM in 13 of 14 settings. The exception is Qwen3-1.7B on LongHealth at 10%.
• Fold − AM intervals span −3.8 to +5.0 points, so only gaps of about 4 to 5 points are resolvable.
• KL to the full cache is lower than dropping β in all 12 comparable settings, and beyond noise in 10.
• Dropping β leaves more questions with KL above 1 in 11 of 12.

Serving. Unmodified vLLM loads folded caches through a KV connector. Answers match Hugging Face on all 51 Qwen3-4B questions, and greedy output matches on 180 of 200 Qwen3-14B LongBench samples. Capacity and decode-speed runs use random weights up to Qwen3-32B.

Limits

Per-head budgets. AM's published per-head budgets produce ragged per-head cache lengths, which the engine's memory manager must support whether or not β is kept. The paper therefore serves only equal budgets. Under per-head budgets, the fold is within noise of AM in 9 of 10 settings but no longer beats dropping β in KL.

Far queries. Nothing holds q·u near 1 on queries far from the context. In the layer study AM itself is worse than keep-without-bias on the median head from Δ = 4k.

Truncated page. The fetched text ends partway through Section 5, so the tables, appendices and Section 6 timings were not visible. The visible text gives no code or connector link.