flashinfer.sampling.sampling_from_probs

flashinfer.sampling.sampling_from_probs(probs: torch.Tensor, uniform_samples: torch.Tensor)

Fused GPU kernel for category sampling from probabilities.

Parameters:

• probs (torch.Tensor) – Probabilities, shape (batch_size, num_classes).

• uniform_samples (torch.Tensor) – The uniform samples used as needle for sampling, shape (batch_size,). Expected to be uniformly distributed in [0, 1).

Returns:

samples – Sampled categories, shape (batch_size,).

Return type:

torch.Tensor

Examples

>>> import torch
>>> import flashinfer
>>> torch.manual_seed(42)
>>> batch_size = 4
>>> vocab_size = 5
>>> pre_norm_prob = torch.rand(batch_size, vocab_size).to(0)
>>> norm_prob = pre_norm_prob / pre_norm_prob.sum(dim=-1, keepdim=True)
>>> norm_prob
tensor([[0.2499, 0.2592, 0.1085, 0.2718, 0.1106],
        [0.2205, 0.0942, 0.2912, 0.3452, 0.0489],
        [0.2522, 0.1602, 0.2346, 0.1532, 0.2000],
        [0.1543, 0.3182, 0.2062, 0.0958, 0.2255]], device='cuda:0')
>>> uniform_samples = torch.rand(batch_size).to(0)
>>> samples = flashinfer.sampling.sampling_from_probs(norm_prob, uniform_samples)
>>> samples
tensor([1, 2, 1, 4], device='cuda:0', dtype=torch.int32)

Notes

This function expects float32 inputs, and the output is int32.