如果你曾盯着闪烁的光标等待大语言模型 (LLM) 完成一个段落，你就已经体会到了自回归生成固有的瓶颈。这些模型一次生成一个 token (词元) ，而对于每一个 token，模型都必须将海量的参数从内存搬运到计算单元。

在 GPU 加速的世界里，这被称为“内存带宽受限 (memory bandwidth bound) ”。生成的速度并非取决于 GPU 做数学运算有多快，而是取决于它移动数据的速度有多快。

针对这个问题，一个流行的解决方案是投机采样 (Speculative Sampling) 。 这项技术使用一个较小的“草稿 (draft) ”模型来猜测未来的 token，然后由较大的“目标 (target) ”模型并行地验证这些猜测。虽然这减少了所需的前向传播次数，但验证步骤本身涉及复杂的概率计算，仍然可能限制性能。

在这篇深度文章中，我们将探讨一篇近期的研究论文 “Optimized Speculative Sampling for GPU Hardware Accelerators” (针对 GPU 硬件加速器优化的投机采样) , 该论文提出了两种优化这一验证阶段的新颖方法。我们将了解重新分配 GPU 工作负载如何产生更快——但在数学上完全相同——的结果，以及如何用 Sigmoid 函数近似概率分布，从而以微小的精度损失换取巨大的速度提升。

背景: 为什么采样很慢？

要理解解决方案，我们首先需要了解硬件限制。现代 GPU，如 NVIDIA A100，拥有分层的内存结构:

  1. HBM (高带宽内存) : 容量巨大 (例如 80GB) ，但访问速度相对较慢。这是模型权重存放的地方。
  2. SRAM (静态随机存取内存) : 容量极小 (每个流式多处理器只有几百 KB) ，但速度极快。这是片上内存，通常用作可编程缓存或“共享内存”。

在标准的自回归解码中，GPU 会花费大量时间在 HBM 和 SRAM 之间读写中间矩阵。投机采样试图通过提前预测 \(K\) 个 token 来缓解这一问题。我们不再按顺序运行大模型 \(K\) 次，而是运行一次来并行验证所有 \(K\) 个 token。

然而，本文的作者注意到，验证逻辑本身——即用于接受或拒绝这些草稿 token 的数学运算——在标准实现中并非最优。它涉及了太多针对 HBM 的读写操作。

背景知识: 投机采样的数学原理

在优化硬件之前，让我们先看看我们要计算的数学公式。

在投机采样中，我们有一个目标分布 \(p(x)\) (大模型) 和一个草稿分布 \(q(x)\) (小模型) 。当草稿模型猜测一个 token \(x_{i+c}\) 时，我们要根据接受概率 \(\tau_c\) 来决定是否保留它。

标准的拒绝标准定义为:

公式 1: 接受概率公式。

这里，\(r_c\) 是从均匀分布 \(U(0,1)\) 中抽取的随机数。如果 \(r_c \le \tau_c\)，则接受该 token。

如果 token 被拒绝，我们必须从一个校正后的分布中重新采样一个新的 token，以确保最终输出严格遵循目标模型的分布。这个校正分布定义为:

公式 2: 重采样分布公式。

函数 max_norm 对分布差值进行归一化，确保概率为正且总和为 1:

公式 3: max_norm 函数细节。

计算这些值需要对整个词表 (通常包含 30,000+ 个 token) 进行逐元素操作。标准方法是从 HBM 加载 \(p\) 和 \(q\) 的概率，计算差值，写回 HBM，再次读取进行归一化，依此类推。这种“内存流量”正是研究人员试图消除的。

方法 1: 通过内核融合进行精确优化

该论文的第一个贡献是精确优化 (Exact Optimization) 。 “精确”意味着数学输出与基准线是位级一致的 (bit-wise identical，或在浮点误差范围内) 。它产生完全相同的文本生成结果，但速度更快。

策略: 分块 (Tiling) 与 SRAM

核心洞察在于，对于词表的各个块 (chunk) ，\(\tau\) (接受率) 和重采样分布的计算在很大程度上是独立的。研究人员没有全局处理整个词表，而是将工作负载分解为多个“分块 (tiles) ”。

他们编写了一个自定义的 GPU 内核 (Kernel) ，将工作负载重新分配给 GPU 线程和线程块。过程如下:

  1. 加载到 SRAM: 一个线程块将一部分 \(p(x)\) 和 \(q(x)\) 的概率从慢速 HBM 加载到快速 SRAM 中。
  2. 并行计算: 线程完全在 SRAM 内计算接受概率 \(\tau\) 和重采样方程 (公式 3) 的分子。
  3. 并行归约: 公式 3 中的分母 (所有概率之和) 是唯一棘手的部分，因为它需要全局求和。内核在 SRAM 内的块中执行部分求和。
  4. 写回: 只有必要的部分结果才会被写回 HBM。

通过将中间值保留在 SRAM 中，GPU 避免了多次昂贵的 HBM 访问。

图 1: 精确优化流程的可视化。注意数据是如何从 HBM (左) 移动到 SRAM (中) 并在那里进行主要处理的。概率分布 \\(p\\) 和 \\(q\\) 被切分为子词表 (\\(V_k\\)) 。接受率 (\\(\\tau\\)) 和重采样分子 (\\(a_k\\)) 的计算并发进行。只有最终聚合的结果会被送回 HBM。

图 1: 我们的精确优化方法概览。我们利用快速 SRAM 并行计算投机采样所需的大部分结果，并在其中读写中间结果。我们最大化每个块的线程数，以便在不耗尽可用片上内存的情况下，对尽可能多的元素进行并行计算。

这种方法最大化了并行性。通过对内核进行分块，GPU 可以饱和其计算核心，而无需等待内存读取。因为它忠实地实现了标准方程 (1、2 和 3) ，所以精度没有任何损失 。

方法 2: 通过 Sigmoid 进行近似优化

研究人员并没有止步于“精确”。他们问: 如果我们用极小的精度代价换取巨大的速度提升会怎样？

精确方法——实际上也是几乎所有 Transformer 操作——的瓶颈在于 Softmax 函数。

公式 6: Softmax 公式显示了对全局求和的依赖。

看看上面 Softmax 方程中的分母。要计算哪怕单个 token 的概率，你也需要对词表中每一个 token 的指数进行求和。这创建了一个跨越所有线程块的依赖关系，使其难以完全并行化。

切换到 Sigmoid

作者建议使用逐元素的 Sigmoid 近似来替换 Softmax。与 Softmax 不同，Sigmoid 独立处理每个 Logit。它不关心其他 Logit 在做什么。

为了实现这一点，Logits (原始模型输出，\(z\)) 首先使用常数 \(\alpha\) 和 \(\beta\) 进行缩放，以适应 Sigmoid 函数的有效范围。近似概率 \(\hat{p}\) 和 \(\hat{q}\) 计算如下:

公式 7: p 和 q 的 Sigmoid 近似公式。

因为 \(\hat{p}\) 和 \(\hat{q}\) 是通过 Sigmoid 计算的，严格来说它们不是有效的概率分布 (它们的和不完全等于 1) 。但是，它们保留了 token 的相对排名，这对于投机采样来说通常已经足够了。

使用这些近似概率，接受标准和重采样逻辑更新为:

公式 8: 近似接受概率。

公式 9: 近似重采样分布。

融合内核

由于 Sigmoid 是逐元素的，它可以直接融合到采样内核中。在检查 token 是否被接受之前，不需要等待全局求和来归一化数值。

图 2: Sigmoid 近似的工作流程。红色突出显示的是关键区别: 内核不再接收预先计算好的概率，而是接收原始 Logits (\\(z_p, z_q\\)) ，对其进行缩放，并在 SRAM 中直接应用 Sigmoid。这创建了一个高度并行、非阻塞的流水线。

图 2: Sigmoid 近似在每个线程块内的计算概览。每组 Logits 由最小常数 \(\alpha\) 和最大常数 \(\beta\) 缩放。然后计算 Sigmoid 激活值 \(\sigma\) 并将草稿和目标 Logits 的每个片段存储在 SRAM 中。随后，中间值 \(\hat{f}_k(x)\)，\(\hat{a}_k(x)\)，\(\hat{b}_k\) 和 \(\hat{\tau}_{c_k}(x)\) 的计算类似于图 1。生成的输出随后用于更新 HBM 中的 \(\hat{\tau}_c(x), \hat{a}(x)\) 和 \(\hat{b}\)。

这种方法明显更快，因为它绕过了 Softmax 所需的昂贵全局同步。然而，它是“非精确的”，意味着生成的文本可能与标准模型不同。

实验与结果

为了验证这些方法，研究人员在两个任务上进行了广泛的基准测试:

  1. 自动语音识别 (ASR): 使用 Whisper (Target) 和 Distil-Whisper (Draft)。
  2. 文本摘要: 使用 Llama 2、Qwen 和 Gemma 模型。

他们使用Profiling 时间 (采样函数的原始 GPU 执行时间) 和准确率 (ASR 的词错误率，摘要的 ROUGE 分数) 来衡量性能。

性能提升

结果显示了明显的速度层级。基准线 (Baseline) 最慢，精确优化 (Exact Optimization) 提供了稳健的提升，而 Sigmoid 近似 (Sigmoid Approximation) 则快得惊人。

表 1: 主要结果。请看“\\(\\Delta\\%\\) Profiling Time”列。精确方法 (中右) 始终提供约 6-13% 的加速，且精度完全相同 (WER/ROUGE 分数基本保持不变) 。Sigmoid 方法 (最右) 提供了 37-94% 的巨大加速，尽管有时精度分数略有下降。

例如，在CNN/DM摘要任务中，使用 Llama2 13B 时，精确方法将 Profiling 时间缩短了 10.1% , 而 Sigmoid 方法缩短了 37.2% 。

Sigmoid 方法足够准确吗？

Sigmoid 方法的权衡在准确率指标中显而易见。

  • ASR: 词错误率 (WER) 略有增加 (例如，在 LibriSpeech Clean 上从 0.08 增加到 0.09) 。这是一种轻微的退化。
  • 摘要: ROUGE-1 分数小幅下降 (例如，从 0.31 降至 0.29) 。

虽然是“非精确的”，但模型通常保持连贯。研究人员发现接受率 (draft token 被接受的频率) 仍然相对较高，表明 Sigmoid 近似密切反映了真实分布的决策。

稳定性与内存

对于系统架构师来说，一个重要的问题是这些优化是否会引入不稳定性或内存开销。

研究人员绘制了每个解码步骤的执行时间与初始草稿 token 数量 (\(\gamma\)) 的关系图。

图 3: 执行时间与草稿 token 数量的关系。绿线 (基准线) 始终最高 (最慢) 。橙线 (精确) 较低，蓝线 (Sigmoid) 最低 (最快) 。关键是，随着我们投机性地生成更多 token，线条保持平坦或线性，表明优化具有良好的扩展性。

关于内存，这些优化非常高效。因为它们依赖 SRAM 进行中间存储，所以不会导致 HBM 使用量膨胀。

图 5: Whisper 模型的 HBM 内存峰值使用量。精确 (圆圈) 和基准线 (方块) 几乎完全重叠。Sigmoid 方法 (三角形) 在某些情况下实际上使用的内存略少，因为它避免了存储完整的概率矩阵，转而进行即时计算。

结论与启示

这篇论文强调了现代 AI 系统中的一个关键教训: 算法必须是硬件感知的。

如果仅仅将投机采样的数学运算视为“黑盒”，标准实现就会错失性能提升的机会。通过理解 GPU 的内存层次结构 (HBM 与 SRAM) 以及函数的属性 (Softmax 与 Sigmoid) ，我们可以榨取出显著的收益。

  • “精确”方法对于需要保证保真度的生产系统来说是必然之选。它在采样步骤上提供了约 10% 的免费加速，且没有任何负面影响。
  • “Sigmoid”方法提供了一个更激进的选择。对于可以接受 1-2% 指标性能下降的应用 (例如，延迟至上的实时聊天机器人) ，采样内核中近 2 倍的加速是具有变革意义的。

随着大语言模型不断增长，“内存墙”只会变得更高。像这样优化数据移动而不仅仅是计算的技术，对于让下一代 AI 变得更易用、响应更迅速至关重要。