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

Attention update

×

tensorx · 607

The attention update added to the residual stream.

shape
Δattn\Delta_{\mathrm{attn}} [B batch, 1 positions, 7168 d]
Sources

Contextual evidence from Attention residual addition

A block applies attention-norm, MLA, ffn-norm and the FFN; the first n_dense_layers blocks use the dense MLP, the rest MoE. Residual additions happen inside the next normalization.

block · deepseek-ai/DeepSeek-V3.2-Exp / inference/model.py
Pinned revision 87e509a2e5a1 ↗ · lines 826–851

Read code excerpt
        self.attn = MLA(args)
        self.ffn = MLP(args.dim, args.inter_dim) if layer_id < args.n_dense_layers else MoE(args)
        self.attn_norm = RMSNorm(args.dim)
        self.ffn_norm = RMSNorm(args.dim)

… lines 830–843 omitted …

        if residual is None:
            x, residual = self.attn_norm(x), x
        else:
            x, residual = self.attn_norm(x, residual)
        x = self.attn(x, start_pos, freqs_cis, mask)
        x, residual = self.ffn_norm(x, residual)
        x = self.ffn(x)
        return x, residual

Contextual evidence from Attention residual addition

RMSNorm owns one gain per feature and computes in float32; with a residual argument it first adds the residual, fusing the residual addition into the next normalization.

rmsnorm · deepseek-ai/DeepSeek-V3.2-Exp / inference/model.py
Pinned revision 87e509a2e5a1 ↗ · lines 280–306

Read code excerpt
    def __init__(self, dim: int, eps: float = 1e-6):
        super().__init__()
        self.dim = dim
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim, dtype=torch.float32))

    def forward(self, x: torch.Tensor, residual: Optional[torch.Tensor] = None):
        """
        Forward pass for RMSNorm.

        Args:
            x (torch.Tensor): Input tensor.

        Returns:
            torch.Tensor: Normalized tensor with the same shape as input.
        """
        dtype = x.dtype
        if residual is None:
            x = x.float()
            var = x.pow(2).mean(-1, keepdim=True)
            x = x * torch.rsqrt(var + self.eps)
            return (self.weight * x).to(dtype)
        else:
            x = residual = x.float() + residual.float()
            var = x.pow(2).mean(-1, keepdim=True)
            x = x * torch.rsqrt(var + self.eps)
            return (self.weight * x).to(dtype), residual.to(dtype)

Contextual evidence from Sparse attention

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:])

Contextual evidence from Sparse attention

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