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. Individual packed requests may be empty, and an all-empty launch returns its empty output without dispatching a GPU kernel. 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 K/V 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 nondecreasing, 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 planned topology 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.

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.

  • workspace_buffer (torch.Tensor) – Caller-owned byte workspace for this planned layout.

  • kv_lora_rank (int) – Latent and RoPE dimensions.

  • qk_rope_head_dim (int) – Latent and RoPE dimensions.

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

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

  • max_seq_len (int) – Static maximum K/V length used for policy selection and JIT caching.

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

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

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

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

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

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

  • out_dtype (torch.dtype) – Output dtype.