Attention, KV Caching & Perplexity: How LLM Inference Really Works

Self-Attention, Briefly

Every token position produces three vectors: a Query (what am I looking for?), a Key (what do I offer?), and a Value (what information do I actually carry?). Attention scores every position against every other position by comparing queries to keys:

The √d_kin the denominator isn’t decorative — without it, dot products grow large enough to push softmax into regions with almost no gradient, which stalls learning. Multi-head attention just runs several of these in parallel with different learned projections and concatenates the results, so the model can track more than one kind of relationship between tokens at once.

Every Symbol, With an Example

Take the same 3-token sentence used throughout this page: “The cat sat.” Say each token’s embedding has already been projected down to a 4-dimensional Query, Key, and Value vector (a real model learns these projections; here they’re just small numbers so the arithmetic stays traceable by hand).

• Q (Query), shape 3×4— one row per token, each row is that token asking “which other tokens are relevant to me?” “cat”’s query row is what gets compared against every token’s key below.
• K (Key), shape 3×4— one row per token, each row is that token advertising “here’s what I contain, ask me if this is relevant.” Queries get dotted against these.
• V (Value), shape 3×4 — one row per token, the actual content that gets mixed into the output once attention decides how much each token matters.
• d_k = 4— the width of each Query/Key row. It’s the same number that goes under the square root in the denominator.

Multiplying Q @ K.T gives a 3×3 matrix — row i, column j is how strongly token i’s query matches token j’s key. Divide every entry by √4 = 2, then run softmax across each row so every token’s attention weights over the 3 tokens (including itself) sum to 1:

attention_weights =
[[0.422, 0.155, 0.422],   # "The" attends 42% to itself, 16% to "cat", 42% to "sat"
 [0.155, 0.422, 0.422],   # "cat" attends mostly to "cat" and "sat"
 [0.274, 0.274, 0.452]]   # "sat" attends most to itself

Multiply those weights by V and each token walks away with a new vector — a weighted blend of every token’s Value, in exactly the proportions attention just computed. That blended vector, not the original embedding, is what the next layer actually sees.

The Same Example, in Code

Here’s the entire formula implemented from scratch with NumPy — no framework, so every line maps directly to a piece of the equation above. Running it reproduces the exact numbers shown above.

import numpy as np

# Toy example: 3 tokens ("The", "cat", "sat"), each already projected into
# 4-dimensional Query, Key, and Value vectors. A real model produces these
# via learned weight matrices (W_Q, W_K, W_V) applied to token embeddings —
# here we skip straight to the projected vectors so the attention math
# itself stays easy to trace by hand.
Q = np.array([
    [1, 0, 1, 0],   # "The" — what this token is looking for
    [0, 1, 0, 1],   # "cat"
    [1, 1, 0, 0],   # "sat"
], dtype=float)

K = np.array([
    [1, 0, 1, 0],   # "The" — what this token offers, to be matched against queries
    [0, 1, 0, 1],   # "cat"
    [1, 1, 1, 1],   # "sat"
], dtype=float)

V = np.array([
    [1, 2, 3, 4],    # "The" — the actual content this token contributes
    [5, 6, 7, 8],    # "cat"
    [9, 10, 11, 12], # "sat"
], dtype=float)

d_k = Q.shape[-1]  # 4 — width of each Query/Key row, used in the sqrt(d_k) scale

def softmax(x):
    # Subtract each row's max before exponentiating (numerically stable —
    # the same trick TemperaturePlayground.jsx uses client-side in JS).
    x = x - np.max(x, axis=-1, keepdims=True)
    exp_x = np.exp(x)
    return exp_x / np.sum(exp_x, axis=-1, keepdims=True)

def scaled_dot_product_attention(Q, K, V):
    # Step 1 — QK^T: score every query against every key with a dot product.
    # scores[i][j] = how strongly token i's query matches token j's key.
    scores = Q @ K.T                       # shape: (3 tokens, 3 tokens)

    # Step 2 — /sqrt(d_k): scale down before softmax. Without this, larger
    # d_k makes dot products grow large, pushing softmax toward a near-
    # one-hot output and killing its gradient during training.
    scaled_scores = scores / np.sqrt(d_k)

    # Step 3 — softmax(...): normalize each row into a probability
    # distribution, so every token's attention weights over all tokens
    # (itself included) sum to exactly 1.
    attention_weights = softmax(scaled_scores)

    # Step 4 — (...) @ V: use those weights to build a weighted average of
    # the Value vectors. This IS the output — each token's new
    # representation is a mixture of every token's Value, weighted by how
    # much attention it received.
    output = attention_weights @ V         # shape: (3 tokens, 4 dims)

    return output, attention_weights

output, weights = scaled_dot_product_attention(Q, K, V)

print("Attention weights (each row sums to 1):")
print(np.round(weights, 3))
# [[0.422 0.155 0.422]
#  [0.155 0.422 0.422]
#  [0.274 0.274 0.452]]

print("Output (each row = one token's new, context-mixed representation):")
print(np.round(output, 3))
# [[5.    6.    7.    8.   ]
#  [6.068 7.068 8.068 9.068]
#  [5.711 6.711 7.711 8.711]]

Why Generation Is Sequential

During training, the model sees the entire correct output sequence up front (this is “teacher forcing”), so every position’s prediction and loss can be computed in one parallel pass across the whole sequence. Inference doesn’t have that luxury: token 50 hasn’t been generated yet when the model is deciding token 10, and token 10’s output directly feeds into what token 11 sees. That dependency chain is the entire reason token-by-token generation runs one forward pass at a time — it’s a mathematical constraint of autoregression, not a hardware limitation.

KV Caching

Left naively, every new token would force the model to recompute Key and Value vectors for every token that came before it — quadratic work for a sequence that’s already been processed once. KV caching stores each token’s Key/Value vectors the moment they’re computed, so generating token N+1 only requires computing Q/K/V for the single new token and reusing the cached K/V for everything earlier. This is standard in every production inference engine and is the single biggest reason token-by-token generation is practical at all at conversational speed.

The Naive Way vs. the Cached Way

Generating 4 tokens one at a time means running attention 4 times — once per new token. The naive way reprocesses the entire sequence-so-far at every single step:

• Step 1: project K/V for 1 token.
• Step 2: project K/V for 2 tokens — the 1st one again, plus the new 2nd.
• Step 3: project K/V for 3 tokens — 2 old ones again, plus the new 3rd.
• Step 4: project K/V for 4 tokens — 3 old ones again, plus the new 4th.

That’s 1+2+3+4 = 10 total Key/Value projections to generate just 4 tokens — growth that gets worse quadratically as the sequence gets longer. The cached way only ever projects the newtoken’s K/V and appends it to a running cache, so the same 4 tokens cost exactly 4 projections, one per step, no matter how long the sequence eventually gets.

The Same Comparison, in Code

Both versions below use the identical scaled_dot_product_attention building block from the section above — KV caching doesn’t change the attention math at all, it only changes how much of K and V gets recomputed before that math runs.

import numpy as np

# 4 token embeddings arriving one at a time, as if being generated in a
# loop. In a real model these come from the embedding lookup table
# (see "From Text to Numbers" on the Intermediate page) — here they're
# just small fixed vectors so the projection counts stay easy to verify.
embeddings = np.array([
    [1, 0, 0, 1],   # "The"
    [0, 1, 0, 1],   # "cat"
    [0, 0, 1, 1],   # "sat"
    [1, 1, 0, 0],   # "down"
], dtype=float)

# Shared projection matrices — every token uses the SAME W_Q/W_K/W_V,
# only the input embedding changes per token.
W_Q = np.eye(4)
W_K = np.eye(4)
W_V = np.array([[0,1,0,0],[0,0,1,0],[0,0,0,1],[1,0,0,0]], dtype=float)

d_k = 4

def softmax(x):
    x = x - np.max(x, axis=-1, keepdims=True)
    e = np.exp(x)
    return e / np.sum(e, axis=-1, keepdims=True)

# --- Naive: recompute K/V for every token in the sequence, every step ---
naive_projection_count = 0
for t in range(1, len(embeddings) + 1):
    seq_so_far = embeddings[:t]
    K = seq_so_far @ W_K   # re-projects ALL t tokens, including old ones
    V = seq_so_far @ W_V
    naive_projection_count += t   # t projections done again this step

# --- Cached: project only the new token, reuse everything else ---
K_cache, V_cache = [], []
cached_projection_count = 0
for t in range(len(embeddings)):
    e_t = embeddings[t]
    q_t = e_t @ W_Q                 # new token's query
    k_t = e_t @ W_K                 # new token's key  -- computed ONCE, ever
    v_t = e_t @ W_V                 # new token's value -- computed ONCE, ever
    K_cache.append(k_t)
    V_cache.append(v_t)
    cached_projection_count += 1    # only the new token gets projected

    K_all = np.stack(K_cache)       # old (reused) + new keys
    V_all = np.stack(V_cache)
    scores = (q_t @ K_all.T) / np.sqrt(d_k)
    weights = softmax(scores)
    output = weights @ V_all        # this step's attention output

print("naive projections across 4 steps:", naive_projection_count)   # 10
print("cached projections across 4 steps:", cached_projection_count) # 4

Running this prints 10 for the naive loop and 4 for the cached one — for only 4 tokens. Extend either loop to a 1,000-token conversation and the gap between 1+2+…+1000 and just 1000 is the entire reason production inference engines cache K/V instead of recomputing it.

Speculative Decoding

Even with KV caching, one token per forward pass is still one token per forward pass. Speculative decoding gets around this by having a small, fast “draft” model propose several tokens ahead in one go. The full-size model then checks all of those candidates in a single parallel pass and keeps whichever prefix of them it actually agrees with, discarding the rest. When the draft model guesses well, several tokens come out for the cost of roughly one large-model forward pass — with no change to the output distribution, since every accepted token still has to pass the large model’s own verification.

Source: Leviathan, Kalman & Matias, “Fast Inference from Transformers via Speculative Decoding,” ICML 2023

Perplexity: Quantifying Next-Token Prediction

Perplexity turns “how good was the model’s next-token prediction” into one number: the exponential of the average negative log-likelihood the model assigned to the tokens that actually appeared.

A perplexity of 1 means the model assigned 100% probability to every correct token — a perfect run. A perplexity of 10 means the model was, on average, about as uncertain as if it had to guess uniformly among 10 equally likely options at each step. Lower is always better.

Source: Hugging Face, “Perplexity of Fixed-Length Models”

Try it on the same “The cat sat on the mat.” example from the earlier levels — drag each token’s assigned probability and watch perplexity respond:

Taming Repetition

Greedy or low-temperature decoding tends to loop: once a phrase becomes likely, repeating it often stays likely too. Production systems fight this by penalizing tokens directly in the logits before softmax runs — a repetition penalty downweights any token already seen in the output, while separate frequency and presence penalties scale that downweighting by how often (or simply whether) a token has already appeared. These are ordinary logit adjustments, applied before the exact same softmax-and-sample step covered earlier — no new mechanism, just another knob on the same pipeline.

What the Paper Says

“

We propose a new simple network architecture, the Transformer, based solely on attention mechanisms, dispensing with recurrence and convolutions entirely.

Fun Fact

Speculative decoding is “lossless” by design: because every draft token still has to pass the large model’s own verification step, the final output distribution is mathematically identical to running the large model alone, token by token. The speedup is free — it changes nothing about what the model would have said.

Test Yourself