flashinfer.attention.prims_ts.get_prims_ts_batch_decode_mla_workspace_size

flashinfer.attention.prims_ts.get_prims_ts_batch_decode_mla_workspace_size(batch_size: int, num_heads: int, kv_lora_rank: int, qk_rope_head_dim: int, page_size: int, max_seq_len: int, *, seq_len_q: int | None = None, max_seq_len_q: int | None = None, q_dtype: dtype = torch.bfloat16, kv_dtype: dtype | None = None, out_dtype: dtype = torch.bfloat16, mask_type: Literal['dense', 'causal'] = 'causal', device=None) → int

Return caller-workspace bytes for one automatic MLA policy.

The arguments define the same semantic JIT key as prims_ts_batch_decode_with_kv_cache_mla(). Policy and private scratch layout are resolved without compiling a kernel. max_seq_len_q is the static per-request Q bound for both fixed and packed-query launches; seq_len_q remains a backward-compatible fixed-Q alias. If neither is supplied, the bound is one. The returned byte count includes both split-KV scratch and the internal FP32 LSE tensor. Allocate a contiguous torch.int8 or torch.uint8 CUDA buffer; MLA does not require its contents to be initialized before first use.