flashinfer.topk

Efficient Top-K selection kernels.

See also

For Top-K based sampling, see flashinfer.sampling which provides top_k_sampling_from_probs(), top_k_top_p_sampling_from_probs(), top_k_renorm_probs(), and top_k_mask_logits().

Top-K Selection

top_k(input, k[, sorted])	Radix-based Top-K selection.
top_k_page_table_transform(input, ...[, ...])	Fused Top-K selection + Page Table Transform for sparse attention.
top_k_ragged_transform(input, offsets, ...)	Fused Top-K selection + Ragged Index Transform for sparse attention.

Utility Functions

topk.can_implement_filtered_topk()

Check if the GPU supports enough shared memory for FilteredTopK algorithm.