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

2024-08-20

arXiv.org

AI Summary

Large language models (LLMs) with long contexts (e.g., 32K tokens) face a fundamental trade-off between low latency and high throughput during batched inference. Speculative decoding (SD), while promising, suffers from diminishing returns under large batch sizes due to quadratic KV-cache memory growth with both batch size and sequence length. Method: This work first empirically identifies the scaling law of SD speedup w.r.t. batch size. It then introduces a sparse KV-cache-aware intelligent draft generation strategy to alleviate memory bottlenecks. Further, it integrates dynamic bottleneck profiling with 8×A100 multi-GPU parallel inference for holistic optimization in mid-to-long sequence regimes (32–256 batch size). Contribution/Results: Our approach achieves 2.0× and 1.84× inference speedup on LLaMA-2-7B-32K and LLaMA-3.1-8B, respectively, without accuracy loss. It overturns the conventional view that SD fails under high-throughput settings and establishes a new paradigm for efficient long-context inference.

Technology Category

• Machine Learning: Large Multimodal Models (LMMs)
• Natural Language Processing: (Large) Language Models
• Search and Optimization: Learning to Search

Application Category

• Search and Retrieval-Augmented AI: Search Tool Learning with LLM: Teaching LLMs to invoke search and make use of retrieved information
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems

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 https://github.com/Infini-AI-Lab/MagicDec/.

Problem

• Reducing latency and increasing throughput in long-context LLM applications
• Optimizing speculative decoding for high batch sizes and long sequences
• Enhancing KV cache efficiency to improve speedup in inference

Innovation

• Uses speculative decoding for high throughput
• Leverages sparse KV cache draft model
• Proposes optimal drafting strategy model