flashinfer.mla_kv_pack.concat_mla_kv_quant_fp8

flashinfer.mla_kv_pack.concat_mla_kv_quant_fp8(kv_nope: Tensor, k_pe: Tensor, key: Tensor | None = None, value: Tensor | None = None, *, nope_dim: int = 128, backend: str = 'auto') → Tuple[Tensor, Tensor]

Fused MLA context K/V pack with fp8 e4m3 quantization.

Turns the per-head kv_b_proj output kv_nope = [k_nope | v] and the shared, already-RoPE’d k_pe into the fp8 key = [k_nope | k_pe] / value = v buffers a ragged MLA prefill kernel consumes (e.g. trtllm_ragged_attention_deepseek() with separate Q/K/V), in one memory pass instead of two casts and two strided copies.

Numerics: the cast is the saturating e4m3 conversion (round to nearest even; finite overflow and +-inf -> +-448; NaN -> 0x7F), which is what Tensor.to(torch.float8_e4m3fn) produces on the GPU since torch 2.13. Output is bit-exact with that cast plus slice copies, on every backend.

Two fused backends serve bf16 inputs with head geometry nope_dim=128, rope_dim=64, v_dim=128 within the head-count and token-count surface of the package’s mla_kv_pack_fp8_workloads.json (1..128 heads, 1..131072 tokens) on compute capability 10.0 and newer:

  • "specialized": the hand-written CUDA kernel (one warp per token, a fully unrolled variant for 12 local heads, runtime head count otherwise), compiled for every compute capability >= 10.0. Needs contiguous, 16-byte-aligned tensors.

  • "cake": generated Cake programs (one per head group, i.e. the number of head pairs a warp owns, planned from the token and head counts) built for the exact compute capabilities 10.0 and 10.3. Needs 32-byte-aligned tensors; k_pe may be a row-strided column slice of a wider row-major workspace (unit last stride, row stride a multiple of 16 elements).

  • "auto" (default): "specialized" for 12 local heads (96 heads over TP8, where the two backends measure within +-1 % of each other on GB300/B200 at every token count), "cake" for every other head count (1.03-1.4x faster there); if the preferred backend cannot serve the call the other one is tried.

Every call no fused backend serves takes the composable torch path. FLASHINFER_SPECIALIZED_KERNEL_DISABLE=1 forces the composable path (read at call time).

Parameters:

  • kv_nope (torch.Tensor) – [num_tokens, num_heads, nope_dim + v_dim] bf16 (contiguous for the fused paths).

  • k_pe (torch.Tensor) – [num_tokens, rope_dim] or [num_tokens, 1, rope_dim] bf16, shared across heads; contiguous, or for the Cake backend a column slice of a row-major workspace whose row stride is a multiple of 16 elements (e.g. latent[:, 512:] of a [num_tokens, 576] buffer).

  • key (Optional[torch.Tensor]) – [num_tokens, num_heads, nope_dim + rope_dim] float8_e4m3fn output; allocated when None.

  • value (Optional[torch.Tensor]) – [num_tokens, num_heads, v_dim] float8_e4m3fn output; allocated when None.

  • nope_dim (int) – Split point of kv_nope’s last dim (default 128).

  • backend (str) – "auto" (default), "specialized" or "cake". An explicit fused backend that cannot serve the call takes the composable path.

Returns:

(key, value).

Return type:

Tuple[torch.Tensor, torch.Tensor]