flashinfer.logits_processor.LogitsProcessor¶
- class flashinfer.logits_processor.LogitsProcessor(**params: Any)¶
LogitsProcessor defines high-level transformations that can be applied to logits or probabilities. Each processor is automatically legalized into low-level
Op
orParameterizedOp
that can be type-checked, validated, and fused for optimal performance. Users can extend this class to implement their own processors.- Parameters:
**params (Any) – Processor-specific parameters at compile-time.
Examples
>>> import torch >>> from flashinfer.logits_processor import LogitsPipe, TopK, Sample, TensorType >>> torch.manual_seed(42) >>> >>> # Create a pipeline that legalizes to a fused op. >>> pipe = LogitsPipe([ ... TopK(), # Top-k filtering on logits ... Sample() # Sample from the filtered distribution ... ], input_type=TensorType.PROBS) # assume the input is probabilities >>> >>> pipe LogitsPipe([TopK -> Sample], ops=[ProbsTopKOp -> ProbsSampleOp], compiled_ops=[FusedProbsTopKSampleOp])
Notes
Subclasses must implement the
legalize()
method to convert the high-level processor into one or more low-level operators with specific input/output types- __init__(**params: Any)¶
Initialize the processor.
- Parameters:
**params (Any) – Processor-specific parameters at compile-time.
Methods
__init__
(**params)Initialize the processor.
legalize
(input_type)Legalize the processor into a list of low-level operators.