Prefill worked example — Q, K, V step by step

The key thing to see is that during prefill, we compute Q, K, and V for every prompt token at once, then form a whole attention matrix in one large parallel operation.

Let's use a deliberately tiny Transformer.

Suppose the prompt is just three tokens:

"The cat sat"

We'll pretend each token has a 2-dimensional representation:

Stack them into a matrix:

Rows = tokens.

1. Compute Q, K, and V¶

A real Transformer has learned matrices \(W_Q,W_K,W_V\).

For our toy example, suppose:

During prefill, the GPU does these as matrix multiplications:

Notice the important bit: it doesn't separately do "The", then "cat", then "sat".

It feeds the whole matrix \(X\) through.

Queries¶

Since \(W_Q\) is the identity:

So:

The → Q = [1,0]
cat → Q = [0,1]
sat → Q = [1,1]

Keys¶

gives:

So:

The → K = [1,1]
cat → K = [0,1]
sat → K = [1,2]

Values¶

gives:

So:

The → V = [1,0]
cat → V = [1,2]
sat → V = [2,2]

At this point prefill has produced:

             Q        K        V

The        [1,0]    [1,1]    [1,0]
cat        [0,1]    [0,1]    [1,2]
sat        [1,1]    [1,2]    [2,2]

2. Every query compares itself against every key¶

Now we calculate:

This is the part that creates the attention scores between tokens.

For example, consider the query belonging to "sat":

Compare it with "The"'s key:

Compare with "cat":

Compare with "sat":

So "sat" gets:

              The   cat   sat

sat query      2     1     3

But the clever thing about prefill is that all nine comparisons happen as one matrix multiplication:

giving:

Interpret it as:

Each row asks:

How much should this token attend to each token?

3. But there's a causal mask¶

An autoregressive LLM isn't allowed to look into the future.

When training or prefilling:

• "The" may see only "The"
• "cat" may see "The" and "cat"
• "sat" may see "The", "cat" and "sat"

So we apply a mask:

The upper-right half is blocked.

This is an important subtlety.

Even though all tokens are being computed in parallel, each token behaves mathematically as though it only knows its past.

So parallelism does not violate autoregressive causality.

4. Scale and softmax¶

Normally the score is divided by:

Here \(d_k=2\), so:

Let's just focus on "sat".

Its raw scores were:

Scaled:

approximately:

Softmax turns those into attention weights of approximately:

So "sat" is effectively saying:

28% attention → The
14% attention → cat
58% attention → sat

Again, these numbers don't literally mean linguistic importance. This is just one tiny fake attention head.

5. Use those weights to mix the Values¶

Remember the value vectors:

The → [1,0]
cat → [1,2]
sat → [2,2]

The output for "sat" becomes approximately:

Calculate each contribution:

giving approximately:

So "sat" entered the attention operation as:

and comes out with a contextual representation based partly upon information from the previous tokens:

That's essentially the magic of self-attention.

Now look at prefill as a whole¶

For "The cat sat", we haven't performed three independent sequential passes.

Instead, the GPU can roughly do:

             ┌─────────────┐
"The" ───────┤             │
"cat" ───────┤   X matrix  │
"sat" ───────┤             │
             └──────┬──────┘
                    │
          ┌─────────┼─────────┐
          ↓         ↓         ↓
         XWq       XWk       XWv
          ↓         ↓         ↓
          Q         K         V
           \        |        /
            \       |       /
             └── QKᵀ ──────┘
                   ↓
              causal mask
                   ↓
                softmax
                   ↓
                  × V
                   ↓
       contextualized vectors
       for ALL prompt tokens

That's why prefill parallelizes so nicely.

And here's where the KV cache appears¶

Once prefill finishes, we have computed the K and V vectors for:

The
cat
sat

We save them:

and:

Now suppose the model chooses:

"down"

as the next token.

During decode, we don't redo Q/K/V for "The cat sat".

We calculate only:

Qdown
Kdown
Vdown

Then:

attends against:

and uses the corresponding cached values.

Afterward, we append:

to the KV cache.

Then generate the next token.

One subtle but very important point¶

You might wonder:

If during prefill we ultimately only need the model's prediction after the last prompt token, why bother computing attention outputs for all earlier tokens?

Because Transformers are layered.

Suppose there are 80 layers.

At layer 1, "cat" computes a contextualized representation based on "The".

At layer 2, "sat" may attend to the layer-1 representation of "cat", which already contains information from "The".

So information propagates through the sequence as it moves upward through layers.

You need the representations—and the K/V entries—for every token at every layer.

A real KV cache therefore isn't just:

token → K,V

It's closer to:

Layer 1:
   token1 K,V
   token2 K,V
   ...
Layer 2:
   token1 K,V
   token2 K,V
   ...
...
Layer 80:
   token1 K,V
   token2 K,V
   ...

That is why a long context across a large model can consume many gigabytes of KV-cache memory.

The mental model¶

Prefill is essentially constructing a searchable memory of your prompt.

For every token and every layer, the model creates:

• K: an address / description of what this token contains
• V: the information available at that address

Then during generation, each new token produces a Q:

"What information do I need right now?"

and searches those cached keys to retrieve a weighted mixture of values.

That's slightly oversimplified, but it's a surprisingly useful way to think about what the KV cache is doing.