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

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

Launch block-wide bitonic top-k for total_seq_len score rows.

For N ≤ PERSISTENT_TOPK_MAX_N (= 2048) a single block sorts the whole row. For N > PERSISTENT_TOPK_MAX_N a streaming variant folds _TILE-wide tiles into a running top-_TILE champion; this requires K ≤ PERSISTENT_TOPK_MAX_N (the champion width). Call sites needing K > PERSISTENT_TOPK_MAX_N must use topk_gpu.

Each row of N float32 scores yields the K highest-scoring column indices (as int32) in descending score order in out_idxs.

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.
• ​K (Int): Top-k count per token (≤ N, and ≤ PERSISTENT_TOPK_MAX_N when N > 2048).
• ​total_seq_len (Int): Number of rows (one block per row).