flashinfer.attention.prims_ts.batch_decode_mla_with_paged_kv_cache

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

One-shot convenience wrapper for fixed or packed-query MLA decode.