flashinfer.msa_ops.msa_icp_decode_index_score_sharded

flashinfer.msa_ops.msa_icp_decode_index_score_sharded(index_q: Tensor, index_k: Tensor, block_table: Tensor, seq_lens: Tensor, score_out: Tensor, *, rank: int, world_size: int, page_size: int, decode_query_len: int, max_decode_query_len: int, 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.

Score decode queries against this rank’s index-K shard.

Writes score_out[t, h, c] for the local page columns c of the request’s ceil(seq_len / page_size) pages rounded up to a multiple of four (P128: column = global B128 block; P256: column = local page, global block 2 c + rank * 2 // world_size); other columns are not written.

index_qtorch.Tensor

float8_e4m3fn [B * decode_query_len, 4, 128], contiguous; token j of a request sits at position seq_len - decode_query_len + j.

index_ktorch.Tensor

float8_e4m3fn [pages, page_size / world_size, 128] shard.

block_tabletorch.Tensor

int32 [B, >= ceil(seq_len / page_size)], unit last stride.

seq_lenstorch.Tensor

int32 [B].

score_outtorch.Tensor

float32 [T, 4, cols], 0 < cols <= 8192, cols % 4 == 0, unit last stride (msa_icp_decode_score_columns()).

rank, world_size, page_sizeint

Layout: world_size 2 or 4, page_size 128 or 256.

decode_query_len, max_decode_query_lenint

1 <= decode_query_len <= max_decode_query_len, 4 * max_decode_query_len <= 32.

enable_pdlbool

Launch with programmatic dependent launch. On compute capability 10.0 and 10.3 with 64 or 128 local K rows per page (P128 TP2, P256 TP2 and TP4), and on 10.7 with 32 or 64 (P128 TP4 and TP2, P256 TP4), the scorer then reads seq_lens and block_table and starts loading the index_k pages that end at or before seq_len - decode_query_len before it waits on the preceding kernel. The kernel launched just before it must not write them, and must either be launched without programmatic dependent launch or wait on its own predecessor before it triggers its dependents, so that every earlier write is complete. enable_pdl=False removes the requirement.

backendstr

Only "cake".