flashinfer.msa_ops.prepare_msa_icp_indexer

flashinfer.msa_ops.prepare_msa_icp_indexer(*, world_size: int, rank: int, num_index_heads: int, page_size: int, init_blocks: int, local_blocks: int, max_decode_tokens: int, max_decode_query_len: int, max_num_batched_tokens: int, max_num_seqs: int, max_model_len: int, num_sparse_layers: int, device: device, backend: str = 'cake') → Any

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.

Build this rank’s ICP indexer from the serving limits.

The returned MSAIcpIndexerPlan runs one layer’s chains with caller-owned buffers:

  • decode(index_q, index_k, block_table, seq_lens, num_valid_pages, scores, output, *, comm, slot, tagged, decode_query_len, enable_pdl=True, pool=None): score plane -> Top16 select and publish -> owner merge (the default; its score plane follows the enable_pdl contract of msa_icp_decode_index_score_sharded()); with pool the fused scorer and publisher of msa_icp_decode_score_publish() replace the first two steps (scores unused), and without pool they do so in the facade’s own pool for batches of plan.fused_decode_min_batch or more requests (compute capability 10.7 at P128 TP2 and TP4: 8; None elsewhere);

  • prefill_window_plan(window, ...) once per query window and step, prefill_local(query, index_k, block_table, window_plan, score_scratch, num_valid_pages, carrier_view, *, pool=None) per query chunk, the caller’s one exchange per layer, then prefill_merge(received, num_valid_pages, output). With pool (uint8 of plan.prefill_pool_shape(), caller-owned) prefill_local runs the fused list scorer and list-to-carrier merge of msa_icp_prefill_score_select() instead (score_scratch may be None; windows of plan.prefill_window_tokens() tokens); plan.fused_prefill tells whether that path is the serving default for the layout (False everywhere for now).

comm holds the receive storage views and counters laid out by msa_icp_comm_layout() (recv, tagged_recv, recv_bases, tagged_recv_bases, publish_generations, expected_generations, errors; check them once with check_comm). Call warmup() on every rank before the first decode or prefill exchange and before CUDA Graph capture: it loads every kernel into the device’s context, so no first launch waits on module loading while a peer’s merge spins. The facade allocates only its prefill plan storage, and the fused chain’s pool where fused_decode_min_batch is set, here; no method allocates or synchronizes.

world_size, rankint

Tensor-parallel group size (2 or 4) and this rank.

num_index_headsint

Index heads across the group; must be 4.

page_sizeint

128 or 256.

init_blocks, local_blocksint

Forced begin / end blocks (sum at most 16).

max_decode_tokens, max_decode_query_lenint

Receive capacity and the longest decode query (4 * len <= 32).

max_num_batched_tokens, max_num_seqs, max_model_lenint

Scheduler limits that size the prefill plan storage.

num_sparse_layersint

Sparse indexer layers per step; tagged decode rows need at least two.

devicetorch.device

The rank’s CUDA device.

backendstr

Only "cake".