flashinfer.topk

高效的 Top-K 选择内核。

参见

对于基于 Top-K 的采样，请参阅 flashinfer.sampling，它提供了 top_k_sampling_from_probs(), top_k_top_p_sampling_from_probs(), top_k_renorm_probs(), 和 top_k_mask_logits()。

Top-K 选择

top_k(input, k[, sorted])	基于基数的 Top-K 选择。
top_k_page_table_transform(input, ...[, ...])	融合的 Top-K 选择 + 页表转换，用于稀疏注意力。
top_k_ragged_transform(input, offsets, ...)	融合的 Top-K 选择 + 锯齿索引转换，用于稀疏注意力。

实用函数

topk.can_implement_filtered_topk()

检查 GPU 是否支持 FilteredTopK 算法所需的足够共享内存。