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, block_tables: 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 page-table 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. block_tables is an int32 CUDA tensor with shape [B, C]. Entries are contiguous within each row while the outer row stride may be greater than C. 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. The table must contain enough columns for every per-run length:

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

Every active page ID must index kv_cache. Inactive tail entries at and beyond ceil(seq_lens[b] / page_size) are ignored.

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. Per-run table, 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.

  • block_tables (torch.Tensor) – CUDA int32 physical page IDs with shape [B, C], unit inner stride, and non-overlapping row stride at least C.

  • 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. In packed-query mode, a non-default value is a backward-compatible alias for max_seq_len_q and must agree with it when both are provided.

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

  • max_seq_len_q (int, optional) – Per-request packed-query length capacity. When omitted for packed Q, a non-default seq_len_q must supply the bound. In fixed-query mode, it must equal seq_len_q.

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