Prompt Cache: Modular Attention Reuse for Low-Latency Inference (2311.04934v2)

Abstract: We present Prompt Cache, an approach for accelerating inference for LLMs (LLM) by reusing attention states across different LLM prompts. Many input prompts have overlapping text segments, such as system messages, prompt templates, and documents provided for context. Our key insight is that by precomputing and storing the attention states of these frequently occurring text segments on the inference server, we can efficiently reuse them when these segments appear in user prompts. Prompt Cache employs a schema to explicitly define such reusable text segments, called prompt modules. The schema ensures positional accuracy during attention state reuse and provides users with an interface to access cached states in their prompt. Using a prototype implementation, we evaluate Prompt Cache across several LLMs. We show that Prompt Cache significantly reduce latency in time-to-first-token, especially for longer prompts such as document-based question answering and recommendations. The improvements range from 8x for GPU-based inference to 60x for CPU-based inference, all while maintaining output accuracy and without the need for model parameter modifications.

• The paper introduces a modular approach to reuse precomputed attention states in LLM inference, reducing redundant computations on repetitive prompt segments.
• It leverages a Prompt Markup Language to uniquely encode reusable modules, achieving significant speed improvements across GPU and CPU architectures.
• Empirical evaluations show that the method maintains output accuracy while optimizing resource usage, enabling efficient deployment in resource-constrained environments.

Modular Attention Reuse in LLMs: A Study on Prompt Cache

The paper "Prompt Cache: Modular Attention Reuse for Low-Latency Inference" introduces an innovative approach to reducing computational overhead in LLM inference through reusable attention states. The authors, affiliated with Yale University and Google, explore the potential of reusing precomputed attention states across different inference requests. Their work centers on a system called Prompt Cache, which efficiently accelerates LLM inference by leveraging text segments shared across prompts.

Overview of Prompt Cache

Prompt Cache capitalizes on the repetitive nature of certain prompt segments, such as system messages, prompt templates, and contextual documents. The authors recognize that many applications, including legal analysis, healthcare, and education, generate prompts with overlapping components. By caching the attention states of these common segments on the server, Prompt Cache reduces the need for recomputation, cutting down on processing time and resource usage.

The system employs a Prompt Markup Language (PML) schema to define modular, reusable text segments, termed as prompt modules. During inference, attention states are precomputed for these modules and stored in memory for future reuse. This ensures efficiency in handling similar prompt sequences, markedly improving time-to-first-token (TTFT) latency. Empirical evaluations with various LLMs demonstrate that Prompt Cache achieves up to 8× improvement for GPU-based inference and 60× for CPU-based inference without sacrificing output accuracy.

Technical Implementation

Prompt Cache incorporates two main ideas to overcome the challenge of position-dependent attention states in transformers. First, it explicitly structures prompts using the PML, allowing for unique positional encoding of each module. Second, it benefits from the empirical observation that attention states with non-contiguous position IDs still preserve semantic meaning in LLMs. This insight allows for dynamic compilation of attention states, maintaining effectiveness even with changes in module order or prompt customization.

The authors present a prototype implemented with HuggingFace Transformers and assess it with models such as Llama2, Falcon, and MPT. The implementation accommodates both CPU and GPU memory architectures for storing cached states, striking a balance between latency reduction and memory overhead.

Numerical Results and Implications

Evaluations conducted using the LongBench suite reveal significant latency reductions across various datasets and platform configurations. The results indicate that Prompt Cache provides substantial efficiency gains, particularly for longer sequences where computational requirements scale quadratically. The paper reports latency reductions consistently demonstrating the quadratic advantage of caching attention states.

Regarding output quality, the authors ensure that Prompt Cache does not compromise the precision of LLM responses. The benchmarking results show comparable, if not slightly improved, performance metrics against baseline methods.

Practical and Theoretical Implications

The proposed Prompt Cache system highlights promising advancements in LLM efficiency, potentially facilitating the deployment of these models in resource-constrained environments. By reducing latency, it enhances user experience, particularly in applications requiring rapid interaction. Additionally, this work opens pathways for further research into attention state optimization, extending beyond standard KV Cache mechanisms.

Future developments could see Prompt Cache underpin enhanced LLM serving systems, equipped with cache management and replacement strategies to further leverage both DRAM and HBM resources. Furthermore, integration with retrieval-augmented models or compression techniques for attention states could yield additional efficiency improvements.

Conclusion

Prompt Cache represents a step forward in the optimization of LLM inference, with its modular approach to attention reuse. The paper presents a comprehensive paper of the architecture, demonstrating its practical efficacy and theoretical potential. This method's impact is a significant enhancement in processing efficiency, setting the stage for future innovations in LLM deployment strategies.