flashinfer.logits_processor¶

用于构建 LLM 输出处理管道的声明式、可插拔框架。

管道构建¶

使用 LogitsPipe 创建处理管道

import torch
from flashinfer.logits_processor import LogitsPipe, Temperature, Softmax, TopP, Sample

# Create a pipeline
pipe = LogitsPipe([
    Temperature(),      # Scale logits by temperature
    Softmax(),          # Convert logits to probabilities
    TopP(),             # Apply top-p filtering
    Sample()            # Sample from the distribution
])

# Apply the pipeline
batch_size = 4
vocab_size = 5
logits = torch.randn(batch_size, vocab_size, device="cuda")
output_ids = pipe(logits, temperature=0.7, top_p=0.9)

管道¶

LogitsPipe(processors[, compile, ...])

提供了一种声明式的方式来构建 LLM 输出的处理管道。

处理器¶

`LogitsProcessor`(**params)	LogitsProcessor 定义了可以应用于 logits 或概率的高级转换。
`Temperature`(**params)	用于 logits 的温度缩放处理器。
`Softmax`([enable_pdl])	将 logits 转换为概率的 Softmax 处理器。
`TopK`([joint_topk_topp])	Top-k 过滤处理器。
`TopP`(**params)	Top-p (nucleus) 过滤处理器。
`MinP`(**params)	Min-p 过滤处理器。
`Sample`([deterministic])	生成 token 索引的采样处理器。

类型¶

`TensorType`(value[, names, module, qualname, ...])	TensorType 表示管道中张量的语义类型。
`TaggedTensor`(data, type)	在管道执行过程中维护语义类型信息的张量包装器。

自定义功能¶

自定义 Logits 处理器¶

您可以通过继承 LogitsProcessor 来创建自己的 logits 处理器

class CustomLogitsProcessor(LogitsProcessor):

    def __init__(self, **params: Any):
        super().__init__(**params)

    def legalize(self, input_type: TensorType) -> List["Op"]:
        return [CustomOp(**self.params)]

class CustomOp(Op):
    # Define the input and output tensor types
    IN = TensorType.LOGITS
    OUT = TensorType.LOGITS

    def __call__(self, tensor: TaggedTensor, **kwargs: Any) -> TaggedTensor:
        pass

pipe = LogitsPipe([CustomLogitsProcessor()])  # The pipe will be compiled into [CustomOp]

自定义融合规则¶

您可以注册自定义融合规则来优化特定的处理器组合

def custom_fusion_guard(window: List[Op]) -> bool:
    # Whether the fusion should be applied
    return True

def build_custom_fusion(window: List[Op]) -> Op:
    # Create a fused operator by setting the parameters etc.
    return CustomOp()

custom_rule = FusionRule(
    pattern=(Temperature, Softmax),
    guard=custom_fusion_guard,
    build=build_custom_fusion,
    prio=20
)

pipe = LogitsPipe(
    [Temperature(), Softmax(), Sample()],
    custom_fusion_rules=[custom_rule]
)   # The compiled ops in the pipeline will be [CustomOp, Sample]