Folding attention biases into keys simplifies compact key value caches
Folding Attention Biases into Keys Makes Compacted KV Caches Run on Standard Kernels
Machine LearningComputation and Language
Summary
Storing data efficiently in large language models can be tricky because some methods need extra information that standard tools don’t handle well. The authors found a way to merge this extra information (called a bias) directly into the key data, allowing existing systems to use compact storage without changes. Their approach keeps accuracy nearly the same as the original method and works well with multiple model types and tasks. This makes working with big models faster and simpler without losing quality.
What this means in practice
- •For ai system engineers: Enable existing inference systems to support compacted key-value caches without modifying attention kernels, improving efficiency in large language model deployment.
- •For natural language processing tool builders: Integrate compact key-value caching methods that maintain output quality across various model sizes and datasets for efficient language model serving.
Authors
Utkarsh Kanwat
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.