LLM Inference Walkthrough — Tensor Shapes and Core Formulas Across the Whole Pipeline

Lay out a Llama 3-style dense decoder-only model end to end from embedding to sampling, embedding the math positions of common variants (MHA / MQA / GQA / MLA, RMSNorm / LayerNorm, SwiGLU / GeGLU, Flash Attention / Paged Attention) along the way — so that after reading you can draw this diagram from memory. One fact frames the whole article: this architecture has barely changed in twenty years; every “inference optimization” is a local surgery somewhere on the same skeleton.

Notation Conventions — Llama 3 8B as the Reference

The same notation is used throughout, with Llama 3 8B as the concrete example:

All shape annotations follow PyTorch convention, written as [B, ..., H]; weight matrices follow “in × out”, written as $W \in \mathbb{R}^{[\text{in}, \text{out}]}$.

Core Formulas Quick Reference — Embedding · Norm · Attn · FFN · LM Head

Embedding

• $\mathbf{x}_i$ — embedding vector at position $i$
• $E \in \mathbb{R}^{V \times H}$ — embedding lookup table ($V$ tokens, each $H$-dim)
• $\text{token\_id}_i \in \{0, 1, \ldots, V-1\}$ — integer ID of the $i$-th input token

Many implementations share $E$ with the LM Head’s $W_{\text{lm}}$ (tied embedding), saving memory and providing mild regularization; Llama models do not share by default.

Normalization: LayerNorm vs RMSNorm

Standard LayerNorm:

• $\mathbf{x} \in \mathbb{R}^{H}$ — single-token hidden state
• $\mu, \sigma^{2} \in \mathbb{R}$ — mean and variance computed across the $H$ dimensions
• $\epsilon$ — small constant to avoid divide-by-zero (typically $10^{-5}$)
• $\boldsymbol{\gamma}, \boldsymbol{\beta} \in \mathbb{R}^{H}$ — learnable scale / shift
• $\odot$ — element-wise multiplication

RMSNorm (the mainstream choice for Llama / Mistral / Qwen):

• $\mathbf{x} \in \mathbb{R}^{H}$ — single-token hidden state
• $\boldsymbol{\gamma} \in \mathbb{R}^{H}$ — learnable per-channel scale (no $\boldsymbol{\beta}$)
• $\epsilon$ — small constant to avoid divide-by-zero
• The denominator is the root mean square of $\mathbf{x}$ — hence “RMSNorm”

RMSNorm drops the mean and $\boldsymbol{\beta}$, roughly halving both compute and parameters; empirically the quality cost is nearly zero. All mainstream inference engines organize this as pre-norm: norm lives inside the residual branch, and the residual main path bypasses it.

Q/K/V Projection + Positional Encoding

• $X \in \mathbb{R}^{B \times S \times H}$ — output of the previous RMSNorm
• $W_Q \in \mathbb{R}^{H \times n_q d}$ — Q projection weight
• $W_K, W_V \in \mathbb{R}^{H \times n_{kv} d}$ — K, V projection weights (narrower under GQA since $n_{kv} < n_q$)
• $Q \in \mathbb{R}^{B \times S \times n_q d}$, $K, V \in \mathbb{R}^{B \times S \times n_{kv} d}$ — projection outputs, later reshaped to expose the head dimension

RoPE (Rotary Positional Embedding) core idea: rather than adding “absolute position $m$” to the embedding (like the original Transformer’s sinusoidal), let position act as a rotation on Q and K, so that when the two vectors are dotted only the relative position $n - m$ survives.

Split each head’s $d$ dimensions into $d/2$ two-dim subspaces; the $k$-th subspace ($k = 0, 1, \ldots, d/2 - 1$) corresponds to coordinates $(q_{2k}, q_{2k+1})$. At position $m$ that pair is multiplied by a 2D rotation matrix with angle $m\theta_k$:

• $m \in \{0, 1, \ldots, S-1\}$ — absolute position of the current token
• $k \in \{0, 1, \ldots, d/2 - 1\}$ — 2D subspace index (every two dims of the per-head dim $d$ form one pair)
• $(q_{2k}, q_{2k+1})$ — the $k$-th 2D pair of the projected Q vector
• $(q'^{(m)}_{2k}, q'^{(m)}_{2k+1})$ — the rotated pair at position $m$
• $\theta_k$ — base angular velocity of the $k$-th subspace
• $\text{base}$ — frequency-decay base (Llama default $10000$; YaRN etc. dynamically enlarge it)

Expanded to scalars:

Same symbols as above — this just unrolls the $2\times 2$ matrix into two scalar identities, easier to map to code.

This is the standard counterclockwise rotation matrix $R(\phi) = \begin{pmatrix} \cos\phi & -\sin\phi \\ \sin\phi & \phantom{-}\cos\phi \end{pmatrix}$; acting on a 2D vector is equivalent to multiplying by $e^{i\phi}$ in the complex plane — which is why the official Llama repo reshapes $(q_{2k}, q_{2k+1})$ to complex64 and multiplies by $\cos(m\theta_k) + i\sin(m\theta_k)$ directly (HuggingFace equivalently pairs $(q_k, q_{k+d/2})$ in a “half-rotation” form; the two differ only by a coordinate permutation). K rotates by the same angle to get $\mathbf{k}'^{(n)}$.

Why does rotation encode relative position? Let $\mathbf{R}_m$ denote the block-diagonal full-$d$ rotation (the $k$-th $2\times 2$ block uses angle $m\theta_k$). It’s orthogonal and satisfies $\mathbf{R}_m^{\top}\mathbf{R}_n = \mathbf{R}_{n-m}$, so:

• $\mathbf{q}, \mathbf{k} \in \mathbb{R}^{d}$ — per-head Q, K vectors (pre-rotation)
• $\mathbf{R}_m, \mathbf{R}_n \in \mathbb{R}^{d \times d}$ — block-diagonal rotation matrices for positions $m$, $n$
• $\langle \cdot, \cdot \rangle$ — standard inner product
• Last step uses $\mathbf{R}_m^{\top} \mathbf{R}_n = \mathbf{R}_{n-m}$: rotations are orthogonal and angles are additive

When attention computes $QK^{\top}$, every $(q_i, k_j)$ score depends only on the difference $j - i$ — absolute position is automatically canceled inside the dot product, leaving relative position. This is RoPE’s key property and the fundamental reason it is more stable than additive positional encodings.

Frequency spectrum design: $\theta_k = \text{base}^{-2k/d}$ ($\text{base}$ usually $10000$; $d$ is the per-head dim, not model hidden dim) assigns the $d/2$ subspaces angular velocities from fast to slow:

• $k = 0$: $\theta_0 = 1$, period $2\pi \approx 6.28$ tokens — carries the “near-neighbor” signal.
• $k = d/2 - 1$: $\theta \approx \text{base}^{-(d-2)/d} \approx 10^{-4}$, period $\approx 2\pi \cdot 10000 \approx 62\text{K}$ tokens — carries the “long-range” signal.

The geometric distribution lets one head carry positional signals at many scales simultaneously — isomorphic to the original Transformer’s sinusoidal, just moved from addition to multiplication.

Long-context scaling: training only saw $m \le L_{\text{train}}$, and low-frequency subspaces can’t even complete one full period inside $L_{\text{train}}$; once inference reaches $m > L_{\text{train}}$, the low-frequency angles fall outside the training distribution and attention immediately degrades. Three common fixes all modify $\theta_k$:

• Position Interpolation (Chen et al. 2023): $m \to m/s$, equivalent to $\theta_k \to \theta_k/s$ — compresses all frequencies uniformly. Simple but wastes high-frequency precision.
• NTK-aware scaling: scales low frequencies while preserving high ones, equivalent to $\text{base} \to \text{base} \cdot s^{d/(d-2)}$.
• YaRN (Peng et al. 2023): band-wise treatment — high frequencies (period $\ll L_{\text{train}}$, full periods seen during training) are left alone, low frequencies (period $\gg L_{\text{train}}$, no full periods seen) are PI-scaled, and the mid band interpolates smoothly; additionally a $1/\sqrt{t}$ temperature correction cancels the attention-entropy drift introduced by the scaling. Llama 3.1 / 3.2 went from 8K → 128K with YaRN.

RoPE acts only on Q and K, not on V — V is the weighted value itself and needs no positional signal.

Scaled Dot-Product Attention

• $Q, K, V$ — projected and head-split tensors, with the last two dims being $[S, d]$ (head dim broadcast implicit)
• $QK^{\top} \in \mathbb{R}^{S \times S}$ — similarity matrix between every $(q_i, k_j)$ pair
• $\sqrt{d}$ — scaling factor that keeps softmax out of the near-zero-gradient region
• $M \in \mathbb{R}^{S \times S}$ — causal mask with upper triangle set to $-\infty$ (position $i$ can only see $\le i$)
• softmax normalizes along the K dim, producing attention weights that then re-weight V back to $[S, d]$

Multi-Head Variants: MHA / MQA / GQA / MLA

All four variants differ only on the K, V side: Q is always $n_q$ independent heads; what changes is how many independent K, V sets exist, and whether K, V are low-rank compressed. Below we write head-level formulas for original MHA and each of the three variants under one notation — dropping the batch and layer dims, looking at head $h$ at position $i$.

MHA — Multi-Head Attention (original, Vaswani et al. 2017)

Each Q head gets its own K, V — $n_q$ independent sets:

• $W_Q^{(h)}, W_K^{(h)}, W_V^{(h)} \in \mathbb{R}^{H \times d}$ — per-head Q, K, V projections (concatenated into a big matrix in practice, mathematically equivalent)
• Each token writes $n_q$ pairs $(\mathbf{k}^{(h)}, \mathbf{v}^{(h)})$ into the cache — $2 n_q d$ scalars in total
• Llama 2 7B: $n_q = 32, d = 128$, per-token cache = $2 \cdot 32 \cdot 128 \cdot 2\text{B} = 16\text{ KB}$ / layer

MQA — Multi-Query Attention (Shazeer 2019)

All $n_q$ Q heads share a single K, V — only 1 set left:

• $W_K, W_V \in \mathbb{R}^{H \times d}$ — one shared K, V projection across all heads
• Cache shrinks to $2d$ — $n_q\times$ smaller than MHA; but capacity-constrained, large models drop quality if MQA is applied directly — rarely used standalone in practice, mostly superseded by GQA

GQA — Grouped-Query Attention (Ainslie et al. 2023)

Split the $n_q$ Q heads into $n_{kv}$ groups, share K, V within each group — a continuous interpolation between MHA and MQA:

• $W_K^{(g)}, W_V^{(g)} \in \mathbb{R}^{H \times d}$, $g = 1, \ldots, n_{kv}$ — one K, V set per group
• $g(h)$ — group index of Q head $h$
• $n_{kv} = n_q$ recovers MHA, $n_{kv} = 1$ recovers MQA
• Kernels materialize only $n_{kv}$ K, V sets; during score computation, K and V are broadcast along the group dim up to $n_q$, not actually replicated
• Llama 3 70B: $n_q = 64, n_{kv} = 8, d = 128$, per-token cache = $2 \cdot 8 \cdot 128 \cdot 2\text{B} = 4\text{ KB}$ / layer — 8× smaller than a same-shape MHA

MLA — Multi-Head Latent Attention (DeepSeek V2 2024)

GQA only shrinks cache linearly in head count; MLA jointly compresses K and V into a low-rank latent $\mathbf{c}^{KV}$, and breaks RoPE off into a shared shallow branch. Four steps:

(1) Content branch — K, V share a down-projection; only the latent $\mathbf{c}^{KV}_i$ goes into the cache:

(2) Q-side low-rank — saves training memory; Q is never cached at inference:

(3) Decoupled RoPE branch — K side computes one shared $d_r$-dim vector reused by all heads:

(4) Concat + attention — content and RoPE parts are concatenated along the head dim before scoring:

• $W^{DKV} \in \mathbb{R}^{H \times d_c}$ — shared KV down-projection (DeepSeek V3 uses $d_c = 512$)
• $W_K^{U,(h)}, W_V^{U,(h)} \in \mathbb{R}^{d_c \times d}$ — per-head up-projections
• $W^{KR} \in \mathbb{R}^{H \times d_r}$ — shared RoPE K branch (DeepSeek V3 uses $d_r = 64$)
• $\mathbf{c}^{KV}_i \in \mathbb{R}^{d_c}, \mathbf{k}^{R}_i \in \mathbb{R}^{d_r}$ — these two are all that MLA actually caches — $d_c + d_r$ scalars
• DeepSeek V3: $n_q = 128, d = 128, d_c = 512, d_r = 64$, per-token cache = $(512 + 64) \cdot 2\text{B} = 1152$ bytes / layer — another order of magnitude below a same-scale GQA

Why does RoPE need its own branch? At inference there’s a “weight-folding” trick — by associativity, fold $W_K^{U,(h)}$ into $W_Q^{U,(h)}$:

K’s content part is never actually reconstructed — attention runs directly on the cached $\mathbf{c}^{KV}$. But RoPE’s rotation angle depends on absolute position $m$, so it cannot be folded into fixed weights — applying RoPE on reconstructed K kills the fold. DeepSeek’s fix: pull RoPE out into an independent $d_r$-dim shallow branch, so “foldable content” and “must-rotate-live position” don’t interfere. Cache compression and relative-position signal both survive.

Overview: Params / FLOPs / KV cache

Below we decompose each variant’s cost step by step. Conventions: one attention sublayer per layer, prefill length $S$, ignoring norm / bias / softmax and other non-matmul terms; a linear projection $X W$ (with $X \in \mathbb{R}^{S \times m}, W \in \mathbb{R}^{m \times n}$) counts as $2Smn$ FLOPs (2 FLOPs per MAC); $QK^{\top}$ and $AV$ do not deduct the causal-mask triangular half. MLA additionally uses $d_q'$ = Q-side latent dim (DeepSeek V3 uses $1536$).

Parameters (per layer)

Prefill FLOPs (per layer, sequence length $S$)

“RoPE branch” in the MHA/MQA/GQA columns means the elementwise rotation applied to Q, K — $\mathcal{O}(S n_q d)$, negligible next to matmuls. In MLA it refers specifically to the extra $W_Q^R$ / $W^{KR}$ projections, which are genuine matmuls and must be counted.

KV cache (per token, per layer, fp16)

Three observations from reading these tables horizontally:

1. All variants only touch the K/V side. Q projection, $W_O$, and $AV$ are identical across the four — the surgery only touches K, V, so total attention params and FLOPs never differ by an order of magnitude at the same model size.
2. MHA → MQA/GQA saves params, FLOPs, and cache together; MLA trades params/FLOPs for cache. GQA shrinks $n_{kv}$, scaling K/V projections, FLOPs, and cache down linearly. MLA does the opposite — adds DKV + two UK/UV stages + a dedicated RoPE branch, leaving params and prefill FLOPs at the same order as a same-size GQA, in exchange for shrinking per-token cache from KB scale down to ~1 KB.
3. The $S^2$ term doesn’t differ across MHA/MQA/GQA. $QK^{\top}$ and $AV$ are both $2 n_q S^2 d$ (K, V sharing is broadcast only, never reducing the quadratic term). So once $S \gg H/d$, the three variants’ prefill FLOPs converge — what truly separates them is how many cache bytes each decoded token must read, which is memory I/O, not compute.

Output Projection + Residual

• $\mathbf{x} \in \mathbb{R}^{H}$ — attention sublayer input (the value before pre-norm)
• $Q', K', V'$ — RoPE-rotated Q and K; $V' = V$ (V is not rotated, the prime is just notational uniformity)
• $W_O \in \mathbb{R}^{n_q d \times H}$ — output projection that maps concatenated multi-head output back to $H$
• $\mathbf{h} \in \mathbb{R}^{H}$ — attention sublayer output + residual

FFN Variants

Classic Bilinear FFN (GPT-2)

• $\mathbf{x} \in \mathbb{R}^{H}$ — FFN input
• $W_1 \in \mathbb{R}^{H \times I}$ — up-projection
• $W_2 \in \mathbb{R}^{I \times H}$ — down-projection
• $\phi$ — element-wise scalar nonlinearity (see below)

Activation Functions

The scalar nonlinearity applied after the up-projection. Input and output are both scalars $x$; it’s applied independently per token, per hidden dim. Four functions cover all the mainstream choices used in the Transformer era.

ReLU (Nair & Hinton 2010)

• Identity on $x > 0$, zero on $x \le 0$
• Cheapest to compute; but zero gradient on the negative side — the “dead neuron” problem
• Used in the original Transformer and early BERT implementations

Sigmoid

• Compresses $\mathbb{R}$ into $(0, 1)$, a natural “gate” signal — the original GLU’s $\phi$ is exactly this
• Largely abandoned as a standalone FFN activation — saturates at both ends, killing gradients

GeLU (Gaussian Error Linear Unit, Hendrycks & Gimpel 2016)

In practice, the OpenAI tanh approximation is used (numerical error $< 10^{-3}$, avoids the erf call):

• $\Phi$ is the standard normal CDF — intuitively “let $x$ through weighted by its tail probability”
• Everywhere differentiable, non-monotonic (a small negative dip on the $x < 0$ side), smoother than ReLU
• Used in GPT-2/3, BERT, ViT

SiLU / Swish (Ramachandran et al. 2017)

• Shape very close to GeLU (also smooth, non-monotonic, passes through origin) but with a simpler closed form — no $\text{erf}$, no cubic term
• Self-gated: uses its own sigmoid to control signal throughput
• Used by PaLM and the entire Llama family (as the gating $\phi$ in SwiGLU)

Of these four, ReLU and sigmoid have largely exited mainstream FFNs; the active ones are GeLU (GPT era) and SiLU (Llama / PaLM and later). The bf16 numerical difference between the two is < 1%; paper choice is mostly path-dependent. The real engineering inflection point was switching activation from “applied to the projection” (classic FFN) to “applied to the gate” (the GLU family below).

GLU Family (used by Llama, PaLM, Mistral)

• $\mathbf{x} \in \mathbb{R}^{H}$ — FFN input

• $W_{\text{gate}}, W_{\text{up}} \in \mathbb{R}^{H \times I}$ — two independent up-projections

• $W_{\text{down}} \in \mathbb{R}^{I \times H}$ — down-projection

• $\odot$ — element-wise multiplication

• $\phi$ — gating activation; choice determines the variant name:

  Variant	$\phi$	Representative model
  GLU	$\sigma$	Dauphin et al. 2017 original
  ReGLU	$\text{ReLU}$	—
  GeGLU	$\text{GeLU}$	T5 v1.1
  SwiGLU	$\text{SiLU}$	PaLM, Llama 1/2/3

SwiGLU has one more projection than the classic GeLU-FFN (three matrices vs two); to align parameter budget, implementations typically set $I = \tfrac{2}{3} \cdot 4H$ ($4H$ is the GPT-2 convention), and Llama 3 8B’s $I = 14336 \approx \tfrac{2}{3}\cdot 4\cdot 4096 \cdot 1.3$.

Mixture-of-Experts (MoE)

In a classic dense FFN, every token passes through the same $W_{\text{up}} / W_{\text{down}}$ pair — all parameters used, all FLOPs paid. MoE (Shazeer et al. 2017) replicates the FFN $N$ times (“experts”) and routes each token through only the top $k$, so total parameters scale linearly while per-token activated params and FLOPs stay almost flat — capacity grows without paying compute.

Formally the FFN sublayer becomes:

The router (gating network) decides which experts each token visits:

• $\mathbf{x} \in \mathbb{R}^{H}$ — current token’s hidden state (FFN sublayer input)
• $W_g \in \mathbb{R}^{H \times N}$ — router projection mapping the hidden state to $N$ expert logits
• $\mathbf{s}(\mathbf{x}) \in \mathbb{R}^{N}$ — router probabilities across all experts
• $\mathcal{T}_k(\mathbf{x})$ — index set of top-$k$ selected experts
• $g_i(\mathbf{x})$ — combine weight; Mixtral and DeepSeek re-normalize the top-$k$ scores so they sum to 1
• $\text{FFN}_i$ — $i$-th expert, typically a SwiGLU FFN with its own (unshared) weights

MoE vs Dense FFN

Key trade-off: decoupling capacity from FLOPs. At the same activated parameter count (i.e. same FLOP budget), MoE can pack 8–32× more total parameters — knowledge capacity goes up for free. The costs are VRAM (must fit all experts), routing stability (avoiding hot experts), and inference batching (per-token paths differ).

Load Balance

A naive router collapses onto a few hot experts during training. GShard / Switch use an auxiliary loss as a soft constraint:

• $f_i$ — fraction of tokens in the current batch routed to expert $i$
• $\bar{s}_i$ — mean router probability assigned to expert $i$ within the batch
• $\alpha$ — auxiliary loss weight (Switch uses $\sim 0.01$)
• Intuition: large $f_i$ AND large $\bar{s}_i$ means the expert is chosen often and with high confidence → penalized → gradient pushes that router logit back down

DeepSeek V3 goes further with auxiliary-loss-free load balance: maintain a bias $b_i$ per expert, base TopK on $s_i + b_i$; raise $b_i$ for under-used experts and lower it for over-used ones — affecting selection without polluting gradients, avoiding the main-task accuracy hit that aux loss can cause.

Modern Implementations

Four patterns:

1. Coarse → fine. Mixtral-era designs are “few-and-fat” (8 experts each near a dense FFN’s size); DeepSeek onward shrinks each expert and pushes count past 100, going “many-and-thin” — the same activated FLOPs now span exponentially more combinations.
2. Shared experts become standard. DeepSeek / Qwen-MoE / Llama 4 all reserve 1–2 “everyone-must-visit” shared experts per layer to absorb generic patterns; the sparse experts handle specialization.
3. MoE pairs with KV-cache compression. Pushing total params to hundreds of billions makes long-context decode equally pressured — it’s no coincidence DeepSeek V3 ships MLA + fine-grained MoE together.
4. top-$k$ goes to the extremes. Switch (k=1) → Mixtral (k=2) → DeepSeek V3 (k=8 with small experts) → Llama 4 (k=1 + shared). Small $k$ makes batching and capacity bounds tractable; fine-graining and shared experts make up the lost expressivity.

Residual

• $\mathbf{h}$ — attention sublayer output (already contains the first residual)
• $\mathbf{x}_{\text{out}} \in \mathbb{R}^{H}$ — full Transformer layer output, fed into the next layer
• RMSNorm sits inside the residual branch (pre-norm); the main path $\mathbf{h}$ goes straight through

LM Head + Sampling

• $\mathbf{x}_{\text{final}} \in \mathbb{R}^{H}$ — hidden state after the final RMSNorm, at the last position
• $W_{\text{lm}} \in \mathbb{R}^{H \times V}$ — output embedding matrix (optionally shared with $E$ — tied embedding)
• $\text{logits} \in \mathbb{R}^{V}$ — unnormalized score per vocab token
• $T > 0$ — temperature; larger flattens the distribution ($T \to \infty$ → uniform, $T \to 0$ → argmax)
• $\mathbf{p} \in \mathbb{R}^{V}$ — normalized probability distribution

Several logits transformations are typically applied before sampling:

Repetition penalties (repetition / frequency / presence penalty):

• $v$ — a token in the vocabulary
• $\text{logits}_v$ — the raw logit for that token
• $\alpha \ge 0$ — presence penalty (subtract once if ever seen)
• $\beta \ge 0$ — frequency penalty (subtract linearly by occurrence count)
• $\mathbb{1}[\cdot]$ — indicator function (1 if the condition holds, else 0)
• $\text{count}(v)$ — number of times token $v$ has appeared in the generated sequence

Top-k: keep the largest $k$ logits, set the rest to $-\infty$.

Top-p (nucleus): sort by probability descending, keep the smallest set whose cumulative probability $\le p$.

Min-p: keep tokens with $p_v \ge p_{\min} \cdot p_{\max}$; friendlier to low-entropy distributions.

Typical-p: truncate by deviation from the conditional entropy, keeping the set where $|-\log p_v - H(\mathbf{p})|$ is small.

All truncations act before/after the probability distribution itself, without changing the skeleton of the formula.

Prefill Stage Shape Transitions — S Tokens Walk Through Once

Input: input_ids [B, S] = [2, 10]. The figure below shows the forward pass of one Transformer layer, wrapped in 32 layers; the shape stays at [B, S, H] = [2, 10, 4096] throughout, with residual edges shown as orange dashed lines.

After prefill, the KV Cache state: each layer holds the first 10 positions.

Decode Stage Shape Transitions — Step t Processes Only 1 Token

Prior state: $\text{cache\_len} = S + t - 1$ positions filled.

Input: input_ids [B, 1] = [2, 1] (the 1 token generated at the previous step).

Prefill vs Decode Shape Comparison — GEMM vs GEMV · Compute vs Bandwidth

This table is the starting point for understanding every inference acceleration effort: prefill is like training’s forward, compute-bound; decode is a chain of GEMVs, memory-bound, with most time spent fetching weights into SMs. The optimization directions of the two are worlds apart.

KV Cache Shape and Growth — Few KB Per Token · Hundreds of MB at Long Context

One pair of caches per layer:

• $S_{\max}$ — pre-allocated max sequence length (typically the model’s context cap or the scheduler’s limit)
• Other symbols follow the top-of-article convention ($B, n_{kv}, d$); the 4 dimensions follow PyTorch’s [B, head, seq, head_dim] ordering

Per-token, per-layer cache size (fp16):

• Leftmost $2$ — one for K, one for V
• $2\ \text{bytes}$ — fp16 element size (fp8 / int8 cuts this to 1/2 or 1/4)
• Right side plugs in Llama 3 8B: $n_{kv} = 8$, $d = 128$

Per-token across the full 32-layer model: $4\ \text{KB} \times 32 = 128\ \text{KB} / \text{token}$. A 4096-token request: $128\ \text{KB} \times 4096 \approx 512\ \text{MB}$.

Several engineering optimizations:

• Paged Attention (vLLM): split the cache into fixed-size blocks (typically 16 tokens), with a block table mapping virtual to physical addresses, eliminating fragmentation. The formulas don’t change; only the tensor layout and access pattern change.
• Sliding Window Attention (Mistral): keep only the most recent $W$ tokens of K and V. Cache cap drops from $S_{\max}$ to $W$, at the cost of information truncation, with long-range dependencies relayed through cross-layer stacking.
• INT8 / FP8 KV Cache: quantize fp16 cache down to int8 or even fp8, per-channel or per-token quantization, with controllable error and cache footprint cut by 1/2 to 1/4. Representative work: KIVI / KVQuant.
• KV compression / eviction (H2O, StreamingLLM, SnapKV): drop unimportant positions based on attention weights; used at very long context lengths.
• MLA: mentioned earlier — modifies cache shape at the model-structure level, not as a postprocess.

Per-Step Compute and Memory Cost — Llama 3 8B fp16 · H100 Knee ~330 FLOPs/byte

The shape diagrams above show shapes but not magnitudes. 90% of inference optimization discussion is about “how many FLOPs does this step cost, how many bytes move,” so let’s lay each step’s cost into tables directly.

Using Llama 3 8B, fp16, $B=1$ as the baseline; for prefill take $S=2048$; for decode take $\text{cache\_len}=2048$ (some step during generation around the 2048th token).

Reference hardware knee: H100 SXM fp16 theoretical compute ~989 TFLOPs, HBM bandwidth ~3 TB/s, roofline knee $\text{AI}^{*} \approx 330\ \text{FLOPs/byte}$. Above it is compute-bound, below is memory-bound.

Weight Distribution

Full-model fp16 weights are ~16 GB; the “floor price” of every forward pass is to scan these 16 GB from HBM. At H100’s 3 TB/s, $= 16/3000 \approx 5.3\ \text{ms}$ — this is the physical lower bound of single-request decode.

Per-Layer, Per-Step Compute / Memory I/O

Compare the same layer’s substeps under prefill ($N=S$) and decode ($N=1$). “Weight HBM” is the weight bytes fetched from VRAM; “KV HBM” is the KV-cache bytes read/written. Intermediate activations are assumed fused into kernels and not counted separately.

A few direct conclusions:

• FFN is the real protagonist. $W_{\text{gate}} + W_{\text{up}} + W_{\text{down}}$ consume ~75% of FLOPs and ~80% of weight bandwidth. MoE, sparse activation, and FFN quantization all target this block.
• The 4 attention projections (Q/K/V/O) account for ~18%; the actual $QK^{\top}$ and $\cdot V$ only ~7% — in prefill, attention isn’t the bottleneck — the projections are.
• Decode’s KV reads are 8 MB per layer; at $\text{cache\_len}=2048$ this is only ~2% of weight reads. But once context stretches to 64K or 128K, it grows tens of times, overtaking weight bandwidth as the new bottleneck (this is why Paged Attention, sliding window, and KV quantization exist).

One Full Forward Pass

Adding 32 layers + embedding + LM head:

Decode’s 1.05 FLOPs/byte is 2.5 orders of magnitude below H100’s knee of 330 — meaning ideal single-request decode compute utilization is only $1.05/330 \approx 0.3\%$. This is the mathematical basis for continuous batching: push $B$ to 32 so the same weight read is amortized across 32 requests; arithmetic intensity scales by 32×, decode throughput grows almost linearly until the attention portion or compute itself becomes the wall.

Mental-Math Rules

Two rules cover 90% of inference performance estimation:

1. FLOPs ≈ $2 P N$: $P$ is the parameter count (~8B); $N$ is the total number of tokens this forward pass processes. Each parameter is used once per token (one MAC = 2 FLOPs). E.g., prefill $S=2048$: $2 \times 8\text{B} \times 2048 \approx 33\ \text{TFLOPs}$, matching the itemized sum of 31 TFLOPs.
2. Weight HBM I/O ≈ $2 P$ bytes (fp16): one forward pass scans the model once, about 16 GB.

Arithmetic intensity is essentially $\frac{2 P N}{2 P} = N$ — the total number of tokens participating in this forward. Prefill has $S \cdot B$ tokens; decode has only $B$. This single number directly determines why prefill and decode have different bottlenecks.

Compute Complexity Overview — With vs Without KV Cache Differ by Three Orders of Magnitude

Prefill (process $S$ tokens at once):

• $L \cdot S \cdot H^{2}$ — per layer, 4 $H \times H$ projections + 3 $H \times I$ FFN projections (with $I \sim 4H$), applied to $S$ tokens
• $L \cdot S^{2} \cdot H$ — attention’s $QK^{\top}$ and $\cdot V$, with the $S \times S$ score matrix
• At short sequences linear layers dominate; once $S \gtrsim H$ the attention quadratic catches up

Decode per step (process 1 token, history $\text{cache\_len}$):

• $\text{cache\_len}$ — number of currently cached positions ($= S + t - 1$)
• Single step processes 1 new token, so linear layers’ $S$ becomes $1$; attention still scans the full cache and grows linearly with it

Total complexity to generate $T$ tokens:

• $T$ — number of generated tokens (top-of-article convention)
• $(S + T)$ — average attention span (prompt + generated segment)
• First term sums linear layers across $T$ decode steps; second term approximates the attention portion summed over $t = 1, \ldots, T$ (exact form is $\sum_t (S + t)$)

Without KV cache: $O(L \cdot (S+T)^{3})$ — a massive difference.

Reading FLOPs and bandwidth together is even more illuminating — that’s what the previous section’s table shows: prefill challenges the compute ceiling; decode challenges the bandwidth ceiling; continuous batching’s point is to fuse $N$ requests’ decodes into one large GEMV, amortizing the weight-fetch cost across $N$ requests, with throughput rising linearly until compute or the attention portion becomes the bottleneck.

How Engineering Optimizations Plug into the Formulas — Flash Attn / Spec Decode / Continuous Batch

Flash Attention: mathematically equivalent to standard attention — the formulas don’t change a single character. Engineering-wise it fuses softmax and matmul into one kernel, updating softmax’s running statistics (max, sum) in a streaming manner over blocks, avoiding writing the $S \times S$ attention matrix back to HBM. Complexity unchanged; memory drops from $O(S^{2})$ to $O(S)$; speedup comes mainly from reduced HBM access. FA-2 shifted the partition granularity from heads to query blocks; FA-3 on H100/H200 adds warpgroup MMA + producer-consumer async pipelining.

Flash Decoding: at decode, Flash Attention’s Q has only one row, so kernel parallelism is too low. Flash Decoding splits the $\text{cache\_len}$ dimension of K and V into chunks for parallelism and then does a final log-sum-exp reduction. The formula is the same softmax, just split into two passes.

Speculative Decoding: a small “draft” model generates $k$ tokens sequentially, then the large model verifies them with one prefill over the $k$ positions. Acceptance rule:

• $x$ — a candidate token produced by the draft model
• $p_{\text{target}}(x)$ — probability the large (target) model assigns to $x$ at that position
• $p_{\text{draft}}(x)$ — probability the small (draft) model assigns to $x$ at the same position
• Combined with “on rejection, resample from $\max(0, p_{\text{target}} - p_{\text{draft}})$”, this rule provably yields the same sampling distribution as direct target-model decoding — zero quality loss