This document is relevant for: Trn2, Trn3

QKV CTE MLA Kernel API Reference#

DeepSeek MLA QKV projection with MX quantization for Context Encoding.

Implements the full QKV projection pipeline for Multi-head Latent Attention (MLA) with MX (fp8) quantization. Includes two-stage low-rank projections for both Q and KV paths, fused RMSNorm, and Rotary Position Embedding (RoPE). Supports LNC sharding on the sequence dimension. Sequence length is tiled at 128, and batch size is 1 (CTE processes a single context).

Background#

The qkv_mla_mx kernel implements the full QKV projection pipeline for DeepSeek Multi-head Latent Attention (MLA) with MX (fp8) quantization during Context Encoding. It performs two-stage low-rank projections for both the Q and KV paths, fused RMSNorm, and RoPE. The qkv_mla_mx_deepseek_v4 entry point provides the DeepSeek v4 variant of the same projection.

API Reference#

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

qkv_mla_mx#

nkilib.experimental.qkv.qkv_mla_mx(x_hbm: nl.ndarray, wqkv_a_hbm: nl.ndarray, wqkv_a_scale_hbm: nl.ndarray, wq_b_hbm: nl.ndarray, wq_b_scale_hbm: nl.ndarray, q_norm_gamma_hbm: nl.ndarray, wkv_b_hbm: nl.ndarray, wkv_b_scale_hbm: nl.ndarray, kv_norm_gamma_hbm: nl.ndarray, cos_cache_hbm: nl.ndarray, sin_cache_hbm: nl.ndarray, n_heads: int, qk_nope_head_dim: int, qk_rope_head_dim: int, v_head_dim: int, kv_lora_rank: int, qk_lora_rank: int, norm_eps: float = 1e-06) → Tuple[nl.ndarray, nl.ndarray, nl.ndarray]#

DeepSeek MLA QKV projection with MX quantization for Context Encoding.

Parameters:

  • x_hbm (nl.ndarray) – [B, S, H] bf16, Input hidden states

  • wqkv_a_hbm (nl.ndarray) – [H//4, qk_lora_rank + kv_lora_rank + qk_rope_head_dim] fp8x4, First combined Q/KV projection weights

  • wqkv_a_scale_hbm (nl.ndarray) – [H//128, ceil((qk_lora_rank + kv_lora_rank + qk_rope_head_dim)/128)] uint8, DeepSeek block-128 compact scales for wqkv_a

  • wq_b_hbm (nl.ndarray) – [qk_lora_rank//4, n_heads * qk_head_dim] fp8x4, Second Q projection weights

  • wq_b_scale_hbm (nl.ndarray) – [qk_lora_rank//128, ceil(n_heads * qk_head_dim / 128)] uint8, DeepSeek block-128 compact scales for wq_b

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

  • wkv_b_hbm (nl.ndarray) – [kv_lora_rank//4, n_heads * (qk_nope_head_dim + v_head_dim)] fp8x4, Second KV projection weights

  • wkv_b_scale_hbm (nl.ndarray) – [kv_lora_rank//128, ceil(n_heads * (qk_nope_head_dim + v_head_dim) / 128)] uint8, MX scales for wkv_b

  • kv_norm_gamma_hbm (nl.ndarray) – [1, kv_lora_rank] bf16, RMSNorm gamma for KV intermediate

  • cos_cache_hbm (nl.ndarray) – [B, S, qk_rope_head_dim] bf16, Cosine RoPE frequencies

  • sin_cache_hbm (nl.ndarray) – [B, S, qk_rope_head_dim] bf16, Sine RoPE frequencies

  • n_heads (int) – Number of attention heads

  • qk_nope_head_dim (int) – Non-RoPE portion of Q/K head dimension

  • qk_rope_head_dim (int) – RoPE portion of Q/K head dimension

  • v_head_dim (int) – Value head dimension

  • kv_lora_rank (int) – Latent dimension for KV compression

  • qk_lora_rank (int) – Latent dimension for Q compression

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

Returns:

[B, S, n_heads, qk_head_dim] bf16, Query projections with RoPE applied

Return type:

nl.ndarray

Returns:

[B, S, n_heads, qk_head_dim] bf16, Key projections with RoPE applied

Return type:

nl.ndarray

Returns:

[B, S, n_heads, v_head_dim] bf16, Value projections

Return type:

nl.ndarray

Notes:

  • Matmul shapes: Combined Q/KV Path Stage 1: x[B,S,H] @ wqkv_a[H, qk_lora_rank + kv_lora_rank + qk_rope_head_dim] -> qkv_a_out[B,S,qk_lora_rank + kv_lora_rank + qk_rope_head_dim] Split: qr[B,S,qk_lora_rank], kv[B,S,kv_lora_rank], k_pe[B,S,qk_rope_head_dim] Q Path Stage 2: norm(qr)[B,S,qk_lora_rank] @ wq_b[qk_lora_rank, n_heads*qk_head_dim] -> q[B,S,n_heads*qk_head_dim] KV Path Stage 2: norm(kv)[B,S,kv_lora_rank] @ wkv_b[kv_lora_rank, n_heads*(qk_nope_head_dim+v_head_dim)] -> kv_out[B,S,n_heads*(qk_nope_head_dim+v_head_dim)] Split: k_nope[B,S,n_heads,qk_nope_head_dim], v[B,S,n_heads,v_head_dim] Final assembly: q_pe = q[…, qk_nope_head_dim:] -> RoPE -> q[…, qk_nope_head_dim:] k_pe -> RoPE -> broadcast to all heads -> concat with k_nope -> K

Dimensions:

  • B: Batch size

  • S: Sequence length

  • H: Hidden dimension (input)

  • n_heads: Number of attention heads

  • qk_lora_rank: Q latent dimension (e.g., 1536)

  • kv_lora_rank: KV latent dimension (e.g., 512)

  • qk_head_dim: qk_nope_head_dim + qk_rope_head_dim (e.g., 128 + 64 = 192)

  • qk_nope_head_dim: Non-rotary part of Q/K (e.g., 128)

  • qk_rope_head_dim: Rotary part of Q/K (e.g., 64)

  • v_head_dim: Value head dimension (e.g., 128)

qkv_mla_mx_deepseek_v4#

nkilib.experimental.qkv.qkv_mla_mx_deepseek_v4(x_hbm: nl.ndarray, wqkv_hbm: nl.ndarray, wqkv_scale_hbm: nl.ndarray, wq_b_hbm: nl.ndarray, wq_b_scale_hbm: nl.ndarray, q_norm_gamma_hbm: nl.ndarray, kv_norm_gamma_hbm: nl.ndarray, cos_cache_hbm: nl.ndarray, sin_cache_hbm: nl.ndarray, n_heads: int, head_dim: int, qk_rope_head_dim: int, kv_lora_rank: int, qk_lora_rank: int, norm_eps: float = 1e-06) → Tuple[nl.ndarray, nl.ndarray]#

DeepSeek v4 MLA QKV projection with MX quantization.

Parameters:

  • x_hbm (nl.ndarray) – [B, S, H] bf16 input hidden states

  • wqkv_hbm (nl.ndarray) – [H//4, qk_lora_rank + kv_dim] fp8x4 fused Q/KV first projection

  • wqkv_scale_hbm (nl.ndarray) – [H//128, ceil((qk_lora_rank + kv_dim)/128)] uint8 DeepSeek block-128 compact scales

  • wq_b_hbm (nl.ndarray) – [qk_lora_rank//4, n_heads * head_dim] fp8x4 second Q projection

  • wq_b_scale_hbm (nl.ndarray) – [qk_lora_rank//128, ceil(n_heads * head_dim / 128)] uint8 DeepSeek block-128 compact scales

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

  • kv_norm_gamma_hbm (nl.ndarray) – [1, kv_dim] bf16 RMSNorm gamma for KV

  • cos_cache_hbm (nl.ndarray) – [B, S, qk_rope_head_dim] bf16

  • sin_cache_hbm (nl.ndarray) – [B, S, qk_rope_head_dim] bf16

  • n_heads (int) – number of attention heads

  • head_dim (int) – full Q/K head dimension (nope + rope)

  • qk_rope_head_dim (int) – RoPE dimension

  • kv_lora_rank (int) – KV latent dimension

  • qk_lora_rank (int) – Q latent dimension

  • norm_eps (float) – RMSNorm epsilon

Returns:

[B, S, n_heads, head_dim] bf16 with RoPE on last rope_dim

Return type:

nl.ndarray

Returns:

[B, S, kv_dim] bf16 with RoPE on last rope_dim

Return type:

nl.ndarray

Notes:

  • Matmul shapes: Fused first projection: x[B,S,H] @ wqkv[H, qk_lora_rank + kv_dim] -> qkv_a_out[B,S,qk_lora_rank + kv_dim] Split: qr[B,S,qk_lora_rank], kv[B,S,kv_dim] Q Path Stage 2: norm(qr)[B,S,qk_lora_rank] @ wq_b[qk_lora_rank, n_heads*head_dim] -> q[B,S,n_heads*head_dim] Final assembly: q -> per-head rsqrt norm -> RoPE on q[…, -rope_dim:] kv -> RoPE on kv[…, -rope_dim:] (kv latent returned directly)

Dimensions:

  • B: Batch size

  • S: Sequence length

  • H: Hidden dimension (input)

  • n_heads: Number of attention heads

  • qk_lora_rank: Q latent dimension

  • kv_lora_rank: KV latent dimension

  • kv_dim: kv_lora_rank + qk_rope_head_dim (full kv output width)

  • head_dim: full Q/K head dimension (nope + rope)

  • qk_rope_head_dim: RoPE portion of Q/K head dimension

This document is relevant for: Trn2, Trn3