ICML 2025

RAPID: Long-Context Inference with Retrieval-Augmented
Speculative Decoding
Guanzheng Chen * 1 2 3 Qilong Feng * 1 Jinjie Ni 1 Xin Li 2 3 Michael Qizhe Shieh 1
Code:
https://github.com/NUS-TRAIL/RAPID
Abstract
The emergence of long-context large language
models (LLMs) offers a promising alternative to
traditional retrieval-augmented generation (RAG)
for processing extensive documents. However, the
computational overhead of long-context inference
presents significant efficiency challenges. While
Speculative Decoding (SD) traditionally acceler-
ates inference using smaller draft models, its effec-
tiveness diminishes substantially in long-context
scenarios due to memory-bound KV cache op-
erations.
We introduce Retrieval-Augmented
SPeculatIve Decoding (RAPID), which leverages
RAG for both accelerating and enhancing gener-
ation quality in long-context inference. RAPID
introduces the RAG drafter—a draft LLM oper-
ating on shortened retrieval contexts—to spec-
ulate on the generation of long-context target
LLMs. Our approach enables a new paradigm
where same-scale or even larger LLMs can serve
as RAG drafters while maintaining computational
efficiency. To fully leverage the potentially supe-
rior capabilities from stronger RAG drafters, we
develop an inference-time knowledge transfer that
enriches the target distribution by RAG. Exten-
sive experiments on the LLaMA-3.1 and Qwen2.5
backbones demonstrate that RAPID effectively
integrates the strengths of both RAG and long-
context LLMs, achieving significant performance
improvements (e.g., from 39.33 to 42.83 on In-
finiteBench for LLaMA-3.1-8B) with more than
2× speedups for long-context inference. Our anal-
yses also reveal the robustness of RAPID across
various context lengths and retrieval quality.
*Equal contribution 1National University of Singapore 2DAMO
Academy, Alibaba Group 3Hupan Lab, 310023, Hangzhou, China.
Correspondence to: Guanzheng Chen <gc.chen@u.nus.edu>,
Michael Qizhe Shieh <michaelshieh@comp.nus.edu.s...