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

Sparse attention

Open inside →×

learnedMLA.forward · 580–607

Attend over each query’s selected positions with 128 heads and project the result to the model width.

Why

This plate draws how the reference runs a whole sequence: ordinary attention with per-head keys and values, where the indexer’s choice arrives as a mask. The authors implement this masked mode to simulate sparse attention and say it can be more efficient on short contexts. DeepSeek-V3.2-Exp report §3

In the other mode one key–value latent per position is shared by all query heads; the sparse kernel needs each entry shared across queries to be efficient. DeepSeek-V3.2-Exp report §1

O=AVO=AV
Δattn=oWO\Delta_{\mathrm{attn}}=oW_{O}
All 9 steps
Attention weights:\text{Attention weights:}
S=QKTS=QK^{\mathsf T}
S′=S⋅m2192S'=S\cdot\frac{m^2}{\sqrt{192}}
m=1+0.1ln⁡40m=1+0.1\ln 40
S~=S′+Msel\tilde S=S'+{\htmlClass{vl-stream-context}{M^{\mathrm{sel}}}}
A[b,i,t,s]=exp⁡(S~[b,i,t,s])∑u=0T−1exp⁡(S~[b,i,t,u])A[b,i,t,s]=\frac{\exp(\tilde S[b,i,t,s])}{\sum_{u=0}^{T-1}\exp(\tilde S[b,i,t,u])}
Mix and project:\text{Mix and project:}
O=AVO=AV
o=reshape⁡(O)o=\operatorname{reshape}(O)
Δattn=oWO\Delta_{\mathrm{attn}}=oW_{O}
in
QQ [B batch, 128 heads, T positions, 192 content + rotary]KK [B batch, 128 heads, T positions, 192 content + rotary]VV [B batch, 128 heads, T positions, 128 value]MselM^{\mathrm{sel}} [B batch, T queries, T keys]
out
Δattn\Delta_{\mathrm{attn}} [B batch, T positions, 7168 d]
parameters
117,440,512 · × 58 MoE blocks = 6,811,549,696
Parameters

117,440,512 parameters per instance, across its learned components

Exact per-instance accounting
ComponentCalculationParameters
Mix and project117,440,512 across its components117,440,512
Per-instance total117,440,512

Across 58 independent MoE blocks: 117,440,512×58=6,811,549,696117{,}440{,}512 \times 58 = 6{,}811{,}549{,}696 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

The reference takes this path when a mask is passed, for more than one position (model.py 574–589); a single decoded token scores the cached latent directly.

Codemodel.py 574–589, 607–608
model.py 574–589
        if mask is not None:    # MHA prefill
            q = torch.cat([q_nope, q_pe], dim=-1)
            kv = self.wkv_b(kv)
            kv = kv.view(bsz, seqlen, self.n_local_heads, self.qk_nope_head_dim + self.v_head_dim)
            k_nope, v = torch.split(kv, [self.qk_nope_head_dim, self.v_head_dim], dim=-1)
            k = torch.cat([k_nope, k_pe.expand(-1, -1, self.n_local_heads, -1)], dim=-1)
            scores = torch.einsum("bshd,bthd->bsht", q, k).mul_(self.softmax_scale)

            # indexer
            topk_indices = self.indexer(x, qr, start_pos, freqs_cis, mask)
            index_mask = torch.full((bsz, seqlen, seqlen), float("-inf"), device=x.device).scatter_(-1, topk_indices, 0)
            index_mask += mask
            scores += index_mask.unsqueeze(2)

            scores = scores.softmax(dim=-1)
            x = torch.einsum("bsht,bthd->bshd", scores, v)
model.py 607–608
        x = self.wo(x.flatten(2))
        return x
Sources2

Direct evidence

In the full-sequence (prefill) path, the latent is up-projected into per-head keys and values; each head’s key joins its non-rotated part with the shared rotary key; scores are scaled, masked by the indexer’s selection and the causal mask, normalized and used to mix values.

mla-prefill · deepseek-ai/DeepSeek-V3.2-Exp / inference/model.py
Pinned revision 87e509a2e5a1 ↗ · lines 574–589

Read code excerpt
        if mask is not None:    # MHA prefill
            q = torch.cat([q_nope, q_pe], dim=-1)
            kv = self.wkv_b(kv)
            kv = kv.view(bsz, seqlen, self.n_local_heads, self.qk_nope_head_dim + self.v_head_dim)
            k_nope, v = torch.split(kv, [self.qk_nope_head_dim, self.v_head_dim], dim=-1)
            k = torch.cat([k_nope, k_pe.expand(-1, -1, self.n_local_heads, -1)], dim=-1)
            scores = torch.einsum("bshd,bthd->bsht", q, k).mul_(self.softmax_scale)

            # indexer
            topk_indices = self.indexer(x, qr, start_pos, freqs_cis, mask)
            index_mask = torch.full((bsz, seqlen, seqlen), float("-inf"), device=x.device).scatter_(-1, topk_indices, 0)
            index_mask += mask
            scores += index_mask.unsqueeze(2)

            scores = scores.softmax(dim=-1)
            x = torch.einsum("bsht,bthd->bshd", scores, v)

Direct evidence

Head outputs are flattened and projected back to the model width.

mla-output · deepseek-ai/DeepSeek-V3.2-Exp / inference/model.py
Pinned revision 87e509a2e5a1 ↗ · lines 607–608

Read code excerpt
        x = self.wo(x.flatten(2))
        return x