flashinfer.msa_ops.msa_icp_index_score_sharded

flashinfer.msa_ops.msa_icp_index_score_sharded(q: Tensor, k: Tensor, block_table: Tensor, plan, *, rank: int, max_score: Tensor, sm_scale: float = 1.0, initialize: bool = True, tp: int = 2, compact: bool = False, 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 a prefill query chunk against this rank’s index-K shard.

For every token t, head h and score column c, writes the raw FP32 maximum of q[t, h] . k_row over this rank’s keys of the column’s B128 block that the token may see (causal), -inf when it sees none.

qtorch.Tensor

float8_e4m3fn [T, 4, 128], contiguous; T is the plan’s total_q.

ktorch.Tensor

float8_e4m3fn [pages, P / tp, 128]: this rank’s shard, dense rows, 16-byte aligned, page pitch a multiple of 16 bytes.

block_tabletorch.Tensor

int32 [B, >= pages per segment], unit last stride: physical page ids, the same on every rank.

planMSAIcpScorePlan

From msa_icp_index_score_sharded_plan_device().

rank, tpint

This rank and the group size (2 or 4).

max_scoretorch.Tensor

float32 [T, 4, K], unit last stride (a strided view is allowed), K at least the plan’s B128 columns (half of them, rounded up, when compact). Columns past the token’s causal prefix are written only when initialize (filled with -inf first).

sm_scalefloat

Accepted for MSA compatibility; must be finite; scores are unscaled.

compactbool

P256 only: one column per local page (global id 2 c + rank * 2 // tp) instead of one per global B128 block.

enable_pdlbool

Launch with programmatic dependent launch.

backendstr

Only "cake".