This document is relevant for: Trn2, Trn3

MLA Sparse Attention CTE Kernel API Reference#

Standalone sparse latent + RoPE attention (S-sharded across cores).

KERNEL A of the split DeepSeek-V3.2 sparse-MLA forward: the un-projected latent value path. Computes, per query, absorbed-latent sparse attention over the topk-selected cache rows. Intended for Context Encoding with DeepSeek-V3.2 dims (L == 512 kv_lora_rank, R == 64, up to 128 heads, topk K a multiple of 128 up to ~2048); pair with the o_proj kernel B for V-up + o_proj. Requires B == 1 and S divisible by the number of cores.

Background#

The mla_sparse_attention_cte_kernel kernel performs standalone sparse latent + RoPE attention (S-sharded across cores).

API Reference#

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

mla_sparse_attention_cte_kernel#

nkilib.experimental.mla.deepseek.mla_sparse_attention_cte_kernel(q_lift_hbm: nl.NkiTensor, q_pe_hbm: nl.NkiTensor, c_kv_hbm: nl.NkiTensor, k_pe_hbm: nl.NkiTensor, topk_indices_hbm: nl.NkiTensor, softmax_scale: float, topk_tiled: bool = False) → nl.NkiTensor#

Standalone sparse latent + RoPE attention (S-sharded across cores).

Parameters:

  • q_lift_hbm (nl.NkiTensor) – [B, S, H, L] bf16, per-head absorbed Q latent.

  • q_pe_hbm (nl.NkiTensor) – [B, S, H, R] bf16, per-head pre-rotated RoPE queries.

  • c_kv_hbm (nl.NkiTensor) – [B, S_kv, L] bf16, latent KV cache.

  • k_pe_hbm (nl.NkiTensor) – [B, S_kv, R] bf16, pre-rotated RoPE key cache.

  • topk_indices_hbm (nl.NkiTensor) – int32 topk cache-row indices. Flat [B, S, K] when topk_tiled is False; partition-tiled [num_s_tiles, NUM_TOPK_BATCHES, P_MAX, K // 16] when topk_tiled is True.

  • softmax_scale (float) – Scaling factor applied to the attention scores. Must be positive. DeepSeek’s scale is head_dim**-0.5 times a squared mscale correction, so it is always positive in practice.

  • topk_tiled (bool) – Select the topk_indices_hbm layout (default False = flat).

Returns:

[B, S, H * L] bf16 latent attention output, row-major h * L + l, with each head’s latent columns pre-permuted into MX 4-pack order.

Return type:

nl.ndarray

Notes:

  • S-sharded across cores; under Context Parallelism the framework gathers the latent KV to the full S_kv before this kernel is called.

  • The MM2 output is written with each head’s latent columns pre-permuted into MX 4-pack order (natural latent l = 4*group + sub at physical column sub*(L//4) + group), the cross-kernel layout contract the o_proj kernel’s DMA-transpose load depends on. Keep in lockstep with the o_proj load, W_uv load, and both refs.

Dimensions:

  • B: Batch size (must be 1)

  • S: Query sequence length (this rank’s S-shard)

  • S_kv: Cache (key/value) sequence length

  • H: Number of attention heads

  • L: Latent (kv_lora_rank) dimension (must be 512 = P_MAX * 4)

  • R: RoPE head dimension

This document is relevant for: Trn2, Trn3