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,520 so the padded 128-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 an argument contributing to the semantic JIT key changes, because the internal section offsets can change with that key. 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 semantic key 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.