How Attention Became Efficient & Scalable: KV Caching, MQA, GQA, MLA, and DSA
-->
Attention mechanisms have evolved considerably to make transformer inference faster and more memory-efficient.
These notes trace that evolution: from vanilla self-attention , through KV caching , to memory-saving variants like MQA , GQA , and MLA , and finally to sparse attention methods like SWA and Deepseek Sparse Attention (DSA) .
Table of Contents
Self-Attention
Masked (Causal) Self-Attention
Multi-Head Attention
Key-Value Caching (KV Caching)
Deepseek R1/V3 memory example
Multi-Query Attention (MQA)
Grouped Query Attention (GQA)
Multihead Latent Attention (MLA)
MLA at Inference Time
Sliding Window Attention (SWA)
Deepseek Sparse Attention (DSA)
Quantization & Rotation in DSA
DSA Training
Appendix
References
Citation
Self-Attention
\[Q = XW_Q, \quad K = XW_K, \quad V = XW_V\]
\[\text{Attention}(Q,K,V) = \text{Softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V = O\]
Masked (Causal) Self-Attention
\[\text{Masked Attention}(Q,K,V) = \text{Softmax}\left(\frac{QK^T}{\sqrt{d_k}} + M\right)V\]
\[\begin{aligned}<br>\text{where} \\<br>M &\equiv \text{lookahead mask}<br>\end{aligned}\]
Multi-Head Attention
\[\text{head}_i = \text{Attention}(Q,K,V)\]
\[\text{MHA} = \text{multi-head attention}(Q,K,V) = \text{Concat}(\text{head}_1, \text{head}_2, \dots, \text{head}_H)W_O = Z\]
OR
\[\begin{aligned}<br>O_i &= \text{Attention}(Q_i, K_i, V_i) \\<br>O &= \text{MultiheadAttention}(Q,K,V) \\<br>O &= \text{Concat}(O_1, O_2, \dots) \\<br>\Delta X &= OW_O<br>\end{aligned}\]
Key-Value Caching (KV Caching)
Due to masking, the transformer model can compute attention for all the tokens at the same time. This is called the prefilling stage .
During inference, the key and value vectors from the previous tokens DO NOT change at all across multiple attention heads & all attention layers. It is therefore wasteful to recompute these vectors every time for each new token.
A simple solution is to store the key & value vectors from previous tokens in memory, and reuse them when needed. This is known as KV caching .
KV caching helps avoid unnecessary computation & therefore speeds up the decoding process.
In KV caching, we DO NOT need to cache any of the previous query vectors because of the masked out entries, which have no effect on the attention output.
KV caching is very effective but memory intensive.
Deepseek R1/V3 memory example
\[\text{Required memory} = (2)(\text{KV dim})(\text{Precision})(\#\text{Heads})(\#\text{Layers})(\text{Sequence length})\]
\[= (2)(128)(2\ \text{bytes/element})(128\ \text{heads})(61\ \text{layers})(32{,}768\ \text{tokens})\]
\[= 131\ \text{GB} \quad (\text{A lot of memory!!!})\]
Multi-Query Attention (MQA)
How do we reduce the memory requirements?
In the MHA mechanism, we cannot adjust the KV dim, precision, or the number of layers, but maybe we can look at the number of heads.
One simple idea is to reduce the number of heads for key and value matrices from $h$ to $1$. Therefore, there is only a single copy of key & value vectors for each token in a layer. This single copy of KV vectors is then shared across all attention heads. This is known as multi-query attention .
\[\begin{aligned}<br>W_Q &\in \mathbb{R}^{d \times (d_k \times h)} \\<br>W_K &\in \mathbb{R}^{d \times (d_k \times 1)} \\<br>W_V &\in \mathbb{R}^{d \times (d_v \times 1)}<br>\end{aligned}\]
KV cache per token
MHA<br>4 MB
MQA<br>31 KB
128x reduction in memory usage
An issue with MQA, however, is that it sacrifices the ability to capture complex relationships between tokens. As a result, MQA’s performance degrades considerably compared to the original MHA.
The drastic reduction of number of heads from $128$ to $1$ is likely the reason for the performance degradation of MQA. We can instead reduce the number of heads to a smaller value denoted as $n_g$.
Grouped Query Attention (GQA)
Instead of the total reduction of number of heads to 1 for the key and value matrices as in MQA, GQA reduces the number of heads for the key and value matrices to a smaller value, $n_g$.
GQA maintains the same overall pattern of attention as MHA, but collapses the number of key-value heads by sharing them across multiple query heads.
MHA : each of head₁, head₂, head₃, head₄ has its own Q, K, V.
GQA : head₁ and head₂ share one K/V pair; head₃ and head₄ share another K/V pair.
MQA : head₁, head₂, head₃, head₄ all share a single K/V pair.
GQA strikes a balance between the memory efficiency of MQA and the expressive power of MHA.
GQA is a popular choice in modern LLMs, including Llama 3 8B (Meta), Qwen 3 4B (Alibaba), Gemma 3 27B (Google), Mistral Small 3.1 24B (Mistral), SmolLM3 3B (Hugging Face), etc.
KV cache per token
MHA<br>4 MB
MQA<br>31 KB
GQA ($n_g = 16$)<br>500 KB
8x reduction in memory usage from MHA to GQA ($n_g = 16$)
Multihead Latent Attention (MLA)
See detailed notes on MLA here.
MLA reduces memory usage & slightly improves model performance compared to others...