flashinfer.attention.prims_ts.prims_ts_batch_decode_with_kv_cache

flashinfer.attention.prims_ts.prims_ts_batch_decode_with_kv_cache(query: Tensor, kv_cache: Tensor | tuple[Tensor, Tensor], workspace_buffer: Tensor, paged_kv_indptr: Tensor, paged_kv_indices: Tensor, seq_lens: Tensor, max_seq_len: int, *, seq_len_q: int = 1, qo_indptr: Tensor | None = None, max_seq_len_q: int | None = None, bmm1_scale: float | None = None, bmm2_scale: float = 1.0, out: Tensor | None = None, out_dtype: dtype | None = None, mask_type: Literal['dense', 'causal'] = 'dense', window_left: int = -1, kv_layout: Literal['HND'] = 'HND') → Tensor

Launch fixed or packed-Q native-CSR FMHA decode with caller scratch.

For seq_len_q=1, query and the returned output both have shape [B, Hq, D]. For seq_len_q>1, both use compact token-major [B, SQ, Hq, D] storage. The kernel writes that layout directly; no layout transpose is performed. When qo_indptr is supplied, Q and O use packed [total_q, Hq, D] storage. Request b owns rows qo_indptr[b]:qo_indptr[b+1]; max_seq_len_q is only the static workspace/JIT bound and is required for this standalone packed interface. To keep this launch path free of device-to-host synchronization, callers must ensure that packed offsets start at zero, are strictly increasing, end at query.shape[0], and have every delta at most max_seq_len_q. For causal masking, every fixed or packed per-request Q length must also be no greater than the corresponding live seq_lens value.

kv_cache is either a combined [pages, 2, Hkv, page_size, D] tensor or a (K, V) tuple of [pages, Hkv, page_size, D] tensors. The metadata uses FlashInfer’s native CSR page-ID ABI; seq_lens is explicit and max_seq_len is the exact static maximum used for automatic policy selection and JIT caching. It must be no larger than 2,147,483,392 so the padded 256-token K/V tile endpoint remains representable as signed Int32. Each request must own enough CSR entries for its live length:

(seq_lens[b] + page_size - 1) // page_size <= (
    paged_kv_indptr[b + 1] - paged_kv_indptr[b]
)

The indptr must start at zero, increase strictly, and end at paged_kv_indices.numel(); every live page ID must index kv_cache.

workspace_buffer must be zero-initialized before its first use and re-zeroed whenever any workspace-layout input, including batch_size, changes because the internal section offsets can move even when the compiled callable is reused. It is exclusive to one in-flight launch or captured graph and must not overlap query, K/V cache, metadata, or output storage. Runtime sequence lengths must remain positive and no larger than max_seq_len; this hot path deliberately does not read device metadata back to the host. Live CSR, length, page-ID, and packed-Q values may change between completed launches or graph replays only while all of their contracts remain valid. They must not be mutated concurrently with a launch or replay that reads them. Warm the planned topology before CUDA graph capture and provide out to avoid an output allocation. Captured graphs must retain stable metadata storage; qo_indptr values may change only while the packed-offset contract remains valid, every delta stays within the compiled bound, and the final offset continues to match the captured query/output extent. window_left=-1 disables the left window; a non-negative value requires causal masking and includes the current token. No backend fallback or scheduling knob is exposed.

Parameters:

  • query (torch.Tensor) – Fixed or packed query tensor.

  • kv_cache (torch.Tensor or tuple[torch.Tensor, torch.Tensor]) – Combined or separate paged K/V storage.

  • workspace_buffer (torch.Tensor) – Zero-initialized caller-owned byte workspace for this planned layout.

  • paged_kv_indptr (torch.Tensor) – Native CSR row offsets and physical page IDs.

  • paged_kv_indices (torch.Tensor) – Native CSR row offsets and physical page IDs.

  • seq_lens (torch.Tensor) – Live K/V sequence lengths for each request.

  • max_seq_len (int) – Static maximum K/V length used for policy selection and JIT caching.

  • seq_len_q (int) – Fixed query length when qo_indptr is omitted.

  • qo_indptr (torch.Tensor, optional) – Cumulative query offsets selecting packed-query mode.

  • max_seq_len_q (int, optional) – Static packed-query length bound.

  • bmm1_scale (float, optional) – QK and value/output scaling factors.

  • bmm2_scale (float, optional) – QK and value/output scaling factors.

  • out (torch.Tensor, optional) – Caller-owned output tensor.

  • out_dtype (torch.dtype, optional) – Output dtype; defaults to out.dtype or the query dtype.

  • mask_type ({"dense", "causal"}) – Attention mask mode.

  • window_left (int) – Left sliding-window extent, or -1 to disable the window.

  • kv_layout ({"HND"}) – Layout of the paged K/V cache.