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

Masked (Causal) Self-Attention

Multi-Head Attention

OR

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

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.

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.

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 using MHA.

  

• MLA compresses $Q, K, V$ representations into a low-rank space to reduce the size of the KV cache using a down-projection matrix which produces latent representations:

  \[C^Q = XW^Q_\downarrow \qquad C^{KV} = XW^{KV}_\downarrow\]
• These compressed latent vectors are then mapped back to $W^Q_\uparrow$, $W^K_\uparrow$, $W^V_\uparrow$ for each attention head using the corresponding up-projection matrices.
• Issue: This low-rank KV compression fails with rotary position embedding (RoPE).
• Fix: A decoupled RoPE technique which introduces extra multi-head queries $W^{QR}$ and a shared key $W^{KR}$ to encode positional information.
• For MLA, the Query, Key & Values are regrouped for each head:
  • The Query is constructed by concatenating the compressed query $Q^C$ with the rotated query $Q^R$.
  • The Key is constructed similarly by concatenating the compressed key $K^C$ with the rotated key $K^R$.

57x reduction in memory usage from MHA to MLA

MLA at Inference Time

Note: the up-projection matrix for the value can be incorporated into the output matrix.

• Due to the associative property, we can merge $W_{\downarrow}^QW_{\uparrow,i}^QW_{\uparrow,i}^{K^T}$ into a single matrix, which means we do not need to explicitly compute the key vectors during inference.
• Also, $W_{\uparrow}^V$ can be absorbed into $W^O$. By incorporating the up-projection matrix into the query and output matrices, we can compute attention efficiently without incurring extra computation.
• MLA has a fundamental limitation: it is incompatible with Rotary Positional Embedding (RoPE):

• We can no longer absorb the up-projection matrix $W_{\uparrow}^K$ during inference. Therefore, during inference, we must recompute the keys for all preceding tokens, which decreases inference efficiency.
• To fix this, we introduce additional multihead queries & a shared key in order to effectively encode position information when using RoPE.
• We concatenate the original query/key vectors and the rotated ones as input to the MHA computation. This approach is known as decoupled RoPE.
• MLA is used in recent LLMs such as Deepseek V3 and R1, Kimi K2, GLM-5, Ling 2.5, Mistral Large 3, Sarvam 105B, etc.
• However, as we generate more tokens, we must calculate attention between the current token and all preceding tokens. This results in slower token-per-second throughput. To fix this, Deepseek proposed a lightning indexer.

Sliding Window Attention (SWA)

• Instead of attending to the entire prefix, each token only attends to a fixed window of recent tokens around its position.
• SWA reduces the memory & compute cost of long-context inference (slower tok/s throughput) by limiting how many previous tokens each position can attend to.

• Because attention is restricted to a local token neighborhood, this mechanism is often referred to as local attention.

  

  Figure: Sliding Window Attention (Window = 2) ¹.
  

• Regular (causal) self-attention mask: each token attends to all preceding tokens.
• Sliding Window Attention mask (window = 2): each token attends only to itself and the previous token.

Deepseek Sparse Attention (DSA)

• Deepseek proposed a lightning indexer to address the slower token-per-second throughput of MLA.
• The key concept of the lightning indexer is to quickly assess the relevance of each token and select only the most relevant ones for attention computation.
• The selected tokens are not determined by a fixed-width local window like in SWA, but by an indexer-plus-selector setup (top-k selector), where a lightning indexer computes index relevance scores, and a token selector keeps only a smaller set of high-scoring past positions.

• Apply partial RoPE to $W^Q$ and $W^K$ before rotation & quantization, then lightning indexing.

Quantization & Rotation in DSA

• How is the attention computation of the indexer between the query vector and all the key vectors made efficient?
• The key step is to first quantize the query & key vectors into 8-bit (BF16/FP32 → FP8) representation.
• Quantization offers a coarse approximation, greatly speeding up the calculation of index scores, and is fine for use since we only want to find the most relevant tokens rather than compute the exact attention scores.
• However, naive quantization of the query & key vectors may result in inaccuracy. To mitigate this issue, we need to apply a rotation.
• Floating point formats with higher precision (float16 or above) can accurately represent the full dynamic range of values of $q^t$ and $k^s$ ($q^t = 0.85, -0.000000017$; $k^s = -0.99, 0.000001$). But this is problematic when we quantize these vectors with lower precision.

• Multiplying the query & key vectors with a random orthogonal matrix $R$ reduces the range of values, but this requires another matrix multiplication: $q^tR, k^sR$.
• A simpler & even more effective approach is to use the Hadamard transform $H$, which yields deterministically computed output vectors that exhibit uniform mixing across all coordinates.
• The Hadamard transform effectively spreads out large spikes across all coordinates, allowing the quantizer to retain more information.
• The Hadamard transform can be implemented efficiently with highly optimized GPU kernels.
• For implementation, a Fast Walsh-Hadamard Transform (FWHT) is used instead of a dense Hadamard matrix, $H_D$ ².

Example

• DSA’s use of a lightning indexer to efficiently identify the most relevant tokens allows for 2 to 3x faster processing of long sequences, while reducing memory consumption by ~30 to 40%. This efficiency boost maintains the same performance level as previous models.

DSA Training

I) Dense warm-up stage: First, we freeze the main MLA layer & only update the parameters of the lightning indexer.

• The target distribution is obtained by summing the mean attention scores across all heads for each query token, followed by applying L1 normalization along the sequence dimension.
• This dense warm-up stage aligns the indexer’s outputs with the main attention distribution.

II) Sparse stage: After warming up the indexer, we add fine-grained token selection and train all the model parameters so that the model can learn the sparse attention patterns of DSA.

• During this stage, we still encourage the indexer to match the main attention distribution, but only for the selected tokens.
• We also detach the indexer inputs from the computational graph, allowing us to optimize it separately.
• The indexer is trained only on its own loss, while the main model is optimized solely on the language modeling loss.

Appendix

Citation

If you found this blog post helpful, please consider citing it:

@article{obasi2026attentionEfficientScalable,
  title   = "How Attention Became Efficient & Scalable",
  author  = "Obasi, Chizoba",
  journal = "chizkidd.github.io",
  year    = "2026",
  month   = "Aug",
  url     = "https://chizkidd.github.io/2026/08/05/attention-efficient-scalable/"
}

References