61 transformer blocks › 58 MoE blocks › Sparse latent attention
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
All 10 steps
- in
- [B batch, 128 heads, 1 positions, 576 latent + rotary] [B batch, 1001 cached positions, 576 latent + rotary] [B batch, 1 queries, 1001 keys]
- out
- [B batch, 1 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
| Component | Calculation | Parameters |
|---|---|---|
| Mix and project | 117,440,512 across its components | 117,440,512 |
| Per-instance total | 117,440,512 | |
Across 58 independent MoE blocks: 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 590–608
else: # MQA decode
if self.dequant_wkv_b is None and self.wkv_b.scale is not None:
self.dequant_wkv_b = weight_dequant(self.wkv_b.weight, self.wkv_b.scale)
wkv_b = self.wkv_b.weight if self.dequant_wkv_b is None else self.dequant_wkv_b
wkv_b = wkv_b.view(self.n_local_heads, -1, self.kv_lora_rank)
q_nope = torch.einsum("bshd,hdc->bshc", q_nope, wkv_b[:, :self.qk_nope_head_dim])
scores = (torch.einsum("bshc,btc->bsht", q_nope, self.kv_cache[:bsz, :end_pos]) +
torch.einsum("bshr,btr->bsht", q_pe, self.pe_cache[:bsz, :end_pos])) * self.softmax_scale
# indexer
topk_indices = self.indexer(x, qr, start_pos, freqs_cis, mask)
index_mask = torch.full((bsz, 1, end_pos), float("-inf"), device=x.device).scatter_(-1, topk_indices, 0)
scores += index_mask.unsqueeze(2)
scores = scores.softmax(dim=-1)
x = torch.einsum("bsht,btc->bshc", scores, self.kv_cache[:bsz, :end_pos])
x = torch.einsum("bshc,hdc->bshd", x, wkv_b[:, -self.v_head_dim:]) x = self.wo(x.flatten(2))
return xSources2
Direct evidence
When decoding one token, the implementation absorbs the key–value up-projection into queries and outputs and attends over the cached latent directly.
mla-decode · deepseek-ai/DeepSeek-V3.2-Exp / inference/model.py
Pinned revision 87e509a2e5a1 ↗ · lines 590–606
Read code excerpt
else: # MQA decode
if self.dequant_wkv_b is None and self.wkv_b.scale is not None:
self.dequant_wkv_b = weight_dequant(self.wkv_b.weight, self.wkv_b.scale)
wkv_b = self.wkv_b.weight if self.dequant_wkv_b is None else self.dequant_wkv_b
wkv_b = wkv_b.view(self.n_local_heads, -1, self.kv_lora_rank)
q_nope = torch.einsum("bshd,hdc->bshc", q_nope, wkv_b[:, :self.qk_nope_head_dim])
scores = (torch.einsum("bshc,btc->bsht", q_nope, self.kv_cache[:bsz, :end_pos]) +
torch.einsum("bshr,btr->bsht", q_pe, self.pe_cache[:bsz, :end_pos])) * self.softmax_scale
# indexer
topk_indices = self.indexer(x, qr, start_pos, freqs_cis, mask)
index_mask = torch.full((bsz, 1, end_pos), float("-inf"), device=x.device).scatter_(-1, topk_indices, 0)
scores += index_mask.unsqueeze(2)
scores = scores.softmax(dim=-1)
x = torch.einsum("bsht,btc->bshc", scores, self.kv_cache[:bsz, :end_pos])
x = torch.einsum("bshc,hdc->bshd", x, wkv_b[:, -self.v_head_dim:])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