Predizione KV per un Miglioramento del Tempo al Primo Token

Abstract

L'inferenza con modelli linguistici basati su trasformatori inizia con una fase di elaborazione del prompt. In questa fase, il modello genera il primo token di output e memorizza la cache KV necessaria per i passaggi di generazione futuri. Questa fase di elaborazione del prompt può essere computazionalmente costosa, richiedendo 10 secondi o più per modelli da miliardi di parametri su dispositivi edge quando le lunghezze dei prompt o le dimensioni dei batch aumentano. Ciò degrada l'esperienza dell'utente introducendo una latenza significativa nei risultati del modello. Per ridurre il tempo impiegato per produrre il primo output (noto come "tempo per il primo token", o TTFT) di un modello preaddestrato, introduciamo un nuovo metodo chiamato Predizione KV. Nel nostro metodo, un piccolo modello ausiliario viene utilizzato per elaborare il prompt e produrre un'approssimazione della cache KV utilizzata da un modello di base. Questa cache KV approssimata viene quindi utilizzata con il modello di base per la generazione autoregressiva senza la necessità di interrogare nuovamente il modello ausiliario. Dimostriamo che il nostro metodo produce un compromesso efficienza-accuratezza pareto-ottimale rispetto alle baselines. Su TriviaQA, mostriamo miglioramenti di accuratezza relativi nell'intervallo del 15%-50% su una serie di budget TTFT FLOPs. Dimostriamo anche miglioramenti di accuratezza fino al 30% nel completamento del codice Python HumanEval a budget fissi di TTFT FLOPs. Inoltre, valutiamo i modelli su una CPU Apple M2 Pro e dimostriamo che il nostro miglioramento nei FLOPs si traduce in un'accelerazione del TTFT sull'hardware. Rilasciamo il nostro codice su 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 .

Predizione KV per un Miglioramento del Tempo al Primo Token

KV Prediction for Improved Time to First Token

Abstract

Support