The math behind how LLMs are trained and served – Reiner Pope
In a Nutshell
Reiner Pope explains LLM training and inference using roofline analysis, revealing that latency and cost stem from balancing compute time (linear in batch size and active parameters) and memory time (dominated by fetching total parameters and KV cache, linear in batch size and context length). Optimal batch sizes (~2K sequences) amortize fixed weight fetches, favoring sparse MoE models like DeepSeek V3 (high sparsity enables larger batches and quality gains), with hardware scale-up (e.g., Blackwell's 72-GPU racks) unlocking bigger models via fast intra-rack all-to-all communication over slower scale-out. API pricing reflects these dynamics—higher costs for fast modes, outputs vs. inputs, and long contexts—while equalizing training, RL, and inference compute predicts over-Chinchilla scaling, and KV cache limits long contexts without bandwidth breakthroughs.
These notes were generated by AI and may contain inaccuracies.
Interview with Reiner Pope, CEO of MatX, a new chip startup. Previously worked on TPU architecture at Google. Blackboard lecture format in new studio on model architecture, ML infra. Understanding training and inference in a cluster explains why AI is the way it is, architectures, API prices, AI progress.
Reiner is an angel investor in MatX, unrelated to podcast.
Companies like Claude, Codex, Cursor offer Fast Mode: 6x price for 2.5x speed token streaming. Mechanically, why pay more for faster latency? Could pay 100x more for much faster? Could have Slow Mode for minutes wait, cheaper prices?
Batch size is big effect on latency and cost. Also speculative decoding or multi-token prediction.
Roofline analysis on Blackwell NVL72 cluster (rack of 72 GPUs): memory bandwidth and compute performance.
Two model factors: time to operate on weights, time on context (KV cache).
Estimate inference time approximation: time ≥ memory fetches time, ≥ compute time.
Compute: multiply by all active parameters, work on attention.
Time_compute ≥ (batch_size * active_parameters) / FLOPs (chip compute throughput). Accounts for weight matrix multiplies; attention compute small, ignored.
Batch: serving many users simultaneously. Without batching, cost/economics 1000x worse.
DeepSeek V3: 37B active parameters, 700B total.
Memory: fetch all total parameters (weights), plus KV cache fetch (depends on batch_size).
KV cache: for each batch element, fetch entire context length * bytes_per_token.
Autoregressive inference (decode): tensor of text tokens (sequence length). Each decode step produces one token: full forward pass on weights, attention to past tokens via KV cache (internal representations). Attention dominated by memory fetches.
Time_memory ≥ max( weight_fetch_time, KV_fetch_time ), where weight_fetch = total_parameters / memory_bandwidth; KV_fetch = batch_size * context_length * bytes_per_token / memory_bandwidth.
Time = max( t_compute (linear in batch_size), t_memory = t_weight_fetch (constant) + t_KV_fetch (linear in batch_size) ).
- t_compute: linear slope.
- t_memory: constant offset (weights) + linear (KV).
Sum: curved slope. Max: thick line. Latency has lower bound from reading total parameters (memory bandwidth limit).
Slopes: compute dominates KV at large batch if KV slope < compute slope; sensitive to context length. Longer context → memory-limited.
Balance (slopes match): equally memory/compute-bound (desirable). Optimal ~100K context; 200K drops MFU significantly (model assumes linear KV fetch; sparse attention scales √context).
DeepSeek uses sparse attention.
Cost = (time * GPU rental, e.g. $2/hour) / tokens_processed = t / batch_size.
- t_compute / B: constant.
- t_KV / B: constant.
- t_weights / B: 1/B parabola.
Max/sum: high cost at batch=1 (unamortized weights), drops to compute limit.
Slow Mode can't amortize much more (KV/compute per batch unique).
Set memory time = compute time (ignore KV for simplicity): total_params / BW = batch * active_params / FLOPs.
Batch ≥ (FLOPs / BW) * (total_params / active_params) ≈ 300 * sparsity.
Sparsity = total/active (DeepSeek: 32/256 experts active → sparsity=8, batch~2400). Practice: double/triple (2-3K tokens/batch, ~2000 sequences).
KV increases optimal batch.
Hardware ratio FLOPs/BW stable (A100→H100→B100); dimensionless via FP4 (0.5 byte/multiply).
Prep with Jane Street engineers Clark, Axel: low-latency trading uses FPGAs for ns latencies (grids touch 100MB SRAM, tens ns response). CPUs collect full 1500-byte packet; FPGAs react to earliest bytes (3ns clock).
Jane Street hiring: JaneStreet.com/Dwarkesh (full prep video there). Liquid cooling, network design.
Train departure analogy: fixed batch interval (e.g. 20ms), new batch every 20ms regardless.
Worst latency: request after departure → wait 20ms + 20ms compute = 40ms.
20ms from memory capacity / bandwidth (~20ms on HBM gens; Rubin: 288GB / 20TB/s = 15ms). Time to read all HBM (weights/KV); mostly reads.
Optimal batch independent of model size (only sparsity).
Tokens/sec = batch_size / interval * 60 ≈ 64 * batch_size (~128K tps for 2K batch).
1/1000th of Gemini's hundreds millions tps worldwide (still competitive scale).
Centralization: 2000 concurrent users not huge.
Higher sparsity → less compute (compute bottleneck at large batch).
Quality vs speed empirical.
Paper: "Unified Scaling Laws for Routed Language Models".
Hold active params fixed, increase sparsity (expert count): quality increases (e.g. 64-expert 370M active = dense 1.3B; 64x params for ~4x active).
Increasing sparsity: compute win, memory amortized by larger batch. Pure win until user limit; sparser with more users.
MoE history: 2017, GShard, Switch Transformer, DeepSeek innovations.
Sparsity consumes memory capacity in addition to memory bandwidth. To spend less time computing, increase sparsity, which requires bigger batch sizes and thus more memory capacity.
Tokens go through a router layer that decides which experts to route to—a small fraction, e.g., 1 in 32. Each expert is a normal MLP with up projection, nonlinearity, and down projection. Results are combined via inverse operation and added to residual connections.
Expert parallelism places different experts on different GPUs. DeepSeek model has 256 experts. On a Blackwell rack with 72 GPUs, simplify to 64 GPUs (ignore 8), so 4 experts per GPU; for diagram, use 2 experts per GPU.
Communication: Tokens routed centrally to experts across GPUs (all-to-all pattern), with cost on input and output. Router has copies on each GPU.
One rack bounds expert layer size due to rack-to-rack communication bottleneck. Blackwell racks fit MoE all-to-all perfectly within rack.
Rack: Physical structure ~few meters tall, 1-2m wide, holds ~64 GPUs/XPUs, constrained by power delivery, weight, cooling.
Nvidia: GPUs on outside, NVSwitches inside; every GPU connects to switches (NVLink, scale-up network), enabling all-to-all in 2 hops at high bandwidth. Rack-to-rack (scale-out network) ~8x slower.
Across two racks, half tokens stay in-rack (fast), half cross (slow bottleneck).
Hierarchy of switches manages cabling congestion. Nvidia progressed: Hopper (8 GPUs), Blackwell (72), Rubin (~500) via form factor (trays to racks) and denser cable designs.
Wire density: Doubling GPUs requires twice cable density; constraints include connector density, backplane, cable bend radius, space/weight/power/cooling.
GPT-4 (2023, rumored >1T parameters). Hopper (2022): 8 GPUs, 640GB. Blackwell (recently deployed): 10-20TB scale-up, enough for 5T model + KV cache.
Larger scale-up domains unlock larger models. Google had large scale-up domains, aiding Gemini pre-training.
DeepSeek MoE: more but finer-grained experts. Limits: active parameters by compute, total by scale-up size.
All-to-all (MoE) favors single rack. Other parallelisms: tensor parallelism (less relevant with smaller experts), data parallelism, pipeline parallelism (better for multiple racks).
Pipeline across layers/racks: e.g., change rack between MoE layers. Compare scale-out vs scale-up bandwidth.
Scale-up 8x faster than scale-out. Ratio time_scale-up / time_scale-out = (1/8) * (activated_experts) * (layers_per_stage) * 2 (up/down).
Want ratio >=1, so product >8. With ~8 activated experts (DeepSeek ~16-32), can increase layers_per_stage.
Can pipeline racks (one layer per rack). Matches model architecture: experts on GPUs, layers on racks.
Model scales in layers, model dim, DFF dim, experts; cut along profitable dimensions (experts, layers).
Ilya: Avoid pipeline parallelism due to bubbles/architectural constraints (e.g., Kimi residuals attending back).
Benefits: Saves memory capacity (not runtime/memory time). Reduces per-rack memory bottleneck, enabling faster scaling.
Axel (GPU engineer, Jane Street): Pipelining requires micro-batches, reducing weight loading amortization.
Inference: Forward pass across racks (e.g., 4 racks). Run sequential inferences to fill bubbles (no idle time).
Training: Forward micro-batches, then all-to-all backward (hard stop for optimal batch size). Smaller batches better for convergence (fresher gradients), larger better for systems efficiency; trade-off.
Techniques: zero-bubble, 1F1B interleave forward/backward.
Inference: Neutral on batch size/latency, saves memory capacity per rack (e.g., quarter model).
Even in inference, not heavily used: racks have surplus memory (Blackwell rack >>1TB for 1T model).
Theoretically allows hardware with less HBM/GPU by spanning racks.
Created animations for a blog write-up on how weights shard and gradients flow depending on parallelism using Cursor. Gave Cursor lecture notes and a sketch made during the lecture to visualize a specific hierarchical collective explained by Horace. The first version was already pretty good, then used design mode to select and tweak specific components. Did this without a clear end state in mind. Cursor's Composer 2 Fast model was quick enough for instantaneous iteration: try an idea, test in the built-in browser, make changes immediately. Went through 10 versions in under 20 minutes. Published the animation with lecture notes in a blog post—link in description. To try iterative design flow, go to cursor.com/dwarkesh.
Everybody's talking about the memory wall: memory is super expensive, not enough memory, smartphone volume will go down 30% because there's not enough memory. Hyperscalers spending 50% of CapEx this year on memory—high hundreds of billions or maybe a trillion total CapEx, half on memory, a huge constraint. That's why no new laptops and phones this year. But at the same time, too much memory in systems—why is Jensen shoving all this memory into racks if not needed? Previously discussed memory time, memory bandwidth, compute bandwidth; now looking at memory capacity without parallelism scheme.
Memory Capacity Demand
Demand on memory is number of total parameters N_total to fit weights, plus KVs as batch size B times context length L times bytes per token. Argued for pipelining to solve this. Consider running on number of GPUs with E (expert parallelism, e.g., 64 GPUs sharding an expert layer) and P (pipelining extent, e.g., 4 racks). Total memory requirement across system, but per-GPU memory c_mem is N_total + B * L * bytes_per_token divided by E * P.
Parameters perfectly divided among GPUs in a rack; layers perfectly divided among racks. Contexts perfectly sharded across GPUs in rack and layers across racks.
Micro-Batching in Pipelining
Analyze batch size B. Micro-batching vs. global batching. In pipelining diagram, one batch goes forward, then loops for decode: generate tokens sequentially, writing to KV cache. Split global batch size B into number of micro-batches times batch size per micro-batch b. Number of micro-batches equals number of pipeline stages (e.g., 4: zero, one, two, three) to avoid idle time—visual proof as it wraps around. Micro-batch size b ~2000 (300 times sparsity), the 20ms train takeoff size. Global B = micro-batches * b. Number of micro-batches minimized to fill pipeline, equals P.
Frontier models use pipelining during massive scale training; can be done for inference but less attractive—useful for weights, not KVs.
Memory Footprint with Pipelining
Micro-batch size equals P, so substitute: P * b * L * bytes in c_mem. Factor: weights term divides by E * P, KV term divides by E only (P cancels). Increasing P reduces weights memory per GPU but KV memory stays constant. After enough pipelining (even P=2), weights term small, KV cache dominant.
KVs not shared between layers, but pipelining across layers means storing fewer KV layers per stage. However, to keep all racks busy, more sequences in flight simultaneously—cancels out, no per-GPU saving. Can't amortize KV caches across batch or pipeline stages.
Inference Parallelism Strategies
For inference, DeepSeek paper: lots of expert parallelism up to scale-up domain size, very little pipelining (maybe none or two, enough for weights). Only expert parallelism and minimal pipelining make sense. Past tensor parallelism (cutting within expert) not profitable now—experts too small. Frontier labs do inference within single scale-up, unless model exceeds rack memory or extremely sparse.
Model size scaling slow until recently. Pipelining solves rack capacity for model parameters (not KV cache). Larger scale-up domains not due to capacity, but other factors.
Latency Costs in Pipelining
Rack-to-rack hop latency: few milliseconds (order of magnitude uncertain), e.g., P=4 realistic. ~10ms per token (from GPU/TPU to network card, top-of-rack switch, other rack, reverse; maybe data center switch). In decode, latencies stack sequentially across stages. Small P not huge impact, but adds up (e.g., 20 to 30ms).
Scale-up size relevant via t_mem latency improvement: t_mem for weights = N_total / memory bandwidth. Bandwidth = scale-up size * GPU bandwidth. Increased 8x from Hopper (vs. 1.5-2x per generation for other terms). Pipelining solves capacity, scale-up solves bandwidth—enables longer contexts (key for agentic models), lower latency (e.g., sparse model on H100 box high latency).
Chinchilla scaling for training compute optimal size vs. data. Now optimize mixture of training and inference compute. With RL, pre-training used for RL generation and final inference. Over-training: smaller model trained longer than Chinchilla (faster learning), cheaper inference.
Heuristic for Cost Equalization
Total cost = training cost + inference cost; minimum where costs equal (heuristic for 1/x + x, power laws, etc.). Equalize pre-training, RL, inference costs.
- Pre-training: 6 * active params N * pre-train data D_pre (6ND FLOPs).
- RL: ~2-6 * N * D_RL (forward+backward 6 if all rollouts trained; decode inefficiency, not all rollouts used; decode MFU < training).
- Inference: 2 * N * D_inf (forward only).
Ballpark: 33% split each (or 50% training/RL + 50% inference). N cancels. D_pre + α * D_RL = β * D_inf.
α ~2-6 (FLOPs) * 30% inefficiency → ~1/10 to 2/10. β ~1/3 * 30% →1/10. Adjust: inference tokens ~ hundreds millions/sec * 2 months deployment.
Estimating Real Numbers
500M tokens/sec * 2 months = 2.6e15 tokens (too large; adjust 5-10x smaller for family → ~50M/sec/model → 200T tokens inference). Pre-training ~150T tokens rumored. Active params ~100B → Chinchilla D_chin = 20 * 100B = 2T tokens. Actual ~100x over Chinchilla (200T / 2T). Equalizing costs first-principles predicts pre-training data matching inference tokens.
Models generate sum of human knowledge in outputs equaling pre-training inputs. Labs err toward less inference if prediction imperfect or risk of non-frontier model discarded (derate inference).
Gemini 3.1: 50% more expensive over 200k tokens vs. below. Cost increases with context length (memory time vs. compute time).
Time: t_mem (weights + KV) vs. t_compute (matmuls). As function of L: compute constant (mild dependence ignored); t_mem starts high (weights), grows linearly with L (KV). Total time = max, inflection point. Two-tier pricing ensures profitability regardless of L.
So we've got a two-tier pricing structure.
Given the bump at 200k tokens, this aligns somewhat with the crossover point between memory and compute time. Solve for bytes per token assuming equalization at 200k tokens, with large batch size making weight memory time negligible. Focus on KV cache memory time: batch × context length × bytes per token / memory bandwidth = activated params / FLOPs per second.
Batch size cancels out. Bytes per token = (activated params / FLOPs per second) / context length. Plugging in numbers: 1/300 (reciprocal from earlier, stable across hardware), activated params ~100 billion, context length 200k → ~1667 bytes per token (~2 kilobytes), plausible.
Dense attention with reuse across layers (e.g., Character AI blog post, Gemma models): global context shared across layers. Bytes = attention layers × 2 × dhead (e.g., 128) × KV heads. In Character AI-like models, unique contexts per layer = 1.
KV heads store previous tokens (1-8 typical); Q heads temporary for attending token. Example: 8 KV heads × dhead 128 = 1667 bytes exactly. Or fewer KV heads, more layers.
Sparse attention: increase numbers but add 1/sparsity term.
API pricing leaks info: incentivized to price near costs.
Output ~5x more expensive than input (past: 3-5x). Decode: compute for next token. Prefill: process all tokens in parallel, length of pass (decode: len=1, prefill: many).
In chat: decode AI response, prefill user input/file/tool call (KV cache pre-computed, flash attention temporary—no main memory cost). Memory time unchanged; compute 5x higher for output → prefill cheaper per token.
Chart time per token (t / len-pass) vs len-pass. Decode (len-pass=1): t_compute flat; t_mem high (memory-bound). Prefill (larger len-pass): t_compute flat/divided (same); t_mem declines (1/len-pass). Prefill compute-limited, decode memory-bandwidth limited. 5x pricing gap shows heavy memory bandwidth bottleneck at decode.
Cache hits cheaper (e.g., 10x if hit vs miss). Miss: recompute KV from tiny token IDs. Hit: load from memory (HBM vs cheaper DDR/flash).
Cost per token to rematerialize: forward pass on whole model (compute time × GPU $/s). No quadratic in approximation (FLOPs: weight multiplies flat vs context; KV cache linear but low slope—negligible until millions tokens).
Long context limits: memory bandwidth > compute (fundamental). No company >1M tokens due to bandwidth/capacity, not compute.
"We don't need continual learning for AGI, in-context learning is enough." (Dario Amodei et al.) → Need 100M-token context for month-equivalent employee. Requires ML infra changes (memory bandwidth). Sparse attention helps (√ scaling) but context lengths stalled at 100-200k (from GPT-3 8k) → cost-prohibitive, memory wall. HBM not improving hugely; sparse too much loses quality.
- Rematerialization: compute FLOPs / GPU speed × GPU $/s.
- Store HBM: bytes/token × (HBM fraction used × GPU $/s) + retrieve (bytes / HBM bandwidth × GPU $/s).
- Store DDR: similar, higher retrieve (bytes / DDR bandwidth) + hold (bytes / DDR capacity × DDR $/s).
Retrieve: time to move to HBM for inference (compute overlaps if already in HBM). Balance retrieve vs hold costs across tiers; hold time × fraction = retrieve time.
Base input $5/M tokens (remat/retrieve). 5-min write $6.25 (HBM?). 1-hour write → lower tier (DDR/flash?). Drain time = capacity / bandwidth.
HBM: 20ms. DDR: ~1-10s. Flash: ~1min. Spinning disk: ~1hr. 5min/1hr → flash/spinning disk? Equalize costs: hold time ≈ drain time (read whole memory once per hold period).
Shocked at spinning disk use despite slowness/old tech.
Architectures similar: both jumble info across inputs. Crypto: scramble structure to random-like. NNs: extract structure from random-looking data (proteins/DNA/garbled text). Convergent evolution.
Other examples: cake batter stirring (alternate directions).
Random init NN ≈ cipher. Gradient descent + residuals/LayerNorm simplify derivatives (interpretable).
Crypto attacks: differential cryptanalysis (small input diff → large output diff; avalanche). NNs opposite (smooth gradients). Backdoors/adversarial attacks: hide backward-pass input, small perturb → big output change (cipher-like undesired in NNs).
NNs as ciphers: dangerous (need 10yr scrutiny; most broken).
Feistel network: non-invertible f (e.g., MLP) → invertible layer (apply f, swap with remembered input). Imported to NNs: RevNets (2018/2019 paper). Residual + remember input → reversible layers, rematerialize activations backward, near-zero training memory footprint.
Or it mixes it in an interesting way. You'd like to build something out of this that is invertible. The construction is a two-input function rather than a one-input function. Apply f[x], remember what x was to work backwards, and can't drop y. Remember y and add them together to form this tuple. To invert, given the output z, recover x directly as it's there. Recover y by z minus f[x] since x is recovered. This construction is invertible.
This was used in ciphers a ton and still is. It's one of the main mechanisms of constructing ciphers. Ciphers often need to be invertible, especially layers, for better cryptographic properties.
This has been ported into neural nets. There's a 2017 paper called RevNets, reversible networks. It makes the entire network invertible. You can apply it to any network, like a transformer network. Do a forward pass, then run the entire pass backwards. The whole neural network is invertible with this construction. The paper applied it to a layer like a transformer layer.
f is the transformer layer. Normally, input with a residual connection added. The variation: two inputs x and y. x goes through the function, added to y, becomes new x (output x). Then this x becomes output y. This does the residual connection from two layers back. y came from the previous layer as the residual there. Because of this, the whole thing is invertible.
Why care about invertibility? For training. In forward pass of training with four layers (zero, one, two, three), write all activations to HBM. HBM footprint linear in number of layers. This can be the largest memory footprint during training. Normal training: forward pass forward, backward pass reads in reverse.
RevNets idea: because invertible, no need to store activations. Completely rematerialize. Run forward pass, then in backward pass, simultaneously undo forward steps in lockstep to have needed activations. Ends up memory saving. Spending more compute to save memory.
Opposite of KV cache: spending more memory to save compute, generally profitable given hardware.
Keep Dwarkesh Patel in your library
Save the videos and channels worth coming back to, and find them again in one place.
