Machine Learning Engineering and AI Systems at Scale

The arithmetic that governs everything

A transformer forward pass costs about two floating point operations per parameter per token. The factor of two comes from the multiply and the add in a matrix multiplication, and the rule holds well enough for planning. Training costs about six operations per parameter per token, which is the forward pass plus a backward pass that costs roughly twice as much.

Those two constants answer most capacity questions before any code runs. A 70 billion parameter model costs about 140 gigaflops per generated token. Training it on 15 trillion tokens costs about 6.3e24 operations. A cluster of 8,000 accelerators at 40 percent of a 2.2 petaflop peak delivers about 7e18 operations per second, which is roughly ten days.

The more useful arithmetic concerns memory rather than compute. Generating one token requires reading every weight that participates. For a dense model at batch size one, that is every weight, once, per token. The generation rate is therefore bounded by memory bandwidth divided by model size in bytes, and not by the arithmetic rate at all.

The table explains four industry behaviors at once. Quantization is popular because halving the weight bytes doubles the token rate. Mixture of experts architectures are popular because only the routed experts get read. Larger batches are essential because they amortize one weight read across many sequences. And each accelerator generation advertises bandwidth as prominently as it advertises flops.

Two workloads wearing one name

Inference has two phases with opposite resource profiles, and treating them as one workload is the most common architectural error in this area.

Prefill processes the prompt. Every token of the prompt goes through the model at once, as a matrix multiplication with a large inner dimension. Arithmetic intensity is high, the accelerator runs near its floating point peak, and the phase is compute bound. A 2,000 token prompt on a 70 billion parameter model costs about 280 teraflops, which is a fraction of a second.

Decode produces one token at a time. Each step reads all the weights and performs a matrix-vector product per sequence. Arithmetic intensity is close to one operation per byte, which sits far below the ridge point of every accelerator, so the phase is memory bandwidth bound and the arithmetic units idle.

  achieved
  FLOP/s
    |                            ______________________ peak compute
    |                         __/
    |                      __/         *  PREFILL
    |                   __/               compute bound
    |                __/                  intensity 100-500 FLOP/byte
    |             __/
    |          __/
    |       __/    *  DECODE, batch 1
    |    __/          memory bound
    | __/             intensity ~1-2 FLOP/byte
    |/
    +------------------------------------------------> arithmetic
      1      10      100     1000                       intensity
             ^ ridge point (peak FLOP/s divided by bandwidth)

  Raising the batch size moves decode to the right: the same weight
  read serves more sequences, so intensity grows with the batch.

Two conclusions follow. First, batching helps decode enormously and helps prefill very little, because prefill already saturates the arithmetic units. Second, running both phases on the same accelerator makes each one worse. A long prefill blocks every decode step behind it, which shows up as a latency spike on requests that were already in progress.

Chunked prefill splits a long prompt into pieces and interleaves them with decode steps, so no single prefill monopolizes the device. It reduces tail latency and slightly reduces prefill throughput.

Disaggregated serving goes further and puts the two phases on different machines. Prefill nodes are sized for compute, decode nodes are sized for bandwidth and memory capacity, and the key and value cache transfers between them over a fast interconnect. Each phase then scales independently to its own demand, and reported throughput improvements are substantial. TensorRT-LLM, SGLang, and vLLM all support the pattern.

The key and value cache

Attention needs the keys and values of every previous token. Recomputing them at each step makes decode quadratic, so the system caches them. That cache is the dominant memory consumer in a production serving system, and its size has a closed form.

Bytes per token equals two, for the keys and the values, multiplied by four more terms. Those terms are the layer count, the number of key and value heads, the head dimension, and the bytes per element.

Take a 70 billion parameter model with 80 layers, 8 key and value heads under grouped query attention, a head dimension of 128, and BF16 elements. Per token that is 2 times 80 times 8 times 128 times 2, which equals 327,680 bytes, or 320 kibibytes. A 32,000 token context therefore costs 10 gibibytes for one sequence. Sixty-four concurrent sequences at that length need 640 gibibytes, which exceeds any single device.

Grouped query attention, which shares one key and value head across several query heads, cuts the cache by the grouping factor with a small quality cost. Multi-query attention takes it to the limit with one shared head. Multi-head latent attention compresses keys and values into a low-rank latent vector and expands them on use, which trades a little compute for a large memory saving. Every recent large model uses one of these, and the reason is Table 13.2.

PagedAttention manages the cache the way an operating system manages memory. Instead of reserving a contiguous block for each sequence's maximum length, it allocates fixed-size blocks of a few tokens and keeps a per-sequence block table that maps logical positions to physical blocks. Fragmentation falls from 60 to 80 percent waste to a few percent, and the memory freed becomes additional batch capacity.

The block table also makes sharing cheap. Two sequences with the same prefix, such as a shared system prompt, point at the same physical blocks with a reference count. Prefix caching turns a repeated 2,000 token system prompt from a per-request cost into a one-time cost, which is the single largest win available to an agent or chat workload.

Batching

Static batching waits for a fixed number of requests, runs them together, and returns when the longest one finishes. Every short request waits for the longest, and the accelerator idles on the finished sequences.

Continuous batching, which is also called iteration-level scheduling, runs one decode step for the whole batch, removes sequences that finished, and admits waiting requests into the free slots. No request waits for an unrelated one to finish. Combined with PagedAttention and chunked prefill, this is the design that carries every serious inference server. Together the three techniques report three to five times the throughput of a naive loop on the same hardware.

The scheduler has one tuning axis that matters. A larger running batch raises throughput and raises the time between tokens for every sequence in it, because each step does more work. Set the maximum running batch from the inter-token latency target, not from the memory that happens to be free.

Goodput is the metric to optimize. A server that reaches high throughput by letting time to first token reach ten seconds produced tokens that no user waited for. Measure throughput subject to the latency constraint and tune against that.

Parallelism, priced by its communication

A model that does not fit one accelerator must split across several. Five strategies exist, and the right combination depends on the interconnect far more than on the model.

Data parallelism replicates the model and splits the batch. Communication is one gradient all-reduce per training step, and none at inference time. It scales well and it requires the whole model to fit on one device.

Tensor parallelism splits each matrix across devices. Every transformer layer needs two all-reduce operations in the forward pass to recombine partial results. The communication is frequent and latency sensitive, so tensor parallelism belongs inside one NVLink domain and degrades badly across ordinary network links.

Pipeline parallelism assigns whole layers to different devices and passes activations between stages. Communication is point-to-point and small, so it tolerates a slower link. The cost is the pipeline bubble: with p stages and m microbatches, the idle fraction is (p minus 1) divided by (m plus p minus 1). Raise the microbatch count to shrink it.

Expert parallelism places different experts of a mixture-of-experts layer on different devices. Each token routes to its experts, so the communication is an all-to-all exchange twice per such layer. Load balance is the problem: a routing function that favors a few experts leaves most devices idle, so training adds an auxiliary loss that encourages balance.

Sequence or context parallelism splits one long sequence across devices, which is what makes very long contexts feasible when one sequence's activations and cache exceed a single device.

The practical recipe follows the hardware hierarchy. Use tensor parallelism inside a node, where NVLink gives 1.8 terabytes per second. Use pipeline or expert parallelism across nodes, where the fabric gives a fraction of that. Use data parallelism across replicas. Chapter 4 gave the NVLink-to-PCIe ratio of 28 to 1, and this table is that ratio applied to a model.

Quantization

Reducing the numeric precision of weights cuts memory and, because decode is bandwidth bound, cuts latency proportionally. The question is how much quality it costs.

BF16 is the training default. It has the exponent range of FP32 with fewer mantissa bits, which suits gradients. FP16 has more mantissa and less range, so it needs loss scaling.

FP8 is the current production inference default on hardware that supports it. Two formats exist, E4M3 for weights and activations and E5M2 for gradients. Quality loss on most models is within measurement noise, and the memory and bandwidth halve.

FP4 arrived with Blackwell and it doubles the rate again. It needs per-block scaling factors and careful calibration, and quality loss is real but small on many workloads. The published peak numbers for recent accelerators are FP4 numbers, so compare like with like when reading a datasheet.

Integer quantization takes a different path. Weight-only methods such as GPTQ and AWQ keep activations in higher precision and store weights at 4 bits, which suits the memory-bound decode phase exactly. SmoothQuant moves the difficulty from activations to weights so that both can be 8 bit integers.

The rule for a production decision is short. Quantize weights first, because that is where the bandwidth goes. Quantize the key and value cache second, because that is where the capacity goes. Quantize activations last, because that is where the quality goes. Measure on your evaluation set, not on a published benchmark, because quantization damage is task dependent.

Speculative decoding

Decode is bandwidth bound, so the arithmetic units sit idle. Speculative decoding spends that idle arithmetic to reduce the number of sequential steps.

A small draft model proposes k tokens. The large model then evaluates all k in one forward pass. That pass costs about the same as generating one token, because it is bandwidth bound either way. A rejection sampling step accepts the longest correct prefix, and it guarantees that the output distribution matches the distribution of the large model alone. Acceptance rates of 60 to 80 percent give speedups of two to three times with no quality change.

Variants remove the separate draft model. Medusa adds extra prediction heads to the main model. EAGLE predicts at the feature level rather than the token level. Lookahead decoding uses n-grams from the generation itself. All of them exploit the same fact: verification is nearly free when the bottleneck is reading weights.

Training at scale, briefly

Training memory per parameter is the number that surprises people. Take the Adam optimizer in mixed precision. Each parameter costs 2 bytes for the BF16 weight, 2 for the gradient, 4 for the FP32 master copy, 4 for the first moment, and 4 for the second moment. That is 16 bytes per parameter before any activation memory. A 70 billion parameter model needs about 1.1 terabytes of optimizer state alone.

ZeRO and Fully Sharded Data Parallel shard that state across the data-parallel group instead of replicating it. Stage one shards optimizer state, stage two adds gradients, and stage three adds parameters. Stage three gathers each layer's parameters just before use and releases them after, which trades communication for memory and makes very large models trainable on modest per-device memory.

Activation checkpointing trades compute for memory in the other direction. It discards intermediate activations during the forward pass and recomputes them during the backward pass, which cuts activation memory by a large factor at roughly 30 percent more compute.

Failure handling dominates the operational experience. A ten thousand accelerator job encounters hardware failures on a timescale of hours. Frequent asynchronous checkpointing, fast detection, and automatic restart from the last checkpoint are not optional features, they are the difference between a training run that finishes and one that does not.

Sources and further reading

Woosuk Kwon and others, "Efficient Memory Management for Large Language Model Serving with PagedAttention", SOSP 2023, which is the vLLM paper. Yaniv Leviathan and others, "Fast Inference from Transformers via Speculative Decoding", ICML 2023. Noam Shazeer, "Fast Transformer Decoding: One Write-Head is All You Need", for multi-query attention. Samyam Rajbhandari and others, "ZeRO", SC 2020. Mohammad Shoeybi and others, "Megatron-LM", for tensor parallelism. Jared Kaplan and others on scaling laws, and Jordan Hoffmann and others, "Training Compute-Optimal Large Language Models". The vLLM, SGLang, and TensorRT-LLM documentation for current serving features.