flashinfer.logits_processor.Softmax

class flashinfer.logits_processor.Softmax(enable_pdl: bool | None = None, **params: Any)

将 logits 转换为概率的 Softmax 处理器。

应用 softmax 函数。

TensorType.LOGITS -> TensorType.PROBS

参数:

enable_pdl (bool, 可选, 编译时) – 是否为内核实现启用 PDL。默认值为 True。

示例

>>> import torch
>>> from flashinfer.logits_processor import LogitsPipe, Softmax, Sample
>>> torch.manual_seed(42)
>>> pipe = LogitsPipe([Softmax()])
>>> logits = torch.randn(2, 2, device="cuda")
>>> logits
tensor([[ 0.1940,  2.1614], [ -0.1721,  0.8491]], device='cuda:0')
>>> probs = pipe(logits)
>>> probs
tensor([[0.1227, 0.8773], [0.2648, 0.7352]], device='cuda:0')

注意

只能在 pipeline 中出现一次。

__init__(enable_pdl: bool | None = None, **params: Any)

Softmax 处理器的构造函数。

参数:

enable_pdl (bool, 可选, 编译时) – 是否为内核实现启用 PDL。默认值为 None，这意味着如果设备支持 PDL，则内核将自动启用。

方法

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