🤖 AI

Qrita: High-performance Top-k and Top-p using Pivot-based Truncation and Selection

本文介绍了 Qrita，这是一种针对大词汇量的高性能确定性 Top-k 和 Top-p 采样算法，它利用基于枢轴的截断与选择机制，相比现有 GPU 内核实现了高达 1.4 倍的吞吐量提升和 50% 的内存占用降低，从而被采纳为 vLLM 的默认采样器。

原作者： Jongseok Park, Sunga Kim, Alvin Cheung, Ion Stoica

发布于 2026-05-27

📖 1 分钟阅读☕ 轻松阅读

CC BY 4.0

原作者： Jongseok Park, Sunga Kim, Alvin Cheung, Ion Stoica

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象你是一位经营着巨大厨房（大型语言模型）的主厨，你的储藏室里存放着10 万种不同的食材（词汇表）。每次你需要烹饪一道菜（生成一个词）时，都必须从这巨大的储藏室中挑选出最好的几种食材，以确保菜肴美味且不单调。

在人工智能领域，这种挑选过程被称为Top-k（挑选前 k 种最佳食材）和Top-p（挑选一组最小的食材，使其累计“风味评分”达到特定阈值）。

问题：缓慢且杂乱的厨房

目前，大多数厨房处理这一问题的方式是：将所有 10 万种食材倾倒在一张巨大的操作台上，按“风味评分”从优到劣进行排序，然后抓取排名靠前的食材。

问题所在：对 10 万项物品进行排序既缓慢又混乱。它占用巨大的操作台空间（内存），并且会让主厨的双手感到疲惫（计算开销）。
替代方案：有些主厨试图直接随机抓取一把“可能”不错的食材。但这存在风险；有时你会错过最好的食材，有时每次烹饪出来的菜肴都会不同（非确定性），如果你需要精确复制食谱，这就很糟糕。

解决方案：Qrita（智能副主厨）

这篇论文介绍了Qrita，一种用于挑选食材的高速新方法。可以将 Qrita 想象为一位超级聪明的副主厨，它利用两个巧妙的技巧完全跳过了繁琐的排序过程。

技巧一：“高斯 Sigma 截断”（噪声过滤器）

想象在你的 10 万种食材储藏室中，有 9.9 万种只是“噪声”（比如盐、糖和面粉，它们的味道大致相同且平淡）。只有几百种是“明星食材”（比如松露或藏红花）。

Qrita 不需要查看每一个罐子，而是进行快速的嗅探测试。它计算储藏室的平均“味道”和“辣度”（标准差）。然后它画出一条线：“低于这条线的任何东西都只是噪声；我们无需查看。”

结果：它瞬间丢弃了 99% 的储藏室，只留下大约 200 个有趣的罐子。这一过程在单次闪电般的遍历中完成。

技巧二：“四元枢轴搜索”（四向分割）

现在，主厨面前有一小堆 200 种有趣的食材。他们仍然需要找出确切的前 50 名。

旧方法：逐个检查（太慢）或将堆栈分成两半（二分搜索）。
Qrita 的方法：Qrita 不是将堆栈分成两半，而是一次性将其分成四个部分。它会问：“最好的食材是在第一部分、第二部分、第三部分还是第四部分？”
额外优势：它还有一个针对重复食材的特殊规则。如果三个“藏红花”罐子的得分完全相同，Qrita 确切知道需要保留多少个，以确保每次食谱都完全一致（确定性）。这防止了主厨在试图决定选择哪个相同的罐子时陷入无限循环。

为何这很重要（结果）

作者使用一种名为Triton（一种用于编程图形处理器/GPU 的语言）的专用工具构建了 Qrita，并将其与 vLLM 和 SGLang 等主要 AI 引擎目前使用的最佳方法进行了测试。

速度：在现实世界的服务场景中，Qrita 的速度快达1.4 倍，在原始速度测试中快达2 倍。
内存：它使用的内存只有一半，因为它不需要存储整个 10 万项物品的排序列表。
准确性：与一些靠猜测的快速方法不同，Qrita 给出的结果与缓慢但完美的排序方法完全一致。它不会改变输出，只是更快地找到它。

核心结论

Qrita 就像是将一位手动对世界上每一种香料进行排序的主厨，升级为一位智能助手，它能瞬间忽略无聊的东西，将有趣的东西一次性分成四堆，并完美处理重复项。它在不改变答案质量的情况下，使 AI 生成更快、更高效。

注：论文提到，Qrita 现已成为 vLLM（一款用于运行 AI 模型的流行工具）GPU 路径的默认方法，且代码已开放供他人使用。

技术摘要：Qrita - 基于枢轴截断与选择的高性能 Top-k 与 Top-p 算法

问题陈述

在 GPU 上为具有大词汇量（通常超过 10 万个 token）的大语言模型（LLM）高效实现 Top-k 和 Top-p 采样算法仍然是一个重大挑战。现有方法在计算效率与输出正确性之间面临权衡：

基于排序的方法： vLLM 等引擎使用此类方法，它们提供精确的确定性结果，但由于全词汇量排序和步长内存访问，导致计算和内存开销高昂。
随机方法： FlashInfer 和 SGLang 等库使用拒绝采样来避免排序，提供高吞吐量。然而，它们仅返回单个采样 token，而非完整的候选集，因此与需要确定性候选集的任务（如推测解码验证和强化学习）不兼容。
统计截断方法： Top-nσ和 Min-p 等方法利用统计阈值（如均值和标准差）消除排序。虽然高效，但它们改变了采样分布和输出，引发了关于生成质量和连贯性的争议。

目前亟需一种算法，能够在不牺牲 Top-k 和 Top-p 精确语义的前提下，提供内存效率和严格的确定性行为（返回完整候选集）。

方法论

作者提出了Qrita，这是一种基于枢轴截断与选择的高性能 GPU Top-k 和 Top-p 算法。Qrita 扩展了 RTop-k（最初为图神经网络设计）的基于枢轴的策略，以处理大型 LLM 词汇量，同时确保确定性。该方法依赖于两项关键技术：

1. 基于高斯的 $\sigma$ -截断

认识到 LLM 的 logit 通常遵循准高斯分布，其中高概率 token 在“噪声”背景中形成异常值，Qrita 采用预过滤步骤。

机制： 该算法将 logit 分布近似为 $Z \sim \mathcal{N}(\mu, \sigma^2)$ 。它从样本块中计算均值（ $\mu$ ）和标准差（ $\sigma$ ），并计算截断阈值 $t = \mu + \delta \cdot \sigma$ 。
查找表： 系数 $\delta$ 通过基于正态分位函数预先计算的查找表确定，针对目标 $k$ 或 $p$ 值进行了定制。
结果： 此步骤单次遍历即可截断词汇量中巨大的“噪声”区域。如果截断后的集合大小 $|Z_t|$ 超过目标 $k$ （即“截断命中”），搜索仅在较小的子集上进行。如果未命中，则回退到完整 logit 集。这显著减少了后续枢轴搜索的搜索空间。

2. 带重复处理的四元枢轴搜索

为了高效且确定性地找到精确的枢轴值 $\tau$ （恰好有 $k$ 个元素 $\ge \tau$ ）：

四元搜索： 与二分搜索不同，Qrita 在每次迭代中评估三个枢轴（当前搜索范围的 25%、50% 和 75% 分位数）。与二分搜索相比，这将搜索迭代次数减少了一半。
重复处理： 标准枢轴搜索在存在重复 logit 值时（在低精度或长尾分布中很常见）可能会失败或变得不确定。Qrita 跟踪大于枢轴的最小 logit 值（ $z_{mn}$ ）及其出现次数（ $N_{zmn}$ ）。
终止条件： 当找到枢轴 $\tau_n$ 时，搜索终止，条件是：严格大于 $\tau_n$ 的元素数量小于 $k$ ，但大于或等于最小重复值 $z_{mn}$ 的元素数量大于或等于 $k$ 。这使得算法能够精确确定必须保留多少个边界值的重复实例以满足 Top-k 约束，从而确保可重现的确定性输出。

实现细节

框架： 使用 Triton（一种 GPU 编程语言）实现。
自动调优： 采用“分层 Triton 块大小自动调优”策略。由于完整 logit 集和截断集之间的向量大小差异巨大，编译器被允许选择最优的块大小（例如，完整集使用较大的块，截断集使用较小的块），以最大化硬件利用率。
流水线： 该算法并行处理批次，依次应用高斯截断、四元搜索，最后生成重复处理掩码以产生最终输出。

主要贡献

确定性高性能采样： Qrita 提供完整的 Top-k/Top-p 候选集（不同于 FlashInfer/SGLang），同时保持基于排序算法的精确语义，使其适用于推测解码和强化学习。
搜索空间缩减： 利用 logit 的统计特性，高斯 $\sigma$ -截断大幅减少了搜索空间，避免了全词汇量排序。
鲁棒的枢轴搜索： 引入四元搜索并结合显式的重复处理，保证了确定性行为并减少了迭代次数，解决了标准枢轴方法在大词汇量上的低效问题。
集成： Qrita 已作为 vLLM GPU 执行路径的默认 Top-k 和 Top-p 采样器实现。

结果

作者在 NVIDIA H100、RTX4090 和 AMD MI355X GPU 上使用各种模型（Llama、Qwen、GPT-OSS、Gemma）评估了 Qrita，基线包括朴素 PyTorch 排序、vLLM 默认采样器、FlashInfer 和 SGLang。

吞吐量： 在 H100 上，Qrita 相比 vLLM 基线实现了高达 1.4 $\times$ 的端到端服务吞吐量提升，部分配置在 RTX4090 上显示出高达 2 $\times$ 的提升。它显著优于基于排序的基线，同时匹配基于排序算法的输出。
内存使用： 与基于排序的方法相比，Qrita 将内存使用减少了约 50%，这主要是通过避免分配排序所需的大型中间缓冲区实现的。
消融研究：
- 高斯截断： 提供了最大的加速比，与搜索完整 logit 集相比，延迟降低了约 74%。
- 四元搜索： 虽然在“截断命中”情况下由于搜索空间减小而略慢于二分搜索，但在必须搜索完整词汇量的回退情况下，它提供了显著的延迟降低（约 2.3 毫秒）。
- 重复处理： 增加了最小的开销（延迟增加约 3.5%），但对于正确性和确定性至关重要。
- 自动调优： 与使用单一固定块大小相比，将执行延迟降低了近 30%。

意义与主张

该论文声称，Qrita 弥合了大词汇量解码中内存效率与严格确定性行为之间的差距。它断言，通过避免全词汇量排序并利用带有统计预过滤的基于枢轴的选择，Qrita 为 LLM 解码提供了一种高性能解决方案，且兼容需要完整候选集的高级任务（如推测解码）。作者将 Qrita 定位为并非 Top-k/Top-p 统计语义的替代品，而是一种能够精确保留这些语义的高效实现，不同于启发式统计截断方法或随机单 token 采样器。将其作为默认采样器集成到 vLLM 中，突显了其在生产环境中的实际效用。