Dense transformer block

61 transformer blocks › 3 dense blocks

Sparse latent attention

Open inside →×

learnedself.attn

Two ideas in one attention. It is latent: each position is compressed into one small vector, cached with a shared rotary key part, and every head’s key content and value are rebuilt from it. It is sparse: a lightning indexer chooses up to 2,048 positions for each query, and attention runs over those alone.

Why

Latent. Standard attention caches a full key and value for every head at every position, and the authors call that cache a large bottleneck that limits batch size and sequence length when the model is served. Here one 512-feature latent per position is cached and keys and values are rebuilt from it. DeepSeek-V2 §2.1.1–2.1.2

The authors report a much smaller cache with quality comparable to standard multi-head attention. DeepSeek-V3 §2.1.1

Sparse. Attending to every earlier position costs O(L²) for a context of L tokens. Choosing k positions for each query brings the main attention to O(Lk) (the indexer that chooses them still costs O(L²)), and the authors report no substantial loss of quality on short or long contexts. DeepSeek-V3.2-Exp report §3

Both at once. Sparse attention was added by continuing the training of DeepSeek-V3.1-Terminus, which already used latent attention, so it is built on it: what the indexer selects for a query are the cached latents of the chosen positions. DeepSeek-V3.2-Exp report §1

O=AVO=AV
Δattn=oWO\Delta_{\mathrm{attn}}=oW_{O}
All 43 steps
Query compression:\text{Query compression:}
cQ=xWDQc^{{\htmlClass{vl-stream-q}{Q}}}=xW_{DQ}
c^Q=RMSNorm⁡(cQ)\hat c^{{\htmlClass{vl-stream-q}{Q}}}=\operatorname{RMSNorm}(c^{{\htmlClass{vl-stream-q}{Q}}})
Key–value compression:\text{Key–value compression:}
cKV=xWDKVc^{KV}=xW_{DKV}
c^KV=RMSNorm⁡(cKV)\hat c^{KV}=\operatorname{RMSNorm}(c^{KV})
Query heads:\text{Query heads:}
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}}}]
Key heads:\text{Key heads:}
kC=c^KVWUK{\htmlClass{vl-stream-k}{k^{C}}}=\hat c^{KV}W_{UK}
KC=reshape⁡(kC){\htmlClass{vl-stream-k}{K^{C}}}=\operatorname{reshape}({\htmlClass{vl-stream-k}{k^{C}}})
kpe=xWKR{\htmlClass{vl-stream-k}{k^{\mathrm{pe}}}}=xW_{KR}
ktR=RoPE⁡(ktpe,Rt){\htmlClass{vl-stream-k}{k^{R}_{t}}}=\operatorname{RoPE}({\htmlClass{vl-stream-k}{k^{\mathrm{pe}}_{t}}},{\htmlClass{vl-stream-context}{R_t}})
Ki=[KiC; kR]{\htmlClass{vl-stream-k}{K_i}}=[{\htmlClass{vl-stream-k}{K^{C}_i}};\,{\htmlClass{vl-stream-k}{k^{R}}}]
Value heads:\text{Value heads:}
v=c^KVWUV{\htmlClass{vl-stream-v}{v}}=\hat c^{KV}W_{UV}
V=reshape⁡(v){\htmlClass{vl-stream-v}{V}}=\operatorname{reshape}({\htmlClass{vl-stream-v}{v}})
Lightning indexer:\text{Lightning indexer:}
qI=c^QWIQ{\htmlClass{vl-stream-q}{q^{I}}}=\hat c^{{\htmlClass{vl-stream-q}{Q}}}W_{IQ}
QI=reshape⁡(qI){\htmlClass{vl-stream-q}{Q^{I}}}=\operatorname{reshape}({\htmlClass{vl-stream-q}{q^{I}}})
QI~=RoPE⁡64(QI,R)\tilde {\htmlClass{vl-stream-q}{Q^{I}}}=\operatorname{RoPE}_{64}({\htmlClass{vl-stream-q}{Q^{I}}},{\htmlClass{vl-stream-context}{R}})
QIˉ=QI~ 1128H128\bar {\htmlClass{vl-stream-q}{Q^{I}}}=\tilde {\htmlClass{vl-stream-q}{Q^{I}}}\,\tfrac{1}{\sqrt{128}}\mathbf{H}_{128}
kI=xWIK{\htmlClass{vl-stream-k}{k^{I}}}=xW_{IK}
kI^=LayerNorm⁡(kI)\hat {\htmlClass{vl-stream-k}{k^{I}}}=\operatorname{LayerNorm}({\htmlClass{vl-stream-k}{k^{I}}})
kI~=RoPE⁡64(kI^,R)\tilde {\htmlClass{vl-stream-k}{k^{I}}}=\operatorname{RoPE}_{64}(\hat {\htmlClass{vl-stream-k}{k^{I}}},{\htmlClass{vl-stream-context}{R}})
kIˉ=kI~ 1128H128\bar {\htmlClass{vl-stream-k}{k^{I}}}=\tilde {\htmlClass{vl-stream-k}{k^{I}}}\,\tfrac{1}{\sqrt{128}}\mathbf{H}_{128}
w=xWww=xW_{w}
St,sI=164128∑h=164wt,h ReLU⁡ ⁣(Qh,tIˉ⋅ksIˉ)S^{I}_{t,s}=\frac{1}{\sqrt{64}\sqrt{128}}\sum_{h=1}^{64}w_{t,h}\,\operatorname{ReLU}\!\big(\bar {\htmlClass{vl-stream-q}{Q^{I}_{h,t}}}\cdot\bar {\htmlClass{vl-stream-k}{k^{I}_{s}}}\big)
S~I=SI+M\tilde S^{I}=S^{I}+{\htmlClass{vl-stream-context}{M}}
Mt,ssel={Mt,ss∈top⁡2,048(S~tI)−∞otherwise{\htmlClass{vl-stream-context}{M^{\mathrm{sel}}_{t,s}}}=\begin{cases}{\htmlClass{vl-stream-context}{M_{t,s}}}&s\in\operatorname{top}_{2,048}(\tilde S^{I}_{t})\\-\infty&\text{otherwise}\end{cases}
Sparse attention:\text{Sparse attention:}
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])}
O=AVO=AV
o=reshape⁡(O)o=\operatorname{reshape}(O)
Δattn=oWO\Delta_{\mathrm{attn}}=oW_{O}
in
xattnx_{\mathrm{attn}} [B batch, T positions, 7168 d]RR [T positions, 32 pairs]MM [T queries, T keys]
out
Δattn\Delta_{\mathrm{attn}} [B batch, T positions, 7168 d]
parameters
201,066,752 · × 3 dense blocks = 603,200,256
Parameters

201,066,752 parameters per instance, across its learned components

Exact per-instance accounting
ComponentCalculationParameters
Query compression11,011,584 across its components11,011,584
Key–value compression3,670,528 across its components3,670,528
Query heads37,748,736 across its components37,748,736
Key heads8,847,360 across its components8,847,360
Value heads8,388,608 across its components8,388,608
Lightning indexer13,959,424 across its components13,959,424
Sparse attention117,440,512 across its components117,440,512
Per-instance total201,066,752

Across 3 independent dense blocks: 201,066,752×3=603,200,256201{,}066{,}752 \times 3 = 603{,}200{,}256 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.

Codemodel.py 523–608
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 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
Sources3

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

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