flashinfer.msa_ops.msa_icp_index_score_sharded_plan_device

flashinfer.msa_ops.msa_icp_index_score_sharded_plan_device(qo_segment_lens: Tensor, kv_segment_lens: Tensor, qo_segment_offsets: Tensor, kv_segment_offsets: Tensor, qo_offset: Tensor, num_qo_heads: int, *, total_q: int, max_query_len: int, max_seq_len: int, packed_work_range: Tensor, packed_work_info: Tensor, workspace: Tensor, 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 on the device into caller storage.

Graph safe: no device value is read on the host. The plan is reused by every layer of the step; keep the storage of each concurrently live plan distinct.

qo_segment_lens, kv_segment_lens, qo_offsettorch.Tensor

int32 CUDA [B], contiguous. Zero-length segments are allowed.

qo_segment_offsets, kv_segment_offsetstorch.Tensor

int32 CUDA [B + 1] prefix sums of the lengths, contiguous.

num_qo_headsint

Index heads; must be 4.

total_q, max_query_len, max_seq_lenint

Host upper bounds of the query tokens, the longest query segment and the longest key length.

packed_work_rangetorch.Tensor

int64 [num_SMs], contiguous.

packed_work_infotorch.Tensor

int64, at least ((total_q + 127) // 128 + B - 1) * num_qo_heads elements (msa_icp_index_score_plan_storage()).

workspacetorch.Tensor

uint8, at least 32 MiB (accepted for MSA compatibility; not written).

backendstr

Only "cake".

MSAIcpScorePlan