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

会议: NeurIPS 2026
arXiv: 2602.01518
OpenReview: https://openreview.net/forum?id=4dnBXjV4Za
代码: https://github.com/vllm-project/vllm/blob/main/vllm/v1/sample/ops/topk_topp_triton.py
领域: LLM 效率
关键词: LLM 采样、Top-k / Top-p、支点选择、高斯截断、GPU 算子优化

一句话总结¶

Qrita 针对超大词表下传统排序采样内存显存开销过大以及拒绝采样无法保留全量候选集的问题，提出结合高斯 \(\sigma\)-截断与四分法支点搜索的精确 Top-k/Top-p 采样算子，在显存占用减半的同时将端到端吞吐提升高达 1.4 倍并保证输出完全确定性。

研究背景与动机¶

在大语言模型（LLM）的自回归解码过程中，Top-\(k\) 与 Top-\(p\)（核采样）截断是控制生成质量与文本多样性最常用的核心机制。Top-\(k\) 保留 Logit 最高的 \(k\) 个候选 Token 以滤除不连贯的尾部噪声，Top-\(p\) 则根据累积概率质量动态截断，使候选池根据模型置信度自适应伸缩。然而，随着现代 LLM 词表规模扩展到 10 万至 25 万以上，在 GPU 上高效实现精确的 Top-\(k\) 与 Top-\(p\) 面临严峻的硬件效率瓶颈。主流推理引擎以往依赖全词表排序（如双调排序 Bitonic Sort）或多次前缀和计算，跨步内存访问频繁且需要庞大的中间结果缓冲区，在显存带宽受限的 GPU 上造成了显著的计算与内存开销。

为了消除排序瓶颈，近期推理引擎（如 FlashInfer）引入了基于对偶支点的迭代拒绝采样，虽然在统计期望上与原采样分布等价，但其接口直接返回单一样本 Token，彻底丢失了完整的候选 Token 集合。这种概率近似对投机解码（Speculative Decoding）候选集验证、强化学习（RL）探索以及语义缓存等需要显式候选集下游任务完全不兼容；同时，浮点前缀和的非结合性也引入了跨硬件非确定性。另一方面，基于分布截断的纯启发式方案（如 Min-\(p\)、Top-\(n\sigma\)）虽然是单遍线性计算，但根本性改变了输出分布与生成文本连贯性，业界生产环境仍严格要求标准的 Top-\(k\) / Top-\(p\) 语义。

图神经网络中用于节点选取的支点算法（RTop-\(k\)）提供了免排序的思路，但其二分搜索面向小向量设计，面对超大词表时多轮全局扫描会迅速饱和 GPU 显存带宽，且缺乏对重复 Logit 的处理机制，容易陷入震荡与非确定性输出。本文的切入角度是：利用 LLM 输出 Logit 近似高斯分布的先验先滤除绝大多数背景噪声，再在极小的保留集合上进行多支点加速收敛与精确重复去重。核心 idea：将高斯 \(\sigma\)-截断前置过滤与四分法支点搜索相结合，在无需全词表排序的前提下通过三支点并行探测减半扫描轮次，并借助前缀和精确处理边界重复 Logit，实现显存减半、完全确定性且语义严格一致的超高性能采样。

方法详解¶

整体框架¶

Qrita 接收批处理 Logit 矩阵 \(\mathbf{Z} \in \mathbb{R}^{B \times \mathcal{V}}\) 与采样阈值参数，针对每个批次样本独立并行执行算子逻辑，整体架构流水线包含三大连续阶段：高斯 \(\sigma\)-截断过滤（Gaussian \(\sigma\)-Truncation）、四分法支点搜索（Quaternary Pivot Search）与重复 Logit 去重处理（Duplication Handling）。在底层实现上，Qrita 采用持久化内核（Persistent Kernel）架构，根据 GPU 的流式多处理器（SM）数量启动 Triton 程序，并在批次维度和块维度执行分级块大小自动调优（Tiered Block Size Autotuning）。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入未排序 Logit 矩阵 <b>Z</b><br/>(词表规模 V ≥ 100k, 批大小 B)"] --> B["高斯 σ-截断过滤<br/>单遍均值方差估计 + 查表阈值裁剪"]
    B -->|保留集合大小 |Z_t| > k| C["候选收缩集合 Z_s = Z_t<br/>(元素数量收敛到数百至数千)"]
    B -->|截断未命中 / 保留元素不足| D["全量回退集合 Z_s = Z<br/>(保证理论下界与严格正确性)"]
    C --> E["四分法支点搜索<br/>三支点 (25%, 50%, 75%) 并行探测收敛"]
    D --> E
    E --> F["重复 Logit 去重处理<br/>前缀和定位最小越界值并精确保留 N_keep 项"]
    F --> G["输出严格 Top-k / Top-p 掩码与候选集 <b>O</b>"]

关键设计¶

1. 高斯 \(\sigma\)-截断过滤：单遍估计大幅削减搜索空间 现代 LLM 的词表即便超过 10 万，最终有效预测概率通常高度集中在前几十到上百个 Token 中，海量背景 Token 在数值上呈现高斯噪声分布。为避免在整个词表上进行多次昂贵的数据扫描，Qrita 首先将 Logit 分布建模为拟高斯分布 \(\mathbf{Z} \sim \mathcal{N}(\mu, \sigma^2)\)。通过采样计算 Logit 的样本均值 \(\mu\) 与标准差 \(\sigma\)，根据目标采样阈值 \(k\) 或 \(p\) 通过预计算的查找表获取系数 \(\delta\)，构造截断阈值：

\[t = \mu + \delta \cdot \sigma - 0.2 \times |t|\]

其中减去 \(0.2 \times |t|\) 作为安全裕度（Safety Margin），以确保真实有效候选集被高置信度完全覆盖。只有当收集到的候选子集满足 \(|\mathbf{Z}_t| > k\)（或对 Top-\(p\) 累积概率和满足 \(\sum p_i > p\)）时触发截断命中，后续支点搜索仅在极小的紧凑子集 \(\mathbf{Z}_t\)（通常仅数百个元素）上进行；若出现未命中则自动安全回退至全量 Logit \(\mathbf{Z}\)。此步骤将后续所有内存访问量降低了数个数量级。

2. 四分法支点搜索：三支点并行探测减半迭代步数 传统二分支点选择在超大词表上需要高达几十次全局内存扫描。Qrita 将每次迭代的探测支点数从 1 个扩展为 3 个，在当前搜索区间 \([l, r]\) 内分别取 25%、50% 与 75% 分位数位置：

\[\tau_n = l + \frac{n}{4} (r - l), \quad n \in \{1, 2, 3\}\]

在单次规约遍历中，算子同时统计大于各支点的元素总数 \(N_{\tau_n} = |\{z_i \mid z_i > \tau_n\}|\)，并动态跟踪大于各支点的最小 Logit 值 \(z_{mn} = \min(\{z_i \mid z_i > \tau_n\})\) 及其出现频次 \(N_{z_{mn}}\)。三支点探测将有效区间收缩速度提升一倍，搜索轮次减少约 50%。同时，算法引入了兼顾重复值的精确终止准则：只要满足 \((N_{\tau_n} \ge k) \wedge (N_{\tau_n} - N_{z_{mn}} < k)\)，即可判定支点 \(\tau_n\) 与重复边界值 \(z_{mn}\) 已跨越第 \(k\) 个边界，直接提前终止搜索并返回，彻底规避了传统方法在重复 Logit 存在时的死循环风险。

3. 重复 Logit 去重处理：前缀和索引保障严格确定性 在低精度（如 FP16 / BF16）推断或特定退化分布下，往往存在多个完全相等的边界 Logit。若不加以区分，会导致选出的 Token 数量超过 \(k\) 或产生随机扰动。Qrita 提出一种全并行掩码修正逻辑：基于搜索得到的支点 \(\tau\) 生成基准布尔掩码 \(\mathbf{M}_\tau = \mathbb{I}(\mathbf{Z} > \tau)\)，并计算需要保留的边界重复项配额：

\[N_{keep} = N_{dup} - (N_\tau - k)\]

随后定位所有等于边界重复值 \(z_{dup}\) 的元素位置构建指示掩码 \(\mathbf{M}_{z_{dup}} = \mathbb{I}(\mathbf{Z} = z_{dup})\)，对其执行并行前缀和（cumsum）计算：

\[\mathbf{C}_{z_{dup}} = \text{cumsum}(\mathbf{M}_{z_{dup}}) \wedge \mathbf{M}_{z_{dup}}\]

通过将重复元素的出现位次与 \(N_{keep}\) 进行比较，标记出超出配额的重复项并予以剔除：\(\mathbf{M}_{remv} = \mathbb{I}(\mathbf{C}_{z_{dup}} > N_{keep})\)。最终掩码 \(\mathbf{M}_{final} = \mathbf{M}_\tau \wedge \neg \mathbf{M}_{remv}\) 保证仅保留词表中最早出现的 \(N_{keep}\) 个重复项。未入选位置被就地置为 \(-\infty\)。该设计在纯并行硬件上实现了与全词表排序完全一致、可跨运行完美复现的确定性语义。

4. 分级 Triton 块大小自动调优：自适应适配异构计算负载 由于高斯 \(\sigma\)-截断命中与未命中时待处理的操作数规模差异极大（截断命中时仅数百至数千元素，回退时则多达数十万元素），单一固定的 GPU 块划分（Block Size）会导致严重的不匹配：小块在全量模式下调度开销巨大，大块在紧凑模式下则产生过度填充（Padding）与资源空转。Qrita 引入分级块配置策略，对全量张量采用较大块大小、对截断子集采用较小块大小，并通过 Triton Autotune 机制由编译器在离线阶段自动搜索最优的流水线深度与线程块参数组合，实现近 30% 的额外延迟降低。

损失函数 / 训练策略¶

Qrita 是纯粹面向推理阶段的底层 GPU 算子优化算法，不改变神经网络权重，亦不需要任何模型重训练或微调。对于查找表中的参数 \(\delta\)，通过标准正态分布分位数函数 \(\Phi_k^{-1}(1 - k/\mathcal{V})\) 以及离线采样经验累积分布预先构建，无需在线动态拟合。

实验关键数据¶

主实验¶

在 NVIDIA H100 GPU 上，结合 vLLM 推理框架测试不同主流开源模型在 WikiText-2 数据集上的端到端推理服务吞吐（输入/输出分别为 256 与 1024 Token），对比 vLLM 原始默认排序采样器与跳过采样内核的理论上限性能：

模型架构	词表规模	采样配置	原始 vLLM (KTok/s)	Qrita (KTok/s)	理论上限 (KTok/s)	吞吐加速比
R1-Distill-Llama-8B	128,256	\(k=50, p=0.9\)	12.39	14.22	14.50	1.15×
Qwen3-8B	151,936	\(k=50, p=0.9\)	11.57	13.57	13.68	1.16×
GPT-OSS-20B (MoE 3B active)	201,088	\(k=50, p=0.9\)	19.32	26.43	28.39	1.41×
Gemma3-4B	262,208	\(k=50, p=0.9\)	11.51	15.27	15.41	1.31×

消融实验¶

在固定词表规模 128,256、批大小 1024、\(k=50, p=0.9\) 条件下，系统评估 Qrita 各核心组件对算子执行时间（ms）的直接贡献：

实验配置编号	截断状态	搜索算法	去重逻辑	自动调优	执行时间 (ms)	说明
Run A (完整 Qrita)	截断命中	四分法	包含	包含	0.919	完整优化算子基线
Run B	截断命中	四分法	包含	排除	1.354	移除分级块调优，延迟恶化 +47.3%
Run C	截断命中	四分法	排除	排除	1.308	移除重复项处理，仅轻微提速 3.5%
Run D	截断回退	四分法	排除	排除	5.373	模拟截断未命中扫描全词表
Run E	无截断	四分法	排除	排除	4.931	彻底禁用截断过滤机制
Run F	截断命中	二分法	排除	排除	1.281	紧凑子集下二分与四分性能相当
Run G	截断回退	二分法	排除	排除	7.773	全量回退时四分法比二分法提速 2.4ms
Run I	截断命中	排序算法	排除	排除	2.099	子集排序，耗时为四分法 1.6 倍
Run J (PyTorch 基线)	无截断	排序算法	排除	排除	16.487	全词表排序，耗时为完整 Qrita 17.9 倍

关键发现¶

高斯截断是加速核心：高斯 \(\sigma\)-截断带来了高达 74% 的延迟下降（Run C 对比 Run E），通过将词表搜索范围从 128k 压缩到数百项，彻底解除了全局显存带宽瓶颈；截断本身开销约 0.4ms，在整体运行时中占比约 44%。
四分法在兜底场景表现稳健：在截断命中时，四分法与二分法差距不明显；但在截断回退或大词表未命中场景下，四分法减少的访存轮次将延迟从 7.773ms 骤降至 5.373ms，提供了关键的性能兜底。
大词表与 MoE 模型收益更显著：在参数量较小但词表巨大的模型（如 Gemma3-4B，词表 262k）或单步计算量较小的 MoE 稀疏激活模型（如 GPT-OSS-20B）上，端到端加速比分别达到 1.31× 与 1.41×，推理性能达到理论跳过采样上限的 95% 以上。

亮点与洞察¶

统计滤波与精确搜索解耦：巧妙利用 Logit 分布的拟高斯特性进行激进空间截断，但不以统计阈值作为硬切分替代采样，而是仅仅将其作为快速候选初筛，后续借助严格支点搜索保证标准数学语义不变。
全并行确定性边界处理：巧妙设计前缀和索引与计数器联动机制，将传统 CPU 串行分支中的重复项选择转化为纯 GPU 向量化位运算与布尔掩码，彻底解决低精度截断非确定性与死循环问题。
高通用性跨硬件移植能力：完全使用 Triton 语言编写，脱离对特定 CUDA 专有闭源库的依赖，在 AMD MI355X 与消费级 RTX4090 上均能开箱即用并取得翻倍性能提升。

局限与展望¶

批次维度硬件扩展上限：当前 Qrita 实现中一个 Triton 程序绑定一个 SM，在 H100（132 个 SM）上当 Batch Size 超过 128 时并行度提升变缓，未来需要设计单个 Triton 程序并发处理多个批次样本的细粒度划分。
Top-p 截断阈值方差较大：相较于 Top-\(k\) 始终能达到 100% 截断命中率，Top-\(p\) 由于 Softmax 指数放大作用对极端离群值敏感，在部分模型上命中率有所波动，针对特定模型引入离线参数 Profiling 能进一步收紧裕度。
中间临时缓冲区占用：目前截断实现预分配了全量词表长度的临时存储空间，未来可直接将保留缓冲区大小硬性限制在几千元素，从而使中间显存消耗进一步收缩。

评分¶

新颖性: ⭐⭐⭐⭐☆ 巧妙将统计截断、多支点搜索与并行前缀和去重融为一体，开创了超大词表免排序精确采样新范式。
实验充分度: ⭐⭐⭐⭐⭐ 涵盖 4 款主流大模型、多款不同厂商 GPU（H100、RTX4090、MI355X），消融与端到端对比扎实详尽。
写作质量: ⭐⭐⭐⭐⭐ 算法逻辑严密清晰，问题背景阐述透彻，系统级权衡与设计洞察深刻。
价值: ⭐⭐⭐⭐⭐ 已合并成为 vLLM 官方默认 GPU 采样算子，对大模型推理服务基础设施具有重大实用价值。