61 transformer blocks › 58 MoE blocks › Sparse latent attention › Sparse attention › Mix and project
Weighted sum of values
×no parameterstorch.einsum("bsht,bthd->bshd", scores, v)
Mix value vectors using each query’s attention weights.
- in
- [B batch, 128 heads, T queries, T keys] [B batch, 128 heads, T positions, 128 value]
- out
- [B batch, 128 heads, T positions, 128 value]
Codemodel.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)Sources1
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)