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

Mojo function

persistent_topk_block_split

def persistent_topk_block_split[ordered: Bool = False, deterministic: Bool = False](ctx: DeviceContext, in_scores: Pointer[Float32, ImmutAnyOrigin], out_idxs: Pointer[Int32, MutAnyOrigin], N: Int, K: Int, total_seq_len: Int, row_bounds: Optional[Pointer[Int32, ImmutAnyOrigin]] = None)

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.

Parameters:

• ​ordered (Bool): Whether slot q must hold the q-th largest score. True is the historical contract: descending score, ties by ascending column. False promises the same K columns, deterministically, in an unspecified order -- which lets the short-row path skip its ranking pass, most of its work. The set does not depend on this, because the tie-break lives in the key the select already compares. Only shapes that have a cheaper path take one; the rest stay ordered, which satisfies the weaker promise too.
• ​deterministic (Bool): Whether one input must always give one output. True, the default, is the guarantee above: the same K columns every run. False drops it, and drops the set guarantee with it -- when more columns tie at the threshold than there are slots left, which of them survive is decided by the order the block's warps reach a cursor. In exchange the cheap contract's tail needs no ordering scan at all. Only meaningful together with ordered=False; the ordered path is deterministic by construction.

Args:

• ​ctx (DeviceContext): Device context.
• ​in_scores (Pointer[Float32, ImmutAnyOrigin]): Flat score buffer [total_seq_len × N] row-major.
• ​out_idxs (Pointer[Int32, MutAnyOrigin]): Output buffer [total_seq_len × K] row-major (int32).
• ​N (Int): Score columns per token (the row stride).
• ​K (Int): Top-k count per token (≤ N, and ≤ PERSISTENT_TOPK_MAX_N when N > 2048).
• ​total_seq_len (Int): Number of rows.
• ​row_bounds (Optional[Pointer[Int32, ImmutAnyOrigin]]): Optional [total_seq_len] int32 per-row live-column counts; see persistent_topk_block. Row r scans only [0, row_bounds[r]) of its N-wide stripe, so under capture-frozen metadata (where N is a worst-case bound) the scan cost tracks each row's real length. Supported on every path this launcher selects for N > 2048 (the histogram-select family) and on the 2048 single-block path.