KV-Vorhersage zur Verbesserung der Zeit bis zum ersten Token.

KV Prediction for Improved Time to First Token

October 10, 2024

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

cs.AI

Zusammenfassung

Die Inferenz mit auf Transformer basierenden Sprachmodellen beginnt mit einem Schritt zur Prompt-Verarbeitung. In diesem Schritt generiert das Modell das erste Ausgabetoken und speichert den KV-Cache, der für zukünftige Generierungsschritte benötigt wird. Dieser Schritt zur Prompt-Verarbeitung kann rechenintensiv sein und bei Milliarden-Parameter-Modellen auf Edge-Geräten, wenn Prompt-Längen oder Batch-Größen zunehmen, 10 Sekunden oder mehr in Anspruch nehmen. Dies beeinträchtigt die Benutzererfahrung, da signifikante Latenz in die Ausgaben des Modells eingeführt wird. Um die Zeit zur Erzeugung des ersten Ausgabewerts (bekannt als "Zeit bis zum ersten Token" oder TTFT) eines vorab trainierten Modells zu reduzieren, führen wir eine neue Methode namens KV-Vorhersage ein. In unserer Methode wird ein kleines Hilfsmodell verwendet, um die Prompt zu verarbeiten und eine Näherung des KV-Caches zu erzeugen, der von einem Basismodell verwendet wird. Dieser approximierte KV-Cache wird dann mit dem Basismodell für die autoregressive Generierung verwendet, ohne dass das Hilfsmodell erneut abgefragt werden muss. Wir zeigen, dass unsere Methode einen pareto-optimalen Effizienz-Genauigkeits-Trade-off im Vergleich zu Baselines erzielt. Auf TriviaQA zeigen wir relative Genauigkeitsverbesserungen im Bereich von 15 % bis 50 % bei verschiedenen TTFT-FLOPs-Budgets. Wir zeigen auch Genauigkeitsverbesserungen von bis zu 30 % bei der Python-Codevervollständigung von HumanEval bei festgelegten TTFT-FLOPs-Budgets. Darüber hinaus benchmarken wir Modelle auf einem Apple M2 Pro CPU und zeigen, dass unsere Verbesserung bei FLOPs zu einer Beschleunigung der TTFT auf der Hardware führt. Wir veröffentlichen unseren Code unter 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 .