Folding Attention Biases into Keys Makes Compacted KV Caches Run on Standard Kernels
Organizations: Independent Researcher
Abstract
Attention Matching (AM) compacts a KV cache by keeping some keys per head, refitting their values and adding a bias to each kept key's logit. FlashAttention and vLLM's paged caches have no input for it. We fold the bias into the keys along one closed-form direction per head, fitted on context queries, giving a plain cache with AM's size and head dimension and the model's own query path; the output error is bounded by how far each query's projection on that direction is from 1. With equal per-head budgets, unmodified vLLM serves folded caches through a KV connector, matching Hugging Face's answers on all 51 Qwen3-4B questions and its greedy output on 180 of 200 Qwen3-14B LongBench samples. With equal budgets on QuALITY (Qwen3-1.7B, 4B, 8B, Llama-3.1-8B) and LongHealth (three Qwen3 models), the fold is within noise of AM's accuracy in 13 of 14 settings, with lower KL to the full cache than dropping the bias in all 12 where both ran. Under the published per-head budgets, on Qwen3-4B, its Instruct-2507 variant and Llama-3.1-8B, it is within noise of AM in 9 of 10, no longer beats dropping the bias in KL, and the two each win once in accuracy.
Figures & tables
| ratio | worse | worst 10% | AM keep (% heads) | |||
|---|---|---|---|---|---|---|
| model | near | far | far (%) | far (%) | k | |
| Qwen3-1.7B | 1.005 | 0.998 | 11 | 90 | 81 | 18 |
| Qwen3-4B | 1.011 | 1.021 | 35 | 92 | 85 | 17 |
| Qwen3-8B | 1.011 | 1.022 | 40 | 96 | 82 | 17 |
| Qwen3-1.7B | Qwen3-4B | Qwen3-8B | Llama-3.1-8B | |||||
| AM fit | AM-style | ours | ours | ours | ||||
| budget | 10% | 5% | 10% | 5% | 10% | 5% | 10% | 5% |
| full cache | 55.8 | 69.1 | 73.0 | 71.3 | ||||
| accuracy (%) | ||||||||
| keep | 45.5 | 42.3 | 62.2 | 58.8 | 70.2 | 62.5 | 65.7 | 62.4 |
| drop | – | – | 63.3 | 57.9 | 68.4 | 61.2 | 67.4 | 61.4 |
| Qwen3-1.7B | Qwen3-4B | Qwen3-8B | ||||
| budget | 10% | 5% | 10% | 5% | 10% | 5% |
| full cache | 68.3 | 77.4 | 80.4 | |||
| accuracy (%) | ||||||
| keep | 36.2 | 30.4 | 46.7 | 39.7 | 44.7 | 37.2 |
| drop | 38.9 | 33.9 | 54.3 | 45.2 | 51.3 | 42.5 |
| AM | 39.9 | 31.9 | 54.5 | 43.5 | 55.8 | 45.0 |
| full | fold 10% | fold 5% | |
| prompt tokens (mean) | 6,533 | 781 | 462 |
| KV per token (KiB) | 144 | 144 | 144 |
| capacity | 4 | 31 | 44 |
| decode (tokens/s) | 29 | 200 | 297 |
| same letter as HF | 51/51 | 51/51 | 51/51 |
| same 128 tokens as HF | 46/51 | 42/51 | 44/51 |
| Qasper | MF-en | HotpotQA | 2WikiMQA | mean | |
| full cache (HF) | 43.8 | 47.5 | 60.5 | 41.7 | 48.4 |
| AM (HF) | 24.1 | 30.7 | 37.8 | 30.4 | 30.8 |
| drop (HF) | 20.0 | 29.8 | 37.4 | 33.1 | 30.1 |
| fold (HF) | 28.8 | 29.9 | 29.1 | 27.0 | 28.7 |
| fold (vLLM) | 28.0 | 30.2 | 28.9 | 29.5 | 29.2 |
| fold AM |
| capacity | decode (tokens/s) | ||||||
|---|---|---|---|---|---|---|---|
| model | context (max) | full | 10% | 5% | full | 10% | 5% |
| Qwen3-32B | 6,969 | 7 | 45 | 64 | 135 (7) | 872 (45) | 1,232 (64) |
| Qwen3-32B | 13,754 ( ) | 3 | 28 | 45 | 59 (3) | 534 (28) | 864 (45) |
| Qwen3-14B | 6,969 | 39 | 255 | 256 | 766 (39) | 3,442 (128) | 3,831 (128) |
| Qwen3-14B | 27,324 ( ) | 10 | 90 | 158 | 196 (8) | 1,000 (32) | 1,182 (32) |
| method | per-entry logit term | fold applies |
| AM [ 33 ] | , NNLS fit, | yes (this paper) |
| KV-Surgeon, ARC-KV, TAM, online AM [ 23 , 21 , 16 , 17 ] | fitted | yes |
| Latent Compaction [ 8 ] | learned per-slot bias | yes, bias only a |
| Still [ 19 ] | optional , off in all results | yes |
| ToMe, SemantiCache, KeepKV, ZSMerge, AsymKV [ 6 , 27 , 22 , 15 , 9 ] | count or size, -scaled in ZSMerge | yes |
| SelKV [ 7 ] | attention mass | yes |
Appendix figures & tables12 assets
Supplementary material from the paper’s appendix.
Appendix
| 10% | 5% | |||||
| acc | agree | KL | acc | agree | KL | |
| Qwen3-4B, QuALITY; full cache 69.1% | ||||||
| keep | 69.5 | 84.5 | 1.18 | 63.3 | 77.5 | 1.93 |
| drop | 68.7 | 86.0 | 1.14 | 64.4 | 79.4 | 1.65 |
| AM | 67.8 | 86.7 | 0.90 | 65.2 | 81.1 | 1.65 |
| fold | 68.5 | 86.5 | 0.96 | 64.8 | 80.9 | 1.80 |
| 10% | 5% | |||||
| acc | agree | KL | acc | agree | KL | |
| Llama-3.1-8B-Instruct, QuALITY; full cache 71.3% | ||||||
| keep | 69.3 | 85.0 | 0.19 | 67.6 | 82.8 | 0.28 |
| drop | 68.9 | 86.7 | 0.16 | 66.9 | 81.8 | 0.28 |
| AM | 69.3 | 89.0 | 0.12 | 66.7 | 84.6 | 0.22 |
| fold | 68.9 | 88.4 | 0.14 | 66.7 | 83.0 | 0.23 |
| accuracy (%) | fold AM | |||||
| run | keep | drop | AM | fold | acc | KL |
| as in the main runs ( ) | 65 | 63 | 62 | 57 | 5 [ 9, 1] | +0.18 [+0.14, +0.23] |
| folded keys in fp32 | 65 | – | 62 | 58 | 4 [ 7, 1] | +0.19 [+0.14, +0.24] |
| AM’s values | 65 | – | 61 | 64 | +3 [ 1, +5] | 0.01 [ 0.05, +0.03] |
| self-study fit queries | 69 | 65 | 64 | 64 | +0 [ 4, +4] | +0.03 [ 0.03, +0.11] |
| 65 | 63 | 71 | 70 | 1 [ 6, +5] | 0.04 [ 0.09, +0.01] | |
| 10% | 5% | |||||
| acc | agree | KL | acc | agree | KL | |
| QuALITY | ||||||
| merge-drop | 69.1 | 84.3 | 1.15 | 62.9 | 76.4 | 1.95 |
| merge-log | 67.8 | 86.0 | 0.97 | 64.2 | 80.0 | 1.72 |
| merge-fold | 67.2 | 85.4 | 0.98 | 64.4 | 80.9 | 1.74 |
| merge-refit | 68.5 | 86.3 | 0.90 | 66.3 | 81.1 | 1.50 |
| accuracy (%) | unparsed | ||||
|---|---|---|---|---|---|
| base compactor | base | fold | fold base | gained/lost | base/fold |
| KV-Surgeon-fast | 68.5 | 71.8 | +3.3 [ 1.7, +8.2] | 14/8 | 5/3 |
| KVMerger, fitted | 68.0 | 65.2 | 2.8 [ 6.8, +1.1] | 11/16 | 5/11 |
| Qwen3-1.7B | Qwen3-4B | Qwen3-8B | ||||
| 10% | 5% | 10% | 5% | 10% | 5% | |
| AM: no refit refit | 0.26 | 0.02 | +0.37 | +0.35 | +0.09 | +0.01 |
| [ 0.68, +0.14] | [ 0.42, +0.39] | [+0.12, +0.64] | [+0.09, +0.59] | [ 0.11, +0.30] | [ 0.27, +0.28] | |
| fold: no refit refit | 0.53 | 0.21 | +0.43 | +0.44 | +0.36 | +0.67 |
| [ 0.97, 0.08] | [ 0.67, +0.24] | [+0.09, +0.78] | [+0.13, +0.75] | [+0.08, +0.67] | [+0.33, +1.00] | |
| accuracy (%) | |||||||||||
| model | budget | keep | drop | AM | fold | keep | drop | AM | fold | acc | KL |
| Qwen3-1.7B | 20% | 50.4 | 45.7 | 48.1 | 46.6 | 1.74 | 3.00 | 2.07 | 2.41 | 1.5 | +0.34 |
| [ 4.4, +1.1] | [+0.13, +0.54] | ||||||||||
| 10% | 48.3 | 45.3 | 43.1 | 47.2 | 2.82 | 3.62 | 2.57 | 2.88 | +4.1 | +0.31 | |
| [+1.7, +6.8] | [+0.08, +0.53] | ||||||||||
| 5% | 45.7 | 42.3 | 40.8 | 38.4 | 3.59 | 4.47 | 2.69 | 2.89 | 2.4 | +0.20 | |
| Qwen3 | fit | budget | keep | AM | fold | KL | acc (points) |
|---|---|---|---|---|---|---|---|
| 1.7B | AM-style + SS | 10% | 3.78 | 2.11 | 2.26 | +0.14 [ 0.00, +0.30] | +0.9 [ 0.7, +2.4] |
| 5% | 4.97 | 2.44 | 2.65 | +0.21 [+0.07, +0.37] | 0.9 [ 3.4, +1.7] | ||
| 4B | AM-style | 10% | 2.10 | 2.40 | 2.38 | 0.02 [ 0.20, +0.17] | 0.6 [ 2.5, +1.5] |
| 5% | 2.88 | 3.20 | 3.21 | +0.01 [ 0.16, +0.16] | +1.1 [ 0.8, +3.0] | ||
| 8B | AM-style | 10% | 1.48 | 1.35 | 1.18 | 0.16 [ 0.30, 0.03] | +0.4 [ 1.5, +2.2] |
| accuracy (%) | |||||||||||
|---|---|---|---|---|---|---|---|---|---|---|---|
| model | budget | keep | drop | AM | fold | keep | drop | AM | fold | acc | KL |
| Qwen3-1.7B | 10% | 44.5 | 39.7 | 40.7 | 37.7 | 3.62 | 3.78 | 3.23 | 3.36 | 3.0 | +0.13 |
| [ 5.5, 0.5] | [ 0.07, +0.35] | ||||||||||
| 5% | 35.9 | 35.2 | 32.4 | 34.4 | 4.95 | 4.72 | 4.04 | 4.25 | +2.0 | +0.21 | |
| [ 1.0, +5.5] | [ 0.02, +0.46] | ||||||||||
| Qwen3-4B | 10% | 48.5 | 45.7 | 46.5 | 47.5 | 4.25 | 4.96 | 4.79 | 4.88 | +1.0 | +0.09 |
| fold AM | drop fold | ||||
|---|---|---|---|---|---|
| benchmark | model | 10% | 5% | 10% | 5% |
| QuALITY | Qwen3-1.7B | +0.7 [ 1.7, +3.3] | +0.6 [ 1.7, +3.1] | – | – |
| Qwen3-4B | 1.3 [ 3.3, +0.7] | 1.1 [ 3.4, +1.1] | 3.6 [ 6.0, 1.1] | 3.0 [ 6.2, +0.0] | |
| Qwen3-8B | 0.7 [ 3.2, +1.7] | 0.6 [ 2.5, +1.3] | 0.6 [ 2.4, +1.1] | 1.3 [ 4.9, +2.1] | |
| Llama-3.1-8B | +0.6 [ 1.3, +2.6] | +0.0 [ 1.5, +1.5] | 0.4 [ 3.0, +2.4] | 0.7 [ 4.9, +3.0] | |
| LongHealth | Qwen3-1.7B | 7.0 [ 12.2, 2.0] | 1.5 [ 4.8, +1.8] | +3.3 [ 0.5, +6.8] | 1.0 [ 5.1, +3.3] |
| fold | AM | stock run | ||||
|---|---|---|---|---|---|---|
| model | acc (%) | same (%) | acc (%) | same (%) | fold AM KL | drop fold KL |
| Qwen3-1.7B | 47.2 / 46.4 | 97.9 | 43.1 / 43.6 | 95.9 | [ , ] | [ , ] |
| Qwen3-4B | 62.7 / 63.3 | 98.7 | 63.7 / 63.9 | 99.6 | [ , ] | [ , ] |
| Qwen3-8B | 69.1 / 69.3 | 98.9 | 68.4 / 68.9 | 99.1 | [ , ] | [ , ] |
| Llama-3.1-8B | 67.6 / 67.2 | 99.4 | 66.1 / 66.5 | 99.1 | [ , ] | [ , ] |
| decode ( ) | prefill ( ) | |||
| , 1024, 2048, 4096 entries | ||||
| biased | 1.06 (1.02–1.29) | 1.58 (1.09–1.90) | 1.27 (1.16–1.50) | 1.26 (1.16–1.39) |
| Flex | 1.00 (1.00–1.23) | 1.29 (1.00–1.45) | 1.16 (1.00–1.43) | 1.16 (1.07–1.24) |
| Flex / Flex | 1.05 (1.00–1.29) | 1.19 (1.09–1.33) | 1.08 (1.06–1.31) | 1.08 (1.06–1.30) |
| flash ( s) | 13–22 | 19–92 | 52–344 | 130–1503 |
KV-Cache Compression
Key-value cache compression reduces the memory and bandwidth needed to retain attention states during long-context inference.
See all 269 papers