NVIDIA TensorRT-LLM Now Supports Recurrent Drafting for Optimizing LLM Inference

Key Takeaways:

I. ReDrafter delivers up to 2.7x throughput improvements on NVIDIA H100 GPUs, significantly accelerating LLM inference.

II. The integration of RNN-based sampling and tree-style attention within TensorRT-LLM represents a paradigm shift in speculative decoding, optimizing both speed and accuracy.

III. ReDrafter's enhanced efficiency and scalability democratize access to powerful LLMs, accelerating AI research and broadening industry adoption.

Recurrent drafting (ReDrafter), a novel speculative decoding technique developed and open-sourced by Apple, is now integrated into NVIDIA TensorRT-LLM, promising significant performance boosts for large language model (LLM) inference on NVIDIA GPUs. This collaboration marks a significant advancement in accelerating LLM workloads, leveraging RNN-based sampling and tree-style attention within a highly optimized framework. This article explores the technical intricacies of ReDrafter, its performance benefits, integration within TensorRT-LLM, and the broader implications for the AI landscape.

Benchmarking ReDrafter: Quantifying the Performance Gains in LLM Inference

ReDrafter significantly boosts LLM inference performance, achieving up to 2.7x throughput improvements on NVIDIA H100 GPUs with Tensor Parallelism (TP8) compared to baseline LLMs. These gains are particularly pronounced in low-traffic scenarios where GPU resources are often underutilized. Apple's internal benchmarks, as well as early tests with models like Vicuna, demonstrate ReDrafter's consistent performance advantage across various LLM architectures and sizes.

The effectiveness of speculative decoding techniques like ReDrafter is influenced by factors such as GPU utilization and average acceptance rate. ReDrafter maximizes GPU utilization by efficiently generating and validating multiple draft tokens in parallel. A high average acceptance rate, influenced by the number of beams, their lengths, and the quality of the beam search, is crucial for realizing performance benefits. ReDrafter's RNN-based sampling and tree-style attention contribute to maintaining high acceptance rates, ensuring that the extra computation involved in speculative decoding translates into tangible speedups.

The performance improvements observed with ReDrafter vary depending on the specific LLM, GPU architecture, precision level (e.g., FP8, FP16), and input/output sequence lengths. While the 2.7x speedup represents a peak performance figure, real-world applications may experience different levels of improvement. Further benchmarking across a wider range of LLMs, including GPT-3 and LLaMA 2, is crucial for providing a more comprehensive understanding of ReDrafter's capabilities in diverse deployment scenarios.

ReDrafter's integration with TensorRT-LLM opens up new possibilities for optimizing LLM inference pipelines. The in-engine validation and drafting, coupled with inflight batching compatibility, minimize overhead and maximize throughput. This streamlined approach allows TensorRT-LLM's kernel selection and scheduling algorithms to further optimize the network for peak performance, delivering substantial improvements in overall efficiency.

Decoding ReDrafter: RNN-based Sampling and Tree-Style Attention

ReDrafter employs RNN-based sampling, referred to as drafting, as a core component of its speculative decoding strategy. This involves using a recurrent neural network (RNN) to predict future tokens, which are then verified by the main LLM. The RNN's recurrent nature allows it to effectively capture temporal dependencies within the generated sequence, leading to more accurate draft token predictions. This contrasts with methods using independent draft heads, which may suffer from reduced accuracy as the prediction horizon increases.

Tree-style attention, previously utilized in techniques like Medusa, plays a crucial role in ReDrafter's efficiency. This mechanism operates on the output of the RNN's beam search, intelligently identifying and removing redundant prefixes from candidate sequences. By focusing the main LLM's attention on the unique portions of the candidate sequences, tree-style attention significantly reduces computational overhead. This optimized approach contributes to ReDrafter's ability to achieve high throughput without sacrificing accuracy.

ReDrafter's integration into TensorRT-LLM involved significant architectural adaptations to minimize overhead and maximize performance. Unlike Medusa, where path acceptance and token sampling occur in the runtime, ReDrafter performs in-engine validation and drafting. This allows for greater optimization freedom within the TensorRT-LLM engine, enabling more efficient kernel selection and scheduling. The in-engine approach also simplifies runtime changes, further enhancing overall efficiency.

The implementation of ReDrafter within TensorRT-LLM involved adding support for numerous new operations, enabling a more direct mapping of PyTorch code into the TensorRT-LLM model definition. This simplifies the integration of complex models and opens up possibilities for further optimizing existing techniques like Medusa within the TensorRT-LLM framework. This enhanced flexibility empowers developers to explore more sophisticated LLM architectures and achieve even greater performance gains.

Strategic Implications: ReDrafter's Impact on the AI Ecosystem

ReDrafter's integration into TensorRT-LLM holds significant strategic importance for NVIDIA. By enhancing LLM inference performance on its GPUs, NVIDIA strengthens its position as a leading provider of AI hardware and software solutions. This is particularly crucial in the rapidly expanding market for AI inference, where efficiency and scalability are paramount. ReDrafter's performance advantages make NVIDIA GPUs even more attractive to developers, researchers, and businesses deploying LLMs, solidifying NVIDIA's leadership in the AI hardware landscape.

Beyond NVIDIA, ReDrafter's benefits extend to the wider AI ecosystem. Developers gain access to a powerful tool for building and deploying high-performance LLM applications with reduced computational costs and complexity. Researchers can leverage ReDrafter to explore larger and more complex LLMs, pushing the boundaries of AI research. Industries relying on high-performance computing for AI applications gain a competitive edge through improved efficiency, faster response times, and enhanced user experiences. ReDrafter's accessibility through TensorRT-LLM fosters innovation and accelerates the development of next-generation AI solutions.

The Future of LLM Inference: Accelerated, Accessible, and Transformative

The integration of ReDrafter into NVIDIA TensorRT-LLM marks a significant step forward in the evolution of LLM inference. By combining cutting-edge research with practical implementation, NVIDIA and Apple have created a powerful tool that democratizes access to high-performance LLM inference. This collaboration not only accelerates the deployment of sophisticated AI models across various applications but also fosters innovation and drives the development of the next generation of intelligent systems. As LLMs continue to grow in size and complexity, ReDrafter's focus on efficiency and scalability will play an increasingly critical role in shaping the future of AI.

----------

Further Reads

I. Turbocharging Meta Llama 3 Performance with NVIDIA TensorRT-LLM and NVIDIA Triton Inference Server | NVIDIA Technical Blog

II. r/LocalLLaMA on Reddit: We've benchmarked TensorRT-LLM: It's 30-70% faster on the same hardware

III. Benchmarking NVIDIA TensorRT-LLM - Jan