For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /max/get-started.md).

Mojo function

sparse_indexer_prefill_score

def sparse_indexer_prefill_score[dtype: DType, KOperand: MHAOperand, num_index_heads: Int, idx_head_dim: Int, block_size: Int](q: TileTensor[dtype, address_space=q.address_space, linear_idx_type=q.linear_idx_type, element_size=q.element_size], k_operand: KOperand, input_row_offsets: TileTensor[DType.uint32, address_space=input_row_offsets.address_space, linear_idx_type=input_row_offsets.linear_idx_type, element_size=input_row_offsets.element_size], prefix_lens: TileTensor[DType.uint32, address_space=prefix_lens.address_space, linear_idx_type=prefix_lens.linear_idx_type, element_size=prefix_lens.element_size], score: TileTensor[DType.float32, address_space=score.address_space, linear_idx_type=score.linear_idx_type, element_size=score.element_size], batch: Int, total_q: Int, max_num_blocks: Int, init_blocks: Int, local_blocks: Int, sm_scale: Float32, ctx: DeviceContext)

Launch the prefill block-scoring kernel into score.

See sparse_indexer_prefill for the argument contract. Exposed separately so tests can drive scoring and selection independently.