Table of Contents
Fetching ...

AMUSD: Asynchronous Multi-Device Speculative Decoding for LLM Acceleration

Bradley McDanel

TL;DR

This work introduces AMUSD (Asynchronous Multi-device Speculative Decoding), a system that further accelerates generation by decoupling the draft and verify phases into a continuous, asynchronous approach that enables both models to perform predictions independently on separate devices.

Abstract

Large language models typically generate tokens autoregressively, using each token as input for the next. Recent work on Speculative Decoding has sought to accelerate this process by employing a smaller, faster draft model to more quickly generate candidate tokens. These candidates are then verified in parallel by the larger (original) verify model, resulting in overall speedup compared to using the larger model by itself in an autoregressive fashion. In this work, we introduce AMUSD (Asynchronous Multi-device Speculative Decoding), a system that further accelerates generation by decoupling the draft and verify phases into a continuous, asynchronous approach. Unlike conventional speculative decoding, where only one model (draft or verify) performs token generation at a time, AMUSD enables both models to perform predictions independently on separate devices (e.g., GPUs). We evaluate our approach over multiple datasets and show that AMUSD achieves an average 29% improvement over speculative decoding and up to 1.96$\times$ speedup over conventional autoregressive decoding, while achieving identical output quality. Our system is open-source and available at https://github.com/BradMcDanel/AMUSD/.

AMUSD: Asynchronous Multi-Device Speculative Decoding for LLM Acceleration

TL;DR

This work introduces AMUSD (Asynchronous Multi-device Speculative Decoding), a system that further accelerates generation by decoupling the draft and verify phases into a continuous, asynchronous approach that enables both models to perform predictions independently on separate devices.

Abstract

Large language models typically generate tokens autoregressively, using each token as input for the next. Recent work on Speculative Decoding has sought to accelerate this process by employing a smaller, faster draft model to more quickly generate candidate tokens. These candidates are then verified in parallel by the larger (original) verify model, resulting in overall speedup compared to using the larger model by itself in an autoregressive fashion. In this work, we introduce AMUSD (Asynchronous Multi-device Speculative Decoding), a system that further accelerates generation by decoupling the draft and verify phases into a continuous, asynchronous approach. Unlike conventional speculative decoding, where only one model (draft or verify) performs token generation at a time, AMUSD enables both models to perform predictions independently on separate devices (e.g., GPUs). We evaluate our approach over multiple datasets and show that AMUSD achieves an average 29% improvement over speculative decoding and up to 1.96 speedup over conventional autoregressive decoding, while achieving identical output quality. Our system is open-source and available at https://github.com/BradMcDanel/AMUSD/.

Paper Structure

This paper contains 11 sections, 5 figures, 1 table, 1 algorithm.

Figures (5)

  • Figure 1: Comparison with synchronous speculative decoding (left) with AMUSD (right). The synchronous approach alternatives between drafting and verifying phases, meaning only model can work at a time. AMUSD uses asynchronous generation running on multiple GPUs enable continuous generation. The draft model must rollback invalid tokens that conflict with the verify model.
  • Figure 2: Synchronous decoding (top) alternates between draft and verify phases, while asynchronous decoding (bottom) runs both processes in parallel, with rollbacks on invalidation.
  • Figure 3: AMUSD system architecture showing the asynchronous interaction between Draft and Verify processes via shared memory. The Main Process coordinates execution while Draft (GPU 0) and Verify (GPU 1) processes maintain local tensors that sync with shared memory only when necessary, enabling efficient parallel execution while ensuring consistency through strategic updates and rollbacks.
  • Figure 4: Token generation progress over time on a RefactorChat sample. AMUSD maintains consistently higher throughput due to parallel execution.
  • Figure 5: GPU resource utilization comparison across decoding strategies. Note that autoregressive and speculative decoding only actively use GPU0, with GPU1 shown to indicate baseline idle power consumption. Top: average power consumption. Middle: GPU utilization demonstrating resource use across devices. Bottom: total energy consumption per sample.