DeepSeek-V3.2-Exp · Key heads

61 transformer blocks › 58 MoE blocks › Sparse latent attention

Key heads

learnedk

Reconstruct 128 key heads of 192 features: per-head content from the latent and one rotated key shared by all heads.

Why

A key has two parts because rotation by position does not go together with the compression. If the rebuilt keys were rotated, the rotation would stand between the query matrix and the key’s up-projection, and the keys of all earlier tokens would have to be recomputed for every new token. So position is carried by a separate rotary part, one per position and shared by all heads, joined to each head’s content part. DeepSeek-V2 §2.1.3

{\htmlClass{vl-stream-k}{k^{R}_{t}}}=\operatorname{RoPE}({\htmlClass{vl-stream-k}{k^{\mathrm{pe}}_{t}}},{\htmlClass{vl-stream-context}{R_t}})

{\htmlClass{vl-stream-k}{K_i}}=[{\htmlClass{vl-stream-k}{K^{C}_i}};\,{\htmlClass{vl-stream-k}{k^{R}}}]

All 5 steps

{\htmlClass{vl-stream-k}{k^{C}}}=\hat c^{KV}W_{UK}

{\htmlClass{vl-stream-k}{K^{C}}}=\operatorname{reshape}({\htmlClass{vl-stream-k}{k^{C}}})

{\htmlClass{vl-stream-k}{k^{\mathrm{pe}}}}=xW_{KR}

{\htmlClass{vl-stream-k}{k^{R}_{t}}}=\operatorname{RoPE}({\htmlClass{vl-stream-k}{k^{\mathrm{pe}}_{t}}},{\htmlClass{vl-stream-context}{R_t}})

{\htmlClass{vl-stream-k}{K_i}}=[{\htmlClass{vl-stream-k}{K^{C}_i}};\,{\htmlClass{vl-stream-k}{k^{R}}}]

in: $\hat c^{KV}$ [B batch, T positions, 512 kv latent] $x$ [B batch, T positions, 7168 d] $R$ [T positions, 32 pairs]
out: $K$ [B batch, 128 heads, T positions, 192 content + rotary]
parameters: 8,847,360 · × 58 MoE blocks = 513,146,880

Parameters

8,847,360 parameters per instance, across its learned components

Exact per-instance accounting
Component	Calculation	Parameters
Key content up-projection	$512 \times 16{,}384$	8,388,608
Rotary key projection	$7{,}168 \times 64$	458,752
Per-instance total		8,847,360

Across 58 independent MoE blocks: $8{,}847{,}360 \times 58 = 513{,}146{,}880$ 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

One matrix, wkv_b, produces the content keys and the values together; the reference splits them after the product (model.py 576–578).

Content keys carry no rotation, so when decoding the reference folds $W_{UK}$ into the queries and scores them against the cached latent directly, which gives the same scores (model.py 591–597).

Codemodel.py 523–589

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 565–568

        kv = self.wkv_a(x)
        kv, k_pe = torch.split(kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)
        kv = self.kv_norm(kv)
        k_pe = apply_rotary_emb(k_pe.unsqueeze(2), freqs_cis)

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)

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

The joint down-projection yields the key–value latent and one rotary key shared by all heads; the latent is normalized and the key rotated.

mla-kv · deepseek-ai/DeepSeek-V3.2-Exp / inference/model.py
Pinned revision 87e509a2e5a1 ↗ · lines 565–568

Read code excerpt

        kv = self.wkv_a(x)
        kv, k_pe = torch.split(kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)
        kv = self.kv_norm(kv)
        k_pe = apply_rotary_emb(k_pe.unsqueeze(2), freqs_cis)

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)