flashinfer.sdsa.top_k_from_candidates

flashinfer.sdsa.top_k_from_candidates(candidates: Tensor, ranges: Tensor, *, top_keys: Tensor) → None

Write each query row’s exact Top-K keys from the candidates and ties that indexer_threshold_select() stored in its range of slots.

The select stores scores that keep the order of the row’s scores, so the keys of the row’s top_k largest stored scores are the keys of its top_k largest scores. A CTA takes a row: it finds the top_k-th largest score from counts of the score bits, then writes the keys above it and those equal to it. 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:

• candidates (torch.Tensor) – int32 [rows, slots, 2]: the select’s (key, score bits) pairs.

• ranges (torch.Tensor) – int32 [2, rows]: the select’s ranges, the first and past-last slot of each row’s pairs.

• top_keys (torch.Tensor) – int32 [rows, top_k], output: the keys of the row’s top_k largest scores, in no set order, and any of those equal to the top_k-th. A row of at most top_k pairs gets all its keys, then -1.