Qrita：使用基于枢轴的截断和选择的高性能Top-k和Top-p
Qrita: High-performance Top-k and Top-p Algorithm for GPUs using Pivot-based Truncation and Selection
AI总结：
提出Qrita算法，通过基于高斯sigma截断和四元枢轴搜索的枢轴方法，高效实现Top-k和Top-p采样，在保持与排序算法相同输出的同时，将端到端服务吞吐量提升至1.4倍并减少一半内存使用。
摘要：
尽管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获取。
英文摘要：
Despite their importance in model sampling, efficient implementation of the Top-k and Top-p algorithms poses a significant challenge as model vocabulary grows with each new model release. Existing approaches often rely on sorting, which incurs significant computation and memory overhead on GPUs, or on stochastic approaches that alter the algorithm's output. In this work, we propose Qrita, an efficient Top-k and Top-p algorithm based on pivot-based truncation and selection. Qrita leverages pivot-based search for both Top-k and Top-p with two key techniques: 1. Gaussian-based sigma-truncation, which limits the complexity of the algorithm to O(n), and 2. Quaternary pivot search with duplication handling, which halves the number of pivot search iterations and guarantees deterministic output. We implement Qrita using Triton and evaluate its performance against the Top-k and Top-p kernels of high-performance LLM execution engines such as SGLang and FlashInfer, improving end-to-end serving throughput up to 1.4 times with half the memory usage, while providing the same output as stable sort-based Top-k and Top-p algorithms. Qrita is now the default Triton Top-k and Top-p sampler of vLLM, and its implementation is available at https://github.com/vllm-project/vllm/blob/main/vllm/v1/sample/ops/topk_topp_triton.py.
机构、总结、中文摘要均由AI分析而来，可能存在错误，请以论文正文为准。
