ML Systems: Quick recap

01 GPU architecture

A GPU trades fast single threads for throughput: an H100 has 132 SMs with 128 fp32 lanes each, and needs many independent pieces of work.
A kernel runs over a grid of thread blocks; each block runs on one SM; the SM schedules warps of 32 threads that execute one instruction together (SIMT). Divergent branches in a warp run one after the other.
Blocks are independent and are not all resident at once; the blocks that run together form a wave, and a launch just above a full wave pays for a mostly idle extra wave (the tail effect).
An SM hides the latency of HBM loads (hundreds of cycles) by switching between resident warps at no cost. Little's law: data in flight = bandwidth × latency, about 1.3 MB (10 KB per SM) on an H100.
Occupancy = resident warps / 64; it is limited by registers (65,536 per SM), shared memory (228 KB), threads (2,048) and blocks (32). It is a means, not a goal.
Memory hierarchy of an H100: registers and shared memory/L1 (256 KB each per SM), L2 (50 MB, ≈ 5.5 TB/s), HBM (80 GB, 3.35 TB/s). Weights and activations live in HBM, so kernels read and write HBM.
A tensor core computes $D = AB + C$ on small low-precision tiles with fp32 accumulation; an H100 does 989 TFLOP/s dense bf16 on tensor cores and 67 TFLOP/s fp32 on CUDA cores, about 15× less.

02 Roofline and arithmetic intensity

Arithmetic intensity $I$ = FLOPs / bytes moved between HBM and the chip; ridge point $R = P/\mathrm{BW}$ (H100 bf16 ≈ 295, fp8 ≈ 591, A100 bf16 ≈ 153 FLOP/byte).
Roofline bound: $t \ge \max(\text{FLOPs}/P,\ \text{bytes}/\mathrm{BW})$ ; a kernel with $I < R$ is memory-bound and reaches at most $I \cdot \mathrm{BW}$ FLOP/s.
Elementwise ops and norms have $I$ below 5: memory-bound on every GPU, so they speed up by the bandwidth ratio (H100/A100 ≈ 1.6×), not the FLOP/s ratio.
A $(B \times D)(D \times F)$ matmul has $I = B/(1 + B/D + B/F)$ : about $B$ for small $B$ , at most $DF/(D+F)$ .
It is compute-bound for $B \ge B^* = R\,DF/(DF - R(D+F))$ , about $R$ for large $D, F$ (345 tokens for $D = F = 4096$ on an H100); fp8 leaves $B^*$ about the same.
$B$ counts the tokens that share one weight read: micro-batch × sequence length in training and prefill, the number of sequences in decode.
Decode attention has $I \approx 1$ (MHA) or $\approx g$ (GQA) at any batch size, because every sequence reads its own KV cache.
The four-step estimate: count FLOPs, count bytes, divide by $P$ and BW, take the larger time.

03 Kernels, tiling and fusion

A CUDA kernel runs a grid of independent blocks; threads of one block share shared memory and synchronize with __syncthreads(), which every thread of the block must reach. Triton programs one tile per program instance, with a mask at the boundary.
Global memory is read in 32-byte sectors. A warp load is coalesced when consecutive threads read consecutive addresses; a column read of a row-major fp32 matrix uses only 12.5% of the fetched bytes.
Shared memory has 32 banks of 4 bytes; word $w$ is in bank $w \bmod 32$ . $k$ different words in one bank cost $k$ passes; padding tile[32][32] to tile[32][33] removes the column conflict.
Each kernel costs about $\max(o, \text{bytes}/\mathrm{BW})$ , with a launch cost $o$ of a few µs in eager PyTorch.
Fusing a chain of $k$ unary elementwise ops cuts the traffic from $2kN$ to $2N$ elements and the launches from $k$ to 1. Large tensors gain the ratio of bytes, tiny tensors the ratio of launches.
A tiled matmul reads $MK\lceil N/T\rceil + KN\lceil M/T\rceil \approx 2MNK/T$ elements instead of $2MNK$ ; the bf16 intensity per step is $T/2$ .
The tile size is limited by shared memory: $2T^2e \le S$ (227 KB per block on an H100 gives $T \le 241$ in bf16). Real GEMMs reach the peak through L2 reuse across blocks and register tiling.

04 FlashAttention and online softmax

The online softmax keeps a running max $m$ , sum $\ell$ and output $\tilde o$ per query; when the max grows, it multiplies $\ell$ and $\tilde o$ by $e^{m_\text{old} - m_\text{new}}$ , and at the end $\tilde o / \ell$ is the exact softmax average.
The rescale is exact because the factor is the same for every old term; every exponent is $\le 0$ , so nothing overflows.
Two partial results over disjoint key sets merge as $\mathrm{lse} = \log(e^{\mathrm{lse}_A} + e^{\mathrm{lse}_B})$ and $o = e^{\mathrm{lse}_A - \mathrm{lse}} o_A + e^{\mathrm{lse}_B - \mathrm{lse}} o_B$ : a weighted, not a plain, average.
FlashAttention keeps a $Q$ tile on chip, streams $K, V$ tiles, and never writes $S$ or $P$ to HBM; it writes only $O$ and one lse per row.
HBM traffic per head: standard attention $4N^2 + 4Nd = \Theta(Nd + N^2)$ ; FlashAttention $2Nd + 2N^2 d / B_r = \Theta(N^2 d^2 / M)$ , with $B_r$ the query tile height.
The backward pass recomputes $P_{ij} = e^{S_{ij} - \mathrm{lse}_i}$ instead of storing it: 25% more matmul FLOPs, far fewer bytes, faster, because attention is memory-bound.
FlashAttention-2 reduces non-matmul FLOPs, parallelises over query tiles and splits query rows across warps; split-KV (flash-decoding) splits the KV cache into chunks and merges the partials, so decoding uses all SMs.

05 Mixed-precision training

Exponent bits give range, mantissa bits give precision: bf16 (1/8/7) has the range of fp32 and $\varepsilon = 2^{-7}$ ; fp16 (1/5/10) has $\varepsilon = 2^{-10}$ but a largest value of 65504 and underflows below about $6 \times 10^{-8}$ .
Stagnation theorem: if $0 \le x < \mathrm{ulp}(S)/2$ , then $\mathrm{fl}(S + x) = S$ . A sum of equal terms stops after about $2^p$ terms (256 in bf16, 2048 in fp16).
Long sums are fixed by fp32 accumulation (tensor cores: 16-bit inputs, fp32 adds), pairwise summation or Kahan summation.
fp32 master weights are needed because an update below about $2^{-p}|w|$ is lost in a $p$ -bit format; the optimizer step runs in fp32.
Softmax, norms, losses and large reductions run in fp32; elementwise ops stay in 16 bits.
Loss scaling multiplies the loss by $S = 2^k$ , so every gradient is multiplied by $S$ ; unscale before clipping and the update. Dynamic scaling skips the step and halves $S$ on inf, and doubles $S$ after 2000 clean steps.
bf16 needs no loss scaling (fp32 range); FP8 needs a scale factor per tensor or block, with E4M3 for weights and activations and E5M2 for gradients.

06 Training memory and checkpointing

Training memory = model states (depend on the model) + activations (depend on the tokens) + logits + buffers. Keep 10–20% of the GPU free.
Mixed-precision Adam keeps 16 bytes per parameter: bf16 weights 2, bf16 gradients 2, fp32 master weights 4, Adam $m$ 4, Adam $v$ 4. So a 7B model needs 112 GB of model states.
A GPT-style layer has $12h^2 + 13h$ parameters, so $\Psi \approx 12Lh^2$ for large $h$ .
Every operation keeps what its backward pass needs. One transformer layer keeps $sbh(34 + 5as/h)$ bytes; the $5as^2b$ attention core grows with $s^2$ .
FlashAttention and selective recomputation remove the core ( $34sbh$ left) at almost no cost. Full checkpointing keeps $2sbh$ per layer for about 33% more FLOPs.
Memory peaks when the backward pass starts. Checkpointing every $k$ layers gives a peak of about $(L/k)c + kA$ , smallest at $k^* = \sqrt{Lc/A}$ ; for transformers $k^*$ is 1 to 3.
The fp32 logits take $4sbV$ bytes. Data parallelism and ZeRO do not reduce activations.

07 Collectives and the α-β model

α-β model: a message of $m$ bytes takes $\alpha + m\beta$ ; $\alpha$ is the latency per message, $\beta$ is one over the bandwidth.
All-reduce = reduce-scatter + all-gather. All-to-all is a block transpose: rank $j$ receives block $j$ of every rank.
Ring all-reduce: $2(p-1)\alpha + \frac{2(p-1)}{p}n\beta$ . Each rank sends about $2n$ bytes for any $p$ ; only the latency term grows with $p$ .
No all-reduce can make its busiest rank send fewer than $\frac{2(p-1)}{p}n$ bytes, so the ring is bandwidth-optimal.
Recursive doubling: $\lg p\,(\alpha + n\beta)$ . It wins below the crossover size $n^*$ (latency regime); the ring wins above it (bandwidth regime).
Pairwise all-to-all: $(p-1)\alpha + \frac{p-1}{p}n\beta$ ; on a real network it is limited by bisection bandwidth and by the largest block.
busbw = algbw × $2(p-1)/p$ for all-reduce (× $(p-1)/p$ for reduce-scatter, all-gather and all-to-all); compare busbw, not algbw, with the link speed.
Hierarchical all-reduce: reduce-scatter on NVLink, all-reduce of $n/g$ across nodes on all network cards in parallel, all-gather on NVLink.

08 Data parallelism, ZeRO and FSDP

Data parallelism: each GPU computes the gradient of its own micro-batch; the all-reduce averages them. With equal micro-batches and a mean loss this equals one-GPU training on the whole batch, and the copies stay identical.
Plain data parallelism keeps $(4+K)\Psi = 16\Psi$ bytes of model states on every GPU.
ZeRO-1, -2 and -3 shard the optimizer states, then the gradients, then the weights: $4\Psi + K\Psi/N_d$ , $2\Psi + (2+K)\Psi/N_d$ , $(4+K)\Psi/N_d$ bytes per GPU.
ZeRO-1 and ZeRO-2 send $2\Psi$ elements per GPU per step, like DDP (reduce-scatter, update, all-gather). ZeRO-3 sends $3\Psi$ (two all-gathers and one reduce-scatter), 1.5 times DDP.
ZeRO and FSDP shard model states only, never activations.
DDP overlaps a bucketed all-reduce with the backward pass; the last bucket (the first layers) is always exposed.
FSDP gathers one unit at a time before forward and backward, frees it after use, and reduce-scatters the gradients. HSDP shards inside a node and replicates across nodes.
The all-reduce hides when the tokens per GPU per step satisfy $T \ge F/W$ (twice that for fp32 gradients). The model size cancels.

09 Tensor and pipeline parallelism

Megatron TP: split the first matrix of a block by columns and the second by rows; then each block needs one all-reduce in the forward pass ( $g$ ) and one in the backward pass ( $f$ ): 2 + 2 per transformer layer.
An elementwise function (GeLU) can sit between a column split and a row split; splitting the first matrix by rows would need an extra all-reduce before it.
TP communication over compute per layer is $(t-1)F/(3hW)$ : independent of the tokens $T$ , worse with $t$ , better with $h$ ; about 0.6 on NVLink and 5.6 on InfiniBand for $h = 8192$ , so TP stays inside a node.
Pipeline bubble of GPipe and 1F1B: $\frac{p-1}{m+p-1}$ of the step time ( $\frac{p-1}{m}$ of the ideal time); interleaving with $v$ chunks gives $\frac{p-1}{vm+p-1}$ at $v$ times the messages.
GPipe keeps $m$ micro-batches per stage; 1F1B keeps at most $p - s$ on stage $s$ : it saves memory, not time.
$m = B/(bN_d)$ : at a fixed global batch, more data parallelism means fewer micro-batches and a larger bubble; aim for $m \ge 4p$ .
Layout rules: TP inside the NVLink domain, PP across nodes with $t \cdot p$ just large enough to fit, DP outermost; the Llama 3 order is [TP, CP, PP, DP].

10 Sequence, context, expert parallel

Only attention mixes tokens; norms, dropout, residual adds, projections, MLPs and MoE routing are token-wise and can be split along the sequence for free.
Megatron sequence parallelism shards the norm and dropout regions along $s$ inside the TP group; each all-reduce becomes a reduce-scatter plus an all-gather with the same bytes, and the activation memory per layer becomes $sbh(34 + 5as/h)/t$ .
Ring attention keeps the queries on each rank and sends the K/V blocks around a ring of $N_c$ ranks; partial results merge exactly with the logsumexp formula, in any order.
Ring attention hides the K/V transfer when the block size is $c \ge F/W$ (multi-head attention, bf16); GQA divides the bound by $g$ , a balanced causal mask multiplies it by 2.
With a causal mask, contiguous blocks give about 50% efficiency; the zig-zag layout (rank $r$ holds chunks $r$ and $2N_c - 1 - r$ ) gives every rank the same work in every round.
DeepSpeed-Ulysses uses two all-to-alls (sequence → heads → sequence) and sends about $8sh/N_c$ bytes per GPU and layer, but its degree is limited by the number of (K/V) heads and its communication is not overlapped.
An MoE layer with top- $k$ routing uses $k$ of $E$ experts per token: FLOPs follow the active parameters, memory follows all parameters.
Expert parallelism sends $2kTh\,e\,(N_e - 1)/N_e$ bytes per GPU per MoE layer (dispatch plus combine, uniform routing); the busiest GPU sets the time, so capacity factors and load-balancing losses keep the load even.

11 Quantization

Uniform affine quantization: $q = \mathrm{clamp}(\mathrm{round}(x/\Delta) + z, q_{\min}, q_{\max})$ , $\hat x = \Delta(q - z)$ . Symmetric: $z = 0$ , $\Delta = \max|x|/(2^{b-1} - 1)$ . Asymmetric: $\Delta = (\max - \min)/(2^b - 1)$ , $z = \mathrm{round}(-\min/\Delta)$ .
Inside the range the error is at most $\Delta/2$ ; if the error is uniform (data spread over many steps), the mean squared error is $\Delta^2/12$ . Outside the range the clipping error has no bound.
Clipping trade-off: a smaller range gives a smaller $\Delta$ but clips more values; the MSE has a minimum (about $2.5\sigma$ for Gaussian data at 4 bits).
When values share a scale, the largest one sets $\Delta$ : one outlier makes the step coarse for all the others.
An integer matmul can use scales on the outer dimensions (per token, per output channel); scales on the reduction dimension (per input channel, group-wise) must be applied inside the sum. int4 with fp16 scales per group of 128 costs 4.125 bits per weight.
SmoothQuant: $XW = (X\,\mathrm{diag}(s)^{-1})(\mathrm{diag}(s)\,W)$ exactly; $s_j = \max|X_j|^\alpha / \max|W_j|^{1-\alpha}$ moves range from activation outliers into the weights, at no run-time cost.
PTQ uses a calibration set (RTN, GPTQ, AWQ); QAT trains through fake quantization with the straight-through estimator (gradient 1 inside the range, 0 outside).
W4A16 removes weight bytes: it speeds up memory-bound decode (up to 4×), not compute-bound prefill. W8A8/FP8 halves the bytes and doubles the FLOP/s: up to 2× in both regimes.

12 LLM inference serving

KV cache per token: $k = 2 L n_{kv} d_h e$ bytes (128 KiB for Llama-3-8B in bf16); the batch that fits is $B_{\max} = \lfloor M_{kv} / (Ck) \rfloor$ .
A decode step reads $W + BCk$ bytes and does about $BW + gBCk$ FLOPs (bf16). Its intensity is a byte-weighted average of $B$ (shared weights) and $g = a/n_{kv}$ (private caches), so decode is memory-bound in practice.
Prefill is compute-bound (intensity about $T$ ); decode is memory-bound. TTFT measures the first, TPOT the second; a larger batch raises throughput and TPOT together.
Continuous batching re-decides the batch before every step: finished sequences leave, waiting ones join, and no slot idles while requests wait.
PagedAttention stores the KV cache in fixed-size blocks mapped by a block table: waste below one block per sequence, sharing by reference counts, copy-on-write, and preemption of the newest request (swap or recompute).
Speculative sampling accepts a draft token with probability $\min(1, p/q)$ and resamples from $(p - q)_+$ on rejection; every token then has exactly the target distribution.
With acceptance rate $\alpha$ and draft length $\gamma$ , one target pass gives $(1 - \alpha^{\gamma+1})/(1 - \alpha)$ tokens on average; the speed-up is that divided by $1 + \gamma c$ .
Long prompts make TPOT spike; chunked prefill hides prompt chunks in the idle compute of decode steps, and disaggregation moves prefill to other GPUs at the cost of a KV transfer.

13 Profiling, MFU and compilation

Training FLOPs per token of a dense decoder: $F_\text{model} = 6\Psi + 12Lhs$ (PaLM convention, $h$ = attention width); full recomputation makes the hardware run $\tfrac{4}{3}F_\text{model}$ .
$\text{MFU} = \tau F_\text{model} / (n_\text{gpu} P)$ with the dense peak $P$ ; HFU also counts recomputation. MFU is the comparable number because it does not reward repeated work.
Training time $= F_\text{model} D / (n_\text{gpu} P \cdot \text{MFU}) = D/\tau$ ; GPU-hours $= n_\text{gpu} \times$ time, nearly independent of $n_\text{gpu}$ at a fixed MFU.
A step splits into compute, exposed communication and idle time; a timeline shows which part is large.
Launch-bound criterion: the GPU is busy at most $t_k / \max(t_k, t_\text{cpu})$ of the time; a host-device sync (.item(), .cpu(), print(tensor)) empties the queue.
Overhead-bound work is fixed by fewer launches (CUDA graphs, fusion, larger batch), memory-bound work by fewer bytes (fusion), compute-bound work only by faster arithmetic or fewer FLOPs.
torch.compile: TorchDynamo captures graphs with guards (graph breaks at untraceable code, recompiles when a guard fails), AOTAutograd builds the backward, TorchInductor fuses ops into generated Triton (GPU) or C++ (CPU) kernels.
Method: estimate, measure steady-state steps, profile, fix the largest gap first, measure again.