mla

MLA attention kernels with fused fake quantization (prefill and decode).

Functions

mla_attention_decode

Decode one absorbed query token per request over a paged latent cache.

mla_prefill_attention

Varlen MLA prefill attention with fused Q/K/P/V fake quantization.
mla_attention_decode(q, latent_cache, block_table, b_seq_len, *, softmax_scale, kv_lora_rank=512, qk_rope_head_dim=64, page_size=None, num_kv_splits=32, k_qdq=None, k_qdq_amax=None, v_qdq=None, v_qdq_amax=None, p_qdq=None, p_qdq_amax=1.0, return_lse=True, out_dtype=None)

Decode one absorbed query token per request over a paged latent cache.

Quantized-BMM model: the latent cache holds RAW (bf16/fp16) values, and each BMM operand is fake-quantized once, on read, along its own contraction axis — K along the feature axis (k_qdq), V along the token axis (v_qdq), independently. So decode honors k_format and v_format separately. Q is fake-quantized by the caller (dynamic NVFP4 uses an FP32 carrier); P is quantized in-kernel after the row-sum (denominator stays unquantized).

Parameters:

  • q (Tensor) – [batch, num_heads, kv_lora_rank + qk_rope_head_dim] absorbed query. Pass FP32 for the dynamic-NVFP4 QDQ carrier (Q is expected to be fake-quantized by the caller); BF16/FP16 otherwise.

  • latent_cache (Tensor) – [num_blocks, page_size, kv_lora_rank + qk_rope_head_dim] paged latent cache holding raw (unquantized) values.

  • block_table (Tensor) – [batch, max_blocks_per_seq] page table.

  • b_seq_len (Tensor) – [batch] KV sequence lengths.

  • softmax_scale (float) – Softmax scale (required; MLA layers fold in mscale).

  • kv_lora_rank (int) – Latent width (V/output width).

  • qk_rope_head_dim (int) – RoPE feature width appended to the latent.

  • page_size (int | None) – Tokens per page; defaults to latent_cache.shape[1].

  • num_kv_splits (int) – Fixed split count. P/V QDQ numerics follow the split-local schedule; kept fixed by default for reproducibility across batch shapes and devices.

  • k_qdq (str | None) – K fake quant-dequant: None, "fp8", "nvfp4" (feature axis, block-16 for NVFP4). Covers both the NOPE and RoPE slices.

  • k_qdq_amax (float | None) – Per-tensor K amax (None = scale 1.0).

  • v_qdq (str | None) – V fake quant-dequant, same modes; token/contraction axis. The open 16-token tail block re-quantizes as the sequence grows.

  • v_qdq_amax (float | None) – Per-tensor V amax (None = scale 1.0).

  • p_qdq (str | None) – Softmax-P fake quant-dequant: None, "fp8", "nvfp4".

  • p_qdq_amax (float) – Per-tensor P amax (default 1.0, the theoretical bound).

  • return_lse (bool) – Also return the natural-log LSE [batch, num_heads].

  • out_dtype (dtype | None) – Output dtype (default latent_cache.dtype, the model compute dtype expected by the V up-projection).

Returns:

(out [batch, num_heads, kv_lora_rank], lse [batch, num_heads] | None).

Return type:

tuple[Tensor, Tensor | None]

mla_prefill_attention(q, k, v, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, softmax_scale=None, causal=True, return_lse=False, *, q_quant=None, k_quant=None, p_quant=None, v_quant=None, q_amax=None, k_amax=None, p_amax=1.0, v_amax=None, sparsity_n=0, sparsity_m=4, dense_sink_tokens=0, dense_recent_tokens=64, block_m=64, block_n=64)

Varlen MLA prefill attention with fused Q/K/P/V fake quantization.

Parameters:

  • q (Tensor) – [total_q, num_heads, qk_head_dim] packed queries.

  • k (Tensor) – [total_k, num_kv_heads, qk_head_dim] packed keys (NOPE ++ RoPE).

  • v (Tensor) – [total_k, num_kv_heads, v_head_dim] packed values.

  • cu_seqlens_q (Tensor) – [batch + 1] cumulative Q sequence lengths.

  • cu_seqlens_k (Tensor) – [batch + 1] cumulative K/V sequence lengths.

  • max_seqlen_q (int) – Maximum Q sequence length (grid sizing).

  • softmax_scale (float | None) – Scale factor (default qk_head_dim ** -0.5).

  • causal (bool) – Causal masking; Q is treated as the suffix of the KV span. Cached-context chunks use causal=False.

  • return_lse (bool) – Also return the natural-log LSE [num_heads, total_q] for chunk-state merging (merge_attn_states).

  • q_quant (str | None) – Q fake quant-dequant: None, "fp8" (per-tensor E4M3), or "nvfp4" (1x16 blocks along the feature/contraction axis).

  • k_quant (str | None) – K fake quant-dequant, same modes; blocks along features.

  • p_quant (str | None) – Softmax-P fake quant-dequant, blocks along the key axis. The softmax denominator stays unquantized.

  • v_quant (str | None) – V fake quant-dequant, blocks along the key/token axis.

  • q_amax (float | None) – Per-tensor amax for Q (None = scale 1.0).

  • k_amax (float | None) – Per-tensor amax for K (None = scale 1.0).

  • p_amax (float) – Per-tensor amax for P; defaults to 1.0, the theoretical upper bound of the unnormalized P’s amax.

  • v_amax (float | None) – Per-tensor amax for V (None = scale 1.0).

  • sparsity_n (int) – N:M score sparsity along the key axis (0 = off).

  • sparsity_m (int) – N:M group size (4 or 8).

  • dense_sink_tokens (int) – Leading KV tokens kept dense (token-exact).

  • dense_recent_tokens (int) – Recent KV tokens kept dense (token-exact).

  • block_m (int) – Q tile size (multiple of 16).

  • block_n (int) – KV tile size (multiple of 16).

Returns:

Output [total_q, num_heads, v_head_dim] in v.dtype; with return_lse a tuple (out, lse).

Return type:

Tensor | tuple[Tensor, Tensor]