flashinfer.msa_ops.msa_icp_prefill_score_select

flashinfer.msa_ops.msa_icp_prefill_score_select(q: Tensor, k: Tensor, block_table: Tensor, plan, num_valid_pages: Tensor, pool: Tensor, carrier: Tensor, *, rank: int, tp: int = 2, compact: bool = False, force_begin_blocks: int = 0, force_end_blocks: int = 0, enable_pdl: bool = False, 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.

Fused prefill scorer and Top16 select of one query chunk: the carrier

rows equal those of msa_icp_index_score_sharded() followed by msa_icp_topk_select_prefill() (index_scale / index_offset of the layout), without the FP32 score rows.

The list scorer keeps, for every (CTA, plan item) cell, each query row’s 16 best candidates as sorted 64-bit keys in pool; the list-to-carrier merge folds the lists of a (token, head) row into its carrier row.

q, k, block_table, plan, rank, tp : as in msa_icp_index_score_sharded(). num_valid_pages : torch.Tensor

int32 [T]: the token’s global B128 block count (0 for padding rows).

pooltorch.Tensor

Contiguous 16-byte aligned uint8 scratch of at least msa_icp_prefill_pool_bytes() bytes for the plan’s window; never needs zeroing, and one pool serves every chunk and layer.

carriertorch.Tensor

int32 [4, T, 33], unit last stride (a token slice of a larger head-major buffer is allowed); every row t < T is written.

compactbool

Must be True exactly for P256 shards (one column per local page).

force_begin_blocks, force_end_blocksint

As in msa_icp_topk_select_prefill(); sum at most 16.

enable_pdlbool

Launch both kernels with programmatic dependent launch.

backendstr

Only "cake".