flashinfer.mla.nvfp4_quantize_pack_sparse_mla_cache

flashinfer.mla.nvfp4_quantize_pack_sparse_mla_cache(latent_kv: Tensor, *, kv_layout: str = 'HND') → Tensor

Quantize complete DeepSeek-V4 latent-KV pages to the NVFP4 cache ABI.

Parameters:

  • latent_kv (torch.Tensor) – Contiguous CUDA BF16 tensor with shape [num_pages, page_size, 512]. A singleton latent-head axis is also accepted in HND or NHD position. The first 448 values are quantized in groups of 16 to packed E2M1 with E4M3 scales; the final 64 BF16 RoPE values are copied bit-for-bit.

  • kv_layout (str) – Output layout, either "HND" or "NHD".

Returns:

Opaque uint8 paged cache with logical shape [num_pages, 1, page_size, 384] for HND or [num_pages, page_size, 1, 384] for NHD. Within each physical page it stores page_size * 352 data bytes followed by page_size * 32 scale bytes. Consumers must not interpret the last dimension as a contiguous per-token record.

Return type:

torch.Tensor