flashinfer.msa_ops.msa_icp_topk_select_publish

flashinfer.msa_ops.msa_icp_topk_select_publish(scores: Tensor, num_valid_pages: Tensor, recv_bases: Tensor, errors: Tensor, *, slot: int, rank: int, world_size: int, heads_per_rank: int, capacity: int, force_begin_blocks: int, force_end_blocks: int, spin_limit: int = 67108864, enable_pdl: bool = False, generations: Tensor | None = None, 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.

Select each decode (token, head) row’s local Top16 and store it into

the receive row of the head’s owner rank.

Row (t, h) goes to peer h // heads_per_rank at row ((slot * world_size + rank) * heads_per_rank + h % heads_per_rank) * capacity + t of its legacy (int32, 33 words: 16 raw scores, 16 ascending global ids, ready flag) or tagged (int64, 32 words of generation << 32 | word) receive rows. Every row is written on every call, padding rows included.

scorestorch.Tensor

float32 [T, world_size * heads_per_rank, K], 0 < K < 12288, unit last stride.

num_valid_pagestorch.Tensor

int32 [T] global B128 block counts.

recv_basestorch.Tensor

int64 [world_size] device addresses of every rank’s receive rows (the tagged rows’ addresses when generations is given).

errorstorch.Tensor

int32 [>= 1]; a legacy publish that waits spin_limit polls for a free row adds one and traps.

slotint

Receive slot set (0, 1 or 2).

capacityint

Receive rows per (slot, source, head); T <= capacity.

generationsOptional[torch.Tensor]

int32 [3, world_size * heads_per_rank, capacity] publish counters: tagged mode. None: legacy mode. Tagged mode needs at least two sparse layers between uses of one slot set; the receive storage and counters are zeroed once, never between steps.

backendstr

Only "cake".

The remaining parameters are as in msa_icp_topk_select_prefill().