flashinfer.quantization.kernels.nvfp4_kv_materialize.nvfp4_kv_materialize_fp8

flashinfer.quantization.kernels.nvfp4_kv_materialize.nvfp4_kv_materialize_fp8(k_cache: Tensor, v_cache: Tensor, k_block_scales: Tensor, v_block_scales: Tensor, block_tables: Tensor, seq_lens: Tensor, max_kv_len: int, uses_shared_paged_kv_idx: bool = True, out: Tensor | None = None, kv_page_indptr: Tensor | None = None, num_pages: int | None = None) → Tuple[Tensor, Tensor]

Materialize the pages referenced by a paged NVFP4 KV cache as FP8.

Requires SM100, SM103, SM107, SM110, SM120 or SM121; query with nvfp4_kv_materialize_fp8.is_compute_capability_supported(cc).

Parameters:

  • k_cache (torch.Tensor) – Packed NVFP4 K/V in HND layout, [num_pages, num_kv_heads, page_size, head_dim // 2], dtype torch.uint8. The page stride may be larger than a compact page.

  • v_cache (torch.Tensor) – Packed NVFP4 K/V in HND layout, [num_pages, num_kv_heads, page_size, head_dim // 2], dtype torch.uint8. The page stride may be larger than a compact page.

  • k_block_scales (torch.Tensor) – E4M3 block scales stored as torch.uint8, [num_pages, num_kv_heads, page_size, head_dim // 16]. K scales are linear; V scales use trtllm-gen’s 4-token-interleaved layout, as produced by flashinfer.fp4_quantization.nvfp4_quantize_paged_kv_cache().

  • v_block_scales (torch.Tensor) – E4M3 block scales stored as torch.uint8, [num_pages, num_kv_heads, page_size, head_dim // 16]. K scales are linear; V scales use trtllm-gen’s 4-token-interleaved layout, as produced by flashinfer.fp4_quantization.nvfp4_quantize_paged_kv_cache().

  • block_tables (torch.Tensor) – int32 page table, [batch_size, max_num_pages] when uses_shared_paged_kv_idx else [batch_size, 2, max_num_pages]. Entries for the first ceil(seq_lens[b] / page_size) columns must be valid page indices; later entries are ignored.

  • seq_lens (torch.Tensor) – int32 KV lengths, [batch_size].

  • max_kv_len (int) – Upper bound on seq_lens. Determines the workspace capacity.

  • uses_shared_paged_kv_idx (bool) – Whether K and V share page indices.

  • out (Optional[torch.Tensor]) – Optional contiguous float8_e4m3fn workspace with the returned shape, reused across calls.

  • kv_page_indptr (Optional[torch.Tensor]) – Optional int32 [batch_size + 1] CUDA prefix sum of the pages each request reserves in the FP8 copy, starting at 0. Request b must get at least ceil(seq_lens[b] / page_size) pages. With mixed lengths this sizes the copy by the active pages instead of batch_size * max_pages_per_seq. Requires num_pages.

  • num_pages (Optional[int]) – Host-side kv_page_indptr[-1]: the number of pages in the FP8 copy.

Returns:

  • kv_cache_fp8 (torch.Tensor) – [num_pages + 1, 2, num_kv_heads, page_size, head_dim]. Without kv_page_indptr, num_pages = batch_size * max_pages_per_seq with max_pages_per_seq = ceil(max_kv_len / page_size). Request b’s page j is at kv_page_indptr[b] + j + 1 (b * max_pages_per_seq + j + 1 by default); page 0 is zero. Pages past each request’s length are not written.

  • fp8_block_tables (torch.Tensor) – int32 [batch_size, max_pages_per_seq] page table into kv_cache_fp8 (padded entries are 0), for use with uses_shared_paged_kv_idx=True.

Note

Only E2M1 values and E4M3 block scales are multiplied. Apply the tensor-level K/V scales in the downstream FP8 attention kernel.