Previsão de KV para Melhoria do Tempo até o Primeiro Token

KV Prediction for Improved Time to First Token

October 10, 2024

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

cs.AI

Resumo

A inferência com modelos de linguagem baseados em transformadores começa com uma etapa de processamento do prompt. Nesta etapa, o modelo gera o primeiro token de saída e armazena o cache KV necessário para futuras etapas de geração. Esta etapa de processamento do prompt pode ser computacionalmente cara, levando dezenas de segundos ou mais para modelos de bilhões de parâmetros em dispositivos de borda quando os comprimentos dos prompts ou os tamanhos dos lotes aumentam. Isso degrada a experiência do usuário ao introduzir uma latência significativa nas saídas do modelo. Para reduzir o tempo gasto na produção do primeiro resultado (conhecido como "tempo para o primeiro token", ou TTFT) de um modelo pré-treinado, introduzimos um método inovador chamado Predição KV. Em nosso método, um pequeno modelo auxiliar é usado para processar o prompt e produzir uma aproximação do cache KV usado por um modelo base. Este cache KV aproximado é então usado com o modelo base para geração autoregressiva sem a necessidade de consultar novamente o modelo auxiliar. Demonstramos que nosso método produz uma troca eficiência-precisão pareto-ótima quando comparado com baselines. No TriviaQA, demonstramos melhorias relativas de precisão na faixa de 15%-50% em uma variedade de orçamentos de FLOPs de TTFT. Também demonstramos melhorias de precisão de até 30% na conclusão de código Python HumanEval em orçamentos de FLOPs de TTFT fixos. Além disso, avaliamos modelos em uma CPU Apple M2 Pro e demonstramos que nossa melhoria em FLOPs se traduz em uma aceleração de TTFT no hardware. Disponibilizamos nosso código em 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 .