flashinfer.attention.prims_ts.get_prims_ts_batch_decode_workspace_size

flashinfer.attention.prims_ts.get_prims_ts_batch_decode_workspace_size(batch_size: int, num_qo_heads: int, num_kv_heads: int, head_dim: int, page_size: int, max_seq_len: int, *, seq_len_q: int = 1, qo_indptr: Tensor | None = None, max_seq_len_q: int | None = None, q_dtype: dtype = torch.float16, kv_dtype: dtype | None = None, out_dtype: dtype | None = None, mask_type: Literal['dense', 'causal'] = 'dense', window_left: int = -1, kv_layout: Literal['HND'] = 'HND', device: int | str | device | None = None) → int

Return caller-workspace bytes for one automatic FMHA policy.

The arguments define the same semantic JIT key as prims_ts_batch_decode_with_kv_cache(). The query resolves policy and scratch layout but does not compile a kernel. Allocate at least the returned number of bytes as a contiguous torch.int8 or torch.uint8 CUDA tensor and zero it before its first FMHA launch. Re-zero a reused buffer whenever an argument contributing to the semantic JIT key changes, because the internal section offsets can change with that key. Fixed-Q launches use seq_len_q. Packed-Q launches provide qo_indptr and the explicit static max_seq_len_q bound used for workspace geometry and JIT policy. max_seq_len must be no larger than 2,147,483,520 so the padded 128-token K/V tile endpoint remains representable as signed Int32. This sizing helper validates that every cumulative-offset delta is positive and no larger than the bound. If device is omitted, it is inferred from qo_indptr for a packed launch.