LLM System InterviewHao Hoang · First edition, 2026Home

Chapter 9 Making a Single GPU Go Fast

9.1 Reducing Memory Movement

9.1.1 Operator Fusion and Kernel Launch Overhead

Nearly every time I have profiled an LLM training run and found an unexpected bottleneck, it was not in the matmuls but in the surrounding glue: activations, layer norms, residual adds, and masking operations, all running as separate CUDA kernels and each one paying a full round-trip to HBM. The fix is kernel fusion: combining sequential operations into a single kernel so intermediate results never leave the compute unit. This concept appears at the mid level as “what is operator fusion,” and at the senior level as “design a fast custom kernel for this attention variant” or “why is this activation function showing up in your profile?”

Think of the GPU as a factory on an island. HBM is the warehouse on the mainland. Every time a PyTorch operation produces an intermediate result in eager mode, the result gets shipped back to the warehouse, and the next operation has to ship it back over again. If five operations run sequentially without fusion, you have paid five round-trips for data that could have stayed in the factory the whole time.

Kernel fusion eliminates those round-trips. A fused kernel reads the input from HBM once, performs all chained operations entirely in registers and shared memory, then writes the final output to HBM once. For memory-bound operations-those with arithmetic intensity (AI) well below the ridge point-this directly reduces runtime. The HBM bandwidth, not compute throughput, is the bottleneck, so cutting the number of passes through HBM cuts the wall-clock time proportionally.

There is a second, smaller benefit: kernel launch overhead. Each CUDA kernel launch incurs a fixed dispatch cost: the CPU traverses the aten C interface, the CUDA runtime queues the kernel, and the GPU drains its pipeline before starting. This overhead is typically 5-20 μ\mus per launch regardless of tensor size. For small tensors or simple operations, this fixed cost dominates actual compute time. Fusing NN kernels into one eliminates N−1N - 1 launch overheads.

Kernel fusion cuts HBM traffic by keeping intermediate results in registers across the fused operations, and removes per-kernel launch overhead as a secondary benefit.

The memory hierarchy makes the arithmetic unavoidable. Registers are zero-latency for the thread that owns them; L1 and shared memory have about 20 clock cycles latency within an SM; HBM has about 300 clock cycles. An unfused chain of elementwise operations reads from and writes to HBM between every step, paying 300 cycles of stall time for each intermediate. A fused kernel holds every intermediate value in registers and only touches HBM at the start and end.

The alternative-letting each intermediate write back to HBM-has one apparent advantage: simplicity. Py-Torch eager mode does exactly this, and it is fine for large matmuls, which are compute-bound and spend little time on relative memory traffic. But activation functions, residual additions, and normalization layers all have AI well below 1 FLOPs/byte. On those operations, every extra HBM round-trip is paid at the full bandwidth cost for zero additional arithmetic.

Modern JIT compilers automate straightforward fusion. torch.compile identifies sequences of elementwise operations with no cross-thread dependencies, emits a single Triton kernel, and eliminates intermediate allocations. This is production-grade fusion at essentially zero engineering effort.

A Worked Example

The approximate GELU is:

GELU(x)≈0.5⋅x⋅(1+tanh⁡(2/π(x+0.044715x3)))\text{GELU} ( x ) \approx 0 . 5 \cdot x \cdot \bigl( 1 + \tanh \bigl( \sqrt{2 / \pi} ( x + 0 . 044715 x^{3} ) \bigr) \bigr)

Written naively in PyTorch as separate tensor expressions, this decomposes into at least five distinct CUDA kernels: compute x3x^{3}; multiply by 0.044715; add xx; compute tanh⁡\tanh; combine with 0.5⋅x0 . 5 \cdot x. Consider the FFN hidden states of a 7 B model at batch size 512 and hidden dimension 11 008, stored in BF16. Each intermediate tensor is 512×11,008×2=11.3512 \times 11 , 008 \times 2 = 11 . 3 MB. Running five kernels means at minimum 5×2×11.3=1135 \times 2 \times 11 . 3 = 113 MB of extra HBM traffic-one read plus one write per intermediate-on top of the 22.6 MB actually needed.

Empirically, on an A100 at this scale:

  • Naive PyTorch (five separate kernels): 8.1 ms

  • PyTorch built-in nn.functional.gelu (fused CUDA kernel): 1.1 ms

  • Hand-written CUDA kernel (all ops in one kernel): 1.8 ms

  • Triton kernel (block-level, vectorized): 1.85 ms

  • torch.compile applied to the naive code: 1.47 ms

The built-in fused kernel is 7.4×\times faster than the naive version. torch.compile on the naive expression reaches 1.47 ms with no kernel code written, capturing 78% of the maximum speedup. The hand-written CUDA and Triton versions are slightly slower than the cuDNN-tuned built-in because CUDA library implementations carry additional hardware-specific tuning.

What You Actually Decide in Practice

  • Reach for torch.compile first. For standard elementwise chains-activations, bias adds, scaling, residuals- torch.compile will auto-fuse them and close most of the gap. Confirm it is actually fusing by checking the profiler for a single kernel launch rather than multiple vectorized_elementwise_kernel calls. Only move to hand-written Triton when the compiler misses a fusion you expect.

  • Fuse activations into adjacent matmul epilogues. The output of a tiled matmul is already in SRAM; applying a bias add or activation in the same kernel avoids writing the pre-activation to HBM. Cutlass and cuBLAS wrappers expose epilogue fusion for this purpose. In Triton, you can add these steps directly at the end of the matmul loop.

  • Kernel launch overhead dominates at small batch sizes. At the batch sizes used in speculative decoding (1-4 sequences), a chain of ten elementwise ops can spend more time in kernel launch overhead than in actual compute. At these scales, fusing into one kernel is especially important.

  • Do not fuse across operations that change memory layout. If fusion forces the matmul output into a non-power-of-2 leading dimension, tile alignment breaks and the resulting performance loss can exceed the fusion gain. Fuse activations; be cautious with dimension-reshaping ops that affect downstream matmul tile selection.

  • Always profile before and after. The profiler trace for the unfused GELU shows five separate vectorized_elementwise_kernel launches. For the fused version you see exactly one. If you are not checking the trace, you do not know whether the compiler has fused your ops or left them separate.

How This Shows Up in Interviews

  • “We have a custom activation function written as a chain of PyTorch ops. It’s showing up in the profiler. What do you do?” Strong answer: confirm in the trace that the activation is running as multiple separate kernels. Apply torch.compile first-it auto-fuses simple chains into a single Triton kernel. If the op requires custom shared memory layout or mask logic, write a Triton kernel. Profile after each change to confirm the launch count dropped.

  • “Why does the unfused GELU run at 8 ms while the built-in runs at 1.1 ms on the same hardware?” Strong answer: the unfused version has five HBM round-trips for intermediate results, each paying full bandwidth cost for zero additional compute. Fusion keeps all intermediates in registers, reducing HBM traffic from roughly ten passes (five reads, five writes) to two (one read, one write).

  • Senior follow-up: “Under what conditions does kernel fusion not help, or even hurt?” Senior answer: if the bottleneck is compute-bound (AI above the ridge point), fusing memory-bound ops around it does not change the binding constraint-you are waiting on compute, not bandwidth, so fewer HBM accesses save nothing. Fusion can also hurt if it increases register pressure within a warp to the point of register spilling to local memory, adding latency that was not present before. Always validate with a profiler; a rule-of-thumb answer about AI first, then empirical confirmation.

Key Takeaways

  • Each unfused HBM round-trip costs ∼\sim300 clock cycles per element; fusion collapses NN elementwise ops to 1 read ++ 1 write regardless of NN.

  • On an A100, fusing a five-op GELU chain yields a 7.4×\times speedup; torch.compile captures roughly 80% of that gain automatically.

  • Kernel launch overhead (∼\sim5-20 μ\mus per launch) is a secondary effect that matters most at small batch sizes.

  • torch.compile is the default tool for fusion; Triton is warranted when the compiler does not auto-fuse or when the op requires custom shared-memory logic.

  • Profile the CUDA kernel trace-not just wall-clock time-to confirm fusion is actually happening.

9.1.2 Recomputation vs. Storage Trade-Offs

Activation checkpointing is well known as a memory-saving technique: you discard most forward-pass activations and recompute them during the backward pass to stay within a GPU’s memory budget. What is less well known, and what comes up frequently in senior kernel-design interviews, is that recomputation can also make execution faster-not just cheaper in memory. The logic is a direct consequence of the roofline model: if a kernel is memory-bound, you have idle compute cycles burning while the kernel waits for HBM. Spending those idle cycles to recompute intermediate values instead of loading them from HBM is free in terms of wall-clock time, and it removes the HBM read that was the bottleneck. FlashAttention’s backward pass is the canonical production example of this technique applied at scale.

The standard backward pass stores every forward-pass intermediate in HBM so gradient computation can read them back. For an operation f:x→→yf : x \rightarrow \rightarrow y with intermediates s1,...,sks_{1} , . . . , s_{k}, this means:

  • Standard: write s1,...,sks_{1} , . . . , s_{k} to HBM during forward; read them back during backward.

  • Recompute: discard s1,...,sks_{1} , . . . , s_{k} after forward; recompute them from xx during backward.

The recompute approach eliminates kk HBM reads at the cost of kk extra arithmetic passes. That exchange is profitable whenever the kernel is memory-bound, because the compute units were already idle waiting for HBM-the recomputed arithmetic consumes cycles that were wasted otherwise.

Concretely, the decision boundary is the roofline ridge point (approximately 295 FLOPs/byte on H100 BF16). Below that arithmetic intensity, compute is excess capacity. Above it, you are compute-bound, and adding recomputation extends the runtime by consuming cycles that were not free.

The same principle, applied at large scale, underlies FlashAttention. The attention score matrix S=S =

QKT/dQ K^{T} / \sqrt{d} has shape T×TT \times T in FP32. Storing and reloading it during the backward pass costs O(T2)O ( T^{2} ) HBM traffic-quadratic in sequence length. FlashAttention instead recomputes SS tile by tile during the backward pass using QQ and KK that are already being read for gradient computation, keeping the O(T2)O ( T^{2} ) materialization entirely out of HBM.

Recomputation turns excess compute into saved HBM bandwidth: for memory-bound kernels, arithmetic that replaces an HBM read is free, because those compute cycles were idle anyway.

The roofline model quantifies the exchange rate. On an H100 at BF16, for every 295 FLOPs you can compute, you are entitled to transfer 1 byte from HBM. Below the ridge point, the FLOPs side of the ratio is underutilized: you are getting fewer than 295 FLOPs per byte because compute finishes faster than the memory can deliver the next byte.

When you recompute a sigmoid or a running maximum rather than loading it from HBM, you are spending FLOPs from the underutilized side to avoid spending bytes from the overutilized side. The exchange is one-sided in your favor for any memory-bound kernel. The more memory-bound the kernel-the lower its AI-the more recomputation you can afford before adding any latency.

This distinguishes recomputation-for-speed from the more familiar gradient checkpointing. Checkpointing is applied to compute-bound operations (matmuls, full attention blocks) to reclaim peak memory at the cost of paying the compute again. That cost is real because those operations are not memory-bound. Recomputation-for-speed is applied to memory-bound operations where the compute cost is zero against the roofline.

A Worked Example

Consider three stacked sigmoid activations: s1=σ(x)s_{1} = \sigma ( x ), s2=σ(s1)s_{2} = \sigma ( s_{1} ), out =σ(s2)= \sigma ( s_{2} ). All values are BF16 (2 bytes per element). Count the HBM operations in each approach.

Standard approach (store intermediates):

  • Forward: read xx (1 read), compute and write s1s_{1}, s2s_{2}, out (3 writes).

  • Backward: read s1s_{1}, s2s_{2}, doutd_{\text{out}} (3 reads), write dxd x (1 write).

  • Total: 8 HBM operations.

Recomputation approach (discard intermediates):

  • Forward: read xx (1 read), write out only-discard s1s_{1}, s2s_{2} (1 write).

  • Backward: read doutd_{\text{out}} and xx (2 reads), recompute s1s_{1}, s2s_{2}, out from xx entirely in SM registers (0 HBM), write dxd x (1 write).

  • Total: 5 HBM operations.

That is 5/8 = 62.5% of the original HBM traffic. The price is three sigmoid evaluations during the backward pass-arithmetic that is free on a memory-bound kernel. For a hidden dimension of 11 008 at batch 512 in BF16 (22.6 MB per tensor), the three saved reads total 3×22.6=67.83 \times 22 . 6 = 67 . 8 MB per backward call through this layer. At H100 HBM3 bandwidth of 3.35 TB/s, that is roughly 20 μ\mus of HBM stall time eliminated per backward call.

The same accounting applied to the FlashAttention backward pass: the T×TT \times T attention matrix in FP32 at T=4096T = 4096 and batch 1 is 40962×4=644096^{2} \times 4 = 64 MB. Without recomputation, loading this matrix takes 64/3350≈1964 / 3350 \approx 19 ms per backward attention call. FlashAttention recomputes it for the cost of a second pass over the QQ and KK tiles-which are already being read for the gradient computation-making the recomputed attention effectively free.

What You Actually Decide in Practice

  • Recomputation for speed requires confirming memory-bound regime first. Profile the kernel with Nsight Systems or torch.profiler. If the kernel’s GPU time is dominated by memory stalls (visible as low SM utilization and high HBM read bandwidth), it is a candidate. If SM utilization is near 100% and CUDA cores are the bottleneck, adding recomputation will extend the runtime linearly.

  • Apply to any activation with AI <1< 1 FLOPs/byte. Sigmoid, ReLU, GELU, SiLU, layer norm statistics (mean and variance), and softmax exponentials all fall into this category. Dropping their intermediate values and recomputing during backward is almost always a net win on H100 or A100 hardware. This is the kernel-level form of the technique; it is separate from the layer-level gradient checkpointing decision.

  • Do not conflate with gradient checkpointing. Checkpointing (e.g., torch.utils.checkpoint) is applied at the boundary of compute-bound layers to reduce peak memory at a throughput cost. Recomputation-for-speed is applied inside memory-bound kernels at no throughput cost. The two levers operate at different granularities and are both usable simultaneously.

  • Fuse the recomputed path. Recomputing s1s_{1}, s2s_{2}, out during backward only avoids HBM traffic if those recomputations happen inside the same kernel as the gradient accumulation. Launching separate kernels for the recomputed intermediates adds HBM traffic back and defeats the purpose entirely. Recomputation and fusion must be applied together.

  • The break-even point depends on the recompute cost. For cheap operations (sigmoid, ReLU), a single sigmoid is ∼\sim3 FLOPs and costs nothing against the 295 FLOPs/byte ridge. For expensive operations (full softmax over a long sequence, layer norm with large hidden dimension), the recompute cost can itself push the kernel toward compute-bound. Measure before assuming the trade is free.

How This Shows Up in Interviews

  • “FlashAttention doesn’t store the full attention weight matrix. How does the backward pass work?” Strong answer: the backward pass recomputes the attention scores S=QKT/dS = Q K^{T} / \sqrt{d} tile by tile in SRAM using stored QQ, KK, the output OO, and gradient dOd O. The attention score matrix has AI ≈1\approx 1 and is therefore memory-bound; recomputing it costs FLOPs from the excess supply and eliminates the O(T2)O ( T^{2} ) HBM reads that would otherwise dominate backward-pass time and memory.

  • “When should you use activation checkpointing vs. activation recomputation?” Strong answer: checkpointing applies at the layer boundary (matmuls, full transformer blocks) to reduce peak memory when the model does not fit in GPU RAM; the runtime cost is real because those layers are compute-bound. Activation recomputation applies inside memory-bound kernels (elementwise activations, normalization intermediates) to speed up the backward pass; the runtime cost is zero because compute was idle. Confirm the target kernel is memory-bound before applying recomputation for speed.

  • Senior follow-up: “Describe a case where recomputation-for-speed makes the kernel slower.” Senior answer: if the operation is compute-bound-AI above the ridge point-the compute units are fully saturated. Recomputing any intermediate requires FLOPs that must wait in the compute pipeline, directly extending latency. A concrete example: applying activation recomputation inside a large tiled matmul (AI ≫295\gg 295) adds FLOPs to an already saturated pipeline and increases kernel time in proportion. The check is always: measure AI first, then decide whether recompute FLOPs are free or costly.

Key Takeaways

  • Recomputation for speed trades excess compute for scarce HBM bandwidth; for memory-bound kernels (AI below the roofline ridge), the arithmetic is free because compute units were already idle.

  • The triple-sigmoid example reduces HBM ops from 8 to 5 (a 37.5% reduction) by recomputing intermediates in SM registers during the backward pass from the already-loaded input.

  • FlashAttention’s backward pass applies this pattern to eliminate O(T2)O ( T^{2} ) HBM traffic; the T×TT \times T attention matrix is recomputed tile by tile at the cost of a second pass over QQ and KK tiles already in SRAM.

  • Recomputation-for-speed and gradient checkpointing operate at different granularities and serve different goals; using both simultaneously is common and correct.

  • Recomputation only works if the recomputed path is fused into the backward kernel; a separate kernel launch for recomputed intermediates adds HBM traffic back, negating the benefit.

9.2 Exploiting Memory Layout

9.2.1 Memory Coalescing and Burst Sections

Implement a naive matrix transpose in CUDA and watch effective HBM bandwidth drop to 1-2% of peak. The computation is trivial - one read, one write per element - but the access pattern is column-major on a row-major matrix, and that single design choice costs two orders of magnitude of bandwidth. Interviewers at GPU-focused companies will ask why a “simple” kernel runs at a fraction of peak bandwidth, and the answer lives entirely in how your warp’s 32 threads hit DRAM. Memory coalescing is the most impactful low-level GPU optimization, and understanding it means understanding burst-mode DRAM at the hardware level. This topic appears at the systems engineer and ML infrastructure engineer level.

When a warp of 32 threads issues memory reads, the hardware does not process them sequentially one by one. DRAM is physically organized into burst sections: contiguous address blocks, typically 128 bytes, that the DRAM array delivers together in a single sensing operation. The physics: accessing DRAM requires activating a row of memory cells (a sense amplifier step), which is the expensive operation. Once a row is open, reading any subset of that row’s bits is nearly free. The hardware exposes this as: request any single address, receive the full 128-byte cache line.

Memory coalescing exploits this directly. When 32 threads in a warp all read addresses that fall within a small number of burst sections, the memory controller merges them into that many transactions rather than 32 separate ones:

  • Fully coalesced: Thread kk reads address base +k×2+ k \times 2 (consecutive bfloat16 elements from a 128-bytealigned base). All 32 reads map to one 64-byte burst section. One HBM transaction returns 64 bytes; all 64 useful bytes are consumed. Utilization: 100%100 \%.

  • Fully non-coalesced: Thread kk reads address base +k×S+ k \times S where SS is the row stride in bytes (e.g., 4096×2=81924096 \times 2 = 8192 bytes for a 4096-wide matrix). Each thread’s address is 8192 bytes from the next: 128 separate burst sections. Hardware issues 32 transactions. Bytes fetched: 32×128=409632 \times 128 = 4096 bytes. Useful bytes: 32×2=6432 \times 2 = 64 bytes. Utilization: 1.6%1 . 6 \%.

The utilization ratio equals useful_bytes / fetched_bytes. For stride-SS access, utilization =2/min⁡(128,S)= 2 / \min ( 128 , S ) for bfloat16. At stride = 4096, that is 2/128≈1.6%2 / 128 \approx 1 . 6 \% - a 64×\times waste.

Coalesced access is not a software trick: it exploits a fixed DRAM hardware property (burst mode) that already exists; non-coalesced access wastes the free bytes that burst mode delivers, cutting effective bandwidth by up to the number of warp threads.

The reason burst mode exists is that the dominant cost in DRAM access is activating a row of cells (the sense amplifier step), not transferring the bits. Once the row is energized, reading 128 bytes costs almost the same as reading 1 byte. The hardware exposes this by delivering a full cache line on every access - whether you asked for 1 byte or 128.

Coalescing turns this into a benefit: 32 threads’ reads for 64 bytes of useful data require one row activation instead of 32. Non-coalescing turns it into a penalty: 32 row activations for 64 bytes of useful data, fetching 4 KB and discarding 63 of every 64 bytes fetched.

CPUs partially compensate for non-coalesced access via hardware stride prefetchers that detect strided patterns and prefetch cache lines proactively. GPUs have no equivalent adaptive prefetcher. The programmer must arrange accesses correctly, or the performance cliff materializes silently. Profiling tools like Nsight Compute expose the coalescing ratio via l1tex__t_sectors / l1tex__t_requests; a ratio near 32 means fully non-coalesced access.

A Worked Example

A∈R4096×4096A \in \mathbb{R}^{4096 \times 4096} in bfloat16, stored row-major. One warp (32 threads) loads one row’s first 32 elements vs. one column’s first 32 elements.

Row access (coalesced). Thread kk reads A[0,k]A [ 0 , k ]. Address: base +k×2+ k \times 2 bytes.

Address span=32×2=64bytes=one HBM cache line.\text{Address span} = 32 \times 2 = 64 \text{bytes} = \text{one HBM cache line} .

Transactions: 1. Bytes fetched: 128 (one cache line). Useful: 64. Utilization: 50% (slightly under if stride≠\neq cache line size; exactly 1 transaction in practice with 128-byte alignment).

Column access (non-coalesced). Thread kk reads A[k,0]A [ k , 0 ]. Address: base+k×4096×2=+ k \times 4096 \times 2 = base+k×8192+ k \times 8192 bytes.

Distance between consecutive threads =8192= 8192 bytes =64×(= 64 \times (cache line size).) .

Each thread touches a different cache line 64 lines apart. Transactions: 32. Bytes fetched: 32×128=409632 \times 128 = 4096. Useful: 32×2=6432 \times 2 = 64. Utilization: 64/4096=1.56%64 / 4096 = 1 . 56 \%.

At 2 TB/s HBM bandwidth on an A100:

Effective row bandwidth ≈2\approx 2 TB/s ×0.5=1\times 0 . 5 = 1 TB/s (full-rate coalesced)

Effective column bandwidth≈2TB/s×0.0156≈31GB/s.\text{Effective column bandwidth} \approx 2 \text{TB/s} \times 0 . 0156 \approx 31 \text{GB/s} .

A kernel doing nothing but column reads of this matrix would run at 31 GB/s instead of 2 TB/s - a 64×\times degradation. For a bfloat16 element-wise operation with 1 FLOP per element, this column-read version would be so memory-starved that even a single arithmetic instruction per cycle cannot fill the gap.

The fix. Stage through shared memory:

  1. Load a T×TT \times T tile from HBM into shared memory in row-major order (coalesced - each block of 32 threads reads 32 consecutive elements from the same row).

  2. Access the tile in shared memory in column order (fast - shared memory is banked, not burst-mode; column access incurs bank conflicts but not the 64×\times bandwidth penalty of HBM column access).

What You Actually Decide in Practice

  • Always map thread index kk to the innermost (fastest-varying) tensor dimension. For a row-major 2D tensor, thread kk reads A[row][threadIdx.x + blockIdx.x * blockDim.x]. Never iterate threads over rows when data is row-major.

  • Pad all matrix dimensions to multiples of 64 for bfloat16, 32 for float32. A row width of 4097 bfloat16 elements (=8194= 8194 bytes) does not align the start of row 2 to a 128-byte boundary. Threads loading row 2 will straddle two cache lines for the first element, doubling HBM transactions for that access. Padding vocabulary size from 50,257 to 50,304 (nearest multiple of 64) delivered a 25% throughput improvement in nanoGPT - entirely a coalescing alignment fix.

  • When a transposed access is unavoidable, tile through shared memory. The canonical case: the KK matrix in attention requires column access (transposed). Do not read columns directly from HBM. Load a T×TT \times T tile of KK into shared memory coalesced (read rows), then transpose within shared memory (fast), then proceed. FlashAttention does exactly this.

  • Profile with Nsight Compute’s sector/request ratio before optimizing. The metric l1tex__t_sectors_pipe_lsu_mem_global_op_ld.sum divided by l1tex__t_requests_pipe_lsu_mem_global_op_ld.sum reports the average number of HBM sectors per warp load request. For fully coalesced access: ratio ≈1\approx 1-22. For fully non-coalesced 32-thread access: ratio ≈32\approx 32. A ratio >4> 4 is a signal to investigate the access pattern before any other optimization.

  • Embedding lookups are structurally non-coalesced; mitigate at the batch level. Token embedding tables are indexed by token ID (random integers up to vocab size). Each lookup hits a different row, with no spatial locality. Effective bandwidth is limited. Two mitigations: (1) sort requests by token ID within each microbatch so nearby GPU threads access nearby rows; (2) use INT8 or FP8 embeddings to double or quadruple the useful bytes per fetched cache line.

How This Shows Up in Interviews

  • “Why does column-major access of a row-major matrix perform so much worse on a GPU?” Each thread in a warp accesses a different row; their addresses are stride-(row_width) apart. Each address falls in a different 128-byte burst section. The hardware must issue one HBM transaction per thread, fetching 128 bytes but using 2 - up to 64×\times more bandwidth consumed for the same useful data.

  • “A kernel loads a 2D tensor and runs at 3% of peak HBM bandwidth. How do you diagnose and fix it?” Diagnose: Nsight Compute sector/request ratio. If it is near 32, the access is non-coalesced. Identify which dimension threads stride over: if threads stride across rows of a row-major tensor (column access), access is non-coalesced. Fix: either stage through shared memory (load tile coalesced, use within shared memory), or restructure the algorithm to access the tensor in row order.

  • Senior follow-up: “An embedding lookup for a 50,000-vocabulary model runs at 5% of peak HBM bandwidth regardless of batch size. Coalescing is structurally impossible because token indices are random. Name three distinct ways to improve effective bandwidth.” (1) Sort tokens by index within each batch: consecutive GPU threads then access nearby rows, partially amortizing burst section waste. (2) Quantize the embedding table to INT8 or INT4: halves or quarters the bytes per row, so each 128- byte burst section delivers 2-4×\times more embedding elements; effective bandwidth in elements/sec doubles even if bytes/sec is the same. (3) Cache hot embeddings in L2 by structuring the training loop to repeat frequent-token batches before evicting: for a Zipfian token distribution where the top 1% of tokens appear 50% of the time, a 40 MB L2 can absorb those 500 frequently-accessed embedding rows (assuming 50K vocab ×\times 4096 ×\times 2 bytes / 50K ≈\approx 8 KB per embedding; top 0.1% ≈\approx 50 rows ×\times 8 KB =400= 400 KB easily fits).

Key Takeaways

  • DRAM burst mode delivers a full 128-byte cache line on any read; coalesced access (32 warp threads reading 64 consecutive bytes) achieves one HBM transaction; non-coalesced access (32 threads reading 32 scattered cache lines) requires 32 transactions, fetching 4 KB to deliver 64 bytes.

  • Effective bandwidth degradation = useful bytes / fetched bytes; stride-SS access in bfloat16 gives utilization 2/min⁡(128,S)2 / \min ( 128 , S ); at stride = 4096 elements, effective bandwidth is ∼1.6%\sim 1 . 6 \% of peak.

  • Column-major access of a row-major matrix is the canonical non-coalescing trap; fix by staging through shared memory (load coalesced, compute in shared memory), not by reading columns directly from HBM.

  • Pad all matrix dimensions to multiples of 64 (bfloat16) or 32 (float32) to maintain 128-byte alignment across rows; the nanoGPT vocabulary padding example delivered 25% speedup from alignment alone.

  • Profile with Nsight Compute’s sector/request ratio; values >4> 4 indicate non-coalesced access worth investigating.

9.2.2 Tiling for Matrix Multiplies

Add one column to your matrix and watch GEMM throughput drop 30%. The FLOP count barely changed, yet the GPU is suddenly running two execution waves instead of one, with the second wave consuming 90% of its time on 12 tiles while 96 of 108 SMs sit idle. This is the “wavy” matrix multiply performance plot that puzzles every GPU engineer who first encounters it, and tiling is the concept that explains every feature on that plot. Understanding tiling is what lets you justify vocabulary size choices, explain FlashAttention’s performance model, and design kernels that sustain high throughput under interview conditions. This topic appears at the ML infrastructure engineer and systems engineer level.

The naive matrix multiply algorithm for C=A×BC = A \times B with A,B∈RN×NA , B \in \mathbb{R}^{N \times N} reads each element of AA and BB exactly NN times from HBM: computing output element C[i,j]C [ i , j ] requires the full row A[i,:]A [ i , : ] and column B[:,j]B [ : , j ], so A[i,k]A [ i , k ] is read once per output column, i.e., NN times total. With N2N^{2} elements in AA, that is N3N^{3} total HBM element reads - entirely memory-bound for any large NN.

Tiling eliminates the redundancy. Partition AA, BB, and CC into square T×TT \times T sub-matrices. For the output tile C[i0:i0+T,j0:j0+T]C [ i_{0} : i_{0} + T , j_{0} : j_{0} + T ]:

C[i0:i0+T,j0:j0+T]=∑k=0N/T−1A[i0:i0+T,kT:(k+1)T]×B[kT:(k+1)T,j0:j0+T].C [ i_{0} : i_{0} + T , j_{0} : j_{0} + T ] = \sum_{k = 0}^{N / T - 1} A [ i_{0} : i_{0} + T , k T : ( k + 1 ) T ] \times B [ k T : ( k + 1 ) T , j_{0} : j_{0} + T ] .

The algorithm loads one T×TT \times T tile of AA and one T×TT \times T tile of BB into shared memory, computes the partial product entirely in registers/tensor cores, then advances to the next kk and repeats. Each element of AA is now read from HBM only N/TN / T times (once per tile iteration along kk), not NN times. HBM reads per matrix: from N3N^{3} to N3/TN^{3} / T. Improvement: a factor of TT.

The algorithm in pseudocode:

for i_tile in range(N // T):        # output tile row
  for j_tile in range(N // T):      # output tile column
    C_acc = zeros(T, T)             # accumulator in registers
    for k_tile in range(N // T):    # reduction tiles
      # Load from HBM -> shared memory (coalesced)
      A_smem[0:T, 0:T] = A[i_tile*T:(i_tile+1)*T, k_tile*T:(k_tile+1)*T]
      B_smem[0:T, 0:T] = B[k_tile*T:(k_tile+1)*T, j_tile*T:(j_tile+1)*T]
      __syncthreads()
      C_acc += A_smem @ B_smem      # tensor core MMA in shared memory
      __syncthreads()
    C[i_tile*T:(i_tile+1)*T, j_tile*T:(j_tile+1)*T] = C_acc  # write HBM

Each output tile is one thread block, assigned to one SM. The total number of blocks is ⌈N/T⌉2\lceil N / T \rceil^{2}. With 108 SMs on an A100, blocks execute in waves of 108. If the total block count is not a multiple of 108, the last wave is partially populated: most SMs idle while the remaining blocks finish. This is wave quantization.

Tiling reduces HBM reads by a factor of TT by staging the inner reduction loop in shared memory; wave quantization causes throughput cliffs whenever the tile count crosses a multiple of the SM count.

The fundamental insight behind tiling is that matrix multiply has reuse: each element of AA participates in NN dot products (one per output column), and each element of BB participates in NN dot products (one per output row). If you can arrange to perform all NN of those dot products while the element is in shared memory, you only pay the HBM latency once. Tiling is the mechanism that groups those NN uses spatially: a single T×TT \times T tile of AA is used in all TT dot products for the corresponding output tile row before being evicted.

The arithmetic intensity increases proportionally to TT. Naive: 2N32 N^{3} FLOPs, 2N32 N^{3} element HBM reads ⇒\Rightarrow 1 FLOP/element read. Tiled (TT): 2N32 N^{3} FLOPs, 2N3/T2 N^{3} / T element HBM reads ⇒T\Rightarrow T FLOPs/element read. For T=128T = 128, arithmetic intensity is 128 FLOP/element. At 2 bytes/element (bfloat16), that is 64 FLOP/byte - well above the A100’s roofline threshold of 156 FLOP/byte for bfloat16, but much closer to the compute-bound regime than the naive algorithm.

Tiling also subsumes coalescing: when a block loads a tile from HBM into shared memory, the 32 threads in a warp each read consecutive elements from the same row (row-major, one row of the tile per warp). This is perfectly coalesced. Subsequent intra-tile computation uses shared memory, where access order does not matter for bandwidth. Tiling thus handles both the HBM redundancy problem and the coalescing problem simultaneously.

A Worked Example

A,B∈R4096×4096A , B \in \mathbb{R}^{4096 \times 4096}, bfloat16, on an A100 (108 SMs). Tile size T=128T = 128.

HBM traffic reduction. Each element of AA is read N/T=4096/128=32N / T = 4096 / 128 = 32 times from HBM (once per tile iteration along kk). Total HBM reads:

2×40962×32×2bytes=2×16,777,216×64≈2.1GB.2 \times 4096^{2} \times 32 \times 2 \text{bytes} = 2 \times 16 , 777 , 216 \times 64 \approx 2 . 1 \text{GB} .

At 2 TB/s HBM bandwidth: 2.12 . 1 GB/(2/ ( 2 TB/s)≈1.0) \approx 1 . 0 ms.

Compare to naive (NN reads per element): 2×40962×4096×2≈2742 \times 4096^{2} \times 4096 \times 2 \approx 274 TB, which at 2 TB/s would take ∼\sim137,000 s. Even with a 40 MB L2 cache providing partial reuse, naive is the wrong framing for large matrices.

Shared memory budget. Two tiles in flight simultaneously:

2×T2×2bytes=2×1282×2=65,536bytes=64KB.2 \times T^{2} \times 2 \text{bytes} = 2 \times 128^{2} \times 2 = 65 , 536 \text{bytes} = 64 \text{KB} .

A100 configurable shared memory: up to 164 KB per SM. With 64 KB per block, two concurrent blocks fit per SM (2×64=1282 \times 64 = 128 KB <164< 164 KB): occupancy is 2 blocks ×\times 256 threads = 512 threads per SM, or 25% of the 2048-thread maximum. Acceptable for a compute-bound kernel.

Tile count and wave analysis.

Output tiles=⌈4096128⌉2=322=1024.\text{Output tiles} = \biggl\lceil \frac{4096}{128} \biggr\rceil^{2} = 32^{2} = 1024 .
Waves=⌈1024108⌉=10.Last wave tiles=1024−9×108=52.\text{Waves} = \biggl\lceil \frac{1024}{108} \biggr\rceil = 10 . \text{Last wave tiles} = 1024 - 9 \times 108 = 52 .

Last-wave SM utilization: 52/108=48%52 / 108 = 48 \%. Execution time ≈10×twave\approx 10 \times t_{\text{wave}} instead of ideal 9.48×twave9 . 48 \times t_{\text{wave}}: overhead ≈5%\approx 5 \%. Acceptable.

Pathological case. N=1792N = 1792 with tile shape 256×128256 \times 128. Output tiles: (1792/256)×(1792/128)=7×14=98( 1792 / 256 ) \times ( 1792 / 128 ) = 7 \times 14 = 98. With 108 SMs: 98<10898 < 108, so all tiles execute in one wave. GPU runs at 100% SM utilization. Now N=1796N = 1796: ⌈1796/256⌉×⌈1796/128⌉=8×15=120\lceil 1796 / 256 \rceil \times \lceil 1796 / 128 \rceil = 8 \times 15 = 120 tiles. First wave: 108 tiles (100% utilization). Second wave: 12 tiles (12/108=11%12 / 108 = 11 \% utilization). Total: 2 waves at cost ≈2×twave\approx 2 \times t_{\text{wave}} instead of ideal 1.11×twave1 . 11 \times t_{\text{wave}}. Throughput drops ≈45%\approx 45 \% from a 4-element increase in NN.

What You Actually Decide in Practice

  • Use tile size T∈{64,128,256}T \in \{ 64 , 128 , 256 \} on A100/H100; 128 is the default safe choice. Tensor core MMA units operate on 16×1616 \times 16 fragments; 128 is 8×168 \times 16, mapping cleanly. Smaller tiles reduce shared memory pressure but also reduce the arithmetic intensity improvement (less reuse per tile load). Larger tiles (T=256T = 256) improve reuse but consume more shared memory, limiting occupancy to 1 block per SM.

  • Pad matrix dimensions to the nearest multiple of TT (minimum multiple of 64 for bfloat16). Fractional boundary tiles leave threads with no useful work. Padding the nanoGPT vocabulary size from 50,257 to 50,304 (nearest multiple of 64) eliminated a fractional last column-tile and delivered a 25% speedup. Compute the padded size: ⌈D/T⌉×T\lceil D / T \rceil \times T, then run the kernel over the padded region and discard padded outputs.

  • Target tile counts that avoid wave quantization cliffs. For an A100 (108 SMs), target ⌈N/T⌉2≡0(mod108)\lceil N / T \rceil^{2} \equiv 0 ( \mathrm{mod} 108 ) or choose NN such that the last wave has >80> 80 tiles. For T=128T = 128: tile counts that are multiples of 108 with 128-element dimensions include N=1536N = 1536 (12 tiles ×\times 9 tiles =108= 108, one wave), N=4096N = 4096 (32=21024=9.48×108=^{2} 1024 = 9 . 48 \times 108, last wave 48% utilized), or N=4224N = 4224 (33=21089=10.1×108=^{2} 1089 = 10 . 1 \times 108, last wave 9% utilized - bad). Check before committing to a model width.

  • Verify shared memory usage fits within the per-SM budget. A100 maximum shared memory: 164 KB (with explicit configuration). Safe rule: 2×T2×2 \times T^{2} \times dtype_bytes <96< 96 KB per block to allow 2 concurrent blocks per SM. For T=128T = 128 and bfloat16: 2×1282×2=642 \times 128^{2} \times 2 = 64 KB <96< 96 KB - safe. For T=192T = 192 and bfloat16: 2×1922×2=1472 \times 192^{2} \times 2 = 147 KB >96> 96 KB - forces single block per SM, halving occupancy.

  • Do not try to manually tune tile sizes: use cuBLAS or Cutlass for standard GEMMs. Both libraries auto-select tile sizes based on matrix shape, SM count, and data type. Manual tuning matters only for custom operations (e.g., fused attention, custom activations) where a library kernel does not exist. When writing a custom kernel, always benchmark tile sizes {64,128,256}\{ 64 , 128 , 256 \} using nvprof or Nsight Compute before committing to a size.

How This Shows Up in Interviews

  • “Explain why GEMM throughput has a wavy pattern as matrix size increases - why does adding one column sometimes drop throughput by 30%?” Two effects. First: a dimension that is not a multiple of TT leaves the last tile mostly padding; those SMs do little useful work (discretization waste). Second: total tile count crossing a multiple of SM count triggers wave quantization - an extra nearly-empty wave runs at ≪100%\ll 100 \% SM utilization, consuming near-full wave time for a small amount of work. Both effects can occur simultaneously when the dimension changes by one.

  • “How does tiling improve matrix multiply performance on a GPU?” Tiling stages the inner reduction loop in shared memory instead of HBM. Each element of AA and BB is read from HBM N/TN / T times instead of NN times - a factor-of-TT reduction in HBM traffic. For T=128T = 128 and N=4096N = 4096, each element is read 32 times from HBM rather than 4096 times, converting an HBM-dominated workload into a near-compute-bound one.

  • Senior follow-up: “A model uses embedding dimension d=4097d = 4097. Profiling shows GEMM throughput dropped 25% vs. d=4096d = 4096. Identify the two distinct performance hazards and describe a fix for each.” Hazard 1: 40974097 is not a multiple of T=128T = 128; the last column-tile is 1/1281 / 128-width, leaving 127 of 128 threads in that tile doing no useful work. Fix: pad dd to 4224=33×1284224 = 33 \times 128. Hazard 2: after padding, tile count =⌈4224/T⌉2=332=1089= \lceil 4224 / T \rceil^{2} = 33^{2} = 1089; ⌈1089/108⌉=11\lceil 1089 / 108 \rceil = 11 waves with last wave =1089−10×108=9= 1089 - 10 \times 108 = 9 tiles, SM utilization 8%8 \%. Total time ≈11×twave\approx 11 \times t_{\text{wave}} instead of ideal 10.1×twave10 . 1 \times t_{\text{wave}}: a 9% quantization penalty. Fix for hazard 2: find the smallest dd that is (a) ≥4097\geq 4097, (b) a multiple of TT, and (c) yields a tile count divisible by 108 or at least with a last-wave utilization >80%> 80 \%. For T=128T = 128: d=4480=35×128d = 4480 = 35 \times 128 gives 352=1225=11.34×10835^{2} = 1225 = 11 . 34 \times 108, last wave =1225−11×108=37= 1225 - 11 \times 108 = 37 tiles (34%34 \% utilization - still suboptimal). Check d=4608=36×128d = 4608 = 36 \times 128: 362=1296=12×10836^{2} = 1296 = 12 \times 108 - a perfect multiple of 108, one clean wave. Recommend d=4608d = 4608 to the team.

Key Takeaways

  • Tiling reduces HBM reads by a factor of TT (tile size): each element of AA and BB is fetched N/TN / T times instead of NN times; for T=128T = 128, HBM traffic drops 128×128 \times vs. naive and the kernel transitions from memory-bound to near-compute-bound.

  • Total tile count =⌈N/T⌉2= \lceil N / T \rceil^{2} must be a multiple of the SM count (108 on A100, 132 on H100) or the last wave executes with most SMs idle; tile count just above a multiple is the worst case - two waves for nearly one wave of useful work.

  • Shared memory budget: both input tiles must fit simultaneously; 2×T2×2 \times T^{2} \times dtype_bytes <96< 96 KB per block allows two concurrent blocks per SM. For bfloat16 and T=128T = 128: 64 KB - safe.

  • Pad all matrix dimensions to multiples of TT (minimum multiples of 64 for bfloat16): fractional tiles waste SM capacity and break burst-section alignment. The nanoGPT vocabulary padding (50,257 →\rightarrow 50,304) is the canonical production example.

  • For custom kernels, the tile size selection involves three simultaneous constraints: fit within shared memory, coalesce HBM tile loads (multiples of 32 or 64 elements), and achieve favorable tile count vs. SM count ratio. Use cuBLAS/Cutlass for standard GEMMs; invest manual tuning effort only in custom operations.

9.3 Lower Precision and Quantization-Aware Training

9.3.1 FP16 and BF16 in Practice

Cutting numerical precision in half is the highest-leverage single change you can make to LLM training throughput, and it requires zero custom kernel code. The choice between FP16, BF16, and FP32 for each tensor in the training graph is one of the first things an interviewer will probe once you claim to understand efficient training. At the mid-level this surfaces as “why is BF16 faster than FP32”; at the senior level it escalates to “walk me through every place you keep FP32 in a mixed-precision run and explain why.”

Three formats matter for LLM training. FP32 uses 32 bits (1 sign, 8 exponent, 23 mantissa) and supports values up to ≈3.4×1038\approx 3 . 4 \times 10^{38}. FP16 uses 16 bits (1 sign, 5 exponent, 10 mantissa) with a maximum representable value of only 65,504-four orders of magnitude smaller. BF16, the format used by almost every production LLM training run today, also uses 16 bits but allocates 8 bits to the exponent (matching FP32) and only 7 bits to the mantissa. BF16 deliberately trades mantissa precision for dynamic range, keeping the same maximum and minimum values as FP32.

The performance argument is direct. For any memory-bound kernel, wall-clock time equals bytes transferred divided by bandwidth. Halving the bits per element halves the bytes transferred, halving the kernel runtime without changing a single arithmetic operation. A ReLU on nn BF16 elements reads 2n2 n bytes and writes 2n2 n bytes-4n4 n total-at AI =0.25= 0 . 25 FLOPs/byte. The same ReLU in FP32 moves 8n8 n bytes at AI =0.125= 0 . 125 FLOPs/byte. The AI ratio is the same (0.25/0.125=2×0 . 25 / 0 . 125 = 2 \times), but the absolute bytes are halved, so the kernel completes twice as fast.

Matrix multiplications gain a second effect. On H100, BF16 inputs unlock the tensor-core path at ≈989\approx 989 TFLOPs/s. FP32 inputs fall back to general ALUs at ≈67\approx 67 TFLOPs/s-a 14.8×14 . 8 \times peak throughput gap. The standard mixed-precision matmul pattern exploits both: BF16 inputs are fed to the tensor cores; each tile accumulates dot-product partial sums in FP32 accumulator registers; the FP32 result is cast back to BF16 before writing to HBM. This is the default behavior of cuBLAS and torch.autocast when the target dtype is bfloat16.

Halving numerical precision halves HBM traffic for memory-bound ops and unlocks tensor cores for compute-bound ones; BF16 is the right default for weights and activations in LLM training, not FP16.

FP16’s 5-bit exponent limits its maximum representable value to 65,504. Gradient norms in large-scale LLM training regularly exceed this threshold during warmup, around sharp loss changes, and in embedding tables for rare tokens. When an FP16 value overflows to Inf or underflows to zero, the optimizer step either corrupts the weight or silently does nothing. The standard workaround is loss scaling: multiply the scalar loss by a large constant SS before the backward pass so that gradient tensors are scaled up by SS and stay in range, then divide the accumulated gradients by SS before the optimizer step. The machinery works but requires dynamic tuning of SS-too small and underflow still occurs; too large and overflow reappears. Practitioners invest real engineering effort managing this.

BF16 eliminates loss scaling entirely. Because it shares the 8-bit exponent with FP32, its representable range matches FP32 exactly. Gradient values that overflow FP16 are representable in BF16 without any adjustment. The trade-off is lower mantissa precision-7 bits instead of 23 in FP32 or 10 in FP16, giving relative rounding errors of roughly ±0.78%\pm 0 . 78 \% of the value. For weights, activations, and gradient tensors in an LLM forward-backward pass, this rounding is empirically inconsequential. For the accumulation step inside a tile of a matmul, it is not-which is why accumulation stays in FP32 even in a BF16 training run.

A Worked Example

Consider the gate projection and SwiGLU activation in a 7 B model: D=4096D = 4096, F=11,008F = 11 , 008, batch size 512 in training.

SwiGLU activation (memory-bound): The kernel reads two input tensors of size 512×11,008512 \times 11 , 008 and writes one output, for three tensors total.

  • FP32: 3×512×11,008×4=67.73 \times 512 \times 11 , 008 \times 4 = 67 . 7 MB; runtime =67.7/3,350≈20.2μ= 67 . 7 / 3 , 350 \approx 20 . 2 \mus.

  • BF16: 3×512×11,008×2=33.83 \times 512 \times 11 , 008 \times 2 = 33 . 8 MB; runtime ≈10.1μ\approx 10 . 1 \mus.

  • Speedup: 2×2 \times from halved bytes; zero quality loss.

Gate projection matmul (compute-bound at B=512B = 512): FLOPs =2×512×4096×11,008≈46.1= 2 \times 512 \times 4096 \times 11 , 008 \approx 46 . 1 GFLOPs.

  • BF16 tensor cores at 989 TFLOPs/s: 46.1×109/989×1012≈46.6μ46 . 1 \times 10^{9} / 989 \times 10^{12} \approx 46 . 6 \mus theoretical.

  • FP32 general ALUs at 67 TFLOPs/s: 46.1×109/67×1012≈688μ46 . 1 \times 10^{9} / 67 \times 10^{12} \approx 688 \mus theoretical.

  • Switching from FP32 ALUs to BF16 tensor cores: 14.8×14 . 8 \times peak throughput advantage for the dominant operation class.

Single decode (B=1B = 1): The bottleneck is weight-loading, not arithmetic. Gate projection weights: 4096×11,008×2=90.24096 \times 11 , 008 \times 2 = 90 . 2 MB in BF16 vs 180.4180 . 4 MB in FP32.

  • BF16 load time: 90.2/3,350≈26.9μ90 . 2 / 3 , 350 \approx 26 . 9 \mus.

  • FP32 load time: ≈53.8μ\approx 53 . 8 \mus.

  • Summed over all projection layers in a 7 B model, the 2×2 \times weight-load savings directly doubles per-token decode throughput.

What You Actually Decide in Practice

  • Use BF16 for weights, activations, and gradients; never use FP16 for LLM training without loss scaling. Wrap the forward and backward pass with torch.autocast(device_type='cuda', dtype=torch.bfloat16); it automatically dispatches BF16 for matmuls and FP32 accumulation for reductions. On hardware pre-dating BF16 support (Volta, Turing), FP16 with dynamic loss scaling is the only option-but on A100 or H100, BF16 is strictly better and simpler.

  • Keep optimizer state in FP32. Adam’s first moment (mtm_{t}) and second moment (vtv_{t}) are running estimates of the gradient mean and variance. The weight update Δw=−η⋅mt/vt+ϵ\Delta w = - \eta \cdot m_{t} / \sqrt{v_{t} + \epsilon} involves dividing two nearzero BF16 quantities, producing a result that BF16’s ±0.78%\pm 0 . 78 \% relative error rounds to zero at small learning rates or near-converged weights. FP32 master weights and optimizer state are the universal standard; the memory overhead is acceptable because they are not on the HBM bandwidth critical path during forward and backward.

  • Keep layer norm and softmax reductions in FP32 accumulation. Computing the mean and variance of a row for layer normalization requires accurate summation of up to several thousand elements; BF16 accumulation introduces a systematic bias that shifts the normalized output. PyTorch’s built-in nn.LayerNorm upcasts to FP32 internally when inputs are BF16. Do not override this with a manual BF16 implementation unless you have measured that it makes no quality difference on your specific training run.

  • Never store embedding gradients in FP16 on large vocabularies. Embedding gradients for rare tokens are sparse and can be large; FP16 overflows without loss scaling. With BF16 this is not an issue. If an older checkpoint or library forces FP16 embeddings, add per-parameter loss scaling for the embedding table and monitor gradient norms separately.

  • Confirm fusion when switching precision. Lower precision only helps if the kernel using it runs. Profile with torch.profiler before and after adding torch.autocast; confirm that the activation kernels and matmuls show a single BF16 launch each. If torch.compile is not applied, some elementwise chains may still dispatch as FP32 scalar ops regardless of autocast.

How This Shows Up in Interviews

  • “Why is BF16 faster than FP32 for LLM training?” Strong answer: two mechanisms. First, halving bytes per element halves HBM traffic for memory-bound kernels (activations, norms, elementwise ops), directly halving their runtime. Second, BF16 inputs unlock H100 tensor cores at ≈989\approx 989 TFLOPs/s vs general FP32 ALUs at ≈67\approx 67 TFLOPs/s-a 14.8×14 . 8 \times peak advantage on matrix multiplications, which dominate total FLOPs.

  • “Why does production LLM training use BF16 instead of FP16?” Strong answer: FP16’s 5-bit exponent limits its maximum representable value to 65,504; gradient norms in large-scale training regularly exceed this, causing overflow without loss scaling. BF16 matches FP32’s 8-bit exponent and hence its full dynamic range, eliminating overflow at the cost of lower mantissa precision-a trade that is empirically safe for LLM weights, activations, and gradients.

  • Senior follow-up: “Walk me through every place in a BF16 training run where you keep FP32 and why.” Senior answer: (1) Matmul accumulation-tensor-core tiles accumulate in FP32 register buffers before writing BF16 to HBM; handled transparently by cuBLAS. (2) Adam optimizer state (mtm_{t}, vtv_{t}, master weights)-weight updates involve small increments that BF16 rounds to zero near convergence. (3) Layer-Norm and softmax reductions-summation over wide dimensions accumulates rounding bias in BF16. (4) Loss scalar-used for logging and numerical debugging. Everything on the HBM bandwidth critical path (weight matrices, activation tensors, gradient tensors) is BF16.

Key Takeaways

  • BF16 (1+8+7 bits) matches FP32’s dynamic range at half the storage; FP16 (1+5+10) has higher mantissa precision but a 65,504 maximum that overflows in LLM training without loss scaling.

  • Halving bits per element halves HBM traffic; for memory-bound kernels (activations, norms), this halves runtime directly at no quality cost.

  • BF16 unlocks H100 tensor cores at ≈989\approx 989 TFLOPs/s vs FP32 general ALUs at ≈67\approx 67 TFLOPs/s-a 14.8×14 . 8 \times peak advantage for matmuls, the dominant operation class.

  • Keep FP32 for Adam optimizer state, master weights, layer norm and softmax reductions; use torch.autocast to dispatch the rest automatically.

  • At B=1B = 1 inference, BF16 halves decode latency vs FP32 by halving weight-load bytes; the next gain comes from INT8 or FP8, which halve again.

9.3.2 FP8 and INT8 for Inference-Oriented Training

BF16 is the floor of precision for training stability, but it is not the floor of what the hardware supports. INT8 has been standard for inference deployment for several years; FP8 is now native to H100 and has moved into production training pipelines. Both formats halve the bytes per element again versus BF16, and both require decisions that go wrong in distinct ways. Senior interviews probe this territory when the topic is serving cost, large-model inference on minimal hardware, or the next-generation training infrastructure that frontier labs are deploying today.

INT8 represents values as signed 8-bit integers in the range [−128,127][ - 128 , 127 ]. Mapping a floating-point tensor X\mathbf{X} into this range requires a calibrated scale ss and zero-point zz:

xq=round(xs)+z,x^=(xq−z)⋅sx_{q} = \text{round} \Bigl( \frac{x}{s} \Bigr) + z , \qquad \hat{x} = ( x_{q} - z ) \cdot s

Per-tensor quantization uses one (s,z)( s , z ) pair for the entire tensor-fast but coarse. Per-channel quantization assigns one pair per output channel of a weight matrix-more bookkeeping, but it preserves the dynamic range of channels with outlier magnitudes. H100 INT8 tensor-core throughput is ≈1,979\approx 1 , 979 INT8 TOPS-exactly 2×2 \times the BF16 rate of 989989 TFLOPs/s. This arithmetic benefit is only realizable when the kernel is compute-bound; at B=1B = 1 inference the value of INT8 is purely in the halved weight-load bytes.

FP8 arrives on H100 in two hardware formats. E4M3 (1 sign, 4 exponent, 3 mantissa bits) has a maximum value of 448 and good per-element resolution; it is suited for weights and activations in the forward pass where the dynamic range is relatively narrow. E5M2 (1 sign, 5 exponent, 2 mantissa bits) has a maximum of 57,344 and wider dynamic range at the cost of lower precision; it is suited for gradients in the backward pass, which have heavier distribution tails. Each FP8 tensor carries a companion FP32 scaling factor that maps the full floatingpoint range into the narrow FP8 window. H100 FP8 GEMM throughput is ≈1,979\approx 1 , 979 TFLOPs/s, matching INT8 TOPS and delivering 2×2 \times the BF16 rate on compute-bound matmuls.

Post-training quantization (PTQ) calibrates scale factors on a small representative dataset after training completes; no re-training is required. Quantization-aware training (QAT) inserts fake-quantize operations during the training loop that round values and pass straight-through gradients, letting the model learn to compensate for rounding errors before deployment.

FP8 and INT8 each halve memory footprint relative to BF16; the engineering challenge is managing scale calibration and knowing which tensors absorb rounding without accuracy loss and which do not.

At B=1B = 1 inference, every decode step requires loading the full weight tensor for each layer from HBM before the arithmetic can begin. Token generation rate is therefore bytes per second, not FLOPs per second. INT8 weights occupy half the bytes of BF16 weights. With everything else equal, loading half the bytes takes half the time, and per-token decode throughput approximately doubles-independent of whether INT8 tensor cores are used for the arithmetic. This is the dominant reason production serving systems quantize weights: the decode-bound regime at B=1B = 1 means you are buying throughput by reducing bytes, not by improving peak FLOP rate.

FP8 training offers a compound gain. Halving weight and activation storage reduces HBM pressure throughout the forward and backward pass, not just during weight loading. FP8 GEMMs run at 2×2 \times the BF16 tensorcore rate on H100, shifting the compute-bound threshold to lower batch sizes and directly cutting per-step wall time. DeepSeek-V3 reported roughly 2×2 \times end-to-end training throughput improvement versus BF16 on the same hardware using FP8 mixed precision, with no measurable quality degradation after careful per-tensor scale management.

The accuracy risk is real but concentrated. Weight quantization with per-channel INT8 scales loses almost nothing on standard benchmarks: weight distributions in trained LLMs are approximately Gaussian and map cleanly into [−128,127][ - 128 , 127 ]. Activation quantization is harder because activations have outlier channels whose magnitudes can be 10×10 \times the median channel value; per-tensor INT8 scales clip these outliers into the quantized range. Gradient quantization is hardest of all, which is why E5M2 (wider range) rather than E4M3 is the H100 hardware format for backward-pass tensors.

A Worked Example

Consider a 70 B parameter model deployed for autoregressive generation at B=1B = 1.

Memory footprint:

  • BF16 weights: 70×109×270 \times 10^{9} \times 2 bytes =140= 140 GB. This requires two H100 SXM5 80 GB to hold weights alone, before KV cache or activation buffers.

  • Weight-only INT8 (W8A16-weights in INT8, activations in BF16): 70×109×170 \times 10^{9} \times 1 byte =70= 70 GB. Fits in a single H100 SXM5 80 GB with ≈10\approx 10 GB available for KV cache.

  • FP8 weights (E4M3): same 70 GB footprint as INT8. One GPU required.

Decode latency at B=1B = 1:

  • BF16: dominant cost is loading 140 GB of weights at H100 HBM3 bandwidth 3.35 TB/s: 140/3,350≈41.8140 / 3 , 350 \approx 41 . 8 ms per forward pass.

  • W8A16 INT8: load 70 GB of INT8 weights. Load time =70/3,350≈20.9= 70 / 3 , 350 \approx 20 . 9 ms. Dequantization to BF16 before each GEMM adds <1< 1 ms across all layers. Net: ≈2×\approx 2 \times faster per-token throughput.

Calibration overhead: Running 512 forward passes on a representative dataset to compute per-channel INT8 scales takes roughly 5 minutes for a 70 B model. This one-time cost is irrelevant at deployment scale and does not require touching the original training code.

FP8 training throughput: For a matmul of 2×512×4096×11,008≈46.12 \times 512 \times 4096 \times 11 , 008 \approx 46 . 1 GFLOPs in BF16 at 989 TFLOPs/s: ≈46.7μ\approx 46 . 7 \mus. The same GEMM in FP8 at 1,979 TFLOPs/s: ≈23.3μ\approx 23 . 3 \mus. Across all layers and all matmuls in a training step, this 2×2 \times reduction compounds to a roughly 1.51 . 5-2×2 \times end-to-end step time improvement, with the remainder of step time consumed by non-GEMM operations that remain in BF16.

What You Actually Decide in Practice

  • Start with weight-only INT8 (W8A16). Quantize weights per output channel after training; keep activations in BF16. Libraries like bitsandbytes implement this in one line. Quality loss is typically unmeasurable on 7 B-70 B checkpoints, the weight-load halving gives 2×2 \times decode throughput at B=1B = 1, and nothing in the serving stack changes except the weight storage format.

  • Upgrade to full INT8 (W8A8) only after running your full eval suite. W8A8 uses INT8 tensor cores and provides arithmetic speedup at higher batch sizes, but activation quantization requires handling outlier channels. Per-tensor activation scales cause clipping; per-channel activation scales add synchronization overhead in the kernel. Measure on the specific tasks your deployment serves-perplexity alone is insufficient.

  • Never quantize these tensors. Embedding lookup tables: rare tokens have large weight magnitudes that get clipped under per-tensor INT8 scales, silently degrading handling of infrequent vocabulary. Layer norm scale and bias parameters: these are tiny tensors where rounding has outsized effect on normalized outputs. Softmax normalizers: numerical stability of the log-sum-exp computation is sensitive to precision. The final unembedding projection output logits if they feed directly into sampling without a temperature rescaling step.

  • FP8 training requires H100 and the Transformer Engine library. NVIDIA’s transformer_engine handles E4M3/E5M2 format dispatch, delayed per-tensor scale updates (updating scales every 16 steps is the standard default), and FP8 GEMM scheduling. Do not attempt to write FP8 GEMMs manually; the hardware requirements and scale management are non-trivial. The key knob is the scale update frequency: more frequent updates improve numerical stability at a small runtime overhead; 16-step delayed scaling is the practical default used by production deployments.

  • PTQ is sufficient for INT8; QAT is for 4-bit and below. For INT8, calibrating per-channel weight scales and per-tensor activation scales on a 512-sample in-domain dataset closes essentially all the quality gap relative to BF16. QAT-simulating quantization noise during the training loop-is worth the engineering cost only when PTQ quality becomes unacceptable, which empirically happens below 4-bit precision (the domain of GPTQ, AWQ, and QLoRA-style methods).

How This Shows Up in Interviews

  • “Why is INT8 inference faster than BF16 inference at batch size 1?” Strong answer: at B=1B = 1, decode is weight-loading-bound, not compute-bound. INT8 weights occupy half the bytes of BF16 weights, so the GPU loads the same parameters in half the time. The 2×2 \times arithmetic throughput advantage of INT8 tensor cores is irrelevant because the operation is memory-bound regardless of peak TOPS-the compute units finish the arithmetic before the next tile arrives from HBM.

  • “Walk me through the memory implications of weight-only INT8 for a 70 B model.” Strong answer: INT8 halves weight bytes from 140 GB to 70 GB, reducing the serving requirement from two H100 80 GB GPUs to one. Activations stay in BF16, so the KV cache and intermediate tensors are unchanged. Pertoken decode latency approximately halves because the dominant runtime at B=1B = 1 is weight-load time = weight bytes / bandwidth.

  • Senior follow-up: “When would you choose W8A8 over W8A16 for production, and what are the risks?” Senior answer: W8A8 is worth investigating when the batch size is high enough that the linear layers are compute-bound-above the W8A16 ridge point of ≈295\approx 295 FLOPs/byte, which at a single 4096×\times11008 projection corresponds to roughly B>295B > 295. Below that threshold, activation quantization adds calibration complexity and accuracy risk without delivering the INT8 tensor-core throughput gain, because the kernel is still memory-bound. The specific risk to measure is outlier-channel clipping: check the per-channel activation distributions before and after quantization; if more than ∼0.1%\sim 0 . 1 \% of values are clipped, per-channel activation scales are required.

Key Takeaways

  • INT8 and FP8 each halve weight bytes versus BF16; at B=1B = 1 inference, this halves decode latency because the bottleneck is weight-load time, not arithmetic throughput.

  • W8A16 (INT8 weights, BF16 activations) is the safest quantization tier: near-zero quality loss on 7 B-70 B models, 2×2 \times decode throughput, calibrated offline with 512 samples.

  • W8A8 uses INT8 tensor cores at ≈1,979\approx 1 , 979 TOPS (H100), delivering 2×2 \times arithmetic speedup versus BF16- but only when the kernel is compute-bound; activation quantization adds accuracy risk that must be measured.

  • FP8 training on H100 uses E4M3 (forward pass) and E5M2 (backward pass) with per-tensor FP32 scale factors; the Transformer Engine handles dispatch and delayed scaling; end-to-end training throughput improvement versus BF16 is roughly 1.51 . 5-2×2 \times.

  • PTQ on a 512-sample calibration set is sufficient for INT8; QAT is warranted only below 4-bit, where PTQ accuracy loss becomes non-negligible.