摘要：众所周知，对于自回归模型而言，内存带宽一直是制约推理性能的关键瓶颈。当模型生成文本时，它需要反复执行前向传播来预测每个 token，这个过程不仅需要大量的计算资源，更重要的是会频繁访问内存来获取模型权重和中间状态。这种内存密集型的特性导致了较高的延迟，影响用户

众所周知，对于自回归模型而言，内存带宽一直是制约推理性能的关键瓶颈。当模型生成文本时，它需要反复执行前向传播来预测每个 token，这个过程不仅需要大量的计算资源，更重要的是会频繁访问内存来获取模型权重和中间状态。这种内存密集型的特性导致了较高的延迟，影响用户体验。

去年，学界兴起了一种称为推测解码（Speculative Decoding）的技术用以缓解这一问题。这种方法使用较小的“草稿”模型来预测可能的未来 token 序列，再由主模型验证，从而实现并行化生成，提高推理效率。

图丨自回归解码和推测解码（来源：arXiv）

今年初，苹果提出了一种新的推测解码方法 Recurrent Drafter （ReDrafter），对原有方法进行了改进。通过将 RNN 草稿模型与动态树注意力机制相结合，ReDrafter 在每个生成步骤中最多可接受 3.5 个 token，超越了此前方法的水平。

最近，苹果又宣布与英伟达展开深度技术合作，将其创新的 ReDrafter 推测解码技术整合到英伟达 TensorRT-LLM 推理加速框架中。这项合作将帮助开发者在英伟达 GPU 上实现更高效的 LLM 推理性能。

ReDrafter 的核心是一个基于循环神经网络（RNN）的草稿模型。与先前的方法（如 Medusa）类似，它使用 Transformer 最后一层的输出作为输入。创新之处在于，它还将历史 token 的嵌入作为循环输入提供给草稿头。

当 LLM 生成一个 token（比如“quiet”）时，草稿模型会获取这个 token 的嵌入向量 et 和 LLM 最后一层的输出 h，然后通过 RNN 更新隐藏状态：st = f(Ust-1 + W et + b)。更新后的状态 gt = [st, h] 被用来预测下一个 token。这种递归设计让模型能够维护动态的上下文表示，大大提升了预测准确性。

图丨草稿模型以 LLM 的最后隐藏状态 h 作为输入来预测接下来的几个标记。（来源：arXiv）

另一方面，ReDrafter 使用束搜索来探索多个可能的序列延续。为了提高效率，它引入了动态树注意力算法来处理候选序列中的共同前缀。

来看一个具体例子：当生成多个候选序列时，比如“morning sipping coffee and watching”和“morning sipping coffee on her”，这些序列往往共享相同的前缀。传统方法会对每个序列独立计算，而 ReDrafter 通过动态构建前缀树来合并重复计算。系统使用张量运算实现了高效的前缀检测和合并处理，在实践中可以减少 30%-60% 的计算量。

此外，为了提高草稿模型的预测质量，ReDrafter 采用了一种新的知识蒸馏训练方法。系统会直接学习 LLM 的 token 生成分布，而不是简单地拟合真实序列。训练目标被定义为最小化 LLM 和草稿模型输出分布之间的 KL 散度：min(pdraft) KL(pllm(y1:T)|pdraft(y1:T))。

在每个训练位置 t，系统从 LLM 采样未来 T 个 token，并优化经验损失：min(pdraft) Σt −log pdraft(ŷt+1:t+T|y1:t)。这种方法确保了草稿模型能够更好地模仿 LLM 的行为。

研究结果显示，在 MT-Bench 基准测试中，ReDrafter 在 Vicuna 7B 和 13B 模型上实现了 2.8 倍的加速。且 ReDrafter 在每个生成步骤平均可接受 4.2 个 token，这意味着相比传统的每步生成一个 token 的方式，ReDrafter 显著提升了并行度。在大规模批处理场景下，系统的最高吞吐量可达到每秒 1636 个 token。

并且，在不同的硬件平台上，ReDrafter 都展现出良好的适应性。在搭载 M2 Ultra 的系统上，13B 参数量的模型达到了 1.94 倍的加速比。即使在计算资源相对有限的 M1 Max 上，系统仍然实现了 1.37 倍的性能提升。

为了让 ReDrafter 的性能优势更好地在 TensorRT-LLM 上发挥，英伟达的工程团队进行了一系列创新。其最重要的改进是将验证逻辑直接集成到引擎内部，而不是依赖运行时处理。这种设计不仅降低了系统开销，还为内核选择和调度优化提供了更大的自由度。

在批处理方面，研究团队采用了一种创新的 inflight-batching（IFB）策略。当处理混合的请求流时，系统会智能地将上下文阶段和生成阶段的请求分开处理，每类请求使用专门优化的计算流水线。这种设计在保持低延迟的同时，显著提升了 GPU 利用率。

图丨 ReDrafter 的 TensorRT-LLM 引擎的在线批处理兼容计算工作流程（来源：英伟达）

尤其值得一提的是系统对前缀树计算的优化。英伟达开发了专门的算子来处理动态树结构，这些算子经过深度优化，能充分利用 Tensor Core 等硬件特性。同时，通过精心设计的内存管理策略，系统能够高效处理不同长度的序列，减少内存碎片。

在实际部署测试中，研究团队使用一个拥有数百亿参数的生成模型进行了全面评估。在贪婪解码模式下，系统实现了 2.7 倍的生成速度提升。这种性能改进不仅降低了用户感知延迟，还减少了所需的 GPU 数量和能源消耗。

这种性能改进将直接转化为实际价值：它能最大限度地减少用户延迟并减少所需的硬件数量。英伟达在其博客中强调，这次合作显著增强了 TensorRT-LLM 框架的能力。通过提供标准化的接口和优化的实现，使得开发者能够更容易地部署复杂的语言模型。

ReDrafter 此前已在 GitHub 开源，项目地址：https://github.com/apple/ml-recurrent-drafter。

参考资料：

1.https://machinelearning.apple.com/research/redrafter-nvidia-tensorrt-llm

2.https://arxiv.org/html/2403.09919v5

3.https://developer.nvidia.com/blog/nvidia-tensorrt-llm-now-supports-recurrent-drafting-for-optimizing-llm-inference/

排版：刘雅坤

来源：DeepTech深科技