flashinfer.decode.trtllm_batch_decode_with_kv_cache

flashinfer.decode.trtllm_batch_decode_with_kv_cache(query: Tensor, kv_cache: Tensor | Tuple[Tensor, Tensor], workspace_buffer: Tensor, block_tables: Tensor, seq_lens: Tensor, max_seq_len: int, bmm1_scale: float | Tensor = 1.0, bmm2_scale: float | Tensor = 1.0, 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, sinks: List[Tensor] | None = None, kv_layout: str = 'HND', enable_pdl: bool | None = None, backend: str = 'auto', q_len_per_req: int | None = 1, o_scale: float | None = 1.0, mask: Tensor | None = None, max_q_len: int | None = None, cum_seq_lens_q: Tensor | None = None, skip_softmax_threshold_scale_factor: float | None = None, kv_cache_sf: Tuple[Tensor, Tensor] | None = None, uses_shared_paged_kv_idx: bool = True, lse: Tensor | None = None, return_lse: bool = False, bmm1_scale_log2: Tensor | None = None, multi_ctas_kv_counter_buffer: Tensor | None = None, enable_block_sparse_attention: bool = False, cp_world: int = 1, cp_rank: int = 0, causal_seqlens_kv_global: Tensor | None = None, bf16q_fp8kv_transform_mode: Literal['k_only', 'separate_kv'] | None = None) → Tensor | FP4Tensor | Tuple[Tensor | FP4Tensor, Tensor]
Parameters:

  • query (torch.Tensor) – query tensor with shape [num_tokens, num_heads, head_dim], num_tokens = total query tokens in the batch.

  • 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, and 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 or xqa 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. When enable_block_sparse_attention is True: shape [num_kv_heads, batch_size, max_num_pages_per_seq] (contiguous), where each row holds the page indices selected for that (kv head, sequence) pair, packed densely at the front of the row in ascending original order; the remaining entries are ignored.

  • seq_lens (torch.Tensor) – A uint32 1D tensor indicating the kv sequence length of each prompt. shape: [batch_size]. When enable_block_sparse_attention is True: shape [num_kv_heads, batch_size] (contiguous), holding the number of surviving kv tokens per (kv head, sequence) pair after dropping the non-selected pages. Since selected pages are packed in ascending order, only the final selected page of a row may be partially filled.

  • max_seq_len (int) – max sequence length for kv_cache

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

  • bmm1_scale_log2 (Optional[torch.Tensor] = None) – Optional precomputed log2-form bmm1 scale for trtllm-gen. When provided, this single-element FP32 device tensor is passed directly to the FFI and takes precedence over tensor bmm1_scale, avoiding the internal bmm1_scale * log2e CUDA kernel.

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

  • window_left (int = -1) – The left (inclusive) window size for the attention window, when set to -1, the window size will be set to the full length of the sequence. Defaults to -1.

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

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

  • kv_layout (str = "HND") – The layout of the input k/v tensors, could be either NHD or HND. Defaults to 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.

  • enable_pdl (Optional[bool] = None) – Whether to enable Programmatic Dependent Launch (PDL). See https://docs.nvidia.com/cuda/cuda-c-programming-guide/#programmatic-dependent-launch-and-synchronization When set to None, the backend will be chosen based on the device architecture and kernel availability.

  • backend (str = "auto") – The implementation backend, could be auto/xqa or trtllm-gen. Defaults to auto. When set to auto, the backend will be chosen based on the device architecture and kernel availability. For sm_100 and sm_103 (blackwell architecture), auto will choose trtllm-gen backend. For sm_90 (hopper architecture) and sm_120/sm_121 (blackwell architecture), auto will choose xqa backend.

  • q_len_per_req (Optional[int] = 1) – Number of query tokens per request (i.e. speculative-decoding / MTP depth). query is expected to have batch_size * q_len_per_req rows along its leading dimension when cum_seq_lens_q / max_q_len are None. Defaults to 1.

  • o_scale (Optional[float] = 1.0) – output scale factor for xqa fp8 output.

  • mask (Optional[torch.Tensor] = None) – causal attention mask for xqa speculative decoding.

  • max_q_len (Optional[int] = None) – The maximum query sequence length across all requests when using variable-length queries. Only supported by trtllm-gen backend. Must be provided together with cum_seq_lens_q. When None, all requests use uniform query length specified by q_len_per_req.

  • cum_seq_lens_q (Optional[torch.Tensor] = None) – Cumulative query sequence lengths for variable-length query support, shape: [batch_size + 1], dtype: torch.int32. Only supported by trtllm-gen backend. Must be provided together with max_q_len. When None, all requests use uniform query length specified by q_len_per_req.

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

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

    Per-block scale factors for NVFP4 KV cache, as a tuple of (k_scales, v_scales). Each scale tensor has shape [num_pages, num_kv_heads, page_size, head_dim // 16] in HND layout, with dtype torch.float8_e4m3fn.

    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.

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

  • lse (Optional[torch.Tensor] = None) – Optional pre-allocated buffer for the Log-Sum-Exp (LSE) output, only supported by the trtllm-gen backend. Must have shape [num_tokens, num_qo_heads] and 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 internally.

  • return_lse (bool = False) – Whether to return the Log-Sum-Exp (LSE) values. Only supported by the trtllm-gen backend. When True, the function returns a tuple (out, lse) where lse has shape [num_tokens, num_qo_heads] and dtype torch.float32.

  • multi_ctas_kv_counter_buffer (Optional[torch.Tensor] = None) – Optional separate counter buffer for trtllm-gen multi-CTA KV mode. When omitted, 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 at the end of each launch, so it does not need to be re-zeroed between calls.

  • enable_block_sparse_attention (bool = False) – Whether to use block-sparse attention with different sparse KV pages per KV head. Only supported by the trtllm-gen backend. When True, block_tables and seq_lens are extended with a leading num_kv_heads dimension (see their docs above); each KV head only attends to the pages listed in its own row. max_seq_len should be the maximum surviving kv length across all (kv head, sequence) pairs (the dense maximum is a safe upper bound). Not compatible with sliding window (window_left != -1), skip_softmax_threshold_scale_factor, or uses_shared_paged_kv_idx=False.

  • cp_world (int = 1) – Decode-context-parallel world size. The DCP speculative path is enabled only when causal_seqlens_kv_global is provided. Supported values are 1, 2, 4, and 8 on SM100/SM103.

  • cp_rank (int = 0) – Rank in the DCP group. Rank r owns global KV positions r, r + cp_world, ... in its compact local paged cache.

  • causal_seqlens_kv_global (Optional[torch.Tensor] = None) –

    Optional contiguous int32 tensor of shape [batch_size] containing each request’s global KV prefix length before speculative query row 0. When provided, query row j sees exactly max(0, floor((S + j - cp_rank) / cp_world) + 1) local keys. This enables DCP + speculative decoding without a per-row length tensor.

    The native Cake FMHA path requires BF16 Q/O, causal HND paging, and return_lse=True (or a caller-owned lse). The D128 profile uses head group ratio in [1,8]. BF16 KV uses page size 16 and q_len_per_req in {1,2,4,5,6,8}; FP8 e4m3 KV uses page size 64 and supports every q_len_per_req from 1 through 8. The D256 production profile is FP8/page64 with q_len_per_req in {1,2,3,4,5,6,7,8}, num_qo_heads=16, num_kv_heads=1, and cp_world 1 or 4. bmm1_scale is the fused QK scale and bmm2_scale is the FP8 V/output scale (BF16 KV requires bmm2_scale=1). LSE is FP32 base-2, matching the existing TRT-LLM backend contract. Explicit enable_pdl=True is unsupported; leave it at its default for this path. Here max_seq_len is the maximum compact rank-local stored length; it may be zero for an entirely empty rank, while block_tables still supplies one masked physical page slot. Long-context BF16 and underfilled FP8 Split-KV routes use workspace_buffer for partials and require a zero-initialized, reusable multi_ctas_kv_counter_buffer; the kernel resets those completion tickets after every launch. The D128 FP8 route specializes split1–4 and short-shard K/V retention in its JIT cache key. D256 uses retain0 split1/2/3/4/8/16 bodies and measured B1/B8/B16/B32+ routing. Size its workspace with get_dcp_spec_workspace_size_bytes(..., head_dim=256). Prewarm a fixed tensor/layout binding before CUDA Graph capture.

  • bf16q_fp8kv_transform_mode (Optional[Literal["k_only", "separate_kv"]] = None) – Transform mode for BF16 query + FP8 E4M3 KV decode. None selects the default separate transformed-K/V cubins and is ignored by other paths. "k_only" selects the optimized K-only transform cubins, and "separate_kv" selects the separate transformed-K/V cubins.

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.