This document is relevant for: Trn2, Trn3

MLA QKV CTE Kernel API Reference#

DeepSeek MLA QKV projection (MX), emitting absorbed latents.

Consumes the packed MX activation from rmsnorm_mx_prefill (pack_scales=True) directly — no in-kernel quantize. Intended for Context Encoding (prefill) with DeepSeek-V3.2 dimensions: absorption requires qk_nope_head_dim == 128, and the kernel is tuned for n_heads up to 128 and hidden dim H up to 7168 with LNC sharding over the sequence dimension. Best used when the packed MX activation is produced upstream so no re-quantization is needed.

Background#

The mla_qkv_cte_kernel kernel performs DeepSeek MLA QKV projection (MX), emitting absorbed latents.

API Reference#

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

mla_qkv_cte_kernel#

nkilib.experimental.mla.deepseek.mla_qkv_cte_kernel(x_hbm_mx: nl.NkiTensor, wqkv_a_hbm: nl.NkiTensor, wqkv_a_scale_hbm: nl.NkiTensor, wq_b_hbm: nl.NkiTensor, wq_b_scale_hbm: nl.NkiTensor, q_norm_gamma_hbm: nl.NkiTensor, kv_norm_gamma_hbm: nl.NkiTensor, wuk_hbm: nl.NkiTensor, cos_cache_hbm: nl.NkiTensor, sin_cache_hbm: nl.NkiTensor, n_heads: int, qk_nope_head_dim: int, qk_rope_head_dim: int, kv_lora_rank: int, qk_lora_rank: int, norm_eps: float = 1e-06, compact_scales: bool = True, precision: MlaPrecision = MlaPrecision.MX) → Tuple[nl.NkiTensor, ...]#

DeepSeek MLA QKV projection (MX), emitting absorbed latents.

Parameters:

  • x_hbm_mx (nl.NkiTensor) – [B, S, H + scale_region] fp8 packed MX activation (see the param comment).

  • wqkv_a_hbm (nl.NkiTensor) – [H // 4, qk_lora_rank + kv_lora_rank + qk_rope_head_dim] fp8x4.

  • wqkv_a_scale_hbm (nl.NkiTensor) – compact block-128 scales for wqkv_a. BF16: ignored, pass None.

  • wq_b_hbm (nl.NkiTensor) – [qk_lora_rank // 4, n_heads * (qk_nope_head_dim + qk_rope_head_dim)] fp8x4.

  • wq_b_scale_hbm (nl.NkiTensor) – compact block-128 scales for wq_b. BF16: ignored, pass None.

  • q_norm_gamma_hbm (nl.NkiTensor) – [1, qk_lora_rank] bf16 RMSNorm gamma for the Q intermediate.

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

  • wuk_hbm (nl.NkiTensor) – [qk_nope_head_dim, n_heads * kv_lora_rank] bf16 absorption weight (W_uk[h] = [nope, kv_lora], contraction = nope; head h owns columns [h * kv_lora_rank, (h + 1) * kv_lora_rank)).

  • cos_cache_hbm (nl.NkiTensor) – [B, S, qk_rope_head_dim] bf16 RoPE cosine cache.

  • sin_cache_hbm (nl.NkiTensor) – [B, S, qk_rope_head_dim] bf16 RoPE sine cache.

  • n_heads (int) – Number of attention heads.

  • qk_nope_head_dim (int) – Non-RoPE per-head Q/K dimension (must be 128).

  • qk_rope_head_dim (int) – RoPE per-head Q/K dimension.

  • kv_lora_rank (int) – Latent KV LoRA rank.

  • qk_lora_rank (int) – Q LoRA rank.

  • norm_eps (float) – RMSNorm epsilon (default 1e-6).

  • compact_scales (bool) – MX only – see wqkv_a_scale_hbm.

  • precision (MlaPrecision) – MX (default) or BF16. BF16 takes unquantized bf16 activations and weights in natural column order and issues plain nisa.nc_matmul: no 4-packing, no MLA output-column swizzle, no scale broadcast, no quantize_mx, no un-swizzle. It also lifts the MX path’s qk_nope_head_dim == 128 restriction (the absorption contraction is tiled at 128 rows, so GLM’s 192 runs as 128 + 64) and relaxes the LoRA-rank multiple from 512 to 128. Pass None for both scale tensors.

Returns:

[B, S, n_heads, kv_lora_rank] bf16, per-head absorbed Q latent.

Returns:

[B, S, n_heads, qk_rope_head_dim] bf16, per-head RoPE queries.

Returns:

[B, S, kv_lora_rank] bf16, shared latent KV (RMSNorm(kv) * gamma).

Returns:

[B, S, qk_rope_head_dim] bf16, shared RoPE key. Then, for the SAI indexer’s wq_b projection, the q-normed qr latent – MX emits it pre-transposed and pre-quantized so the indexer skips that work, BF16 emits it plain:

Returns:

qr_qtz [num_s_tiles, P_MAX, qk_lora_rank // 512, P_MAX] uint32 (fp8x4) plus qr_scale, the paired block-32 MX scales, same shape uint8.

Returns:

qr [B, S, qk_lora_rank] bf16.

Return type:

Tuple[nl.NkiTensor, ...]

Notes:

  • Absorption requires qk_nope_head_dim == 128 (full-partition bf16 contraction).

  • The absorption matmul (q_nope @ W_uk) runs in bf16, not MX.

Dimensions:

  • B: Batch size

  • S: Sequence length (tokens)

  • H: Hidden dimension size

  • n_heads: Number of attention heads

  • qk_nope_head_dim: Non-RoPE portion of the per-head Q/K dimension (must be 128)

  • qk_rope_head_dim: RoPE portion of the per-head Q/K dimension (R)

  • kv_lora_rank: Latent KV LoRA rank (L)

This document is relevant for: Trn2, Trn3