The Logic Chain of LLM Infra

Scope

The training and inference systems of LLMs (large language models), and where their main designs come from.

Architecture

Each layer has two main modules: attention (lets each token gather information from the tokens before it; split internally into several heads computed independently) and FFN (a two-layer fully connected network, computed for each token separately). $$L$$ layers are stacked, with modules such as tokenization (splitting text into tokens, which are words or parts of words) and embedding added before and after. Most parameters are in the matrices of these $$L$$ layers.

Data

During training, the data on the GPU is of four kinds: parameters (the matrices above, with count $$\Phi$$), gradients and optimizer states (one gradient per parameter, plus two statistics per parameter for Adam). These three together are the training state, and the model alone sets their size. Activations (intermediate results of each module's forward pass, needed in the backward pass) vary in size with the batch (the number of sequences processed at once) and the sequence length. The amount of input data is small, so it can be counted with the activations.

The three costs are of different kinds. Memory cost is a capacity: past the memory capacity, the program cannot run. Compute cost and data-movement cost are times, equal to operations ÷ achieved compute and bytes ÷ achieved bandwidth; when compute units sit idle or bandwidth is not fully used, the time grows. HBM is the GPU memory; SRAM is small, fast on-chip storage. H100 figures are for the SXM version: compute is the dense value without sparsity; NVLink and the cross-machine network are counted in one direction, and the cross-machine network assumes one 400 Gb/s NIC per GPU.

Problem

With FP32 throughout, the matrix-multiply peak is only about 1/15 of the BF16 Tensor Core peak, and each activation value takes 4 bytes; with all activations saved, they may exceed the memory capacity.

Solution

Mixed precisiontrade-off

Matrix-multiply inputs and activations use BF16 (2 bytes per number), with products accumulated in FP32; reductions such as softmax and normalization, and the parameter update, use FP32. The update stays in FP32 because an update smaller than about 1/256 of the parameter is rounded away in BF16. The training state is BF16 parameters and gradients at $$2\Phi$$ each, plus FP32 parameters and Adam's two statistics at $$12\Phi$$, for $$16\Phi$$ bytes in total.

Improvescompute: about 15 times the matrix-multiply peakImprovesmemory: activations halvedSacrificesnumerical accuracyTraining state still $$16\Phi$$: FP32 parameters kept

Activation checkpointingtrade-off

Saves one layer's activations every $$\sqrt{L}$$ layers or so; when the backward pass needs the others, it reruns the forward pass from the nearest saved point.

Improvesmemory: activations from $$L$$ layers down to about $$2\sqrt{L}$$ layersSacrificescompute: one extra forward pass, $$R$$ drops to about 3/4

Problem

For some kernels, memory bandwidth sets the time, independent of peak compute; matrix multiplies also wait for data to be read from GPU memory into the chip.

Solution

Kernel fusionwaste removal

Merges several consecutive kernels limited by memory bandwidth into one kernel; intermediate results stay on chip and are not written back to GPU memory.

Improvesdata movement: bytes read and written to GPU memoryImprovesfixed cost of kernel launches

FlashAttentiontrade-off

Attention computes a query, a key and a value vector for each token, stacks them into matrices $$Q$$, $$K$$ and $$V$$, then computes in turn $$S = QK^\top$$ (match scores between every pair of tokens), $$P = \text{softmax}(S)$$ and the output $$O = PV$$. $$S$$ and $$P$$ are both $$n \times n$$ matrices ($$n$$ is the sequence length), about 32 MB per head at $$n = 4096$$. Softmax needs the maximum and sum of a whole row; the standard implementation uses three kernels and writes $$S$$ and $$P$$ back to GPU memory.

FlashAttention uses online softmax to update each row's maximum and sum block by block, merges the three steps into one kernel and computes them in blocks on chip, so $$S$$ and $$P$$ are not written back to GPU memory; the backward pass recomputes $$S$$ and $$P$$ from $$Q$$, $$K$$ and these two per-row statistics. In the paper's example (GPT-2 medium, sequence length 1024, A100), the forward plus backward time of attention drops from 41.7 ms to 7.3 ms.

Improvesdata movement: GPU memory reads and writes 40.3 → 4.4 GBImprovesmemory: removes the $$n^2$$ term in activationsSacrificescompute: backward recomputes $$S$$ and $$P$$, 66.6 → 75.2 GFLOP

Tiling and pipeliningwaste removal

In an $$n \times n$$ matrix multiply, each number takes part in $$n$$ multiply-adds; if each number (BF16) is read or written only once, the arithmetic intensity is $$n/3$$ (about 1365 at $$n = 4096$$). But the whole matrix does not fit on chip. With tiling, data read into the chip is reused many times; the next tile is read while the current one is computed.

Improvesdata movement: bytes read repeatedly from GPU memoryImprovescompute: time compute units wait for data

Problem

One GPU's compute has a limit; frontier-scale training would take over a thousand years on one GPU.

Solution

Each of $$N$$ GPUs holds a full copy of the model and processes different data. Each step averages the gradients with all-reduce (each GPU contributes one piece of data, and at the end every GPU has their sum), and all copies apply the same update. Communication can overlap with the backward computation; when communication takes less time than computation, $$\eta$$ is close to 1.

Problem

$$16\Phi$$ exceeds one GPU's memory, while the $$N$$ GPUs store $$N$$ identical copies.

Solution

Shards in stages, from the least to the most frequently used data.

ZeRO-1waste removal

Shards the optimizer states ($$12\Phi$$). All-reduce can be split into two steps: a reduce-scatter (each GPU gets $$1/N$$ of the gradient sum) and an all-gather (each GPU sends its $$1/N$$ to all GPUs). ZeRO-1 has each GPU update only its own $$1/N$$ of the parameters between the two steps and then all-gathers the updated parameters, so each GPU needs to store only $$1/N$$ of the optimizer states, and the communication volume is unchanged.

Improvesmemory: $$16\Phi \to 4\Phi + 12\Phi/N$$Data movement unchanged: $$2\Phi$$ communicated per step

ZeRO-2waste removal

Also shards the gradients ($$2\Phi$$): after the reduce-scatter, each GPU needs to keep only its own $$1/N$$ of the gradient sum and can free the rest.

Improvesmemory: $$\to 2\Phi + 14\Phi/N$$Data movement unchanged

ZeRO-3trade-off

Also shards the parameters ($$2\Phi$$): the full parameters of each layer are fetched before it is computed, once in the forward pass and once in the backward pass. The FULL_SHARD mode of PyTorch FSDP corresponds to ZeRO-3.

Improvesmemory: $$\to 16\Phi/N$$Sacrificesdata movement: communication per step $$2\Phi \to 3\Phi$$

Memory is counted in bytes and communication in elements, following the convention of the ZeRO paper.

Problem

With many GPUs, syncing parameters takes longer than computing; with long sequences, the activations of one sequence exceed one GPU's memory capacity.

Solution

TPtrade-off

Tensor parallelism: splits each matrix multiply into $$t$$ parts across $$t$$ GPUs; the $$t$$ GPUs compute one copy of the model together, and $$N$$ in the formula above becomes the number of groups $$N/t$$.

Improvesmemory: training state and the activations inside attention and FFN drop to $$1/t$$Improvesdata movement: across machines, each GPU syncs only $$1/t$$ of the parametersSacrificesdata movement: within a machine, 2 all-reduces per layer in each pass; the forward pass waits on them

SPwaste removal

Sequence parallelism, here meaning the Megatron-LM approach: TP does not split LayerNorm and Dropout; each GPU stores its own copy of their activations, about 3/4 of a layer's activations at $$t = 8$$ (not counting the attention score matrices). SP splits them along the sequence into $$t$$ parts and rewrites the original all-reduce as a reduce-scatter and an all-gather of the same total volume.

Improvesmemory: these activations drop to $$1/t$$Data movement unchanged

CPtrade-off

Context parallelism, the other class of sequence parallelism, for long sequences (Li et al. call it sequence parallelism; Ring Attention and DeepSpeed-Ulysses came later): the activations of the whole layer are split along the sequence into $$c$$ parts, and the K and V of other positions that attention needs are obtained through communication. Ring Attention passes them block by block around a ring, overlapped with computation; DeepSpeed-Ulysses switches to splitting by head with an all-to-all (each GPU sends different data to every other GPU); Llama-3 first all-gathers all K and V.

Improvesmemory: activations of the whole layer drop to $$1/c$$Sacrificesdata movement: K and V sent in every layer

Problem

TP is limited to one machine; TP alone cannot hold a large model, and combined with ZeRO-3, syncing parameters across machines still takes about as long as computing.

Solution

Splits the $$L$$ layers into $$p$$ stages, which pass only activations between them. The batch is split into $$m$$ micro-batches fed in one after another, and the stages process different micro-batches at the same time; at the start and end, some GPUs wait, which is called the bubble, a fraction $$(p-1)/(m+p-1)$$ of the time. A larger $$m$$ shrinks the bubble, but with micro-batches that are too small, arithmetic intensity drops and fixed costs take a larger share.

Problem

TP is limited to one machine, PP has the bubble, and DP is limited by the global batch; none is enough alone.

Solution

$$N = t \times p \times d$$ ($$d$$ is the number of DP groups). One configuration of Llama-3 405B pretraining (sequence length 8K) uses $$8 \times 16 \times 128 = 16384$$ H100s, with FSDP for DP; the 128K-sequence stage changes to TP 8, CP 16, PP 16 and DP 8, which the report calls 4D parallelism. By the formula above, communication time ÷ computation time $$= R \cdot d/(B_\text{net} \cdot G)$$:

Llama-3's FSDP does not free parameters after the forward pass and communicates $$2\Phi$$ elements per step; gradients are sent in FP32, for $$6\Phi$$ bytes in total, the same as the $$3\Phi \times 2$$ bytes used in the table.

Problem

In a dense model, the operations per token are proportional to the parameter count; with MoE, one GPU cannot hold all the experts, and with ZeRO-3 each layer has to fetch the parameters of all experts.

Solution

MoEtrade-off

The FFN in each layer is replaced with $$E$$ FFNs of the same structure with separately trained parameters, called experts; a router picks $$k$$ of them for each token. The parameter count grows with $$E$$, and the operations per token grow only with $$k$$. DeepSeek-V3 has 671B parameters in total, and each token passes through 37B of them.

Improvescompute: $$\Phi$$ in $$6\Phi D$$ counts the parameters each token passes through, 671B → 37BSacrificesload balance: the router may concentrate tokens on a few experts, so extra balancing is needed

EPtrade-off

Spreads the experts across GPUs, so each GPU stores only some of them; each token is sent to the GPU holding the expert it picked and sent back after the computation. Compared with ZeRO-3, the parameters stay in place and tokens are sent instead. In DeepSeek-V3, each MoE layer has 256 routed experts (experts chosen by the router), spread over 64 GPUs in training, 4 per GPU; in decode, each GPU holds only 1.

Improvesmemory: each GPU stores only some of the expertsSacrificesdata movement: 2 all-to-alls per MoE layer in each of the forward and backward passes

The comparison here is with a dense model of the same parameter count, and the two use the same total GPU memory; compared with a dense model of the same operations per token, MoE stores about 634B more parameters, which EP spreads over more GPUs.

Problem

Generating the $$n$$-th token recomputes the previous $$n - 1$$ tokens.

Solution

Stores the K and V of each layer; each step computes only the new token, feeding in 1 token per request per step.

Problem

Each decode step reads all the weights once, only to generate 1 token for each of $$b$$ requests.

Limit

$$b$$ is limited by memory capacity: $$Q_\text{w} + b \cdot n \cdot q_\text{KV} \le$$ memory capacity. Llama 2 7B, a 4096-token context, an 80 GB H100: weights 13.5 GB, KV cache 2.1 GB per request, $$b$$ at most 30. Each step then reads about 78 GB, 64 GB of it KV cache; each request reads its own KV cache, so a larger $$b$$ does not amortize these reads.

GQA (several query heads share one set of K and V) and MLA (K and V compressed into short vectors) both reduce $$q_\text{KV}$$.

Solution

Of the four techniques, PagedAttention raises the $$b$$ that continuous batching can reach; chunked prefill and prefix caching are each independent.

Continuous batchingwaste removal

Requests vary in length, and starting and ending a whole batch together leaves empty slots. Instead, at every step, finished requests leave and new requests join.

Improvesdata movement: more useful tokens per weight readImprovesqueueing time of new requests

PagedAttentionwaste removal

The limit of $$b$$ also depends on how well the KV cache memory is used. When memory is reserved for the maximum length, about 20% holds useful data (only about 38% even when the output length is known in advance). Instead, memory is allocated in fixed-size blocks, and the blocks need not be contiguous. Addressing by block makes the attention kernel about 20% to 26% slower, but $$b$$ can grow and end-to-end throughput rises 2 to 4 times, so this is counted as an implementation overhead.

Improvesmemory: about 96% useful data

Chunked prefilltrade-off

When a new request's prefill is inserted, requests in decode wait for it to finish. Chunked prefill splits the prefill into chunks and computes one chunk together with decode each step, using the compute that decode leaves spare.

Improvesdecode stallsSacrificeslatency of a new request's first tokenSacrificesdata movement: each chunk rereads the KV cache of the chunks before it

Prefix cachingwaste removal

Many requests share the same prefix (for example, a system prompt). Prefix caching keeps the KV cache of computed prefixes in GPU memory for reuse; the cache uses only spare GPU memory, and when the batch needs to grow, the least recently used parts are evicted first.

Improvescompute: repeated prefillImprovesmemory: one copy of each shared prefix

Problem

Larger batches raise throughput but do not shorten the generation time of a single request: that time is set by the number of steps, the bytes read per step and the achieved memory bandwidth.

Solution

Flash-Decodingwaste removal

H100 has 132 SMs (groups of compute units that run independently); the more SMs read GPU memory at once, the closer the achieved bandwidth gets to the peak. FlashAttention divides work among SMs by batch, head and query block; in decode the query has only 1 position, and with a small batch there are fewer pieces of work than SMs. Flash-Decoding also splits the KV cache into chunks along its length, spreads them over more SMs, and at the end merges the chunks' results with a small kernel. The Flash-Decoding blog post reports that on A100 with long contexts, attention is up to about 50 times faster than with FlashAttention.

Improvesdata movement: idle SMs also read the KV cache, raising achieved bandwidth

Quantizationtrade-off

Quantizes weights from BF16 to a lower precision such as INT4; when only weights are quantized, they are read in low precision and converted back to BF16 before the computation.

Improvesdata movement: weight bytes read per step about 1/4Sacrificescompute: converting back to BF16, using spare computeSacrificesnumerical accuracy: GPTQ and similar methods use calibration data to reduce the loss of accuracy

Speculative decodingtrade-off

Generation yields 1 token per step, while verifying several given tokens takes only one forward pass. A small model first generates $$k$$ candidates, the original model verifies $$k+1$$ positions in one forward pass, and a rejection-sampling rule decides how many are accepted; the output distribution is the same as generating one by one. The gain depends on the fraction of candidates accepted; with a large batch, compute is no longer spare and the gain shrinks.

Improvesdata movement: the original model reads its weights less often (the small model also reads its weights, but they are far smaller)Sacrificescompute: verifying $$k+1$$ positions

Problem

Prefill is limited by compute and decode by memory bandwidth; sharing the same GPUs, they interfere with each other.

Solution

Prefill and decode run on two sets of GPUs; after prefill, the KV cache is sent to the decode set, once per request. DistServe reports that, with 90% of requests meeting latency targets, the requests served per second per GPU are up to 7.4 times those of the compared systems, or latency targets 12.6 times tighter can be met.

Problem

The training engine lacks inference optimizations such as KV cache management and continuous batching, so generation is slow.

Solution

Rollout uses an inference engine (vLLM, SGLang) and the update a training engine (FSDP, Megatron-LM); each step, the new parameters are synced to the inference engine and resharded to its partitioning. When run synchronously, $$T_\text{RL} = T_\text{rollout} + T_\text{update} + T_\text{sync}$$.

Two deployments: colocated, the two engines take turns on the same GPUs; disaggregated, the two engines are on different GPUs. When disaggregated, the next batch of rollouts can also start early with the previous version of the parameters; the parameters that generate the data are then one or more steps behind (off-policy), and the training algorithm must tolerate this difference.