flashinfer.msa_ops.msa_icp_index_score_sharded_plan

flashinfer.msa_ops.msa_icp_index_score_sharded_plan(qo_segment_lens: Tensor, kv_segment_lens: Tensor, num_qo_heads: int, *, qo_offset: Tensor, device: device | None = None, backend: str = 'cake')

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.

Plan the sharded prefill scorer from host segment lengths (eager only).

qo_segment_lens, kv_segment_lens, qo_offsettorch.Tensor

int32 CPU [B]: query tokens, key length and query start position (causal prefix) of every segment. The lengths are read on the host.

num_qo_headsint

Index heads; must be 4.

deviceOptional[torch.device]

Device of the plan storage (default: the current CUDA device).

backendstr

Only "cake".

MSAIcpScorePlan

The plan; its storage is allocated here. Use msa_icp_index_score_sharded_plan_device() for CUDA Graphs.