LongSpec: Long-Context Lossless Speculative Decoding with Efficient Drafting and Verification

Fuente: arXiv
Salvato in:
Dettagli Bibliografici
Autori principali: Yang, Penghui, Du, Cunxiao, Zhang, Fengzhuo, Wang, Haonan, Pang, Tianyu, Du, Chao, An, Bo
Natura: Preprint
Pubblicazione: 2025
Soggetti:
Accesso online:
Tags: Aggiungi Tag
Nessun Tag, puoi essere il primo ad aggiungerne!!
_version_ 1866917390260895744
author Yang, Penghui
Du, Cunxiao
Zhang, Fengzhuo
Wang, Haonan
Pang, Tianyu
Du, Chao
An, Bo
author_facet Yang, Penghui
Du, Cunxiao
Zhang, Fengzhuo
Wang, Haonan
Pang, Tianyu
Du, Chao
An, Bo
contents As Large Language Models (LLMs) can now process extremely long contexts, efficient inference over these extended inputs has become increasingly important, especially for emerging applications like LLM agents that highly depend on this capability. Speculative decoding (SD) offers a promising lossless acceleration technique compared to lossy alternatives such as quantization and model cascades. However, most state-of-the-art SD methods are trained on short texts (typically fewer than 4k tokens), making them unsuitable for long-context scenarios. Specifically, adapting these methods to long contexts presents three key challenges: (1) the excessive memory demands posed by draft models due to large Key-Value (KV) cache; (2) performance degradation resulting from the mismatch between short-context training and long-context inference; and (3) inefficiencies in tree attention mechanisms when managing long token sequences. This work introduces LongSpec, a framework that addresses these challenges through three core innovations: a memory-efficient draft model with a constant-sized KV cache; novel position indices that mitigate the training-inference mismatch; and an attention aggregation strategy that combines fast prefix computation with standard tree attention to enable efficient decoding. Experimental results confirm the effectiveness of LongSpec, achieving up to a 3.26x speedup over strong Flash Attention baselines across five long-context understanding datasets, as well as a 2.25x reduction in wall-clock time on the AIME24 long reasoning task with the QwQ model, demonstrating significant latency improvements for long-context applications. The code is available at https://github.com/sail-sg/LongSpec.
format Preprint
id arxiv_https___arxiv_org_abs_2502_17421
institution arXiv
publishDate 2025
record_format arxiv
spellingShingle LongSpec: Long-Context Lossless Speculative Decoding with Efficient Drafting and Verification
Yang, Penghui
Du, Cunxiao
Zhang, Fengzhuo
Wang, Haonan
Pang, Tianyu
Du, Chao
An, Bo
Computation and Language
Artificial Intelligence
Machine Learning
As Large Language Models (LLMs) can now process extremely long contexts, efficient inference over these extended inputs has become increasingly important, especially for emerging applications like LLM agents that highly depend on this capability. Speculative decoding (SD) offers a promising lossless acceleration technique compared to lossy alternatives such as quantization and model cascades. However, most state-of-the-art SD methods are trained on short texts (typically fewer than 4k tokens), making them unsuitable for long-context scenarios. Specifically, adapting these methods to long contexts presents three key challenges: (1) the excessive memory demands posed by draft models due to large Key-Value (KV) cache; (2) performance degradation resulting from the mismatch between short-context training and long-context inference; and (3) inefficiencies in tree attention mechanisms when managing long token sequences. This work introduces LongSpec, a framework that addresses these challenges through three core innovations: a memory-efficient draft model with a constant-sized KV cache; novel position indices that mitigate the training-inference mismatch; and an attention aggregation strategy that combines fast prefix computation with standard tree attention to enable efficient decoding. Experimental results confirm the effectiveness of LongSpec, achieving up to a 3.26x speedup over strong Flash Attention baselines across five long-context understanding datasets, as well as a 2.25x reduction in wall-clock time on the AIME24 long reasoning task with the QwQ model, demonstrating significant latency improvements for long-context applications. The code is available at https://github.com/sail-sg/LongSpec.
title LongSpec: Long-Context Lossless Speculative Decoding with Efficient Drafting and Verification
topic Computation and Language
Artificial Intelligence
Machine Learning
url https://arxiv.org/abs/2502.17421