本文介绍了自回归模型在推理性能方面的瓶颈，以及为解决这一问题而兴起的推测解码技术。苹果提出的ReDrafter方法结合了草稿模型和动态树注意力机制，通过递归设计和束搜索提高预测准确性和效率。ReDrafter与英伟达TensorRT-LLM推理加速框架整合，在GPU上实现更高效的LLM推理性能。文章还介绍了ReDrafter的具体技术细节，包括其基于RNN的草稿模型、动态树注意力算法、知识蒸馏训练方法等，并在不同硬件平台上展示了其性能优势。
自回归模型在生成文本时需要反复执行前向传播，导致大量的计算资源和频繁的内存访问，从而引起较高的延迟。推测解码技术通过使用草稿模型预测未来token序列，再由主模型验证，实现并行化生成，提高推理效率。
ReDrafter结合了循环神经网络（RNN）的草稿模型和动态树注意力机制，使用束搜索探索多个可能的序列延续。其创新之处在于使用递归设计和动态树注意力算法提高预测准确性和效率。
苹果与英伟达展开深度技术合作，将ReDrafter整合到英伟达TensorRT-LLM推理加速框架中。通过提供标准化的接口和优化的实现，使得开发者能够更容易地部署复杂的语言模型。这一合作显著增强了TensorRT-LLM框架的能力。
ReDrafter在MT-Bench基准测试中实现了显著的性能提升。在每个生成步骤中，它可接受多个token，提高了并行度。在大规模批处理场景下，其吞吐量可达到每秒1636个token。此外，在不同的硬件平台上，ReDrafter都展现出良好的适应性。
众所周知，对于自回归模型而言，内存带宽一直是制约推理性能的关键瓶颈。当模型生成文本时，它需要反复执行前向传播来预测每个 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 与自回归相比，每秒令牌速度加快。（来源：苹果）
这种性能改进将直接转化为实际价值：它能最大限度地减少用户延迟并减少所需的硬件数量。英伟达在其博客中强调，这次合作显著增强了 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/
排版：刘雅坤
