flashinfer.page.nvfp4_quantize_append_paged_kv_cache_with_slot_mapping

flashinfer.page.nvfp4_quantize_append_paged_kv_cache_with_slot_mapping(append_key: Tensor, append_value: Tensor, slot_mapping: Tensor, paged_kv_cache: Tensor | Tuple[Tensor, Tensor], kv_cache_sf: Tensor | Tuple[Tensor, Tensor], k_scale: float | Tensor, v_scale: float | Tensor, kv_layout: str = 'NHD') → None

Quantize and write K/V rows into an NVFP4 paged KV cache by slot mapping.

This variant is intended for runtimes that already assign each token to a flat cache slot. slot_mapping[i] is interpreted as page_id * page_size + entry_idx. Negative slots are padding and are ignored. append_key and append_value may contain additional padded rows; only slot_mapping.shape[0] rows are considered.

Parameters:

  • append_key (torch.Tensor) – The key tensor to quantize and write, shape [num_rows, num_kv_heads, head_dim] with dtype torch.float16 or torch.bfloat16. num_rows must be at least slot_mapping.shape[0].

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

  • slot_mapping (torch.Tensor) – Flat cache slot for each row to write, shape [nnz] with dtype torch.int32 or torch.int64. Negative entries are ignored.

  • 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.

  • k_scale (Union[float, torch.Tensor]) – Positive finite global decode scale for K. During CUDA graph capture, this must be a contiguous scalar torch.float32 CUDA tensor on the same device as append_key.

  • v_scale (Union[float, torch.Tensor]) – Positive finite global decode scale for V. During CUDA graph capture, this must be a contiguous scalar torch.float32 CUDA tensor on the same device as append_key.

  • 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