You do not need to derive backpropagation to serve a transformer. You need to know which tensors exist, which ones grow, and which ones get read on every single token.

One block at a time

A decoder-only transformer is a stack of identical blocks. Each block has two main parts:

1. Self-attention. Project the hidden state into queries, keys and values. Compare each query with the keys of current and previous positions. Use the resulting weights to mix the values. Project the result back to the model dimension.
2. MLP. A position-wise feed-forward network, usually two linear layers with a nonlinearity between them. For a model with hidden size $d$ and intermediate size $4d$, the MLP holds the majority of the parameters.

Each part has a residual connection and a normalization layer. At the end of the stack, a language-model head maps the final hidden state to logits over the vocabulary.

For serving, the important tensors are:

The weights are fixed. The KV cache is the only tensor that grows while the request runs.

The attention computation

For a single attention head, with query matrix $Q$, key matrix $K$ and value matrix $V$:

The scaling by $\sqrt{d_k}$ keeps the dot products from growing with head dimension. A causal mask prevents each position from attending to future positions.

In prefill, $Q$, $K$ and $V$ all contain $T$ prompt tokens. The attention matrix is $T \times T$. Computing it costs roughly $O(T^2 d)$ and, if materialized, $O(T^2)$ memory. For a 32K-token prompt, that is a billion entries per head if you write the full matrix to HBM. FlashAttention avoids materializing it by tiling the computation and keeping blocks in SRAM.

In decode, $Q$ has one token, while $K$ and $V$ contain the entire prefix. The attention matrix is $1 \times T$. The math per step is linear in context length, but the KV cache must be read from HBM. That read is a bandwidth cost, and it recurs for every token.

Why the cache exists

Without a cache, generating token $t$ would require recomputing keys and values for all $t-1$ previous tokens. The total work for $N$ output tokens would be $O(N^2)$ in the number of forward passes, not just attention FLOPs. The KV cache stores the keys and values after they are computed, so each decode step only needs to compute the new token’s $K$ and $V$ and append them.

The price is memory. For a model with $L$ layers, $H_{kv}$ KV heads, head dimension $d_h$, and $b$ bytes per element, one token costs

For Llama 3.1 8B in BF16 — 32 layers, 8 KV heads, head dimension 128 — that is

At 8K context, one request needs about 1 GiB of KV cache. At 32K, about 4 GiB. Multiply by concurrency and the cache becomes a first-class capacity constraint. This is the subject of the next lesson.

Grouped-query attention changes the budget

Standard multi-head attention (MHA) has one key/value head per query head. Multi-query attention (MQA) shares a single key/value head across all query heads. Grouped-query attention (GQA) sits between them: query heads are divided into groups, and each group shares one key/value head.

GQA reduces the KV cache by the ratio of query heads to KV heads. Llama 3.1 8B has 32 query heads and 8 KV heads, so its cache is one quarter the size of an MHA equivalent. That is why modern long-context models can serve reasonable concurrency at all.

The tradeoff is quality and implementation complexity. Fewer KV heads mean less representational capacity in the cache. MQA is the most aggressive; GQA is the common compromise. Some architectures go further: multi-head latent attention (MLA) compresses keys and values into a latent vector, shrinking the cache substantially at the cost of more involved kernels. DeepSeek-V2 and V3 use this approach.

For a serving engineer, the practical rule is simple: the KV cache formula is part of the model card you should care about. Parameter count alone does not tell you how many concurrent users fit.

Counting FLOPs, roughly

For a dense decoder with $N$ parameters:

• Prefill over $T$ tokens costs roughly $2NT$ FLOPs for the linear layers, plus attention terms that grow with $T^2$.
• Decode for one token costs roughly $2N$ FLOPs for the linear layers, plus attention work proportional to the cached context.

The factor of 2 is multiply-add. These are order-of-magnitude estimates; they ignore embeddings, normalization, softmax and masking. For MoE models, replace $N$ with the number of active parameters per token, not the total parameter count.

At batch size $B$, decode linear-layer FLOPs become roughly $2NB$, while the weight bytes read remain roughly $2N$ in BF16. That is the arithmetic-intensity argument from the previous lesson, expressed in transformer terms.

Two shapes, one set of weights

It is tempting to treat prefill and decode as separate systems. They are not. They share weights, kernels and often the same GPU. But their shapes differ enough that optimizations rarely transfer:

This is why a serving engine schedules them differently, and why disaggregated serving — running prefill and decode on separate pools — is an active design. The two phases want different hardware ratios and different parallelism.

What this means for a real deployment

If your workload is long prompts and short outputs (RAG, classification, code review), prefill dominates. Invest in compute, prompt caching, and avoiding duplicate work. If your workload is short prompts and long outputs (chat, creative writing, chain-of-thought), decode dominates. Invest in bandwidth, batching, KV-cache efficiency, and possibly speculative decoding.

If you have both, measure the mix. A single “average latency” number hides the fact that your RAG requests and your chat requests are stressing different parts of the machine.

Next: the KV cache in detail — how it grows, how it fragments, and why managing it is most of what a modern serving engine does.