flashinfer.logits_processor.TopK

class flashinfer.logits_processor.TopK(joint_topk_topp: bool = False, **params: Any)

Top-k 过滤处理器。

仅保留概率最高的 k 个 token，并屏蔽其余的。

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

参数:

• joint_topk_topp (bool, optional, 编译时) – 是否在后续跟随 TopP 时启用联合 top-k 和 top-p 过滤。默认值为 False。

• top_k (int 或 torch.Tensor, 运行时) – 要保留的 top token 的数量。可以是一个标量或每个批次的张量。

示例

>>> import torch
>>> from flashinfer.logits_processor import LogitsPipe, TopK, Sample, TensorType
>>> torch.manual_seed(42)
>>>
>>> # Top-k filtering on logits
>>> pipe = LogitsPipe([TopK()], input_type=TensorType.LOGITS)
>>> logits = torch.randn(2, 2, device="cuda")
>>> logits
tensor([[ 0.1940,  2.1614], [ -0.1721,  0.8491]], device='cuda:0')
>>> topk_logits = pipe(logits, top_k=1)
>>> topk_logits
tensor([[  -inf, 2.1614], [  -inf, 0.8491]], device='cuda:0')
>>>
>>> # Top-k filtering on probabilities
>>> pipe = LogitsPipe([TopK()], input_type=TensorType.PROBS)
>>> probs = torch.randn(2, 2, device="cuda")
>>> probs_normed = probs / probs.sum(dim=-1, keepdim=True)
>>> probs_normed
tensor([[  4.4998,  -3.4998], [-18.2893,  19.2893]], device='cuda:0')
>>> topk_probs = pipe(probs_normed, top_k=1)
>>> topk_probs
tensor([[1., 0.], [0., 1.]], device='cuda:0')

注意

应用于 TensorType.LOGITS 时，将非 top-k 值设置为 -inf。应用于 TensorType.PROBS 时，将非 top-k 值归零并重新归一化。

参见

top_k_mask_logits(), top_k_renorm_probs()

__init__(joint_topk_topp: bool = False, **params: Any)

TopK 处理器的构造函数。

参数:

joint_topk_topp (bool, optional, 编译时) – 是否在后续跟随 TopP 时启用联合 top-k 和 top-p 过滤。默认值为 False。

方法

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