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, qr_qtz_hbm: nl.NkiTensor = None, qr_scale_hbm: nl.NkiTensor = None) → Tuple[nl.NkiTensor, nl.NkiTensor, nl.NkiTensor, 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.

  • 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.

  • 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).

  • qr_qtz_hbm (nl.NkiTensor) – Optional pre-allocated buffer for the exported MX-quantized qr latent (SAI indexer fast path). None (default) means no qr export.

  • qr_scale_hbm (nl.NkiTensor) – Optional pre-allocated scale buffer paired with qr_qtz_hbm.

Returns:

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

Return type:

nl.ndarray

Returns:

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

Return type:

nl.ndarray

Returns:

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

Return type:

nl.ndarray

Returns:

[B, S, qk_rope_head_dim] bf16, shared RoPE key.

Return type:

nl.ndarray

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