개선된 첫 번째 토큰까지의 시간을 위한 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초 이상 소요될 수 있습니다. 이는 모델의 출력에 상당한 지연을 도입하여 사용자 경험을 저하시킵니다. 사전 훈련된 모델의 첫 번째 출력 생성에 소요되는 시간을 줄이기 위해, 우리는 KV 예측이라는 새로운 방법을 소개합니다. 우리의 방법에서는 작은 보조 모델을 사용하여 프롬프트를 처리하고 기본 모델이 사용하는 KV 캐시의 근사치를 생성합니다. 이 근사된 KV 캐시는 보조 모델을 다시 쿼리할 필요 없이 기본 모델과 함께 자기회귀 생성에 사용됩니다. 우리의 방법이 베이스라인과 비교했을 때 파레토-최적 효율성-정확도 트레이드오프를 보여준다는 것을 입증합니다. TriviaQA에서는 TTFT FLOPs 예산 범위에서 15%~50%의 상대적 정확도 향상을 보여줍니다. 또한, HumanEval 파이썬 코드 완성에서는 고정된 TTFT FLOPs 예산에서 최대 30%의 정확도 향상을 보여줍니다. 추가로, Apple M2 Pro CPU에서 모델을 벤치마킹하고 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 .