flashinfer.msa_ops.msa_topk_select

flashinfer.msa_ops.msa_topk_select(max_score: Tensor, topk: int, num_valid_pages: int | Tensor | 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 or torch.Tensor, 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).

    May instead be an int32 tensor of shape (total_qo_len,) giving each query token its own valid-page count. Callers whose tokens have differing causal KV extents (decode batches, chunked prefill) want this: with a batch-wide scalar the only way to recover per-token semantics is to post-process the output on the host side, which costs several extra kernel launches per call and is easy to get wrong. force_end_blocks then denotes each token’s own trailing local window, and no selected index can exceed that token’s extent, so no masking pass is needed. Entries are clamped in-kernel to [0, max_k_tiles] (checking them on the host would sync), so an over-large count degrades to the full block range rather than reading out of bounds.

  • 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