Setup

• Model: L layers, hidden dimension \(d\) = 7168, num_heads \(h\) = 128, head_dim \(d_h\) = 128 (so \(d = h × d_h\), for deepseek V3 like setting)
• MLA latent dimension: \(d_c\) = 512 (compressed KV latent)
• Sequence: tokens t₁, t₂, ..., t₁₀₀
• Assume t₁ to t₉₉ are already cached (prefill done). Now we process t₁₀₀ (decode step).

---

Standard MHA: Step-by-step for t₁₀₀

Step 1: Embedding

x₁₀₀ = Embedding(t₁₀₀)    # shape: (d,) = (7168,)

Step 2: Layer 1

Compute Q, K, V for the new token:

q₁₀₀⁽¹⁾ = x₁₀₀ · W_Q⁽¹⁾     # (7168,) × (7168, 7168) → (7168,) i.e. (h × d_h)
k₁₀₀⁽¹⁾ = x₁₀₀ · W_K⁽¹⁾     # (7168,) × (7168, 7168) → (7168,)
v₁₀₀⁽¹⁾ = x₁₀₀ · W_V⁽¹⁾     # (7168,) × (7168, 7168) → (7168,)

Concat with cached K, V from tokens 1~99:

K⁽¹⁾ = concat(KV_Cache_K⁽¹⁾, k₁₀₀⁽¹⁾)   # (99, 7168) + (1, 7168) → (100, 7168)
V⁽¹⁾ = concat(KV_Cache_V⁽¹⁾, v₁₀₀⁽¹⁾)   # (99, 7168) + (1, 7168) → (100, 7168)

Update cache:

KV_Cache_K⁽¹⁾ ← K⁽¹⁾    # now stores (100, 7168)
KV_Cache_V⁽¹⁾ ← V⁽¹⁾    # now stores (100, 7168)

Compute attention (per head, d_h = 128):

# Reshape to heads: Q → (h, 1, d_h), K → (h, 100, d_h), V → (h, 100, d_h)
attn_weights = softmax(q₁₀₀⁽¹⁾ · K⁽¹⁾ᵀ / √d_h)   # (h, 1, 100)
attn_output = attn_weights · V⁽¹⁾                     # (h, 1, d_h) → concat → (7168,)

FFN:

h₁₀₀⁽¹⁾ = FFN(LayerNorm(attn_output + x₁₀₀))   # (7168,)

Step 3: Layer 2

Same as Layer 1, but input is h₁₀₀⁽¹⁾:

q₁₀₀⁽²⁾ = h₁₀₀⁽¹⁾ · W_Q⁽²⁾
k₁₀₀⁽²⁾ = h₁₀₀⁽¹⁾ · W_K⁽²⁾
v₁₀₀⁽²⁾ = h₁₀₀⁽¹⁾ · W_V⁽²⁾

K⁽²⁾ = concat(KV_Cache_K⁽²⁾, k₁₀₀⁽²⁾)   # (100, 7168)
V⁽²⁾ = concat(KV_Cache_V⁽²⁾, v₁₀₀⁽²⁾)   # (100, 7168)
... (same attention + FFN) ...

Step 4: Repeat for Layers 3 ~ L

Each layer independently caches its own K and V.

Step 5: Output

logits = h₁₀₀⁽ᴸ⁾ · W_head    # → vocabulary distribution
t₁₀₁ = argmax(logits)

Cache Size (MHA)

Total cache = L × n × 2 × d
            = L × 99 × 2 × 7168

For L=60, n=99:

= 60 × 99 × 2 × 7168 = 85,155,840 parameters (per sequence)
≈ 85M × 2 bytes (FP16) ≈ 162 MB

---

MLA (Multi-head Latent Attention): Step-by-step for t₁₀₀

Step 1: Embedding

x₁₀₀ = Embedding(t₁₀₀)    # shape: (d,) = (7168,)

Step 2: Layer 1

Down-project to compressed latent (for KV):

c₁₀₀⁽¹⁾ = x₁₀₀ · W_DKV⁽¹⁾    # (7168,) × (7168, 512) → (512,)

Compute Q separately:

q₁₀₀⁽¹⁾ = x₁₀₀ · W_Q⁽¹⁾      # (7168,) × (7168, 7168) → (7168,)

Update latent cache (only store c, NOT full K,V):

Latent_Cache⁽¹⁾ ← concat(Latent_Cache⁽¹⁾, c₁₀₀⁽¹⁾)   # (99, 512) + (1, 512) → (100, 512)

Up-project ALL cached latents to K, V (on-the-fly):

K⁽¹⁾ = Latent_Cache⁽¹⁾ · W_UK⁽¹⁾    # (100, 512) × (512, 7168) → (100, 7168)
V⁽¹⁾ = Latent_Cache⁽¹⁾ · W_UV⁽¹⁾    # (100, 512) × (512, 7168) → (100, 7168)

Compute attention (same as MHA from here):

attn_weights = softmax(q₁₀₀⁽¹⁾ · K⁽¹⁾ᵀ / √d_h)   # (h, 1, 100)
attn_output = attn_weights · V⁽¹⁾                     # → (7168,)

FFN:

h₁₀₀⁽¹⁾ = FFN(LayerNorm(attn_output + x₁₀₀))   # (7168,)

Step 3: Layer 2

Same structure, input is h₁₀₀⁽¹⁾:

c₁₀₀⁽²⁾ = h₁₀₀⁽¹⁾ · W_DKV⁽²⁾    # → (512,)
q₁₀₀⁽²⁾ = h₁₀₀⁽¹⁾ · W_Q⁽²⁾      # → (7168,)

Latent_Cache⁽²⁾ ← concat(Latent_Cache⁽²⁾, c₁₀₀⁽²⁾)  # (100, 512)

K⁽²⁾ = Latent_Cache⁽²⁾ · W_UK⁽²⁾    # (100, 7168)
V⁽²⁾ = Latent_Cache⁽²⁾ · W_UV⁽²⁾    # (100, 7168)
... (attention + FFN) ...

Step 4: Repeat for Layers 3 ~ L

Step 5: Output

Same as MHA.

Cache Size (MLA)

Total cache = L × n × d_c
            = L × 99 × 512

For L=60, n=99:

= 60 × 99 × 512 = 3,041,280 parameters (per sequence)
≈ 3M × 2 bytes (FP16) ≈ 5.8 MB

---

Side-by-side Comparison

	Standard MHA	MLA
What is cached per token per layer	k (7168) + v (7168) = 14,336	c (512)
Cache size per token per layer	2 × d = 14,336	d_c = 512
Compression ratio	1× (baseline)	d_c / (2d) = 512/14336 ≈ 1/28
Total cache (L=60, n=99)	~162 MB	~5.8 MB
Extra compute at decode	None	Up-project c → K, V per step
Accuracy trade-off	Full rank	Low-rank approximation (minimal loss in practice)

---

Key Insight

MHA caches the full K and V vectors for every token at every layer. This grows linearly with sequence length and is the main memory bottleneck for long-context inference.

MLA compresses the KV information into a much smaller latent vector c via a learned down-projection \(W_{DKV}\) (see blog before). At attention time, it reconstructs K and V by up-projecting c with \(W_{UK}\) and \(W_{UV}\). The trade-off is a small amount of extra compute (matrix multiply for up-projection) in exchange for ~28× less cache memory.

This is why DeepSeek V3 can handle much longer contexts with the same GPU memory budget.

---

Note on Actual Implementation

In practice, DeepSeek V3's MLA also handles the Q side with compression (down-project then up-project for Q as well), and uses RoPE (Rotary Position Embedding) on a separate small portion of the key. The above is simplified to show the core KV cache difference.