A new framework using draft models enhances approximate inference for long-context LLMs by better predicting token and key-value pair importance, improving accuracy while maintaining memory and compute efficiency.
Optimizing inference for long-context Large Language Models (LLMs) is
increasingly important due to the quadratic compute and linear memory
complexity of Transformers. Existing approximation methods, such as key-value
(KV) cache dropping, sparse attention, and prompt compression, typically rely
on rough predictions of token or KV pair importance. We propose a novel
framework for approximate LLM inference that leverages small draft models to
more accurately predict the importance of tokens and KV pairs. Specifically, we
introduce two instantiations of our proposed framework: (i) SpecKV, which
leverages a draft output to accurately assess the importance of each KV pair
for more effective KV cache dropping, and (ii) SpecPC, which uses the draft
model’s attention activations to identify and discard unimportant prompt
tokens. To the best of our knowledge, this is the first work to use draft
models for approximate LLM inference acceleration, extending their utility
beyond traditional lossless speculative decoding. We motivate our methods with
theoretical and empirical analyses, and show a strong correlation between the
attention patterns of draft and target models. Extensive experiments on
long-context benchmarks show that our methods consistently achieve higher
accuracy than existing baselines, while preserving the same improvements in
memory usage, latency, and throughput. Our code is available at
https://github.com/furiosa-ai/draft-based-approx-llm.