flashinfer.logits_processor.Sample

class flashinfer.logits_processor.Sample(deterministic: bool = True, **params: Any)

用于生成 token 索引的采样处理器。

从 logits 或概率分布中采样 token。

TensorType.LOGITS -> TensorType.INDICES | TensorType.PROBS -> TensorType.INDICES

参数:

• deterministic (bool, 可选, 编译时) – 是否使用确定性内核实现。默认值为 True。

• indices (torch.Tensor, 可选, 运行时) – 当概率张量共享时，用于批量采样的索引。

• generator (torch.Generator, 可选, 运行时) – 用于可重现采样的随机数生成器。

示例

>>> import torch
>>> from flashinfer.logits_processor import LogitsPipe, Sample, TensorType
>>> torch.manual_seed(42)
>>>
>>> # Sampling from logits
>>> pipe = LogitsPipe([Sample(deterministic=True)], input_type=TensorType.LOGITS)
>>> logits = torch.randn(2, 5, device="cuda")
>>> logits
tensor([[ 0.1940,  2.1614, -0.1721,  0.8491, -1.9244],
        [ 0.6530, -0.6494, -0.8175,  0.5280, -1.2753]], device='cuda:0')
>>> tokens = pipe(logits, top_k=1)
>>> tokens
tensor([0, 1], device='cuda:0')
>>>
>>> # Sampling from probabilities
>>> pipe = LogitsPipe([Sample(deterministic=True)], input_type=TensorType.PROBS)
>>> probs = torch.randn(2, 5, device="cuda")
>>> probs_normed = probs / probs.sum(dim=-1, keepdim=True)
>>> probs_normed
tensor([[ 2.8827,  0.0870,  0.2340, -3.2731,  1.0694],
        [ 0.3526,  0.0928,  0.1601, -0.1737,  0.5683]], device='cuda:0')
>>> tokens = pipe(probs_normed, top_k=1)
>>> tokens
tensor([0, 0], device='cuda:0')

注意

输出 TensorType.INDICES - 没有运算符可以跟随

参见

sampling_from_logits(), sampling_from_probs()

__init__(deterministic: bool = True, **params: Any)

Sample 处理器的构造函数。

参数:

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

方法

__init__([deterministic])	Sample 处理器的构造函数。
legalize(input_type)	将处理器合法化为低级运算符列表。