flashinfer.logits_processor.LogitsPipe

class flashinfer.logits_processor.LogitsPipe(processors: List[LogitsProcessor], compile: bool = True, input_type: TensorType | None = None, custom_fusion_rules: List[FusionRule] | None = None, custom_validity_checks: List[Callable[[List[Op]], None]] | None = None)

提供了一种声明式的方式来构建 LLM 输出的处理流水线。

参数:

• processors (List[LogitsProcessor]) – 按顺序应用的处理器列表。

• compile (bool, optional) – 是否使用融合优化编译流水线。默认值为 True。可以在流水线实例化后调用 LogitsPipe.compile() 来执行编译。

• input_type (Optional[TensorType], optional) – 期望的输入张量类型。它可以是 TensorType.LOGITS 或 TensorType.PROBS。如果第一个处理器可以接受这两种类型，则需要它。在其他情况下，将从第一个处理器自动推断它。默认值为 None。

• custom_fusion_rules (Optional[List[FusionRule]], optional) – 编译期间应用的额外融合规则。默认值为 None。

• custom_validity_checks (Optional[List[ValidityCheck]], optional) – 编译期间应用的额外有效性检查。默认值为 None。

示例

>>> import torch
>>> from flashinfer.logits_processor import LogitsPipe, Temperature, Softmax, TopK, Sample
>>> torch.manual_seed(42)
>>>
>>> # Basic pipeline with temperature, top-k, and sampling
>>> pipe = LogitsPipe([
...     Temperature(),
...     Softmax(),
...     TopK(),
...     Sample(deterministic=True)
... ])
>>>
>>> # Execute the pipeline
>>> logits = torch.randn(4, 32000, device="cuda")  # [batch_size, vocab_size]
>>> token_ids = pipe(logits, temperature=0.9, top_k=40)
>>> token_ids
tensor([15806,  8154, 13923, 20311], device='cuda:0')
>>>
>>> # Pipeline starting from probabilities
>>> from flashinfer.logits_processor import TensorType, TopK
>>>
>>> prob_pipe = LogitsPipe(
...     [TopK(), Sample()],
...     input_type=TensorType.PROBS
... )
>>> probs = torch.softmax(logits, dim=-1)
>>> token_ids = prob_pipe(probs, top_k=40)
>>> token_ids
tensor([  346, 14846,  1517,  9006], device='cuda:0')

注意

• 流水线自动验证操作之间的类型兼容性。

• 操作在可能的情况下会被融合

• 运行时参数（如 temperature、top_k）通过 pipe.call() 传递。

• 输出始终是纯 torch.Tensor，而不是 TaggedTensor。

__init__(processors: List[LogitsProcessor], compile: bool = True, input_type: TensorType | None = None, custom_fusion_rules: List[FusionRule] | None = None, custom_validity_checks: List[Callable[[List[Op]], None]] | None = None)

LogitsPipe 的构造函数。

方法

__init__(processors[, compile, input_type, ...])	LogitsPipe 的构造函数。
compile([custom_fusion_rules, ...])	编译流水线。

属性

initial_type

流水线的初始输入张量类型。