Prédiction KV pour une amélioration du temps jusqu'au premier jeton

KV Prediction for Improved Time to First Token

October 10, 2024

Auteurs: Maxwell Horton, Qingqing Cao, Chenfan Sun, Yanzi Jin, Sachin Mehta, Mohammad Rastegari, Moin Nabi

cs.AI

Résumé

L'inférence avec les modèles de langage basés sur les transformers commence par une étape de traitement de l'invite. Dans cette étape, le modèle génère le premier jeton de sortie et stocke le cache KV nécessaire pour les futures étapes de génération. Cette étape de traitement de l'invite peut être coûteuse en termes de calcul, prenant 10 secondes ou plus pour les modèles à un milliard de paramètres sur des appareils périphériques lorsque les longueurs des invites ou les tailles des lots augmentent. Cela dégrade l'expérience utilisateur en introduisant une latence significative dans les sorties du modèle. Pour réduire le temps passé à produire le premier jeton de sortie (connu sous le nom de "temps jusqu'au premier jeton", ou TTFT) d'un modèle pré-entraîné, nous introduisons une nouvelle méthode appelée Prédiction KV. Dans notre méthode, un petit modèle auxiliaire est utilisé pour traiter l'invite et produire une approximation du cache KV utilisé par un modèle de base. Ce cache KV approximé est ensuite utilisé avec le modèle de base pour une génération autorégressive sans avoir besoin d'interroger à nouveau le modèle auxiliaire. Nous démontrons que notre méthode produit un compromis efficacité-précision optimal de Pareto par rapport aux références. Sur TriviaQA, nous démontrons des améliorations de précision relatives dans une fourchette de 15 % à 50 % sur une gamme de budgets de FLOPs TTFT. Nous démontrons également des améliorations de précision allant jusqu'à 30 % sur la complétion de code Python HumanEval à des budgets FLOPs TTFT fixes. De plus, nous évaluons les modèles sur un processeur Apple M2 Pro et démontrons que notre amélioration en FLOPs se traduit par une accélération du TTFT sur le matériel. Nous mettons notre code à disposition sur https://github.com/apple/corenet/tree/main/projects/kv-prediction.

English

Inference with transformer-based language models begins with a prompt processing step. In this step, the model generates the first output token and stores the KV cache needed for future generation steps. This prompt processing step can be computationally expensive, taking 10s of seconds or more for billion-parameter models on edge devices when prompt lengths or batch sizes rise. This degrades user experience by introducing significant latency into the model's outputs. To reduce the time spent producing the first output (known as the ``time to first token'', or TTFT) of a pretrained model, we introduce a novel method called KV Prediction. In our method, a small auxiliary model is used to process the prompt and produce an approximation of the KV cache used by a base model. This approximated KV cache is then used with the base model for autoregressive generation without the need to query the auxiliary model again. We demonstrate that our method produces a pareto-optimal efficiency-accuracy trade-off when compared to baselines. On TriviaQA, we demonstrate relative accuracy improvements in the range of 15%-50% across a range of TTFT FLOPs budgets. We also demonstrate accuracy improvements of up to 30% on HumanEval python code completion at fixed TTFT FLOPs budgets. Additionally, we benchmark models on an Apple M2 Pro CPU and demonstrate that our improvement in FLOPs translates to a TTFT speedup on hardware. We release our code at https://github.com/apple/corenet/tree/main/projects/kv-prediction .