flashinfer.sampling.sampling_from_probs

flashinfer.sampling.sampling_from_probs(probs: Tensor, indices: Tensor | None = None, deterministic: bool = True, generator: Generator | None = None, check_nan: bool = False, seed: int | None = None, offset: int | None = None) → Tensor

用于从概率中进行类别采样的融合 GPU 内核。

参数:

• probs (torch.Tensor) – 采样的概率。当未提供 indices 时，形状应为 (batch_size, num_classes)，第 i 个输出将从概率的第 i 行中采样。当提供 indices 时，形状应为 (unique_batch_size, num_classes)，其中 unique_batch_size 是唯一概率分布的数量。

• indices (Optional[torch.Tensor]) – 可选的 indices 张量，形状为 (batch_size,)，dtype 为 torch.int32 或 torch.int64，它将每个输出映射到 probs 中的一行。输出张量的 dtype 将与 indices 相同。例如，如果 indices[i] = j，则第 i 个输出将从 probs[j] 中采样。这允许为多个输出重用相同的概率分布。如果未提供 indices，则第 i 个输出将从 probs 的第 i 行中采样，并且输出 dtype 默认为 torch.int32。

• deterministic (bool) – 是否使用确定性内核实现，默认值为 True。

• generator (Optional[torch.Generator]) – 操作的随机数生成器。

• check_nan (bool) – 是否检查 probs 中的 nan，默认值为 False。

• seed (Optional[int]) – 采样操作期间用于 rng 的种子值。

• offset (Optional[int]) – 采样操作期间用于 rng 的偏移值。

返回值:

samples – 采样的类别，形状为 (batch_size,)。

返回值类型:

torch.Tensor

示例

>>> 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')
>>> samples = flashinfer.sampling.sampling_from_probs(norm_prob)
>>> samples
tensor([1, 2, 1, 4], device='cuda:0', dtype=torch.int32)

注意

此函数期望 float32 输入，输出为 int32。