flashinfer.msa_ops.msa_topk_select

flashinfer.msa_ops.msa_topk_select(max_score: Tensor, topk: int, num_valid_pages: int | None = None, output: Tensor | None = None, force_begin_blocks: int = 0, force_end_blocks: int = 0) → Tensor

Select the top-K KV blocks per query token based on attention scores.

Implements the block-scoring pass of Minimax Sparse Attention: given the per-block maximum attention scores from a cheap proxy prefill, selects the topk most important KV blocks for each (query token, head) pair and returns their sorted indices.

Parameters:

  • max_score (torch.Tensor) – Shape (num_qo_heads, max_k_tiles, total_qo_len), dtype float32. Per-KV-block maximum attention scores produced by the proxy prefill pass. Entries for invalid tiles (beyond the actual KV length) must be set to -inf by the caller.

  • topk (int) – Number of KV blocks to select per (query token, head). Must be 16.

  • num_valid_pages (int, optional) – Actual number of valid KV pages (<= max_k_tiles). Indices >= num_valid_pages are replaced with -1 and sorted to the tail. Defaults to max_k_tiles (disables clamping).

  • output (torch.Tensor, optional) – Pre-allocated output tensor of shape (total_qo_len, num_qo_heads, topk), dtype int32. Allocated internally if not provided.

  • force_begin_blocks (int) – Number of KV blocks at the beginning (sink tokens) to always include.

  • force_end_blocks (int) – Number of KV blocks at the end (local window) to always include.

Returns:

Shape (total_qo_len, num_qo_heads, topk), dtype int32. Ascending KV-block indices; -1 entries are tail-padded invalid slots.

Return type:

torch.Tensor