flashinfer.attention.prims_ts.prims_ts_batch_decode_with_kv_cache_mla

flashinfer.attention.prims_ts.prims_ts_batch_decode_with_kv_cache_mla(query: Tensor, kv_cache: Tensor, workspace_buffer: Tensor, kv_lora_rank: int, qk_rope_head_dim: int, block_tables: Tensor, seq_lens: Tensor, max_seq_len: int, *, qo_indptr: Tensor | None = None, max_seq_len_q: int | None = None, out: Tensor | None = None, bmm1_scale: float = 1.0, bmm2_scale: float = 1.0, mask_type: Literal['dense', 'causal'] = 'causal', out_dtype: dtype = torch.bfloat16) → Tensor

Launch fixed or packed-query paged MLA decode with caller-owned scratch.

With qo_indptr=None, query has fixed shape [B, SQ, H, 576]. Otherwise query has compact shape [total_q, H, 576] and qo_indptr contains the B + 1 cumulative Q offsets. Runtime Q lengths are exclusively qo_indptr[b + 1] - qo_indptr[b]; max_seq_len_q is only the static policy, JIT, and workspace bound and is required for compact launches. The last query dimension concatenates the 512 latent and 64 RoPE dimensions. kv_cache accepts compact rank-3 [pages, page_size, 576] or rank-4 [pages, 1, page_size, 576] storage. block_tables and seq_lens follow FlashInfer’s native dense paged-cache ABI; max_seq_len is the exact static policy/JIT maximum. Causal masking is bottom-right aligned: query row i can attend through KV row seq_lens[b] - q_len[b] + i for request b.

The workspace is exclusive to one in-flight launch or captured graph and must not overlap query, K/V cache, metadata, or output storage. Runtime lengths must remain positive and no larger than max_seq_len; this hot path deliberately performs no device-to-host metadata reads. For packed launches, callers must ensure that offsets start at zero, are strictly increasing, end at query.shape[0], and have every delta no larger than 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. Warm the semantic key before CUDA graph capture and provide out to avoid an output allocation. Captured graphs must retain stable qo_indptr storage; its values may change only while that packed-offset contract and the captured query/output extent remain valid. No backend fallback or scheduling knob is exposed.