Transformer Inference
FLOPs
Prefill / Decode
Systems
Transformer Inference Math
A systems-level explanation of where the computation comes from, why attention contains an \(N^2\) term, why most dense-layer work scales as \(d^2\), and why prefill and decode behave very differently on accelerators.
1. Notation
\(N\)
Sequence lengthNumber of tokens processed together.\(d\)
Hidden dimensionWidth of each token representation.\(h\)
Number of headsUsually \(d_h=d/h\).\(d_{ff}\)
MLP widthOften approximately \(4d\), architecture dependent.\(L\)
LayersNumber of transformer blocks.\(B\)
Batch sizeIndependent sequences processed together.
FLOP convention used here: one multiply + one add is counted as 2 FLOPs. Therefore multiplying an \(m\times k\) matrix by a \(k\times n\) matrix costs approximately \(2mkn\) FLOPs.
2. What Happens in One Transformer Layer?
Input X : [N × d]
│
├── Linear projections ──> Q, K, V
│ [N×d] each
│
├── Q Kᵀ
│ ↓
│ attention scores [N×N]
│
├── softmax
│
├── scores × V
│ ↓
│ context [N×d]
│
├── output projection
│
├── residual + norm
│
├── MLP / FFN
│ d → d_ff → d
│
└── residual + norm
The expensive operations are overwhelmingly matrix multiplications. LayerNorm, activation functions, rotary embeddings, and softmax matter for latency and bandwidth, but contribute much fewer arithmetic FLOPs than the dense GEMMs.
Dense projections
For an input \(X\in\mathbb{R}^{N\times d}\) and weight \(W\in\mathbb{R}^{d\times d}\):
\[
XW:\quad (N\times d)(d\times d)\rightarrow(N\times d)
\]
\[
\text{FLOPs} \approx 2Nd^2
\]
Q, K, and V require three such projections:
\[
\text{QKV FLOPs}\approx 3(2Nd^2)=6Nd^2
\]
The attention output projection adds another:
\[
\text{Output projection}\approx 2Nd^2
\]
So attention's dense projection work is approximately:
\[
8Nd^2
\]
3. Why is \(QK^T\) Approximately \(2N^2d\) FLOPs?
This is one of the most important pieces of transformer math.
Assume:
\[
Q\in\mathbb{R}^{N\times d},\qquad K\in\mathbb{R}^{N\times d}
\]
Then:
\[
QK^T:\quad (N\times d)(d\times N)\rightarrow(N\times N)
\]
The output has \(N^2\) elements. Each output element is the dot product of two vectors of length \(d\):
\[
S_{ij}=\sum_{k=1}^{d}Q_{ik}K_{jk}
\]
A length-\(d\) dot product performs roughly \(d\) multiplies and \(d\) adds, or approximately \(2d\) FLOPs.
\[
N^2\text{ output elements}\times 2d
=\boxed{2N^2d}
\]
Intuition: every query token compares itself with every key token. There are \(N\times N=N^2\) token pairs, and each comparison is a \(d\)-dimensional dot product.
Multi-head attention does not remove this term
With \(h\) heads and head dimension \(d_h=d/h\):
\[
h\cdot 2N^2d_h
=h\cdot 2N^2(d/h)
=\boxed{2N^2d}
\]
Splitting into heads changes how the work is partitioned, not the leading total arithmetic count.
The second attention GEMM
After softmax, the attention matrix \(A\in\mathbb{R}^{N\times N}\) multiplies \(V\in\mathbb{R}^{N\times d}\):
\[
AV:\quad (N\times N)(N\times d)\rightarrow(N\times d)
\]
\[
\text{FLOPs}\approx 2N^2d
\]
Therefore the two major attention matrix multiplications together cost:
\[
\boxed{4N^2d}
\]
4. Why Does Compute Often Grow Quadratically With Hidden Size?
The key is that most learned transformer matrices are roughly \(d\times d\).
Take one linear layer:
\[
X:[N\times d],\qquad W:[d\times d]
\]
Its cost is:
\[
2Nd^2
\]
Now double hidden size:
\[
d\rightarrow2d
\]
The cost becomes:
\[
2N(2d)^2=8Nd^2
\]
Original cost was \(2Nd^2\), so:
\[
\frac{8Nd^2}{2Nd^2}=4
\]
Doubling model width causes about 4× dense GEMM FLOPs at the same token count.
This affects Q/K/V projections, the output projection, and the MLP. The MLP is especially important because its intermediate dimension is also proportional to \(d\).
MLP cost
For a conventional FFN with \(d_{ff}\approx4d\):
\[
[N\times d][d\times4d] \Rightarrow 2N d(4d)=8Nd^2
\]
\[
[N\times4d][4d\times d] \Rightarrow 8Nd^2
\]
Total:
\[
\boxed{16Nd^2}
\]
For gated MLPs such as SwiGLU, there is typically an additional input-side projection, so exact constants differ, but the dominant scaling still behaves as \(Nd^2\).
Approximate transformer-layer FLOPs
Using a simple \(4d\) FFN:
| Operation | Approx. FLOPs |
|---|
| Q, K, V projections | \(6Nd^2\) |
| \(QK^T\) | \(2N^2d\) |
| Attention × V | \(2N^2d\) |
| Output projection | \(2Nd^2\) |
| MLP | \(16Nd^2\) |
| Total | \(\boxed{24Nd^2+4N^2d}\) |
The exact coefficient changes across architectures, especially with SwiGLU, grouped-query attention, mixture-of-experts, etc. The useful scaling rule is:
dense layers ∝ \(Nd^2\), while full self-attention contains a term ∝ \(N^2d\).
5. What Happens When Sequence Length Increases?
Two different kinds of work scale differently with \(N\).
| Part | Scaling | Why |
|---|
| Q/K/V and MLP GEMMs | \(O(Nd^2)\) | Each additional token goes through the same dense matrices. |
| Attention score computation | \(O(N^2d)\) | Each token compares against every token. |
| Attention × V | \(O(N^2d)\) | The \(N\times N\) attention matrix is applied to values. |
If \(N\) doubles while \(d\) stays fixed:
\[
Nd^2\rightarrow2Nd^2\qquad \text{(2×)}
\]
\[
N^2d\rightarrow(2N)^2d=4N^2d\qquad \text{(4×)}
\]
At long enough context lengths, the attention term grows increasingly important.
Where do the two terms become comparable?
Compare:
\[
24Nd^2 \quad\text{vs.}\quad 4N^2d
\]
Setting them equal:
\[
24Nd^2=4N^2d
\]
\[
6d=N
\]
So for this simplified layer model, the full-attention arithmetic term becomes comparable to dense-layer arithmetic around \(N\approx6d\). In practice, implementation details, FlashAttention, memory traffic, architecture constants, and hardware utilization shift the real crossover.
6. Prefill vs Decode: Same Model, Very Different Hardware Behavior
Prefill
Process the whole prompt together.
Input shape: roughly \([B,N,d]\)
Large GEMMs, lots of parallelism.
Usually much easier to drive tensor cores efficiently.
Decode
Generate one new token per sequence per step.
Input shape: roughly \([B,1,d]\)
Small token dimension, repeated many times.
Often constrained by moving weights and KV-cache data rather than peak arithmetic.
Prefill attention
For a prompt of \(N\) tokens, attention contains the familiar:
\[
QK^T\approx2N^2d
\]
\[
AV\approx2N^2d
\]
Decode attention
At one decode step, there is only one new query token. If the KV cache already contains \(T\) previous tokens:
\[
Q:[1\times d],\qquad K^T:[d\times T]
\]
\[
QK^T\approx2Td
\]
Then attention probabilities multiply cached V:
\[
[1\times T][T\times d]\Rightarrow2Td
\]
So attention computation for one generated token is approximately:
\[
\boxed{4Td}
\]
The KV cache converts autoregressive decoding from recomputing all prior K/V representations into reusing them. Without the cache, generation would repeatedly redo enormous amounts of work.
7. KV Cache Math
For standard multi-head attention, each cached token stores a key and a value for every layer.
Approximate bytes per token:
\[
\text{KV bytes/token}
\approx
2\times L\times d\times \text{bytes per element}
\]
The factor 2 is for K and V.
For FP16/BF16, bytes per element ≈ 2:
\[
\text{KV bytes/token}\approx4Ld
\]
Example
Suppose \(L=32\), \(d=4096\), BF16:
\[
2\times32\times4096\times2
=524{,}288\text{ bytes}
\]
That is about 512 KiB per token per sequence for standard MHA.
For 8K tokens:
\[
512\text{ KiB}\times8192 \approx 4\text{ GiB}
\]
This is one major reason modern models use MQA/GQA: they reduce the number of distinct K/V heads and therefore KV-cache storage and bandwidth.
8. FLOPs vs FLOP/s: Why More FLOPs Does Not Necessarily Mean Better GPU Utilization
FLOPs means total arithmetic work. FLOP/s means how quickly the hardware performs that work.
A workload can contain more total FLOPs while achieving either higher or lower realized FLOP/s, depending on arithmetic intensity, matrix shapes, memory traffic, kernel efficiency, and parallelism.
Arithmetic intensity
\[
\text{Arithmetic Intensity}
=
\frac{\text{FLOPs}}{\text{Bytes moved}}
\]
Roofline intuition:
\[
\text{Attainable FLOP/s}
\approx
\min\left(
\text{Peak Compute},
\text{Memory Bandwidth}\times\text{Arithmetic Intensity}
\right)
\]
Increasing hidden dimension
Increasing \(d\) makes dense GEMMs larger. Their FLOPs grow roughly as \(d^2\), while data reuse often improves. This tends to increase arithmetic intensity and can improve accelerator utilization—until peak compute, capacity, or another bottleneck is reached.
So when asked:
“What happens to realized FLOP/s when hidden dimension increases?”
Usually it can increase toward peak because GEMMs get larger and hardware is used more efficiently. But total FLOPs also grow rapidly, approximately quadratically with \(d\) for dense layers.
Increasing sequence length
During prefill, increasing \(N\) makes matrices larger and creates more parallel work, which can improve utilization at first. But attention work grows quadratically and memory requirements increase, so very long sequences can eventually become constrained by memory capacity/bandwidth and attention efficiency.
“What happens to realized FLOP/s when sequence length increases?”
For small sequences, realized FLOP/s often rises because kernels become larger and more efficient. At large \(N\), it may plateau near peak or fall if memory, attention, synchronization, or capacity becomes the bottleneck.
Why decode has low FLOP/s even though inference can be slow
Decode repeatedly executes narrow matrix operations for one/few tokens while reading a huge amount of model weights. That makes the workload comparatively low in arithmetic intensity.
A simplified linear layer for one token is:
\[
[1\times d][d\times d]
\]
It performs about:
\[
2d^2\text{ FLOPs}
\]
but may need to stream much of a \(d\times d\) weight matrix from memory. Reusing the weights across a larger batch improves this significantly.
9. Disaggregated Serving: How Would You Split 8 + 8 Accelerators?
Suppose you have separate pools of 8 accelerators for prefill and 8 for decode.
Prefill pool
- Optimize for high compute throughput.
- Batch prompts aggressively.
- Exploit large GEMMs.
- Tensor parallelism can work well for large models.
- Long prompts create high, bursty compute demand.
Decode pool
- Optimize for memory bandwidth and latency.
- Continuous batching is crucial.
- Need large KV-cache capacity.
- Weight reuse across concurrent sequences matters.
- Inter-token latency is a key metric.
The real challenge: KV transfer
After prefill, the decode workers need the prompt's KV cache.
PREFILL ACCELERATORS
Prompt → compute prompt tokens
→ create K/V cache
→ first-token state
│
│ KV transfer
▼
DECODE ACCELERATORS
token t+1
token t+2
token t+3
...
This means the interconnect between prefill and decode pools matters. Moving large KV caches can erase some of the benefit of disaggregation if transfer bandwidth or latency is inadequate.
Should the split always be 8 / 8?
No. The optimal ratio depends on workload.
| Traffic pattern | Likely implication |
|---|
| Long prompts, short outputs | More prefill capacity. |
| Short prompts, very long outputs | More decode capacity. |
| High concurrency | Decode pool and KV capacity become especially important. |
| Very long context | Prefill compute + KV transfer/storage both grow substantially. |
A useful capacity model is:
\[
\text{required prefill capacity}
\propto
\lambda \times E[\text{prompt work}]
\]
\[
\text{required decode capacity}
\propto
\lambda \times E[\text{generated tokens}]
\times E[\text{decode work/token}]
\]
where \(\lambda\) is request arrival rate.
For an interview answer, do not say “always 4 GPUs for prefill and 4 for decode” or any fixed ratio without workload data. Start from prompt-length distribution, output-length distribution, request rate, TTFT target, TPOT target, KV-cache footprint, and interconnect bandwidth.
10. Interview Cheat Sheet
| Question | Compact answer |
|---|
| Why is \(QK^T\) \(2N^2d\)? |
The output is \(N\times N\). Each of its \(N^2\) entries is a length-\(d\) dot product, costing about \(2d\) FLOPs. |
| Why does hidden dimension cause quadratic compute growth? |
Most transformer weights are \(d\times d\). Multiplying \(N\times d\) by \(d\times d\) costs \(2Nd^2\). Double \(d\) → about 4× the FLOPs. |
| What happens when sequence length doubles? |
Dense projections/MLP roughly double; full attention's \(N^2d\) terms roughly quadruple. |
| Why is prefill compute-friendly? |
Many prompt tokens are processed together, producing large GEMMs with high parallelism and good weight reuse. |
| Why is decode bandwidth-sensitive? |
Each step processes one/few new tokens, so there is relatively little arithmetic per byte of weights/KV data moved. |
| What does KV cache save? |
It avoids recomputing K/V for all prior tokens during each autoregressive step. |
| Decode attention cost at context length \(T\)? |
Approximately \(4Td\) FLOPs per layer per generated token for score + value aggregation, ignoring projection constants. |
| What matters in disaggregated serving? |
Different prefill/decode bottlenecks, workload-aware resource split, KV-cache transfer, TTFT, TPOT, batching, memory capacity, and interconnect bandwidth. |
One useful mental model
Dense transformer work:
tokens × model-width²
~ N d²
Full attention work:
token-pairs × model-width
~ N² d
Prefill:
large matrices → compute-friendly
Decode:
tiny token dimension + huge weights/KV
→ bandwidth / latency sensitive