Slim attention: cut your context memory in half without loss -- K-cache is all you need for MHA

Nils Graef; Kayvan Zahiri; Andrew Wasielewski

cs.LGMar 7, 2025

Slim attention: cut your context memory in half without loss -- K-cache is all you need for MHA

Authors: Nils Graef, Kayvan Zahiri, Andrew Wasielewski

Organizations: OpenMachine

Abstract

Slim attention shrinks the context memory size by 2x for transformer models with MHA (multi-head attention), which can speed up inference by up to 2x for large context windows. Slim attention is an exact, mathematically identical implementation of the standard attention mechanism and therefore doesn't compromise model accuracy. In other words, slim attention losslessly compresses the context memory by a factor of 2. For encoder-decoder transformers, the context memory size can be reduced even further: For the Whisper models for example, slim attention reduces the context memory by 8x, which can speed up token generation by 5x for batch size 64 for example. And for the T5-11B model for example, the memory can be reduced by 32x because its MHA projection dimension is larger than the embedding dimension. See https://github.com/OpenMachine-ai/transformer-tricks for code and more transformer tricks, and https://www.youtube.com/watch?v=uVtk3B6YO4Y for this paper's YouTube video.

Figures & tables

Figure 1: Mathematically identical implementations of multi-headed self-attention with square weight matrices

\in\mathbb{R}^{d\times d}

. Left : vanilla version. Right : proposed version where values V are computed from keys K with

W_{KV}=W_{K}^{-1}W_{V}

. The

\triangleq

symbol denotes mathematical identity.

Year	Publisher	Model	Params	$d$	layers	context length	context memory
2024	Meta	CodeLlama-7B [ 8 ]	7B	4,096	32	16k	4.3B
	Meta	CodeLlama-13B [ 8 ]	13B	5,120	40	16k	6.7B
	Google	CodeGemma-7B [ 9 ]	8.5B	3,072	28	8k	1.9B
	Cohere	Aya-23-35B [ 10 ]	35B	8,192	40	8k	5.4B
	HuggingFace	SmolLM2-1.7B [ 11 ]	1.7B	2,048	24	8k	0.8B
	HuggingFace	SmolVLM [ 12 ]	2.3B	2,048	24	16k	1.6B

Table 1: Various transformers with MHA (instead of MQA or GQA) and their maximum KV-cache sizes (in number of activations) based on their respective maximum context length.

d

is the embedding dimension (aka hidden size). See appendix for more MHA models.

Figure 2: Standard MHA with KV-cache during (a) prompt-phase and (b) generate-phase.

Figure 3: Slim attention without V-cache during (a) prompt-phase; (b) unoptimized and (c) optimized generate-phase.

	Projection complexity			MHA complexity
	OPs	reads	intensity	OPs	reads	intensity
Vanilla, see Fig. 2 (b)	$8d^{2}$	$4d^{2}$	2	$4nd$	$2nd$	2
Unoptimized, Fig. 3 (b)	$(2n+6)d^{2}$	$4d^{2}$	$(n+3)/2$	$4nd$	$nd$	4
Optimized, Fig. 3 (c)	$8d^{2}$	$4d^{2}$	2	$2nd(h+1)$	$nd$	$2h+2$

Table 2: Complexity per token per layer during the generate-phase for batch size 1

Chip	TOPS (int8)	Theoretical memory bandwidth (GB/s)	Arithmetic intensity (OPs per byte)
Apple A18 [ 26 ]	35	60	583
Apple M4 Max [ 26 ]	38	410	93
Google TPU v4 [ 27 ]	275	1,200	229
Google TPU v5p [ 27 ]	918	2,765	332
NVIDIA H200 [ 28 ]	1,980	4,800	413
NVIDIA B200 [ 28 ]	4,500	8,000	563

Table 3: TOPS 3 3 3 tera-operations-per-second , memory bandwidth, and arithmetic intensity of popular chips

Model	$d$	$d_{k}$	$h$	$e=d_{k}h$	aspect ratio $r=e/d$
CodeGemma-7B	3,072	256	16	4,096	1.3
T5-3B	1,024	128	32	4,096	4
T5-11B	1,024	128	128	16,384	16

Table 4: Models with non-square KV weight matrices where

e>d

	Baseline	Option 1	Option 2
Cache reduction factor	1	2	$2r$
Size of $W_{V}$ or $W_{KV}$	$de$	$e^{2}$ ( $r$ -times larger)	$de$
Computational complexity	baseline	higher	even higher
Support for RoPE?	Yes	Yes	No

Table 5: Comparison of options for supporting non-square KV weight matrices

Figure 4: Slim attention with X-cache (instead of KV or V-cache) for the generate-phase of transformers with non-square weight matrices with

e>d

	Baseline	Option 1	Option 2
Self KV-cache size	100%	50%	50%
Cross KV-cache size	100%	50%	0
Need to read cross- $W_{K}$ and $W_{V}$ during the generate-phase?	No	$W_{KV}$ only	$W_{K}$ and $W_{V}$
Support for RoPE?	Yes	Yes	No
Complexity of cross-phase	full	half	0

Table 6: Comparison of options for supporting encoder-decoder transformers

Figure 5: Slim attention with a single E-cache (encoder cache) that is shared among all cross-attention layers for the generate-phase of encoder-decoder transformers.

	Whisper models
	tiny	base	small	medium	large	Notes
Params	38M	73M	242M	764M	1.5B	number of parameters
$l$	4	6	12	24	32	number of layers
$d$	384	512	768	1,024	1,280	embedding dimension
$d_{\text{ffn}}$	1,536	2,048	3,072	4,096	5,120	hidden dimension of FFN
Cache sizes (in M):

Table 7: Cache size savings and speedups of generate phase for Whisper models

	tiny	small	medium	large-v3
Largest $\kappa(W_{K})$	$3.7\cdot 10^{4}$	$3.3\cdot 10^{7}$	$1.1\cdot 10^{7}$	$5.2\cdot 10^{9}$
Largest $\kappa(W_{V})$	$5.4\cdot 10^{3}$	$8.0\cdot 10^{4}$	$6.2\cdot 10^{5}$	$4.5\cdot 10^{6}$
Largest $\lvert\det W_{K}\rvert$	$<10^{-231}$	$<10^{-473}$	$<10^{-760}$	$<10^{-744}$
Layers with $\kappa(W_{V})<\kappa(W_{K})$	3 of 4	9 of 12	21 of 24	31 of 32

Table 8: Condition numbers of

W_{K}

and

W_{V}

and determinant of

W_{K}

in the cross-attention of the Whisper decoder

	tiny	small	medium	large-v3
K-cache in every layer	$4.98\cdot 10^{-4}$	$2.24\cdot 10^{-2}$	$4.43\cdot 10^{-3}$	$1.09$
V-cache in every layer	$9.85\cdot 10^{-6}$	$5.14\cdot 10^{-5}$	$7.01\cdot 10^{-5}$	$1.54\cdot 10^{-3}$
Per-layer choice	$7.62\cdot 10^{-5}$	$9.22\cdot 10^{-5}$	$1.62\cdot 10^{-4}$	$2.56\cdot 10^{-3}$
Per-layer choice, float64	$1.95\cdot 10^{-13}$	$1.99\cdot 10^{-13}$	$1.40\cdot 10^{-12}$	$1.39\cdot 10^{-11}$

Table 9: Relative logit error of slim attention in the decoder cross-attention on LibriSpeech, in float32 unless noted

Appendix figures & tables5 assets

Supplementary material from the paper’s appendix.

Appendix

	Term	Dimensions of MatMul	OPs per attention-head
Softmax argument	$Q_{i}K_{i}^{\top}$	$(1\times d_{k})(d_{k}\times n)$	$2nd_{k}$
Weighted sum of $V_{i}$	$\text{softmax}()\cdot V_{i}$	$(1\times n)(n\times d_{k})$	$2nd_{k}$
Weighted sum of $K$	$\text{softmax}()\cdot K$	$(1\times n)(n\times d)$	$2nd$

Appendix

Table 10: OPs per attention head for MHA during the generate phase

	OPs per attention-head
	softmax argument	weighted sum	Total OPs
Vanilla, see Fig. 2 (b)	$2nd_{k}$	$2nd_{k}$	$h(2nd_{k}+2nd_{k})=4nd$
Unoptimized, Fig. 3 (b)	$2nd_{k}$	$2nd_{k}$	$4nd$
Optimized, Fig. 3 (c)	$2nd_{k}$	$2nd$	$h(2nd_{k}+2nd)=2nd(h+1)$

Appendix

Table 11: Total OPs for MHA during the generate phase

Figure 6: Alternative scheme during generate-phase

Figure 7: Computing V from K for the case where the KV-cache width

2e

is larger than

d

Year	Publisher	Model	Params	$d$	layers
2025	Metagene	METAGENE-1 [ 46 ]	6.5B	4096	32
2025	Ai2	OLMoE-1B-7B-0125 [ 47 ]	6.9B	2048	16
2025	Meta	DRAMA-base [ 48 ]	212M	768	12
2024	Meta	Chameleon-7B [ 49 ]	7B	4096	32
2024	Meta	llm-compiler-13B [ 50 ]	13B	5120	40
2024	DeepSeek	DeepSeek-vl2-tiny	3.4B	1280	12

Appendix

Table 12: Additional transformer models with MHA

Memory-Efficient Attention

See all 32 papers

Sliding-window beats linear attention
Aug 28, 2026Alexia Jolicoeur-Martineau, Rhea Sanjay Sukthanker, Pashmina Cameron +1Linear Attention KV Caching
Microsoft Applied Sciences Group (ASG) · Independent
Compressing Long Context into Answer-Aligned Memory Embeddings for LLM Inference
Sep 22, 2026Md Mostafizer Rahman, Md Faizul Ibne Amin, Md Shahajada Mia +2LLM Inference Efficiency Long-Context Language Model Inference
Lucy Family Institute for Data & Society, University of Notre Dame, IN, USA · The University of Aizu, Aizuwakamatsu, Japan · Department of Applied and Computational Mathematics and Statistics, University of Notre Dame, IN, USA
You Do Not Fully Utilize Transformer's Representation Capacity
Feb 13, 2025Gleb Gerasimov, Yaroslav Aksenov, Nikita Balagansky +2Representation Collapse Memory-Augmented Language Models
1T-Tech · 2Moscow Institute of Physics and Technology · 3HSE University

Aug 28, 2026J/K move · Enter open · S save

Sliding-window beats linear attention

Alexia Jolicoeur-Martineau, Rhea Sanjay Sukthanker, Pashmina Cameron, Emy Gervais

Microsoft Applied Sciences Group (ASG) · Independent

Linear Attention KV Caching LLM Inference Acceleration Long-Context Inference Attention Sinks Sliding-Window Attention+1

Open paper AskPDF arXiv

Due to the nature of quadratic attention, Large Language Models (LLMs) consume a lot of memory and energy: every new token costs more than the previous one, and its keys and values must be stored in memory indefinitely, which is unsustainable. Two main lines of work address this: compressing the KV cache, e.g., by evicting or quantizing keys and values, and retrofitting LLMs to use Linear Attention, which replaces the KV cache with a fixed-size state. Retrofitting has attracted a lot of attention, given its promise to solve the quadratic scaling problem with state-of-the-art performance at low cost. However, it has not been properly compared to the simplest form of KV-cache eviction: Sliding Window Attention (SWA) with attention sinks. In this work, we show that SWA with sinks performs as well or better than most retrofitted Linear Attention models across multiple LLMs and downstream tasks, with the largest gains on long-context and generative tasks. On long-context reasoning tasks (Needle-in-a-Haystack and BABILong), SWA achieves massively higher performance (2 to 10 times higher than linear attention). SWA requires no additional training, is extremely fast, and requires little memory, making it an extremely cheap and reliable solution. When the training budget is limited, switching to SWA is a much more effective way to reduce inference memory cost than retrofitting linear attention. Linear attention models have shown promise, but they require training from scratch or extensive retrofitting to reap their architectural benefits and come close to SWA.

Abstract

Memory-Efficient Attention

Explore similar work

Sliding-window beats linear attention

Compressing Long Context into Answer-Aligned Memory Embeddings for LLM Inference

You Do Not Fully Utilize Transformer's Representation Capacity