最初のトークンまでの時間を改善するための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

要旨

Transformerベースの言語モデルによる推論は、プロンプト処理ステップから始まります。このステップでは、モデルが最初の出力トークンを生成し、将来の生成ステップで必要となるKVキャッシュを保存します。プロンプト処理ステップは計算コストが高く、プロンプトの長さやバッチサイズが増加すると、エッジデバイス上で数十秒以上かかることがあります。これにより、モデルの出力にかなりの遅延が生じ、ユーザーエクスペリエンスが低下します。「最初のトークンまでの時間」（TTFTとして知られる）を削減するために、事前学習済みモデルの最初の出力を生成する時間を短縮するために、KV予測という新しい手法を導入します。この手法では、小さな補助モデルを使用してプロンプトを処理し、基本モデルが使用するKVキャッシュの近似値を生成します。この近似されたKVキャッシュは、補助モデルに再度クエリを送信する必要なく、基本モデルと共に自己回帰生成に使用されます。我々の手法が、ベースラインと比較してパレート最適な効率-精度のトレードオフを生み出すことを示します。TriviaQAでは、TTFT FLOPs予算の範囲で、15％〜50％の相対精度向上を実証します。また、HumanEvalのPythonコード補完では、固定された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 .