PaperlayerView exact revision on arXiv

Source-linked AI summary

Prompt Cache: Modular Attention Reuse for Low-Latency Inference

In Gim, Guojun Chen, Seung-seob Lee, Nikhil Sarda, Anurag Khandelwal, Lin Zhong

arXiv:2311.04934v2cs.CLcs.AI

TL;DR

LLM prompts often repeat system messages, templates, and documents, raising the question of how to reuse attention states across requests while handling positional dependence and segment recognition. Prompt Cache defines reusable prompt modules in a schema, precomputes their states, and reuses them in derived prompts. Evaluations report TTFT reductions of 1.5× to 10× for GPU inference with GPU-memory modules and 20× to 70× for CPU inference, with comparable accuracy to baseline.

  • Problem

    Repeated system messages, templates, and documents create an opportunity to reuse attention states across prompts, but reuse must handle position dependence and efficiently recognize cached segments.

  • Method

    Prompt Cache uses Prompt Markup Language schemas to define prompt modules, assign position IDs, precompute their attention states, and reuse them in prompts derived from the schema.

  • Results

    1.5× to 10× TTFT reductions are reported for GPU inference with GPU-memory modules and 20× to 70× for CPU inference, with no significant accuracy loss.

  • Takeaways & Limitations

    Prompt Cache extends KV Cache from single-prompt reuse to modular reuse across multiple prompts while preserving comparable output accuracy to baseline.

  • Takeaways & Limitations

    The study mainly focuses on techniques for modular attention reuse rather than a broader LLM serving system.

Abstract

from arXiv · show

We present Prompt Cache, an approach for accelerating inference for large language models (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.

1 INTRODUCTION

Prompt Cache targets repeated text in LLM prompts by modularizing reusable segments and reusing their attention states across prompts. Its schema addresses positional reuse and segment recognition, while prototype evaluations report substantial TTFT reductions without significant accuracy loss.

  • Motivation and approach: Prompt overlap commonly arises from system messages, documents, and reusable prompt-engineering templates.These patterns occur in applications including legal analysis, healthcare, education, robotics, and tool learning.
  • Motivation and approach: Prompt Cache precomputes attention states for frequently revisited prompt segments and reuses them when those segments recur.This extends attention-state reuse beyond a single prompt.
  • Motivation and approach: Prompt Cache uses Prompt Markup Language to define reusable text segments as prompt modules and assign them unique position IDs.The schema makes segments recognizable and supports positional accuracy during reuse.
  • System operation: Prompt Cache computes schema-module attention states when needed, then reuses them for prompts derived from the same schema.The prototype is implemented on HuggingFace transformers and supports Transformer architectures compatible with KV Cache.
  • Evaluation: Prompt-module storage trades GPU capacity against CPU host-to-device copying overhead.CPU memory can scale to terabyte levels, whereas GPU memory has limited capacity.
  • Evaluation: 1.5× to 10× TTFT reductions occur for GPU inference with GPU-memory modules, while CPU inference reductions range from 20× to 70×.The evaluation uses LongBench recommendation and question-answering tasks based on multiple documents, without significant accuracy loss.

2 BACKGROUND AND RELATED WORK

LLM generation repeatedly applies self-attention, motivating KV Cache to reuse states within one request. Prompt Cache is positioned as a complementary strategy for modular reuse across requests and inference systems.

  • Autoregressive generation: Autoregressive generation appends each predicted token to the prompt and continues until a stopping condition is met.Stopping may follow a predetermined token count, an end-of-sequence token, or another criterion.
  • Key-Value Cache: KV Cache computes input-token attention states during prefill and reuses them while generating subsequent tokens.For causal language models, this avoids recomputing cached key-value states at every generation step.
  • Key-Value Cache: 1/n reduction in per-step matrix-operation computation is reported for KV Cache, from approximately 6nd2 + 4n2d to 6d2 + 4nd.Here, n is the input length and d is the hidden dimension size.
  • Related work: Related KV Cache work studies memory management, pruning, compression, similarity-based reuse, and prefix sharing across requests.These approaches include paged attention and methods that reuse states based on embedding similarity.
  • Related work: Prompt Cache is an orthogonal optimization that can augment multi-GPU inference and high-performance attention kernels, with potential throughput benefits.The supplied passage describes these as compatible system-level optimizations.

3 DESIGN OF PROMPT CACHE

Reusing attention states across prompts requires handling position dependence and recognizing cached segments efficiently. These requirements define the design problem addressed by Prompt Cache.

  • Design challenges: Attention states can be reused across prompts only when a text segment appears at the same position.This constraint arises from positional encoding in Transformers.
  • Design motivation: Prompt overlap includes identical system messages and repeated document sets in legal and medical applications.These recurring segments motivate reuse across multiple inference requests.

3.1 Overview

Prompt Cache addresses cross-prompt attention-state reuse by explicitly structuring reusable modules and preserving their positional relationships. It precomputes module states, reuses them across schema-derived prompts, and assembles cached and newly computed states for inference.

  • Overview: At serving time, cached module states are combined with states computed for parameters and new text to form the full prompt state.Figure 2 illustrates module import, parameter substitution, and concatenation of the resulting states.
  • Overview: Cross-prompt reuse requires handling position-dependent attention states and recognizing cached text segments in incoming prompts.Transformer positional embeddings constrain reuse when shared segments occur at different input positions.
  • Overview: PML explicitly represents reusable text as modules and assigns them position IDs, enabling recognition and positional alignment during reuse.The approach also relies on LLMs operating with discontinuous position IDs when relative token positions are preserved.
  • Overview: Prompt Cache precomputes attention states for modules in a schema, then reuses them in prompts derived from that schema.The schema defines the reusable structure and provides the basis for sharing cached states across prompts.

3.2 Prompt Markup Language (PML)

PML defines schemas and schema-derived prompts around reusable modules, allowing users to import cached content while adding parameters, alternatives, nested structures, and new instructions.

  • 3.2.1 Schema vs. Prompt: PML defines schemas that specify prompt modules, their relative positions, and their hierarchies.Schema-unspecified text is treated as anonymous content that remains included in derived prompts.
  • 3.2.1 Schema vs. Prompt: Prompts reference a schema, import selected modules, and add additional instructions whose attention states are computed rather than cached.The prompt interface uses the schema attribute and module tags to identify reusable content.
  • 3.2.2 Maximizing Reuse with Parameters: Parameterized modules use named placeholders with specified lengths, allowing runtime values to customize reusable prompt content.Parameter values are not cached; the module’s reusable structure is retained around them.
  • 3.2.2 Maximizing Reuse with Parameters: Parameters can create buffers at module boundaries, while also supporting inline modifications that preserve reuse for templated prompts.They allow modules differing at defined locations to share cached attention states.
  • 3.2.3 Other Features: Union modules encode mutually exclusive alternatives that share a starting position ID and can support organization, position-ID conservation, and prefetching.A prompt selects one module from the union, such as one reader-profile description.
  • 3.2.3 Other Features: Parameters customize modules inline, whereas unions organize alternatives and use position IDs more efficiently.The two constructs are also encoded differently.
  • 3.2.3 Other Features: PML supports nested modules and unions, allowing hierarchical prompt composition.Nested modules are imported as components within larger modules.
  • 3.2.3 Other Features: Dedicated tags support instruction-tuned LLM templates, and a Python API can convert prompt programs into PML schemas.These features reduce manual formatting and schema-writing effort across supported prompt-program constructs.

3.3 Encoding Schema

Prompt module encoding computes and stores module attention states with schema-derived position IDs, while parameter handling and attention masking determine reuse behavior and output consistency.

  • 3.3 Encoding Schema: Prompt module encoding extracts module tokens, assigns position IDs from schema locations, and computes the corresponding attention states for storage.A module beginning after preceding modules of lengths 50 and 60 receives starting position ID 110.
  • 3.3 Encoding Schema: Modules in unions share starting positions, so position allocation uses the largest child module length.This preserves the shared-position layout required by union alternatives.
  • 3.3 Encoding Schema: Parameterized modules replace arguments with fixed-length <unk> tokens during encoding and later reuse their recorded position IDs for supplied values.Shorter supplied values are permitted because trailing whitespace does not alter semantics.
  • 3.3 Encoding Schema: Prompt Cache masks attention across independently encoded modules, which can improve or degrade quality depending on their semantic independence.Semantically dependent modules may require scaffolding to share attention spans.
  • 3.3 Encoding Schema: Scaffolding stores jointly encoded module states in addition to individual states, trading extra memory for output consistency when modules are jointly imported.This option is presented for applications that need deterministic results.

3.4 Cached Inference

Cached inference parses schema-derived prompts, retrieves attention states for imported modules, computes states for uncached content, and concatenates them to replace prefill computation.

  • 3.4 Cached Inference: Prompt Cache validates schema alignment and imported modules before serving a prompt.This parsing step ensures that the prompt conforms to its claimed schema.
  • 3.4 Cached Inference: The serving path retrieves cached states for imported modules and computes states for parameters and new text segments.Cached inference separates reusable content from prompt-specific content.
  • 3.4 Cached Inference: Prompt Cache concatenates imported module states with newly computed states to produce the attention states for the entire prompt.This assembled state replaces the standard prefill operation.
  • 3.4 Cached Inference: Concatenating cached KV tensors for imported modules requires only memory copying, and transformer permutation invariance makes concatenation order irrelevant.For modules A and B, the combined tensors are formed by concatenating their respective key and value states.
  • 3.4 Cached Inference: Paged attention can share prompt-module pointers across batched prompts derived from one schema, reducing duplicated KV-cache storage.The optimization targets repeated modules such as shared system prompts.

4 IMPLEMENTATION

Prompt Cache is implemented as a prototype supporting CPU and GPU memory, discontinuous position IDs, and efficient attention-state concatenation.

  • The prototype uses HuggingFace Transformers and PyTorch, reuses existing model weights, and supports prompt modules in CPU or GPU memory.It comprises approximately 3K lines of Python code.
  • GPU deployments can load prompt modules from CPU memory as needed, trading host-to-device copy overhead for access to larger memory capacity.CPU memory can scale to terabyte levels, whereas GPU memory has limited capacity.
  • Prompt Cache requires discontinuous position-ID support, integrated through minor model-specific modifications.The implementation estimates approximately 20 additional lines of code for each LLM.
  • Embedding-table models require no positional changes, while RoPE and ALiBi use lookup tables keyed by position IDs.The lookup tables retrieve rotation matrices for RoPE and adjust bias matrices for ALiBi.
  • A buffered concatenation operator reuses memory when combining attention states, avoiding redundant allocations from PyTorch’s default contiguous-tensor behavior.This optimization improves the memory footprint of prompt-module concatenation.

5 EVALUATION

The evaluation measures latency, output quality, memory overhead, and application fit across LongBench tasks, CPUs, GPUs, and several LLM architectures. Prompt Cache substantially reduces TTFT while preserving comparable accuracy to the baseline.

  • The evaluation compares Prompt Cache with regular KV Cache on TTFT, output quality, memory overhead, and application suitability.TTFT is used because both methods have the same decoding latency after the first token.
  • The study uses LongBench recommendation and question-answering tasks, several open-source LLMs, two CPUs, and three NVIDIA GPUs.The tested models include Llama2, CodeLlama, MPT, and Falcon, with models fitting within a single 40 GB GPU.
  • CPU inference achieves up to 70× latency reduction on the Intel CPU and 20× on the AMD CPU.The reported disparity is associated with the two systems’ memory-bandwidth differences.
  • Prompt Cache produces higher latency for datasets with a larger proportion of uncached prompts, such as TriviaQA.CPU inference benefits more because attention computation is substantially slower on CPUs for longer sequences.
  • Across LongBench datasets and the Llama2, MPT, and Falcon architectures, output accuracy remains comparable to the baseline.The comparison uses deterministic sampling to make outputs with and without Prompt Cache comparable.

5.4 Understanding Latency Improvements

Prompt Cache’s advantage grows with sequence length and model size because memory-copy overhead scales linearly while attention computation scales quadratically. It reduces TTFT but does not change per-token generation latency, and cached modules can reduce memory footprint for shared prompts.

  • Prompt Cache theoretically offers quadratic TTFT improvement over regular KV Cache because memcpy overhead is linear while self-attention computation is quadratic in sequence length.The synthetic evaluation assumes all prompts are cached and compares TTFT using Llama2 7B on one CPU and two GPUs.
  • KV Cache latency increases quadratically with sequence length, whereas Prompt Cache’s memory-copy cost grows linearly.The latency gap expands quadratically and is more pronounced on CPUs than GPUs.
  • At 5K tokens, host-to-host, host-to-device, and device-to-device memcpy latencies are 3.79 ms, 5.34 ms, and 0.23 ms, respectively.
  • Moving from a 7B to 13B model at 3K tokens adds 220 ms for KV Cache but 30 ms for Prompt Cache.The paper attributes this difference to attention complexity scaling with hidden dimension size.
  • Prompt Cache reduces TTFT but leaves token-generation time unchanged, so its end-to-end benefit diminishes as the number of generated tokens increases.On RTX 4090 with Llama 7B and 3K context, TTFT falls from 900 ms to 90 ms while generation remains 32 ms per token.

5.5 Memory Overhead

Prompt Cache supports complex prompt structures while reducing latency as more tokens are cached, but its memory cost scales with cached content and model size.

  • Memory overhead: A 1K-token document requires approximately 180 MB for Falcon 1B and 2.5 GB for Llama 70B.The estimates assume 16-bit floating-point precision.
  • Memory overhead: Hundreds of prompt modules can collectively consume tens of gigabytes, within server-grade GPU memory limits for compact models.
  • Expressive prompt structures: Prompt Cache supports multiple modules, union operations, and parameterized prompts in more complicated prompt structures.These use cases extend beyond the LongBench benchmarks.
  • Code generation: 4× GPU TTFT improvement was achieved for multi-source code generation while the output remained identical.Individual classes such as Unit, Map, Game, and Player were treated as prompt modules.
  • Trip planning: Parameterized trip-planning prompts achieved lower TTFT latency with the same response quality.The schema used an adjustable trip-duration parameter and union modules for destination selection.

6 CONCLUSIONS AND FUTURE WORK

Prompt Cache reuses attention states across LLM prompts through positionally coherent prompt modules. Evaluations report TTFT reductions of up to 8× on GPUs and 60× on CPUs, while future work targets broader serving-system integration.

  • Conclusion: Prompt Cache uses a prompt schema to define reused text segments as modular, positionally coherent prompt modules.Users can incorporate these modules into prompts to reuse attention states.
  • Conclusion: Evaluations indicate TTFT latency reductions of up to 8× on GPUs and 60× on CPUs.
  • Future work: Future work includes GPU cache replacement, memory-overhead reduction, concurrent-request sharing, and applications to latency-sensitive retrieval-augmented generation.The paper identifies these directions as building blocks for future LLM serving systems.
Loading 2311.04934v2…