flashinfer.msa_ops.msa_icp_merge_candidates

flashinfer.msa_ops.msa_icp_merge_candidates(recv_slot: Tensor, num_valid_pages: Tensor, output: Tensor, errors: Tensor, *, init_blocks: int, local_blocks: int, spin_limit: int = 67108864, enable_pdl: bool = False, expected_generations: Tensor | None = None, prefill: 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.

Merge every rank’s Top16 candidates of the heads this rank owns.

Per (token, owned head): ids outside [0, min(num_valid_pages, 8192)) are dropped, a duplicate id keeps its largest score (NaN counts as -inf), ids < init_blocks or >= num_valid_pages - local_blocks rank first, exact ties go to the lower id, and the Top16 ids are written ascending with a -1 tail. A prefill row with num_valid_pages == 0 is written [0, -1, ...]; a decode one all -1.

recv_slottorch.Tensor

One slot set of the receive rows: legacy int32 [W, hpr, capacity, 33] (decode: waits for each ready flag, then clears it; prefill: the exchanged carrier, read without waiting) or tagged int64 [W, hpr, capacity, 32] (with expected_generations). Unit last stride.

num_valid_pagestorch.Tensor

int32 [T].

outputtorch.Tensor

int32 [T, hpr, 16], unit last stride, T <= capacity.

errorstorch.Tensor

int32 [>= 1]; a wait of spin_limit polls adds one and traps.

expected_generationsOptional[torch.Tensor]

int32 [hpr, capacity] merge counters of the slot set (tagged mode).

prefillbool

Prefill rows; incompatible with tagged mode and enable_pdl.

backendstr

Only "cake".