Inférence approximative basée sur des ébauches pour les LLM

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

Résumé

L'optimisation de l'inférence pour les modèles de langage à grand contexte (LLMs) devient de plus en plus cruciale en raison de la complexité quadratique en calcul et linéaire en mémoire des Transformers. Les méthodes d'approximation existantes, telles que l'élimination du cache clé-valeur (KV), l'attention sparse et la compression de prompt, reposent généralement sur des prédictions approximatives de l'importance des tokens ou des paires KV. Nous proposons un nouveau cadre pour l'inférence approximative des LLMs qui exploite de petits modèles de brouillon pour prédire plus précisément l'importance des tokens et des paires KV. Plus précisément, nous introduisons deux instanciations de notre cadre proposé : (i) SpecKV, qui utilise une sortie de brouillon pour évaluer avec précision l'importance de chaque paire KV afin d'améliorer l'élimination du cache KV, et (ii) SpecPC, qui utilise les activations d'attention du modèle de brouillon pour identifier et éliminer les tokens de prompt non importants. À notre connaissance, il s'agit du premier travail à utiliser des modèles de brouillon pour accélérer l'inférence approximative des LLMs, étendant ainsi leur utilité au-delà du décodage spéculatif traditionnel sans perte. Nous justifions nos méthodes par des analyses théoriques et empiriques, et montrons une forte corrélation entre les modèles d'attention des modèles de brouillon et des modèles cibles. Des expériences approfondies sur des benchmarks à long contexte montrent que nos méthodes atteignent systématiquement une précision supérieure à celle des méthodes de référence existantes, tout en conservant les mêmes améliorations en termes d'utilisation de la mémoire, de latence et de débit. Notre code est disponible à l'adresse 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.