KV-voorspelling voor een verbeterde tijd tot de eerste token.

Samenvatting

Het afleiden met op transformer-gebaseerde taalmodellen begint met een promptverwerkingsstap. In deze stap genereert het model het eerste uitvoertoken en slaat de KV-cache op die nodig is voor toekomstige generatiestappen. Deze promptverwerkingsstap kan rekenintensief zijn en 10-tallen seconden of meer in beslag nemen voor modellen met miljarden parameters op randapparaten wanneer de promptlengtes of batchgroottes toenemen. Dit vermindert de gebruikerservaring door aanzienlijke latentie in de uitvoer van het model te introduceren. Om de tijd die wordt besteed aan het produceren van het eerste uitvoerresultaat (bekend als de "tijd tot het eerste token", of TTFT) van een voortraind model te verminderen, introduceren we een nieuwe methode genaamd KV Voorspelling. In onze methode wordt een kleine hulpmodel gebruikt om de prompt te verwerken en een benadering van de KV-cache te produceren die door een basismodel wordt gebruikt. Deze benaderde KV-cache wordt vervolgens gebruikt met het basismodel voor autoregressieve generatie zonder dat het hulpmodel opnieuw hoeft te worden geraadpleegd. We tonen aan dat onze methode een pareto-optimale efficiëntie-nauwkeurigheidshandel oplevert in vergelijking met baselines. Op TriviaQA laten we relatieve nauwkeurigheidsverbeteringen zien in het bereik van 15%-50% over een reeks TTFT FLOPs-budgetten. We laten ook nauwkeurigheidsverbeteringen zien tot 30% bij het aanvullen van Python-code in HumanEval bij vaste TTFT FLOPs-budgetten. Daarnaast benchmarken we modellen op een Apple M2 Pro CPU en tonen aan dat onze verbetering in FLOPs zich vertaalt naar een versnelling van de TTFT op hardware. We publiceren onze code op 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 .

KV-voorspelling voor een verbeterde tijd tot de eerste token.

KV Prediction for Improved Time to First Token

Samenvatting

Support