flashinfer.page.nvfp4_quantize_append_paged_mla_kv_cache

flashinfer.page.nvfp4_quantize_append_paged_mla_kv_cache(append_ckv: Tensor, append_kpe: Tensor, batch_indices: Tensor, positions: Tensor, ckv_cache: Tensor, ckv_sf_cache: Tensor, kpe_cache: Tensor, kv_indices: Tensor, kv_indptr: Tensor, kv_last_page_len: Tensor, ckv_scale: float, kpe_scale: float) → None

Quantize and append MLA (ckv, kpe) rows into an NVFP4 MLA paged KV cache.

The compressed-kv part is quantized to NVFP4 (packed E2M1 with one FP8 E4M3 scale per 16 channels); the rope part is kept in FP8 E4M3 because rope channels are much hotter than ckv channels in DeepSeek-family latents. Dequantization reconstructs e2m1_value * block_scale * ckv_scale for the compressed-kv part and fp8_value * kpe_scale for the rope part. Inputs are quantized directly from fp16/bf16 with no intermediate FP8 requantization.

Currently only ckv_dim=512 and kpe_dim=64 are supported.

Parameters:

  • append_ckv (torch.Tensor) – The compressed kv tensor to append in ragged tensor format, shape [append_indptr[-1], 512], dtype torch.float16 or torch.bfloat16.

  • append_kpe (torch.Tensor) – The rope part to append in ragged tensor format, shape [append_indptr[-1], 64], same dtype as append_ckv. Rope must already be applied by the caller (pass the unrotated values for NoPE-style models).

  • batch_indices (torch.Tensor) – The batch indices of each entry in the appended rows, shape [append_indptr[-1]].

  • positions (torch.Tensor) – The positions of each entry in the appended rows, shape [append_indptr[-1]].

  • ckv_cache (torch.Tensor) – Packed NVFP4 compressed-kv cache, dtype torch.uint8, shape [num_pages, page_size, 256] (two E2M1 values per byte).

  • ckv_sf_cache (torch.Tensor) – Compressed-kv block scale cache, dtype torch.float8_e4m3fn, shape [num_pages, page_size, 32] (one scale per 16 channels).

  • kpe_cache (torch.Tensor) – FP8 rope cache, dtype torch.float8_e4m3fn, shape [num_pages, page_size, 64].

  • kv_indices (torch.Tensor) – The page indices of the paged kv-cache, shape [kv_indptr[-1]].

  • kv_indptr (torch.Tensor) – The indptr of the paged kv-cache, shape [batch_size + 1].

  • kv_last_page_len (torch.Tensor) – The number of entries in the last page of each request in the paged kv cache, shape [batch_size].

  • ckv_scale (float) – Global decode scale for the compressed-kv part; must lie in the positive normal float32 range (the kernel computes in float32).

  • kpe_scale (float) – Global decode scale for the rope part, same range requirement.

Returns:

This function updates ckv_cache, ckv_sf_cache, and kpe_cache in place.

Return type:

None

See also

append_paged_mla_kv_cache, nvfp4_quantize_append_paged_kv_cache