flashinfer.sampling.top_k_top_p_sampling_from_logits¶
- flashinfer.sampling.top_k_top_p_sampling_from_logits(logits: Tensor, top_k: Tensor | int, top_p: Tensor | float, indices: Tensor | None = None, filter_apply_order: str = 'top_k_first', deterministic: bool = True, generator: Generator | None = None, check_nan: bool = False, seed: int | None = None, offset: int | None = None) Tensor¶
用于从预 softmax logits 进行 top-k 和 top-p 采样的融合 GPU 内核，
此算子实现基于 GPU 的拒绝采样，无需显式排序。请查看 博客文章 以获取更多详细信息。
多次拒绝采样的轮次在单个 CUDA 内核中实现，这比启动一系列内核的朴素实现更有效。
- 参数:
logits (torch.Tensor) – 用于采样的预 softmax logits。当未提供 indices 时，形状应为
(batch_size, num_classes)，第 i 个输出将从 logits 的第 i 行中采样。当提供 indices 时，形状应为(unique_batch_size, num_classes)，其中 unique_batch_size 是唯一概率分布的数量。top_k (Union[torch.Tensor, int]) – 一个标量或形状为
(batch_size,)的张量，表示 top-k 采样的阈值。如果是一个标量，则对所有请求使用相同的阈值。如果是一个张量，则每个请求都有自己的阈值。top_p (Union[torch.Tensor, float]) – 一个标量或形状为
(batch_size,)的张量，表示 top-p 采样的阈值。如果是一个标量，则对所有请求使用相同的阈值。如果是一个张量，则每个请求都有自己的阈值。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。filter_apply_order (str) – 应用 top-k 和 top-p 采样的顺序，应为
"top_k_first"或"joint"。如果"top_k_first"，我们首先应用 top-k 过滤器，然后在 top-k 结果上应用 top-p 采样。如果"joint"，我们同时在每一轮中应用 top-k 和 top-p 过滤器。默认值为"top_k_first"。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 >>> top_p = 0.5 >>> top_k = 3 >>> logits = torch.rand(batch_size, vocab_size).to(0) >>> logits tensor([[ 1.9269, 1.4873, 0.9007, -2.1055, -0.7581], [ 1.0783, 0.8008, 1.6806, 0.3559, -0.6866], [-0.4934, 0.2415, -0.2316, 0.0418, -0.2516], [ 0.8599, -0.3097, -0.3957, 0.8034, -0.6216]], device='cuda:0') >>> samples = flashinfer.sampling.top_k_top_p_sampling_from_logits(logits, top_k, top_p) >>> samples tensor([0, 2, 1, 3], device='cuda:0', dtype=torch.int32 >>> probs = torch.softmax(logits, dim=-1) >>> probs tensor([[0.4788, 0.3085, 0.1716, 0.0085, 0.0327], [0.2358, 0.1787, 0.4307, 0.1145, 0.0404], [0.1358, 0.2831, 0.1764, 0.2318, 0.1729], [0.3613, 0.1122, 0.1029, 0.3415, 0.0821]], device='cuda:0') >>> samples tensor([0, 2, 1, 3], device='cuda:0', dtype=torch.int32)
注意
此函数期望 float32 输入，输出为 int32。