This document is relevant for: Trn2, Trn3

MLA Decode Attention Kernel API Reference#

Computes one DeepSeek-V3.2 MLA decode step’s attention over the prior keys an indexer selected, then projects the result out of the cache’s latent space and through the o-projection, returning this rank’s attention-block output.

This is the back half of an MLA decode layer, and the counterpart of the QKV front-fold. It is designed to be traced inline after that front-fold rather than run standalone: the two query tiles arrive in SBUF, already feature-major in the layout the score matmuls index, so a standalone caller has to stage them itself. Selection has already happened, so there is no sparse score bias here. What the selection leaves behind is a per-token map from key slot to absolute position, and that map drives the causal mask. Rather than read a pre-gathered feature-major prior out of HBM, the kernel gathers each token’s selected rows straight from kv_cache in the cache’s native token-major layout, and then transposes once to feature-major for the score matmul. The weighted sum contracts over keys and so wants the token-major tile, which is what the gather already produced, and reads it directly. That deletes a 33.5 MB HBM round trip and one of two mirror-image transposes.

Background#

The mla_decode_attention kernel computes DeepSeek-V3.2 MLA decode attention over the prior keys selected by an indexer, then applies the V up-projection and output projection. It is designed to be traced inline after mla_decode_qkv.

API Reference#

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

mla_decode_attention#

nkilib.experimental.deepseek_v32.mla_decode_attention.mla_decode_attention(q_absorbed_sbuf: nl.NkiTensor, q_pe_sbuf: nl.NkiTensor, kv_cache: nl.NkiTensor, kv_row_cur: nl.NkiTensor, prior_rows: nl.NkiTensor, prior_key_pos: nl.NkiTensor, prior_rows_natural: nl.NkiTensor, positions: nl.NkiTensor, out_absorb_w: nl.NkiTensor, o_proj_w: nl.NkiTensor, n_prior: int, softmax_scale: float) → nl.NkiTensor#

Computes one DeepSeek-V3.2 MLA decode step’s attention over the prior keys an indexer selected, then projects the result out of the cache’s latent space and through the o-projection, returning this rank’s attention-block output.

Parameters:

  • q_absorbed_sbuf (nl.NkiTensor) – [P_MAX, n_latent_tiles, n_tokens, n_heads]. The query already mapped into the cache’s latent space by the front-fold’s absorb.

  • q_pe_sbuf (nl.NkiTensor) – [P_MAX, n_tokens, n_heads]. The query’s RoPE’d positional half.

  • kv_cache (nl.NkiTensor) – [n_blocks, 1, block_size, kv_row_dim]. Paged cache, read only here. The current token’s row must already have been scattered into it by the front-fold.

  • kv_row_cur (nl.NkiTensor) – [n_tokens, kv_row_dim]. This step’s cache row per token, laid out [k_pe | kv_latent].

  • prior_rows (nl.NkiTensor) – [P_MAX, n_tokens, n_prior_tiles] int32. Flat cache row per selected key, in the batched gather’s consumption order. From build_prior_index.

  • prior_key_pos (nl.NkiTensor) – [P_MAX, n_tokens, n_prior_tiles] fp32. Absolute position of each selected key, key-major. From build_prior_index.

  • prior_rows_natural (nl.NkiTensor) – [P_MAX, n_tokens, n_prior_tiles] uint32. The same rows in natural key order, which is what the indirect dma_transpose needs. From build_prior_index.

  • positions (nl.NkiTensor) – [n_tokens] int32. This step’s decode position per sequence.

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

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

  • n_prior (int) – Prior keys the indexer selected per token.

  • softmax_scale (float) – Scale applied to the scores before masking and softmax.

Notes:

  • Scores and softmax weights are held key-major: key tile * P_MAX + p of head h sits at [p, tile, h]. Swapping which matmul operand is stationary transposes the output for free, and putting keys on partitions pays off three times over. The softmax then runs 128 lanes wide instead of on the single partition a [n_heads, n_keys] layout leaves at n_heads == 1; the causal mask becomes a plain per-partition column instead of a broadcast matmul; and the weighted sum’s stationary operand is already oriented correctly, deleting a transpose. This was the single largest win in this kernel’s history, 521.5 -> 402.2 us, and most of it came from lane utilization rather than from the transpose removal. The out-absorb and o-projection weights do not depend on the token, so they are loaded once above the token loop. Both tails are also batched across tokens: they are per-token GEMVs against shared weights, and a 128x1 matmul costs about 360 ns against 179 ns for a 128x128, because the cost is set by the moving operand and a width-1 stationary pays full price for 1/128 of the array. Emitting the weighted sum for the previous token, so the PE could run ahead of a softmax, was measured and is exactly neutral: the post-scheduler reorders globally, so NKI emission order does not control per-engine issue order. Do not expect a win from hand-scheduling here. Giving each latent tile its own 32-partition band via tile_position, so the tiles issue concurrently instead of chaining through one PSUM bank, was measured at 21% slower (433.8 -> 525.7 us). It band-sums per token and key chunk rather than once per kernel, and each sum is three extra full-width Vector operations on what is already the bottleneck engine.

Dimensions:

  • n_tokens: Decode tokens in the batch, one per sequence. Must be <= P_MAX.

  • n_heads: Attention heads on this rank. Must be <= P_MAX.

  • n_prior: Prior keys the indexer selected per token.

  • n_keys: Keys scored per token, n_prior + 1, the selected prior plus the current token.

  • kv_lora_rank: LoRA rank of the KV path, and the width of a cache row’s latent field. Must be a multiple of P_MAX and fit one PSUM bank.

  • qk_rope_dim: Width of a cache row’s k_pe field. Must be <= P_MAX.

  • v_head_dim: Per-head width the out-absorb produces. Must be <= P_MAX.

  • hidden: Model hidden width, which the o-projection emits.

This document is relevant for: Trn2, Trn3