Qrita：使用基于枢轴的截断和选择的 GPU 高性能 Top-k 和 Top-p 算法

Top-k 和 Top-p 是大语言模型采样中的主要截断运算符。尽管它们被广泛使用，但在大词汇量上有效地实施它们仍然是一个重大挑战。现有方法通常依赖于排序，这会在 GPU 上产生大量计算和内存开销；或者依赖于随机方法，这会改变算法输出。在这项工作中，我们提出了 Qrita，一种基于枢轴选择策略的高效 Top-k 和 Top-p 算法。 RTop-k（在图神经网络中使用基于枢轴的搜索进行节点选择）基于 RTop-k，Qrita 通过两项关键技术将基于枢轴的搜索概念扩展到 Top-k 和 Top-p：1. 基于高斯的 sigma 截断，大大减少了目标元素的搜索空间；2. 具有重复处理的四元枢轴搜索，可将枢轴搜索迭代减半并保证确定性输出。我们使用流行的 GPU 编程语言 Triton 提供 Qrita 的完整实现。我们针对 vLLM、SGLang 和 Flashinfer 等高性能 LLM 执行引擎的 Top-k 和 Top-p 内核对 Qrita 进行的评估表明，Qrita 实现了高达 2 倍的吞吐量和一半的内存使用量，同时为基于排序的算法提供相同的输出。