DeepSeek V4.1 · Compressor, pairs pooled

40 layers › 2 groups of 6 layers › Source layer, pairs pooled › Compressed sparse attention

Compressor, pairs pooled

learnedself.compressor(x, start_pos) · 747

One 512-feature vector for every two tokens: a projection, a learned pooling of each pair and a norm.

Why

Simpler. The earlier compressor built each entry from overlapping groups of tokens and gave them an absolute position embedding. This one drops both; the report’s reason is a simpler implementation and more efficient training. DeepSeek V4.1 report §2.3

{\htmlClass{vl-stream-k}{k^{p}_{j}}}=\sum_{i\in\{2j,\,2j+1\}}\operatorname{softmax}_{i}\big(g^{c}_{i}\big)\odot {\htmlClass{vl-stream-k}{k^{c}_{i}}}

c=\gamma_{C}\odot\frac{{\htmlClass{vl-stream-k}{k^{p}}}}{\sqrt{\operatorname{mean}\big(({\htmlClass{vl-stream-k}{k^{p}}})^{2}\big)+\varepsilon}}

All 4 steps

{\htmlClass{vl-stream-k}{k^{c}}}=x_{\mathrm{attn}}W_{C}

g^{c}=x_{\mathrm{attn}}W_{G}

{\htmlClass{vl-stream-k}{k^{p}_{j}}}=\sum_{i\in\{2j,\,2j+1\}}\operatorname{softmax}_{i}\big(g^{c}_{i}\big)\odot {\htmlClass{vl-stream-k}{k^{c}_{i}}}

c=\gamma_{C}\odot\frac{{\htmlClass{vl-stream-k}{k^{p}}}}{\sqrt{\operatorname{mean}\big(({\htmlClass{vl-stream-k}{k^{p}}})^{2}\big)+\varepsilon}}

in: $x_{\mathrm{attn}}$ [B batch, T positions, 5120 d]
out: $c$ [B batch, T/2 compressed positions, 512 d_h]
parameters: 5,243,392 · × 2 groups = 10,486,784

Parameters

5,243,392 parameters per instance, across its learned components

Exact per-instance accounting
Component	Calculation	Parameters
Compressed key–value projection	$5{,}120 \times 512$	2,621,440
Pooling scores	$5{,}120 \times 512$	2,621,440
Compressed RMSNorm	$512$	512
Per-instance total		5,243,392

Across 2 independent groups: $5{,}243{,}392 \times 2 = 10{,}486{,}784$ parameters.

Derived from the pinned reference implementation and configuration; equal to the element counts of the checkpoint’s own tensors on the text path. The quantisation scales, the vision encoder and the DSpark blocks are not counted.

Explanation

Only layers 2, 8, 14, 20 build compressed keys and values; every other layer with a compression ratio reads what the nearest earlier of them wrote (model.py 618–619, 746–748).

Codemodel.py 281–293, 429–485

model.py 429–485

class Compressor(nn.Module):
    """Pools `compress_ratio` consecutive tokens into one KV latent with a learned softmax gate.

    Returns the latent before RoPE, or None while a group is still filling up -- so during decode it
    only yields every `compress_ratio` steps, holding the partial group in `kv_state`/`score_state`.
    Pre-RoPE is deliberate: the indexer needs the unrotated form, so Attention rotates afterwards.
    """

    def __init__(self, args: ModelArgs, layer_id: int):
        super().__init__()
        compress_ratio = args.compress_ratios[layer_id]
        head_dim = args.head_dim
        self.compress_ratio = compress_ratio
        self.head_dim = head_dim
        self.norm = RMSNorm(head_dim, args.norm_eps)
        # ratio 1 is a plain projection, so it stays in the checkpoint's bf16; the softmax pooling
        # above ratio 1 runs in fp32, so those weights are promoted to fp32 to match
        self.wkv = Linear(args.dim, head_dim, dtype=torch.float32 if compress_ratio > 1 else torch.bfloat16)
        if compress_ratio > 1:
            self.wgate = Linear(args.dim, head_dim, dtype=torch.float32)
            # tail of an incomplete group, carried across decode steps
            self.kv_state: torch.Tensor
            self.score_state: torch.Tensor
            state_shape = (args.max_batch_size, compress_ratio, head_dim)
            self.register_buffer("kv_state", torch.zeros(state_shape, dtype=torch.float32), persistent=False)
            self.register_buffer(
                "score_state", torch.full(state_shape, -torch.inf, dtype=torch.float32), persistent=False
            )

    def forward(self, x: torch.Tensor, start_pos: int) -> torch.Tensor | None:
        bsz, seqlen, _ = x.size()
        ratio, dtype = self.compress_ratio, x.dtype
        if ratio == 1:  # one token per group: nothing to pool, so no gate and no fp32
            return self.norm(self.wkv(x))

        x = x.float()
        kv, score = self.wkv(x), self.wgate(x)
        if start_pos == 0:
            should_compress = seqlen >= ratio
            remainder = seqlen % ratio
            cutoff = seqlen - remainder
            if remainder:  # trailing partial group waits in the state
                kv, self.kv_state[:bsz, :remainder] = kv.split([cutoff, remainder], dim=1)
                score, self.score_state[:bsz, :remainder] = score.split([cutoff, remainder], dim=1)
            kv = kv.unflatten(1, (-1, ratio))
            score = score.unflatten(1, (-1, ratio))
            kv = (kv * score.softmax(dim=2)).sum(dim=2)
        else:  # one token per step: fill a slot, and pool only when the group just completed
            should_compress = (start_pos + 1) % ratio == 0
            slot = start_pos % ratio
            self.kv_state[:bsz, slot] = kv.squeeze(1)
            self.score_state[:bsz, slot] = score.squeeze(1)
            if should_compress:
                kv = (self.kv_state[:bsz] * self.score_state[:bsz].softmax(dim=1)).sum(dim=1, keepdim=True)
        if not should_compress:
            return None
        return self.norm(kv.to(dtype))

model.py 281–293

class RMSNorm(nn.Module):
    def __init__(self, dim: int, eps: float = 1e-6):
        super().__init__()
        self.dim = dim
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))

    def forward(self, x: torch.Tensor):
        dtype = x.dtype
        x = x.float()
        var = x.square().mean(-1, keepdim=True)
        x = x * torch.rsqrt(var + self.eps)
        return (self.weight * x).to(dtype)

Sources2

Direct evidence

The compressor makes the compressed key–value latents: at ratio 1 one projection and a norm per token; above 1 a per-feature softmax over each group of tokens, with an unfinished group kept for the next call.

compressor · deepseek-ai/DeepSeek-V4.1-Flash / inference/model.py
Pinned revision 2cba9e42aa02 ↗ · lines 429–485

Read code excerpt

class Compressor(nn.Module):
    """Pools `compress_ratio` consecutive tokens into one KV latent with a learned softmax gate.

    Returns the latent before RoPE, or None while a group is still filling up -- so during decode it
    only yields every `compress_ratio` steps, holding the partial group in `kv_state`/`score_state`.
    Pre-RoPE is deliberate: the indexer needs the unrotated form, so Attention rotates afterwards.
    """

    def __init__(self, args: ModelArgs, layer_id: int):
        super().__init__()
        compress_ratio = args.compress_ratios[layer_id]
        head_dim = args.head_dim
        self.compress_ratio = compress_ratio
        self.head_dim = head_dim
        self.norm = RMSNorm(head_dim, args.norm_eps)
        # ratio 1 is a plain projection, so it stays in the checkpoint's bf16; the softmax pooling
        # above ratio 1 runs in fp32, so those weights are promoted to fp32 to match
        self.wkv = Linear(args.dim, head_dim, dtype=torch.float32 if compress_ratio > 1 else torch.bfloat16)
        if compress_ratio > 1:
            self.wgate = Linear(args.dim, head_dim, dtype=torch.float32)
            # tail of an incomplete group, carried across decode steps
            self.kv_state: torch.Tensor
            self.score_state: torch.Tensor
            state_shape = (args.max_batch_size, compress_ratio, head_dim)
            self.register_buffer("kv_state", torch.zeros(state_shape, dtype=torch.float32), persistent=False)
            self.register_buffer(
                "score_state", torch.full(state_shape, -torch.inf, dtype=torch.float32), persistent=False
            )

    def forward(self, x: torch.Tensor, start_pos: int) -> torch.Tensor | None:
        bsz, seqlen, _ = x.size()
        ratio, dtype = self.compress_ratio, x.dtype
        if ratio == 1:  # one token per group: nothing to pool, so no gate and no fp32
            return self.norm(self.wkv(x))

        x = x.float()
        kv, score = self.wkv(x), self.wgate(x)
        if start_pos == 0:
            should_compress = seqlen >= ratio
            remainder = seqlen % ratio
            cutoff = seqlen - remainder
            if remainder:  # trailing partial group waits in the state
                kv, self.kv_state[:bsz, :remainder] = kv.split([cutoff, remainder], dim=1)
                score, self.score_state[:bsz, :remainder] = score.split([cutoff, remainder], dim=1)
            kv = kv.unflatten(1, (-1, ratio))
            score = score.unflatten(1, (-1, ratio))
            kv = (kv * score.softmax(dim=2)).sum(dim=2)
        else:  # one token per step: fill a slot, and pool only when the group just completed
            should_compress = (start_pos + 1) % ratio == 0
            slot = start_pos % ratio
            self.kv_state[:bsz, slot] = kv.squeeze(1)
            self.score_state[:bsz, slot] = score.squeeze(1)
            if should_compress:
                kv = (self.kv_state[:bsz] * self.score_state[:bsz].softmax(dim=1)).sum(dim=1, keepdim=True)
        if not should_compress:
            return None
        return self.norm(kv.to(dtype))

Direct evidence

RMSNorm divides by the root mean square over the features, computed in float32, and multiplies by a learned gain.

rmsnorm · deepseek-ai/DeepSeek-V4.1-Flash / inference/model.py
Pinned revision 2cba9e42aa02 ↗ · lines 281–293

Read code excerpt

class RMSNorm(nn.Module):
    def __init__(self, dim: int, eps: float = 1e-6):
        super().__init__()
        self.dim = dim
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))

    def forward(self, x: torch.Tensor):
        dtype = x.dtype
        x = x.float()
        var = x.square().mean(-1, keepdim=True)
        x = x * torch.rsqrt(var + self.eps)
        return (self.weight * x).to(dtype)