NVIDIA TensorRT-LLM Now Supports Recurrent Drafting for Optimizing LLM Inference
Recurrent drafting
(referred as ReDrafter) is a novel speculative decoding technique developed and open-sourced by Apple for large language model (LLM) inference now available with
NVIDIA TensorRT-LLM
. ReDrafter helps developers significantly boost LLM workload performance on NVIDIA GPUs.
NVIDIA TensorRT-LLM
is a library for optimizing LLM inference. It provides an easy-to-use Python API to define LLMs and build
NVIDIA TensorRT
engines that contain state-of-the-art optimizations to perform inference efficiently on NVIDIA GPUs. Optimizations include custom attention kernels, inflight batching, paged KV caching, quantization (FP8, INT4 AWQ, INT8 SmoothQuant), and much more.
Speculative decoding
is a technique that accelerates LLM inference by generating multiple tokens in parallel. It uses smaller “draft” modules to predict future tokens, which are then verified by the main model. This method maintains output quality while significantly reducing response times, especially during low traffic periods, by better utilizing available resources for low-latency inference.
ReDrafter employs recurrent neural network (RNN)-based sampling, referred to as
drafting
, combined with tree-style attention previously used in other techniques like Medusa to predict and verify draft tokens from multiple possible paths for better accuracy and to potentially accept more than one token in each iteration of the decoder. NVIDIA collaborated with Apple to add support for this technique in TensorRT-LLM, making it accessible to the broader developer community.
The integration of ReDrafter into TensorRT-LLM expanded its reach, unlocked new optimization potential, and improved on previous methods such as Medusa. For Medusa, the path acceptance and token sampling happens in the TensorRT-LLM runtime, introducing some overhead inside the engine for processing all possible future paths without knowing the accepted path, most of which are ultimately discarded. To reduce such overhead, ReDrafter requires the token validation and acceptance of the best path before drafting future tokens for the next iteration.
TensorRT-LLM has been updated to incorporate drafting and validation logic inside a single engine, rather than relying on the runtime or separate engines to further minimize overhead. This approach provides TensorRT-LLM kernel selection and scheduling more freedom to optimize the network for maximum performance.
To better illustrate ReDrafter improvements, Figure 1 highlights the key differences between its implementation and that of Medusa in TensorRT-LLM. Most of the components related to speculative decoding have been done in-engine for ReDrafter. This significantly simplifies the runtime changes needed for ReDrafter.
Figure 1. Comparison of Medusa (left) and ReDrafter (right) implementations in NVIDIA TensorRT-LLM
The following sections delve into some of the changes that help enable ReDrafter in TensorRT-LLM.
Inflight-batching compatible engine
Inflight-batching (IFB) is a strategy that significantly improves the throughput by batching context-phase and generation-phase requests. Speculative decoding, coupled with IFB, introduces more complexity to the pipeline as context-phase requests need to be handled differently than generation-phase requests, which require draft token validation. Since ReDrafter moves the validation logic inside the model definition, the engine needs that logic as well during validation. Similar to the attention plugin, the batch is split into two smaller batches: one for context requests and another for generation requests. Each smaller batch then enters its computational workflow, and at the end they are combined back to a single batch for drafting.
Figure 2. ReDrafter’s computational workflow for inflight-batching compatible TensorRT-LLM engine
Note that this approach requires that all operators on either path support empty tensors, which could happen if a batch consists of all context requests or all generation requests. This capability adds flexibility to TensorRT-LLM APIs, enabling the definition of more complicated models in the future.
Implementing in-engine validation and drafting
To validate and draft inside the engine, TensorRT-LLM is updated with support for numerous new operations so that PyTorch code can be easily translated into a definition of the TensorRT-LLM model.
The following PyTorch code excerpt is Apple’s
PyTorch implementation of ReDrafter
. The TensorRT-LLM implementation is almost a straightforward line-by-line mapping of the PyTorch version.
PyTorch
def unpack(
    packed_tensor: torch.Tensor,
    unpacker: torch.Tensor,
) -> torch.Tensor:
    assert len(packed_tensor.shape) == 3
    last_dim_size = packed_tensor.shape[2]
    batch_size, beam_width, beam_length = unpacker.shape
    unpacked_data_indices = unpacker.view(
        batch_size, beam_width * beam_length, 1).expand(
        -1, -1, last_dim_size
    )
    unpacked_tensor = torch.gather(
        packed_tensor, 1, unpacked_data_indices).reshape(
        batch_size, beam_width, beam_length, -1
    )
    return unpacked_tensor
TensorRT-LLM
def _unpack_beams(
    x: Tensor,
    indices: Tensor,
    num_beams: int,
    beam_length: int
) -> Tensor:
    assert x.rank() == 3
    d0 = shape(x, 0, INT_DTYPE_STR)
    dl = shape(x, -1, INT_DTYPE_STR)
    indices = view(
        indices, [-1, num_beams * beam_length, 1], False)
    res_shape = concat([d0, num_beams, beam_length, dl])
    res = view(gather_nd(x, indices), res_shape, False)
    return res
This, of course, is a very simple example. For a more complex example, see the
beam search implementation
. With the new functionalities added for ReDrafter, it might be possible to improve the Medusa implementation in TensorRT-LLM to further increase its performance.
ReDrafter performance in TensorRT-LLM
As
benchmarked by Apple
, ReDrafter with TensorRT-LLM can provide up to 2.7x throughput improvements on
NVIDIA H100
GPUs with TP8 over the base LLM.
Note that the performance improvement of any speculative decoding technique can be heavily impacted by many factors, including:
GPU utilization:
Speculative decoding is commonly used for low-traffic scenarios, where GPU resources are typically underutilized due to small batch sizes.
Average acceptance rate:
The latency of each decoding step is increased since speculative decoding must perform extra computation, where a significant portion of it is ultimately wasted after validation. As a result, to see any performance benefits from speculative decoding, the average acceptance rate must be high enough to pay for that extra latency. This is affected by the number of beams, their lengths, and the quality of the beam search itself (which is impacted by the training data).
Task:
It is easier to predict future tokens for some tasks (code completion, for example), which leads to a higher acceptance rate, and thus improved performance.
Summary
This collaboration between NVIDIA and Apple, has made TensorRT-LLM more powerful and more flexible, enabling the LLM community to innovate more sophisticated models and easily deploy them with TensorRT-LLM to achieve unparalleled performance on NVIDIA GPUs. These new features open exciting possibilities, and we eagerly anticipate the next generation of advanced models from the community that leverage TensorRT-LLM capabilities, driving further improvements in LLM workloads.
Explore
NVIDIA TensorRT-LLM
to unlock the full potential of your models on NVIDIA GPUs.
NVIDIA TensorRT-LLM 现支持 Recurrent Drafting，实现 LLM 推理优化
Recurrent Drafting (简称 ReDrafter) 是苹果公司为大语言模型 (LLM) 推理开发并开源的一种新型推测解码技术，该技术现在可与
NVIDIA TensorRT-LLM
一起使用。ReDrafter 帮助开发者大幅提升了 NVIDIA GPU 上的 LLM 工作负载性能。
NVIDIA TensorRT-LLM
是一个 LLM 推理优化库，提供了一个易于使用的 Python API 来定义 LLM 和构建
NVIDIA TensorRT
引擎，这些引擎具有顶尖的优化功能，可在 GPU 上高效执行推理。优化功能包括自定义 Attention Kernel、Inflight Batching、Paged KV Caching、量化技术 (FP8、INT4 AWQ、INT8 SmoothQuant) 等。
推测解码 (Speculative decoding)
是一种通过并行生成多个 token 来加速 LLM 推理的技术。它使用较小的“draft”模块预测未来的 token，然后由主模型进行验证。该方法通过更好地利用可用资源实现低延迟推理，在保持输出质量的同时大大缩短了响应时间，尤其是在低流量时段。
ReDrafter 运用基于循环神经网络 (RNN) 的采样 (称为
Drafting
) 并结合之前在 Medusa 等其他技术中使用的树状注意力，预测和验证来自多个可能路径的 draft token 以提高准确性，并在解码器的每次迭代中接受一个以上 token。NVIDIA 与苹果公司合作，在 TensorRT-LLM 中添加了对该技术的支持，使更加广泛的开发者社区能够使用该技术。
ReDrafter 与 TensorRT-LLM 的集成扩大了该技术的覆盖范围，解锁了新的优化潜力，并改进了 Medusa 等先前的方法。Medusa 的路径接受和 token 采样发生在 TensorRT-LLM 运行时，需要在接受路径未知的情况下处理所有可能的未来路径，而且其中大部分路径最终都会被丢弃，这就给引擎内部带来了一些开销。为了减少这种开销，ReDrafter 要求在 drafting 下一次迭代的未来 token 之前，先验证 token 并接受最佳路径。
为了进一步减少开销，TensorRT-LLM 更新后在单个引擎中整合了 drafting 和验证逻辑，不再依赖运行时或单独的引擎。这种方法为 TensorRT-LLM 内核选择和调度提供了更大的自由度，通过优化网络实现了性能的最大化。
为了更好地说明 ReDrafter 的改进，图 1 展示了 TensorRT-LLM 中 ReDrafter 实现与 Medusa 实现的主要区别。大多数与推测解码相关的组件都在 ReDrafter 的引擎内完成，这大大简化了 ReDrafter 所需的运行时更改。
图 1. NVIDIA TensorRT-LLM 中 Medusa（左）和 ReDrafter（右）实现的比较
下面将深入探讨有助于在 TensorRT-LLM 中启用 ReDrafter 的一些变化。
兼容 Inflight-batching 批处理的引擎
Inflight-batching (IFB) 是一种通过批量处理上下文阶段和生成阶段请求，来显著提高吞吐量的策略。鉴于上下文阶段请求与生成阶段请求的处理方式不同（生成阶段请求需要 draft token 验证），因此结合 IFB 的推测解码会给管线带来更大的复杂性。ReDrafter 将验证逻辑移至模型定义内部，因此引擎在验证过程中也需要该逻辑。与注意力插件类似，该批处理被分成两个较小的批处理：一个用于上下文请求，另一个用于生成请求。然后，每个较小的批处理进入计算工作流，最后再合并成一个批处理进行 drafting 流程。
图 2. ReDrafter 兼容 TensorRT-LLM 引擎的 Inflight-batching 批处理计算工作流
请注意，这种方法要求任一路径上的所有运算符都支持空张量。如果一个批处理由所有上下文请求或所有生成请求组成，就可能出现空张量。该功能增加了 TensorRT-LLM API 的灵活性，使未来定义更复杂的模型成为可能。
实现引擎内验证和
Drafting
为了在引擎内进行验证和 draft，TensorRT-LLM 更新时加入了对许多新操作的支持，这样 PyTorch 代码就可以轻松地转化成一个 TensorRT-LLM 模型的定义。
以下 PyTorch 代码摘录是苹果公司的
PyTorch 实现的 ReDrafter
。TensorRT-LLM 实现几乎就是 PyTorch 版本的直接逐行映射。
PyTorch
def unpack(
    packed_tensor: torch.Tensor,
    unpacker: torch.Tensor,
) -> torch.Tensor:
    assert len(packed_tensor.shape) == 3
    last_dim_size = packed_tensor.shape[2]
    batch_size, beam_width, beam_length = unpacker.shape
    unpacked_data_indices = unpacker.view(
        batch_size, beam_width * beam_length, 1).expand(
        -1, -1, last_dim_size
    )
    unpacked_tensor = torch.gather(
        packed_tensor, 1, unpacked_data_indices).reshape(
        batch_size, beam_width, beam_length, -1
    )
    return unpacked_tensor
TensorRT-LLM
def _unpack_beams(
    x: Tensor,
    indices: Tensor,
    num_beams: int,
    beam_length: int
) -> Tensor:
    assert x.rank() == 3
    d0 = shape(x, 0, INT_DTYPE_STR)
    dl = shape(x, -1, INT_DTYPE_STR)
    indices = view(
        indices, [-1, num_beams * beam_length, 1], False)
    res_shape = concat([d0, num_beams, beam_length, dl])
    res = view(gather_nd(x, indices), res_shape, False)
    return res
当然，这只是一个非常简单的例子。如要了解更复杂的示例，请参见
束搜索实现
。借助为 ReDrafter 添加的新功能，就可以改进 TensorRT-LLM 中的 Medusa 实现，从而进一步提高其性能。
ReDrafter 在 TensorRT-LLM 中的性能
根据
苹果公司的基准测试
，在采用 TP8（Tensor Parallelism with 8 GPUs，8 卡 GPU 张量并行） 的
NVIDIA
GPU 上使用 TensorRT-LLM 的 ReDrafter 最多可将吞吐量提高至基础 LLM 的 2.7 倍。
请注意，任何推测解码技术的性能提升幅度都会受到诸多因素的大幅影响，包括：
GPU 利用率：
推测解码通常用于低流量场景，由于批量较小，GPU 资源的利用率通常较低。
平均接受率：
由于推测解码必须执行额外的计算，而其中很大一部分计算最终会在验证后被浪费，因此每个解码步骤的延迟都会增加。所以要想通过推测解码获得任何性能上的优势，平均接受率必须高到足以弥补增加的延迟。这受到束数量、束长度和束搜索本身质量（受训练数据影响）的影响。
任务：
在某些任务（例如代码完成）中预测未来的 token 更容易，使得接受率更高，性能也会因此而提升。
总结
NVIDIA 与苹果公司的合作让 TensorRT-LLM 变得更加强大和灵活，使 LLM 社区能够创造出更加复杂的模型并通过 TensorRT-LLM 轻松部署，从而在 NVIDIA GPU 上实现无与伦比的性能。这些新特性带来了令人兴奋的可能性，我们热切期待着社区使用 TensorRT-LLM 功能开发出新一代先进模型，进一步改进 LLM 工作负载。
探索
NVIDIA TensorRT-LLM
，在 NVIDIA GPU 上充分释放模型潜能。