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.

Parameters:

  • query (torch.Tensor) – Fixed or packed query tensor with concatenated latent and RoPE heads.

  • kv_cache (torch.Tensor) – Compact paged latent K/V cache.

  • block_tables (torch.Tensor) – Dense physical-page table for each request.

  • seq_lens (torch.Tensor) – Live K/V sequence lengths.

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

  • max_seq_len_q (int, optional) – Static packed-query length bound.

  • kv_lora_rank (int) – Latent and RoPE dimensions.

  • qk_rope_head_dim (int) – Latent and RoPE dimensions.

  • mask_type ({"dense", "causal"}) – Attention mask mode.

  • max_kv_len (int, optional) – Static K/V length bound; defaults to the metadata maximum.

  • bmm1_scale (float) – QK and value/output scaling factors.

  • bmm2_scale (float) – QK and value/output scaling factors.

  • out (torch.Tensor, optional) – Caller-owned output tensor.

  • out_dtype (torch.dtype) – Output dtype.