flashinfer.mla.dsv41_fp8_quantize_pack_sparse_mla_cache

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

Quantize complete DeepSeek-V4.1 latent-KV pages to the 528-byte DSV4_1 FP8 ABI.

This is the device writer for the main (SWA) cache of the DeepSeek-V4.1 formats (kv_cache_format="fp8_dsv41" / "fp8_dsv41_fp4_ca" and SparseMLASm120Wrapper(kv_scale_format="ue8m0_g32")); the companion dsv41_fp4_quantize_pack_sparse_mla_cache() writes the 288-byte compressed (extra) cache.

Parameters:

  • latent_kv (torch.Tensor) – Contiguous CUDA BF16/FP16 tensor with shape [num_pages, page_size, 512]. A singleton latent-head axis is also accepted in HND or NHD position. All 512 values per token (RoPE dims included) are quantized in groups of 32 to E4M3 with one UE8M0 (power-of-two) scale per group: scale = 2**ceil(log2(max(amax / 448, 1e-4))) with amax clamped at 1e-4, values divided by the scale and rounded to nearest-even E4M3 saturating at 448 (the FlashMLA FP8 FOOTER trajectory; bit-identical to the torch reference quantizer used by the FlashInfer tests). A NaN in a group writes a 0xFF scale byte and 0x7F codes; an infinite group writes a 0xFF scale byte, 0x7F for the infinite values and signed zero codes for the finite ones.

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

Returns:

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

Return type:

torch.Tensor