Conceptgebaseerde Benaderende Inferentie voor LLM's

Draft-based Approximate Inference for LLMs

June 10, 2025

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

cs.AI

Samenvatting

Het optimaliseren van inferentie voor Large Language Models (LLMs) met lange context wordt steeds belangrijker vanwege de kwadratische rekenkracht en lineaire geheugencomplexiteit van Transformers. Bestaande benaderingsmethoden, zoals het laten vallen van de key-value (KV) cache, sparse attention en promptcompressie, vertrouwen doorgaans op ruwe voorspellingen van het belang van tokens of KV-paren. Wij stellen een nieuw raamwerk voor voor benaderende LLM-inferentie dat gebruikmaakt van kleine draft-modellen om het belang van tokens en KV-paren nauwkeuriger te voorspellen. Specifiek introduceren we twee uitvoeringen van ons voorgestelde raamwerk: (i) SpecKV, dat een draft-uitvoer gebruikt om het belang van elk KV-paar nauwkeurig te beoordelen voor effectiever laten vallen van de KV-cache, en (ii) SpecPC, dat de aandachtactivaties van het draft-model gebruikt om onbelangrijke prompt-tokens te identificeren en te verwijderen. Voor zover wij weten, is dit het eerste werk dat draft-modellen gebruikt voor benaderende LLM-inferentieversnelling, waardoor hun nut wordt uitgebreid voorbij traditioneel verliesvrij speculatief decoderen. We motiveren onze methoden met theoretische en empirische analyses en tonen een sterke correlatie aan tussen de aandachtpatronen van draft- en doelmodellen. Uitgebreide experimenten op benchmarks met lange context laten zien dat onze methoden consistent een hogere nauwkeurigheid bereiken dan bestaande baselines, terwijl dezelfde verbeteringen in geheugengebruik, latentie en doorvoer behouden blijven. Onze code is beschikbaar op 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.