For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /max/get-started.md).

Mojo function

persistent_topk_block_split

def persistent_topk_block_split(ctx: DeviceContext, in_scores: Pointer[Float32, ImmutAnyOrigin, _safe=False], out_idxs: Pointer[Int32, MutAnyOrigin, _safe=False], N: Int, K: Int, total_seq_len: Int)

Launch bitonic top-k, splitting the N dimension when rows under-fill GPU.

Same contract and output as persistent_topk_block. When the row count is small relative to the SM count and N spans many tiles (the long-context decode regime — a handful of blocks would otherwise each fold the whole row serially), the streaming fold is split across rows * S blocks (phase 1), each producing a sorted top-_TILE partial, then merged per row (phase 2). All other shapes fall back to persistent_topk_block unchanged.

Args:

ctx (DeviceContext): Device context.
in_scores (Pointer[Float32, ImmutAnyOrigin, _safe=False]): Flat score buffer [total_seq_len × N] row-major.
out_idxs (Pointer[Int32, MutAnyOrigin, _safe=False]): Output buffer [total_seq_len × K] row-major (int32).
N (Int): Score columns per token.
K (Int): Top-k count per token (≤ N, and ≤ PERSISTENT_TOPK_MAX_N when N > 2048).
total_seq_len (Int): Number of rows.