DEV Community: Sirajuddin Shaik

Megatron Distributed Training On 16 RTX PRO 6000 GPUs

Sirajuddin Shaik — Mon, 06 Jul 2026 16:57:17 +0000

This is a standalone technical writeup about building and testing a >Megatron distributed training setup across a 2-node, 16-GPU RTX PRO >6000 Blackwell cluster.

The goal was not to produce a final high-quality model. The goal was to understand the infrastructure layer that makes larger training possible: model import, distributed launch, parallelism layout, forward/backward correctness, checkpointing, resharding, and failure modes.

Scope

The focus was training infrastructure, not final model quality.

Quality depends on data, recipes, hyperparameters, evals, token budget, and repeated ablations. These tests focused on whether the distributed training stack works:

launch
model construction
checkpoint import
checkpoint resharding
forward/backward
gradient reduction
no NaNs/skips
checkpoint save
repeatable config recording

In other words, the question was:

Can a model be converted, split across GPUs, trained for a few steps, synchronized correctly, saved, and then used as a base for longer continued-pretraining or supervised fine-tuning runs?

Hardware / Cluster

2 nodes.
8x RTX PRO 6000 Blackwell GPUs per node.
16 GPUs total.
Each node used node-local NVMe storage for datasets, checkpoints, and logs.
Private internode networking was used for NCCL/distributed training.
Most larger model smoke tests were run on one free 8-GPU node.
The public WikiText 5D run used both nodes and all 16 GPUs.

Stack

Megatron-LM: training runtime, launch scripts, model-parallel groups, distributed optimizer, and checkpoint flow.
Megatron-Bridge: Hugging Face to Megatron conversion, model providers, recipes, and training entrypoints.
Megatron-Core: tensor, pipeline, context, expert, and sequence-parallel primitives used by the models.
Transformer Engine: optimized transformer kernels, FP8 paths, attention kernels, and communication overlap.
NCCL: GPU collectives for tensor parallelism, data parallelism, expert routing, and multinode communication.
Torch distributed checkpointing: sharded checkpoint save/load and resharding across different layouts.
Hugging Face integration: public model/tokenizer download and checkpoint import.
BF16, FP8, and MXFP8 precision paths.

How The Pieces Fit Together

The stack has three main layers:

Model source format

Public model checkpoints usually start as Hugging Face snapshots. That format is convenient for inference and ecosystem compatibility, but it is not the same layout Megatron uses for distributed training.

Conversion and model provider layer

Megatron-Bridge maps the Hugging Face model into Megatron's model-provider and checkpoint format. This is where architecture-specific details matter: MoE routing, expert layout, attention variants, Mamba blocks, tokenizer files, precision config, and the target parallelism layout.

Distributed training layer

Megatron-LM/Megatron-Core then train the model using explicit parallel groups. The same model can be loaded with different runtime layouts, such as TP=2, PP=2, or EP=4, as long as the model architecture and checkpoint layout support it.

The practical flow looked like this:

Hugging Face model/tokenizer
        |
        v
Megatron-Bridge import / model provider
        |
        v
Megatron-format checkpoint
        |
        v
Megatron distributed launch
        |
        v
train -> validate -> save checkpoint -> inspect logs

How A Megatron Run Works

A useful Megatron smoke run is more than "the script started." The minimum proof is:

The distributed process group initializes across all ranks.
The model builds with the requested TP, PP, CP, DP, EP, ETP, EDP, and SP layout.
The dataset and tokenizer paths resolve on every participating node.
A forward pass runs without backend or shape errors.
Backward pass and gradient synchronization complete.
No NaNs or skipped iterations appear.
A checkpoint is written in the expected distributed layout.
Validation/test can run when enabled.

This is why short runs are still valuable. They test the infrastructure contract before spending time on long training.

How Megatron-Bridge Was Used

Megatron-Bridge was useful because the experiments involved public Hugging Face models, but the training target was Megatron. Bridge handled the important middle layer:

Load public model configuration and tokenizer assets.
Map Hugging Face parameters into Megatron parameter names and tensor layouts.
Create a Megatron-compatible checkpoint.
Support runtime resharding when the training layout differs from the import layout.
Provide model-specific recipes and providers for architectures like DeepSeek MoE and Nemotron hybrid models.

The practical learning: conversion is not a side task. For large-model training, conversion and checkpoint layout are part of the training system.

Parallelism Mental Model

Megatron is useful because it exposes several parallelism axes directly:

Axis	What It Splits	Why It Matters
`TP`	Dense tensor operations inside layers	Reduces per-GPU memory and compute for large matrix operations
`PP`	Model layers across pipeline stages	Lets different GPUs hold different layer blocks
`CP`	Long sequence/context tokens	Enables longer-context training by splitting the token axis
`DP`	Data batches over replicated model groups	Improves throughput while synchronizing gradients
`EP`	MoE experts across GPUs	Prevents every GPU from storing every expert
`ETP`	Weights inside each expert	Helps when individual experts are too large
`EDP`	Expert-parallel groups over data	Adds data-parallel throughput for expert groups
`SP`	Temporary activations inside tensor-parallel groups	Saves activation memory; does not add a new world-size dimension

The important point is that these axes are not just multiplication. A mathematically valid world size can still fail because of layer counts, attention backend support, kernel availability, memory pressure, or checkpoint layout.

For example:

PP=2 failed on DeepSeek-V2-Lite until uneven pipeline splitting was used, because the model has 27 layers.
CP=2 worked on the custom MoE GPT and Nemotron paths, but DeepSeek-V2-Lite's MLA attention path hit a backend issue.
SP helped memory with TP, but it did not add a new rank dimension.

Checkpointing And Resharding

Checkpointing is part of the infrastructure test. A run that finishes forward/backward but cannot save is not a complete distributed training proof.

Two checkpoint ideas mattered:

Import checkpoint: the Megatron-format checkpoint created from a public Hugging Face model.
Runtime checkpoint: the checkpoint written after distributed training with a specific parallelism layout.

Megatron can reshard between layouts, but the surrounding storage and metadata still need to be clean. The 16-GPU public WikiText run proved fresh train/save/eval. Resume from a temporary node-local checkpoint layout still needs cleanup because it timed out after metadata load.

Glossary

Term	Meaning
CPT	Continued pretraining: next-token training on text data
SFT	Supervised fine-tuning: prompt/response style training
Smoke test	Short run to prove launch, distributed communication, forward/backward, checkpoint save, and basic stability
MoE	Mixture of Experts: model layers route tokens to selected expert networks
Checkpoint resharding	Loading/saving checkpoints across different tensor, pipeline, expert, or data-parallel layouts

Why Megatron For This Work

FSDP is useful for many dense-model and fine-tuning jobs. Megatron was the better target here because the goal was explicit topology control:

TP: tensor parallelism, splitting dense tensor operations.
PP: pipeline parallelism, splitting layers across stages.
CP: context parallelism, splitting long sequence/context tokens.
DP: data parallelism, replicating model groups across data batches.
EP: expert parallelism, distributing MoE experts across GPUs.
ETP: expert tensor parallelism, sharding weights inside each expert.
EDP: expert data parallelism, replicating expert-parallel groups.
SP: sequence parallelism, reducing activation memory inside TP groups.

For MoE and hybrid models, the bottleneck is not only parameter memory. It can be dense weights, routed experts, sequence length, optimizer state, checkpoint layout, or communication. Megatron exposes those axes directly.

Experiments Run

Model / Task	Run name	Config	Precision	Iters	Seq	Result
Public WikiText 5D continued pretraining	Public WikiText 5D	dense `TP=2 PP=2 CP=2 DP=2`, expert `ETP=1 EP=4 EDP=2`, `PP=2`	BF16 + TE + FlashAttention	20	-	passed train/save/eval on 16 GPUs
DeepSeek-V2-Lite import	`deepseek_v2_lite_ep8_import`	`TP=1 PP=1 EP=8 ETP=1`	import	-	-	passed
DeepSeek-V2-Lite continued pretraining	`deepseek_v2_lite_ep8_bf16_smoke`	`TP=1 PP=1 EP=8 CP=1`	BF16	4	512	passed
DeepSeek-V2-Lite continued pretraining	`deepseek_v2_lite_ep8_fp8_smoke`	`TP=1 PP=1 EP=8 CP=1`	FP8 delayed	4	512	passed
DeepSeek-V2-Lite continued pretraining	`deepseek_v2_lite_tp2_ep4_sp_fp8_smoke`	`TP=2 PP=1 EP=4 CP=1 SP=on`	FP8 delayed	2	512	passed
DeepSeek-V2-Lite continued pretraining	`deepseek_v2_lite_tp4_ep2_sp_fp8_smoke`	`TP=4 PP=1 EP=2 CP=1 SP=on`	FP8 delayed	2	512	passed
DeepSeek-V2-Lite continued pretraining	`deepseek_v2_lite_tp2_pp2_ep2_uneven_sp_fp8_smoke`	`TP=2 PP=2 EP=2 CP=1 SP=on`	FP8 delayed	2	512	passed
DeepSeek-V2-Lite supervised fine-tuning	`sft_deepseek_v2_lite_ep8_smoke_20260701_retry1`	`TP=1 PP=1 EP=8 CP=1`	BF16	2	512	passed
DeepSeek-V2-Lite supervised fine-tuning	`sft_deepseek_v2_lite_tp2_pp2_ep2_uneven_sp_smoke_20260701`	`TP=2 PP=2 EP=2 CP=1 SP=on`	BF16	2	512	passed
DeepSeek-V2-Lite continued pretraining	`cpt_deepseek_v2_lite_tp2_pp2_ep2_uneven_sp_bf16_20260701`	`TP=2 PP=2 EP=2 CP=1 SP=on`	BF16	2	512	passed
Custom MoE GPT continued pretraining	`cpt_wikitext_tp2_cp2_ep4_edp2_smoke_20260701`	dense `TP=2 PP=1 CP=2 DP=2`, expert `ETP=1 EP=4 EDP=2`	BF16	2	-	passed train/save/validation/test
Nemotron-3-Nano continued pretraining	`nemotron_3_nano_ep8_bf16_smoke_20260701_retry2`	`TP=1 PP=1 EP=8 ETP=1 CP=1`	BF16	2	512	passed
Nemotron-3-Nano continued pretraining	`nemotron_3_nano_tp2_ep4_sp_bf16_smoke_20260701_retry1`	`TP=2 PP=1 EP=4 ETP=1 CP=1 SP=on`	BF16	1	256	passed
Nemotron-3-Nano continued pretraining	`nemotron_3_nano_pp2_ep4_bf16_smoke_20260701`	`TP=1 PP=2 EP=4 ETP=1 CP=1`	BF16	1	256	passed
Nemotron-3-Nano continued pretraining	`nemotron_3_nano_pp2_cp2_ep4_bf16_smoke_20260701`	`TP=1 PP=2 EP=4 ETP=1 CP=2`	BF16	1	256	passed
Nemotron-3-Nano continued pretraining	`nemotron_3_nano_pp2_ep2_etp2_bf16_smoke_20260701`	`TP=1 PP=2 EP=2 ETP=2 CP=1`	BF16	1	256	passed

Runtime Observations From Training Logs

These are rank-0 memory observations from the training logs. They are not full performance benchmarks, but they show why different parallelism layouts mattered.

Model / Layout	Logged Rank-0 Memory	What It Showed
DeepSeek-V2-Lite `EP=8` BF16	about `42.6-43.2 GB` reserved	EP-only baseline fit comfortably on one 8-GPU node
DeepSeek-V2-Lite `TP=2 PP=2 EP=2 SP` BF16/FP8	about `37-41 GB` reserved depending precision/config	uneven PP plus TP/EP gave better memory shape than the EP-only layout
Nemotron-3-Nano `EP=8` BF16	about `82.9-85.9 GB` reserved	the hybrid Mamba/MoE model fit, but used much more memory
Nemotron-3-Nano `PP=2 EP=4` BF16	about `73.5 GB` reserved	pipeline parallelism gave useful memory headroom
Nemotron-3-Nano `TP=2 EP=4 SP` BF16	about `98.7 GB` reserved	this configuration worked but was too tight to treat as a comfortable production layout

The strongest practical lesson was that PP was not just a throughput knob. For the hybrid Nemotron model, PP=2 was a major memory tool. It made some layouts much more comfortable than pure TP+EP.

Model Notes And Learnings

Public WikiText 16-GPU 5D Run

Dataset: Hugging Face Salesforce/wikitext, config wikitext-2-raw-v1.
Tokenizer files: local Qwen3 tokenizer files from Qwen/Qwen3-30B-A3B.
Nodes: 2.
Total GPUs: 16.
Dense dimensions: TP=2, PP=2, CP=2, DP=2.
Expert dimensions: ETP=1, EP=4, EDP=2, with PP=2.
Experts: 128 total, 32 local experts per EP shard.
Control model: 4 layers, hidden size 256, FFN 512, 8 attention heads, 2 query groups.
Precision/path: BF16, Transformer Engine, FlashAttention.
Result:
- trained iterations 1 through 20
- saved checkpoints at iterations 10 and 20
- ran validation/test
- memory stayed small, under about 0.9 GB reserved on shown ranks
Important caveat: fresh train/save/eval worked; resume from the node-local checkpoint layout hit a post-load NCCL timeout and remains unresolved.

DeepSeek-V2-Lite

HF model: deepseek-ai/DeepSeek-V2-Lite
Imported checkpoint: internal Megatron-format checkpoint created from the public Hugging Face snapshot.
Snapshot/checkpoint size: about 30G
Public scale:
- Total parameters: 16B
- Activated parameters per token: 2.4B
Important fix: PP=2 needs uneven layer split because the model has 27 layers.
- FIRST_PP_LAYERS=14
- LAST_PP_LAYERS=13
What this validated: checkpoint import, BF16, FP8 delayed scaling, expert parallelism, tensor parallelism, sequence parallelism, supervised fine-tuning flow, and continued-pretraining flow.

Nemotron-3-Nano

Model/recipe: nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16
Architecture: Mamba2-Transformer Hybrid Mixture of Experts.
HF tokenizer/small files: local copy of the public Hugging Face model assets.
Public scale:
- Total parameters: 30B
- Activated parameters per token: about 3B
Mamba fast path was disabled because causal-conv1d was not installed.
TP=2 needed UB_SKIPMC=1 due Transformer Engine CUDA multicast overlap.
What this validated: Megatron training on a Mamba2-Transformer Hybrid MoE architecture across EP, TP, PP, CP, and ETP combinations.

DeepSeek-V4-Flash

Target: deepseek-ai/DeepSeek-V4-Flash
Probe checkpoint: internal Megatron-format development checkpoint.
Result: reached distributed Megatron setup but did not fit full training on one 8x96 GB node.
What this validated: launch, config construction, distributed setup, and MXFP8 path exploration before hitting the expected memory wall.

Issues Faced And Fixes

Area	Issue	Fix / Status
Optional imports	Bridge recipe imports pulled optional packages like `diffusers`, `nvidia_resiliency_ext`, `megatron.energon`	Used narrower recipe loading and small stubs where safe
Python helper build	Dataset helper could pick the wrong Python environment	Added launcher shim so helper builds use the intended virtual environment and pybind11 include
Checkpoint import	Needed HF to Megatron import before training	DeepSeek-V2-Lite import passed
Checkpoint resharding	Runtime TP/PP differed from imported checkpoint	Megatron resharding worked for passing DeepSeek-V2-Lite runs
Multinode dataset cache	GPTDataset cache files existed on one node but not the other	Staged dataset/tokenizer under identical absolute paths on both nodes and built local cache on each node
Multinode checkpoint resume	Fresh 16-GPU train/save/eval worked, but resume from temporary node-local checkpoint layout timed out after metadata load	Treat resume as unresolved; likely needs shared/synced persistent checkpoint view and clean tracker metadata
Pipeline layers	DeepSeek-V2-Lite has 27 layers	Used uneven PP: 14/13
Context parallelism	DeepSeek-V2-Lite `CP=2` failed in MLA attention path	CP validated on custom MoE GPT and Nemotron; DeepSeek MLA remains blocked
Mamba fast path	Missing `causal-conv1d`	Disabled memory-efficient Mamba path; fallback worked
TE TP overlap	CUDA multicast unsupported for Nemotron TP run	`UB_SKIPMC=1` fixed it
Memory	DeepSeek-V4-Flash full training did not fit one node	Needs more GPUs, more parallelism, PEFT/freezing, or smaller model
Memory	Nemotron `TP=2 EP=4` was tight	Passed, but not a comfortable production config

What Was Proven

Axis	Status
`TP`	Proven on DeepSeek-V2-Lite and Nemotron-3-Nano
`PP`	Proven on DeepSeek-V2-Lite and Nemotron-3-Nano
`CP`	Proven on 16-GPU public WikiText, custom MoE GPT, and Nemotron-3-Nano
`DP`	Proven explicitly in the 16-GPU public WikiText run and custom MoE GPT
`EP`	Proven across 16-GPU public WikiText, DeepSeek-V2-Lite, custom MoE GPT, and Nemotron-3-Nano
`ETP`	Proven on Nemotron-3-Nano
`EDP`	Proven on 16-GPU public WikiText and custom MoE GPT
`SP`	Proven on TP configs

Takeaways

Megatron is valuable because it makes training topology explicit. The work moved from "can I launch a model?" to "can I control how a large model is split, trained, checkpointed, and debugged?"

The main practical lessons:

Megatron-Bridge conversion and checkpoint layout are part of the training system.
A distributed training proof should include checkpoint save, not only forward/backward.
MoE models need separate reasoning about dense weights, routed experts, expert routing, and optimizer state.
Context parallelism is model-path dependent; it can work in one architecture and fail in another backend path.
Pipeline parallelism is not only for throughput. It can be the cleanest way to create memory headroom.
Short smoke runs are useful when they validate the full infrastructure path and expose the next bottleneck.

Mamba/SSM Basics

Sirajuddin Shaik — Sun, 31 May 2026 13:50:41 +0000

State Space Models offer linear-time sequence modeling with content-aware selective filtering, challenging Transformers for long-context inference.

Why This Matters

State Space Models (SSMs) provide a principled alternative to Transformers for long-sequence modeling. In production systems handling long contexts (e.g., code generation, genomic analysis), Transformer attention's quadratic cost becomes a bottleneck. Mamba achieves linear-time inference with constant-memory state, making it viable for million-token contexts where attention-based models are prohibitively expensive.

Core Idea

SSMs originate from continuous-time control theory: a latent state evolves over time driven by input, and observations are linear projections of that state. Mamba's key innovation is making the SSM parameters input-selective — the model learns to gate which information enters and exits the state, mimicking attention's ability to focus on relevant tokens without the $O (n^{2})$ cost.

Technical Details

The continuous-time SSM is defined as:

x^{'} (t) = A x (t) + B u (t), y (t) = C x (t) + D u (t)

where $x (t) \in R^{N}$ is latent state, $u (t)$ is input, and $A \in R^{N \times N}$ , $B \in R^{N \times 1}$ , $C \in R^{1 \times N}$ . Using zero-order hold discretization with step $Δ$ :

A ˉ = exp (Δ A), B ˉ = (Δ A)^{- 1} (exp (Δ A) - I) \cdot Δ B

The recurrent update becomes:

x_{k} = A ˉ x_{k - 1} + B ˉ u_{k}, y_{k} = C x_{k}

Mamba's selective mechanism makes $B$ , $C$ , and $Δ$ input-dependent:

B_{k} = Linear B (x_{k}), C_{k} = Linear_{C} (x_{k}), Δ_{k} = softplus (Linear Δ (x_{k}))

The parallel scan algorithm computes this recurrence in $O (n lo g n)$ during training. Inference is $O (1)$ O(1) per token with fixed state size N $N$ , yielding constant-memory decoding regardless of sequence length.

How It Works

Project input: Map token $u_{k}$ to expanded dimension $D \cdot N$ .
Generate selective parameters: Compute input-dependent $B_{k}$ , $C_{k}$ , $Δ_{k}$ from $u_{k}$ .
Discretize: Convert continuous $(A, B)$ to discrete $(A ˉ_{k}, B ˉ_{k})$ using $Δ_{k}$ .
Recurrent scan: Apply parallel scan (training) or sequential update (inference) to compute hidden states $x_{k}$ .
Output projection: Compute $y_{k} = C_{k} x_{k}$ , then project through gating (SiLU) to output dimension.

Key Insights

Selectivity is essential: Non-selective SSMs (S4) cannot do in-context retrieval; making $B, C, Δ$ input-dependent enables content-aware filtering.
Diagonal + low-rank structure on $A$ enables $O (n)$ recurrence; Mamba uses diagonal $A$ matrices exclusively.
Hardware-aware design: The scan kernel is IO-bound, not compute-bound — Mamba's CUDA kernel fuses discretization, scan, and output projection to minimize memory reads.
Linear decoding cost: Unlike KV-cache which grows linearly, SSM state is fixed-size $O (N D)$ , making generation memory-constant.

Sources

Gu, A. & Dao, T. "Mamba: Linear-Time Sequence Modeling with Selective State Spaces." arXiv:2312.00752 (2023). https://arxiv.org/abs/2312.00752
Gu, A. & Dao, T. "Mamba-3: Improved Sequence Modeling using State Space Principles." arXiv:2603.15569 (2026). https://arxiv.org/abs/2603.15569

Multi-Head Latent Attention (MLA)

Sirajuddin Shaik — Sat, 23 May 2026 14:33:01 +0000

Compressing KV cache via low-rank projections - the attention mechanism behind DeepSeek-V2/V3 and Kimi K2.x

Why This Matters

Multi-Head Latent Attention (MLA) is the attention variant that replaces standard Multi-Head Attention (MHA) in DeepSeek-V2, DeepSeek-V3, and Kimi K2.x models. Instead of caching full KV pairs per head, MLA projects them into a low-dimensional latent space, achieving 3-5× KV cache compression with minimal quality loss.

MLA changes how prefix caching, chunked prefill, and paged attention must be implemented

Formal Definition

Standard Multi-Head Attention (MHA)

For input $X \in R^{n \times d}$ , MHA computes per-head projections:

Q_{h} = X W_{Q (h)}, K_{h} = X W_{K (h)}, V_{h} = X W_{V (h)}

where $W_{Q (h)} \in R^{d \times d_{k}}$ , $W_{K (h)} \in R^{d \times d_{k}}$ , $W_{V (h)} \in R^{d \times d_{v}}$ .

KV cache size per token: $2 \times n_{h} \times d_{k}$ elements.

MLA: Low-Rank Latent Projection

MLA replaces the per-head KV projections with a shared low-rank latent compression:

Compression (KV → Latent):

c^{K V} = X W_{DK V} \in R^{n \times d_{c}}

where $W_{DK V} \in R^{d \times d_{c}}$ is the down-projection matrix and $d_{c} ≪ n_{h} \times d_{k}$ .

Decompression (Latent → KV):

K h = c^{K V} W U K^{(h)}, V h = c^{K V} W U V^{(h)}

where $W U K^{(h)} \in R^{d_{c} \times d_{k}}$ and $W U V^{(h)} \in R^{d_{c} \times d_{v}}$ are up-projection matrices.

KV cache per token: Only $c^{K V} \in R^{d_{c}}$ is stored - a single vector of dimension $d_{c}$ .

Query Compression (Optional)

MLA also compresses queries for training efficiency:

c^{Q} = XW D Q \in R^{n \times d_{c'}}

Q_{h} = c^{Q} W U Q^{(h)}

This doesn't affect the KV cache but reduces the activation memory during training.

Rotary Position Embedding (RoPE) Handling

RoPE is applied to the decompressed queries and keys. To keep the KV cache small, MLA applies RoPE to a separate "absorbed" key projection:

K^h = RoPE (c^{K V} W K R^{(h)})

where $W_{K R (h)} \in R^{d_{c} \times d_{r}}$ with $d_{r} ≪ d_{k}$ is a narrow projection that carries positional information. The cached representation remains $c^{K V}$ (position-agnostic), and the RoPE key $K^_{h}$ is recomputed at attention time from the cached latent.

Compression Ratio

For a model with $n_{h}$ heads and head dimension $d_{k}$ :

The total MLA KV cache per token stores:

The latent vector $c^{K V}$ of dimension $d_{c}$
The decoupled RoPE keys for all heads: $n_{h} \times d_{r}$

Compression Ratio = \frac{2 \cdot n _{h} \cdot d _{k}}{d _{c} + n _{h} \cdot d _{r}}

Approximating DeepSeek-V3 dimensions (see paper Table 1 for exact values):
$n_{h} = 128$ , $d_{k} = 128$ , $d_{c} = 512$ , $d_{r} = 64$ :

\frac{2 \times 128 \times 128}{512 + 128 \times 64} = \frac{32 , 768}{8 , 704} \approx 3.76 \times compression

Real-world impact: With 3.76× less cache per layer, a 60-layer model at 64K context drops from ~1TB to ~266GB of KV cache. That's the difference between needing a DGX H200 cluster versus a single 8×H100 node.

Note: Reported "5-10×" compression figures in literature often come from models with smaller $d_{r}$ (e.g., 32 or fewer) or measure only the $c^{K V}$ component. The practical ratio depends on the specific dimension choices.

Core Concepts

1. Weight Absorption (The Key Trick)

The critical insight in MLA is that the up-projection matrices $W_{U K (h)}$ can be absorbed into the query projection during attention computation:

Attention (Q, K, V) = softmax (\frac{Q _{h} K _{h T}}{d _{k}}) V_{h}

Substituting the decompressed forms:

Q h K_{h T} = (c^{Q} W U Q^{(h)}) (c^{K V} W U K^{(h)})^{T} = c^{Q} (W U Q^{(h)} W_{U K (h)}^{T}) c^{K V}^{T}

If we define $W ab sor b e d^{(h)} = W U Q^{(h)} W_{U K (h)}^{T} \in R^{d_{c'} \times d_{c}}$ , then:

Q h K_{h T} = c^{Q} W ab sor b e d^{(h)} c^{K V}^{T}

This means the attention score can be computed directly from the latent representations, avoiding explicit decompression of K and V for the score computation. However, the V decompression is still needed for the output.

Practical implication: During decoding, we can compute attention scores without materializing the full K matrix. Only V needs decompression after softmax.

2. Decoupled RoPE Strategy

RoPE requires position-dependent keys, which conflicts with caching a position-agnostic latent. MLA solves this with a decoupled key:

Content key: $K h^{co n t e n t} = c^{K V} W U K^{(h)}$ - cached in latent form
Position key: $K h^{ro p e} = RoPE (c^{K V} W K R^{(h)})$ - small, position-aware, must be cached separately

The attention score becomes:

score (q, k) = \frac{Q _{h co n t e n t} \cdot K _{h co n t e n t} ^{T}}{d _{k}} + \frac{Q _{h ro p e} \cdot K _{h ro p e} ^{T}}{d _{r}}

Practical implication: The KV cache stores both $c^{K V}$ (latent) and $K_{h ro p e}$ (decoupled rope key). Total cache per token: $d_{c} + n_{h} \times d_{r}$ .

3. MLA vs GQA vs MHA

Property	MHA	GQA	MLA
KV groups	$n_{h}$	$n_{h} / g$	1 (latent)
Cache per token	$2 n_{h} d_{k}$	$2 (n_{h} / g) d_{k}$	$d_{c} + n_{h} d_{r}$
Quality	Baseline	Slight drop	Comparable
Attention score	$Q K^{T}$	$Q K^{T}$ (shared K)	Latent $Q K^{T}$
RoPE compatibility	Native	Native	Decoupled

GQA reduces cache by sharing KV heads across query groups. MLA reduces cache more aggressively by projecting to a shared latent. The quality difference is minimal because the up-projection matrices are learned and can reconstruct head-specific information.

4. Impact on Batched Serving

MLA dramatically changes the memory-vs-compute tradeoff in serving:

Memory-bound decoding phase: With MHA, long contexts exhaust GPU HBM due to KV cache. MLA's compression allows:

Longer context windows (~4× more tokens in same memory)
Larger batch sizes (more concurrent requests)
Better prefix caching hit rates (smaller cache entries)

Compute-bound prefill phase: MLA adds decompression overhead, but this is amortized:

Prefill is already compute-heavy (O(n²) attention)
The additional matmuls for up-projection are O(n × d_c × d_v) per layer
Net effect: minor prefill slowdown, massive decoding speedup

Implementation

import torch
import torch.nn as nn
import math

class MultiHeadLatentAttention(nn.Module):
    """
    MLA attention layer matching DeepSeek-V2/V3 and Kimi K2.x architecture.

    Key features:
    - Low-rank KV compression (cache only c_KV latent vector)
    - Decoupled RoPE for position-aware attention
    - Weight absorption for efficient score computation
    """

    def __init__(
        self,
        d_model: int = 4096,
        n_heads: int = 128,
        d_k: int = 128,
        d_v: int = 128,
        d_c: int = 512,          # KV latent dimension (compression target)
        d_c_prime: int = 1536,   # Query latent dimension
        d_r: int = 64,           # Decoupled RoPE key dimension per head
        max_seq_len: int = 8192,
        rope_base: float = 10000.0,
    ):
        super().__init__()
        self.d_model = d_model
        self.n_heads = n_heads
        self.d_k = d_k
        self.d_v = d_v
        self.d_c = d_c
        self.d_c_prime = d_c_prime
        self.d_r = d_r

        # === Down-projections (compression) ===
        self.w_dkv = nn.Linear(d_model, d_c, bias=False)        # KV latent
        self.w_dq = nn.Linear(d_model, d_c_prime, bias=False)   # Q latent

        # === Up-projections (decompression) ===
        # KV up-projections: latent -> per-head K and V
        self.w_uk = nn.Linear(d_c, n_heads * d_k, bias=False)
        self.w_uv = nn.Linear(d_c, n_heads * d_v, bias=False)

        # Q up-projection: latent -> per-head Q
        self.w_uq = nn.Linear(d_c_prime, n_heads * d_k, bias=False)

        # === Decoupled RoPE projections ===
        self.w_kr = nn.Linear(d_c, n_heads * d_r, bias=False)   # Rope key from latent
        self.w_qr = nn.Linear(d_c_prime, n_heads * d_r, bias=False)  # Rope query from latent

        # === Output projection ===
        self.w_o = nn.Linear(n_heads * d_v, d_model, bias=False)

        # RoPE frequencies
        inv_freq = 1.0 / (rope_base ** (torch.arange(0, d_r, 2).float() / d_r))
        self.register_buffer('inv_freq', inv_freq)

    def _apply_rope(self, x: torch.Tensor, seq_len: int) -> torch.Tensor:
        """Apply rotary position embedding to tensor of shape [batch, seq, n_heads, d_r]."""
        t = torch.arange(seq_len, device=x.device, dtype=self.inv_freq.dtype)
        freqs = torch.outer(t, self.inv_freq)              # [seq, d_r//2]
        cos = freqs.cos().unsqueeze(0).unsqueeze(2)        # [1, seq, 1, d_r//2]
        sin = freqs.sin().unsqueeze(0).unsqueeze(2)

        x1, x2 = x[..., ::2], x[..., 1::2]
        rotated = torch.stack([
            x1 * cos - x2 * sin,
            x1 * sin + x2 * cos,
        ], dim=-1).flatten(-2)
        return rotated

    def forward(
        self,
        x: torch.Tensor,
        kv_cache: torch.Tensor = None,
        start_pos: int = 0,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """
        Args:
            x: Input tensor [batch, seq_len, d_model]
            kv_cache: Cached c_KV from previous tokens [batch, cache_len, d_c]
            start_pos: Position offset for RoPE

        Returns:
            output: [batch, seq_len, d_model]
            new_kv_cache: Updated cache [batch, cache_len + seq_len, d_c]
        """
        B, S, _ = x.shape

        # === Step 1: Compress to latent space ===
        c_kv = self.w_dkv(x)     # [B, S, d_c] - THIS is what gets cached
        c_q = self.w_dq(x)       # [B, S, d_c']

        # === Step 2: Decompress for attention computation ===
        # K, V up-projection from latent
        k_content = self.w_uk(c_kv)    # [B, S, n_heads * d_k]
        v = self.w_uv(c_kv)            # [B, S, n_heads * d_v]
        q_content = self.w_uq(c_q)     # [B, S, n_heads * d_k]

        # Reshape to multi-head format
        q_content = q_content.view(B, S, self.n_heads, self.d_k)
        k_content = k_content.view(B, S, self.n_heads, self.d_k)
        v = v.view(B, S, self.n_heads, self.d_v)

        # === Step 3: Decoupled RoPE ===
        # Project to rope-specific dimensions and apply RoPE
        k_rope = self.w_kr(c_kv).view(B, S, self.n_heads, self.d_r)
        q_rope = self.w_qr(c_q).view(B, S, self.n_heads, self.d_r)

        k_rope = self._apply_rope(k_rope, start_pos + S)
        q_rope = self._apply_rope(q_rope, start_pos + S)

        # Concatenate content + rope for full key and query
        q = torch.cat([q_content, q_rope], dim=-1)   # [B, S, n_heads, d_k + d_r]
        k = torch.cat([k_content, k_rope], dim=-1)   # [B, S, n_heads, d_k + d_r]

        # === Step 4: KV cache management ===
        if kv_cache is not None:
            # Append new latent to cache
            new_kv_cache = torch.cat([kv_cache, c_kv], dim=1)
            # Decompress full cache for attention
            k_cache = self.w_uk(kv_cache).view(B, -1, self.n_heads, self.d_k)
            k_cache_rope = self._apply_rope(
                self.w_kr(kv_cache).view(B, -1, self.n_heads, self.d_r),
                start_pos  # cache already has positions 0..start_pos-1
            )
            k = torch.cat([
                torch.cat([k_cache, k_cache_rope], dim=-1),
                k
            ], dim=1)
            v_cache = self.w_uv(kv_cache).view(B, -1, self.n_heads, self.d_v)
            v = torch.cat([v_cache, v], dim=1)
        else:
            new_kv_cache = c_kv

        # === Step 5: Compute attention ===
        # Transpose for attention: [B, n_heads, seq, dim]
        q = q.transpose(1, 2)
        k = k.transpose(1, 2)
        v = v.transpose(1, 2)

        d_attn = self.d_k + self.d_r
        attn_weights = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(d_attn)
        attn_weights = torch.softmax(attn_weights, dim=-1)
        attn_output = torch.matmul(attn_weights, v)   # [B, n_heads, S, d_v]

        # === Step 6: Output projection ===
        attn_output = attn_output.transpose(1, 2).contiguous().view(B, S, -1)
        output = self.w_o(attn_output)

        return output, new_kv_cache


# === Example: Compare MLA vs MHA cache sizes ===
def compare_cache_sizes():
    """Demonstrate the KV cache savings of MLA over MHA."""
    n_heads = 128
    d_k = 128
    d_c = 512     # DeepSeek-V3 latent dim
    d_r = 64      # Decoupled rope dim per head
    seq_len = 65536   # 64K context
    bytes_per_element = 2   # FP16

    # MHA: cache K and V for all heads
    mha_cache_per_token = 2 * n_heads * d_k   # K + V
    mha_total = mha_cache_per_token * seq_len * bytes_per_element / (1024**3)

    # MLA: cache c_KV latent + decoupled rope keys for all heads
    mla_cache_per_token = d_c + n_heads * d_r
    mla_total = mla_cache_per_token * seq_len * bytes_per_element / (1024**3)

    print(f"MHA KV cache (64K ctx): {mha_total:.2f} GB per layer")
    print(f"MLA KV cache (64K ctx): {mla_total:.2f} GB per layer")
    print(f"Compression ratio: {mha_cache_per_token / mla_cache_per_token:.1f}x")
    print(f"\nFor 60 layers:")
    print(f"  MHA: {mha_total * 60:.1f} GB")
    print(f"  MLA: {mla_total * 60:.1f} GB")
    print(f"  Savings: {(mha_total - mla_total) * 60:.1f} GB")


if __name__ == "__main__":
    # Test MLA forward pass
    mla = MultiHeadLatentAttention(
        d_model=4096, n_heads=8, d_k=64, d_v=64,
        d_c=128, d_c_prime=256, d_r=32,
    )

    x = torch.randn(2, 10, 4096)   # batch=2, seq=10
    output, cache = mla(x)
    print(f"Output shape: {output.shape}")   # [2, 10, 4096]
    print(f"Cache shape: {cache.shape}")     # [2, 10, 128] - only d_c!

    # Autoregressive step
    x2 = torch.randn(2, 1, 4096)
    output2, cache2 = mla(x2, kv_cache=cache, start_pos=10)
    print(f"Output2 shape: {output2.shape}")   # [2, 1, 4096]
    print(f"Cache2 shape: {cache2.shape}")     # [2, 11, 128] - grew by 1

    print("\n--- Cache Comparison ---")
    compare_cache_sizes()

Prerequisites

kv-cache - You must understand standard KV caching before understanding why MLA compresses it
paged-attention - MLA changes what gets paged (latent vectors, not KV pairs)
flash-attention - MLA's absorbed attention can be fused into FlashAttention-style kernels
attention-mechanism - Foundation for understanding attention computation

References

DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model - Liu et al., 2024. arxiv:2405.04434 - Original MLA paper introducing the latent compression and decoupled RoPE strategy.
DeepSeek-V3 Technical Report - DeepSeek-AI, 2024. arxiv:2412.19437 - Scales MLA to 671B MoE with auxiliary-loss-free routing. Details the multi-token prediction (MTP) that inspired EAGLE-style draft heads.
vLLM MLA Implementation - github.com/vllm-project/vllm - Production MLA kernel with weight absorption and FlashAttention integration.
FlashInfer MLA Attention - github.com/flashinfer-ai/flashinfer - Custom CUDA kernels for MLA that support both prefill and decode phases with batched latent cache.

# Multi-Head Latent Attention (MLA)

Sirajuddin Shaik — Sat, 23 May 2026 14:33:01 +0000

Compressing KV cache via low-rank projections - the attention mechanism behind DeepSeek-V2/V3 and Kimi K2.x

Why This Matters

MLA changes how prefix caching, chunked prefill, and paged attention must be implemented

Formal Definition

Standard Multi-Head Attention (MHA)

For input $X \in R^{n \times d}$ , MHA computes per-head projections:

Q_{h} = X W_{Q (h)}, K_{h} = X W_{K (h)}, V_{h} = X W_{V (h)}

where $W_{Q (h)} \in R^{d \times d_{k}}$ , $W_{K (h)} \in R^{d \times d_{k}}$ , $W_{V (h)} \in R^{d \times d_{v}}$ .

KV cache size per token: $2 \times n_{h} \times d_{k}$ elements.

MLA: Low-Rank Latent Projection

MLA replaces the per-head KV projections with a shared low-rank latent compression:

Compression (KV → Latent):

c^{K V} = X W_{DK V} \in R^{n \times d_{c}}

where $W_{DK V} \in R^{d \times d_{c}}$ is the down-projection matrix and $d_{c} ≪ n_{h} \times d_{k}$ .

Decompression (Latent → KV):

K h = c^{K V} W U K^{(h)}, V h = c^{K V} W U V^{(h)}

where $W U K^{(h)} \in R^{d_{c} \times d_{k}}$ and $W U V^{(h)} \in R^{d_{c} \times d_{v}}$ are up-projection matrices.

KV cache per token: Only $c^{K V} \in R^{d_{c}}$ is stored - a single vector of dimension $d_{c}$ .

Query Compression (Optional)

MLA also compresses queries for training efficiency:

c^{Q} = XW D Q \in R^{n \times d_{c'}}

Q_{h} = c^{Q} W U Q^{(h)}

This doesn't affect the KV cache but reduces the activation memory during training.

Rotary Position Embedding (RoPE) Handling

RoPE is applied to the decompressed queries and keys. To keep the KV cache small, MLA applies RoPE to a separate "absorbed" key projection:

K^h = RoPE (c^{K V} W K R^{(h)})

Compression Ratio

For a model with $n_{h}$ heads and head dimension $d_{k}$ :

The total MLA KV cache per token stores:

The latent vector $c^{K V}$ of dimension $d_{c}$
The decoupled RoPE keys for all heads: $n_{h} \times d_{r}$

Compression Ratio = \frac{2 \cdot n _{h} \cdot d _{k}}{d _{c} + n _{h} \cdot d _{r}}

Approximating DeepSeek-V3 dimensions (see paper Table 1 for exact values):
$n_{h} = 128$ , $d_{k} = 128$ , $d_{c} = 512$ , $d_{r} = 64$ :

\frac{2 \times 128 \times 128}{512 + 128 \times 64} = \frac{32 , 768}{8 , 704} \approx 3.76 \times compression

Core Concepts

1. Weight Absorption (The Key Trick)

The critical insight in MLA is that the up-projection matrices $W_{U K (h)}$ can be absorbed into the query projection during attention computation:

Attention (Q, K, V) = softmax (\frac{Q _{h} K _{h T}}{d _{k}}) V_{h}

Substituting the decompressed forms:

Q h K_{h T} = (c^{Q} W U Q^{(h)}) (c^{K V} W U K^{(h)})^{T} = c^{Q} (W U Q^{(h)} W_{U K (h)}^{T}) c^{K V}^{T}

If we define $W ab sor b e d^{(h)} = W U Q^{(h)} W_{U K (h)}^{T} \in R^{d_{c'} \times d_{c}}$ , then:

Q h K_{h T} = c^{Q} W ab sor b e d^{(h)} c^{K V}^{T}

Practical implication: During decoding, we can compute attention scores without materializing the full K matrix. Only V needs decompression after softmax.

2. Decoupled RoPE Strategy

RoPE requires position-dependent keys, which conflicts with caching a position-agnostic latent. MLA solves this with a decoupled key:

Content key: $K h^{co n t e n t} = c^{K V} W U K^{(h)}$ - cached in latent form
Position key: $K h^{ro p e} = RoPE (c^{K V} W K R^{(h)})$ - small, position-aware, must be cached separately

The attention score becomes:

score (q, k) = \frac{Q _{h co n t e n t} \cdot K _{h co n t e n t} ^{T}}{d _{k}} + \frac{Q _{h ro p e} \cdot K _{h ro p e} ^{T}}{d _{r}}

Practical implication: The KV cache stores both $c^{K V}$ (latent) and $K_{h ro p e}$ (decoupled rope key). Total cache per token: $d_{c} + n_{h} \times d_{r}$ .

3. MLA vs GQA vs MHA

Property	MHA	GQA	MLA
KV groups	$n_{h}$	$n_{h} / g$	1 (latent)
Cache per token	$2 n_{h} d_{k}$	$2 (n_{h} / g) d_{k}$	$d_{c} + n_{h} d_{r}$
Quality	Baseline	Slight drop	Comparable
Attention score	$Q K^{T}$	$Q K^{T}$ (shared K)	Latent $Q K^{T}$
RoPE compatibility	Native	Native	Decoupled

4. Impact on Batched Serving

MLA dramatically changes the memory-vs-compute tradeoff in serving:

Memory-bound decoding phase: With MHA, long contexts exhaust GPU HBM due to KV cache. MLA's compression allows:

Longer context windows (~4× more tokens in same memory)
Larger batch sizes (more concurrent requests)
Better prefix caching hit rates (smaller cache entries)

Compute-bound prefill phase: MLA adds decompression overhead, but this is amortized:

Prefill is already compute-heavy (O(n²) attention)
The additional matmuls for up-projection are O(n × d_c × d_v) per layer
Net effect: minor prefill slowdown, massive decoding speedup

Implementation

import torch
import torch.nn as nn
import math

class MultiHeadLatentAttention(nn.Module):
    """
    MLA attention layer matching DeepSeek-V2/V3 and Kimi K2.x architecture.

    Key features:
    - Low-rank KV compression (cache only c_KV latent vector)
    - Decoupled RoPE for position-aware attention
    - Weight absorption for efficient score computation
    """

    def __init__(
        self,
        d_model: int = 4096,
        n_heads: int = 128,
        d_k: int = 128,
        d_v: int = 128,
        d_c: int = 512,          # KV latent dimension (compression target)
        d_c_prime: int = 1536,   # Query latent dimension
        d_r: int = 64,           # Decoupled RoPE key dimension per head
        max_seq_len: int = 8192,
        rope_base: float = 10000.0,
    ):
        super().__init__()
        self.d_model = d_model
        self.n_heads = n_heads
        self.d_k = d_k
        self.d_v = d_v
        self.d_c = d_c
        self.d_c_prime = d_c_prime
        self.d_r = d_r

        # === Down-projections (compression) ===
        self.w_dkv = nn.Linear(d_model, d_c, bias=False)        # KV latent
        self.w_dq = nn.Linear(d_model, d_c_prime, bias=False)   # Q latent

        # === Up-projections (decompression) ===
        # KV up-projections: latent -> per-head K and V
        self.w_uk = nn.Linear(d_c, n_heads * d_k, bias=False)
        self.w_uv = nn.Linear(d_c, n_heads * d_v, bias=False)

        # Q up-projection: latent -> per-head Q
        self.w_uq = nn.Linear(d_c_prime, n_heads * d_k, bias=False)

        # === Decoupled RoPE projections ===
        self.w_kr = nn.Linear(d_c, n_heads * d_r, bias=False)   # Rope key from latent
        self.w_qr = nn.Linear(d_c_prime, n_heads * d_r, bias=False)  # Rope query from latent

        # === Output projection ===
        self.w_o = nn.Linear(n_heads * d_v, d_model, bias=False)

        # RoPE frequencies
        inv_freq = 1.0 / (rope_base ** (torch.arange(0, d_r, 2).float() / d_r))
        self.register_buffer('inv_freq', inv_freq)

    def _apply_rope(self, x: torch.Tensor, seq_len: int) -> torch.Tensor:
        """Apply rotary position embedding to tensor of shape [batch, seq, n_heads, d_r]."""
        t = torch.arange(seq_len, device=x.device, dtype=self.inv_freq.dtype)
        freqs = torch.outer(t, self.inv_freq)              # [seq, d_r//2]
        cos = freqs.cos().unsqueeze(0).unsqueeze(2)        # [1, seq, 1, d_r//2]
        sin = freqs.sin().unsqueeze(0).unsqueeze(2)

        x1, x2 = x[..., ::2], x[..., 1::2]
        rotated = torch.stack([
            x1 * cos - x2 * sin,
            x1 * sin + x2 * cos,
        ], dim=-1).flatten(-2)
        return rotated

    def forward(
        self,
        x: torch.Tensor,
        kv_cache: torch.Tensor = None,
        start_pos: int = 0,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """
        Args:
            x: Input tensor [batch, seq_len, d_model]
            kv_cache: Cached c_KV from previous tokens [batch, cache_len, d_c]
            start_pos: Position offset for RoPE

        Returns:
            output: [batch, seq_len, d_model]
            new_kv_cache: Updated cache [batch, cache_len + seq_len, d_c]
        """
        B, S, _ = x.shape

        # === Step 1: Compress to latent space ===
        c_kv = self.w_dkv(x)     # [B, S, d_c] - THIS is what gets cached
        c_q = self.w_dq(x)       # [B, S, d_c']

        # === Step 2: Decompress for attention computation ===
        # K, V up-projection from latent
        k_content = self.w_uk(c_kv)    # [B, S, n_heads * d_k]
        v = self.w_uv(c_kv)            # [B, S, n_heads * d_v]
        q_content = self.w_uq(c_q)     # [B, S, n_heads * d_k]

        # Reshape to multi-head format
        q_content = q_content.view(B, S, self.n_heads, self.d_k)
        k_content = k_content.view(B, S, self.n_heads, self.d_k)
        v = v.view(B, S, self.n_heads, self.d_v)

        # === Step 3: Decoupled RoPE ===
        # Project to rope-specific dimensions and apply RoPE
        k_rope = self.w_kr(c_kv).view(B, S, self.n_heads, self.d_r)
        q_rope = self.w_qr(c_q).view(B, S, self.n_heads, self.d_r)

        k_rope = self._apply_rope(k_rope, start_pos + S)
        q_rope = self._apply_rope(q_rope, start_pos + S)

        # Concatenate content + rope for full key and query
        q = torch.cat([q_content, q_rope], dim=-1)   # [B, S, n_heads, d_k + d_r]
        k = torch.cat([k_content, k_rope], dim=-1)   # [B, S, n_heads, d_k + d_r]

        # === Step 4: KV cache management ===
        if kv_cache is not None:
            # Append new latent to cache
            new_kv_cache = torch.cat([kv_cache, c_kv], dim=1)
            # Decompress full cache for attention
            k_cache = self.w_uk(kv_cache).view(B, -1, self.n_heads, self.d_k)
            k_cache_rope = self._apply_rope(
                self.w_kr(kv_cache).view(B, -1, self.n_heads, self.d_r),
                start_pos  # cache already has positions 0..start_pos-1
            )
            k = torch.cat([
                torch.cat([k_cache, k_cache_rope], dim=-1),
                k
            ], dim=1)
            v_cache = self.w_uv(kv_cache).view(B, -1, self.n_heads, self.d_v)
            v = torch.cat([v_cache, v], dim=1)
        else:
            new_kv_cache = c_kv

        # === Step 5: Compute attention ===
        # Transpose for attention: [B, n_heads, seq, dim]
        q = q.transpose(1, 2)
        k = k.transpose(1, 2)
        v = v.transpose(1, 2)

        d_attn = self.d_k + self.d_r
        attn_weights = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(d_attn)
        attn_weights = torch.softmax(attn_weights, dim=-1)
        attn_output = torch.matmul(attn_weights, v)   # [B, n_heads, S, d_v]

        # === Step 6: Output projection ===
        attn_output = attn_output.transpose(1, 2).contiguous().view(B, S, -1)
        output = self.w_o(attn_output)

        return output, new_kv_cache


# === Example: Compare MLA vs MHA cache sizes ===
def compare_cache_sizes():
    """Demonstrate the KV cache savings of MLA over MHA."""
    n_heads = 128
    d_k = 128
    d_c = 512     # DeepSeek-V3 latent dim
    d_r = 64      # Decoupled rope dim per head
    seq_len = 65536   # 64K context
    bytes_per_element = 2   # FP16

    # MHA: cache K and V for all heads
    mha_cache_per_token = 2 * n_heads * d_k   # K + V
    mha_total = mha_cache_per_token * seq_len * bytes_per_element / (1024**3)

    # MLA: cache c_KV latent + decoupled rope keys for all heads
    mla_cache_per_token = d_c + n_heads * d_r
    mla_total = mla_cache_per_token * seq_len * bytes_per_element / (1024**3)

    print(f"MHA KV cache (64K ctx): {mha_total:.2f} GB per layer")
    print(f"MLA KV cache (64K ctx): {mla_total:.2f} GB per layer")
    print(f"Compression ratio: {mha_cache_per_token / mla_cache_per_token:.1f}x")
    print(f"\nFor 60 layers:")
    print(f"  MHA: {mha_total * 60:.1f} GB")
    print(f"  MLA: {mla_total * 60:.1f} GB")
    print(f"  Savings: {(mha_total - mla_total) * 60:.1f} GB")


if __name__ == "__main__":
    # Test MLA forward pass
    mla = MultiHeadLatentAttention(
        d_model=4096, n_heads=8, d_k=64, d_v=64,
        d_c=128, d_c_prime=256, d_r=32,
    )

    x = torch.randn(2, 10, 4096)   # batch=2, seq=10
    output, cache = mla(x)
    print(f"Output shape: {output.shape}")   # [2, 10, 4096]
    print(f"Cache shape: {cache.shape}")     # [2, 10, 128] - only d_c!

    # Autoregressive step
    x2 = torch.randn(2, 1, 4096)
    output2, cache2 = mla(x2, kv_cache=cache, start_pos=10)
    print(f"Output2 shape: {output2.shape}")   # [2, 1, 4096]
    print(f"Cache2 shape: {cache2.shape}")     # [2, 11, 128] - grew by 1

    print("\n--- Cache Comparison ---")
    compare_cache_sizes()

Prerequisites

kv-cache - You must understand standard KV caching before understanding why MLA compresses it
paged-attention - MLA changes what gets paged (latent vectors, not KV pairs)
flash-attention - MLA's absorbed attention can be fused into FlashAttention-style kernels
attention-mechanism - Foundation for understanding attention computation

References

DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model - Liu et al., 2024. arxiv:2405.04434 - Original MLA paper introducing the latent compression and decoupled RoPE strategy.
DeepSeek-V3 Technical Report - DeepSeek-AI, 2024. arxiv:2412.19437 - Scales MLA to 671B MoE with auxiliary-loss-free routing. Details the multi-token prediction (MTP) that inspired EAGLE-style draft heads.
vLLM MLA Implementation - github.com/vllm-project/vllm - Production MLA kernel with weight absorption and FlashAttention integration.
FlashInfer MLA Attention - github.com/flashinfer-ai/flashinfer - Custom CUDA kernels for MLA that support both prefill and decode phases with batched latent cache.