flashinfer.sampling

用于 LLM 采样的内核。

参见

对于高效的 Top-K 选择（无需采样），请参阅 flashinfer.topk，它提供了 top_k()、top_k_page_table_transform() 和 top_k_ragged_transform()。

sampling_from_probs(probs[, indices, ...])	用于从概率中进行类别采样的融合 GPU 内核。
top_p_sampling_from_probs(probs, top_p[, ...])	用于从概率中进行 Top-P 采样（核采样）的融合 GPU 内核，此运算符实现基于 GPU 的拒绝采样，无需显式排序。
top_k_sampling_from_probs(probs, top_k[, ...])	用于从概率中进行 Top-K 采样的融合 GPU 内核，此运算符实现基于 GPU 的拒绝采样，无需显式排序。
min_p_sampling_from_probs(probs, min_p[, ...])	用于从概率中进行 min_p 采样的融合 GPU 内核，
top_k_top_p_sampling_from_logits(logits, ...)	用于从预 softmax logits 中进行 Top-K 和 Top-P 采样的融合 GPU 内核，
top_k_top_p_sampling_from_probs(probs, ...)	用于从概率中进行 Top-K 和 Top-P 采样的融合 GPU 内核，
top_p_renorm_probs(probs, top_p)	用于通过 Top-P 阈值重新归一化概率的融合 GPU 内核。
top_k_renorm_probs(probs, top_k)	用于通过 Top-K 阈值重新归一化概率的融合 GPU 内核。
top_k_mask_logits(logits, top_k)	用于通过 Top-K 阈值屏蔽 logits 的融合 GPU 内核。
chain_speculative_sampling(draft_probs, ...)	用于序列生成进行推测采样的融合 GPU 内核（论文 Accelerating Large Language Model Decoding with Speculative Sampling 中提出），其中草稿模型为每个请求生成一系列（链）的 token。