DeepSeek-V3.2-Exp · Lightning indexer

61 transformer blocks › 58 MoE blocks › Sparse latent attention

Lightning indexer

learnedself.indexer

Score every position up to its own for each query cheaply and keep the 2,048 best; attention then uses only those.

Why

To choose positions the model needs a score for every pair of positions, which is what attention itself computes. The indexer computes a far cheaper score for every pair: few heads, one key per position shared by them, ReLU chosen for throughput, and it can run in FP8. DeepSeek-V3.2-Exp report §1

Its scores can stand in for attention because it is trained to match it: the target is the main attention’s scores summed over heads and normalized, and the loss is a KL divergence. DeepSeek-V3.2-Exp report §2.1

The indexer still costs O(L²). The authors report that it needs much less computation than the attention of the model before it, and that selecting k positions takes the main attention to O(Lk). DeepSeek-V3.2-Exp report §3

\tilde S^{I}=S^{I}+{\htmlClass{vl-stream-context}{M}}

{\htmlClass{vl-stream-context}{M^{\mathrm{sel}}_{s}}}=\begin{cases}0&s\in\operatorname{top}_{1,001}(S^{I})\\-\infty&\text{otherwise}\end{cases}

All 14 steps

\text{Indexer queries:}

{\htmlClass{vl-stream-q}{q^{I}}}=\hat c^{{\htmlClass{vl-stream-q}{Q}}}W_{IQ}

{\htmlClass{vl-stream-q}{Q^{I}}}=\operatorname{reshape}({\htmlClass{vl-stream-q}{q^{I}}})

\tilde {\htmlClass{vl-stream-q}{Q^{I}}}=\operatorname{RoPE}_{64}({\htmlClass{vl-stream-q}{Q^{I}}},{\htmlClass{vl-stream-context}{R}})

\bar {\htmlClass{vl-stream-q}{Q^{I}}}=\tilde {\htmlClass{vl-stream-q}{Q^{I}}}\,\tfrac{1}{\sqrt{128}}\mathbf{H}_{128}

\text{Indexer key:}

{\htmlClass{vl-stream-k}{k^{I}}}=xW_{IK}

\hat {\htmlClass{vl-stream-k}{k^{I}}}=\operatorname{LayerNorm}({\htmlClass{vl-stream-k}{k^{I}}})

\tilde {\htmlClass{vl-stream-k}{k^{I}}}=\operatorname{RoPE}_{64}(\hat {\htmlClass{vl-stream-k}{k^{I}}},{\htmlClass{vl-stream-context}{R}})

\bar {\htmlClass{vl-stream-k}{k^{I}}}=\tilde {\htmlClass{vl-stream-k}{k^{I}}}\,\tfrac{1}{\sqrt{128}}\mathbf{H}_{128}

w=xW_{w}

S^{I}_{t,s}=\frac{1}{\sqrt{64}\sqrt{128}}\sum_{h=1}^{64}w_{t,h}\,\operatorname{ReLU}\!\big(\bar {\htmlClass{vl-stream-q}{Q^{I}_{h,t}}}\cdot\bar {\htmlClass{vl-stream-k}{k^{I}_{s}}}\big)

\tilde S^{I}=S^{I}+{\htmlClass{vl-stream-context}{M}}

{\htmlClass{vl-stream-context}{M^{\mathrm{sel}}_{s}}}=\begin{cases}0&s\in\operatorname{top}_{1,001}(S^{I})\\-\infty&\text{otherwise}\end{cases}

in: $x$ [B batch, 1 positions, 7168 d] $\hat c^{Q}$ [B batch, 1 positions, 1536 q latent] $R$ [1 positions, 32 pairs] $M$ [1 queries, 1001 keys]
out: $M^{\mathrm{sel}}$ [B batch, 1 queries, 1001 keys]
parameters: 13,959,424 · × 58 MoE blocks = 809,646,592

Parameters

13,959,424 parameters per instance, across its learned components

Exact per-instance accounting
Component	Calculation	Parameters
Indexer queries	12,582,912 across its components	12,582,912
Indexer key	917,760 across its components	917,760
Head weight projection	$7{,}168 \times 64$	458,752
Per-instance total		13,959,424

Across 58 independent MoE blocks: $13{,}959{,}424 \times 58 = 809{,}646{,}592$ 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

This is DeepSeek Sparse Attention, the one architectural change from V3.1: each query attends to at most 2,048 positions, the ones chosen here. With 64 heads of 128 features and one key shared by all heads, scoring every pair of positions costs far less than attention itself.

With 2,048 or fewer positions in the sequence, every permitted position is kept and attention is dense.

Codemodel.py 436–452, 457–487 · kernel.py 260–273

model.py 436–452

    def __init__(self, args: ModelArgs):
        super().__init__()
        self.dim: int = args.dim
        self.n_heads: int = args.index_n_heads
        self.n_local_heads = args.index_n_heads // world_size
        self.head_dim: int = args.index_head_dim
        self.rope_head_dim: int = args.qk_rope_head_dim
        self.index_topk: int = args.index_topk
        self.q_lora_rank: int = args.q_lora_rank
        self.wq_b = Linear(self.q_lora_rank, self.n_heads * self.head_dim)
        self.wk = Linear(self.dim, self.head_dim)
        self.k_norm = LayerNorm(self.head_dim)
        # weights_proj in the checkpoint is stored in bf16, while the parameters here are stored in fp32 for convenient.
        self.weights_proj = Linear(self.dim, self.n_heads, dtype=torch.float32)
        self.softmax_scale = self.head_dim ** -0.5
        self.scale_fmt = args.scale_fmt

model.py 457–487

    def forward(self, x: torch.Tensor, qr: torch.Tensor, start_pos: int, freqs_cis: torch.Tensor, mask: Optional[torch.Tensor]):
        bsz, seqlen, _ = x.size()
        end_pos = start_pos + seqlen
        q = self.wq_b(qr)
        q = q.view(bsz, seqlen, self.n_heads, self.head_dim)
        q_pe, q_nope = torch.split(q, [self.rope_head_dim, self.head_dim - self.rope_head_dim], dim=-1)
        # rope in indexer is not interleaved
        q_pe = apply_rotary_emb(q_pe, freqs_cis, False)
        q = torch.cat([q_pe, q_nope], dim=-1)
        k = self.wk(x)
        k = self.k_norm(k)
        k_pe, k_nope = torch.split(k, [self.rope_head_dim, self.head_dim - self.rope_head_dim], dim=-1)
        # rope in indexer is not interleaved
        k_pe = apply_rotary_emb(k_pe.unsqueeze(2), freqs_cis, False).squeeze(2)
        k = torch.cat([k_pe, k_nope], dim=-1)
        q = rotate_activation(q)
        k = rotate_activation(k)
        q_fp8, q_scale = act_quant(q, block_size, self.scale_fmt)
        k_fp8, k_scale = act_quant(k, block_size, self.scale_fmt)
        self.k_cache[:bsz, start_pos:end_pos] = k_fp8
        self.k_scale_cache[:bsz, start_pos:end_pos] = k_scale
        weights = self.weights_proj(x.float()) * self.n_heads ** -0.5
        weights = weights.unsqueeze(-1) * q_scale * self.softmax_scale
        index_score = fp8_index(q_fp8.contiguous(), weights, self.k_cache[:bsz, :end_pos].contiguous(), self.k_scale_cache[:bsz, :end_pos].contiguous())
        if mask is not None:
            index_score += mask
        topk_indices = index_score.topk(min(self.index_topk, end_pos), dim=-1)[1]
        topk_indices_ = topk_indices.clone()
        dist.broadcast(topk_indices_, src=0)
        assert torch.all(topk_indices == topk_indices_), f"{topk_indices=} {topk_indices_=}"
        return topk_indices

kernel.py 260–273

    """
    Perform index score using FP8 precision.

    Args:
        q (torch.Tensor): The Q tensor, must be contiguous.
        q_s (torch.Tensor): The scaling factor for Q (float), must be contiguous.
        k (torch.Tensor): The K tensor, must be contiguous.
        k_s (torch.Tensor): The scaling factor for K (e8m0 here), must be contiguous.

        fp8 q @ fp8 k -> fp32 logits
        relu(fp32 logits) * q_s (weights) -> fp32 logits
        fp32 logits -> fp32 logits_sum
        fp32 logits_sum * k_s (e8m0) -> fp32 index_score
    """

Sources3

Direct evidence

The lightning indexer owns a query up-projection from the query latent, a single key projection with LayerNorm, and a per-head weight projection.

indexer-init · deepseek-ai/DeepSeek-V3.2-Exp / inference/model.py
Pinned revision 87e509a2e5a1 ↗ · lines 436–452

Read code excerpt

    def __init__(self, args: ModelArgs):
        super().__init__()
        self.dim: int = args.dim
        self.n_heads: int = args.index_n_heads
        self.n_local_heads = args.index_n_heads // world_size
        self.head_dim: int = args.index_head_dim
        self.rope_head_dim: int = args.qk_rope_head_dim
        self.index_topk: int = args.index_topk
        self.q_lora_rank: int = args.q_lora_rank
        self.wq_b = Linear(self.q_lora_rank, self.n_heads * self.head_dim)
        self.wk = Linear(self.dim, self.head_dim)
        self.k_norm = LayerNorm(self.head_dim)
        # weights_proj in the checkpoint is stored in bf16, while the parameters here are stored in fp32 for convenient.
        self.weights_proj = Linear(self.dim, self.n_heads, dtype=torch.float32)
        self.softmax_scale = self.head_dim ** -0.5
        self.scale_fmt = args.scale_fmt

Direct evidence

Indexer scores combine per-head query–key products with learned per-head weights, add the causal mask and keep the top-k positions for each query.

indexer-forward · deepseek-ai/DeepSeek-V3.2-Exp / inference/model.py
Pinned revision 87e509a2e5a1 ↗ · lines 457–487

Read code excerpt

    def forward(self, x: torch.Tensor, qr: torch.Tensor, start_pos: int, freqs_cis: torch.Tensor, mask: Optional[torch.Tensor]):
        bsz, seqlen, _ = x.size()
        end_pos = start_pos + seqlen
        q = self.wq_b(qr)
        q = q.view(bsz, seqlen, self.n_heads, self.head_dim)
        q_pe, q_nope = torch.split(q, [self.rope_head_dim, self.head_dim - self.rope_head_dim], dim=-1)
        # rope in indexer is not interleaved
        q_pe = apply_rotary_emb(q_pe, freqs_cis, False)
        q = torch.cat([q_pe, q_nope], dim=-1)
        k = self.wk(x)
        k = self.k_norm(k)
        k_pe, k_nope = torch.split(k, [self.rope_head_dim, self.head_dim - self.rope_head_dim], dim=-1)
        # rope in indexer is not interleaved
        k_pe = apply_rotary_emb(k_pe.unsqueeze(2), freqs_cis, False).squeeze(2)
        k = torch.cat([k_pe, k_nope], dim=-1)
        q = rotate_activation(q)
        k = rotate_activation(k)
        q_fp8, q_scale = act_quant(q, block_size, self.scale_fmt)
        k_fp8, k_scale = act_quant(k, block_size, self.scale_fmt)
        self.k_cache[:bsz, start_pos:end_pos] = k_fp8
        self.k_scale_cache[:bsz, start_pos:end_pos] = k_scale
        weights = self.weights_proj(x.float()) * self.n_heads ** -0.5
        weights = weights.unsqueeze(-1) * q_scale * self.softmax_scale
        index_score = fp8_index(q_fp8.contiguous(), weights, self.k_cache[:bsz, :end_pos].contiguous(), self.k_scale_cache[:bsz, :end_pos].contiguous())
        if mask is not None:
            index_score += mask
        topk_indices = index_score.topk(min(self.index_topk, end_pos), dim=-1)[1]
        topk_indices_ = topk_indices.clone()
        dist.broadcast(topk_indices_, src=0)
        assert torch.all(topk_indices == topk_indices_), f"{topk_indices=} {topk_indices_=}"
        return topk_indices

Direct evidence

The index kernel applies ReLU to each head’s query–key product, weights it per head and sums over heads.

index-kernel · deepseek-ai/DeepSeek-V3.2-Exp / inference/kernel.py
Pinned revision 87e509a2e5a1 ↗ · lines 260–273

Read code excerpt

    """
    Perform index score using FP8 precision.

    Args:
        q (torch.Tensor): The Q tensor, must be contiguous.
        q_s (torch.Tensor): The scaling factor for Q (float), must be contiguous.
        k (torch.Tensor): The K tensor, must be contiguous.
        k_s (torch.Tensor): The scaling factor for K (e8m0 here), must be contiguous.

        fp8 q @ fp8 k -> fp32 logits
        relu(fp32 logits) * q_s (weights) -> fp32 logits
        fp32 logits -> fp32 logits_sum
        fp32 logits_sum * k_s (e8m0) -> fp32 index_score
    """