Preloading Knowledge Into a Model Instead of Retrieving It (Part A)

Recap

In Part 12, we found that RAG latency does not live where most people optimize it.

Why RAG Latency Is a Prefill Problem, Not a Retrieval Problem

Part 12: Why prefill dominates RAG latency, and how to reuse KV caches that share no prefix, with implementations

Daily Dose of Data ScienceAvi Chawla

Retrieval takes tens of milliseconds. The model then reads the retrieved chunks, and that reading step can take seconds.

We separated the two phases of inference and saw why.

• Prefill processes every input token in parallel and is compute-bound, producing the key-value cache.
• Decode then generates one token at a time, reusing that cache, and is memory-bandwidth-bound.

We then examined prefix caching, which is considered the most standard fix.

Serving engines split the cache into fixed-size blocks and hash each block together with its preceding blocks. That rule guarantees any reused block is bit-identical to a freshly computed one. It also means two requests retrieving the same documents in a different order share nothing at all.

Finally, we took apart the obvious repair, which is precomputing each chunk's cache once and stitching the pieces together.

It fails for three separate reasons. Position is encoded into the key vectors and must be corrected by rotation. Cross-attention is missing entirely, because a chunk prefilled alone never saw the chunks before it. Attention sinks appear at the start of every chunk instead of once at the start of the input.

We built all of this from scratch and watched it appear in real tensors. Layer 0 matched the joint prefill exactly, while deeper layers diverged, and 12 percent of tokens carried 67.3 percent of the total error.

That work assumed retrieval happens, then asked how to make reading cheaper.

In this chapter, we will explore what happens when we move the work earlier still. We'll understand how a knowledge base can be processed once, before any query arrives, and stored as a reusable artifact.

This is the first of three parts on preloading. Here we cover the idea itself, along with the two things that bound it, which are how much context a model can actually use and what a cache costs when you rent it from a provider.

Part 14 takes up compression, and why most published methods cannot run before a query exists.

Part 15 covers the two approaches that require training, and what running any of this in production involves.

As always, every notion will be explained through clear examples and walkthroughs, and each part ends with runnable experiments on a real model.

Let's begin!

The focus of this part

Part 12 optimized the reuse of work already done. A request arrived, the model read some chunks and processed them, and we tried to avoid processing them again.

There is an earlier point to intervene at this stage.

In several cases, a knowledge base often does not change between queries.

For instance, a product manual, a legal contract, a codebase, or an internal policy document is the same text for every user who asks about it.

Nothing forces the pipeline design to wait for a query before processing that text.

If the model reads the corpus once, before any user shows up, the resulting cache can be stored and reused. Every query then skips straight to decoding. Retrieval, chunking, and embedding disappear from the request path entirely.

This idea has a name in the literature, cache-augmented generation, and we will get to it shortly.

But treating it as a single technique misses what is actually going on. Preloading involves a family of approaches that differ in one respect, which is how much processing you do offline and what you keep afterwards.

The Mooncake team, who built the serving platform behind the Kimi chatbot, titled their FAST 2025 paper around exactly this trade. They describe it as trading more storage for less computation.

The spectrum

Before going into any single approach, it helps to see the whole range laid out. Each step down this list moves more work offline and keeps a different kind of artifact.

At the top, we have ordinary retrieval, where nothing is precomputed and every query costs a full prefill.

Below that, we have prefix caching, which keeps whatever happened to be computed for an earlier request. We covered this in Part 12 and saw its limitations.

Next comes naive preloading, where the entire corpus is prefilled once and its raw cache is stored. This is where cache-augmented generation belongs.

Then modular preloading, where each passage is prefilled separately so it can be assembled in any order. This requires either accepting the cross-attention loss from Part 12 or changing the model so the loss no longer matters.

Then compressed preloading, where the stored cache is shrunk before it is saved. This turns out to have a requirement that rules out most published compression methods, and that requirement is the most interesting finding in this chapter.

At the bottom is trained preloading, where the corpus is distilled into a small cache through an actual training run.

With the map in place, let's work through it from the top.

Naive preloading

The simplest version of the idea simply takes the entire knowledge source, run one prefill pass over it, and keep the resulting cache.

The approach was formalized in the following paper:

The authors describe preloading all relevant resources into the model's context and caching the runtime parameters, so that inference can proceed directly from the preloaded cache without retrieval.

The three phases

The lifecycle has three steps, and each one maps onto the details we discussed in Part 12.

Step 1: Preload

The corpus is formatted as a single input and passed through the model once. Every layer produces keys and values for every token, exactly as in any prefill.

The cost here is a normal prefill cost, and because attention is quadratic in sequence length, it is not small. The benefit is that it is paid once rather than per query since all queries use the same cache generated from prefill.

Step 2: Persist

The resulting cache is serialized and written to storage. The artifact is a set of tensors, one key tensor and one value tensor per layer, and its size follows directly from the cache memory formula we derived in Part 12.

Let's revisit the formula we discussed in the previous part because every storage decision in this chapter rests on it. It tells you how many bytes a preloaded corpus becomes.

Let's consider the symbols one at a time:

• $M$ is the size of the stored cache in bytes.
• The leading 2 counts the key tensor and the value tensor, since every cached token needs both.
• $L$ is the number of transformer layers, because each layer keeps its own cache.
• $H$ is the number of key-value heads and $d_{h}$ is the dimension of each head, so their product is how many numbers one token contributes at one layer.
• $T$ is the number of tokens in the corpus.
• Finally, $b$ is the bytes used per number, which is 4 for float32 and 2 for float16.

Looking at the structure, everything except $T$ is fixed the moment you choose a model.

All those terms collapse into a single per-token constant, and the total is that constant multiplied by corpus length. The formula is linear in how much text you preload, with no dependence on what the text says.

The intuition is that a preloaded cache charges a fixed cost per token of corpus.

Doubling the corpus doubles the artifact. Halving the precision halves it. A page of dense legal text and a page of filler will cost exactly the same, because only the token count matters.

There are two practical consequences of this formulation:

• The first is that the artifact can be sized before any code is written, by reading $L$, $H$, and $d_{h}$ off the model's config file.
• The second is that models using grouped-query attention, where many query heads share a smaller set of key-value heads, produce far smaller caches, because $H$ counts key-value heads rather than query heads.

We will put real numbers through this formula in the hands-on section, where a 512-token corpus produces an artifact the formula predicts to within 0.2 percent of its actual size on disk.

Step 3: Infer

When a query arrives, the stored cache is loaded back into memory, the query tokens are appended, and decoding proceeds. The corpus is never read again.

How preload can improve answers

Speed is the most obvious benefit when you preload the entire context. But the increase in answer quality is that's evident so let's understand this now.

A retrieval pipeline can fail by fetching the wrong passage.

That failure mode does not exist when the whole corpus is present in the context, because there is no selection step to get wrong.

The CAG authors position this directly against retrieval-based methods, noting that precomputed caching approaches which still rely on retrieval remain vulnerable to the retrieval failures inherent to those systems.

There is another effect that connects to the Part 12 of this course.

Chunking a document destroys the relationships between its parts.

But a preloaded corpus was processed in a single joint prefill step, so every token attended to every earlier token normally.

None of the cross-attention loss we measured in Part 12 applies here, because nothing was ever prefilled in isolation.

That makes preloading particularly suited to questions whose answers span distant parts of a document, which is exactly where chunk-based retrieval struggles.

The trade-off, however, is that all of this holds only while the corpus fits comfortably inside the context window.

Cost economics of preloading

At this stage, let's see how many queries a preloaded corpus needs to see before the offline work is worth doing.