MSA

Use the fmha_sm100 wrapper first when your project already targets the MSA package surface. Use the direct MATE APIs below when you need the native plan and runtime contract.

MSA (MiniMax Sparse Attention) covers dense, paged, and sparse attention on MUSA.

The mate package also exports these entrypoints at the top level.

Dense and paged paths

mate.msa_interface.msa_plan(batch_size: int, max_seqlen_q: int, max_seqlen_k: int, num_qo_heads: int, *, total_seqlen_k: int | None = None, num_kv_heads: int | None = None, num_kv_splits: int = -1, page_size: int = -1, sparse_block_size: int = 128, output_maxscore: bool = False, kv_block_num: int = -1, force_begin_blocks: int = 0, force_end_blocks: int = 0, force_blocks_count_in_topk: bool = True, causal: bool = True, sparse_kernel_mode: str = 'auto', use_fp8_kvcache: bool = False) → tuple[bool, int, int, MsaPlan, MsaPlan | None]

Build a capture-safe MSA plan from static capacity information.

This function never inspects live sequence-length values. max_seqlen_q and max_seqlen_k determine launch geometry. total_seqlen_k is the physical K-cache token capacity; paged callers should pass the full cache capacity rather than the sum of current request lengths. Per-request lengths and page indirection must be supplied to msa() through MsaRuntimeMetadata.

mate.msa_interface.msa(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, plan_info: MsaPlan | tuple[bool, int, int, MsaPlan, MsaPlan | None], kv_indices: torch.Tensor | None = None, kv_block_indexes: torch.Tensor | None = None, q_offset_override: int | torch.Tensor | None = None, out: torch.Tensor | None = None, max_score: torch.Tensor | None = None, **kwargs) → tuple[torch.Tensor | None, torch.Tensor | None]

Run the initial MATE-side implementation of MSA.

class mate.msa_interface.MsaPlan(batch_size: 'int', qo_lens: 'torch.Tensor', kv_lens: 'torch.Tensor', qo_offset: 'torch.Tensor', kv_page_indptr: 'Optional[torch.Tensor]', num_qo_heads: 'int', num_kv_heads: 'int', page_size: 'int', sparse_block_size: 'int', num_kv_splits: 'int', kv_block_num: 'int', force_begin_blocks: 'int', force_end_blocks: 'int', force_blocks_count_in_topk: 'bool', causal: 'bool', output_maxscore: 'bool', sparse_kernel_mode: 'str', use_fp8_kvcache: 'bool', mode: 'MsaPlanMode', prefill_plan: 'MsaPrefillPlan', decode_plan: 'Optional[MsaDecodePlan]' = None, is_static: 'bool' = False)

class mate.msa_interface.MsaPrefillPlan(cu_seqlens_q: 'torch.Tensor', cu_seqlens_k: 'torch.Tensor', total_seqlen_k: 'int', max_seqlen_q: 'int', max_seqlen_k: 'int')

class mate.msa_interface.MsaDecodePlan(cache_seqlens: 'torch.Tensor', page_size: 'int', sparse_block_size: 'int', scheduler_metadata: 'Optional[torch.Tensor]' = None)

class mate.msa_interface.MsaRuntimeMetadata(qo_lens: torch.Tensor, kv_lens: torch.Tensor, qo_offset: torch.Tensor, cu_seqlens_q: torch.Tensor | None = None, cu_seqlens_k: torch.Tensor | None = None, kv_page_indptr: torch.Tensor | None = None, page_table: torch.Tensor | None = None, seqused_k: torch.Tensor | None = None)

Device-resident, per-call metadata for a static MSA plan.

msa_plan() captures only static capacity and launch information; this object supplies the live sequence metadata on every invocation. All populated tensors must stay at fixed addresses/shapes for graph replay; MATE does not copy them to another device when this object is used.

cu_seqlens_q and cu_seqlens_k are optional for eager use and are derived from qo_lens/kv_lens when omitted. Graph callers should provide them (and seqused_k/page_table) explicitly so the call performs no metadata allocation or host synchronization. page_table is the preferred representation for paged MSA: it has a fixed [batch, max_pages] shape and can therefore be captured safely. The flat kv_page_indptr representation remains available for legacy callers.

The page table helper build_page_table_from_flat_kv_indices turns flat page indices into a fixed-shape page table for paged MSA.

Sparse paths

mate.msa_interface.sparse_msa_plan(*args, **kwargs)

mate.msa_interface.sparse_msa(*args, **kwargs)

mate.msa_interface.sparse_topk_select(max_score: torch.Tensor, topk: int, num_valid_pages: int | None = None, output: torch.Tensor | None = None, force_begin_blocks: int = 0, force_end_blocks: int = 0, force_blocks_count_in_topk: bool = True, query_positions: torch.Tensor | None = None) → torch.Tensor

mate.msa_interface.sparse_decode_atten_func(*args, **kwargs)