.. meta:: :description: DeepSeek MLA QKV projection (MX), emitting absorbed latents. :date-modified: 10/19/2026 .. currentmodule:: nkilib.experimental.mla.deepseek 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 ^^^^^^^^^^^^^^^^^^ .. py:function:: 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. :param x_hbm_mx: ``[B, S, H + scale_region]`` fp8 packed MX activation (see the param comment). :type x_hbm_mx: ``nl.NkiTensor`` :param wqkv_a_hbm: ``[H // 4, qk_lora_rank + kv_lora_rank + qk_rope_head_dim]`` fp8x4. :type wqkv_a_hbm: ``nl.NkiTensor`` :param wqkv_a_scale_hbm: compact block-128 scales for ``wqkv_a``. BF16: ignored, pass ``None``. :type wqkv_a_scale_hbm: ``nl.NkiTensor`` :param wq_b_hbm: ``[qk_lora_rank // 4, n_heads * (qk_nope_head_dim + qk_rope_head_dim)]`` fp8x4. :type wq_b_hbm: ``nl.NkiTensor`` :param wq_b_scale_hbm: compact block-128 scales for ``wq_b``. BF16: ignored, pass ``None``. :type wq_b_scale_hbm: ``nl.NkiTensor`` :param q_norm_gamma_hbm: ``[1, qk_lora_rank]`` bf16 RMSNorm gamma for the Q intermediate. :type q_norm_gamma_hbm: ``nl.NkiTensor`` :param kv_norm_gamma_hbm: ``[1, kv_lora_rank]`` bf16 RMSNorm gamma for the KV latent. :type kv_norm_gamma_hbm: ``nl.NkiTensor`` :param wuk_hbm: ``[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)``). :type wuk_hbm: ``nl.NkiTensor`` :param cos_cache_hbm: ``[B, S, qk_rope_head_dim]`` bf16 RoPE cosine cache. :type cos_cache_hbm: ``nl.NkiTensor`` :param sin_cache_hbm: ``[B, S, qk_rope_head_dim]`` bf16 RoPE sine cache. :type sin_cache_hbm: ``nl.NkiTensor`` :param n_heads: Number of attention heads. :type n_heads: ``int`` :param qk_nope_head_dim: Non-RoPE per-head Q/K dimension (must be 128). :type qk_nope_head_dim: ``int`` :param qk_rope_head_dim: RoPE per-head Q/K dimension. :type qk_rope_head_dim: ``int`` :param kv_lora_rank: Latent KV LoRA rank. :type kv_lora_rank: ``int`` :param qk_lora_rank: Q LoRA rank. :type qk_lora_rank: ``int`` :param norm_eps: RMSNorm epsilon (default 1e-6). :type norm_eps: ``float`` :param compact_scales: MX only -- see ``wqkv_a_scale_hbm``. :type compact_scales: ``bool`` :param precision: ``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. :type precision: ``MlaPrecision`` :return: ``[B, S, n_heads, kv_lora_rank]`` bf16, per-head absorbed Q latent. :return: ``[B, S, n_heads, qk_rope_head_dim]`` bf16, per-head RoPE queries. :return: ``[B, S, kv_lora_rank]`` bf16, shared latent KV (RMSNorm(kv) * gamma). :return: ``[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: :return: ``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. :return: ``qr`` ``[B, S, qk_lora_rank]`` bf16. :rtype: ``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)