title: KV Prediction for Improved Time to First Token
summary: 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 .
paper_authors._id: 670c99dda88de61073ac57ce
paper_authors.name: Maxwell Horton
paper_authors.status: admin_assigned
paper_authors.statusLastChangedAt: 2024-10-14T08:35:23.709Z
paper_authors._id: 647a5afc732a0cd866eb512f
paper_authors.avatarUrl: /avatars/fbd0ba2a7ba70414ac903e583ce62933.svg
paper_authors.fullname: Maxwell Horton
paper_authors.type: user
paper_authors.user: mchorton
paper_authors._id: 670c99dda88de61073ac57cf
paper_authors.name: Qingqing Cao
paper_authors.status: admin_assigned
paper_authors.statusLastChangedAt: 2024-10-14T08:42:56.210Z
paper_authors._id: 6619aa776e681d5e88f9d95a
paper_authors.avatarUrl: https://cdn-avatars.huggingface.co/v1/production/uploads/6619aa776e681d5e88f9d95a/fJ_HBq87VfdtJCu6RAUEw.jpeg
paper_authors.fullname: Qingqing Cao
paper_authors.type: user
paper_authors.user: qicao-apple
paper_authors._id: 670c99dda88de61073ac57d0
paper_authors.name: Chenfan Sun
paper_authors.status: claimed_verified
paper_authors.statusLastChangedAt: 2024-10-17T11:27:46.152Z
paper_authors._id: 658895a1c4b20046639a7e59
paper_authors.avatarUrl: /avatars/f80c3f58efa5d22464590bee7a6041a4.svg
paper_authors.fullname: Chenfan Sun
paper_authors.type: user
paper_authors.user: erpang007
paper_authors._id: 670c99dda88de61073ac57d1
paper_authors.name: Yanzi Jin
paper_authors.status: admin_assigned
paper_authors.statusLastChangedAt: 2024-10-14T08:39:07.212Z
paper_authors._id: 6627307fa39494ebb3a39194
paper_authors.avatarUrl: /avatars/99290642fd6d4a6162f52cfe5d46c648.svg
paper_authors.fullname: Yanzi Jin
paper_authors.type: user
paper_authors.user: yjin25
paper_authors._id: 670c99dda88de61073ac57d2
paper_authors.name: Sachin Mehta
paper_authors._id: 670c99dda88de61073ac57d3
paper_authors.name: Mohammad Rastegari
paper_authors.status: admin_assigned
paper_authors.statusLastChangedAt: 2024-10-14T08:38:36.380Z
paper_authors._id: 65d4c7734c2163d05eab4110
paper_authors.avatarUrl: /avatars/b116a9aca7b0a2a4836a29cb4f069d42.svg
paper_authors.fullname: Mohammad Rastegari
paper_authors.type: user
paper_authors.user: mrastegari
paper_authors._id: 670c99dda88de61073ac57d4
paper_authors.name: Moin Nabi
paper_authors.status: admin_assigned
paper_authors.statusLastChangedAt: 2024-10-14T08:35:29.716Z
paper_authors._id: 5fe36e9297b72d62ab9b90cd
paper_authors.avatarUrl: /avatars/edfea7e8c86dcf355241f5a48365b35a.svg
paper_authors.fullname: Moin Nabi
paper_authors.type: user
paper_authors.user: moinnabi
paper_title: KV Prediction for Improved Time to First Token
paper_summary: 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 .
paper_ai_keywords: transformer-based language models prompt processing step KV cache time to first token (TTFT) KV Prediction auxiliary model autoregressive generation pareto-optimal efficiency-accuracy trade-off TriviaQA HumanEval FLOPs Apple M2 Pro CPU