KV Prediction for Improved Time to First Token (2410.08391v1)

Abstract: Inference with transformer-based LLMs 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 .

• The paper presents KV Prediction, which uses a smaller auxiliary transformer to quickly approximate the KV cache and reduce time to first token.
• The method achieves significant efficiency gains with 15%-50% accuracy improvement on TriviaQA and 30% better results on HumanEval tasks.
• The approach enhances performance on resource-constrained devices and paves the way for future integration with quantized or pruned models.

KV Prediction for Improved Time to First Token

The paper, "KV Prediction for Improved Time to First Token," addresses a prevalent challenge in transformer-based LLMs, particularly concerning the latency induced during the initial token generation phase, also known as "time to first token" (TTFT). This delay is pronounced on edge devices, where computational resources are limited. The authors propose a methodology termed KV Prediction, designed to mitigate this issue.

Methodology

KV Prediction leverages a smaller auxiliary transformer model to swiftly process input prompts and generate an approximated KV cache, which the larger base model subsequently uses for token generation. This approach allows for significant reductions in TTFT by effectively offloading initial computations to a less resource-intensive model. Once the KV cache is predicted, the base model operates without further dependency on the auxiliary model, thereby minimizing any runtime overhead.

Results and Observations

The authors report notable efficiency gains without substantial sacrifices in prediction accuracy. For instance, on the TriviaQA benchmark, the proposed method demonstrates a 15%-50% accuracy improvement relative to TTFT FLOPs budgets. The improvement is also evident in the HumanEval python code completion task, where up to 30% better accuracy is achieved over fixed TTFT FLOPs budgets. These enhancements are also translated into empirical speedups on an Apple M2 Pro CPU, demonstrating the practical applicability of the proposed solution.

Implications

From a theoretical perspective, this approach highlights the potential of leveraging auxiliary models to predict and approximate intricate model components, such as the KV cache, which traditionally require extensive computational resources. On a practical level, the reduction in TTFT translates to enhanced user experiences, particularly on devices with constrained processing capabilities.

Future Directions

Looking ahead, the implications of this research suggest several avenues for exploration. There is potential for further refinement of the linear prediction models for the KV cache, possibly exploring non-linear mappings to enhance accuracy retention. Additionally, the compatibility of this approach with quantized models or pruned architectures could yield even greater efficiency gains, broadening the scope of its applicability.

In summary, KV Prediction offers a compelling strategy for addressing TTFT in transformer models, providing a balanced trade-off between computational efficiency and model accuracy. The release of their codebase further facilitates replication and exploration, underscoring the paper's contribution to the growing discourse on efficient AI inference.