Multi-Head Latent Attention and KV Cache Compression

1. The Core Idea: Compressing the Cache Itself

In our previous lessons on MQA and GQA, the strategy was to reduce the number of distinct Key and Value vectors we store. MLA takes a different approach. The central question it asks is:

What if we could compress the Key and Value representations themselves into a much smaller, combined vector before caching?

This is the essence of MLA. It employs a technique known as low-rank compression. Instead of caching the full, separate Key and Value vectors for each head (or group of heads), MLA projects the input's hidden state down into a single, low-dimensional latent vector. This compressed latent vector is the only thing that gets stored in the KV cache.

This diagram provides an excellent high-level view of how MLA's architecture differs from the others we've discussed.

To get a concise and well-animated overview of this core idea, let's watch a short segment from Welch Labs.

2. The MLA Mechanism: Down-Projection and On-the-Fly Up-Projection

As the video explained, MLA is a two-step process that cleverly manipulates the attention computation to work with this compressed latent vector.

1. Down-Projection: The input hidden state $x$ is projected into a low-dimensional latent vector, let's call it $C_{KV}$, using a learned weight matrix $W_{DKV}$. This matrix is shared across all attention heads in the layer.

   $$C_{KV}=x\cdot W_{DKV}$$

   This much smaller $C_{KV}$ is what gets stored in the KV cache.

2. Up-Projection: During the actual attention calculation, this cached latent vector $C_{KV}$ is projected back up to generate the final Key ($K_i$) and Value ($V_i$) vectors for each head $i$. This is done using separate, per-head weight matrices, $W_{UK,i}$ and $W_{UV,i}$.

   $$K_i=C_{KV}\cdot W_{UK,i}$$ $$V_i=C_{KV}\cdot W_{UV,i}$$

This is the crucial innovation. By using unique up-projection matrices for each head, MLA allows every head to learn its own interpretation of the shared latent information. This preserves the representational diversity and performance of MHA, a capability that MQA and GQA sacrifice by forcing heads to share identical K/V vectors.

To dive deeper into this workflow and understand the mathematical elegance behind it, we'll watch a more detailed explanation.

The "Absorption Trick"

As the video highlighted, a naive implementation would require expanding the entire latent cache at every step, which would be computationally expensive. The "absorption trick" is a clever rearrangement of the matrix multiplications. It combines the query projection with the key's up-projection, allowing the attention scores to be computed directly between the query and the compressed latent cache. This makes MLA highly efficient and compatible with optimized kernels like FlashAttention, without adding extra compute during inference.

3. A Deeper Dive and Quantifying the Savings

Now that we have the conceptual and mechanistic understanding, let's solidify it with a detailed reading and a concrete calculation of the memory savings.

Let's use the numbers from the article to make the savings tangible. For a model at the scale of DeepSeek-V2 with a 128,000 token context:

• MHA Cache: Would require ~512 GB. This is unfeasible on a single GPU.
• GQA Cache (8 groups): Reduces this to a more manageable 32 GB.
• MLA Cache (latent dim 512): Slashes this to just ~8 GB.

The formula for the MLA cache size is:
$\text{Cache Size (MLA)}=B\times S\times L\times d_c\times \text{precision\_bytes}$
Where:

• $B$ = Batch size
• $S$ = Sequence length
• $L$ = Number of layers
• $d_c$ = Dimension of the latent cache

Compared to MHA's 2 \times n_h \times d_h factor inside the formula, MLA's d_c is vastly smaller. This is the source of the ~57x memory reduction that DeepSeek reports.

Conclusion

In this lesson, we have deconstructed Multi-head Latent Attention, a key architectural innovation that makes extremely long context windows practical.

Key Takeaways:

• MLA tackles the KV cache bottleneck by applying low-rank compression. It projects hidden states into a single, shared, low-dimensional latent vector that is then cached.
• This latent vector is projected back up on-the-fly for each attention head using unique weights, preserving the model's representational capacity and performance, unlike MQA or GQA which share heads.
• The result is a dramatic reduction in KV cache memory—up to ~57x in DeepSeek's case—without a corresponding drop in model quality. This makes MLA a true improvement to the Transformer architecture.

Preview of the Next Lesson:
We've now seen two powerful ways to optimize the attention mechanism: sharing K/V heads (GQA) and compressing the K/V representations (MLA). The other major computational component of a Transformer block is the Feed-Forward Network (FFN). In our next lesson, we will begin the "Scaling with Mixture-of-Experts" module and explore how Mixture-of-Experts (MoE) layers allow models to scale their parameter counts into the trillions without a proportional increase in compute cost per token.