flashinfer.prefill.trtllm_batch_context_with_kv_cache

flashinfer.prefill.trtllm_batch_context_with_kv_cache(query: Tensor, kv_cache: Tensor | Tuple[Tensor, Tensor], workspace_buffer: Tensor, block_tables: Tensor, seq_lens: Tensor, max_q_len: int, max_kv_len: int, bmm1_scale: float | Tensor, bmm2_scale: float | Tensor, batch_size: int, cum_seq_lens_q: Tensor, cum_seq_lens_kv: Tensor, window_left: int = -1, out: Tensor | FP4Tensor | None = None, out_dtype: str | dtype | None = None, o_sf_scale: float | None = None, o_sf_vec_size: int | None = None, kv_layout: str = 'HND', enable_pdl: bool | None = None, sinks: List[Tensor] | None = None, kv_cache_sf: Tensor | Tuple[Tensor, Tensor] | None = None, skip_softmax_threshold_scale_factor: float | None = None, uses_shared_paged_kv_idx: bool = True, causal: bool = True, lse: Tensor | None = None, return_lse: bool = False, multi_ctas_kv_counter_buffer: Tensor | None = None, use_fp16_softmax: bool | None = None, uses_spcompress: bool | None = None, backend: str = 'trtllm-gen', window_right: int = -1, variable_window_token_starts: Tensor | None = None, variable_window_token_ends: Tensor | None = None) → Tensor | FP4Tensor | Tuple[Tensor | FP4Tensor, Tensor]

Parameters:

• query (torch.Tensor) – query tensor with shape [num_tokens, num_heads, head_dim]

• kv_cache (Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]) –

  If kv_cache is a single tensor, it should be a tensor with shape [num_pages, 1 or 2, num_kv_heads, page_size, head_dim] if kv_layout is “HND”, or [num_pages, 1 or 2, page_size, num_kv_heads, head_dim] if kv_layout is “NHD”. If kv_cache is a tuple of two tensors, it should be a tuple of two tensors with shape [num_pages, num_kv_heads, page_size, head_dim] if kv_layout is “HND”, or [num_pages, page_size, num_kv_heads, head_dim] if kv_layout is “NHD”. The first tensor is the key cache, the second tensor is the value cache.

  Contiguity requirements (trtllm-gen backend):

  • The head_dim (last dim) must have stride 1. This is a TMA hardware constraint

  • The head and batch/page dims can have arbitrary strides.

• workspace_buffer (torch.Tensor) – Workspace used for trtllm-gen softmax stats/scratch.

• block_tables (torch.Tensor) – Page table of kv cache. When uses_shared_paged_kv_idx is True (default): shape [batch_size, max_num_pages_per_seq]. When uses_shared_paged_kv_idx is False: shape [batch_size, 2, max_num_pages_per_seq] where dim 1 distinguishes K (0) and V (1) page indices.

• seq_lens (torch.Tensor) – A uint32 1D tensor indicating the kv sequence length of each prompt. shape: [batch_size]

• max_q_len (int) – max sequence length for query

• max_kv_len (int) – max sequence length for kv_cache

• bmm1_scale (Union[float, torch.Tensor]) – Fused scale for bmm1 input. When using the trtllm-gen backend, it can be a torch.Tensor with dtype torch.float32.

• bmm2_scale (Union[float, torch.Tensor]) – fused scale for bmm2 input. when using trtllm-gen backend, it can be a torch.Tensor with dtype torch.float32.

• batch_size (int) – batch size

• cum_seq_lens_q (torch.Tensor) – cumulative sequence length for query. shape: [batch_size + 1]

• cum_seq_lens_kv (torch.Tensor) – cumulative sequence length for kv_cache. shape: [batch_size + 1]

• window_left (int = -1) – The left (inclusive) window size for the attention window. A value of N includes the current position and keys up to N positions to its left. -1 means unbounded. Defaults to -1.

• window_right (int = -1) – The right (inclusive) window size for the attention window. A value of N includes the current position and keys up to N positions to its right. For backward compatibility, causal attention with window_left >= 0 treats window_right=-1 as 0. A non-causal fixed window requires a non-negative value. Defaults to -1.

• variable_window_token_starts (Optional[torch.Tensor]) – Contiguous int32 tensor with one entry per packed query token. Together with variable_window_token_ends, supplies the inclusive K/V interval visible to every query token. Bounds use positions within each batch item’s full K/V sequence and must be nondecreasing within a sequence. These value invariants are the caller’s responsibility; they are checked with a device-to-host synchronization only when FLASHINFER_VALIDATE_INPUTS=1 and CUDA graph capture is not active.

• variable_window_token_ends (Optional[torch.Tensor]) – Inclusive interval ends paired with variable_window_token_starts. Variable windows cannot be combined with fixed window_left/window_right bounds; when supplied, they define the complete attention mask and the causal flag is ignored.

• out (Optional[Union[torch.Tensor, FP4Tensor]] = None) – output tensor, if not provided, will be allocated with out_dtype, if out_dtype is not provided, will use the type of query.

• out_dtype (Optional[Union[torch.dtype, str]] = None) – output dtype, if not provided, will use the type of out. For nvfp4, use string nvfp4.

• o_sf_scale (Optional[float] = None) – scale for nvfp4 output tensor scale factor.

• o_sf_vec_size (Optional[int] = None) – vector size for nvfp4 output tensor scale factor.

• enable_pdl (Optional[bool] = None) – Whether to enable Programmatic Dependent Launch (PDL). See https://docs.nvidia.com/cuda/cuda-programming-guide/04-special-topics/programmatic-dependent-launch.html Defaults to None, which means it will be enabled if the device supports PDL.

• kv_layout (str = "HND") – Layout of kv-cache, can be “HND” or “NHD”, default is “HND”. For the trtllm-gen backend with NVFP4 KV cache, using NHD will trigger an automatic transpose and .contiguous() copy of both the KV data and block scale tensors to convert them to HND layout. This incurs extra memory allocation and data copy overhead. Use HND for better performance.

• sinks (Optional[List[torch.Tensor]] = None) – additional value per head in the denominator of the softmax.

• kv_cache_sf (Optional[Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]] = None) –

  Per-block scale factors for NVFP4 KV cache. Accepts the same formats as kv_cache:

  • a tuple (k_scales, v_scales) of 4-D tensors, each following kv_layout: [num_pages, num_kv_heads, page_size, head_dim // 16] for HND, or [num_pages, page_size, num_kv_heads, head_dim // 16] for NHD.

  • a single 5-D tensor with shape [num_pages, 2, ...] matching the layout of kv_cache, split on dim 1 to yield k (index 0) and v (index 1) scales.

  Both tensors have dtype torch.float8_e4m3fn. k_scales uses a linear (row-major) layout, while v_scales must use TRT-LLM’s 4-token interleaved layout within each [page_size, head_dim // 16] tile. Use flashinfer.fp4_quantization.nvfp4_quantize_paged_kv_cache() to produce correctly formatted scale factors.

  When kv_layout is NHD, scale tensors are transposed to HND internally (incurring a .contiguous() copy). Use HND for better performance.

  For a head_dim without an NVFP4 KV context kernel, convert the pages to FP8 with flashinfer.quantization.kernels.nvfp4_kv_materialize_fp8() and pass the FP8 cache instead.

  Contiguity requirements (trtllm-gen backend):

  • The last two dims (page_size, head_dim // 16) must be contiguous (i.e., stride[-1] == 1 and stride[-2] == head_dim // 16). This is because the kernel reshapes them into (16, page_size * head_dim / 16 / 16) to satisfy TMA’s 16-byte box width minimum.

  • The head and batch/page dims can have arbitrary strides.

• skip_softmax_threshold_scale_factor (Optional[float] = None) – threshold scale factor for skipping softmax operations. Providing a value for this parameter enables skip-softmax sparsity as described in: https://arxiv.org/abs/2512.12087 If no value is provided, then standard attention is used. Setting the threshold to a higher value generally increases kernel performance at the cost of accuracy degradation. The actual threshold value equals the provided threshold_scale_factor divided by the context length.

• uses_shared_paged_kv_idx (bool = True) – Whether the K and V page indices are shared as a unified index. True (default) uses vLLM/FlashInfer layout with a 2D page table. False uses TRT-LLM layout with a 3D page table [batch_size, 2, max_num_pages_per_seq].

• use_fp16_softmax (Optional[bool] = None) – Select the …Fp16Softmax… cubin variant (FP16 softmax accumulator). Currently only shipped for BF16 Q/KV/O kernels.

• uses_spcompress (Optional[bool] = None) – Select the …Spcomp… cubin variant (sparse compression). Currently only shipped for FP8 Q kernels.

• causal (bool = True) – Whether to apply a causal mask. Set to False to request dense / bidirectional attention. For the TRTLLM-gen paged context path, non-causal currently requires window_left == -1.

• lse (Optional[torch.Tensor] = None) – Optional pre-allocated buffer for Log-Sum-Exp values. Must have shape [num_tokens, num_qo_heads] with dtype torch.float32. If provided, the tensor is filled regardless of return_lse. If return_lse is True and this is None, a buffer will be allocated.

• return_lse (bool = False) – Whether to return Log-Sum-Exp values. When True, returns (out, lse).

• multi_ctas_kv_counter_buffer (Optional[torch.Tensor] = None) – Optional separate trtllm-gen multi-CTA KV counter buffer. If not provided, FlashInfer creates a fresh zeroed internal buffer. When provided, it must be zero-initialized at allocation (e.g. via torch.zeros); the kernel self-resets the counters after each launch, so it does not need to be re-zeroed between calls.

• backend (str = "trtllm-gen") – The implementation backend, either trtllm-gen or cake. cake selects the separately versioned Cake FMHA product. The default remains the conventional TRTLLM implementation.

Returns:

• out (Union[torch.Tensor, FP4Tensor]) – output torch.Tensor or FP4Tensor.

• lse (torch.Tensor, optional) – Only returned when return_lse is True. Shape [num_tokens, num_qo_heads] with dtype torch.float32.