flashinfer.mla.trtllm_prefill_with_kv_cache_mla

flashinfer.mla.trtllm_prefill_with_kv_cache_mla(query: Tensor, kv_cache: Tensor, workspace_buffer: Tensor, qk_nope_head_dim: int, kv_lora_rank: int, qk_rope_head_dim: int, block_tables: Tensor, seq_lens: Tensor | None, max_seq_len: int, sparse_mla_top_k: int = 0, out: Tensor | None = None, bmm1_scale: float | Tensor = 1.0, bmm2_scale: float | Tensor = 1.0, sinks: List[Tensor] | None = None, skip_softmax_threshold_scale_factor: float | None = None, enable_pdl: bool | None = None, backend: str = 'auto', is_var_seq: bool = True, uses_shared_paged_kv_idx: bool = True, lse: Tensor | None = None, return_lse: bool = False, cute_dsl_impl: str = 'auto', kv_scale_format: str = 'auto', cum_seq_lens_q: Tensor | None = None, max_q_len: int | None = None, multi_ctas_kv_counter_buffer: Tensor | None = None, sparse_mla_top_k_lens: Tensor | None = None, enable_dcp: bool = False, cp_world: int = 1, cp_rank: int = 0, causal_seqlens_kv_global: Tensor | None = None, use_fp16_softmax: bool | None = None, return_lse_base: Literal['basee', 'base2'] | None = None) → Tensor | Tuple[Tensor, Tensor]

Run MLA with decode-compatible semantics for prefill callers.

This function shares its implementation, argument and return contract, exceptions, and backend selection with trtllm_batch_decode_with_kv_cache_mla(). Use q_len_per_request == 1 for decode, or use a compatible backend with q_len_per_request > 1 for incremental prefill and multi-token prediction (MTP).

XQA only supports q_len_per_request == 1. Because backend="auto" may select XQA, callers that require multi-token prefill must select a compatible backend explicitly.

This entrypoint is also available as flashinfer.prefill.trtllm_prefill_with_kv_cache_mla(). See trtllm_batch_decode_with_kv_cache_mla() for complete parameter and return-value documentation.