flashinfer.msa_ops.msa_icp_prefill_pool_bytes

flashinfer.msa_ops.msa_icp_prefill_pool_bytes(*, total_q: int, batch: int, device: device, num_heads: int = 4) → int

Bytes of the fused prefill path’s pool for a query window of at most total_q tokens in batch segments on device: one 32 KiB piece per (scorer CTA, plan item) cell, ceil(total_q / 64) + batch - 1 + num_sms - 1 pieces.