Original Investigation: This article was originally published with interactive benchmarks and hardware telemetry at EyesTech Systems Research.
1. The 128k Context Memory Wall
Autoregressive transformer inference is split into two radically different computational regimes:
- The Prefill Phase: Processing input prompt tokens simultaneously. This phase is compute-bound, achieving high arithmetic intensity on Tensor Cores via dense General Matrix Multiply (GEMM) operations.
- The Decode Phase: Generating output tokens sequentially one-by-one. Each new token must attend to the Key and Value vectors of every preceding token. Arithmetic intensity collapses to $\approx 1$ FLOP per byte streamed. As a result, decoding is strictly memory-bandwidth bound.
To avoid recomputing Keys and Values at every autoregressive step $t$, inference runtimes store these vectors in High Bandwidth Memory (HBM). The memory footprint of the Key-Value (KV) cache scales linearly with sequence length $L$, batch size $B$, number of layers $n_l$, number of KV heads $n_{kv}$, and head dimension $d_h$:
$$\text{Memory}{\text{KV}} = 2 \times n_l \times n{kv} \times d_h \times p_{\text{bytes}} \times B \times L$$
Where $p_{\text{bytes}}$ is the precision in bytes ($2$ for FP16/BF16, $1$ for FP8).
The Per-Token Memory Footprint Across Modern LLMs
| Architecture | Model Baseline | Layers ($n_l$) | Query Heads ($n_h$) | KV Heads ($n_{kv}$) | Head Dim ($d_h$) | Precision | KV Cache / Token |
|---|---|---|---|---|---|---|---|
| Standard MHA | DeepSeek 67B Baseline | 60 | 128 | 128 | 128 | FP16 (2 B) | 3,932,160 Bytes (3.84 MB) |
| Standard MHA | Llama 2 70B (Hypothetical MHA) | 80 | 64 | 64 | 128 | FP16 (2 B) | 2,621,440 Bytes (2.50 MB) |
| GQA (8:1) | Llama 3 70B / 405B | 80 | 64 | 8 | 128 | FP16 (2 B) | 327,680 Bytes (320.0 KB) |
| GQA (4:1) | Mistral Large | 88 | 64 | 8 | 128 | FP16 (2 B) | 360,448 Bytes (352.0 KB) |
| DeepSeek MLA | DeepSeek-V2 / DeepSeek-V3 | 60 | 128 | — (Latent) | 576 scalars | FP16 (2 B) | 138,240 Bytes (135.0 KB) |
| DeepSeek MLA | DeepSeek-V2 / DeepSeek-V3 | 60 | 128 | — (Latent) | 576 scalars | FP8 (1 B) | 69,120 Bytes (67.5 KB) |
Total KV Cache at Scale ($B=1$, Sequence Length Scaling)
| Context ($L$) | DeepSeek 67B (MHA, FP16) | Llama 3 70B (GQA 8:1, FP16) | DeepSeek MLA (FP16) | DeepSeek MLA (FP8) |
|---|---|---|---|---|
| 8,192 (8k) | 30.72 GB | 2.56 GB | 1.08 GB | 0.54 GB |
| 32,768 (32k) | 122.88 GB | 10.24 GB | 4.32 GB | 2.16 GB |
| 65,536 (64k) | 245.76 GB | 20.48 GB | 8.64 GB | 4.32 GB |
| 131,072 (128k) | 503.32 GB | 40.96 GB | 17.28 GB | 8.64 GB |
On an 80GB NVIDIA H100 SXM5 GPU (3.35 TB/s peak HBM3 bandwidth):
- For Llama 3 70B in FP16, weights consume $\approx 140$ GB across tensor-parallel ranks. A single 128k context stream consumes 40.96 GB of KV cache.
- At a concurrent batch size of $B=8$, streaming $8 \times 40.96\text{ GB} = 327.68\text{ GB}$ per token generation step requires $\frac{327.68\text{ GB}}{3,350\text{ GB/s}} = \mathbf{97.8\text{ ms per token}}$ speed-of-light transfer time alone, capping generation to $\approx 10.2$ tokens/second while leaving 90%+ of Tensor Core FLOPS starved.
Multi-Query Attention (MQA) attempts to fix this by sharing a single KV head across all query heads ($n_{kv}=1$), but suffers severe representational degradation (3.8% to 6.2% score drops on GSM8k and multi-document recall).
DeepSeek solved this with Multi-Head Latent Attention (MLA).
2. Low-Rank Latent Compression
Instead of saving full Key and Value tensors for all 128 heads in memory, MLA projects the hidden state $h_t \in \mathbb{R}^d$ into a compressed latent coordinate space:
$$c_t^{KV} = W^{DKV} h_t$$
Where:
- $h_t \in \mathbb{R}^{5120}$ is the layer's input representation.
- $W^{DKV} \in \mathbb{R}^{d_c \times d}$ is the down-projection matrix.
- $d_c = 512$ is the compressed latent dimension.
During training, Keys and Values for all $n_h = 128$ attention heads are up-projected from this single latent vector:
$$k_{t,i}^C = W_{(i)}^{UK} c_t^{KV} \quad \in \mathbb{R}^{d_h}$$
$$v_{t,i}^C = W_{(i)}^{UV} c_t^{KV} \quad \in \mathbb{R}^{d_v}$$
Where $d_h = d_v = 128$ and $i \in [1, n_h]$.
Content Compression Ratio
In standard Multi-Head Attention ($n_h = 128, d_h = 128, d_v = 128$):
- Per-token Key dimension: $128 \times 128 = 16,384$ scalars.
- Per-token Value dimension: $128 \times 128 = 16,384$ scalars.
- Total uncompressed scalars: $16,384 + 16,384 = \mathbf{32,768\text{ scalars}}$.
In MLA:
- Compressed content latent $c_t^{KV}$: $\mathbf{512\text{ scalars}}$.
- Content compression ratio: $\frac{512}{32,768} = \frac{1}{64} = \mathbf{1.56\%}$ (a 98.44% reduction).
Because each head $i$ uses a distinct slice $W_{(i)}^{UK}$, the heads can express independent attention patterns across the 512-dimensional manifold, preventing the catastrophic rank collapse of MQA.
3. The RoPE Dilemma & Decoupled Positional Embeddings
Standard Rotary Positional Embeddings (RoPE) apply a position-dependent rotation matrix $\mathcal{R}_t \in \mathbb{R}^{d_h \times d_h}$ to each key vector:
$$\mathcal{R}t = \operatorname{diag}\left(R{\theta_1, t}, R_{\theta_2, t}, \dots, R_{\theta_{d_h/2}, t}\right)$$
If we naively apply RoPE to the up-projected keys:
$$k_{t,i} = \mathcal{R}t (W{(i)}^{UK} c_t^{KV})$$
When computing the dot product between query $q_t$ and historical key $k_s$:
$$\text{Score}{t,s,i} = (\mathcal{R}_t q{t,i})^T (\mathcal{R}s W{(i)}^{UK} c_s^{KV})$$
Notice the critical roadblock: $\mathcal{R}s$ and $W{(i)}^{UK}$ do not commute ($\mathcal{R}s W{(i)}^{UK} \neq W_{(i)}^{UK} \mathcal{R}s$). If you store only $c_s^{KV}$, your inference runtime would have to multiply $W{(i)}^{UK} c_s^{KV}$ and then apply $\mathcal{R}_s$ for all previous tokens at every single decoding step! This would completely erase any memory bandwidth savings.
DeepSeek's Solution: Decoupled RoPE
DeepSeek bifurcates keys and queries into separate content and positional streams:
- Content Stream ($k_{t,i}^C$): Dimension $d_h = 128$. Carries semantic information, derived from the compressed latent $c_t^{KV}$.
- RoPE Stream ($k_t^R$): Dimension $d_h^R = 64$. Carries positional information, generated directly via $W^{KR} h_t$ and rotated by $\mathcal{R}_t$.
Crucially, $k_t^R$ is shared across all 128 attention heads:
$$k_{t,i} = \begin{bmatrix} k_{t,i}^C \ k_t^R \end{bmatrix} \in \mathbb{R}^{128 + 64} = \mathbb{R}^{192}$$
$$q_{t,i} = \begin{bmatrix} q_{t,i}^C \ q_{t,i}^R \end{bmatrix} \in \mathbb{R}^{128 + 64} = \mathbb{R}^{192}$$
The attention dot product splits into two additive terms:
$$\text{Score}{t,s,i} = q{t,i}^T k_{s,i} = (q_{t,i}^C)^T k_{s,i}^C + (q_{t,i}^R)^T k_s^R$$
4. The Matrix Absorption Trick: Zero-Decompression Inference
Because $k_{s,i}^C$ is purely linear ($k_{s,i}^C = W_{(i)}^{UK} c_s^{KV}$) without any rotary matrix, we apply the associative property of matrix multiplication:
$$(q_{t,i}^C)^T k_{s,i}^C = (q_{t,i}^C)^T (W_{(i)}^{UK} c_s^{KV}) = \left( (W_{(i)}^{UK})^T q_{t,i}^C \right)^T c_s^{KV}$$
We define the absorbed query:
$$\tilde{q}{t,i}^C = (W{(i)}^{UK})^T q_{t,i}^C \quad \in \mathbb{R}^{512}$$
This absorption is computed once for the current token $t$. Then, the attention dot product is performed directly between $\tilde{q}_{t,i}^C$ and the cached 512-dim latent $c_s^{KV}$!
Similarly, for Values:
$$o_{t,i} = \sum_{s=1}^t \alpha_{t,s,i} (W_{(i)}^{UV} c_s^{KV}) = W_{(i)}^{UV} \left( \sum_{s=1}^t \alpha_{t,s,i} c_s^{KV} \right)$$
The attention weights $\alpha_{t,s,i}$ are multiplied directly against the 512-dimensional cached vectors $c_s^{KV}$. The up-projection $W_{(i)}^{UV}$ is fused offline into the output projection matrix $W_{(i)}^O$:
$$W_{(i)}^{OV} = W_{(i)}^O W_{(i)}^{UV} \in \mathbb{R}^{d \times d_c}$$
Total Cache Footprint per Token:
In MLA, the KV cache stores exclusively:
- $c_s^{KV} \in \mathbb{R}^{512}$ (compressed content latent)
- $k_s^R \in \mathbb{R}^{64}$ (shared RoPE key)
$$\text{Total Scalars per Token} = 512 + 64 = \mathbf{576\text{ scalars}}$$
- Standard 32-head MHA ($2 \times 32 \times 128 = 8,192$ scalars): $$\frac{8,192 - 576}{8,192} = \mathbf{92.97\% \approx 93\%\text{ reduction}}$$
- Full 128-head MHA ($2 \times 128 \times 128 = 32,768$ scalars): $$\frac{32,768 - 576}{32,768} = \mathbf{98.24\%\text{ reduction}}$$
- 8-head GQA ($2 \times 8 \times 128 = 2,048$ scalars): $$\frac{2,048 - 576}{2,048} = \mathbf{71.88\%\text{ reduction}}$$
5. Production PyTorch Reference Implementation
Here is the complete, runnable single-token decoding kernel with query absorption and decoupled RoPE:
import math
import torch
import torch.nn as nn
from typing import Tuple
class MultiHeadLatentAttentionDecode(nn.Module):
"""
Production-grade Multi-Head Latent Attention (MLA) Autoregressive Decoding Kernel
Demonstrating Query Absorption, Decoupled RoPE, and Zero-Decompression KV Cache Streaming.
Reference: https://github.com/abhishek2512mishra/deepseek-mla-kvcache
"""
def __init__(
self,
d_model: int = 5120,
n_heads: int = 128,
d_head: int = 128,
d_latent: int = 512,
d_rope: int = 64
):
super().__init__()
self.d_model = d_model
self.n_heads = n_heads
self.d_head = d_head
self.d_latent = d_latent # d_c (512 scalars)
self.d_rope = d_rope # d_h^R (64 scalars)
self.scale = 1.0 / math.sqrt(d_head + d_rope)
# 1. KV Down-Projection: Projects hidden state to shared latent space
self.W_DKV = nn.Linear(d_model, d_latent, bias=False)
# 2. KV Up-Projection Matrices (absorbed during decode)
self.W_UK = nn.Parameter(torch.empty(n_heads, d_head, d_latent))
self.W_UV = nn.Parameter(torch.empty(n_heads, d_head, d_latent))
# 3. Decoupled RoPE Key Projection (shared across all heads)
self.W_KR = nn.Linear(d_model, d_rope, bias=False)
# 4. Query Compression & Projections
self.W_DQ = nn.Linear(d_model, 1536, bias=False)
self.W_UQ = nn.Linear(1536, n_heads * d_head, bias=False)
self.W_QR = nn.Linear(1536, n_heads * d_rope, bias=False)
# 5. Output Projection
self.W_O = nn.Linear(n_heads * d_head, d_model, bias=False)
nn.init.normal_(self.W_UK, std=0.02)
nn.init.normal_(self.W_UV, std=0.02)
def apply_rope(self, x: torch.Tensor, pos: int) -> torch.Tensor:
"""Applies 1D Rotary Position Embedding to 2D coordinates."""
half_dim = x.shape[-1] // 2
freqs = torch.exp(-math.log(10000.0) * torch.arange(0, half_dim, device=x.device) / half_dim)
angles = pos * freqs
cos = torch.cos(angles).repeat(2)
sin = torch.sin(angles).repeat(2)
x_rot = torch.cat([-x[..., half_dim:], x[..., :half_dim]], dim=-1)
return (x * cos) + (x_rot * sin)
def forward_decode(
self,
h_t: torch.Tensor,
current_pos: int,
kv_cache_latent: torch.Tensor,
kv_cache_rope: torch.Tensor
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
"""
Single-token decode step with matrix absorption.
kv_cache_latent: [Batch, SeqLen, d_latent]
kv_cache_rope: [Batch, SeqLen, d_rope]
"""
B = h_t.shape[0]
# STEP 1: Compute Current Token Cache Entries (Only 576 scalars stored!)
c_t_kv = self.W_DKV(h_t) # [B, 1, 512]
k_t_rope = self.apply_rope(self.W_KR(h_t), current_pos) # [B, 1, 64]
# Append to persistent KV cache
kv_cache_latent = torch.cat([kv_cache_latent, c_t_kv], dim=1)
kv_cache_rope = torch.cat([kv_cache_rope, k_t_rope], dim=1)
# STEP 2: Ephemeral Query Processing
c_t_q = self.W_DQ(h_t) # [B, 1, 1536]
q_content = self.W_UQ(c_t_q).view(B, self.n_heads, self.d_head) # [B, 128, 128]
q_rope = self.W_QR(c_t_q).view(B, self.n_heads, self.d_rope) # [B, 128, 64]
q_rope = self.apply_rope(q_rope, current_pos)
# STEP 3: MATRIX ABSORPTION
# Project active Query into latent space: q_absorbed = q_content @ W_UK
# W_UK: [128, 128, 512] -> q_absorbed: [B, 128, 512]
q_absorbed = torch.einsum('bhd,hdm->bhm', q_content, self.W_UK)
# STEP 4: Direct Latent Attention Dot Product
# Content score computed in 512-dim latent space
score_content = torch.einsum('bhm,bsm->bhs', q_absorbed, kv_cache_latent)
# Positional score computed in 64-dim RoPE space
score_rope = torch.einsum('bhr,bsr->bhs', q_rope, kv_cache_rope)
attention_scores = (score_content + score_rope) * self.scale
attention_weights = torch.softmax(attention_scores, dim=-1) # [B, 128, SeqLen]
# STEP 5: Value Aggregation in Latent Space
# Sum attention weights directly against 512-dim cached latents
u_latent = torch.einsum('bhs,bsm->bhm', attention_weights, kv_cache_latent) # [B, 128, 512]
# Final projection via fused Value-Output matrix
v_projected = torch.einsum('bhm,hdm->bhd', u_latent, self.W_UV)
output = self.W_O(v_projected.reshape(B, 1, self.n_heads * self.d_head))
return output, kv_cache_latent, kv_cache_rope
6. Companion Open-Source Repository
The complete implementation, automated unit tests, and performance benchmarking harness are available in our open-source companion repository:
👉 github.com/abhishek2512mishra/deepseek-mla-kvcache
You can clone and run the verification suite locally:
git clone https://github.com/abhishek2512mishra/deepseek-mla-kvcache.git
cd deepseek-mla-kvcache
python3 mla_decode.py
Key Takeaways for Systems Engineers
- Memory Bandwidth Governs Decode Throughput: The KV cache memory footprint is the single biggest bottleneck for long-context inference ($L \ge 32\text{k}$).
- Latent Representation Over Head Pruning: Instead of discarding heads as in MQA or GQA, MLA retains 128 heads of expressive attention while compressing the stored representations by 93%.
- Associativity Is Free Speed: By absorbing the up-projection weight matrix $W^{UK}$ into the ephemeral query $q_t$, you eliminate the need to decompress historical KV pairs in HBM.
- Decoupled RoPE Is Required: Applying positional encodings to a separate 64-dimensional channel preserves spatial awareness without compromising low-rank matrix absorption.
For complete mathematical proofs, hardware profiling, and interactive latency curves, visit the original publication at EyesTech Systems Research.
Top comments (0)