Predicción de KV para Mejorar el Tiempo hasta el Primer 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

Resumen

La inferencia con modelos de lenguaje basados en transformadores comienza con un paso de procesamiento de la indicación. En este paso, el modelo genera el primer token de salida y almacena la caché KV necesaria para futuros pasos de generación. Este paso de procesamiento de la indicación puede ser computacionalmente costoso, tardando 10 segundos o más para modelos de mil millones de parámetros en dispositivos periféricos cuando las longitudes de las indicaciones o los tamaños de lote aumentan. Esto degrada la experiencia del usuario al introducir una latencia significativa en las salidas del modelo. Para reducir el tiempo dedicado a producir el primer resultado (conocido como el "tiempo hasta el primer token", o TTFT) de un modelo preentrenado, presentamos un método novedoso llamado Predicción KV. En nuestro método, se utiliza un pequeño modelo auxiliar para procesar la indicación y producir una aproximación de la caché KV utilizada por un modelo base. Esta caché KV aproximada se utiliza luego con el modelo base para generación autoregresiva sin necesidad de consultar nuevamente al modelo auxiliar. Demostramos que nuestro método produce un equilibrio óptimo entre eficiencia y precisión en comparación con los baselines. En TriviaQA, demostramos mejoras relativas de precisión en el rango del 15% al 50% en una variedad de presupuestos de FLOPs de TTFT. También demostramos mejoras de precisión de hasta el 30% en la finalización de código Python HumanEval en presupuestos fijos de FLOPs de TTFT. Además, evaluamos los modelos en una CPU Apple M2 Pro y demostramos que nuestra mejora en FLOPs se traduce en una aceleración de TTFT en el hardware. Publicamos nuestro código en 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 .