All posts

SCIENCE · · 6 min read

Can We Use the KV Cache for Retrieval?

By Idan Schwartz

Without a KV cache, every generation would begin by processing the entire context again: the system prompt, tool definitions, conversation history, documents, and images. Thousands of unchanged tokens would be recomputed before the model reached the user's new question.

The KV cache avoids this repetition. It stores the intermediate attention state of the existing prefix, allowing the model to process only the tokens added afterward.

What a KV cache is

At every transformer layer, each token produces two vectors: a key and a value.

Later tokens compare their query vectors with the stored keys to decide which earlier tokens matter. They then read the corresponding values.

The KV cache stores these key-value pairs for every processed token and every layer.

Each prompt token produces a key and a value vector in every layer. Those pairs are cached and stay valid as new tokens arrive, so only the new question is fresh work. For granite-3.2-2b in fp16, that is 80 KB per token, or 938 MB for a 12,000-token prompt.

Under ordinary causal attention, adding a new token does not change the keys and values of the tokens before it. Once a prefix has been processed, its cached state can be reused as more tokens are appended.

This is an important constraint. A KV cache is primarily a prefix cache. The representation of a token depends on everything that appears before it and, in many models, on its position. Change the prefix or move the token, and the cached state may no longer be valid.

Keeping it is worth 69x

We took a 12,000-token prefix and ran four queries against it.

First, we discarded the cache between queries, forcing the model to rebuild the entire prefix each time. Then we processed the prefix once and reused its cached state.

Four requests, one 12,000-token prefix Total time
Rebuilding the cache each time 159.6s
Building it once and keeping it 2.3s

At this length, building the cache pays for itself after roughly one additional query. The code and every raw result are public.

None of this is new. Reusing attention state across separate requests, rather than only while generating one response, was already demonstrated in 2023.

Images are an interesting case

We placed one product image at the beginning of the prompt and asked Qwen2-VL-2B four different questions about it.

The image occupied 576 of the 597 prompt tokens, or 96.5% of everything the model had to process. Compared with the image, the questions were almost free.

One image, four queries Per query
Re-encoding the image every question 1,237 ms
Encoding once, reusing it 217 ms

Because the image appears before the question, it forms a stable prefix. The image does not change when the shopper asks about its price, color, or material, so its visual representation and prefix state can be computed once and reused.

That invites an obvious question: if a product image remains unchanged, why wait for a shopper to encode it?

Index the whole catalog

For a single-product query, the basic setup is straightforward:

  1. Encode each product image ahead of time.
  2. Store its reusable representation under the product ID.
  3. Retrieve the relevant product when a question arrives.
  4. Place the question after the stored image prefix.

The cache stops being only a per-conversation optimization. It becomes an index of products that the model has already processed.

The arithmetic determines how far this can scale.

In our setup, one image produces 576 tokens, and Qwen2-VL-2B uses 28 KB of KV memory per token. That puts one product at roughly 15.8 MB:

Catalog size KV at fp16 At int8
1,000 products 15.4 GB 7.7 GB
10,000 products 154 GB 77 GB
100,000 products 1.5 TB 769 GB

These figures cover only the KV state. Model weights, active requests, temporary activations, and memory-management overhead still need space.

A thousand products can fit on one high-memory GPU at int8. Ten thousand require a node rather than a single card. Beyond that, the cache needs multiple storage tiers.

That infrastructure is already being developed. Recent work keeps part of a large KV pool in CPU memory and streams the required pages back to the GPU, overlapping transfers with computation. On the vision side, work on large-scale reranking precomputes and compresses vision tokens before online inference.

The two lines of work address the same pressure: encoding is expensive, the resulting state is reusable, and keeping all of it on the GPU is not always practical.

There is one limit to the database analogy. Independently cached KV blocks cannot generally be concatenated in arbitrary orders. The state of a later token normally depends on the tokens and positions before it.

A single retrieved product works naturally as an initial prefix. Combining several independently cached products requires more care, such as fixed prompt slots, position-aware caching, or partial recomputation.

Hosted APIs solve only part of this

Hosted providers support prompt caching, and applications should use it whenever requests share a large stable prefix.

The details differ by provider, but the general pattern is the same: place static material first, identify or mark the reusable prefix, and send changing content afterward.

# OpenAI: keep static content first and use a stable cache key
prompt_cache_key = "tenant-42"

# Anthropic: mark the end of the reusable prefix
"cache_control": {"type": "ephemeral", "ttl": "1h"}

# Google: create the cache object once, then reference it
cache = client.caches.create(...)
# Later: cached_content=cache.name

Note that you should also read the cached-token counters in the usage response. A provider doing nothing can otherwise look exactly like one successfully reusing the prefix.

Hosted caching is still designed around prompt reuse. The provider controls retention, minimum cacheable length, routing, eviction, and the physical representation. You do not receive an independently addressable KV object that can be kept permanently or moved between memory tiers.

A catalog index has different requirements. It may contain thousands of product states, remain useful for months, serve every visitor, and retrieve entries in an arbitrary order. That is not the usual abstraction offered by a hosted prompt cache.

Open weights change the shape of the problem

With an open-weight model, the cache is an object you control.

You choose how long it remains available, whether it sits in GPU or CPU memory, when it is evicted, and whether it is stored at 16, 8, or fewer bits. One cached prefix can be shared across many concurrent requests rather than rebuilt separately for each conversation.

Model architecture also becomes a deployment decision. Models using grouped-query or multi-query attention store fewer key-value heads than models using standard multi-head attention. The grouped-query attention configuration can therefore have a large effect on cache size.

You own the memory management in return.

That is less exotic than it sounds. Modern serving stacks already divide KV memory into blocks, allocate those blocks dynamically, and share them across requests. PagedAttention is one well-known example.

The catalog case extends the same idea across longer periods and larger collections of reusable content.

Aigency's take

Treating the KV cache only as a temporary optimization for one conversation is too narrow.

A site contains large amounts of information that change slowly: product images, descriptions, policies, page content, and tool definitions. Processing all of it only after a shopper asks a question repeats work that could have been completed earlier.

At Aigency, we read the site before shoppers arrive. We crawl its pages, build retrieval indexes, preprocess stable text and images, and assemble reusable parts of the prompt ahead of time. When a question arrives, the system can focus its online compute on the new part of the interaction.

Not everything should be stored as KV. Some information belongs in a conventional retrieval index. Some is better kept as compressed visual features. Some prefixes are worth retaining on the GPU, while others should stay in CPU memory or be reconstructed when needed.

The practical problem is deciding what to precompute, which representation to preserve, and where to keep it.

The direction is clear: encode stable content once, retain what is useful, and avoid making the model reread the same site for every visitor.

https://aigency.ai/blog/kv-cache-for-retrieval

Talk to us.

Own where the AI answer lands. Drop your email and we'll be in touch to book a demo.

[email protected] · Run the free scan