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.

MSA improves indexer performance and skips metadata computation on short queries.

The mate package also exports the primary planning and execution 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.

Sparse MSA accepts a preallocated out whose supported dtype may differ from Q/K/V. Omitting out preserves the Q dtype.

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', is_varlen: 'bool')

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, maxscore_schedule: torch.Tensor | None = None, maxscore_schedule_ready: bool = False)

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. maxscore_schedule is an optional fixed-address int32 workspace with shape [num_mps, 2]. Its rows contain the Q-work range assigned to each logical MP slot; the CTA’s 1-D grid coordinate derives the K partition. By default the device metadata kernel refreshes it on the same stream before the persistent max-score launch. Callers may set maxscore_schedule_ready when that exact workspace has already been prepared for the current sequence metadata on the same stream.

mate.msa_interface.MsaPlanInfo

Built-in immutable sequence.

If no argument is given, the constructor returns an empty tuple. If iterable is specified the tuple is initialized from iterable’s items.

If the argument is a tuple, the return value is the same object.

alias of tuple[bool, int, int, MsaPlan, MsaPlan | None]

Use build_page_table_from_flat_kv_indices to turn 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)