SonicSampler：用于LLM采样和推测性验证的统一Tile感知内核

# SonicSampler：用于LLM采样和投机验证的统一分片感知内核

来源：https://arxiv.org/html/2607.20475

\\contribution \[\*\]贡献相等\\correspondence\\codeTBA（2026年5月）

###### 摘要

LLM推理中的采样包含一组组合操作，包括logits处理、令牌选择和投机解码的验证。然而，现有的实现要么仅加速该流水线的子集，要么依赖多次内核启动，要么假设整个批次内采样行为同质，限制了对动态服务工作负载的支持，并阻碍了高效的CUDA Graph执行。

我们提出SonicSampler，一套统一的分片感知Triton内核，该内核将完整的采样流水线垂直融合到一个固定的、工作负载感知的执行模型中。我们的内核支持动态的每请求采样行为，包括语法约束解码、重复惩罚、频率和存在惩罚、logits偏置、温度缩放、top-k/ top-p/ min-p过滤以及投机验证——所有这些都在单个批处理内核中完成，同时保持完全兼容CUDA Graph。

我们方法的核心是一种新颖的分层两级top-k算法，该算法相比竞争基线实现了高达**10倍**的加速，并利用LLM输出的低熵结构在大词汇表上实现高效选择。在异构投机解码工作负载上，SonicSampler相对于最先进的基线实现了高达**16倍**的采样加速，同时保持了灵活的批处理执行。

## 1 引言

大语言模型（LLM）推理的性能在延迟敏感型应用中日益关键，从实时语音接口到智能体工作流，其中较小的模型必须在严格的时间预算内执行（defossez2024moshi; belcak2025small）。虽然大量系统工作集中在加速模型前向传播上，但LLM生成的有效性也依赖于采样的**效率**和**灵活性**——即将模型logits转换为离散令牌决策的过程。采样不仅仅是一个后处理细节；相反，它控制着输出的多样性、可控性和结构有效性（holtzman2020curious; zhu2023penalty; willard2023outlines; dong2025xgrammarflexibleefficientstructured）。它还使得诸如best-of-N解码等质量改进策略成为可能，其中抽取多个多样化候选并筛选以提升最终输出质量（wang2022self; chen2023universal; li2025selfmoa; wang2025effect）。因此，改进LLM服务不仅需要更快的模型执行，还需要一个在现实部署约束下既高效又富有表现力的采样流水线。

这种采样瓶颈在投机解码中尤为突出（leviathan2023fast; chen2023accelerating），其中轻量级的草稿模型提出候选续写，这些候选必须根据目标分布进行采样和验证。随着草稿模型变得越来越轻量级，例如Medusa头（cai2024medusa）、EAGLE（li2024eagle）和多令牌预测（gloeckle2024mtp），下游采样流水线的相对成本也相应增长。在实践中，该流水线并非单一原语，而是操作的异构组合，包括语法约束掩码、重复惩罚、频率和存在惩罚、logits偏置、温度缩放、top-k/ top-p/ min-p过滤、随机或贪婪令牌选择以及投机验证（fan2018hierarchical; holtzman2020curious; nguyen2024minp; hewitt2022truncation）。

现有的内核实现无法充分解决这种复杂性。有的加速了孤立组件，如语法掩码、惩罚应用或概率过滤，但将整体流水线碎片化为多次内核启动，引入了额外的开销和中间内存流量。另一些虽然部分融合了采样，但假设批次内行为同质，因此无法支持连续批处理（yu2022orca）下出现的混合贪婪和随机配置。关键的是，许多top-k实现也不兼容CUDA Graph，丧失了在生产服务引擎中日益重要的启动开销摊销（yu2022orca; kwon2023vllm; nvidia2023tensorrtllm; zheng2024sglang）。

在这项工作中，我们提出SonicSampler，一套统一的分片感知Triton（tillet2019triton）内核，该内核将完整的采样流水线（从logits处理到投机验证）垂直融合到一个兼容CUDA Graph的执行模型中。我们的内核支持**单步模式**（标准和草稿模型推理）和**多步模式**（验证），在批处理调度内通过紧凑的比特级指示器实现完全动态的每请求配置。

我们方法的核心是对top-k瓶颈的重构。虽然过滤名义上需要对词汇表（$2^{17}$–$2^{18}$个令牌）进行全局归约，但实际的下一个令牌分布高度集中；因此，与采样相关的有效支持集远小于完整词汇表。这激发了一个有界的候选集，其中$k=128$，只要保留集合之外的概率质量对于所选截断规则可忽略不计则足够。我们在第4.5节对此进行了经验验证，表明top-128保留了几乎所有相关质量并保持了下游准确性。利用这一特性，我们引入了一个**两级分层top-k算法**。第1阶段对词汇表块进行分片，在将每个块归约为$k$个候选之前应用完整的logits处理前奏，通过自适应基数或双调选择，并在保持单调性和稳定性的字典序方案下编码。第2阶段跨块收集并合并候选，形成最终的top-k。这种map-reduce公式通过将计算密集型的前奏和后记与归约阶段融合，最大化算术强度。因此，我们将贡献总结如下：

- • 我们提出了一套统一的CUDA Graph兼容内核套件，融合了logits处理、采样和投机验证，同时支持单个批处理工作负载内的混合贪婪和随机解码。
- • 我们引入了一种两级分片感知分层top-k算法，相对于现有的基于基数和双调的替代方案，实现了高达**10倍**的加速。
- • 我们证明了在投机解码工作负载上，相对于FlashInfer和其他最先进基线，采样加速高达**16倍**。

## 2 相关工作

现有的采样系统可以沿三个轴进行分类：内核融合、异构采样支持以及CUDA Graph兼容性。虽然先前的工作在这个设计空间中探索了不同的点，但没有一个能同时满足所有三个要求。我们在此提供对代表性方法的重点讨论，并将更全面的回顾推迟到附录A，同时在附录B中提供功能比较。

#### 独立采样内核。
内核库如FlashInfer（flashinfer2024）和XGrammar（dong2025xgrammarflexibleefficientstructured）提供了采样流水线中关键组件的高效实现，包括top-k/top-p选择和语法约束解码。这些工作为生产系统建立了强大的构建块，尽管它们作为独立内核运行时组合在一起。

#### 融合采样内核。
一些近期方法（modular2025max）引入了单内核内的部分采样融合。然而，这些设计假设采样配置同质，无法在批处理工作负载内原生支持混合贪婪和随机解码，限制了它们在生产服务中的适用性。其他并行的努力，如MegaKernel（cheng2025miragepersistentkernelcompiler），将模型前向传播融合到单个执行单元中，减少了内核启动开销，它们以模型执行为目标而非采样，因此是互补的。FlashSampling（ruiz2026flashsampling）通过将LM头与下游采样融合来减少内存流量，针对的是后logits处理互补的推理阶段。它还简要指出了未来工作中的两级top-k分解。我们的工作则侧重于系统设计，以在实践中有高效实现这种分解。

#### 高效top-k选择。
另一条独立的工作线专注于加速top-k选择本身，涵盖多遍基数选择（radik2024; wang2025tilelangcomposabletiledprogramming）、基于双调的流式内核（tillet2019triton）以及混合或分布感知方案（flashinfer2024; park2026qritahighperformancetopktopp）。这些方法将选择视为孤立原语，并未考虑其与上游logits变换和下游采样或验证步骤的紧密耦合，错失了垂直融合的机会。

综合来看，这些工作为LLM推理和采样提供了高效的原语和部分融合策略。SonicSampler在这些洞察的基础上，将logits处理、采样和验证统一到一个兼容CUDA Graph的执行模型中，该模型支持异构批处理工作负载。

## 3 SonicSampler

本节详细介绍SonicSampler。我们的设计遵循三个原则：（1）利用采样工作负载的固有结构以最大化融合机会，（2）为算法简化提供理论基础，（3）在单个批处理调度内支持生产服务的完整异构性。

### 3.1 背景

为了奠定SonicSampler设计的基础，我们首先介绍LLM采样流水线及其组成操作。这些操作定义了我们的执行模型旨在优化的计算结构。

LLM采样通过一系列logits处理、截断、采样和（可选）投机验证步骤，将模型logits转换为离散令牌。给定logits $\mathbf{x} \in \mathbb{R}^V$，这些操作产生概率分布并选择下一个令牌。

#### Logits处理。
我们将logits处理视为应用于原始logits的一系列变换，包括语法掩码、重复和频率惩罚、logits偏置以及温度缩放（keskar2019ctrl; dong2025xgrammarflexibleefficientstructured; hinton2015distilling）。这些操作独立作用于每个令牌，因此形成一种**映射风格**的计算，在词汇表上完全可并行化。

#### 概率截断。
截断方法如top-k、top-p和min-p选择候选令牌的子集（fan2018hierarchical; holtzman2020curious; nguyen2024minp）。与logits处理相比，这些操作依赖于跨令牌的比较（例如，排序或累积概率），因此需要对词汇表进行**全局归约**。值得注意的是，我们在附录中的等式12中展示了min-p可以重新表述以避免显式softmax。

#### Top-k选择原语。
GPU上基于排序的截断通常通过通用top-k选择原语实现。*基于基数的选择*通过从最高到最低有效位逐组划分键，在每次遍历中递归到包含第$k$个元素的桶（alabi2012fast）。*基于双调的选择*建立在双调排序网络之上（batcher1968sorting），其规则、无分支的结构自然地映射到SIMT执行，特别适用于分片局部、固定大小的归约（satish2009designing; shanbhag2018efficient）。

#### 采样。
令牌选择可以通过Gumbel-Max重参数化（jang2016categorical）表示为对扰动logits的$\arg\max$，将随机和贪婪采样统一到单一公式：$x^* = \arg\max_i (x_i + \epsilon_i)$，其中$\epsilon_i$引入随机性。

#### 投机验证。
投机解码通过让草稿模型提出候选令牌来加速推理，当$u \cdot q(d) \leq p(d)$且$u \sim \mathcal{U}(0,1)$时接受候选，或者在拒绝情况下从残差分布$\tilde{p}(x) \propto \max\{0, p(x) - q(x)\}$中采样（leviathan2023fast）。

这些组件表现出一种特征结构，即逐元素变换后跟全局归约，这激发了接下来介绍的工作负载感知执行模型。详细公式化推迟到附录C。

图1：SonicSampler针对单步（非投机）工作负载的两级分片感知融合采样流水线。**第1阶段**（顶部，从左到右）将语法掩码、重复惩罚、logits偏置和温度缩放与字典序编码及通过优化双调或自适应基数-双调选择（右侧面板）进行的每分片局部Top-k垂直融合，仅将每分片的$k$个打包结果写入HBM暂存区作为跨内核的唯一内存边界。**第2阶段**（底部，从右到左）跨分片重新索引，通过双调归并网络进行全局Top-k归约（采用升序/降序交替分片布局），应用概率截断，并通过Gumbel噪声和ArgMax发射所选令牌，完成通过共享HBM频带的顺时针数据循环。自适应基数-双调策略（右侧）通过单遍优化归约树筛选最显著的6位，分支到低10位的2遍基数选择，或回退到优化双调选择。

### 3.2 工作负载感知执行模型

现有实现将每个采样阶段分派为单独的内核，引入了启动开销和中间物化，而top-k归约本身作为对完整词汇表的单程序流式传递执行，从而将占用率限制为$B$个并发程序并排除分片执行。我们将全局归约分解为对词汇表分片（大小为$B_N$）的**两级分层**操作：

1.  **第1阶段**启动$B \cdot Z_v$个程序（其中$Z_v = \lceil V / B_N \rceil$），每个程序将logits处理前奏与分片局部top-k归约融合，将$k$个打包候选发射到全局内存中的暂存区。
2.  **第2阶段**启动$B$个程序，每个程序收集$Z_v \cdot k$个打包候选，执行跨分片合并，然后进行垂直融合的后记，包括概率截断、Gumbel扰动和$\arg\max$选择。

与单阶段公式相比，这种分解将第1阶段的并行度提高了$Z_v$倍，同时将所有分片间通信限制在紧凑的$B \cdot Z_v \cdot k$个打包`uint32`条目的暂存区内。融合边界的选取确保了在两个阶段之间，中间logits或概率向量永远不会以词汇表规模被物化。

### 3.3 有界Top-k的熵充分性

两级分解需要限制从每个分片保留的候选数量。我们确定对于经过良好训练的LLM，一个适度的界$k=128$足以在实际top-p阈值下捕获有效概率质量。

#### 理论分析。
考虑一个在$V$个令牌上的概率分布，其中质量$P$集中在$K$个候选（top-k集）中，质量$(1-P)$分布