Multi-head latent attention

Reference code deepseek-ai/DeepSeek-V3.2-Exp @ 87e509a · checkpoint, tokenizer and tensors not verified · How to read

61 transformer blocks › 58 MoE blocks › Sparse latent attention

learnedq · 575

Expand the query latent into 128 query heads of 192 features: 128 content features and 64 rotated features.

QtR=RoPE⁡(Qtpe,Rt){\htmlClass{vl-stream-q}{Q^{R}_{t}}}=\operatorname{RoPE}({\htmlClass{vl-stream-q}{Q^{\mathrm{pe}}_{t}}},{\htmlClass{vl-stream-context}{R_t}})
Q=[QC; QR]{\htmlClass{vl-stream-q}{Q}}=[{\htmlClass{vl-stream-q}{Q^{C}}};\,{\htmlClass{vl-stream-q}{Q^{R}}}]
All 6 steps
qC=c^QWUQ{\htmlClass{vl-stream-q}{q^{C}}}=\hat c^{{\htmlClass{vl-stream-q}{Q}}}W_{UQ}
QC=reshape⁡(qC){\htmlClass{vl-stream-q}{Q^{C}}}=\operatorname{reshape}({\htmlClass{vl-stream-q}{q^{C}}})
qpe=c^QWQR{\htmlClass{vl-stream-q}{q^{\mathrm{pe}}}}=\hat c^{{\htmlClass{vl-stream-q}{Q}}}W_{QR}
Qpe=reshape⁡(qpe){\htmlClass{vl-stream-q}{Q^{\mathrm{pe}}}}=\operatorname{reshape}({\htmlClass{vl-stream-q}{q^{\mathrm{pe}}}})
QtR=RoPE⁡(Qtpe,Rt){\htmlClass{vl-stream-q}{Q^{R}_{t}}}=\operatorname{RoPE}({\htmlClass{vl-stream-q}{Q^{\mathrm{pe}}_{t}}},{\htmlClass{vl-stream-context}{R_t}})
Q=[QC; QR]{\htmlClass{vl-stream-q}{Q}}=[{\htmlClass{vl-stream-q}{Q^{C}}};\,{\htmlClass{vl-stream-q}{Q^{R}}}]
in
c^Q\hat c^{Q} [B batch, T positions, 1536 q latent]RR [T positions, 32 pairs]
out
QQ [B batch, 128 heads, T positions, 192 content + rotary]
parameters
37,748,736 · × 58 MoE blocks = 2,189,426,688
Parameters

37,748,736 parameters per instance, across its learned components

Exact per-instance accounting
ComponentCalculationParameters
Query content up-projection1,536×16,3841{,}536 \times 16{,}38425,165,824
Query rotary up-projection1,536×8,1921{,}536 \times 8{,}19212,582,912
Per-instance total37,748,736

Across 58 independent MoE blocks: 37,748,736×58=2,189,426,68837{,}748{,}736 \times 58 = 2{,}189{,}426{,}688 parameters.

Derived from the pinned reference configuration and inference code. Checkpoint tensors have not been independently counted; FP8 block scales and the multi-token-prediction module are not counted.

Explanation

One matrix, wq_b, produces both parts; the reference splits each head’s 192 features after the product (model.py 561–563).

Codemodel.py 523–539, 560–564
model.py 523–539
        self.qk_head_dim = args.qk_nope_head_dim + args.qk_rope_head_dim
        self.v_head_dim = args.v_head_dim

        self.wq_a = Linear(self.dim, self.q_lora_rank)
        self.q_norm = RMSNorm(self.q_lora_rank)
        self.wq_b = ColumnParallelLinear(self.q_lora_rank, self.n_heads * self.qk_head_dim)
        self.wkv_a = Linear(self.dim, self.kv_lora_rank + self.qk_rope_head_dim)
        self.kv_norm = RMSNorm(self.kv_lora_rank)
        self.wkv_b = ColumnParallelLinear(self.kv_lora_rank, self.n_heads * (self.qk_nope_head_dim + self.v_head_dim))
        self.wo = RowParallelLinear(self.n_heads * self.v_head_dim, self.dim)
        self.softmax_scale = self.qk_head_dim ** -0.5
        self.scale_fmt = args.scale_fmt
        if args.max_seq_len > args.original_seq_len:
            mscale = 0.1 * args.mscale * math.log(args.rope_factor) + 1.0
            self.softmax_scale = self.softmax_scale * mscale * mscale

        self.indexer = Indexer(args)
model.py 560–564
        qr = self.q_norm(self.wq_a(x))
        q = self.wq_b(qr)
        q = q.view(bsz, seqlen, self.n_local_heads, self.qk_head_dim)
        q_nope, q_pe = torch.split(q, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)
        q_pe = apply_rotary_emb(q_pe, freqs_cis)
Sources2

Direct evidence

MLA owns a query down-projection, query RMSNorm and up-projection, a joint key–value down-projection that also emits the shared rotary key, a latent RMSNorm, a key–value up-projection and an output projection.

mla-init · deepseek-ai/DeepSeek-V3.2-Exp / inference/model.py
Pinned revision 87e509a2e5a1 ↗ · lines 523–539

Read code excerpt
        self.qk_head_dim = args.qk_nope_head_dim + args.qk_rope_head_dim
        self.v_head_dim = args.v_head_dim

        self.wq_a = Linear(self.dim, self.q_lora_rank)
        self.q_norm = RMSNorm(self.q_lora_rank)
        self.wq_b = ColumnParallelLinear(self.q_lora_rank, self.n_heads * self.qk_head_dim)
        self.wkv_a = Linear(self.dim, self.kv_lora_rank + self.qk_rope_head_dim)
        self.kv_norm = RMSNorm(self.kv_lora_rank)
        self.wkv_b = ColumnParallelLinear(self.kv_lora_rank, self.n_heads * (self.qk_nope_head_dim + self.v_head_dim))
        self.wo = RowParallelLinear(self.n_heads * self.v_head_dim, self.dim)
        self.softmax_scale = self.qk_head_dim ** -0.5
        self.scale_fmt = args.scale_fmt
        if args.max_seq_len > args.original_seq_len:
            mscale = 0.1 * args.mscale * math.log(args.rope_factor) + 1.0
            self.softmax_scale = self.softmax_scale * mscale * mscale

        self.indexer = Indexer(args)

Direct evidence

Queries are produced through a low-rank latent, split into a non-rotated and a rotated part, and rotated.

mla-queries · deepseek-ai/DeepSeek-V3.2-Exp / inference/model.py
Pinned revision 87e509a2e5a1 ↗ · lines 560–564

Read code excerpt
        qr = self.q_norm(self.wq_a(x))
        q = self.wq_b(qr)
        q = q.view(bsz, seqlen, self.n_local_heads, self.qk_head_dim)
        q_nope, q_pe = torch.split(q, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)
        q_pe = apply_rotary_emb(q_pe, freqs_cis)