flashinfer.msa_ops.msa_icp_topk_select_prefill

flashinfer.msa_ops.msa_icp_topk_select_prefill(scores: Tensor, num_valid_pages: Tensor, carrier: Tensor, *, force_begin_blocks: int = 0, force_end_blocks: int = 0, index_scale: int = 1, index_offset: int = 0, backend: str = 'cake') → None

Warning

MiniMax-M3 ICP indexer kernels (sharded index scores, Top16 select/publish, merge) is experimental: it provides no compatibility guarantees and may change or be removed without deprecation.

Write each (head, token) row’s local Top16 into a prefill carrier.

scorestorch.Tensor

float32 [T, 4, K], 0 < K < 12288, unit last stride.

num_valid_pagestorch.Tensor

int32 [T]: the token’s global B128 block count (0 for padding rows). Local column c is global block c * index_scale + index_offset; blocks at or past min(num_valid_pages, 8192) are not candidates.

carriertorch.Tensor

int32 [4, T, 33], unit last stride (a token slice of a larger head-major buffer is allowed). Row [h, t] receives 16 raw score words, 16 global ids ascending ((-inf, -1) padding) and 1.

force_begin_blocks, force_end_blocksint

Blocks < force_begin_blocks and >= num_valid_pages - force_end_blocks are always selected; sum at most 16.

index_scale, index_offsetint

(1, 0) for P128, (2, rank * 2 // tp) for compact P256 scores.

backendstr

Only "cake".