MagicDec: Breaking the Latency-Throughput Tradeoff for Long Context Generation with Speculative Decoding

Jian Chen,Vashisth Tiwari,Ranajoy Sadhukhan,Zhuoming Chen,Jinyuan Shi,Ian En-Hsu Yen,Beidi Chen

2024-08-24

Abstract:Large Language Models (LLMs) have become more prevalent in long-context applications such as interactive chatbots, document analysis, and agent workflows, but it is challenging to serve long-context requests with low latency and high throughput. Speculative decoding (SD) is a widely used technique to reduce latency without sacrificing performance but the conventional wisdom suggests that its efficacy is limited to small batch sizes. In MagicDec, we show that surprisingly SD can achieve speedup even for a high throughput inference regime for moderate to long sequences. More interestingly, an intelligent drafting strategy can achieve better speedup with increasing batch size based on our rigorous analysis. MagicDec first identifies the bottleneck shifts with increasing batch size and sequence length, and uses these insights to deploy speculative decoding more effectively for high throughput inference. Then, it leverages draft models with sparse KV cache to address the KV bottleneck that scales with both sequence length and batch size. This finding underscores the broad applicability of speculative decoding in long-context serving, as it can enhance throughput and reduce latency without compromising accuracy. For moderate to long sequences, we demonstrate up to 2x speedup for LLaMA-2-7B-32K and 1.84x speedup for LLaMA-3.1-8B when serving batch sizes ranging from 32 to 256 on 8 NVIDIA A100 GPUs. The code is available at <a class="link-external link-https" href="https://github.com/Infini-AI-Lab/MagicDec/" rel="external noopener nofollow">this https URL</a>.

Computation and Language

What problem does this paper attempt to address?

The paper aims to address the trade-off between low latency and high throughput that large language models (LLMs) face when handling long-context requests. Specifically:

1. **Challenge Existing Notions**: The paper challenges the traditional view that speculative decoding (SD) is only suitable for small batch sizes to improve throughput. The authors demonstrate through research that, in the case of medium to long sequences, speculative decoding can simultaneously improve throughput, reduce latency, and maintain accuracy.

2. **Address Bottleneck Issues**: The paper proposes a new strategy that analyzes the changes in bottlenecks under different batch sizes and sequence lengths, using sparse KV caching to solve the KV bottleneck problem. This method effectively enhances system performance when processing large-scale batch data.

3. **Experimental Validation**: Through extensive experiments on different hardware platforms, the paper showcases the advantages of speculative decoding technology in handling long sequences, particularly achieving significant speed improvements on larger batch data.

In summary, the core objective of the paper is to achieve low latency and high throughput in long-context generation tasks by optimizing speculative decoding methods without sacrificing accuracy.