This document is relevant for: Trn2, Trn3

DeepSeek-V3.2 Sparse Indexer Kernel API Reference#

Score every cached key slot against the current query and select the highest scoring ones.

This is the DSA lightning indexer for one decode step, optimized for the token-generation regime where each sequence contributes a single token: n_tokens <= P_MAX, with 16 at the shipping shape. The query projection is always MX-quantized, so it requires trn3 or later.

Background#

The deepseek_v32_sparse_indexer kernel implements the DeepSeek-V3.2 sparse attention (DSA) lightning indexer for one decode step. It scores every cached key slot against the current query and selects the highest-scoring slots. It requires Trn3 or later.

API Reference#

Source code for this kernel API can be found at: sparse_indexer.py

deepseek_v32_sparse_indexer#

nkilib.experimental.deepseek_v32.sparse_indexer.deepseek_v32_sparse_indexer(hidden_states: nl.NkiTensor, qr: nl.NkiTensor, wq_b: nl.NkiTensor, wq_b_scale: nl.NkiTensor, wk: nl.NkiTensor, k_norm_weight: nl.NkiTensor, k_norm_bias: nl.NkiTensor, weights_proj: nl.NkiTensor, hadamard: nl.NkiTensor, cos: nl.NkiTensor, sin: nl.NkiTensor, key_cache: nl.NkiTensor, positions: nl.NkiTensor, index_topk: int, valid_len: int | None = None) → DeepseekV32SparseIndexerOutput#

Score every cached key slot against the current query and select the highest scoring ones.

Parameters:

  • hidden_states (nl.NkiTensor) – [n_tokens, hidden], Hidden states for this step

  • qr (nl.NkiTensor) – [n_tokens, q_lora_rank], q_norm(wq_a(x)) from the MLA block

  • wq_b (nl.NkiTensor) – [q_lora_rank // K_PACK, n_heads * head_dim], Query projection weight, fp8x4 MX-packed along the contraction dimension

  • wq_b_scale (nl.NkiTensor) – [q_lora_rank // MX_BLOCK, n_heads * head_dim], uint8 e8m0 MX scale paired with wq_b

  • wk (nl.NkiTensor) – [hidden, head_dim], Key projection weight

  • k_norm_weight (nl.NkiTensor) – [1, head_dim], fp32 LayerNorm gain for the key

  • k_norm_bias (nl.NkiTensor) – [1, head_dim], fp32 LayerNorm bias for the key

  • weights_proj (nl.NkiTensor) – [hidden, n_heads], Per-head weight projection

  • hadamard (nl.NkiTensor) – [head_dim, head_dim], Scaled Hadamard matrix with 1/sqrt(head_dim) pre-baked

  • cos (nl.NkiTensor) – [n_tokens, rope_dim // 2], RoPE cosine table

  • sin (nl.NkiTensor) – [n_tokens, rope_dim // 2], RoPE sine table

  • key_cache (nl.NkiTensor) – [n_tokens, max_seq_len, head_dim], Caller-owned key cache, written in place at each token’s current position

  • positions (nl.NkiTensor) – [n_tokens], int32 absolute position of each sequence’s current token

  • shapes – Input shapes and derived tiling constants

Returns:

[n_tokens, max_seq_len], fp32 per-slot relevance, with slots the token cannot attend to clamped to the mask bound

Return type:

nl.NkiTensor

Returns:

[n_tokens, index_topk], int32 selected slot positions

Return type:

nl.NkiTensor

Notes:

  • topk_indices is an unordered set: column order carries no meaning.

  • The mask bound is finite rather than -inf because the top-k consumer reads bfloat16, where -inf makes the hardware max-select produce NaN.

  • q_lora_rank must be a multiple of P_MAX * K_PACK, the MX K-tile.

Dimensions:

  • n_tokens: Decode tokens in the batch, one per sequence

  • n_heads: Indexer heads

  • head_dim: Per-head feature width, which must equal P_MAX

  • hidden: Model hidden width, the contraction dimension of the key projection

  • q_lora_rank: LoRA rank of the query path, the contraction dimension of the query projection

  • rope_dim: Leading feature span that RoPE rotates

  • max_seq_len: Key-cache capacity per sequence, in slots

This document is relevant for: Trn2, Trn3