flashinfer.sdsa.indexer_threshold_select

flashinfer.sdsa.indexer_threshold_select(query: Tensor, key_cache: Tensor, row_ends: Tensor, tile_requests: Tensor, block_table: Tensor, *, thresholds: Tensor, counters: Tensor, candidates: Tensor, max_context: int, top_k: int, query_scales: Tensor | None = None, key_scales: Tensor | None = None, arrivals: Tensor | None = None, ranges: Tensor | None = None) → None

Select the keys that score above each query row’s threshold, in one pass over the row’s paged keys and without writing its scores.

A key that scores above its row’s threshold is a candidate, and one that scores neither above nor below it (equal, or NaN) is a tie. Each is counted and stored as a (key index, score bits) pair in the row’s candidate or tie slots, its score less the threshold where that is exact (scores at most twice it), which keeps their order. A row fits when its candidates fit capacity and, with at most top_k ties, reach min(row_ends[r], top_k). Its exact Top-top_k keys are then among its slots, and a threshold at or a little below its top_k-th largest score makes it fit.

With arrivals and ranges, and a capacity of at least top_k - 1, every row fits. Any threshold then gives the same result, and only decides how often a repair runs: the last CTA of a query tile to finish sweeps the tile again for each row that does not fit, and selects it again at its radix-selected min(row_ends[r], top_k)-th largest score. It runs as one launch without host synchronization, so it can be captured in a CUDA graph (run it once before capture to compile it).

Parameters:

  • query (torch.Tensor) – As in indexer_scores(), in tiles of 128 rows.

  • key_cache (torch.Tensor) – As in indexer_scores().

  • row_ends (torch.Tensor) – int32 [rows]: row r sees keys below row_ends[r].

  • tile_requests (torch.Tensor) – int32 [tiles]: each query tile’s request, a row of block_table.

  • block_table (torch.Tensor) – int32 [requests, blocks]: each request’s pages of 64 keys.

  • thresholds (torch.Tensor) – float32 [rows]: each row’s threshold. A repaired row gets its new threshold.

  • counters (torch.Tensor) – int32 [rows, 2], zeroed by the caller: each row’s candidate and tie counts.

  • candidates (torch.Tensor) – int32 [rows, top_k + capacity, 2], output: the n-th candidate at top_k + n if n < capacity, the n-th tie at top_k - 1 - n if n < top_k. Slot order is unspecified.

  • max_context (int) – The most keys a row sees. It sizes the grid, and the block table must cover it.

  • top_k (int) – The number of tie slots.

  • query_scales (torch.Tensor, optional) – NVFP4 scales, as in indexer_scores().

  • key_scales (torch.Tensor, optional) – NVFP4 scales, as in indexer_scores().

  • arrivals (torch.Tensor, optional) – int32 [tiles]: zeroed once, and left zeroed by every call.

  • ranges (torch.Tensor, optional) – int32 [2, rows], output: the range of slots holding each row’s candidates and ties. Other slots are unspecified.