flashinfer.attention.prims_ts.batch_mla_decode_with_paged_kv_cache

flashinfer.attention.prims_ts.batch_mla_decode_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.

This helper reads seq_lens and, for packed Q, qo_indptr on the host to derive plan bounds, then constructs a temporary wrapper. Invoke it outside CUDA Graph capture. Capture-sensitive callers should pre-plan BatchMLADecodePagedTSWrapper and use run(validate=False).

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. Rows must be inner contiguous and non-overlapping, but may have padding between them.

  • 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) – Per-request packed-query length capacity. For packed Q it defaults to the maximum delta in qo_indptr; an explicit value may be larger. For fixed Q it defaults to the query’s sequence extent and, when provided, must equal that extent.

  • 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.

Returns:

The fixed or packed MLA attention output.

Return type:

torch.Tensor