Entwurfsbasierte approximative Inferenz für LLMs

Draft-based Approximate Inference for LLMs

June 10, 2025

Autoren: Kevin Galim, Ethan Ewer, Wonjun Kang, Minjae Lee, Hyung Il Koo, Kangwook Lee

cs.AI

Zusammenfassung

Die Optimierung der Inferenz für Large Language Models (LLMs) mit langem Kontext wird zunehmend wichtiger, da Transformers eine quadratische Rechenkomplexität und eine lineare Speicherkomplexität aufweisen. Bestehende Approximationsmethoden, wie das Weglassen des Key-Value (KV)-Caches, sparsame Aufmerksamkeit und Prompt-Kompression, basieren typischerweise auf groben Vorhersagen der Bedeutung von Tokens oder KV-Paaren. Wir schlagen ein neuartiges Framework für die approximative LLM-Inferenz vor, das kleine Draft-Modelle nutzt, um die Bedeutung von Tokens und KV-Paaren genauer vorherzusagen. Konkret führen wir zwei Instanzen unseres vorgeschlagenen Frameworks ein: (i) SpecKV, das einen Draft-Output verwendet, um die Bedeutung jedes KV-Paares genau zu bewerten und so das Weglassen des KV-Caches effektiver zu gestalten, und (ii) SpecPC, das die Aufmerksamkeitsaktivitäten des Draft-Modells nutzt, um unwichtige Prompt-Tokens zu identifizieren und zu verwerfen. Unseres Wissens ist dies die erste Arbeit, die Draft-Modelle für die Beschleunigung der approximativen LLM-Inferenz verwendet und damit ihren Nutzen über das traditionelle verlustfreie spekulative Decodieren hinaus erweitert. Wir untermauern unsere Methoden mit theoretischen und empirischen Analysen und zeigen eine starke Korrelation zwischen den Aufmerksamkeitsmustern von Draft- und Zielmodellen. Umfangreiche Experimente auf Benchmarks mit langem Kontext zeigen, dass unsere Methoden durchweg eine höhere Genauigkeit als bestehende Baselines erreichen, während die gleichen Verbesserungen in Bezug auf Speicherverbrauch, Latenz und Durchsatz erhalten bleiben. Unser Code ist verfügbar unter https://github.com/furiosa-ai/draft-based-approx-llm.

English

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.