Qrita：基于枢轴截断与选择的高性能 Top-k 与 Top-p
人工智能
2026-05-27 v2
摘要
尽管 Top-k 和 Top-p 在模型采样中至关重要，但其大词汇表的高效实现仍是一个重大挑战。现有方法通常依赖排序，导致在 GPU 上计算和内存开销显著，或依赖随机方法会改变算法的输出。在本工作中，我们提出 Qrita，一种基于枢轴截断与选择的高效 Top-k 和 Top-p 算法。Qrita 利用枢轴搜索实现 Top-k 和 Top-p，包含两个关键技术：1. 基于高斯的 sigma 截断，大幅减少词汇表的搜索空间；2. 四分枢轴搜索带重复处理，减半枢轴搜索迭代次数并保证确定性输出。我们使用 Triton 实现 Qrita，并对比高性能 LLM 执行引擎如 SGLang 和 FlashInfer 的 Top-k 和 Top-p 内核，显示在端到端服务吞吐量提升最高可达 1.4 倍，内存使用仅为排序方法的一半，同时提供相同输出。Qrita 现为 vLLM GPU 执行路径的默认 Top-k 和 Top-p 采样器，亦提供基于三角实现的 Qrita，地址为 https://github.com/vllm-project/vllm/blob/main/vllm/v1/sample/ops/topk_topp_triton.py。
引用
@article{arxiv.2602.01518,
title = {Qrita: High-performance Top-k and Top-p using Pivot-based Truncation and Selection},
author = {Jongseok Park and Sunga Kim and Alvin Cheung and Ion Stoica},
journal= {arXiv preprint arXiv:2602.01518},
year = {2026}
}