flashinfer.sampling.top_p_renorm_probs¶
- flashinfer.sampling.top_p_renorm_probs(probs: Tensor, top_p: Tensor | float) Tensor¶
用于通过 top-p 阈值进行概率重归一化的融合 GPU 内核。
- 参数:
probs (torch.Tensor) – 概率，形状
(batch_size, num_classes)。top_p (Union[torch.Tensor, float]) – 要用于重新归一化概率的 top-p 阈值，可以是标量或形状为
(batch_size,)的张量，应在(0, 1)范围内。如果为标量，则对所有请求使用相同的阈值。如果为张量，则每个请求都有自己的阈值。我们屏蔽掉小于 threshold 的概率，其中probs[probs >= threshold]的累积和为 top_p，并重新归一化概率。
- 返回值:
renorm_probs – 重新归一化的概率，形状
(batch_size, num_classes)。- 返回值类型:
torch.Tensor
示例
>>> import torch >>> import flashinfer >>> torch.manual_seed(42) >>> batch_size = 4 >>> vocab_size = 5 >>> top_p = 0.3 >>> pre_norm_prob = torch.rand(batch_size, vocab_size).to(0) >>> prob = pre_norm_prob / pre_norm_prob.sum(dim=-1, keepdim=True) >>> 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') >>> renormed_probs = flashinfer.sampling.top_p_renorm_probs(prob, top_p) >>> renormed_probs tensor([[0.0000, 0.4882, 0.0000, 0.5118, 0.0000], [0.0000, 0.0000, 0.0000, 1.0000, 0.0000], [0.5181, 0.0000, 0.4819, 0.0000, 0.0000], [0.0000, 1.0000, 0.0000, 0.0000, 0.0000]], device='cuda:0')
注意
将
top_p_renorm_probs和sampling_from_probs的组合应等效于top_p_sampling_from_probs。