flashinfer.page.nvfp4_quantize_append_paged_kv_cache

flashinfer.page.nvfp4_quantize_append_paged_kv_cache(append_key: Tensor, append_value: Tensor, batch_indices: Tensor, positions: Tensor, paged_kv_cache: Tensor | Tuple[Tensor, Tensor], kv_cache_sf: Tensor | Tuple[Tensor, Tensor], kv_indices: Tensor, kv_indptr: Tensor, kv_last_page_len: Tensor, k_scale: float, v_scale: float, kv_layout: str = 'NHD') → None

Quantize and append K/V rows into an NVFP4 paged KV cache.

append_key and append_value must be fp16/bf16 tensors with shape [nnz, num_kv_heads, head_dim]. The function writes packed E2M1 data into uint8 paged K/V cache tensors with last dimension head_dim // 2 and writes FP8 E4M3 block scales into kv_cache_sf tensors with last dimension head_dim // 16.

k_scale and v_scale are the global decode scales consumed by the NVFP4 attention kernels, i.e. dequantization reconstructs values as e2m1_value * block_scale * global_scale.

Parameters:

  • append_key (torch.Tensor) – The key tensor to quantize and append, shape [nnz, num_kv_heads, head_dim] with dtype torch.float16 or torch.bfloat16.

  • append_value (torch.Tensor) – The value tensor to quantize and append, with the same shape and dtype as append_key.

  • batch_indices (torch.Tensor) – The batch index for each appended row, shape [nnz].

  • positions (torch.Tensor) – The logical token position for each appended row, shape [nnz].

  • paged_kv_cache (Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]) – Caller-owned packed NVFP4 K/V cache. For tuple input, each tensor has shape [max_num_pages, page_size, num_kv_heads, head_dim // 2] when kv_layout="NHD" and [max_num_pages, num_kv_heads, page_size, head_dim // 2] when kv_layout="HND". A stacked 5-D cache is also accepted with K/V on the second dimension.

  • kv_cache_sf (Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]) – Caller-owned FP8 E4M3 scale cache with the same tuple or stacked cache format as paged_kv_cache, replacing head_dim // 2 with head_dim // 16.

  • 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, shape [batch_size].

  • k_scale (float) – Positive finite global decode scale for K.

  • v_scale (float) – Positive finite global decode scale for V.

  • kv_layout (str) – Layout of the paged KV cache, either "NHD" or "HND".

Returns:

This function updates paged_kv_cache and kv_cache_sf in place.

Return type:

None

Note

The function assumes that the space for appended K/V rows has already been allocated and described by kv_indices, kv_indptr, and kv_last_page_len.