Draft-based Approximate Inference for LLMs
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.
684ae1f3dbd21a9cc27b0f32
Kevin Galim
684ae1f3dbd21a9cc27b0f33
Ethan Ewer
684ae1f3dbd21a9cc27b0f34
Wonjun Kang
684ae1f3dbd21a9cc27b0f35
Minjae Lee
684ae1f3dbd21a9cc27b0f36
Hyung Il Koo
684ae1f3dbd21a9cc27b0f37
Kangwook Lee
Draft-based Approximate Inference for LLMs
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.
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.
Large Language Models Transformers key-value cache dropping sparse attention prompt compression draft models SpecKV SpecPC attention activations long-context benchmarks