GPUs & Kernels

Almost every performance question in modern ML ("why is decode slow?", "why did FlashAttention help?", "why does FP8 not give 2× end to end?") comes down to one picture: a chip that can do arithmetic hundreds of times faster than it can move bytes from memory. This page builds the mental model an AI engineer needs: how a GPU executes work, where the bytes live, how to tell whether an op is compute-, memory- or overhead-bound, and how the important kernels (GEMM, FlashAttention, decode attention) are engineered around those limits. Interviewers ask about this because it separates people who can use models from people who can make them cheap and fast.

1. GPU architecture for ML engineers

A CPU spends its transistors making one thread fast: big caches, branch prediction, out-of-order execution. A GPU spends them on throughput: thousands of simple lanes, and enough concurrent threads that when some are waiting on memory, others have work ready. That one design decision explains most GPU programming advice.

Streaming multiprocessors (SMs)

The GPU is a grid of identical SMs. An H100 SXM has 132 of them NVIDIA 2022. Each SM is split into 4 sub-partitions, and each sub-partition has a warp scheduler, a register-file slice, FP32/INT32 lanes, special-function units (exp, sin, rsqrt) and one tensor core. Per SM you also get a 256 KB block of on-chip SRAM, split between L1 cache and programmer-managed shared memory (up to 228 KB of shared memory), plus a 256 KB register file NVIDIA 2022. One way to picture it: an SM is a small vector processor with a scratchpad, and the GPU is 132 of those sharing an L2 and HBM.

Threads, warps, blocks, grids (and clusters)

You launch a kernel over a grid of thread blocks (also called CTAs, cooperative thread arrays). The hardware assigns each block to one SM, where it stays until it finishes. Threads within a block can synchronize (__syncthreads()) and share data through shared memory. Threads in different blocks cannot, except through global memory and atomics, or on Hopper and later, within a thread block cluster: a group of blocks scheduled together on neighbouring SMs that can read each other's shared memory ("distributed shared memory") NVIDIA 2022.

The warp is the real unit of execution: 32 threads that share an instruction stream. NVIDIA calls this SIMT (single instruction, multiple threads). It's SIMD with per-lane addressing and the illusion of independent control flow. Each cycle, each scheduler picks a warp whose operands are ready and issues its next instruction.

Latency hiding and occupancy

A global-memory load takes hundreds of cycles (roughly 400–800). The GPU doesn't stall. It switches to another resident warp at zero cost, because every resident warp's registers stay live in the register file. Occupancy is the ratio of resident warps to the maximum the SM supports (64 warps = 2048 threads per SM on H100 NVIDIA docs). Three resources limit it:

• Registers. 65,536 32-bit registers per SM. A kernel using 128 registers/thread fits \(65536/128 = 512\) threads = 16 warps → 25% occupancy. Using 255 registers (the per-thread max) gives ~8 warps.
• Shared memory. If each block needs 100 KB, only 2 blocks fit in 228 KB.
• Block/thread limits. Max blocks per SM and max threads per block.

The key point: high occupancy is a means, not a goal. Good GEMM and attention kernels often run at low occupancy (1–2 big blocks per SM, ~8–16 warps). They get latency tolerance from instruction-level parallelism and explicit asynchronous copy pipelines instead of from many warps. Low-occupancy, high-ILP kernels are the norm at the top end. Memory-bound elementwise kernels, on the other hand, usually do need enough warps (or enough bytes in flight per warp) to saturate bandwidth. By Little's law, bytes in flight = bandwidth × latency: \(3.35\,\text{TB/s} \times {\sim}600\,\text{ns} \approx 2\,\text{MB}\) in flight across the chip, or roughly 15 KB per SM.

Warp divergence

If threads in one warp take different branches of an if, the warp runs both paths one after the other, with lanes masked off. A 50/50 split halves throughput for that region. Since Volta, threads have independent program counters (so intra-warp locking and the like is legal), but divergent paths still serialize. Divergence between warps costs nothing. In ML kernels divergence mostly appears at boundaries (masking the last partial tile, causal masks). The standard trick is to handle full tiles on a branch-free fast path and only mask the edge tiles, which is exactly what FlashAttention does for causal attention.

Tensor cores: where the FLOPs are

Tensor cores are matrix-multiply-accumulate units: each instruction computes \(D = A B + C\) on a small tile (for example 16×8×16 per warp on Ampere-style mma.sync). On Hopper the instruction is WGMMA, issued asynchronously by a whole warpgroup (128 threads) with operands that can come straight from shared memory, on tiles up to 64×256×16. On Blackwell, the tcgen05 MMA instructions write accumulators into a new per-SM tensor memory (TMEM) instead of registers, and a "2-CTA" mode lets a pair of SMs cooperate on one larger MMA Zadouri+ 2026.

Sources for H100: NVIDIA H100 NVIDIA 2022. A100 numbers as quoted in He 2022. Vendor peaks are marketing numbers measured at boost clocks. Real GEMMs typically reach 60–80% of peak, and sustained clocks drop under power limits. One measurement found an H100 matmul runs measurably faster on all-zeros input than on random data, because switching power triggers throttling He 2024.

2. The memory hierarchy

Performance engineering on GPUs is mostly data-movement engineering. Each level of the hierarchy is roughly an order of magnitude bigger and slower than the one above it, and the art is to load each byte from HBM once and reuse it many times from faster levels.

Coalescing: how global memory wants to be read

HBM is accessed in 32-byte sectors (grouped into 128-byte cache lines). When the 32 threads of a warp load 32 consecutive 4-byte words, the hardware merges them into 4 sector transactions: a coalesced access. If they stride by 1 KB, every thread touches a different sector and you waste up to 8× bandwidth, because 32 bytes are fetched for every 4 used Harris 2013. Rules of thumb: map the fastest-varying thread index (threadIdx.x) to the contiguous dimension, use vectorized 128-bit loads (float4, 8 BF16 values) per thread, and keep row starts aligned (padding dimensions to multiples of 8/16/64 helps both alignment and tensor-core tile shapes).

Shared-memory bank conflicts

Shared memory is split into 32 banks, each 4 bytes wide. Address \(a\) maps to bank \((a/4) \bmod 32\). If the threads of a warp hit distinct banks (or the same word, which broadcasts), the access takes one cycle. If \(k\) threads hit different words in the same bank, it replays \(k\) times. The textbook case is reading a column of a 32×32 float tile: every element of the column sits in the same bank, giving a 32-way conflict. The fixes are padding (declare tile[32][33] so columns shift banks) or swizzling, which XORs the column index with bits of the row index so that both row-wise and column-wise accesses spread across banks. CUTLASS/CuTe and TMA use swizzled layouts (32/64/128-byte swizzle modes) because padding wastes space and breaks the layouts tensor cores expect Harris 2013.

3. The roofline model and arithmetic intensity

The roofline model Williams+ 2009 is the single most useful tool for reasoning about kernel performance. Define the arithmetic intensity of a kernel as the FLOPs it performs per byte it moves to or from (some level of) memory:

The ridge point \(I^* = P_{\text{peak}}/B_{\text{mem}}\) is the intensity at which a kernel switches from memory-bound to compute-bound. For H100 SXM in dense BF16: \(I^* \approx 989\times10^{12} / 3.35\times10^{12} \approx 295\) FLOP/byte. In FP8 it's about 590. A kernel has to do roughly 300 multiply-adds' worth of work for every byte it reads from HBM before the tensor cores become the limit. Equivalently, as a time bound:

This "max of compute time and memory time" form is what you use for back-of-envelope estimates Austin+ 2025 (Scaling Book).

Worked examples

Assume BF16 (2 bytes per element) and count only compulsory HBM traffic (each input read once, each output written once).

(a) Elementwise add, \(z = x + y\), \(n\) elements

FLOPs \(= n\). Bytes \(= 3 \times 2n = 6n\). So \(I = 1/6 \approx 0.17\). At 3.35 TB/s that's ~0.56 TFLOP/s, about 0.06% of tensor peak. Even GELU (~10–20 FLOPs/element, on CUDA cores rather than tensor cores) stays far below the ridge. For a 64M-element GELU: read 128 MB + write 128 MB = 256 MB → ~76 µs. The math takes ~1–10 µs. Every elementwise and normalization op is memory-bound. Its cost is its bytes.

(b) GEMM, \(C_{M\times N} = A_{M\times K} B_{K\times N}\)

For a square \(n^3\) GEMM this is \(n/3\). At \(n=4096\), \(I \approx 1365\), well past the ridge, so compute-bound. Time ≈ \(2 \cdot 4096^3 / 989\text{T} \approx 139\,\mu s\) at peak, realistically ~170–200 µs. Now consider a skinny GEMM, where \(M\) is the number of tokens: with \(M \ll K, N\), \(I \approx \frac{MNK}{KN} = M\). Intensity equals the number of tokens in the batch. That's the whole story of LLM decode.

(c) Decode GEMV: one token through a weight matrix

For batch \(B\) tokens through a \(d\times d\) BF16 weight: FLOPs \(= 2Bd^2\), bytes ≈ \(2d^2\) (weights dominate), so \(I \approx B\). At \(B=1\), \(I \approx 1\), about 300× below the ridge. A model with \(P\) parameters in BF16 must stream \(2P\) bytes per decode step:

Batching is nearly free until \(B\) approaches the ridge (~150–300 tokens): you stream the same weights once and do \(B\times\) the work. That's why serving systems batch aggressively, and why quantizing weights to 8 or 4 bits speeds up decode almost proportionally (fewer bytes), even when the math still runs in BF16. Speculative decoding also exploits this: verifying \(k\) draft tokens costs about the same as one memory-bound step. (See the inference page.)

(d) Attention

Per head, with sequence length \(N\) and head dim \(d\): FLOPs ≈ \(4N^2 d\) (\(QK^\top\) and \(PV\)).

• Naive attention writes the \(N\times N\) score matrix \(S\) to HBM, reads it for softmax, writes \(P\), reads \(P\) again: roughly \(4 \times 2N^2\) extra bytes. Intensity ≈ \(4N^2 d / 8N^2 = d/2 \approx 64\) for \(d=128\). Memory-bound, and \(O(N^2)\) memory besides.
• FlashAttention keeps \(S, P\) on chip, so bytes ≈ \(4Nd \times 2\) (read Q, K, V, write O), ignoring re-reads of K/V per Q-tile. Intensity ≈ \(4N^2 d / 8Nd = N/2\). At \(N = 4096\), ~2000: compute-bound. (Re-reading K/V once per Q tile multiplies traffic by \(N/B_r\) but still leaves prefill compute-bound.)
• Decode attention (1 query token vs \(N\) cached keys): FLOPs ≈ \(4Nd\), bytes ≈ \(2Nd \times 2\) (K and V cache), so \(I \approx 1\). With grouped-query attention, \(g\) query heads share each KV head, so \(I \approx g\) (e.g. 8). Still memory-bound. This is why KV-cache size, GQA/MLA and KV quantization matter so much.

4. Compute-bound, memory-bound, overhead-bound

Horace He's "Making Deep Learning Go Brrrr From First Principles" splits runtime into three buckets: compute (time on FLOPs), memory (time moving tensors) and overhead (everything else: Python interpreter, framework dispatch, kernel launches) He 2022. The practical value is diagnosis: each regime has a different fix, and optimizing the wrong one does nothing.

The "doubling test" from the same essay is a great diagnostic to quote: if doubling the batch size increases runtime by much less than 2×, you're overhead-bound. Once runtime grows proportionally, you've left the overhead regime. Similarly, if switching an op from FP32 to BF16 halves its time, it was memory-bound (half the bytes). If it speeds up 8–15×, it moved onto tensor cores and was compute-bound.

Overhead magnitudes: a PyTorch eager op costs on the order of several µs of CPU-side dispatch, and a raw CUDA kernel launch costs a few µs. Meanwhile an H100 can do ~1 GFLOP per µs of tensor math, so a small decode-step kernel can finish faster than the CPU can enqueue the next one. Overhead dominates small-batch inference, small models, RL rollouts with tiny batches, and anything with lots of tiny ops (optimizers over many parameter tensors, which is why "foreach"/fused optimizers exist).

5. Kernel fusion, torch.compile and CUDA graphs

Why fusion matters

Consider y = gelu(x @ W + b) followed by dropout and a residual add, in eager mode. The matmul writes its output to HBM. The bias add reads it and writes it back. GELU reads and writes again. Dropout reads, writes, and also writes a mask. The residual add reads two tensors and writes one. Each memory-bound step costs a full HBM round trip of the activation. Fusing all the elementwise work into the GEMM's epilogue (applied to the accumulator tile while it's still in registers) leaves exactly one write. For a chain of \(k\) elementwise ops on a tensor of \(S\) bytes, unfused traffic is ~\(2kS\) and fused traffic is ~\(2S\): a \(k\times\) speedup for the memory-bound part. He's essay shows x.cos().cos() fused running about as fast as a single cos He 2022.

Fusion has limits. Fusing a reduction (softmax, layernorm) with something that needs the finished reduction requires either the whole row to fit on chip or a clever recurrence (online softmax, which section 7 covers). Fusing two GEMMs back to back is hard because the intermediate is large and the second GEMM needs all of it. FlashAttention is precisely "fuse GEMM → softmax → GEMM", which is why it was a research result rather than a compiler pass.

torch.compile and TorchInductor

PyTorch 2's torch.compile has two stages Ansel+ 2024:

1. TorchDynamo captures the Python program into FX graphs by intercepting bytecode at runtime. It installs guards (shapes, dtypes, Python values) that trigger recompiles when violated, and inserts graph breaks around untraceable code (data-dependent control flow, some Python side effects, .item()).
2. TorchInductor lowers the graph, decides fusion groups, and generates Triton kernels for GPUs (C++/OpenMP for CPUs). Matmuls usually go to cuBLAS/CUTLASS, or optionally autotuned Triton templates with fused epilogues (mode="max-autotune").

Typical gains are 1.3–2× for training on memory-bound-heavy models. Practical gotchas: compile time (minutes for large models), recompilation storms with dynamic shapes (mark dims dynamic, pad or bucket sequence lengths), and graph breaks that silently split the graph (inspect with TORCH_LOGS="graph_breaks" or torch._dynamo.explain). Treat exact flag names as something to check in the current docs PyTorch docs.

CUDA graphs: removing launch overhead

A CUDA graph records a sequence of kernel launches (with their arguments and memory addresses) once, then replays the whole DAG with a single launch NVIDIA 2019. The per-kernel CPU cost disappears, and the driver can launch dependent kernels back to back. Constraints: static shapes and static memory addresses (inputs are copied into fixed buffers), no CPU synchronization or data-dependent control flow inside the graph. LLM inference servers capture one graph per batch-size bucket for the decode step, which is a big part of why decode with small batches isn't launch-bound. In PyTorch, torch.cuda.graphs / torch.cuda.CUDAGraph provide manual capture, and torch.compile(mode="reduce-overhead") applies CUDA graphs automatically PyTorch 2021.

6. GEMM: tiling all the way down

Matrix multiplication is the kernel that matters most and the canonical example of hierarchical tiling. The naive kernel (one thread per output element, loop over \(K\) reading A and B from global memory) has an intensity of about 0.25 FLOP/byte and reaches a few percent of peak. Simon Boehm's step-by-step walkthrough goes from ~1% to ~90%+ of cuBLAS through a sequence of changes, and it's worth knowing that sequence Boehm 2022.

Step 1: block tiling in shared memory

Each block loads a \(BM\times BK\) slab of A and a \(BK\times BN\) slab of B into shared memory, synchronizes, then every thread computes from SMEM. Each element loaded from global memory is now reused \(BN\) times (for A) or \(BM\) times (for B). Intensity against global memory per K-step:

A 128×128 tile gives 64 FLOP/byte and 128×256 gives ~85. That's below the HBM ridge of ~295, which is why the order in which blocks are scheduled matters. Blocks that run at the same time and share A rows or B columns get L2 hits, so actual HBM traffic is several times lower. The "grouped" or "swizzled" block ordering in the Triton matmul tutorial and CUTLASS's rasterization options exist for exactly this L2 reuse. Hopper's TMA multicast within a cluster goes further: one load from L2 is delivered to the shared memory of several SMs.

Step 2: register tiling

Shared memory is fast but finite (see the bank discussion). So each thread (or with tensor cores, each warp) computes a \(TM\times TN\) micro-tile of outputs held in registers. Per \(k\) step it loads \(TM + TN\) values from SMEM and does \(TM\cdot TN\) FMAs, giving a SMEM-level intensity of \(TM\cdot TN/(TM+TN)\). With tensor cores, the "register tile" is the MMA fragment, and the hardware dictates its shape.

Step 3: coalesced, vectorized, conflict-free loads

Global → shared copies use 128-bit vector loads with consecutive threads on consecutive addresses. The SMEM layout is swizzled so that the tensor-core fragment loads (ldmatrix on Ampere) are bank-conflict-free. Transposing on the fly (e.g. when B is stored row-major but the MMA wants column-major) is done in the layout, not with scalar shuffles.

Step 4: software pipelining (double/multi-buffering)

Without pipelining, a block alternates between "load tile k" (tensor cores idle) and "compute tile k" (memory idle). With a ring of \(S\) SMEM stages, the block issues loads for tiles \(k+1 \ldots k+S-1\) while computing tile \(k\). Ampere introduced cp.async: an asynchronous global → shared copy that bypasses registers, completion-tracked by commit groups. Hopper introduced the Tensor Memory Accelerator (TMA): a single thread issues a bulk copy of a whole multi-dimensional tile, described by a tensor map (base, shape, strides, swizzle), and the hardware computes addresses and signals completion through an mbarrier in shared memory NVIDIA 2022. This frees the other threads from address arithmetic and enables multicast.

Step 5: warp specialization and WGMMA (Hopper)

With TMA and async WGMMA, the best Hopper kernels split warps into roles. Producer warps (often one warp, with registers deallocated via setmaxnreg) issue TMA loads into the stage ring. Consumer warpgroups wait on the "full" barrier for a stage, issue WGMMA, and release the stage through the "empty" barrier. This is a classic bounded producer–consumer queue implemented in hardware barriers.

Two further patterns you'll hear about: persistent kernels (launch exactly one CTA per SM, each looping over many output tiles, so tile setup and epilogue overlap with the next tile's mainloop; the scheduler is in software) and Stream-K / split-K (split the K dimension across CTAs when \(M\times N\) is too small to fill the GPU, then reduce the partials).

Tile and wave quantization

Two shape effects explain "why is my GEMM slower with N=4097 than N=4096?" Tile quantization: if a dimension isn't a multiple of the tile size, the edge tiles do wasted work. Wave quantization: with 132 SMs and one tile per SM per wave, 133 tiles take two waves, the second nearly empty, so efficiency drops by about half. Keep dims multiples of 64/128, and pad vocabulary sizes (a known trick is padding vocab to a multiple of 64 or 128).

CUTLASS and CuTe

CUTLASS is NVIDIA's open-source C++ template library for GEMMs and convolutions at every level of this hierarchy NVIDIA CUTLASS. CUTLASS 3.x is built on CuTe, an algebra of layouts: a layout is a (shape, stride) pair mapping logical coordinates to offsets, and layouts compose, tile (local_tile, local_partition) and swizzle. This lets one express "partition this tile across these threads in this tensor-core fragment layout" uniformly NVIDIA blog. In 2025 NVIDIA shipped the CuTe DSL, a Python-embedded way to write CuTe kernels with much faster compile times than C++ templates NVIDIA docs. FlashAttention-4 is written in it Zadouri+ 2026.

7. FlashAttention and FlashDecoding

Standard attention for one head computes \(S = QK^\top/\sqrt{d}\) (\(N\times N\)), \(P = \text{softmax}(S)\), \(O = PV\). Implemented as three kernels, it writes and reads \(S\) and \(P\) through HBM: \(O(N^2)\) memory and traffic. FlashAttention made attention IO-aware. It computes exact attention (no approximation) in tiles that live in SRAM and never materializes \(S\) or \(P\) in HBM Dao+ 2022.

The online softmax recurrence

The obstacle is that softmax normalizes over the whole row, but we want to process K/V in blocks. The fix comes from the online normalizer trick Milakov & Gimelshein 2018, also used for memory-efficient attention Rabe & Staats 2021. Keep, for each query row, a running max \(m\), a running denominator \(\ell\), and an unnormalized output accumulator \(\tilde O\). When a new block of scores \(S^{(j)} = Q_i K_j^\top\) arrives:

Why it's exact: every term ever added to \(\ell\) and \(\tilde O\) carries a factor \(e^{-m_{\text{old}}}\). Multiplying by \(e^{m_{\text{old}} - m_{\text{new}}}\) re-bases everything to the new max, so at the end both numerator and denominator are referenced to the true row max and the ratio equals the standard softmax. Subtracting the max keeps every \(\exp\) argument ≤ 0, so there's no overflow. For the backward pass, store just the per-row log-sum-exp \(L = m + \log \ell\) (\(N\) floats, not \(N^2\)).

FlashAttention (v1, 2022): tiling + recomputation + IO analysis

• Tiling. The forward pass uses the recurrence above. Block sizes are chosen so that \(Q_i\), \(K_j\), \(V_j\) and \(S_{ij}\) fit in on-chip SRAM of size \(M\).
• Recomputation in backward. Instead of storing \(P\) (\(N^2\)), the backward pass recomputes \(S_{ij}\) and \(P_{ij} = \exp(S_{ij} - L)\) tile by tile from Q, K and the saved LSE. This adds FLOPs but saves far more HBM traffic, a trade that's clearly right when you're memory-bound. It's essentially selective activation checkpointing inside the kernel.
• IO complexity. Standard attention needs \(\Theta(Nd + N^2)\) HBM accesses, and FlashAttention needs \(\Theta(N^2 d^2 / M)\). For typical \(d\) (64–128) and \(M\) (~100–200 KB), \(d^2/M\) is well below 1, so traffic drops by many times. The paper also proves this is asymptotically optimal over a range of SRAM sizes Dao+ 2022.
• Memory is linear in \(N\), which is what made long context practical to train.

The reported results were roughly 3× faster GPT-2 training and 15% faster BERT-large than the baselines of the time Dao+ 2022.

FlashAttention-2 (2023): parallelism and work partitioning

FA1 reached only ~25–40% of A100 peak. FA2 made three changes Dao 2023:

1. Fewer non-matmul FLOPs. Don't rescale \(\tilde O\) by \(1/\ell\) on every iteration. Keep it unnormalized and divide once at the end (as in the equations above). Non-matmul FLOPs cost up to ~16× more per FLOP than tensor-core FLOPs, so this matters.
2. Parallelize over sequence length. FA1 parallelized over batch × heads only, which underfills 108–132 SMs at long context and small batch. FA2 also launches one CTA per Q block, with the Q-block loop as the outer, parallel loop.
3. Better warp partitioning within a block. FA1 split K/V across warps ("split-K"), so warps had to exchange partial results through shared memory. FA2 splits Q across warps, all sharing K/V. Each warp owns complete rows and no inter-warp reduction is needed.

The result was about 2× over FA1 and 50–73% of A100 theoretical peak in the forward pass Dao 2023.

FlashAttention-3 (2024): Hopper asynchrony and FP8

On H100, FA2 reached only ~35% of peak because it didn't use TMA/WGMMA, and because softmax became a big fraction of time. H100 has ~989 TFLOPS of matmul but only ~3.9 TFLOPS of special-function (exp) throughput, a 256× gap Dao+ 2024 (blog). FA3's ideas Shah+ 2024:

• Warp specialization: producer warps issue TMA loads of K/V; consumer warpgroups do WGMMA + softmax (the GEMM pipeline pattern from section 6).
• Ping-pong scheduling between warpgroups: while warpgroup 1 does softmax (exp on the SFUs), warpgroup 2 runs its GEMMs on the tensor cores, and vice versa, enforced with named barriers.
• Intra-warpgroup pipelining: overlap the softmax of block \(j\) with the \(QK^\top\) GEMM of block \(j+1\), at the cost of extra registers.
• FP8: block quantization plus "incoherent processing" (multiply Q and K by a random orthogonal/Hadamard matrix to spread outliers before quantizing), giving lower error than naive per-tensor FP8.

Reported results: up to 740 TFLOPS in FP16 (~75% utilization) and close to 1.2 PFLOPS in FP8 on H100 Shah+ 2024.

FlashAttention-4 (2025–26): Blackwell

FA4 exists. It was first released in the flash-attention repo in 2025 and described in a March 2026 paper, "FlashAttention-4: Algorithm and Kernel Pipelining Co-Design for Asymmetric Hardware Scaling" Zadouri+ 2026. The premise is that on Blackwell, tensor-core throughput doubled while exp throughput and shared-memory bandwidth scaled much less, so the non-matmul work must shrink further. Its main ideas, per the abstract:

• Pipelines redesigned around fully asynchronous tcgen05 MMAs with accumulators in tensor memory, and larger tiles.
• Software-emulated exponential (a polynomial approximation on the FMA units) to offload part of the exp work from the SFUs.
• Conditional softmax rescaling: skip rescaling \(\tilde O\) unless the running max has grown enough to threaten numerical stability. The math stays exact because the final \(\tilde O/\ell\) uses a consistent reference max.
• TMEM and 2-CTA MMA in the backward pass to reduce shared-memory traffic and atomics.
• Written entirely in CuTe DSL (Python), with 20–30× faster compilation than C++ templates.

Reported: up to ~1613 TFLOPS BF16 forward on B200 (~71% utilization), up to 1.3× over cuDNN 9.13 and 2.7× over Triton Zadouri+ 2026.

FlashDecoding and split-K for decode

At decode time the query has length 1 per sequence. FA2's grid is batch × heads × (Q blocks = 1). With batch 1 and 32 heads that's only 32 CTAs on 132 SMs, each streaming a long KV cache serially, so most of the GPU idles. Flash-Decoding also splits the KV sequence into chunks processed by different CTAs, each producing a partial output \(O_s\) and its log-sum-exp \(\text{lse}_s\). A small second kernel then merges them Dao+ 2023 (CRFM blog):

This is the same associativity that makes online softmax work, applied across CTAs instead of across loop iterations. The authors reported up to 8× faster end-to-end generation for very long sequences, and attention itself up to ~50× faster than FlashAttention at 64k context and batch 1 CRFM 2023. Choosing the number of splits is a heuristic: enough CTAs to fill the SMs, but not so many that the reduction dominates. Follow-ups such as FlashDecoding++ Hong+ 2023 attack the remaining synchronization. Production engines (FlashInfer, vLLM, SGLang, TensorRT-LLM) combine split-KV with paged KV layouts and GQA packing, where all query heads in a group share one pass over the KV block so intensity rises to ~\(g\). DeepSeek's FlashMLA does the analogous thing for multi-head latent attention DeepSeek FlashMLA.

A related tool is FlexAttention in PyTorch, which lets you write a score_mod/mask_mod in Python (causal, sliding window, ALiBi, document masking) and compiles it into a fused FlashAttention-style Triton kernel, skipping fully masked blocks via a block mask Dong+ 2024.

8. Writing kernels: CUDA C++, Triton and the new DSLs

CUDA C++: the thread-level model

In CUDA you write the program for one thread. You compute your global index from blockIdx, blockDim and threadIdx, and manage shared memory, synchronization, vectorization, tensor-core fragments and async copies yourself (or through CUTLASS/CuTe). You get maximum control, including inline PTX and every new hardware feature on day one. The cost is a lot of code and room for subtle bugs (races, bank conflicts, misaligned TMA descriptors).

// CUDA: one thread per element, grid-stride not shown
__global__ void add(const float* x, const float* y, float* z, int n) {
  int i = blockIdx.x * blockDim.x + threadIdx.x;
  if (i < n) z[i] = x[i] + y[i];      // coalesced: lane i touches address i
}
// launch: add<<<(n + 255) / 256, 256>>>(x, y, z, n);

Triton: the block-level model

Triton Tillet+ 2019 moves the unit of programming from the thread to the block (program). You write code that operates on whole tiles (tl.arange vectors and 2D blocks). The compiler decides how to map them onto threads, coalesce loads, use shared memory, insert synchronization, pipeline loads (num_stages) and emit tensor-core instructions for tl.dot. You still choose the tiling, the grid and the fusion strategy, which is where most of the performance comes from. It's the language TorchInductor emits, and it runs on NVIDIA and AMD GPUs.

import torch, triton
import triton.language as tl

@triton.jit
def add_kernel(x_ptr, y_ptr, out_ptr, n, BLOCK: tl.constexpr):
    pid = tl.program_id(axis=0)                 # which block am I?
    offs = pid * BLOCK + tl.arange(0, BLOCK)    # a vector of BLOCK indices
    mask = offs < n                             # guard the ragged last block
    x = tl.load(x_ptr + offs, mask=mask)
    y = tl.load(y_ptr + offs, mask=mask)
    tl.store(out_ptr + offs, x + y, mask=mask)

def add(x, y):
    out = torch.empty_like(x)
    n = out.numel()
    grid = lambda meta: (triton.cdiv(n, meta["BLOCK"]),)
    add_kernel[grid](x, y, out, n, BLOCK=1024)
    return out

A fused row-wise softmax, one program per row, with the whole row held on chip (the pattern from the official tutorial Triton docs):

@triton.jit
def softmax_kernel(out_ptr, in_ptr, in_stride, out_stride, n_cols,
                   BLOCK: tl.constexpr):
    row = tl.program_id(0)
    cols = tl.arange(0, BLOCK)                  # BLOCK = next_power_of_2(n_cols)
    mask = cols < n_cols
    x = tl.load(in_ptr + row * in_stride + cols, mask=mask, other=-float("inf"))
    x = x.to(tl.float32)                        # do the math in fp32
    x = x - tl.max(x, axis=0)                   # numerical stability
    num = tl.exp(x)
    y = num / tl.sum(num, axis=0)
    tl.store(out_ptr + row * out_stride + cols, y, mask=mask)

# grid = (n_rows,);  BLOCK = triton.next_power_of_2(n_cols)
# Traffic: 1 read + 1 write per element (vs ~5 passes for eager max/sub/exp/sum/div).

When a row doesn't fit on chip (e.g. a 256k vocabulary), you loop over column chunks using the online-softmax recurrence: two passes (max and sum fused online, then normalize) instead of three.

The newer tile DSLs

Between "Triton hides too much for Hopper/Blackwell peak" and "CUDA C++ templates are painful", a crop of tile-level languages appeared in 2024–2026:

9. Profiling: finding where the time goes

PyTorch profiler (torch.profiler.profile with ProfilerActivity.CPU/CUDA, record_function labels, export_chrome_trace) gives an op-level view mapped back to your Python code, plus memory profiling. It's the right first tool for model code PyTorch docs.

10. Numerics: formats, accumulation and rounding

Key ideas

• Accumulate in higher precision. Tensor cores multiply low-precision inputs and accumulate in FP32 (or FP16 optionally). A dot product of length \(K = 8192\) accumulated in BF16 loses precision as the running sum grows, because each add rounds to 8 mantissa bits. Hopper's FP8 tensor-core accumulation turned out to keep only ~14 bits of precision. DeepSeek-V3 worked around this by periodically promoting partial sums to FP32 CUDA-core registers (every 128 elements along K) DeepSeek-AI 2024.
• Scaling is what makes FP8/FP4 work. A format with 3 or 1 mantissa bits and narrow range needs a scale so values fill the representable range: \(x \approx s \cdot q\). The granularity is a trade-off. Per-tensor (cheap, but one outlier wrecks it; "delayed scaling" uses the amax history), per-channel/per-token, or fine-grained blocks: DeepSeek-V3 used 1×128 activation tiles and 128×128 weight blocks, and MX/NVFP4 bake 32-/16-element block scales into the hardware format so the tensor core applies them.
• Reductions and softmax in FP32. Norm statistics, softmax max/sum, loss and logits are commonly kept in FP32 even in "BF16" models. The Triton softmax above upcasts for this reason.
• Stochastic rounding rounds up with probability equal to the fractional distance, so the rounding is unbiased in expectation. It matters when many tiny updates are added to a large value (e.g. BF16 weights + small gradient updates, where round-to-nearest silently drops them) Gupta+ 2015. Some accelerators support it in hardware. On GPUs it's usually done in software inside optimizer kernels, or avoided by keeping FP32 master weights.
• Non-determinism. Floating-point addition isn't associative, so split-K, atomics and different tile configs produce bitwise-different results run to run. This shows up in RL and evals as train/inference mismatch. "Batch-invariant" or deterministic kernels trade some speed for reproducibility.

11. Communication kernels and compute–comm overlap

Multi-GPU training and inference add a fourth cost: moving bytes between GPUs. Collectives (all-reduce, all-gather, reduce-scatter, all-to-all) are implemented by libraries such as NCCL NVIDIA NCCL as GPU kernels that read and write HBM and drive NVLink/InfiniBand.

• Bandwidth cost. A ring all-reduce of \(S\) bytes over \(n\) GPUs sends and receives \(2\frac{n-1}{n}S\) per GPU (a reduce-scatter, then an all-gather). With 1 GB of BF16 gradients over 8 H100s at an effective ~350–450 GB/s per direction: about 1.75 GB / 400 GB/s ≈ 4–5 ms.
• Latency cost. Small messages (TP all-reduces in decode) are latency-bound (µs per hop), so tree algorithms, fused "one-shot" all-reduce kernels and NVLink-SHARP-style in-switch reduction matter more than bandwidth there.
• Overlap. Run comm on a separate stream concurrent with independent compute: DDP buckets gradients and all-reduces them while backward continues, and FSDP prefetches the next layer's all-gather. The catch is that NCCL kernels occupy SMs and HBM bandwidth, so "overlapped" compute slows down. That's why some systems restrict comm to a few SMs or use copy engines. DeepSeek's DeepEP is a good example of custom all-to-all kernels for MoE expert parallelism with tunable SM usage DeepSeek DeepEP.
• Fused compute–comm kernels (e.g. GEMM + reduce-scatter tiled so each tile is sent as soon as it's ready, as in async tensor parallelism) push overlap into the kernel itself.

12. Beyond NVIDIA: AMD, TPU, Trainium

Systolic arrays, briefly. A TPU MXU is a 2D grid of multiply-accumulate cells. In the weight-stationary scheme, weights are preloaded into the cells, activations flow in from one side, and partial sums flow down. Each value is reused across a whole row or column of cells without touching memory, so a 128×128 array does 16K MACs per cycle with very little data movement Jouppi+ 2017. The trade-off: it's superb for large dense matmuls with friendly shapes (dims multiples of 128/256) and less flexible for irregular work. More of the optimization burden sits in the compiler (XLA), which explains the JAX-plus-compiler culture around TPUs. Roofline reasoning carries over unchanged: same equations, different peak and bandwidth numbers.

13. Common kernel interview tasks (with sketched solutions)

Task 1: "Estimate the runtime of this kernel"

Recipe: count compulsory bytes and FLOPs, compute both times, take the max, then sanity-check overhead. Example: RMSNorm over a \(16384 \times 8192\) BF16 activation. Bytes = read 256 MB + write 256 MB (+ tiny weight) ≈ 512 MB, so 512 MB / 3.35 TB/s ≈ 153 µs at peak. Well-written kernels reach ~80–90% of peak bandwidth, so ~170–190 µs. FLOPs ≈ 4 per element ≈ 0.5 GFLOP, negligible. If the measured time is 600 µs, the kernel is making multiple passes or reads uncoalesced. Quote achieved bandwidth = bytes / time as your efficiency metric.

Task 2: Fused softmax in Triton

See section 8. Talking points: one program per row, BLOCK = next_power_of_2(n_cols), other=-inf for masked lanes (so they don't affect the max and exp to 0), subtract the max, FP32 math, a single read and write. For very wide rows, loop over chunks with online max/sum and do a second pass to write normalized outputs. Expected performance: near peak bandwidth. The eager version makes ~4–5 passes.

Task 3: LayerNorm forward in Triton

@triton.jit
def layernorm_fwd(X, Y, W, B, Mean, Rstd, stride, N, eps, BLOCK: tl.constexpr):
    row = tl.program_id(0)
    cols = tl.arange(0, BLOCK)
    mask = cols < N
    x = tl.load(X + row * stride + cols, mask=mask, other=0.).to(tl.float32)
    mean = tl.sum(x, axis=0) / N
    xc = tl.where(mask, x - mean, 0.)
    var = tl.sum(xc * xc, axis=0) / N
    rstd = 1 / tl.sqrt(var + eps)
    w = tl.load(W + cols, mask=mask); b = tl.load(B + cols, mask=mask)
    y = xc * rstd * w + b
    tl.store(Y + row * stride + cols, y, mask=mask)
    tl.store(Mean + row, mean); tl.store(Rstd + row, rstd)   # saved for backward

Talking points: masked lanes must not pollute the variance (tl.where), save mean and rstd for backward, and the backward pass needs a reduction across rows for \(dW\) and \(dB\) (partial sums per program, then a second reduction kernel, or atomics; the Triton tutorial uses a lock-based accumulation). Welford's algorithm is the numerically stable alternative when you can't hold the row.

Task 4: Tiled matmul in Triton

@triton.jit
def matmul(A, B, C, M, N, K, sam, sak, sbk, sbn, scm, scn,
           BM: tl.constexpr, BN: tl.constexpr, BK: tl.constexpr):
    pid_m = tl.program_id(0); pid_n = tl.program_id(1)
    rm = pid_m * BM + tl.arange(0, BM)
    rn = pid_n * BN + tl.arange(0, BN)
    rk = tl.arange(0, BK)
    a_ptrs = A + rm[:, None] * sam + rk[None, :] * sak
    b_ptrs = B + rk[:, None] * sbk + rn[None, :] * sbn
    acc = tl.zeros((BM, BN), dtype=tl.float32)          # FP32 accumulator
    for k in range(0, tl.cdiv(K, BK)):
        a = tl.load(a_ptrs, mask=(rm[:, None] < M) & (rk[None, :] + k * BK < K), other=0.)
        b = tl.load(b_ptrs, mask=(rk[:, None] + k * BK < K) & (rn[None, :] < N), other=0.)
        acc += tl.dot(a, b)                             # tensor cores
        a_ptrs += BK * sak; b_ptrs += BK * sbk
    # epilogue: fuse bias / activation here before the cast
    c = acc.to(C.dtype.element_ty)
    tl.store(C + rm[:, None] * scm + rn[None, :] * scn, c,
             mask=(rm[:, None] < M) & (rn[None, :] < N))

Talking points: FP32 accumulator, masks for ragged edges, @triton.autotune over (BM, BN, BK, num_warps, num_stages), grouped program ordering (a 1D pid remapped so consecutive programs share A rows) for L2 reuse, and epilogue fusion. Compare to cuBLAS at large sizes, where you should reach a large fraction of it, and at skinny shapes, where split-K helps.

Task 5: "Why is this kernel slow?" (diagnosis)

Common planted bugs: column-major access by consecutive threads (uncoalesced, fix by swapping index roles or transposing through SMEM), a reduction done with global atomics per element (fix with a warp-shuffle reduction __shfl_down_sync, then a block reduction in SMEM, then one atomic per block), a 32-way bank conflict on transposed SMEM reads (pad to 33), too few blocks for the GPU (grid of 16 on 132 SMs, so split the work), and register spills from huge per-thread arrays.

Task 6: Parallel reduction / sum

The standard hierarchy: each thread sums a strided chunk in registers (grid-stride loop with vector loads, so bandwidth-bound), then a warp reduction via shuffles (5 steps for 32 lanes, no SMEM), then a block reduction through SMEM (one value per warp), then either a second kernel or one atomic per block. The target is ~peak HBM bandwidth, because a sum is \(I \approx 0.25\)–\(0.5\) FLOP/byte.

Task 7: Back-of-envelope serving questions

"Max decode tokens/s for a 70B BF16 model on 8×H100 at batch 1?" ≈ 1 / 5.2 ms ≈ 190 tokens/s upper bound, realistically less. "What's the KV cache per token for Llama-3-70B-like settings (80 layers, 8 KV heads, d=128, BF16)?" \(2 \times 80 \times 8 \times 128 \times 2\,\text{B} = 327{,}680\,\text{B} \approx 320\,\text{KB/token}\), so a 32k-token sequence needs ~10 GB. At batch 32 with 32k contexts the KV cache is ~340 GB, so every decode step reads more than twice the 140 GB of weights (and it won't even fit next to the weights on 8×80 GB). Long-context batch decode is KV-bandwidth-bound, which is why GQA/MLA, KV quantization and paged attention exist.

Interview question bank