This document is relevant for: Trn2, Trn3

MLA Attention Block Gather Decode Kernel API Reference#

Computes one DeepSeek-V3.2 MLA decode step end to end, from the residual stream to this rank’s attention-block output, attending over only the prior keys the Lightning Indexer selected.

This is the whole attention block: it traces the QKV front-fold and the attention core into a single NEFF rather than calling them as separate kernels. That is what makes the composition worth having, because the front-fold’s two query tiles stay in SBUF and are handed straight to the core. Run standalone the tiles would have to round-trip through HBM, and the core’s signature is built for the inline hand-off, not for a caller staging them itself. The block owns no numerics of its own. Every value it returns is computed by one of the two kernels it composes, so its correctness is theirs plus the wiring, and its reference composes their references for the same reason.

Background#

The mla_attention_block_gather_decode kernel computes one DeepSeek-V3.2 MLA decode step end to end, from the residual stream to the rank’s attention-block output, attending only over the prior keys selected by the lightning indexer. It composes mla_decode_qkv and mla_decode_attention in a single kernel so the query tiles stay in SBUF.

API Reference#

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

mla_attention_block_gather_decode#

nkilib.experimental.deepseek_v32.mla_attention_block_gather_decode.mla_attention_block_gather_decode(hidden_states: nl.NkiTensor, input_norm_w: nl.NkiTensor, wq_a: nl.NkiTensor, q_norm_w: nl.NkiTensor, wq_b: nl.NkiTensor, wkv_a: nl.NkiTensor, kv_norm_w: nl.NkiTensor, q_absorb_w: nl.NkiTensor, out_absorb_w: nl.NkiTensor, o_proj_w: nl.NkiTensor, cos: nl.NkiTensor, sin: nl.NkiTensor, kv_cache: nl.NkiTensor, slot_mapping: nl.NkiTensor, block_table: nl.NkiTensor, positions: nl.NkiTensor, gather_indices: nl.NkiTensor, n_heads: int, qk_nope_head_dim: int, qk_rope_head_dim: int, kv_lora_rank: int, softmax_scale: float, norm_eps: float = 1e-06) → nl.NkiTensor#

Computes one DeepSeek-V3.2 MLA decode step end to end, from the residual stream to this rank’s attention-block output, attending over only the prior keys the Lightning Indexer selected.

Parameters:

  • hidden_states (nl.NkiTensor) – [n_tokens, hidden]. This step’s residual-stream hidden states.

  • input_norm_w (nl.NkiTensor) – [1, hidden] fp32. RMSNorm gamma for the layer input, folded in by the front-fold rather than applied by the caller.

  • wq_a (nl.NkiTensor) – [hidden, q_lora_rank]. Query down-projection.

  • q_norm_w (nl.NkiTensor) – [1, q_lora_rank] fp32. RMSNorm gamma for the query latent.

  • wq_b (nl.NkiTensor) – [q_lora_rank, qk_packed_dim]. Query up-projection, emitting every head’s [q_nope | q_pe] concatenated.

  • wkv_a (nl.NkiTensor) – [hidden, kv_row_dim]. KV down-projection, emitting one cache row per token.

  • kv_norm_w (nl.NkiTensor) – [1, kv_lora_rank] fp32. RMSNorm gamma for the KV latent.

  • q_absorb_w (nl.NkiTensor) – [n_heads, qk_nope_dim, kv_lora_rank]. The nope half of wkv_b, which maps a head’s nope query into the cache’s latent space.

  • out_absorb_w (nl.NkiTensor) – [n_heads, v_head_dim, kv_lora_rank]. The value half of wkv_b, which maps the latent attention output back out to a head’s value width.

  • o_proj_w (nl.NkiTensor) – [n_heads * v_head_dim, hidden]. This rank’s shard of the o-projection.

  • cos (nl.NkiTensor) – [n_tokens, qk_rope_dim // 2]. Cosine table for interleaved RoPE.

  • sin (nl.NkiTensor) – [n_tokens, qk_rope_dim // 2]. Sine table for interleaved RoPE.

  • kv_cache (nl.NkiTensor) – [n_blocks, 1, block_size, kv_row_dim]. Paged cache, mutated in place: the front-fold writes this step’s row, and the core reads the selected prior rows back.

  • slot_mapping (nl.NkiTensor) – [n_tokens] int32. Flat cache row each token’s own key is written to.

  • block_table (nl.NkiTensor) – [n_tokens, max_blocks_per_seq] int32. Per-sequence logical-to-physical block map, which resolves a selected position to a cache row.

  • positions (nl.NkiTensor) – [n_tokens] int32. This step’s absolute position per token, which the causal mask compares each selected key against.

  • gather_indices (nl.NkiTensor) – [n_tokens, n_prior] int32. The absolute prior positions the indexer selected, unsorted. Its width fixes n_prior.

  • n_heads (int) – Attention heads on this rank.

  • qk_nope_head_dim (int) – Per-head query/key width that RoPE does not rotate.

  • qk_rope_head_dim (int) – Per-head query/key width that RoPE rotates.

  • kv_lora_rank (int) – LoRA rank of the KV path, and the latent width of a cache row.

  • softmax_scale (float) – Score scale, conventionally the inverse square root of the full per-head query/key width.

  • norm_eps (float) – Epsilon for every RMSNorm in the front-fold.

Notes:

  • The index build is issued first, ahead of the front-fold that does not consume it. Its two inputs are under 1 MB together, but queued behind the projection weights they sit behind roughly 31 MB on the same DMA queue and do not land until about 97 us, and no per-token gather in the core can start until they do. Fusing the two halves is worth about 45 us at the shipping shape, measured as 385.8 us for the two run separately against 340.2 us here. The saving is the HBM round trip the query tiles would otherwise make, which is why the core takes them in SBUF.

Dimensions:

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

  • n_heads: Attention heads on this rank.

  • n_prior: Prior keys the indexer selected per token, taken from gather_indices.

This document is relevant for: Trn2, Trn3