Прогнозирование KV для улучшения времени до первого токена.

KV Prediction for Improved Time to First Token

October 10, 2024

Авторы: Maxwell Horton, Qingqing Cao, Chenfan Sun, Yanzi Jin, Sachin Mehta, Mohammad Rastegari, Moin Nabi

cs.AI

Аннотация

Вывод с моделями языка на основе трансформеров начинается с этапа обработки подсказки. На этом этапе модель генерирует первый токен вывода и сохраняет кэш KV, необходимый для будущих шагов генерации. Этот этап обработки подсказки может быть вычислительно затратным, занимая 10 секунд или более для моделей с миллиардом параметров на периферийных устройствах при увеличении длины подсказки или размеров пакетов. Это снижает пользовательский опыт за счет введения значительной задержки в выходные данные модели. Для сокращения времени, затраченного на создание первого вывода (известного как "время до первого токена", или TTFT) предварительно обученной модели, мы предлагаем новый метод, называемый Прогнозирование KV. В нашем методе используется небольшая вспомогательная модель для обработки подсказки и создания приближения кэша KV, используемого базовой моделью. Этот приближенный кэш KV затем используется с базовой моделью для авторегрессивной генерации без необходимости повторного запроса к вспомогательной модели. Мы демонстрируем, что наш метод обеспечивает оптимальный компромисс между эффективностью и точностью по сравнению с базовыми показателями. На наборе данных TriviaQA мы показываем относительное улучшение точности в диапазоне от 15% до 50% при различных бюджетах TTFT FLOPs. Мы также демонстрируем улучшение точности до 30% при заведомо заданных бюджетах TTFT FLOPs для завершения кода на Python в HumanEval. Кроме того, мы проводим бенчмаркинг моделей на процессоре Apple M2 Pro и демонстрируем, что наше улучшение в FLOPs приводит к ускорению TTFT на аппаратном обеспечении. Мы выкладываем наш код по адресу 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 .