KV Prediction for Improved Time to First Token
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 .
670c99dda88de61073ac57ce
Maxwell Horton
admin_assigned
2024-10-14T08:35:23.709Z
647a5afc732a0cd866eb512f
/avatars/fbd0ba2a7ba70414ac903e583ce62933.svg
Maxwell Horton
user
mchorton
670c99dda88de61073ac57cf
Qingqing Cao
admin_assigned
2024-10-14T08:42:56.210Z
6619aa776e681d5e88f9d95a
https://cdn-avatars.huggingface.co/v1/production/uploads/6619aa776e681d5e88f9d95a/fJ_HBq87VfdtJCu6RAUEw.jpeg
Qingqing Cao
user
qicao-apple
670c99dda88de61073ac57d0
Chenfan Sun
claimed_verified
2024-10-17T11:27:46.152Z
658895a1c4b20046639a7e59
/avatars/f80c3f58efa5d22464590bee7a6041a4.svg
Chenfan Sun
user
erpang007
670c99dda88de61073ac57d1
Yanzi Jin
admin_assigned
2024-10-14T08:39:07.212Z
6627307fa39494ebb3a39194
/avatars/99290642fd6d4a6162f52cfe5d46c648.svg
Yanzi Jin
user
yjin25
670c99dda88de61073ac57d2
Sachin Mehta
670c99dda88de61073ac57d3
Mohammad Rastegari
admin_assigned
2024-10-14T08:38:36.380Z
65d4c7734c2163d05eab4110
/avatars/b116a9aca7b0a2a4836a29cb4f069d42.svg
Mohammad Rastegari
user
mrastegari
670c99dda88de61073ac57d4
Moin Nabi
admin_assigned
2024-10-14T08:35:29.716Z
5fe36e9297b72d62ab9b90cd
/avatars/edfea7e8c86dcf355241f5a48365b35a.svg
Moin Nabi
user
moinnabi
KV Prediction for Improved Time to First Token
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 .
KV Prediction reduces time to first token by using a small auxiliary model to approximate the KV cache, improving efficiency with minimal loss in accuracy.
transformer-based language models KV cache time to first token TTFT autoregressive generation KV Prediction efficiency-accuracy trade-off TriviaQA HumanEval python code completion Apple M2 Pro CPU FLOPs