flashinfer.sampling

Kernels for LLM sampling.

Important

Batch-length seed and offset tensors now apply one value per output row. Previously, all rows used element zero, even when a batch-length tensor was supplied. Calls with distinct values in these tensors can therefore produce different samples after this fix. Scalar values, length-one tensors, and calls using a shared torch.Generator retain their existing behavior.

The output row index still feeds the Philox subsequence. Per-row seeds do not make sampling invariant to batch position, even with the same seed and offset.

See also

For efficient Top-K selection (without sampling), see flashinfer.topk which provides top_k(), top_k_page_table_transform(), and top_k_ragged_transform().

sampling_from_probs(probs[, indices, ...])	Fused GPU kernel for category sampling from probabilities.
sampling_from_logits(logits[, indices, ...])	Fused GPU kernel for category sampling from logits.
softmax(logits[, temperature, enable_pdl])	Fused GPU kernel for online safe softmax with temperature scaling.
top_p_sampling_from_probs(probs, top_p[, ...])	Fused GPU kernel for top-p (nucleus) sampling from probabilities.
top_k_sampling_from_probs(probs, top_k[, ...])	Fused GPU kernel for top-k sampling from probabilities.
min_p_sampling_from_probs(probs, min_p[, ...])	Fused GPU kernel for min_p sampling from probabilities.
top_k_top_p_sampling_from_logits(logits, ...)	Fused GPU kernel for top-k and top-p sampling from pre-softmax logits.
top_k_top_p_sampling_from_probs(probs, ...)	Fused GPU kernel for top-k and top-p sampling from probabilities.
top_p_renorm_probs(probs, top_p[, ...])	Fused GPU kernel for renormalizing probabilities by top-p thresholding.
top_k_renorm_probs(probs, top_k[, ...])	Fused GPU kernel for renormalizing probabilities by top-k thresholding.
top_k_mask_logits(logits, top_k)	Fused GPU kernel for masking logits by top-k thresholding.
chain_speculative_sampling(draft_probs, ...)	Fused-GPU kernel for speculative sampling for sequence generation (proposed in paper Accelerating Large Language Model Decoding with Speculative Sampling), where the draft model generates a sequence (chain) of tokens for each request.