Prompt Cache: Modular Attention Reuse for Low-Latency Inference

Before a language model writes its first token, it computes attention keys and values for every token in the prompt. The standard KV cache saves those states within one request so they are not recomputed for each generated token. But across requests, the same system message, document or template is recomputed every time. In Gim, Lin Zhong and colleagues at Yale and Google asked how to reuse those states across prompts. Their Prompt Cache appeared at MLSys 2024.

Two problems stand in the way. Attention states depend on position, because positional encodings are folded into keys and values, so a document's states computed at one position cannot be dropped in at another. And the system has to recognize which parts of a new prompt it has already seen.

The answer is a Prompt Markup Language (PML). A schema declares prompt modules, reusable text segments such as a city guide, a document or a template, and fixes each module's range of position IDs. Modules can have parameters, placeholders of declared length filled at request time, and can be grouped into unions of alternatives that share a starting position. A prompt names its schema, imports modules, fills parameters and adds new text. Prompt Cache precomputes and stores the attention states of every module when the schema is loaded. Serving a prompt concatenates the cached states of imported modules and computes only the parameters and new text.

The method rests on an empirical finding: models still work when position IDs have gaps between modules, as long as relative positions within each module are preserved. The price is an approximation: each module was encoded alone, so its tokens never attended to other modules, much like a local attention mask. On LongBench with Llama 2, MPT and Falcon, accuracy stayed close to full prefill on average, while time to first token dropped sharply, especially for CPU inference.

Customer prompts are repetitive: the same long system prompt, the same policy manual, the same product catalog section, the same few-shot examples. Reusing their computation is one of the cheapest wins in serving.

Prefix caching is the safe default. Serving engines and API providers commonly reuse states for an identical prompt prefix, which is exact. Put the stable, shared parts at the start of the prompt and the variable parts at the end to benefit.

Modular reuse goes further and is approximate. Prompt Cache reuses modules anywhere in the prompt and in any combination, at the cost of modules not attending to each other. Fine for independent context such as a style guide plus a document; risky when the answer needs to connect facts across documents.

Memory is the binding constraint. Cached states take 0.5 MB per token for a 7B model at FP16. A few long documents fit in GPU memory; hundreds need CPU memory, where the speedup shrinks to 1.5 to 3 times because of copy time.

Time to first token is not total latency. Generation speed per token is unchanged. For short answers to long prompts, the gain is large; for long answers, it is a smaller share of the total.