Module block_sparse 

MiniMax-M3’s block-sparse attention selection: which 128-token KV blocks a query may look at.

Ported from FreeToken’s models/minimax_m3/args.py selection rule. This is the decision half – which blocks are visible. The attention itself is then an ordinary masked pass over the positions those blocks cover, which crate::attention::causal_mla_attention_sparse already does.

Four rules, and one of them is what stops a NaN

• A block’s score is the MAX over its positions, not the mean and not the sum. A block earns its place on its single best match: one strongly-related token in a block of 128 is exactly the case sparse attention exists to catch, and a mean would average it away against 127 unrelated neighbours.

• No softmax scale. The raw dot product is used directly, because only the ORDER of the scores is consumed. Dividing by sqrt(d) would scale every score by the same positive constant and change nothing, so the reference does not, and neither does this.

• Selection is per KV head, with no cross-head reduction. One index head scores for one KV head, and that KV head’s whole GQA group reads the blocks it picked. Reducing across heads first – by summing or maxing the scores – would give every group the same block set, which is the opposite of what per-head selection is for.

• The newest local_blocks and the first init_blocks are force-included, before any scoring. This is not a quality heuristic bolted on top: it is what guarantees the selection is never empty. A query early in a sequence, or one whose scores are all equally poor, would otherwise select zero blocks, and attention over zero positions is a softmax over an empty set – which is a NaN that propagates through the whole forward pass and surfaces as garbage output, not as an error.

The block size is an ABI, not a tuning knob

MINIMAX_BLOCK_SIZE is 128 and also pins the KV page size: the selection hands back block indices, and a pager whose page is a different size cannot honour them without splitting or merging pages, which is exactly the bookkeeping block-sparse attention exists to avoid.

Structs

BlockSparseConfig

How many blocks a query may see, and which are free.

Constants

MINIMAX_BLOCK_SIZE

MiniMax-M3’s KV block, in tokens. Also the KV page size – see the module docs.

Functions

block_sparse_select

The blocks each KV head may read for one query position, ascending.

positions_of_blocks

The positions a block selection covers, ascending and causally clipped – the form crate::attention::causal_mla_attention_sparse takes.