Sparse 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

Attention weights

Open inside →×

no parametersscores · 580–588

Compare queries with keys, scale, restrict each query to its selected positions and normalize.

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=01000exp⁡(S~[b,i,t,u])A[b,i,t,s]=\frac{\exp(\tilde S[b,i,t,s])}{\sum_{u=0}^{1000}\exp(\tilde S[b,i,t,u])}
All 4 steps
Si=QiC~ c^KVT+QiR kRTS_i=\tilde {\htmlClass{vl-stream-q}{Q^{C}_i}}\,\hat c^{KV\mathsf T}+{\htmlClass{vl-stream-q}{Q^{R}_i}}\,{\htmlClass{vl-stream-k}{k^{R\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=01000exp⁡(S~[b,i,t,u])A[b,i,t,s]=\frac{\exp(\tilde S[b,i,t,s])}{\sum_{u=0}^{1000}\exp(\tilde S[b,i,t,u])}
in
Q~\tilde Q [B batch, 128 heads, 1 positions, 576 latent + rotary]K^\hat K [B batch, 1001 cached positions, 576 latent + rotary]MselM^{\mathrm{sel}} [B batch, 1 queries, 1001 keys]
out
AA [B batch, 128 heads, 1 queries, 1001 keys]
Explanation

The reference adds the causal mask to the selection before this step (model.py 584–585).

Codemodel.py 590–608
model.py 590–606
        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:])
model.py 607–608
        x = self.wo(x.flatten(2))
        return x
Sources2

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