Break the Sequential Dependency of LLM Inference Using Lookahead Decoding

Autoregressive decoding of large language models (LLMs) is memory bandwidth bounded, resulting in high latency and significant wastes of the parallel processing power of modern accelerators. Existing methods for accelerating LLM decoding often require a draft model (e.g., speculative decoding), which is nontrivial to obtain and unable to generalize. In this paper, we introduce Lookahead decoding, an exact, parallel decoding algorithm that accelerates LLM decoding without needing auxiliary models or data stores. It allows trading per-step log(FLOPs) to reduce the number of total decoding steps, is more parallelizable on single or multiple modern accelerators, and is compatible with concurrent memory-efficient attention (e.g., FlashAttention). Our implementation of Lookahead decoding can speed up autoregressive decoding by up to 1.8x on MT-bench and 4x with strong scaling on multiple GPUs in code completion tasks. Our code is avialable at https://github.com/hao-ai-lab/LookaheadDecoding

한 줄 요약

Lookahead Decoding은 보조 모델 없이 LLM 추론 속도를 1.8x~4x까지 향상시키는 정확한 병렬 디코딩 알고리즘이다.

핵심 기여도

Lookahead Decoding이라는 새로운 정확한 병렬 디코딩 알고리즘을 제안.
기존 추측-검증 방식과 달리 보조 모델 없이도 작동하며, 최대 1.8x (MT-bench) 및 4x (다중 GPU)의 가속 성능 달성.
log(FLOPs) 증가에 따라 디코딩 단계 수가 선형 감소하며, 병렬화 가능성을 확장.
FlashAttention과 호환되며, 분산 CUDA 구현을 통해 다중 GPU 환경에서 Lookahead Parallelism을 지원.

핵심 아이디어

기존의 자동 회귀 디코딩은 단일 토큰만 생성하며, 병렬 처리 능력을 낭비한다. Lookahead Decoding은 이 문제를 해결하기 위해 자동 회귀 디코딩을 비선형 시스템의 고정점 반복(Jacobi decoding) 문제로 재구성한다. 이 방식은 단계별로 여러 토큰을 병렬 생성하고, 이들을 검증하는 두 가지 브랜치(lookahead branch, verification branch)를 사용한다. 이는 기존 추측-검증 방식과 달리, 별도의 보조 모델 없이도 작동하며, n-gram 풀을 활용해 과거 생성 결과를 재사용함으로써 효율성을 높인다. 특히, log(FLOPs)가 증가할수록 디코딩 단계 수가 줄어들어, 계산 자원과 속도 간의 트레이드오프를 가능하게 한다.

기술적 접근법

**Lookahead Decoding**: Jacobi decoding 기반으로, 고정점 반복을 통해 병렬 토큰 생성.
**n-gram pool**: 과거 생성된 n-gram을 캐싱하여 반복 검증 없이 재사용.
**CUDA 및 분산 구현**: Lookahead Parallelism을 통해 다중 GPU 환경에서 병렬 처리 가능.
**FlashAttention 호환**: 메모리 효율적인 어텐션 알고리즘과 결합 가능.
**Sampling 방법 지원**: greedy, top-K, top-P 등 다양한 샘플링 방법 유지.

주요 결과

**LLaMA-2 모델에서 MT-bench 데이터셋**: 1.8x의 추론 속도 향상.
**코드 완성 작업에서 8개 GPU 사용 시**: Lookahead Parallelism으로 최대 4x의 속도 향상.
**log(FLOPs) 증가에 따른 디코딩 단계 수 감소**: 단계 수가 log(FLOPs)에 비례하여 줄어들며, 계산 효율성 향상.

의의 및 한계

Lookahead Decoding은 기존 추측-검증 방식의 주요 단점인 보조 모델 의존성을 제거하고, 병렬 처리를 가능하게 함으로써 LLM 추론의 효율성을 크게 향상시킨다. 특히, FlashAttention과 결합된 CUDA 구현은 실용적 성능 향상에 기여하며, 다중 GPU 환경에서의 확장성도 입증되었다. 그러나, Jacobi decoding 자체는 토큰 위치 오류가 발생할 수 있어, 실제 성능 향상은 검증 단계에서의 정확도에 크게 의존한다. 또한, 특정 작업(예: 긴 텍스트 생성)에서는 n-gram 풀의 효과가 제한될 수 있다.

실용적 활용

Lookahead Decoding은 코드 생성, 대화형 챗봇, 실시간 검색 등 지연에 민감한 응용 분야에서 유용하게 사용될 수 있다. 특히, 고성능 GPU 클러스터를 활용한 대규모 LLM 서비스에서 추론 속도 향상과 병렬화를 통해 비용 효율성을 높일 수 있다.