Parallel computing, from Amdahl's law to CUDA kernels that saturate an H100

Parallel computing is the discipline of trading one long dependency chain for many short ones, and then paying honestly for the communication that trade creates. This page derives the speedup laws and the work-depth framework, analyzes scan, reduce, and sort as parallel algorithms, then descends into the CUDA execution model, covering warps, coalescing, occupancy, shared-memory tiling, warp shuffles, and the transformer's building blocks written as CUDA and Triton kernels, each step checked against kernels written for this page and measured on an NVIDIA H100 80GB in this repository. By the end, a reduction that started at 358.0 GB/s runs at 3034.7 GB/s, a matmul climbs from 5.9 to 25.1 TFLOPS on the way to cuBLAS at 51.2, a full causal multi-head attention block drops from 24.044 ms to 1.352 ms, and attention at sequence length 8192 drops from 98.946 ms to 3.575 ms without changing a single flop.

Why this subject matters now

Single-thread performance stopped scaling around 2005, but the economic pressure that used to push clock frequency now pushes parallel width, and the numbers have become extreme. The H100 measured throughout this page exposes 132 streaming multiprocessors, each holding up to 2048 resident threads, so a fully occupied chip is juggling 270,336 threads at once. Its HBM3 stack delivers a measured 3063.5 GB/s and its tensor cores a measured 744.6 bf16 TFLOPS. No sequential abstraction survives contact with a machine like this. The programmer who treats it as a fast serial processor gets a fraction of a percent of what the silicon can do. The one who understands decomposition, communication cost, and the memory hierarchy gets factors of ten to a hundred, and this page measures those factors rather than asserting them.

The subject has also changed shape. Twenty years ago parallel computing meant MPI clusters and OpenMP loops for scientific codes, and the canonical textbook problems were stencils and sparse solvers. Today the dominant parallel workload is the training and serving of neural networks, and the practitioner's daily objects are CUDA kernels, Triton programs, collective communication over NVLink, and attention kernels whose design is constrained by memory traffic rather than arithmetic. The theory did not change. Amdahl's argument from 1967, Brent's scheduling bound from 1974, and Blelloch's scan constructions from 1990 predict the behavior of every kernel measured below. What changed is that the theory is now applied at a scale where getting it wrong costs real money, and where a single well-understood kernel, FlashAttention being the canonical example, can shift the economics of an entire industry. The interview version of this subject is quantitative. Given a kernel, say what bounds it, predict its runtime from first principles, and name the transformation that moves the bound.

The end of frequency scaling

Parallelism is not a preference. It is what was left after a physical argument closed. Dennard scaling, stated in 1974, observed that shrinking a transistor's linear dimension by \(1/\kappa\) also shrinks its voltage and current by \(1/\kappa\), so dynamic power per transistor, \(P \propto C V^2 f\), falls exactly as fast as transistor density rises. Power per unit area stays constant while both density and frequency go up, which meant free performance, every generation, for thirty years. It ended when threshold voltage stopped following, because subthreshold leakage grows exponentially as \(V_{th}\) falls, so \(V_{dd}\) had to stall near 1 volt around 2005. With \(V\) fixed and \(C\) still falling only linearly, the cubic term in the classic frequency-voltage relationship vanishes and power becomes roughly linear in \(f\) at constant area, which at a fixed cooling budget of order 100 to 700 watts pins clocks near 3 to 5 GHz. Everything since has spent the still-growing transistor budget on width instead of speed, on more cores, wider vector units, more independent memory channels, and finally the fixed-function matrix datapaths of a tensor core. The H100's arithmetic peak is roughly four orders of magnitude above a single scalar CPU core's, and essentially none of that came from clock rate. A programmer's only access to it is through explicit parallel structure.

Throughput versus latency, strong versus weak scaling

Two distinctions decide which of the laws below applies to a given problem. The first is architectural. A latency-optimized processor minimizes the time to finish one operation and spends its area on the machinery that shortens a single dependent chain, such as deep out-of-order windows, branch predictors, aggressive prefetchers, and large private caches. A throughput-optimized processor maximizes operations completed per second across many independent chains and spends its area on the operations themselves, tolerating latency by having other work ready. That is the H100's bargain. A memory access costs several hundred nanoseconds, roughly what it costs on a CPU, but the SM holds up to 64 resident warps and switches among them for free, so the latency is never on the critical path as long as parallelism exists. The consequence is asymmetric. Give a GPU a single dependent chain and it is slower than a laptop. Give it \(10^5\) independent chains and nothing else competes.

The second distinction is about how the problem grows. Strong scaling holds the problem size fixed and adds processors, asking how much the wall clock shrinks. It is the regime of interactive latency, of a single inference request, and of any deadline that does not move. Weak scaling grows the problem in proportion to the machine and asks whether the wall clock stays constant. It is the regime of training runs, simulations, and search. Strong scaling is governed by Amdahl's law and is brutally hard, because the serial residue and the communication term stay fixed while the parallel work per processor shrinks toward them. Weak scaling is governed by Gustafson's accounting and is comparatively easy, because the parallel work per processor stays constant and only the communication term grows, usually logarithmically. Both laws are derived next. The practical skill is knowing which question is being asked before quoting either.

                        strong scaling            weak scaling
 problem size           fixed                     grows with p
 question               how much faster?          can we stay on time?
 governing law          Amdahl (asymptote 1/f)    Gustafson (linear in p)
 what breaks first      serial residue, latency   bandwidth, collectives
 canonical workload     one inference request     a training run
 GPU analogue           small batch, launch gaps  big batch, saturated SMs

Speedup, efficiency, and the two laws

Definitions and the honest baseline

Let \(T_1\) be the runtime of the best serial implementation and \(T_p\) the runtime on \(p\) processors. Speedup and efficiency are

$$ S(p) = \frac{T_1}{T_p}, \qquad E(p) = \frac{S(p)}{p} . $$

The baseline matters more than the parallel code. A parallel program timed against itself running on one thread hides the overhead the parallelization added, such as partitioning logic, atomic operations, and restructured data layouts. Published speedups that beat \(p\) (superlinear speedup) are almost always a cache effect, the aggregate L2 and L3 of \(p\) cores holding a working set that thrashed on one core, and occasionally a search algorithm getting lucky on work order. They are not evidence that the law below is wrong, because the law is about a fixed amount of work at fixed per-operation cost.

Amdahl's law, derived

Split the serial runtime into a fraction \(f\) that cannot be parallelized and a fraction \(1-f\) that parallelizes perfectly. Normalizing \(T_1 = 1\),

$$ T_p = f + \frac{1-f}{p} \quad\Longrightarrow\quad S(p) = \frac{1}{f + \dfrac{1-f}{p}}, \qquad \lim_{p \to \infty} S(p) = \frac{1}{f}. $$

The derivation is three lines, but the consequences are brutal and worth internalizing numerically. The serial fraction is a hard asymptote. With \(f = 0.05\), infinite hardware yields at most \(20\times\). Worse, the approach to the asymptote is slow, because the parallel term \((1-f)/p\) must shrink well below \(f\) before the asymptote is felt, which happens around \(p \approx (1-f)/f\). Past that point, additional processors buy almost nothing, and efficiency \(E(p)\) collapses toward zero. In practice \(f\) is not a constant of the program but of the program plus its environment, such as kernel launch latency, the serial part of an optimizer step, a data loader, or a Python interpreter dispatching work. On GPUs, Amdahl's law usually arrives disguised as "the GPU is idle between kernels."

S(p) for f = 0.05, drawn to scale.  Asymptote 1/f = 20 is the dashed line.

 20 ┤- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -  1/f = 20
 16 ┤                                          ●━━━━━━━━━●  p=1024: 19.3
 12 ┤                        ●  p=64: 15.4
  8 ┤          ●  p=16: 9.14
  4 ┤   ●  p=4: 3.48
  0 ┼───┴──────┴─────────────┴────────────────────────────┴──────►  p
     1   4     16            64                          1024

 p_(1/2) = (1-f)/f = 19 processors reach S = 10, half the asymptote.
 Doubling 64 -> 128 buys 15.4 -> 17.4.  Doubling 512 -> 1024 buys 19.0 -> 19.3.
 Efficiency E(p) = S(p)/p:  87% at p=4, 57% at p=16, 24% at p=64, 1.9% at p=1024.

The efficiency row is the one to internalize. Efficiency is the fraction of purchased processor-seconds that do useful work, \(E(p) = S(p)/p\), and under Amdahl it equals \(1/(fp + 1 - f)\), decaying like \(1/(fp)\) once \(fp \gg 1\). At \(f = 0.05\) a 1024-way machine returns 1.9 percent of its capacity. Nothing about the code is wrong. The workload simply ran out of parallelism, and the only remedies are shrinking \(f\) or growing the problem, which is the door Gustafson walks through next.

Problem 1

A renderer spends 92 percent of its serial runtime in a ray tracing loop that parallelizes perfectly and 8 percent in scene setup that does not. Compute the speedup on 8, 64, and 1024 processors, the efficiency at 64, and the asymptotic speedup. At what processor count does the machine reach half of the asymptotic speedup?

Solution. With \(f = 0.08\), \(S(8) = 1/(0.08 + 0.92/8) = 1/0.195 = 5.13\). \(S(64) = 1/(0.08 + 0.92/64) = 1/0.094375 = 10.60\), so efficiency at 64 is \(10.60/64 = 16.6\) percent, meaning five of every six processor-seconds are wasted. \(S(1024) = 1/(0.08 + 0.92/1024) = 1/0.080898 = 12.36\). The asymptote is \(1/0.08 = 12.5\), so going from 64 to 1024 processors, a \(16\times\) increase in hardware, buys a factor of \(12.36/10.60 = 1.17\).

Half the asymptote is \(S = 6.25\), reached when \(f + (1-f)/p = 0.16\), i.e. \(0.92/p = 0.08\), giving \(p = 11.5\), so the machine is half saturated at just 12 processors. This "half-performance point" at \(p_{1/2} = (1-f)/f\) is the single most useful number to extract from a serial profile before buying hardware, and the same functional form reappears below as the half-bandwidth message size of a communication link.

Gustafson's rebuttal, scaling the problem with the machine

Amdahl fixes the problem size and asks how runtime shrinks. Gustafson (1988) observed that real users fix the runtime and grow the problem. Suppose a run on \(p\) processors spends, of its own wall clock, a fraction \(s\) in serial code and \(1-s\) in parallel code. Normalize that wall clock to 1. A single processor executing the same (scaled) workload would need the serial part unchanged plus \(p\) times the parallel part.

$$ T_1^{scaled} = s + p\,(1-s) \quad\Longrightarrow\quad S_{scaled}(p) = s + p\,(1-s), $$

which is linear in \(p\) with slope \(1-s\), with no asymptote. The two laws do not contradict each other. They answer different questions with different definitions of the serial fraction (Amdahl's \(f\) is a fraction of the serial run, Gustafson's \(s\) a fraction of the parallel run). The deeper point is that parallel machines are worth building because workloads scale. Weather models add resolution, and language models add parameters and tokens. Training a modern network across thousands of GPUs is a Gustafson regime, which is why data parallelism works at all. Inference latency on a single prompt, in contrast, is an Amdahl regime, which is why serving is so much harder to parallelize than training.

Problem 2

A simulation running on 128 processors spends 5 percent of its wall clock in serial code. (a) Compute the Gustafson scaled speedup. (b) Find the Amdahl serial fraction \(f\) that would give the same speedup at \(p = 128\), and explain the gap between \(f\) and \(s = 0.05\).

Solution. (a) \(S_{scaled} = 0.05 + 128 \times 0.95 = 121.65\), an efficiency of \(121.65/128 = 95\) percent.

(b) Set \(1/(f + (1-f)/128) = 121.65\) and solve. Then \(f + (1-f)/128 = 1/121.65 = 0.008220\), so \(f\,(1 - 1/128) = 0.008220 - 0.0078125\), giving \(f = 0.000411\). The check \(1/(0.000411 + 0.999589/128) = 121.65\) confirms it. The two fractions differ by a factor of about 120 because they are measured against different denominators. Five percent of a short parallel run is the same absolute serial time as 0.04 percent of the \(\sim\!122\times\) longer serial run of the scaled problem. Stating a serial fraction without stating the denominator is how speedup claims mislead.

Decomposition, assignment, orchestration, mapping

Between "this problem has parallelism" and "this program runs fast" sit four decisions, and naming them separately is what keeps a performance discussion from becoming a guessing game. Decomposition exposes concurrency by breaking the computation into tasks whose dependency graph has small depth. The constraint is Amdahl's, so the target is enough tasks to keep every unit busy, typically several times the processor count so that load imbalance can be absorbed. Assignment distributes tasks to workers, statically (a compile-time or launch-time partition, cheap and predictable, correct when task costs are uniform) or dynamically (a shared queue or work-stealing deque, robust to irregular costs at the price of synchronization). Orchestration is the communication and synchronization that assignment created. It decides which barriers exist, which data moves, how messages are batched, and how communication overlaps computation. It is where most performance is won or lost, because it is the only stage that touches the memory system directly. Mapping binds workers to hardware, deciding which core, which socket, which NUMA node, which SM, and which GPU in which node of which rack.

A CUDA kernel makes the four visible in a single launch. Decomposition is the choice of what one thread computes, one output element for the naive matmul and an \(8 \times 8\) register tile for the fast one. Assignment is the grid and block shape, which is static. The grid-stride loop is a static assignment deliberately coarsened so that each thread's share is long enough to amortize setup. Orchestration is __syncthreads, shared-memory staging, and the atomics at the end of a reduction. Mapping is the hardware's job at block granularity, but the programmer still controls the part that matters, which blocks share an SM and therefore its shared memory and L1, through occupancy. When a kernel underperforms, the diagnosis is almost always naming which of the four is wrong, and the measured progressions later on this page are exactly this. reduce0 to reduce4 changes decomposition and orchestration without touching the algorithm and gains \(8.5\times\).

The orthogonal axis is what is being split. Data parallelism applies the same operation to many elements. The tasks are homogeneous, the dependency graph is shallow and regular, and the code is a loop over an index space. It is the form GPUs, vector units, and collective libraries are built for, and nearly everything in deep learning is a data-parallel program. Task parallelism runs different operations concurrently, a decode stage feeding a compute stage feeding a write stage, or the independent branches of a divide-and-conquer recursion. Its parallelism is bounded by the number of distinct tasks rather than by the data size, so it saturates early, and its natural expression is a thread pool or a fork-join scheduler rather than a kernel launch. Real systems are hybrids. A training step is data-parallel inside each kernel, task-parallel across the data loader, the forward compute, and the gradient all-reduce, and the pipelining of those three is a task-parallel schedule wrapped around data-parallel work. Pipeline parallelism is the special case that deserves its own name because its efficiency has a closed form, derived in the practice section below.

Communication cost and the alpha-beta model

Decomposing a computation creates messages, and messages have a fixed cost and a marginal cost. The standard model charges

$$ T(m) = \alpha + \frac{m}{\beta} $$

for a message of \(m\) bytes, where \(\alpha\) is latency (software stack, link traversal, synchronization) and \(\beta\) is asymptotic bandwidth. Two derived quantities do most of the work in practice. The half-bandwidth message size \(m_{1/2} = \alpha\beta\) is the size at which half the link's bandwidth is achieved (setting \(m/T(m) = \beta/2\) gives \(m = \alpha\beta\)). Messages far below it are latency-bound and should be batched. And for collectives, the relevant metric is bus bandwidth, which normalizes algorithm traffic so it can be compared to the link speed. A ring all-reduce of \(m\) bytes on \(N\) devices moves \(2(N-1)/N \cdot m\) bytes through each link, so \(\text{busbw} = 2(N-1)/N \cdot m / T\). Richer models exist and are worth knowing by name. BSP (Valiant 1990) charges supersteps with global barriers, and LogP (Culler et al. 1993) separates latency, overhead, and gap. The alpha-beta form is the one used daily to size messages.

The numbers below are measured on this machine, a torch.distributed all-reduce (NCCL) between two H100s over NVLink, whose hardware limit is 450 GB/s per direction, median of 20 runs.

Message sizeTime (ms)Bus bandwidth (GB/s)Fraction of NVLink peak
1 MB0.03628.86.4%
16 MB0.097173.538.6%
256 MB0.924290.764.6%
1024 MB3.271328.373.0%
Problem 3

Fit the alpha-beta model to the measured 1 MB and 1024 MB all-reduce points above (with \(N=2\), each byte crosses the link once, so treat \(T(m) = \alpha + m/\beta\) directly). Report \(\alpha\), \(\beta\), and the half-bandwidth message size, then use the fit to predict the 16 MB point and compare with the measurement.

Solution. Two points determine the line, giving \(\beta = (m_2 - m_1)/(T_2 - T_1) = (1024 - 1) \times 2^{20} \,/\, (3.271 - 0.036)\,\text{ms} = 331.6\) GB/s, and \(\alpha = T_1 - m_1/\beta = 36\,\mu s - 1\,\text{MiB}/331.6\,\text{GB/s} = 36 - 3.2 = 32.8\) \(\mu s\). The half-bandwidth size is \(m_{1/2} = \alpha\beta = 32.8\,\mu s \times 331.6\,\text{GB/s} \approx 10.4\) MiB, so any gradient bucket much smaller than about 10 MB is paying mostly latency, which is exactly why data-parallel trainers coalesce gradients into buckets of tens of megabytes before reducing.

The prediction at 16 MiB is \(T = 32.8\,\mu s + 16\,\text{MiB} / 331.6\,\text{GB/s} = 32.8 + 50.6 = 83.4\,\mu s\), versus 97 \(\mu s\) measured. The model underpredicts by 14 percent because NCCL switches protocols and channel counts with message size, so \(\alpha\) and \(\beta\) are not truly constant. A two-parameter fit is a design tool, not a simulator. Note also that the asymptotic \(\beta\) is 74 percent of the 450 GB/s link. Protocol overhead and the fact that all-reduce must also add the numbers keep even giant messages off the hardware roof.

Message passing and the collectives

When memory is not shared, communication becomes explicit, and the interface that standardized it is MPI, with point-to-point MPI_Send and MPI_Recv, their nonblocking Isend/Irecv forms that return immediately so computation can overlap the transfer, and above them a small algebra of collectives that every rank in a communicator enters together. Six carry almost all of the traffic in practice. Broadcast sends one rank's buffer to all. Reduce combines all ranks' buffers with an associative operator into one rank, while Allreduce leaves the result on every rank. Scatter cuts one buffer into pieces, one per rank, Gather inverts it, and Allgather leaves the concatenation everywhere. ReduceScatter reduces and then leaves rank \(i\) holding only the \(i\)-th slice of the result, which is the piece that makes the fast all-reduce work. The value of the collective interface is that the algorithm is the library's choice. The same call becomes a tree at small sizes and a ring at large ones, and the topology detection that picks between them is why nobody should hand-roll these.

Tree all-reduce. Reduce up a binary tree of \(N\) ranks, then broadcast down. Each phase has \(\log_2 N\) rounds and each round moves \(m\) bytes on the links that are active, so under the alpha-beta model

$$ T_{\text{tree}} = 2\log_2 N \left( \alpha + \frac{m}{\beta} \right) + \log_2 N \cdot \gamma m , $$

with \(\gamma\) the per-byte cost of the reduction arithmetic. The latency term is only \(2\alpha\log_2 N\), the best available, but the bandwidth term carries a factor \(\log_2 N\), since every byte crosses \(\log_2 N\) links on the way up and \(\log_2 N\) on the way down. Trees win when \(m\) is small.

Ring all-reduce. Arrange the ranks in a ring and split each buffer into \(N\) chunks of \(m/N\) bytes. Phase one is a reduce-scatter of \(N-1\) steps, in each of which every rank simultaneously sends one chunk to its successor and receives one from its predecessor, adding what it receives into its own copy. After \(N-1\) steps, rank \(i\) holds the fully reduced chunk \(i\) and nobody holds any other complete chunk. Phase two is an allgather with the identical communication pattern and no arithmetic. Another \(N-1\) steps circulate the completed chunks until every rank has all of them. Counting bytes on any one link,

$$ T_{\text{ring}} = 2(N-1)\left( \alpha + \frac{m/N}{\beta} \right) + (N-1)\,\gamma\,\frac{m}{N} = 2(N-1)\alpha + \frac{2(N-1)}{N}\cdot\frac{m}{\beta} + \frac{N-1}{N}\gamma m . $$

The factor \(2(N-1)/N\) is the whole point. It is bounded above by 2 for every \(N\), so a ring all-reduce moves at most twice the buffer size through each link no matter how many ranks participate, while the tree's cost grows like \(\log_2 N\). This is Patarasuk and Yuan's bandwidth-optimal result, and it is also the reason the community reports bus bandwidth rather than raw throughput. Dividing the achieved \(m/T\) by that same \(2(N-1)/N\) yields a number directly comparable to the link speed, which is how the measured table above was computed. The trade is latency. Ring pays \(2(N-1)\alpha\), linear in \(N\), against the tree's \(2\alpha\log_2 N\). The crossover is where \(2(N-1)\alpha\) exceeds the tree's extra bandwidth term, which for the measured \(\alpha \approx 33\,\mu s\) and \(\beta \approx 332\) GB/s lands in the low megabytes. NCCL implements exactly this dichotomy, plus a double binary tree construction that keeps every link busy in both directions, and selects among them by message size and detected topology. The same algorithms appear in Open MPI's coll components.

ring all-reduce, N = 4, buffer split into 4 chunks per rank

 reduce-scatter, 3 steps        allgather, 3 steps
 r0: [a0 a1 a2 a3]              each completed chunk circulates
 r1: [b0 b1 b2 b3]              r0 ends with chunk 0 complete
 r2: [c0 c1 c2 c3]              r1 ends with chunk 1 complete   ──►  then
 r3: [d0 d1 d2 d3]              ...                                  all-to-all
                                                                     completion
 each step: rank i sends chunk (i-step) mod N to rank i+1, adds what arrives
 bytes on one link = 2(N-1)/N * m = 1.5m at N=4, 1.75m at N=8, -> 2m as N grows

In distributed training this maps onto NCCL one-for-one. Data-parallel gradient synchronization is an all-reduce, bucketed to tens of megabytes so it sits on the bandwidth side of the crossover and overlapped with the backward pass so the communication hides under compute. Fully-sharded data parallelism replaces it with reduce-scatter on gradients plus all-gather on parameters, which is the same two phases the ring already runs, exposed separately so that the parameter copy need not be replicated. Tensor parallelism issues an all-reduce or reduce-scatter after each partitioned matmul, at a frequency high enough that it is only viable inside one NVLink domain. Expert-parallel mixture-of-experts layers use all-to-all, whose cost model has no \(\log\) at all and whose performance is set by the worst-case pairing in the topology.

Problem 4

A 1.5-billion-parameter model is trained with bf16 gradients (2 bytes each) across \(N = 64\) GPUs. The interconnect has \(\alpha = 5\,\mu s\) and \(\beta = 200\) GB/s per link. Ignore the reduction arithmetic term. (a) Compute the ring all-reduce time for the whole gradient in one message. (b) Compute the tree all-reduce time for the same message. (c) Repeat both for a single 512 KB bucket and state which algorithm each size prefers.

Solution. The gradient is \(m = 1.5 \times 10^9 \times 2 = 3.0\) GB.

(a) For the ring, \(T = 2(N-1)\alpha + \frac{2(N-1)}{N}\frac{m}{\beta} = 2 \times 63 \times 5\,\mu s + \frac{126}{64} \times \frac{3.0\, \text{GB}}{200\,\text{GB/s}} = 630\,\mu s + 1.969 \times 15{,}000\,\mu s = 630 + 29{,}531 = 30.2\) ms. The latency term is 2 percent of the total.

(b) For the tree, \(T = 2\log_2 N (\alpha + m/\beta) = 12 \times (5 + 15{,}000)\, \mu s = 180.1\) ms, \(6.0\times\) worse. The ratio is \(2\log_2 N \big/ \frac{2(N-1)}{N} = 12/1.969 = 6.1\), which matches. At large \(m\) the two algorithms differ by exactly the ratio of their bandwidth coefficients.

(c) At \(m = 512\) KB \(= 5.24 \times 10^5\) bytes, \(m/\beta = 2.62\,\mu s\). The ring takes \(630 + 1.969 \times 2.62 = 630 + 5.2 = 635\,\mu s\), of which 99 percent is latency. The tree takes \(12 \times (5 + 2.62) = 91.4\,\mu s\), \(6.9\times\) better. The crossover, where \(2(N-1)\alpha = (2\log_2 N - 2(N-1)/N)\,m/\beta\), is at \(m = 630\,\mu s \times 200\,\text{GB/s} / 10.03 = 12.6\) MB. Below that use a tree, above it use a ring, and this is the switch NCCL makes automatically. It is also why gradient bucketing has two opposing pressures. Buckets must be large enough to escape the latency regime but small enough that the last bucket's transfer still overlaps real backward-pass compute.

Work, depth, and parallel algorithms

The work-depth model and Brent's bound

Machine-independent analysis of a parallel algorithm uses two numbers. The work \(W\) is the total operation count, what a serial machine would execute. The depth (or span) \(D\) is the length of the longest chain of dependent operations, the runtime on infinitely many processors. Both are properties of the algorithm's dependency DAG, not of any machine. Two bounds are immediate, \(T_p \ge W/p\) (the processors cannot do more than \(p\) operations per step) and \(T_p \ge D\) (dependencies must be respected). Brent's theorem (1974) says greedy scheduling gets within a factor of two of both.

$$ T_p \le \frac{W}{p} + D . $$

The proof is short enough to give in full. Level the DAG by longest path from an input, so that level \(t\) holds the \(m_t\) operations whose longest dependency chain has length \(t\), for \(t = 1, \dots, D\), with \(\sum_t m_t = W\). A greedy scheduler can execute all of level \(t\) in \(\lceil m_t / p \rceil\) steps, because every operation in the level has all its inputs ready. Summing and using \(\lceil x \rceil < x + 1\) gives

$$ T_p \le \sum_{t=1}^{D} \left\lceil \frac{m_t}{p} \right\rceil < \sum_{t=1}^{D} \left( \frac{m_t}{p} + 1 \right) = \frac{W}{p} + D . $$

The ratio \(W/D\) is the algorithm's average parallelism, the largest \(p\) for which the \(W/p\) term still dominates and processors remain useful. A GPU makes the model concrete. The H100's 132 SMs at 2048 threads each want \(p\) on the order of \(10^5\), so an algorithm needs \(W/D \gg 10^5\) before the chip is even in its operating regime. The design tension of the whole field lives in these two numbers. Transformations that cut depth usually add work, and the right trade depends on how much parallel hardware is actually available.

Reduction, the canonical tree

Summing \(n\) numbers serially is \(W = n - 1\), \(D = n - 1\), with no parallelism, because each partial sum depends on the last. Associativity is the license to re-parenthesize. Pair the inputs into a balanced binary tree, \(n/2\) additions in the first round, \(n/4\) in the second, and so on, giving

$$ W = \frac{n}{2} + \frac{n}{4} + \cdots + 1 = n - 1, \qquad D = \log_2 n . $$

Same work, exponentially smaller depth. The parallelism \(W/D = (n-1)/\log_2 n\) at \(n = 2^{28}\) is about \(9.6 \times 10^6\), comfortably above what the H100 needs. For floating point, re-parenthesization changes rounding, so a parallel sum is not bitwise equal to a serial sum. The tree order is actually more accurate on average (error grows like \(O(\log n)\) rather than \(O(n)\) in the naive bound), but it is a different number, which is one root cause of run-to-run nondeterminism in floating-point training jobs when atomics decide the order instead of a fixed tree. The measured seven-kernel reduction progression later on this page is this tree meeting real hardware.

Scan, Hillis-Steele versus Blelloch

The prefix sum (scan) of \(x_1, \dots, x_n\) under an associative operator \(\oplus\) produces all prefixes \(y_k = x_1 \oplus \cdots \oplus x_k\). It looks inherently serial, \(y_k = y_{k-1} \oplus x_k\), and it is the standard demonstration that "looks serial" means nothing. It is also, following Blelloch (1990), the workhorse primitive of data-parallel programming. Stream compaction, radix sort, sparse-matrix layouts, recurrences \(z_k = a_k z_{k-1} + b_k\) (scan over \(2\times 2\) matrix products, the same trick that parallelizes linear state-space models), and quicksort partitioning all reduce to scan.

Hillis-Steele (1986), step-efficient. In round \(d = 1, 2, \dots, \log_2 n\), every element at index \(k \ge 2^{d-1}\) updates \(y_k \leftarrow y_k \oplus y_{k - 2^{d-1}}\). After round \(d\), each \(y_k\) holds the sum of the \(\min(k, 2^d)\) elements ending at \(k\), and induction on \(d\) makes this precise. Depth is \(\log_2 n\) rounds. For work, round \(d\) performs \(n - 2^{d-1}\) additions, so

$$ W_{HS} = \sum_{d=1}^{\log_2 n} \big(n - 2^{d-1}\big) = n \log_2 n - (n - 1) . $$

Blelloch (1990), work-efficient. Two sweeps over a conceptual balanced tree. The up-sweep is the reduction tree, \(n - 1\) additions, leaving each internal node holding the sum of its leaves. The down-sweep replaces the root with the identity and, at each node, passes its own value to its left child and (its value \(\oplus\) the left child's stored sum) to its right child, another \(n - 1\) operations. What arrives at each leaf is exactly the sum of everything to its left, the exclusive scan. The total is

$$ W_{B} = 2(n-1) = O(n), \qquad D_B = 2\log_2 n . $$

Blelloch does asymptotically optimal work at twice the depth. Hillis-Steele does \(\log_2 n\) times more work at half the depth. Which wins on a real GPU depends on whether the machine is saturated. Below saturation, extra work is free and lower depth wins, which is why warp-level scans in production libraries (CUB's warp scan) are Hillis-Steele across 32 lanes, while the block and device levels are work-efficient structures glued together in the three-phase pattern of scanning blocks locally, scanning the block sums, and adding the carries. Merrill and Garland's decoupled look-back (2016) turns the middle phase into a single pass, reaching memcpy-rate scans. It is the algorithm inside CUB's DeviceScan today.

Scanning \(n = 2^{26}\) floats on this machine measures that difference. A three-phase Triton scan written for this page runs in 0.405 ms (1324.9 GB/s counting one read and one write), while torch.cumsum, which dispatches to CUB's single-pass decoupled look-back, runs in 0.270 ms (1987.1 GB/s), agreeing to a maximum relative error of \(1.1 \times 10^{-6}\). Both are far below the 3034.7 GB/s the reduction reaches, and the reason is structural rather than a coding defect. A three-phase scan touches the data twice, once to scan blocks and once to add carries, so its floor is four traversals of memory rather than two. Decoupled look-back exists precisely to get back to two, by having each block publish its aggregate and then inspect its predecessors' published state instead of waiting for a separate global pass. Scan is the primitive where the gap between an obvious implementation and the state of the art is a factor, not a percentage, and the factor is entirely memory traffic.

Problem 5

For \(n = 2^{20}\), (a) count the additions performed by Hillis-Steele and by Blelloch scan, (b) using Brent's bound with \(p = 2^{15}\) lanes, estimate the step counts of both and decide which is faster in that regime, and (c) repeat with \(p = 2^{20}\) (one lane per element).

Solution. (a) \(W_{HS} = n\log_2 n - (n-1) = 20 \times 1{,}048{,}576 - 1{,}048{,}575 = 19{,}922{,}945\) additions. \(W_B = 2(n-1) = 2{,}097{,}150\). Hillis-Steele does \(9.5\times\) more work.

(b) Brent gives \(T_p \le W/p + D\). Hillis-Steele takes \(19{,}922{,}945/32{,}768 + 20 = 608 + 20 = 628\) steps. Blelloch takes \(2{,}097{,}150/32{,}768 + 40 = 64 + 40 = 104\) steps. With the machine oversubscribed (\(n \gg p\)), the work term dominates and the work-efficient algorithm wins by roughly \(6\times\).

(c) With \(p = n = 2^{20}\), Hillis-Steele takes \(\le 19 + 20 = 39\) steps (the work term is \(19{,}922{,}945 / 1{,}048{,}576 = 19\)), Blelloch \(\le 2 + 40 = 42\) steps, and the true step counts are just the depths, 20 versus 40. At full saturation the step-efficient algorithm wins despite doing \(9.5\times\) the work, because the extra work rides in otherwise idle lanes. This crossover is the whole design argument in miniature, and it is why production scans are hybrids, Hillis-Steele inside a warp with work-efficient composition across warps and blocks.

Sorting and the scan-centric toolkit

Parallel sorting shows how primitives compose. The split operation sends elements with bit 0 before elements with bit 1, stably, using two scans. An exclusive scan of the zero-flags gives each zero its destination, and the total zero count plus an exclusive scan of the one-flags places the ones. Radix sort is \(b\) successive splits for \(b\)-bit keys, so its work is \(O(bn)\) with depth \(O(b \log n)\), and in practice GPUs process 4 to 8 bits per pass with per-block histograms and a scan over histogram columns. CUB's radix sort is the practical standard for keys of moderate width. Comparison-based alternatives matter when keys are wide or comparisons are custom. Bitonic networks do \(O(n \log^2 n)\) work at \(O(\log^2 n)\) depth and are the standard choice inside a block where their obliviousness (fixed compare-exchange pattern, no data-dependent control flow) fits SIMD lanes perfectly, while sample sort, which splits the input by \(p-1\) sampled pivots into independently sorted buckets, wins at device scale. The theoretical benchmark is Cole's merge sort at \(O(n\log n)\) work and \(O(\log n)\) depth, rarely implemented but the proof that comparison sorting has no inherent depth barrier above \(\log n\).

The bitonic network deserves its depth counted, because the count explains where it is used. A bitonic sequence rises then falls. Batcher's observation is that a compare-exchange between element \(i\) and element \(i + n/2\) of a bitonic sequence of length \(n\) splits it into two bitonic halves with every element of the lower half no greater than every element of the upper. Sorting a bitonic sequence of length \(n\) therefore takes \(\log_2 n\) rounds of \(n/2\) compare-exchanges each, and building a bitonic sequence of length \(n\) requires bitonically sorting two halves of length \(n/2\) in opposite directions first. Writing \(D(n)\) for the depth of the full sort, \(D(n) = D(n/2) + \log_2 n\), which unrolls to \(D(n) = \sum_{k=1}^{\log_2 n} k = \tfrac{1}{2}\log_2 n(\log_2 n + 1)\) rounds, so \(W = \tfrac{n}{4}\log_2 n(\log_2 n+1)\) compare-exchanges. For \(n = 1024\) that is 55 rounds and 28,160 comparisons against \(n \log_2 n = 10{,}240\) for an optimal comparison sort, a \(2.75\times\) work premium bought for a completely oblivious, branch-free, index-computable schedule. Inside a block or a warp, where lanes would otherwise idle and any data-dependent control flow costs divergence, that is a good trade. Across a device, where the work premium is paid in real bandwidth, radix sort wins. For the measured comparison on this machine, torch.sort on \(2^{26}\) random 32-bit keys, which dispatches to CUB's radix sort, takes 4.473 ms, or 15.0 billion keys per second, which is roughly 8 passes over 268 MB at close to streaming rate.

Stream compaction is the third member of the scan family and the one that shows up most in real pipelines. Given \(n\) elements and a predicate, it produces the surviving elements packed contiguously, in order. The algorithm is two lines of the algebra above. Compute the flag vector \(f_i \in \{0,1\}\), take its exclusive scan to get each survivor's destination index, and scatter. Work is \(O(n)\), depth is the scan's \(O(\log n)\), and the output length is the scan's total. Everything that filters on a GPU takes this form, removing terminated rays in a path tracer, gathering active particles, selecting the tokens routed to one expert in a mixture-of-experts layer, or collecting the non-zeros of a sparse operation. The measured version on this machine, a Triton kernel that computes per-tile popcounts, scans the tile counts, and then scatters survivors in one pass over \(n = 2^{26}\) floats with a 50 percent predicate, runs in 0.976 ms against 0.522 ms for torch's fused masked_select, with exact agreement. Both are well under the streaming ceiling because the scatter writes are only half-dense and the output addresses are data-dependent, which is the permanent tax on compaction. Reads coalesce, writes do not.

Shared memory, with threads, locks, and the memory model

Shared-memory parallelism is treated compactly here because it has its own page. See concurrent systems programming for the full treatment of memory models, linearizability, condition variables, and the lock-free landscape. What follows is the subset a performance-minded parallel programmer must have loaded, the cost model rather than the correctness theory.

Locks, barriers, and what they cost

A thread is a schedulable instruction stream sharing an address space with its siblings, which is what makes communication free and correctness hard. A mutex restores serialization over a region. Its cost has three parts, and only the third usually matters. The uncontended acquire is a single successful atomic compare-and-swap, tens of cycles. The contended acquire adds a futex-style descheduling round trip, a few microseconds. The real cost is Amdahl's. A critical section of length \(c\) entered by \(p\) threads at rate \(\lambda\) each serializes at utilization \(\rho = p\lambda c\), and as \(\rho \to 1\) the queue grows without bound. The practical rule follows directly, and it is the only lock rule worth memorizing. Shrink the critical section or shard the lock, because making the lock itself faster moves a term that was never dominant.

A barrier makes every thread wait for the slowest, which converts load imbalance into wall-clock time with no discount. If thread \(i\) takes time \(T_i\) in a phase, the phase costs \(\max_i T_i\), and the expected maximum of \(p\) independent draws grows with \(p\) even when the mean does not. With \(p = 1000\) workers whose times are normal with mean \(\mu\) and standard deviation \(\sigma\), the expected maximum is about \(\mu + 3.24\sigma\). This is the quantitative reason bulk-synchronous designs lose to asynchronous ones at scale, and the reason a straggler in a data-parallel training step costs the whole cluster. CUDA's __syncthreads is a barrier over one block only, which is why it is fast, tens of cycles when warps are already co-resident, and why the CUDA model refuses to offer a cheap grid-wide barrier, which would force all blocks to be resident simultaneously and destroy the scheduling freedom that makes kernels portable across GPU sizes.

Atomics and the acquire/release model

An atomic read-modify-write does two separable jobs, and conflating them is the most common source of both bugs and needless slowness. The first is atomicity, meaning the read, the modify, and the write are indivisible with respect to other atomics on the same location. The second is ordering, meaning which other memory operations, on other locations, are visible to be before or after it. C++ and the hardware expose these separately. memory_order_relaxed buys atomicity alone and is the right choice for a statistics counter or a histogram bin, where no other data is being published. release on a store guarantees that everything the storing thread wrote before it becomes visible to any thread that performs an acquire load reading that value. The pair is a one-way fence, and together they implement the publish-then-observe idiom that every lock, every queue, and every reference count is built from. Sequential consistency, the default, additionally imposes a single global order on all such operations, which on x86 costs an mfence and on weakly ordered machines like ARM or the GPU costs considerably more.

The mapping to GPUs is direct and worth stating because CUDA's atomics are usually taught without it. atomicAdd in CUDA is relaxed, ordering nothing. Publishing data between blocks requires __threadfence(), which is the release/acquire machinery spelled out, and the correct modern spelling is cuda::atomic_ref with an explicit scope (thread_scope_block, thread_scope_device, thread_scope_system). Scope is the extra dimension GPUs add. A block-scoped atomic can be satisfied in the SM's shared memory or L1 and never leave the SM, which is exactly why the privatized histogram measured later is \(144.6\times\) faster than the device-scoped version.

False sharing, a cache-line calculation

Coherence is maintained in units of cache lines, 64 bytes on x86 and 128 on some ARM parts, not in units of variables. Two threads writing to different variables that happen to share a line will ping-pong that line between their private caches, each write invalidating the other's copy, with no logical sharing at all. The arithmetic is simple. A 64-byte line holds sixteen 4-byte counters, so eight threads writing eight adjacent int counters occupy a single line and every one of their writes is a coherence miss. Under MESI, a write to a line held Shared or Modified elsewhere requires a read-for-ownership, invalidating every other copy. The round trip through the last-level cache or the interconnect costs on the order of 100 nanoseconds against roughly 1 nanosecond for an L1 hit, a factor of about 100 per write. Padding each counter to alignas(64) costs 60 wasted bytes per counter and removes the traffic entirely.

Measured on this machine's 26-core Xeon Platinum 8480+, eight threads each performing 50 million relaxed atomic increments take 4.245 s when the eight counters share one 64-byte line and 0.304 s when each is padded to its own line, a \(14.0\times\) penalty. Per increment that is 10.6 ns versus 0.76 ns, which brackets the coherence round trip predicted above once the atomic's own cost is included. No line of algorithm pseudocode shows this. It is a property of the layout, and it is the single most common reason a correctly-parallelized loop fails to scale. The same phenomenon appears at every level of the hierarchy. On the GPU it is called a bank conflict inside shared memory and sector overfetch in HBM, and the measured transpose and strided-read numbers below are its GPU-side equivalents.

Lock-free structures and ABA

A lock-free structure guarantees system-wide progress, in that some thread always completes in a bounded number of steps, so a descheduled thread cannot block the others. The canonical construction is Treiber's stack, whose push loads the head, points the new node at it, and installs the new node with a compare-and-swap, retrying if the head moved. It is correct, simple, and carries a trap. Suppose thread A reads head \(= X\) and is descheduled. Threads B and C pop \(X\), pop \(Y\), and push \(X\) back. Thread A resumes and its compare-and-swap succeeds because the head is again \(X\), but the rest of the list behind \(X\) is now different, and A has spliced a freed node back into the structure. This is the ABA problem. Compare-and-swap tests a value, not a history. The standard remedies are a tagged pointer that packs a monotone counter beside the address so the compare covers a version, hazard pointers that publish which nodes a thread is reading so no one reclaims them, and epoch-based reclamation that defers frees until every thread has passed a grace point. Every one of these is a memory-reclamation scheme, which is the real lesson. The hard part of lock-free programming is not the algorithm, it is deciding when memory may be freed. On GPUs the question mostly evaporates, because kernels are bulk-synchronous and allocation happens on the host, which is one reason GPU code gets away with far cruder synchronization than CPU code.

SPMD, vectorization, and the CPU runtime

The model, and why compilers cannot do this alone

Nearly all parallel code today is written in the single program, multiple data style, with one program text executed by many workers, each discovering its identity through an index (an MPI rank, a thread id, a CUDA thread coordinate) and using it to claim a slice of the data. SPMD is a contract about divergence. Workers may take different branches, and the model defines what that costs. On multicore CPUs the workers are threads and divergence is free. Inside a vector unit, the workers are SIMD lanes and divergence is emulated with masks. Both sides of a branch execute, with lanes disabled on the side they did not take. The ISPC compiler (Pharr and Mark, 2012) made this explicit and ergonomic on CPUs. A "gang" of program instances maps to the lanes of an AVX register, varying values live one-per-lane, uniform values are shared, and the compiler inserts the masks. CUDA's SIMT model is the same idea with the mask management moved into hardware, which is the correct one-sentence relationship between ISPC and CUDA. Automatic vectorization of scalar loops fails not because compilers are weak but because the semantics of scalar C demand proofs (no aliasing, no cross-iteration dependence) that the SPMD contract simply grants by construction.

A measured CPU reduction, and the false-sharing tax

Two CPU measurements from this machine (a 26-core Xeon Platinum 8480+) calibrate the model before the GPU sections. A scalar OpenMP sum over \(2^{28}\) floats runs at 5.6 GB/s on one thread and 111.3 GB/s on 26 threads, a \(19.9\times\) scaling that stops at the socket's memory bandwidth, not at its arithmetic. Adding hyperthreads makes it worse (74.1 GB/s at 52 threads) because sibling hyperthreads share load ports on an already bandwidth-bound loop. And a false-sharing microbenchmark, 8 threads doing relaxed atomic increments to per-thread counters, runs in 4.245 s when the counters share one 64-byte cache line and 0.304 s when each is padded to its own line, a \(14.0\times\) penalty for a layout detail that no line of algorithm pseudocode would ever show. Communication cost on a shared-memory machine is implicit, levied by the coherence protocol in units of cache lines, and the first skill of multicore performance work is making that invisible traffic visible.

// Scalar OpenMP reduction over 2^28 floats (1 GiB).
// Measured on the 26-core Xeon 8480+ in this repository:
//   1 thread: 5.6 GB/s    26 threads: 111.3 GB/s    52 threads: 74.1 GB/s
// build: g++ -O3 -fopenmp -o sum sum.cpp
#include <omp.h>
#include <cstddef>

float parallel_sum(const float *x, size_t n) {
    float s = 0.0f;
    // The reduction clause gives each thread a private accumulator and
    // combines them in a tree at the end: W = n-1, D = O(n/p + log p).
    // Without it, a shared "s += x[i]" would be both a data race and a
    // serialization point.
    #pragma omp parallel for reduction(+:s) schedule(static)
    for (size_t i = 0; i < n; ++i)
        s += x[i];
    return s;
}

Gang width, SIMD width, and the cost of divergence

The width is a hardware constant and the programmer's unit of waste. On this machine's Xeon, AVX-512 gives 16 fp32 lanes, so an ISPC gang of 16 program instances occupies one register. AVX2 gives 8, NEON and SVE at 128 bits give 4, and a CUDA warp gives 32. Divergence costs the same way in all of them. If a branch splits a gang so that a fraction \(\phi\) of lanes take the \(t\)-cycle side and \(1-\phi\) take the \(e\)-cycle side, a masked implementation executes both, costing \(t + e\) instead of the \(\phi t + (1-\phi)e\) a scalar machine would pay. The worst case is a balanced two-way branch, where efficiency drops to 50 percent, and an \(n\)-way switch in which every lane picks a different arm drops it to \(1/n\). Two corollaries follow. First, divergence is only a problem within a gang. Thirty-two warps that each take a different branch, uniformly within themselves, cost nothing at all, which is why sorting or bucketing work so that a warp's lanes agree is a standard and highly effective optimization. Second, the cost is bounded by the sum of the arms, so short divergent regions are cheap and long ones are not. Hoisting a long branch out of a kernel and launching two kernels is frequently the right answer.

Thread pools, work stealing, and the Cilk bound

Task-parallel CPU code needs a scheduler, and the argument for which one is settled. A thread pool with a single shared queue is correct and contends on one lock at every task boundary, which caps throughput at roughly one task per lock round trip. Work stealing, introduced in this form by Blumofe and Leiserson for Cilk in 1994, gives each worker its own double-ended queue. The owner pushes and pops at the bottom with no synchronization in the common case, and an idle worker steals from the top of a random victim's deque. Stealing from the opposite end matters twice over. It minimizes contention, since owner and thief touch different ends, and it steals the oldest task, which in a divide-and-conquer recursion is the largest remaining subtree, so one steal buys a lot of work.

The result that justifies the design is a probabilistic bound. A work-stealing scheduler executes a fully strict computation with work \(T_1\) and span \(T_\infty\) on \(p\) processors in expected time

$$ T_p \le \frac{T_1}{p} + O(T_\infty) , $$

which is Brent's bound recovered by a decentralized, randomized scheduler with no global knowledge. The proof idea is an accounting argument on steal attempts. Every time a processor is idle it makes a steal attempt, and a potential function on the remaining critical path shows that \(\Theta(p)\) steal attempts suffice in expectation to reduce the potential by a constant factor, so the total number of steal attempts is \(O(pT_\infty)\), and dividing the processor-cycles spent stealing by \(p\) gives the \(O(T_\infty)\) additive term. The same paper bounds the space, \(S_p \le p S_1\), and the communication. This is the theoretical backing for essentially every modern task runtime. Intel's TBB, Java's ForkJoinPool, Go's goroutine scheduler, Rust's Rayon, OpenMP's task construct, and the async runtimes built on Tokio all implement per-worker deques with randomized stealing.

OpenMP is the loop-parallel counterpart and remains the shortest path from a serial loop to a scaling one. Its scheduling clauses are the assignment decision made explicit. static partitions the index range at entry with zero runtime overhead and is right when iterations cost the same. dynamic hands out chunks from a shared counter and is right when they do not, at the cost of an atomic per chunk. guided starts with large chunks and shrinks them, amortizing overhead early and balancing late. The measured OpenMP reduction below shows what the ceiling is on a CPU once the assignment is right. It is memory, not arithmetic.

Processes versus threads, in practice

Threads share an address space while processes do not, and the choice is usually forced by something other than performance. Threads win on communication, which is a pointer, and lose on isolation and on any global interpreter state. The second clause is why Python data pipelines historically used multiprocessing. With a global interpreter lock, threads interleave bytecode rather than executing it concurrently, so CPU-bound Python gets no scaling from threads and full scaling from processes, at the price of pickling every argument across a pipe. That constraint is loosening, with free-threaded builds landing as an option, but the operational shape of the ecosystem still assumes it. The costs are quantitative and stable. Creating a thread is tens of microseconds, and creating a process is hundreds. A shared-memory handoff between threads is nanoseconds, while a pipe or socket handoff between processes is single-digit microseconds plus a copy unless a shared-memory segment is used explicitly. In GPU work the distinction resolves itself. The standard distributed-training layout is one process per GPU, so that each has its own CUDA context and its own Python interpreter, with NCCL rather than shared memory carrying the data between them. The GIL then stops mattering, because the Python thread's job is only to enqueue asynchronous kernels.

The CUDA execution model

Grids, blocks, and warps, the three-level contract

A CUDA kernel launch names a grid of thread blocks, each block up to 1024 threads. The hardware adds a third level the programmer does not declare but must design around, the warp, 32 threads that execute in lockstep on a SIMD datapath. The three levels have distinct guarantees, and the entire model is those guarantees. Threads in one block run on one SM, can share its shared memory, and can barrier with __syncthreads. Blocks are independent and schedulable in any order, which is the property that makes a kernel scale across any number of SMs and is why inter-block cooperation within a kernel requires either atomics or the newer cooperative-launch machinery. Warps are where SIMT meets physics. Divergence within a warp serializes the branch paths, and memory requests are issued per warp, which is what makes coalescing (next subsection) the dominant performance rule.

grid (kernel launch)
 ├── block (0,0) ──► scheduled on some SM, shares 228 KB smem/L1,
 │     ├── warp 0: threads 0-31    ─┐  __syncthreads() barrier scope
 │     ├── warp 1: threads 32-63    ├─ resident together, zero-cost
 │     └── ...                     ─┘  context switch between warps
 ├── block (1,0) ──► any other SM, no ordering vs block (0,0)
 └── ...                H100: 132 SMs x up to 64 resident warps each

The reason GPUs tolerate a 500-nanosecond memory latency without the machinery a CPU needs (out-of-order windows, prefetchers) is oversubscription. Each SM holds up to 64 resident warps and switches among them every cycle at zero cost, so while one warp waits on memory, others issue. This is Little's law in silicon. To sustain bandwidth \(B\) at latency \(L\), the machine must keep \(B \times L\) bytes in flight, and at \(B = 3\) TB/s, \(L = 500\) ns that is 1.5 MB of outstanding requests, far more than any single thread can generate. Warps are the mechanism that generates it.

The memory hierarchy in numbers

Every optimization on this page is a move between two levels of the following table, so the constants are worth holding. Registers are per-thread and private, 65,536 32-bit registers per SM, addressable only by index known at compile time. A register access is free in the sense that it is part of the instruction. Shared memory is per-block, carved out of a 256 KB unified block with L1 on Hopper, with 48 KB available per block by default and up to 227 KB with an explicit opt-in. Latency is roughly 20 to 30 cycles and bandwidth per SM is a few thousand bytes per cycle across 32 banks. L1 is the same physical array, hardware-managed. L2 is 50 MB, device-wide, shared by all SMs, and is the coherence point for global memory. HBM3 is 80 GB at a measured 3063.5 GB/s and roughly 500 to 700 nanoseconds of latency. Beyond that lies host memory over PCIe at tens of GB/s, and other GPUs over NVLink at a measured 328.3 GB/s of bus bandwidth.

level        scope        size (H100)      latency        who manages it
 registers    thread       256 KB/SM        ~0             compiler
 shared/L1    block        256 KB/SM        ~25 cyc        programmer / hw
 L2           device       50 MB            ~200 cyc       hardware
 HBM3         device       80 GB            ~500-700 ns    hardware
 NVLink       node         peer GPUs        ~2-10 us       NCCL
 PCIe/host    system       host DRAM        ~10 us         driver

 capacity x ~200 per step down, bandwidth / ~10, latency x ~20.
 Every kernel optimization below is: move a reuse pattern one level up.

The ratios matter more than the absolutes. An optimization is worth roughly a factor of ten per level climbed, and the measured progressions confirm it. Tiling moves matmul reuse from HBM to shared memory for \(1.4\times\) on this kernel shape, register blocking moves it from shared memory to registers for \(3.1\times\), and fusing attention moves the entire score matrix from HBM to registers for \(27.7\times\).

Coalescing, measured

When a warp issues a load, the hardware merges the 32 addresses into as few 32-byte sectors as possible. Thirty-two consecutive floats occupy \(32 \times 4 = 128\) bytes, four sectors, perfectly coalesced with every fetched byte used. A strided access pattern touches more sectors for the same useful bytes. At stride 8 floats and beyond, every lane's 4-byte load drags in its own 32-byte sector and the useful fraction is at most \(4/32 = 12.5\) percent. The kernel measured here reads \(2^{28}\) floats with a template stride on the H100 80GB, counting only useful bytes in the bandwidth number.

Stride (floats)Effective bandwidth (GB/s)Fraction of stride-1Sector-model prediction
12380.51.001.00
21856.30.780.50
41218.30.510.25
8674.30.280.125
16352.90.150.125
32325.20.140.125

The measurements sit above the pure sector model at small strides because the L2 (50 MB on H100) catches part of the overfetch when neighboring warps want the skipped words, and they approach the \(1/8\) floor from above once the stride exceeds the sector. The engineering rule falls out directly. Structure-of-arrays layouts keep each field contiguous per warp and coalesce, while array-of-structures layouts are a built-in stride equal to the struct size. A \(7.3\times\) bandwidth difference (2380.5 versus 325.2 GB/s) for identical useful bytes is the largest single factor a data-layout decision controls on this machine.

Shared memory has the analogous rule with banks instead of sectors, 32 banks, 4 bytes wide, and a warp conflicts when two lanes hit different addresses in one bank, serializing the access. The measured transpose in the implementation section walks that lesson end to end, from 451.1 GB/s to 2694.0 against a 2837.4 GB/s copy.

Occupancy, the resource-partitioning arithmetic

Each SM has a fixed budget of 65,536 32-bit registers, up to 228 KB of shared memory (48 KB per block without opt-in), a maximum of 2048 resident threads (64 warps), and a maximum number of resident blocks. A kernel's per-thread register count and per-block shared memory determine how many blocks fit, and occupancy is the resulting fraction of the 2048-thread ceiling. The calculation is integer division three ways, taking the minimum, and it is worked in full in Problem 6. Occupancy matters only insofar as it feeds Little's law, enough resident warps to cover latency. Volkov's classic argument (2010) showed that past that coverage point, lower occupancy is often faster, because fewer resident threads means more registers per thread, which buys instruction-level parallelism and register-blocked data reuse worth more than additional warps. The measured matmul progression below lands exactly on this point. The fastest custom kernel on this page runs at 25 percent occupancy.

Reductions on the GPU, a measured seven-step progression

The parallel reduction is the standard vehicle for learning GPU optimization, following the structure of Harris's NVIDIA reduction notes, and this page's version was written and measured fresh on the H100, summing \(n = 2^{28}\) floats (1.074 GB read), with correctness checked against torch.sum and a median of 30 runs. A reduction reads \(n\) floats and writes one, so its arithmetic intensity is \(1/4\) flop per byte and the only meaningful score is bandwidth. Time is bytes over GB/s, and the finish line is the memcpy rate.

KernelKey changeTime (ms)GB/svs first
reduce0interleaved tree, modulo indexing3.000358.01.0x
reduce1sequential addressing, no divergence/conflicts1.229873.62.4x
reduce2first add during load (halve blocks)0.6581631.54.6x
reduce3warp shuffles replace shared memory0.4782245.66.3x
reduce4grid-stride loop, few blocks, more per thread0.3573009.78.4x
reduce5vectorized float4 loads0.3543034.78.5x
torch.sumlibrary baseline2974.48.3x

Each row is one idea. reduce0 pairs elements at interleaved offsets with a modulo test, so within every warp half the lanes fail the predicate each round (divergence) and the shared-memory access pattern conflicts. It manages 12 percent of memory bandwidth. reduce1 renumbers so that active threads are contiguous, the same tree with no divergence, \(2.4\times\). reduce2 notes that a tree over a block leaves half the threads idle after the first step, so each thread should first add two (or more) global elements while loading. The tree shrinks and the load loop, which is the bandwidth-bound part, does more of the work. reduce3 replaces the last five tree levels with register-to-register warp shuffles. reduce4 is the structural insight worth generalizing. Launch only enough blocks to fill the machine and let each thread loop over a grid-sized stride, so the bandwidth-critical loop is long, fully coalesced (consecutive threads read consecutive addresses each iteration), and amortizes all setup. It reaches 99 percent of the 3034.7 GB/s ceiling that reduce5's float4 loads then touch. The final kernel beats torch.sum by 2 percent, and the whole progression is \(8.5\times\) without changing the algorithm's work or depth at all. Every factor came from the memory system.

Warp shuffles

The __shfl_down_sync(mask, v, d) intrinsic returns lane \(i+d\)'s value of \(v\) to lane \(i\) within a warp, register to register, no shared memory and no barrier, because the warp's lockstep execution is the synchronization. Five shuffle-adds with offsets 16, 8, 4, 2, 1 reduce a warp. The mask argument (0xffffffff for a full warp) exists because post-Volta GPUs execute divergent lanes independently, and the sync suffix forces reconvergence of the named lanes first. The same primitive family (shfl_up for scans, shfl_xor for butterfly exchanges) makes the warp a 32-wide register file visible to the programmer, and warp-level scan plus shuffle reduction is the base case of essentially every CUB block primitive.

Streams, asynchronous copy, and cooperative groups

A CUDA stream is an ordered queue of work. Operations in one stream execute in issue order, operations in different streams have no ordering relative to each other, and cudaEvent objects insert explicit cross-stream dependencies. Streams are what turn a GPU program into a pipeline. Copy the next batch on one stream while computing the current batch on another, and the transfer time disappears under the compute as long as the host memory is page-locked, since only pinned memory can be DMA'd without a staging copy. The same mechanism underlies overlapping the gradient all-reduce with the backward pass in data-parallel training, which is a two-stream pipeline with events at the bucket boundaries. Because kernel launch is asynchronous, a host thread can run far ahead of the device, and this is exactly what hides launch latency. It is also what makes naive timing wrong, since without a synchronization the host measures the enqueue, not the work. CUDA graphs go further and record a whole dependency DAG of launches once, then replay it with a single submission, which matters when kernels are short enough that the per-launch overhead of a few microseconds is a real fraction of the step. The small-batch inference regime is exactly an Amdahl problem in launch overhead.

Inside a kernel, cuda::memcpy_async and the underlying cp.async instruction (Ampere onward) copy global memory into shared memory without routing the data through registers and without blocking the issuing warp, so a kernel can prefetch the next tile while computing the current one. This is software pipelining expressed in the memory system, and it is the difference between the tiled matmul on this page and the library kernels it loses to. Cooperative groups generalize the synchronization vocabulary beyond __syncthreads. tiled_partition<32> names a warp explicitly so shuffle-based reductions can be written generically, coalesced_threads() names the currently active lanes after divergence, and a cooperative launch plus grid_group::sync() gives a genuine grid-wide barrier, at the cost of requiring that the entire grid be co-resident, which caps the grid at what the device can hold and gives up the portability that ordinary kernels enjoy.

Hopper's additions, TMA, thread block clusters, and warpgroup matmul

The H100 measured throughout this page adds three mechanisms that break the classical grid-block-thread abstraction, and current high-performance kernels are organized around them. The Tensor Memory Accelerator is a dedicated unit that performs bulk asynchronous copies between global and shared memory given a multi-dimensional tensor descriptor. One thread issues the copy, and the hardware generates all the addresses, handles boundary clamping, and signals completion through a shared-memory barrier. Compared with cp.async, it removes the address arithmetic from the instruction stream entirely, which matters because in a well-tuned matmul that arithmetic was competing with the actual math for issue slots. Thread block clusters add a level between grid and block. A cluster is a set of blocks (up to 8 on the H100) guaranteed to be co-resident on the same GPU processing cluster, able to synchronize with each other and to read each other's shared memory through a distributed shared memory address space. That is a genuinely new capability, since before it the only inter-block channel was global memory. Warpgroup matrix multiply-accumulate (wgmma) issues a matmul from a warpgroup of four warps, 128 threads, asynchronously, reading its operands directly from shared memory rather than requiring them staged into registers first.

The three combine into a structure that looks nothing like the SPMD kernels earlier on this page, warp specialization, in which some warpgroups act as producers issuing TMA loads into a circular buffer of shared-memory tiles, and others act as consumers issuing wgmma against those tiles, with asynchronous barriers carrying the handoff. It is a dataflow pipeline running inside one thread block. FlashAttention-3 is the canonical published example, and it is where the measured 640.1 TFLOPS at \(L = 16384\) comes from, overlapping the softmax, which runs on the ordinary vector units, with the two matmuls, which run on the tensor cores, so that neither datapath idles. CUTLASS 3 exposes the same structures generically through its CuTe layout algebra. The honest summary for a practitioner is that these features are currently reached through libraries and compilers rather than written by hand. Triton emits wgmma and TMA operations from ordinary tl.dot and tl.load code, which is why the Triton matmul measured on this page reaches 511.6 TFLOPS without a single Hopper-specific line.

Atomics and privatization, the measured histogram

Reductions to many cells, a histogram being the pure case, add contention to the analysis. One measured kernel bins \(2^{28}\) bytes into 256 bins with a global atomicAdd per element. It takes 57.902 ms, because popular bins serialize. The privatized version gives each block a 256-entry shared-memory histogram (1 KB), accumulates locally with shared-memory atomics, then merges 256 values per block into the global histogram. It takes 0.401 ms, a measured \(144.6\times\), running at 670.1 GB/s of input. The general principle is that atomics are cheap when contention is low, so restructure to make contention low. Privatization trades a factor of (blocks) more merge traffic for a factor of (elements per block) less contended traffic, the same shape of trade as the reduction tree itself. The same restructuring appears one level up in distributed systems. Tsinghua's Gemini graph-processing work makes locality and contention, rather than partitioning cleverness, the object of design for exactly this reason, and ETH Zurich's concurrency analysis of distributed deep learning organizes the entire training-parallelism literature around where the contended communication lands.

Matmul, the tiling progression, measured

Arithmetic intensity is a property you choose

An \(n \times n\) matmul does \(2n^3\) flops on \(3n^2\) values (read \(A\), \(B\), write \(C\)), so the algorithm has \(O(n)\) intrinsic flops per byte and can be compute-bound. Whether a given kernel is depends entirely on how much of \(A\) and \(B\) it re-reads. A kernel that loads both operands from global memory for every multiply-add moves 8 bytes per 2 flops, an intensity of \(1/4\), hopelessly memory-bound. Tiling is the act of buying intensity with on-chip memory, and Problem 7 derives the exchange rate. A \(T \times T\) shared-memory tile yields \(T/4\) flops per global byte in fp32. The measured progression below is all fp32 at \(n = 4096\) (137.4 GFLOP per multiply), each kernel verified against cuBLAS output.

KernelTime (ms)TFLOPSptxas regs/threadsmem/blockOccupancy
naive (1 output/thread, global only)23.485.9320100%
tiled32 (32x32 smem tiles)17.048.1328192 B100%
blocked128 (128x128 block tile, 8x8 register tile/thread)5.4825.11128192 B25%
cuBLAS SGEMM2.6851.2library (register double-buffering, wider tiles)
cuBLAS fp16 (tensor cores)0.19702.4wgmma tensor-core path, same tiling logic

The naive kernel is instructive precisely because it is not as slow as pure arithmetic-intensity accounting predicts. At intensity \(1/4\), 3063.5 GB/s of HBM would cap it at 0.77 TFLOPS, yet it reaches 5.9 because adjacent threads read the same row of \(A\) and consecutive elements of \(B\), and the L1 and L2 caches convert that incidental locality into reuse. Hardware caches are the reason lower bounds need care. They are not a substitute for tiling, just a discount on its absence. tiled32 makes the reuse explicit, 32 flops per global byte pair loaded, but only improves \(1.4\times\). With one output per thread, each thread issues two shared-memory reads per FMA, and shared-memory bandwidth, not HBM, becomes the binding resource. This is the step most tutorials stop at and the reason their matmuls stall near 10 percent of peak.

blocked128 attacks the shared-memory bound the same way tiled32 attacked the HBM bound, one level down the hierarchy. Each thread computes an \(8 \times 8\) register tile of \(C\), so each value fetched from shared memory into a register participates in 8 FMAs. Per iteration a thread reads 8+8 shared floats and performs 64 FMAs, 4 flops per shared-memory byte instead of \(1/4\). The cost is registers, 112 per thread by ptxas count, which caps residency at 2 blocks (512 threads) per SM, 25 percent occupancy, and the kernel is \(3.1\times\) faster than the full-occupancy tiled32 anyway, Volkov's argument made concrete. The remaining \(2\times\) to cuBLAS SGEMM is double-buffering (prefetching the next tile into registers while computing the current one, hiding shared-memory latency), wider and asymmetric tiles, and vectorized 128-bit loads. The further \(13.7\times\) from cuBLAS fp32 to fp16 is not software at all but the tensor-core datapath, which exists precisely because matmul's buyable intensity is high enough to justify dedicated silicon.

per-thread work at each level (one k-slice iteration)

 tiled32:    smem As[ty][k], Bs[k][tx] ──► 1 FMA      (2 smem reads / 2 flops)
 blocked128: regs a[0..7] ◄── smem As    8 reads
             regs b[0..7] ◄── smem Bs    8 reads
             acc[i][j] += a[i]*b[j]      64 FMAs      (16 smem reads / 128 flops)

 HBM traffic        /32  (tile reuse)      ── tiled32 buys this
 smem traffic       /8   (register reuse)  ── blocked128 buys this
 issue slots        /2   (tensor core mma) ── hardware buys this

The teaching-kernel gap, stated honestly

At the large square size where the library is at its best, \(n = 8192\), the measured cuBLAS numbers on this H100 are 51.4 TFLOPS in fp32, 409.7 in tf32, and 728.7 in bf16, against 700.8 in fp16. Set the page's best hand-written kernel, 25.1 fp32 TFLOPS at \(n = 4096\), beside them. It is roughly half of cuBLAS in the same precision and one twenty-ninth of what the same chip does in bf16. That gap is not a failure of the derivation, and pretending otherwise is the most common dishonesty in GPU tutorials. It decomposes into four pieces, each identifiable.

First, datapath, worth \(8\times\) to \(14\times\) and unavailable to any kernel written in scalar FMA form. The tf32 and bf16 columns are the tensor cores, reached only through mma/wgmma instructions or a compiler that emits them, which is why the Triton matmul on this page reaches 511.6 TFLOPS from ordinary tl.dot while the CUDA kernel cannot. Second, latency hiding, worth roughly \(2\times\), meaning double-buffered shared-memory tiles and cp.async or TMA prefetch, so the FMA pipeline never waits on a load. Third, shape specialization. cuBLAS carries hundreds of kernels and selects tile shapes, split-k factors, and epilogue fusions per problem size at runtime, which is why its own throughput swings from 30.6 TFLOPS at \(n = 1024\) to 51.4 at \(n = 8192\) in fp32 and from 108.9 to 728.7 in bf16, a \(6.7\times\) range on the same operation. Fourth, instruction economy. Address arithmetic, predication, and bounds checks compete with the math for issue slots, and the library's generated code amortizes them across much larger register tiles.

The useful conclusion is not that hand-written matmuls are pointless but that matmul specifically is the wrong place to spend the effort, because the shape is fixed, the libraries are excellent, and the datapath is behind an instruction the derivation does not reach. The value of walking the tiling progression is that every other kernel on this page is a shape the library does not have, and the reasoning transfers exactly. Buy intensity with on-chip memory, then buy it again with registers, then check whether the bottleneck moved. That is how FlashAttention was derived, and there the same reasoning was worth \(27.7\times\) because no library kernel existed to lose to.

The roofline model, calibrated on this machine

The roofline (Williams, Waterman, and Patterson, 2009) reduces a machine to two numbers and a kernel to one. For a kernel with arithmetic intensity \(I\) flops per byte of memory traffic, on a machine with peak compute \(P\) and memory bandwidth \(B\),

$$ \text{attainable flops/s} = \min\big(P,\, I \times B\big). $$

The derivation is a one-line application of the bottleneck principle. The kernel needs \(F\) flops and \(F/I\) bytes, and with perfect overlap of compute and transfer, time is \(\max(F/P,\, F/(IB))\), which inverts to the throughput above. The ridge point \(I^* = P/B\) is the machine balance, the intensity at which the roofs meet. Using this page's measured numbers rather than datasheet values, \(B = 3063.5\) GB/s (streaming add), and \(P\) depends on the datapath, which is the H100's defining feature. For bf16 tensor cores, \(P = 744.6\) measured TFLOPS, so \(I^* = 744.6 \times 10^{12} / 3063.5 \times 10^9 = 243\) flops per byte. For fp32 on the vector units, cuBLAS's 51.2 TFLOPS gives \(I^* = 16.7\). One machine, two ridges an order of magnitude apart, so "is this kernel compute-bound" has no answer until the precision and datapath are named.

log-log roofline, this H100 (measured)

 744.6 TF ┤ bf16 tensor ─────────────────────────●━━━━━━━━━━
          │                                  ridge I*=243
  51.2 TF ┤ fp32 ────────────●━━━━━━━━━━━━ (I*=16.7)
          │                ╱   ╲ sgemm fp32 n=4096: I=341, bound by P
          │              ╱
          │  slope = 3063.5 GB/s
 255 GF   ┤    ╱● vadd I=1/12          measured 232 GF ─ 91% of roof
 766 GF   ┤   ╱ ● reduce I=1/4         measured 758 GF ─ 99% of roof
          └──┴────┴────┴────┴────┴────┴────┴──── I (flops/byte)
            1/12  1/4   1    8   16.7  32   243

Reading kernels onto this chart is the daily use. The elementwise and reduction kernels of this page live at \(I \le 1/4\), three orders of magnitude left of the bf16 ridge. No code change short of fusion moves them, and their only honest metric is fraction of \(B\). The measured fusion experiment makes the point. A four-op elementwise chain launched as separate PyTorch kernels takes 0.911 ms, and the same chain compiled into one kernel takes 0.214 ms, a \(4.26\times\) speedup that is purely the roofline's memory term, four round trips of intermediates collapsed into one. Matmul at \(n = 4096\) has intrinsic intensity \(I = 2n^3 / (3n^2 \times 4) = n/6 = 683\) in fp32 accounting (341 if counting 8 bytes per element in and out), far right of the fp32 ridge. It is compute-bound, and indeed cuBLAS hits its fp32 peak. Attention, next section, is the interesting case, because its intensity is a design choice.

Classifying every kernel measured on this page against the two ridges, \(I^*_{\text{fp32}} = 16.7\) and \(I^*_{\text{bf16}} = 243\) flops per byte, gives the whole page as one table. Intensity is computed from first-principles byte and flop counts, not from a profiler.

KernelArithmetic intensityBound byMeasuredFraction of its roof
vector add1/12 flop/bytememory2785.3 GB/s91% of 3063.5
reduction (reduce5)1/4memory3034.7 GB/s99%
fused elementwise chain~1memory2616.6 GB/s85%
transpose (tiled, padded)0 (pure movement)memory2694.0 GB/s95% of copy
histogram (privatized)~1/1 on 1-byte inputmemory670.1 GB/s22%, atomics-limited
softmax (fused, fp32)~5 (exp counted as 1)memory1480.3 GB/s48%
LayerNorm (Welford, fp32)~5memory1928.6 GB/s63%
RMSNorm (fp32)~4memory2137.6 GB/s70%
scan (three-phase)1/8 over 4 traversalsmemory1324.9 GB/s43%
stream compaction0memory, scattered writes481.2 GB/s16%
SGEMM fp32 n=4096 (naive)1/4 realizedmemory + issue5.9 TFLOPSsee text
SGEMM fp32 (blocked128)32 realizedcompute (smem/issue)25.1 TFLOPS49% of 51.2
cuBLAS SGEMM n=8192683compute51.4 TFLOPS100%
naive attention L=8192~1 realizedmemory22.2 TFLOPSat the memory roof
fused attention L=16384~L/4, unboundedcompute (tensor core)640.1 TFLOPS86% of 744.6

Two rows repay a second look. The histogram at 670.1 GB/s is memory-bound in form but atomics-bound in fact. Its input is one byte per element and its traffic ought to run at streaming rate, but shared-memory atomic throughput on 256 bins caps it well short. The normalization kernels sit at 48 to 70 percent of the streaming roof because they must read each row twice, once to compute statistics and once to apply them, unless the row fits in registers. The measured RMSNorm is fastest of the three because it computes one statistic instead of two and skips the bias. Every entry in the table except the last two is memory-bound, which is the empirical form of the claim that modern accelerators are memory machines with arithmetic attached.

Attention, the kernel-design capstone

Why naive attention is memory-bound at any scale

Single-head attention computes, for \(Q, K, V \in \R^{L \times d}\),

$$ O = \softmax\!\Big(\frac{QK^{\T}}{\sqrt{d}}\Big) V , $$

which is \(4L^2d\) flops (two \(L \times L \times d\) matmuls) on \(O(Ld)\) input bytes, a very high intrinsic intensity. The naive implementation throws that intensity away by materializing the \(L \times L\) score matrix \(S\) and the probability matrix \(P\) in HBM. It writes \(S\), reads it for softmax, writes \(P\), and reads \(P\) for the second matmul. Those four passes cost \(4 \times 2L^2\) bytes in half precision, and since \(L \gg d\), realized intensity collapses to \(O(d)\) divided by a constant, pinning the kernel to the memory roof while the quadratic memory footprint grows until it does not fit at all. The measured numbers below (batch \(\times\) heads \(= 128\), \(d = 64\), half precision, H100 80GB) show naive attention stuck at 22 to 26 TFLOPS, 3 percent of the machine's measured bf16 peak, at every length until it runs out of memory, while the fused kernel gains throughput with length as its inner loops grow.

Lnaive msnaive TFLOPSnaive peak memflash msflash TFLOPSflash peak memspeedup
5120.37123.20.20 GB0.050170.20.07 GB7.42x
10241.29726.50.64 GB0.110313.10.10 GB11.79x
20485.30925.92.32 GB0.299460.00.17 GB17.76x
409623.07023.88.91 GB0.976563.00.30 GB23.64x
819298.94622.235.00 GB3.575615.10.57 GB27.68x
16384OOM, the 68.72 GB score matrix does not fit13.742640.11.12 GB

Online softmax, the enabling identity

Fusing attention requires computing softmax over a row of scores without having the whole row. Softmax needs the row max \(m\) for numerical stability and the normalizer \(\ell = \sum_j e^{x_j - m}\). Milakov and Gimelshein (2018) showed both can be maintained in one pass. Process the row in chunks. After some chunks, hold the running pair \((m, \ell)\). A new chunk \(c\) with local max \(\tilde m\) and local sum \(\tilde\ell = \sum_{j \in c} e^{x_j - \tilde m}\) merges via the new max \(m' = \max(m, \tilde m)\), giving

$$ \ell' = \ell\, e^{\,m - m'} + \tilde\ell\, e^{\,\tilde m - m'} , $$

which is exact, because multiplying by \(e^{m - m'}\) rebases every previously accumulated exponential from the old max to the new one, since \(e^{x_j - m} \cdot e^{m - m'} = e^{x_j - m'}\). FlashAttention (Dao, Fu, Ermon, Rudra, and Ré, 2022) extends the same rebase to the output accumulator. Maintain \(o = \sum_j e^{x_j - m}\, v_j\) unnormalized, scale the old \(o\) by \(e^{m - m'}\) when the max moves, add the new chunk's \(\tilde P V\) contribution, and divide by \(\ell\) once at the end. Every intermediate lives in registers and shared memory, so \(S\) and \(P\) never touch HBM. The measured fused softmax on this machine shows the standalone win. An online-softmax CUDA kernel over an 8192x8192 matrix runs one read and one write, 0.363 ms at 1480.3 GB/s, versus 0.863 ms for the unfused three-pass version, with torch.softmax at 0.301 ms confirming the library does the same fusion.

FlashAttention, the tiling and rescaling, derived in full

The online-softmax identity handles one row's normalizer. FlashAttention's contribution is to carry the same rebase through the value accumulation so that an entire attention row can be finished without ever holding the row. Fix one query row \(q \in \R^{d}\) and partition the \(L\) keys into blocks \(B_1, \dots, B_T\) of \(B_c\) keys each. Define, after processing blocks \(1..t\), the three state variables

$$ m^{(t)} = \max_{j \in B_1 \cup \cdots \cup B_t} s_j , \qquad \ell^{(t)} = \!\!\sum_{j \in B_1 \cup \cdots \cup B_t}\!\! e^{\,s_j - m^{(t)}} , \qquad o^{(t)} = \!\!\sum_{j \in B_1 \cup \cdots \cup B_t}\!\! e^{\,s_j - m^{(t)}} v_j , $$

where \(s_j = q^{\T}k_j/\sqrt{d}\). The claim is that these can be updated from block \(t+1\) alone. Let \(\tilde m = \max_{j \in B_{t+1}} s_j\), \(m' = \max(m^{(t)}, \tilde m)\), and write \(\alpha = e^{\,m^{(t)} - m'}\), \(\beta = e^{\,\tilde m - m'}\). Then, splitting the sum at the block boundary and multiplying each part by one in the form \(e^{m - m'}e^{m' - m}\),

$$ \ell^{(t+1)} = \sum_{j \le B_t} e^{\,s_j - m'} + \sum_{j \in B_{t+1}} e^{\,s_j - m'} = \alpha \underbrace{\sum_{j \le B_t} e^{\,s_j - m^{(t)}}}_{\ell^{(t)}} + \beta \underbrace{\sum_{j \in B_{t+1}} e^{\,s_j - \tilde m}}_{\tilde\ell} , $$ $$ o^{(t+1)} = \alpha\, o^{(t)} + \beta \sum_{j \in B_{t+1}} e^{\,s_j - \tilde m} v_j = \alpha\, o^{(t)} + \beta\, \tilde P_{t+1} V_{t+1} , $$

because \(e^{s_j - m'} = e^{s_j - m^{(t)}} e^{m^{(t)} - m'}\) is an exact identity for every \(j\) in the first group and correspondingly for the second. Nothing is approximated. The rescaling is multiplication by one, written in a form that changes which max the accumulator is referenced to. After the last block, \(o = o^{(T)}/\ell^{(T)}\) is exactly \(\softmax(s)^{\T}V\). The numerical property that makes it usable is that every exponential argument \(s_j - m'\) is at most zero, so no term can overflow, and the largest is exactly zero, so no term underflows to insignificance relative to the largest.

flash forward, one query block Q_i of B_r rows, T key blocks

 load Q_i  (B_r x d)  into registers/smem       once, stays resident
 m = -inf (B_r), l = 0 (B_r), O = 0 (B_r x d)   online-softmax state
 for t = 1..T:
     load K_t, V_t  (B_c x d) into shared memory        <-- the only HBM reads
     S  = Q_i K_t^T * scale            (B_r x B_c)      tensor cores, on chip
     m' = max(m, rowmax(S))                              merge maxima
     P  = exp(S - m')                  (B_r x B_c)       never leaves the SM
     a  = exp(m - m')                                    rebase factor
     l  = a*l + rowsum(P);  O = a*O + P V_t              tensor cores again
     m  = m'
 store O / l  (B_r x d)                                 <-- the only HBM write

 HBM traffic per query block: d*B_r (read Q) + 2*L*d (stream K,V) + d*B_r (write O)
 HBM traffic materialized:    the same, plus 4 * L * B_r for S and P

The traffic arithmetic makes the measured table exact rather than suggestive. With \(L/B_r\) query blocks each streaming all of \(K\) and \(V\), total HBM traffic is \(\Theta\!\left(\frac{L}{B_r}\cdot Ld\right) = \Theta(L^2 d / B_r)\), and since the shared-memory budget \(M\) must hold a \(K\) tile, a \(V\) tile, and the \(Q\) tile, \(B_r\) and \(B_c\) scale like \(M/d\), giving the paper's \(\Theta(L^2d^2/M)\) against the materialized algorithm's \(\Theta(L^2 + Ld)\). The ratio is \(d^2/M\), a factor of tens at \(d = 64\) and \(M \approx 100\) KB. Concretely, at the measured configuration (\(BH = 128\), \(L = 8192\), \(d = 64\), 2 bytes per element), streaming \(Q, K, V, O\) once costs \(4 \times 128 \times 8192 \times 64 \times 2 = 0.54\) GB, while the two score tensors cost \(4 \times 128 \times 8192^2 \times 2 = 68.7\) GB, a \(127\times\) difference in bytes for identical arithmetic. At \(3063.5\) GB/s those are 0.18 ms and 22.4 ms of unavoidable traffic respectively, and the measured 3.575 ms fused against 98.946 ms naive sits above both floors by the same factor, as it must, because the fused kernel is compute-bound at 615.1 TFLOPS while the naive one is traffic-bound at 22.2.

The backward pass follows by storing only \((m, \ell)\) per row (2 floats) and recomputing \(S\) tiles from \(Q\) and \(K\) during the gradient computation. Recomputation is cheaper than the memory traffic it replaces, the signature trade of the whole post-2022 kernel era. FlashAttention-2 (Dao, 2023) reorganized the loops so the outer loop over \(Q\) blocks is the parallel grid dimension and warps split the head dimension rather than the sequence, cutting non-matmul overhead and inter-warp communication. FlashAttention-3 (Shah, Bikshandi, Zhang, Thakkar, Ramani, and Dao, 2024) rebuilt it around Hopper's asynchronous TMA copies and warpgroup matmuls with software pipelining between the two matmuls and the softmax.

Two smaller measured points complete the picture. A minimal fp32 CUDA flash-style kernel written for this page (one block per query tile, 167 registers per thread, 32 KB shared memory by ptxas) reaches 9.6 TFLOPS versus 16.4 for PyTorch's fp32 SDPA. It is correct and \(O(L)\) in memory, but far from the tensor cores. The Triton version at fp16, 60 lines in the implementation section below, reaches 340.2 TFLOPS on batch \(\times\) heads \(= 64\), \(L = 4096\), \(d = 64\), beating that configuration's SDPA baseline at 297.0. A tile-level language with a good compiler is enough to be production-competitive on this kernel, which is the argument Triton's designers made.

Multi-head attention as a whole block

Attention in a real model is never the bare kernel. A multi-head attention block takes \(X \in \R^{B \times L \times d_{\text{model}}}\), applies three projections (usually fused into one \(d_{\text{model}} \times 3d_{\text{model}}\) matmul), reshapes into \(H\) heads of width \(d_h = d_{\text{model}}/H\), runs attention per head, concatenates, and applies an output projection. The parallel structure is worth stating precisely because it is the reason the transformer trains at all. The head axis is pure data parallelism with no communication, the batch axis likewise, and the only sequential dependence is within a head's softmax. That means the natural grid is \((L/B_r, B \cdot H)\), exactly the two-dimensional launch the Triton kernel below uses, and it is why attention scales to \(10^5\)-way parallelism without any algorithmic cleverness.

The flop accounting splits into two pieces with different intensities. The projections cost \(2BLd_{\text{model}}(3d_{\text{model}}) + 2BLd_{\text{model}}^2 = 8BLd_{\text{model}}^2\) flops and are ordinary matmuls, compute-bound and handled well by cuBLAS. The attention itself costs \(4BHL^2d_h = 4BL^2 d_{\text{model}}\) flops, halved by causal masking. Their ratio is \(2d_{\text{model}}/L\) (before the causal factor), so the crossover where attention starts to dominate the block is \(L \approx 2 d_{\text{model}}\). Below it the projections are the cost, above it attention is, and this single ratio explains why long-context work is dominated by attention kernel engineering while short-context serving is dominated by GEMM shapes.

The measurements below are for a full causal block on this H100 at \(B = 8\), \(L = 4096\), \(d_{\text{model}} = 1024\), \(H = 16\) heads of width 64, bf16, including both projections and verified to a relative error of \(3.7 \times 10^{-3}\) between the two paths.

Multi-head attention block, causalTime (ms)TFLOPSPeak memory
naive, materialized (B,H,L,L) scores + masked_fill + softmax24.04422.99.29 GB
fused, same projections, flash SDPA backend1.352406.60.68 GB
ratio17.78x17.8x13.7x

The block does 274.9 GFLOP of projection work and 274.9 GFLOP of causal attention work, an even split at exactly the \(L = 2d_{\text{model}}\) crossover derived above. That split is what makes the fused number instructive. At 406.6 TFLOPS the block is below the 728.7 TFLOPS that a pure bf16 matmul reaches on this machine, and the reason is not inefficiency in either half but the mix, since the causal kernel wastes half its tiles on masked-out regions unless it skips them, and the projections at \(L d_{\text{model}} = 4096 \times 1024\) are a smaller, less efficient GEMM shape than the \(n = 8192\) square case. The naive path's 22.9 TFLOPS matches the standalone naive attention numbers almost exactly, confirming that in that path the projections are invisible. Essentially all of the 24 ms is spent writing and reading the \(8 \times 16 \times 4096^2 \times 2 = 8.6\) GB score tensor, plus the masking pass over it, plus the softmax pass over it.

A separate measurement isolates the causal case for the hand-written kernel. A causal Triton flash kernel written for this page, which skips key blocks entirely beyond the diagonal, runs at \(BH = 64\), \(L = 4096\), \(d = 64\), fp16 in 0.577 ms and 238.3 TFLOPS, against PyTorch's flash backend at 0.533 ms and 257.9 TFLOPS, with a maximum absolute error of \(4.9 \times 10^{-4}\). Ninety-two percent of a production kernel from forty lines of Triton is the headline, but the more useful observation is the shape of the remaining gap. The loop bound \(\lceil (i+1)B_r / B_c \rceil\) skips whole blocks but still computes the full \(B_r \times B_c\) tile on the diagonal, wasting slightly under half of one block-row, and production kernels additionally balance the triangular work across the grid so that programs handling early query blocks do not finish long before those handling late ones. Load imbalance in a triangular loop is a decomposition problem, and it is the same problem the assignment discussion opened with.

Worked problems

Problem 6

The blocked128 matmul kernel compiles to 112 registers per thread and 8192 bytes of shared memory per block, launched with 256 threads per block. For an H100 SM (65,536 registers, 228 KB shared memory available to the kernel, 2048-thread and 32-block ceilings), compute the occupancy and identify the limiting resource. Then verify the shared memory figure from the tile shape, a block tile of \(128 \times 128\), \(k\)-slice depth 8, fp32.

Solution. Registers give \(112 \times 32 = 3584\) per warp, and \(\lfloor 65536 / 3584 \rfloor = 18\) warps fit by registers. A block is \(256/32 = 8\) warps, so \(\lfloor 18/8 \rfloor = 2\) blocks fit, 16 warps, 512 threads. Shared memory allows \(\lfloor 228 \times 1024 / 8192 \rfloor = 28\) blocks, not binding. The thread ceiling allows \(2048/256 = 8\) blocks, not binding. Occupancy \(= 512 / 2048 = 25\) percent, limited by registers.

For the shared memory check, per k-slice the block stages a \(128 \times 8\) panel of \(A\) and an \(8 \times 128\) panel of \(B\), so \((128 \times 8 + 8 \times 128) \times 4 = 2048 \times 4 = 8192\) bytes, matching ptxas. For the work check, \(128 \times 128 = 16384\) outputs over 256 threads is 64 outputs per thread, the \(8 \times 8\) register tile, which with two 8-wide operand vectors requires on the order of \(64 + 16\) accumulator and operand registers before addressing arithmetic, hence the 112-register count. This kernel at 25 percent occupancy measures 25.1 TFLOPS against 8.1 for the 100-percent-occupancy tiled32. Occupancy is a constraint, not an objective.

Problem 7

Derive the arithmetic intensity of a \(T \times T\) shared-memory-tiled fp32 matmul as a function of \(T\) (counting global memory traffic only), then use the roofline with this machine's measured \(B = 3063.5\) GB/s to predict the memory-roof ceiling for \(T = 32\) and explain why the measured tiled32 kernel (8.1 TFLOPS) does not reach it. Finally, compute the minimum time for the \(2^{28}\)-element vector add and compare with the measured 1.157 ms.

Solution. Each block computes a \(T \times T\) tile of \(C\) by marching over \(n/T\) k-slices, loading one \(T \times T\) tile of \(A\) and one of \(B\) per slice. Global loads per block are \(2 T^2 \cdot (n/T) = 2nT\) floats. Blocks number \((n/T)^2\), so total traffic is \((n/T)^2 \cdot 2nT \cdot 4 = 8n^3/T\) bytes, against \(2n^3\) flops, so \(I(T) = 2n^3 / (8n^3/T) = T/4\) flops per byte. Doubling the tile doubles intensity, which is the entire logic of tiling in one formula. For \(T = 32\), \(I = 8\), and the memory roof is \(8 \times 3063.5 = 24.5\) TFLOPS.

The measured 8.1 TFLOPS is a third of that ceiling, so HBM is not the binding constraint. The kernel has moved its bottleneck on-chip. At one output per thread it issues two 4-byte shared-memory reads per FMA, demanding 8 bytes per 2 flops from a shared-memory system whose bandwidth per SM is a small multiple of the FMA rate. The register-tiled kernel exists to fix exactly this, and its measured 25.1 TFLOPS sits above tiled32's shared-memory bound but below the cuBLAS 51.2, which needs double-buffering to cover latency. Rooflines nest. HBM, L2, shared memory, and registers each have one.

The vector add moves \(3 \times 4 \times 2^{28} = 3.221\) GB of traffic (two reads, one write), and at 3063.5 GB/s the floor is \(3.221/3.0635 = 1.051\) ms. The measured 1.157 ms, i.e. 2785.3 GB/s, is 91 percent of the streaming ceiling. The gap is write-allocate and DRAM page effects on the mixed read/write stream. The measured Triton version of the same kernel (1.078 ms, 2988.4 GB/s) and torch.add (3059.2 GB/s) bracket the same floor, confirming the model. For \(I = 1/12\), every correct implementation lands on the memory roof and nothing else matters.

Problem 8

For the naive attention configuration measured above (batch \(\times\) heads \(= 128\), \(L = 8192\), \(d = 64\), 2-byte elements), (a) compute the size of the materialized score matrix and check it against the measured 35.00 GB peak memory, (b) lower-bound the HBM traffic of the naive implementation and convert it to a time floor at the measured 3063.5 GB/s, and (c) compare with the measured naive (98.946 ms) and flash (3.575 ms) times and state the conclusion.

Solution. (a) One score matrix is \(128 \times 8192^2\) elements \(\times\) 2 bytes \(= 128 \times 67{,}108{,}864 \times 2 = 17.18\) GB. Naive attention holds \(S\) and \(P\) simultaneously (softmax cannot safely overwrite in place while both are live across kernel boundaries), so \(2 \times 17.18 = 34.36\) GB, plus \(Q, K, V, O\) at \(4 \times 128 \times 8192 \times 64 \times 2\) bytes \(= 0.54\) GB, giving 34.90 GB, matching the measured 35.00 GB peak within allocator rounding.

(b) The score tensors must be written once and read once each. \(S\) is written by the first matmul and read by softmax, and \(P\) is written by softmax and read by the second matmul. That is \(4 \times 17.18 = 68.7\) GB minimum, ignoring \(Q,K,V\) entirely. The time floor is \(68.7 / 3063.5\) GB/s \(= 22.4\) ms.

(c) The measured naive time, 98.946 ms, is \(4.4\times\) the floor (softmax makes extra passes for max and sum, and the matmuls' own tiling re-reads operands), while flash at 3.575 ms is \(6.3\times\) below the naive floor. No schedule of the materializing algorithm, however perfect, can come within a factor of six of the fused kernel on this machine. The 2.199 Pflop of work is identical in both. The entire \(27.68\times\) measured speedup is memory traffic. This is the cleanest available demonstration that on modern accelerators the algorithm's IO complexity, not its flop count, is the object of design.

Problem 9

Eight threads on a 64-byte-cache-line machine increment eight int32 counters stored in one array. (a) How many cache lines do the counters occupy, and how many coherence transactions does one round of eight increments cost? (b) The measured run is 4.245 s for \(8 \times 5 \times 10^7\) increments unpadded and 0.304 s padded. Convert both to nanoseconds per increment and infer the cost of the coherence round trip. (c) A colleague proposes padding to 16 bytes instead of 64 to save memory. Evaluate.

Solution. (a) Eight int32 counters are 32 bytes, so they occupy one 64-byte line (assuming they do not straddle a boundary). Under MESI, a write requires the line in Modified state exclusively, so every increment by a different core issues a read-for-ownership that invalidates the other seven copies, eight transactions per round, each a full line transfer, when the logical sharing is zero. Padded to 64 bytes each, the counters occupy eight lines, every core keeps its own line in Modified state permanently, and the steady-state coherence traffic is zero.

(b) Total increments \(= 8 \times 5 \times 10^7 = 4 \times 10^8\). Unpadded, \(4.245\,\text{s} / 4\times10^8 = 10.6\) ns per increment. Padded, \(0.304 / 4\times10^8 = 0.76\) ns. The difference, 9.85 ns, is the amortized coherence cost per increment. Since the threads contend on one line, the eight cores serialize on it, so the per-transaction latency is roughly \(8 \times 9.85 = 79\) ns of line ownership transfer, consistent with a last-level-cache or mesh round trip on a large server part. The padded 0.76 ns is close to one L1-resident atomic per cycle at \(\sim\!2.6\) GHz, which is the correct floor.

(c) It fails on this machine and may pass on another, which is the worst outcome. Sixteen-byte padding puts four counters per 64-byte line, so sharing is reduced fourfold, not eliminated. The expected time lands near \(0.76 + 9.85/ (8/4) \approx 5.7\) ns, most of the penalty intact. Worse, hardware prefetchers on several x86 generations fetch adjacent line pairs, making the effective coherence granularity 128 bytes, so even 64-byte padding can under-deliver. The portable answer is alignas(std::hardware_destructive_interference_size), and the structural answer is better. Give each thread a private accumulator in a register or on its own stack and combine once at the end, which is what OpenMP's reduction clause and this page's per-block privatized histogram both do. Padding treats the symptom, while privatization removes the sharing.

Problem 10

A causal multi-head attention block has \(B = 8\), \(L = 4096\), \(d_{\text{model}} = 1024\), \(H = 16\), bf16. (a) Compute the projection and attention flops separately and confirm the measured 274.9 GFLOP each. (b) Compute the score-tensor bytes the naive path must write and read, and the resulting time floor at 3063.5 GB/s. (c) The measured naive time is 24.044 ms and the fused time is 1.352 ms. Place both on the roofline and say what a further optimization would have to change.

Solution. (a) For the projections, the fused QKV matmul is \(2 \cdot BL \cdot d_{\text{model}} \cdot 3d_{\text{model}} = 2 \times (8 \times 4096) \times 1024 \times 3072 = 2.061 \times 10^{11}\) flops, and the output projection is \(2 \times 32768 \times 1024 \times 1024 = 6.87 \times 10^{10}\). Total \(2.749 \times 10^{11} = 274.9\) GFLOP. The attention is \(4BHL^2d_h \times \tfrac{1}{2}\) for causal \(= 4 \times 8 \times 16 \times 4096^2 \times 64 \times 0.5 = 2.749 \times 10^{11} = 274.9\) GFLOP. They match, which is the \(L = 2d_{\text{model}}\) crossover. Here \(2 \times 1024 = 2048\), and with the causal halving the balance point doubles to \(L = 4096\), exactly this configuration.

(b) The score tensor is \(B \times H \times L \times L \times 2\) bytes \(= 8 \times 16 \times 4096^2 \times 2 = 4.295\) GB. The naive path writes \(S\), reads it to apply the mask, writes the masked copy, reads it for softmax, writes \(P\), and reads \(P\) for the second matmul. Even assuming the mask fuses into the softmax, that is four traversals, \(4 \times 4.295 = 17.18\) GB, a floor of \(17.18/3.0635 = 5.61\) ms. Measured 24.044 ms is \(4.3\times\) the floor, the extra coming from the separate masked_fill pass, the softmax's own multiple passes, and the fact that the two batched matmuls re-read their operands during tiling.

(c) Total work is \(549.8\) GFLOP. The naive path runs at \(549.8\,\text{GFLOP} / 24.044\, \text{ms} = 22.9\) TFLOPS, which is 3.1 percent of the measured 728.7 bf16 peak and sits on the memory roof. The fused path runs at \(549.8/1.352 = 406.6\) TFLOPS, 56 percent of peak, on the compute roof. Because the fused block is compute-bound, further memory optimization is worthless. The remaining headroom is in the arithmetic itself. Three levers exist and all are algorithmic rather than IO. Skip the fully-masked upper-triangular tiles rather than computing and discarding them, which is worth up to \(2\times\) on the attention half. Use a larger or better-shaped GEMM for the projections, since \(32768 \times 1024 \times 3072\) is a less efficient shape than a square \(n = 8192\). And overlap the softmax's vector-unit work with the tensor cores, which is what FlashAttention-3's warp specialization does. The roofline does not tell you what to do next. It tells you which half of the machine to look at.

Implementation

Every kernel in this section is a working program measured for this page on the H100 80GB (PyTorch 2.7, CUDA 12.8, Triton 3.3). The quoted numbers are medians of 30 timed runs after warmup, with correctness checked against the PyTorch reference. The CUDA and Triton tabs implement the same algorithm at two abstraction levels. CUDA states the thread's view and manages shared memory by hand, while Triton states the tile's view and lets the compiler place data and schedule lanes.

Vector add, the bandwidth baseline

The measured numbers are CUDA one-thread-per-element at 1.157 ms (2785.3 GB/s), CUDA grid-stride at 1.179 ms (2733.3 GB/s), Triton at 1.078 ms (2988.4 GB/s), and torch.add at 3059.2 GB/s. Any kernel this simple that does not land within about 10 percent of memcpy has a launch-configuration or coalescing bug, which is what makes it the right first benchmark on new hardware.

// vadd.cu: c = a + b over n = 2^28 floats. Measured: 1.157 ms, 2785.3 GB/s.
// ptxas: 12 registers, 0 bytes smem. AI = 1 flop / 12 bytes: pure bandwidth.
#include <cuda_runtime.h>

__global__ void vadd(const float* __restrict__ a,
                     const float* __restrict__ b,
                     float* __restrict__ c, int n) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;   // global thread id
    if (i < n)                                       // tail guard
        c[i] = a[i] + b[i];                          // warp reads 128 B of a,
}                                                    // 128 B of b: coalesced

// Grid-stride variant: launch ~2 blocks/SM and loop. Same bandwidth here
// (measured 1.179 ms) but composes with any n and amortizes setup.
__global__ void vadd_gridstride(const float* __restrict__ a,
                                const float* __restrict__ b,
                                float* __restrict__ c, int n) {
    int stride = gridDim.x * blockDim.x;
    for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n; i += stride)
        c[i] = a[i] + b[i];   // consecutive lanes -> consecutive addresses
}

int main() {
    int n = 1 << 28;
    float *a, *b, *c;
    cudaMalloc(&a, n * sizeof(float));
    cudaMalloc(&b, n * sizeof(float));
    cudaMalloc(&c, n * sizeof(float));
    int block = 256;
    vadd<<<(n + block - 1) / block, block>>>(a, b, c, n);
    vadd_gridstride<<<132 * 2, block>>>(a, b, c, n);
    return cudaDeviceSynchronize() != cudaSuccess;
}
# vadd in Triton. Measured on the H100: 1.078 ms, 2988.4 GB/s.
# One "program" owns a BLOCK-sized tile; the compiler maps it to warps.
import torch
import triton
import triton.language as tl

@triton.jit
def vadd_kernel(a_ptr, b_ptr, c_ptr, n, BLOCK: tl.constexpr):
    pid = tl.program_id(axis=0)
    offs = pid * BLOCK + tl.arange(0, BLOCK)   # (BLOCK,) tile of indices
    mask = offs < n                            # tail guard, vectorized
    a = tl.load(a_ptr + offs, mask=mask)       # coalescing is the compiler's
    b = tl.load(b_ptr + offs, mask=mask)       # job: contiguous tile ->
    tl.store(c_ptr + offs, a + b, mask=mask)   # contiguous sectors

def vadd(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
    c = torch.empty_like(a)                    # (n,)
    n = a.numel()
    grid = (triton.cdiv(n, 1024),)
    vadd_kernel[grid](a, b, c, n, BLOCK=1024)
    return c

x = torch.randn(1 << 28, device="cuda")        # (2^28,) = 1 GiB
y = torch.randn(1 << 28, device="cuda")
assert torch.equal(vadd(x, y), x + y)

Reduction, grid-stride loads and warp shuffles

This is the reduce4 kernel from the measured progression, 0.357 ms and 3009.7 GB/s on \(2^{28}\) floats, 99 percent of this machine's streaming ceiling and slightly ahead of torch.sum (2974.4 GB/s). ptxas reports 13 registers and 128 bytes of shared memory, exactly the 32-float array of per-warp partials. The Triton tab is the same two-level shape (2837.4 GB/s measured) with the tree implicit in tl.sum.

// reduce4: grid-stride + warp shuffles. Measured: 0.357 ms, 3009.7 GB/s.
// Structure: long coalesced serial loop (bandwidth phase), then a
// log-depth tree entirely in registers (compute phase, ~0 cost).
#include <cuda_runtime.h>

__global__ void reduce_gridstride(const float* __restrict__ x,
                                  float* __restrict__ out, size_t n) {
    float sum = 0.0f;
    size_t stride = (size_t)gridDim.x * blockDim.x;
    for (size_t i = (size_t)blockIdx.x * blockDim.x + threadIdx.x;
         i < n; i += stride)
        sum += x[i];                         // each thread: ~n/(grid) adds,
                                             // consecutive lanes coalesce
    // warp tree: 5 shuffle-adds, register to register, no barrier needed
    for (int off = 16; off > 0; off >>= 1)
        sum += __shfl_down_sync(0xffffffffu, sum, off);

    __shared__ float warp_partial[32];       // 128 B, matches ptxas
    int lane = threadIdx.x & 31, warp = threadIdx.x >> 5;
    if (lane == 0) warp_partial[warp] = sum; // one value per warp
    __syncthreads();

    if (warp == 0) {                         // first warp reduces the partials
        int nwarps = blockDim.x >> 5;
        sum = (lane < nwarps) ? warp_partial[lane] : 0.0f;
        for (int off = 16; off > 0; off >>= 1)
            sum += __shfl_down_sync(0xffffffffu, sum, off);
        if (lane == 0) atomicAdd(out, sum);  // few blocks -> low contention
    }
}
// launch: reduce_gridstride<<<132 * 4, 256>>>(x, out, n);  out zeroed first
# Reduction in Triton. Measured: 0.378 ms, 2837.4 GB/s on 2^28 floats.
# tl.sum lowers to the same shuffle tree the CUDA tab writes by hand.
import torch
import triton
import triton.language as tl

@triton.jit
def reduce_kernel(x_ptr, out_ptr, n, BLOCK: tl.constexpr):
    pid = tl.program_id(0)
    offs = pid * BLOCK + tl.arange(0, BLOCK)
    x = tl.load(x_ptr + offs, mask=offs < n, other=0.0)
    partial = tl.sum(x, axis=0)              # block-level tree in registers
    tl.atomic_add(out_ptr, partial)          # one atomic per program

def gpu_sum(x: torch.Tensor) -> torch.Tensor:
    out = torch.zeros(1, device=x.device, dtype=torch.float32)
    n = x.numel()
    reduce_kernel[(triton.cdiv(n, 4096),)](x, out, n, BLOCK=4096)
    return out

x = torch.randn(1 << 28, device="cuda")
# atomics reorder float adds across runs: compare with a tolerance
assert torch.allclose(gpu_sum(x), x.sum(), rtol=1e-4)

Transpose, shared memory and one column of padding

A transpose computes nothing. It only moves \(n^2\) elements, so its ceiling is the copy rate and every deficit is a memory-system fault. The naive kernel cannot coalesce both ends at once. If the read is contiguous the write has stride \(n\), and at \(n = 8192\) that is 32 KB between consecutive lanes, one sector per lane. Staging a \(32 \times 32\) tile through shared memory fixes both ends, since the transpose happens on chip. That introduces the second fault. Reading a column of a \(32 \times 32\) shared array puts all 32 lanes in the same bank, a 32-way conflict serializing the access. Padding the row to 33 floats shifts each row by one bank, so column \(c\) of row \(r\) lands in bank \((33r + c) \bmod 32 = (r + c) \bmod 32\), all distinct across \(r\). Measured at \(n = 8192\) (536.9 MB read, 536.9 MB written), each version verified bit-exact against x.t().contiguous(), the naive kernel runs at 451.1 GB/s, tiled at 1702.9 GB/s, and tiled with padding at 2694.0 GB/s, against 2837.4 GB/s for a straight copy. One byte of padding per row is worth \(1.58\times\).

// transpose.cu, n = 8192 fp32. Measured on this H100:
//   naive 451.1 GB/s | tiled 1702.9 | tiled+pad 2694.0 | copy 2837.4
// ptxas (padded): 30 regs, 4224 B smem (= 32 * 33 * 4), 0 spills.
#include <cuda_runtime.h>

// Naive: reads coalesce (consecutive x = consecutive addresses),
// writes do not (consecutive x = addresses n floats apart).
__global__ void transpose_naive(const float* __restrict__ in,
                                float* __restrict__ out, int n) {
    int x = blockIdx.x * 32 + threadIdx.x;      // column
    int y = blockIdx.y * 32 + threadIdx.y;      // row (block is 32 x 8)
    for (int j = 0; j < 32; j += 8)             // 4 rows per thread
        if (x < n && y + j < n)
            out[(long)x * n + (y + j)] = in[(long)(y + j) * n + x];
}

// Tiled: both global accesses coalesce; the transpose happens in smem.
// PAD = 1 makes the stride 33, so a column of the tile spans 32 banks.
template <int PAD>
__global__ void transpose_tiled(const float* __restrict__ in,
                                float* __restrict__ out, int n) {
    __shared__ float tile[32][32 + PAD];
    int x = blockIdx.x * 32 + threadIdx.x;
    int y = blockIdx.y * 32 + threadIdx.y;
    for (int j = 0; j < 32; j += 8)             // coalesced read of a tile
        if (x < n && y + j < n)
            tile[threadIdx.y + j][threadIdx.x] = in[(long)(y + j) * n + x];
    __syncthreads();
    x = blockIdx.y * 32 + threadIdx.x;          // swap block indices so the
    y = blockIdx.x * 32 + threadIdx.y;          // write is also coalesced
    for (int j = 0; j < 32; j += 8)
        if (x < n && y + j < n)                 // column read of the tile:
            out[(long)(y + j) * n + x] = tile[threadIdx.x][threadIdx.y + j];
}
// launch: dim3 grid(n/32, n/32), block(32, 8);
//         transpose_tiled<1><<<grid, block>>>(in, out, n);
# Triton transpose. Measured: 0.216 ms, 2488.8 GB/s at n = 8192.
# The compiler owns the shared-memory staging and the swizzle that
# replaces the hand-written padding column.
import torch
import triton
import triton.language as tl

@triton.jit
def transpose_kernel(in_ptr, out_ptr, n,
                     BM: tl.constexpr, BN: tl.constexpr):
    pid_m = tl.program_id(0)
    pid_n = tl.program_id(1)
    rm = pid_m * BM + tl.arange(0, BM)          # (BM,) rows of input
    rn = pid_n * BN + tl.arange(0, BN)          # (BN,) cols of input
    src = in_ptr + rm[:, None] * n + rn[None, :]        # (BM, BN)
    x = tl.load(src, mask=(rm[:, None] < n) & (rn[None, :] < n))
    dst = out_ptr + rn[:, None] * n + rm[None, :]       # (BN, BM)
    tl.store(dst, tl.trans(x), mask=(rn[:, None] < n) & (rm[None, :] < n))

def transpose(x):
    n = x.shape[0]
    o = torch.empty_like(x)
    grid = (triton.cdiv(n, 64), triton.cdiv(n, 64))
    transpose_kernel[grid](x, o, n, BM=64, BN=64)
    return o

x = torch.randn(8192, 8192, device="cuda")
assert torch.equal(transpose(x), x.t().contiguous())
import torch

# The reference. x.t() is free: it only rewrites the stride metadata,
# producing a non-contiguous view with no data movement at all.
x = torch.randn(8192, 8192, device="cuda")
v = x.t()                          # (8192, 8192) view, stride (1, 8192)
assert v.data_ptr() == x.data_ptr()
assert not v.is_contiguous()

# .contiguous() is the kernel that actually moves 1.07 GB. This is what
# the CUDA and Triton tabs reimplement, and it is what a "free" transpose
# costs the moment something downstream demands a contiguous layout.
y = v.contiguous()                 # measured baseline for the tabs above

# The practical lesson: prefer operators that accept strides. cuBLAS and
# tl.dot both take a transpose flag, so a matmul against x.t() never
# needs this kernel; a reshape after a transpose always does.
q = torch.randn(8, 4096, 16, 64, device="cuda")   # (B, L, H, d)
q = q.transpose(1, 2)                             # (B, H, L, d) view
# .reshape here would force a copy; .contiguous() makes the cost explicit.
assert not q.is_contiguous()

Scan, the three-phase decomposition in tensor code

Rather than reproduce a full CUDA scan (CUB's implementation is the reference, in the open-source section), this block shows the decomposition every GPU scan uses, expressed in PyTorch, and the same computation through JAX's built-in associative machinery, which applies the Blelloch construction to any associative operator. The second JAX example scans \(2\times 2\) matrix products to solve the recurrence \(z_k = a_k z_{k-1} + b_k\) in \(O(\log n)\) depth, the trick behind parallelized linear state-space layers.

import torch

x = torch.randn(1 << 20, device="cuda")          # (n,), n = 2^20
inc = torch.cumsum(x, dim=0)                     # library scan
exc = inc - x                                    # exclusive from inclusive

# The three-phase structure of every device-wide GPU scan:
#   1) scan each block locally, 2) scan the block totals,
#   3) add each block's carry to its elements.
B = 1024
blocks = x.view(-1, B)                           # (n/B, B)
local  = torch.cumsum(blocks, dim=1)             # phase 1: (n/B, B)
totals = local[:, -1]                            # (n/B,) block sums
carry  = torch.cumsum(totals, dim=0) - totals    # phase 2: exclusive scan
out    = (local + carry[:, None]).reshape(-1)    # phase 3: apply carries
assert torch.allclose(out, inc, atol=1e-2)

# Scan as sort primitive: stable split by bit -> radix sort pass
keys  = torch.randint(0, 2, (16,), device="cuda")
zeros = (keys == 0).long()
dst0  = torch.cumsum(zeros, 0) - zeros           # exclusive scan of 0-flags
dst1  = zeros.sum() + torch.cumsum(1 - zeros, 0) - (1 - zeros)
dst   = torch.where(keys == 0, dst0, dst1)       # scatter destinations
import jax
import jax.numpy as jnp
from jax import lax

x = jax.random.normal(jax.random.PRNGKey(0), (1 << 20,))  # (n,)
inc = jnp.cumsum(x)
# Blelloch-style construction for any associative op, O(n) work:
inc2 = lax.associative_scan(jnp.add, x)
run_max = lax.associative_scan(jnp.maximum, x)   # running maximum

# Linear recurrence z_k = a_k z_{k-1} + b_k as a scan over the
# associative composition of affine maps (a, b) ∘ (a', b'):
def affine_compose(f, g):
    a1, b1 = f
    a2, b2 = g
    return (a2 * a1, a2 * b1 + b2)               # apply f, then g

key_a, key_b = jax.random.split(jax.random.PRNGKey(1))
a = jax.random.uniform(key_a, (1 << 20,)) * 0.9  # (n,) decay terms
b = jax.random.normal(key_b, (1 << 20,))         # (n,) inputs
_, z = lax.associative_scan(affine_compose, (a, b))
# z[k] == a[k] * z[k-1] + b[k], computed at depth O(log n):
z_ref = b[0]
for k in range(1, 5):
    z_ref = a[k] * z_ref + b[k]
assert jnp.allclose(z[4], z_ref, rtol=1e-4)

Matmul, shared-memory tiles in CUDA, tl.dot in Triton

The CUDA tab is the measured tiled32 kernel (17.04 ms, 8.1 TFLOPS fp32 at \(n = 4096\)) with comments marking where blocked128 (5.48 ms, 25.1 TFLOPS) departs from it. The Triton tab is the measured fp16 kernel (0.269 ms, 511.6 TFLOPS against cuBLAS at 702.4). tl.dot compiles to tensor-core instructions, and the compiler handles shared-memory staging, swizzling, and pipelining that CUDA leaves to the programmer.

// tiled32: measured 17.04 ms / 8.1 TFLOPS fp32 at n = 4096.
// ptxas: 32 regs, 8192 B smem, full occupancy -- and still 6x off cuBLAS,
// because 1 output/thread leaves it bound on shared-memory bandwidth.
#define TILE 32

__global__ void sgemm_tiled(const float* __restrict__ A,
                            const float* __restrict__ B,
                            float* __restrict__ C, int n) {
    __shared__ float As[TILE][TILE];          // 4 KB
    __shared__ float Bs[TILE][TILE];          // 4 KB
    int row = blockIdx.y * TILE + threadIdx.y;
    int col = blockIdx.x * TILE + threadIdx.x;
    float acc = 0.0f;
    for (int t = 0; t < n / TILE; ++t) {      // march over k in tile slices
        As[threadIdx.y][threadIdx.x] = A[row * n + t * TILE + threadIdx.x];
        Bs[threadIdx.y][threadIdx.x] = B[(t * TILE + threadIdx.y) * n + col];
        __syncthreads();                      // tile staged before use
        #pragma unroll
        for (int k = 0; k < TILE; ++k)        // 32 FMAs per global element:
            acc += As[threadIdx.y][k] * Bs[k][threadIdx.x];   // AI = T/4 = 8
        __syncthreads();                      // done before overwrite
    }
    C[row * n + col] = acc;
}
// blocked128 (measured 25.1 TFLOPS) changes exactly one thing: each thread
// owns an 8x8 register tile of C. Inner iteration becomes
//   float a[8], b[8];             // read 16 floats from smem
//   acc[i][j] += a[i] * b[j];     // do 64 FMAs: 4 flops per smem byte
// cost: 112 regs/thread -> 25% occupancy -> 3.1x faster anyway.
# Triton matmul, fp16 in / fp32 accumulate. Measured at n = 4096:
# 0.269 ms, 511.6 TFLOPS (cuBLAS fp16: 702.4). Dims assumed multiples of BM/BN/BK.
import torch
import triton
import triton.language as tl

@triton.jit
def matmul_kernel(a_ptr, b_ptr, c_ptr, 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)          # (BM,) rows of C
    rn = pid_n * BN + tl.arange(0, BN)          # (BN,) cols of C
    rk = tl.arange(0, BK)                       # (BK,) k-slice
    a_ptrs = a_ptr + rm[:, None] * sam + rk[None, :] * sak   # (BM, BK)
    b_ptrs = b_ptr + rk[:, None] * sbk + rn[None, :] * sbn   # (BK, BN)
    acc = tl.zeros((BM, BN), dtype=tl.float32)  # register tile of C
    for _ in range(0, K, BK):
        a = tl.load(a_ptrs)                     # compiler stages via smem,
        b = tl.load(b_ptrs)                     # swizzles, pipelines
        acc = tl.dot(a, b, acc)                 # tensor-core MMA
        a_ptrs += BK * sak
        b_ptrs += BK * sbk
    c_ptrs = c_ptr + rm[:, None] * scm + rn[None, :] * scn
    tl.store(c_ptrs, acc.to(tl.float16))

def matmul(a, b):
    M, K = a.shape
    K2, N = b.shape
    c = torch.empty(M, N, device=a.device, dtype=torch.float16)
    grid = (triton.cdiv(M, 128), triton.cdiv(N, 128))
    matmul_kernel[grid](a, b, c, M, N, K,
                        a.stride(0), a.stride(1), b.stride(0), b.stride(1),
                        c.stride(0), c.stride(1), BM=128, BN=128, BK=64)
    return c

a = torch.randn(4096, 4096, device="cuda", dtype=torch.float16)
b = torch.randn(4096, 4096, device="cuda", dtype=torch.float16)
assert torch.allclose(matmul(a, b), a @ b, atol=1e-1, rtol=1e-2)

Softmax, one pass instead of three

The textbook stable softmax makes three passes over a row, finding the max, accumulating \(\sum_j e^{x_j - m}\), and dividing. Three passes over a row that does not fit in registers means three trips to HBM, and since softmax does a handful of flops per element it is squarely memory-bound, so the pass count is the runtime. The online-softmax recurrence collapses the first two into one. Maintain \((m, \ell)\) and merge each new element with \(\ell \leftarrow \ell e^{m - m'} + e^{x - m'}\), which is the identity derived in the attention section applied to a chunk of size one. The kernel below runs one block per row, has each thread walk the row with a grid-stride loop accumulating its own \((m, \ell)\), merges the per-thread pairs with __shfl_xor_sync butterfly exchanges inside each warp and a 32-entry shared array across warps, then makes a second pass to write. Two passes, not three, and the second one hits L2 for rows small enough to survive there.

Measured over an \(8192 \times 8192\) fp32 matrix, verified to a maximum absolute error of \(2.8 \times 10^{-9}\) against torch.softmax, the fused CUDA kernel runs in 0.363 ms (1480.3 GB/s), the Triton version in 0.294 ms (1828.0 GB/s), torch.softmax in 0.301 ms, and an unfused three-pass implementation in 0.863 ms. The \(2.4\times\) between three passes and two-plus-a-cached-one is the entire content of the optimization, and the fact that PyTorch's own kernel lands with the hand-written ones confirms the library already does it.

// softmax_online.cu: one block per row, online (max, sum) in one pass.
// Measured: 0.363 ms, 1480.3 GB/s on 8192 x 8192 fp32; max abs err 2.8e-9
// vs torch.softmax. ptxas: 20 regs, 256 B smem, 0 spills.
#include <cuda_runtime.h>
#include <math_constants.h>

// Merge two (max, sum) accumulators exactly: rebase both to the joint max.
__device__ __forceinline__ void merge(float& m, float& d, float m2, float d2) {
    float mn = fmaxf(m, m2);
    if (mn == -CUDART_INF_F) return;            // both accumulators empty
    d = d * expf(m - mn) + d2 * expf(m2 - mn);  // exact: exp(x-m)exp(m-mn)
    m = mn;                                     //        = exp(x-mn)
}

__global__ void softmax_online(const float* __restrict__ x,
                               float* __restrict__ y, int cols) {
    int row = blockIdx.x;
    const float* xr = x + (long)row * cols;
    float* yr = y + (long)row * cols;

    float m = -CUDART_INF_F, d = 0.f;           // this thread's (max, sum)
    for (int j = threadIdx.x; j < cols; j += blockDim.x) {
        float v = xr[j];                        // coalesced: lane i reads j+i
        float mn = fmaxf(m, v);
        d = d * expf(m - mn) + expf(v - mn);    // online merge, chunk of 1
        m = mn;
    }
    // butterfly reduce the (max, sum) pair across the warp: after 5 rounds
    // every lane holds the warp's merged accumulator.
    for (int off = 16; off > 0; off >>= 1)
        merge(m, d, __shfl_xor_sync(0xffffffffu, m, off),
                    __shfl_xor_sync(0xffffffffu, d, off));

    __shared__ float ms[32], ds[32];            // 256 B, matches ptxas
    int lane = threadIdx.x & 31, warp = threadIdx.x >> 5;
    if (lane == 0) { ms[warp] = m; ds[warp] = d; }
    __syncthreads();
    if (warp == 0) {                            // one warp merges the warps
        int nw = blockDim.x >> 5;
        m = lane < nw ? ms[lane] : -CUDART_INF_F;
        d = lane < nw ? ds[lane] : 0.f;
        for (int off = 16; off > 0; off >>= 1)
            merge(m, d, __shfl_xor_sync(0xffffffffu, m, off),
                        __shfl_xor_sync(0xffffffffu, d, off));
        if (lane == 0) { ms[0] = m; ds[0] = d; }
    }
    __syncthreads();
    m = ms[0]; d = ds[0];                       // broadcast the row statistics
    for (int j = threadIdx.x; j < cols; j += blockDim.x)
        yr[j] = expf(xr[j] - m) / d;            // second and final pass
}
// launch: softmax_online<<<rows, 256>>>(x, y, cols);
# Triton softmax, one program per row. Measured: 0.294 ms, 1828.0 GB/s
# on 8192 x 8192 fp32 (torch.softmax: 0.301 ms).
import torch
import triton
import triton.language as tl

@triton.jit
def softmax_kernel(x_ptr, y_ptr, stride, cols, BLOCK: tl.constexpr):
    row = tl.program_id(0)
    offs = tl.arange(0, BLOCK)                  # BLOCK >= cols, power of two
    mask = offs < cols
    x = tl.load(x_ptr + row * stride + offs, mask=mask, other=float("-inf"))
    x = x - tl.max(x, axis=0)                   # stability shift
    e = tl.exp(x)
    tl.store(y_ptr + row * stride + offs, e / tl.sum(e, axis=0), mask=mask)

def softmax(x):
    rows, cols = x.shape
    y = torch.empty_like(x)
    BLOCK = triton.next_power_of_2(cols)
    softmax_kernel[(rows,)](x, y, x.stride(0), cols, BLOCK=BLOCK,
                            num_warps=8 if BLOCK >= 2048 else 4)
    return y

x = torch.randn(8192, 8192, device="cuda")
assert torch.allclose(softmax(x), torch.softmax(x, dim=1), atol=1e-6)
# Note the tradeoff: this version holds the whole row in registers, so it
# needs one pass but caps cols at what BLOCK * num_warps can hold. The CUDA
# tab streams the row instead and works at any width.
import torch

x = torch.randn(8192, 8192, device="cuda")

# The reference. torch.softmax is already a fused two-pass kernel;
# measured 0.301 ms / 1782.8 GB/s on this H100.
y = torch.softmax(x, dim=1)

# The unfused three-pass version the fusion is measured against:
# 0.863 ms, i.e. 2.4x slower for identical arithmetic.
def softmax_unfused(x):
    m = x.max(dim=1, keepdim=True).values     # pass 1: read x
    e = torch.exp(x - m)                      # pass 2: read x, write e
    return e / e.sum(dim=1, keepdim=True)     # pass 3: read e twice

assert torch.allclose(softmax_unfused(x), y, atol=1e-6)

# Why this matters beyond softmax: any row-wise statistic followed by a
# row-wise application has the same 3-pass/2-pass structure. LayerNorm,
# RMSNorm, and the attention softmax are all this pattern.

LayerNorm and RMSNorm, Welford in a fused kernel

Normalization has the same shape as softmax and one extra subtlety. LayerNorm needs the row mean and variance, and the textbook one-pass formula \(\Var = \E[x^2] - \E[x]^2\) is numerically hazardous. When the mean is large relative to the spread, it subtracts two nearly equal large numbers and loses most of the significant digits, in the worst case returning a negative variance. Welford's recurrence avoids the cancellation by tracking the mean and the sum of squared deviations \(M_2\) directly. For one new sample \(x\) with running count \(n\), mean \(\mu\), and \(M_2\),

$$ n' = n+1, \qquad \delta = x - \mu, \qquad \mu' = \mu + \frac{\delta}{n'}, \qquad M_2' = M_2 + \delta\,(x - \mu') . $$

The parallel version needs a merge rather than an increment, since each thread accumulates its own strided slice. For two accumulators \((n_a, \mu_a, M_{2,a})\) and \((n_b, \mu_b, M_{2,b})\) with \(\delta = \mu_b - \mu_a\) and \(n = n_a + n_b\), Chan, Golub, and LeVeque's combination is

$$ \mu = \mu_a + \delta\frac{n_b}{n}, \qquad M_2 = M_{2,a} + M_{2,b} + \delta^2 \frac{n_a n_b}{n} . $$

The derivation of the \(M_2\) term is a direct expansion. Take \(M_2 = \sum_i (x_i - \mu)^2\) over the union, split into the two groups, with each group's deviation rewritten around its own mean, \(\sum_{i \in a}(x_i - \mu)^2 = M_{2,a} + n_a(\mu_a - \mu)^2\), and then \(\mu_a - \mu = -\delta n_b/n\) and \(\mu_b - \mu = \delta n_a/n\) substituted, giving \(n_a \delta^2 n_b^2/n^2 + n_b \delta^2 n_a^2/n^2 = \delta^2 n_a n_b/n\). The merge is associative, which is exactly what a warp-shuffle tree requires. RMSNorm drops the mean entirely, normalizing by \(\sqrt{\E[x^2] + \epsilon}\), so it needs one accumulator instead of three and no cancellation risk at all. The reason it is both cheaper and, empirically, no worse in transformers is the same reason, that the mean subtraction was never doing much work.

Measured over a \(16384 \times 4096\) fp32 matrix, each verified against the PyTorch reference, the fused Welford LayerNorm runs in 0.278 ms (1928.6 GB/s) with a maximum absolute error of \(2.9 \times 10^{-6}\), against F.layer_norm at 0.283 ms. The Triton LayerNorm runs at 0.291 ms and the Triton RMSNorm at 0.275 ms (1949.8 GB/s), and the CUDA RMSNorm at 0.251 ms (2137.6 GB/s, 16 registers, 128 bytes of shared memory, no spills) with a maximum absolute error of \(1.9 \times 10^{-6}\), which is \(2.9\times\) faster than F.rms_norm at 0.731 ms on this shape.

// layernorm.cu: fused Welford LayerNorm, one block per row.
// Measured: 0.278 ms, 1928.6 GB/s on 16384 x 4096 fp32; max abs err 2.9e-6.
// ptxas: 22 regs, 384 B smem, 0 spills.
#include <cuda_runtime.h>

// Chan-Golub-LeVeque parallel merge of two Welford accumulators.
__device__ __forceinline__ void wmerge(float& mean, float& m2, float& cnt,
                                       float mean2, float m22, float cnt2) {
    float tot = cnt + cnt2;
    if (cnt2 == 0.f || tot == 0.f) return;
    float delta = mean2 - mean;
    mean += delta * cnt2 / tot;                    // mu_a + delta * n_b / n
    m2   += m22 + delta * delta * cnt * cnt2 / tot;// + delta^2 n_a n_b / n
    cnt   = tot;
}

__global__ void layernorm(const float* __restrict__ x,
                          const float* __restrict__ gamma,
                          const float* __restrict__ beta,
                          float* __restrict__ y, int cols, float eps) {
    int row = blockIdx.x;
    const float* xr = x + (long)row * cols;
    float* yr = y + (long)row * cols;

    float mean = 0.f, m2 = 0.f, cnt = 0.f;         // per-thread accumulator
    for (int j = threadIdx.x; j < cols; j += blockDim.x) {
        float v = xr[j];                           // coalesced stride-blockDim
        cnt += 1.f;
        float delta = v - mean;
        mean += delta / cnt;                       // Welford increment
        m2   += delta * (v - mean);
    }
    for (int off = 16; off > 0; off >>= 1)         // warp tree over triples
        wmerge(mean, m2, cnt, __shfl_xor_sync(0xffffffffu, mean, off),
                              __shfl_xor_sync(0xffffffffu, m2,   off),
                              __shfl_xor_sync(0xffffffffu, cnt,  off));

    __shared__ float sm[32], s2[32], sc[32];       // 384 B, matches ptxas
    int lane = threadIdx.x & 31, warp = threadIdx.x >> 5;
    if (lane == 0) { sm[warp] = mean; s2[warp] = m2; sc[warp] = cnt; }
    __syncthreads();
    if (warp == 0) {
        int nw = blockDim.x >> 5;
        mean = lane < nw ? sm[lane] : 0.f;
        m2   = lane < nw ? s2[lane] : 0.f;
        cnt  = lane < nw ? sc[lane] : 0.f;
        for (int off = 16; off > 0; off >>= 1)
            wmerge(mean, m2, cnt, __shfl_xor_sync(0xffffffffu, mean, off),
                                  __shfl_xor_sync(0xffffffffu, m2,   off),
                                  __shfl_xor_sync(0xffffffffu, cnt,  off));
        if (lane == 0) { sm[0] = mean; s2[0] = m2; }
    }
    __syncthreads();
    mean = sm[0];
    float rstd = rsqrtf(s2[0] / cols + eps);       // M2 / n is the variance
    for (int j = threadIdx.x; j < cols; j += blockDim.x)
        yr[j] = (xr[j] - mean) * rstd * gamma[j] + beta[j];
}

// RMSNorm: one statistic, no mean, no cancellation risk.
// Measured: 0.251 ms, 2137.6 GB/s, max abs err 1.9e-6; 16 regs, 128 B smem.
__global__ void rmsnorm(const float* __restrict__ x,
                        const float* __restrict__ gamma,
                        float* __restrict__ y, int cols, float eps) {
    int row = blockIdx.x;
    const float* xr = x + (long)row * cols;
    float* yr = y + (long)row * cols;
    float ss = 0.f;
    for (int j = threadIdx.x; j < cols; j += blockDim.x) {
        float v = xr[j]; ss += v * v;
    }
    for (int off = 16; off > 0; off >>= 1)
        ss += __shfl_down_sync(0xffffffffu, ss, off);
    __shared__ float s[32];
    int lane = threadIdx.x & 31, warp = threadIdx.x >> 5;
    if (lane == 0) s[warp] = ss;
    __syncthreads();
    if (warp == 0) {
        int nw = blockDim.x >> 5;
        ss = lane < nw ? s[lane] : 0.f;
        for (int off = 16; off > 0; off >>= 1)
            ss += __shfl_down_sync(0xffffffffu, ss, off);
        if (lane == 0) s[0] = rsqrtf(ss / cols + eps);
    }
    __syncthreads();
    float r = s[0];
    for (int j = threadIdx.x; j < cols; j += blockDim.x)
        yr[j] = xr[j] * r * gamma[j];
}
# Triton LayerNorm and RMSNorm, one program per row.
# Measured on 16384 x 4096 fp32: layernorm 0.291 ms / 1842.7 GB/s,
# rmsnorm 0.275 ms / 1949.8 GB/s.
import torch
import triton
import triton.language as tl

@triton.jit
def layernorm_kernel(x_ptr, g_ptr, b_ptr, y_ptr, stride, cols, eps,
                     BLOCK: tl.constexpr):
    row = tl.program_id(0)
    offs = tl.arange(0, BLOCK)
    mask = offs < cols
    x = tl.load(x_ptr + row * stride + offs, mask=mask, other=0.0)
    mean = tl.sum(x, axis=0) / cols
    xc = tl.where(mask, x - mean, 0.0)           # centered, zeros outside
    var = tl.sum(xc * xc, axis=0) / cols         # E[(x-mu)^2] directly:
    rstd = 1.0 / tl.sqrt(var + eps)              # no E[x^2]-E[x]^2 cancellation
    g = tl.load(g_ptr + offs, mask=mask, other=0.0)
    b = tl.load(b_ptr + offs, mask=mask, other=0.0)
    tl.store(y_ptr + row * stride + offs, xc * rstd * g + b, mask=mask)

@triton.jit
def rmsnorm_kernel(x_ptr, g_ptr, y_ptr, stride, cols, eps,
                   BLOCK: tl.constexpr):
    row = tl.program_id(0)
    offs = tl.arange(0, BLOCK)
    mask = offs < cols
    x = tl.load(x_ptr + row * stride + offs, mask=mask, other=0.0)
    rstd = 1.0 / tl.sqrt(tl.sum(x * x, axis=0) / cols + eps)
    g = tl.load(g_ptr + offs, mask=mask, other=0.0)
    tl.store(y_ptr + row * stride + offs, x * rstd * g, mask=mask)

rows, cols = 16384, 4096
x = torch.randn(rows, cols, device="cuda")
g = torch.randn(cols, device="cuda"); b = torch.randn(cols, device="cuda")
y = torch.empty_like(x)
BLOCK = triton.next_power_of_2(cols)
layernorm_kernel[(rows,)](x, g, b, y, x.stride(0), cols, 1e-5,
                          BLOCK=BLOCK, num_warps=8)
ref = torch.nn.functional.layer_norm(x, (cols,), g, b, 1e-5)
assert torch.allclose(y, ref, atol=1e-4)
import torch
import torch.nn.functional as F

rows, cols = 16384, 4096
x = torch.randn(rows, cols, device="cuda")
g = torch.randn(cols, device="cuda")
b = torch.randn(cols, device="cuda")

# References the kernels above are checked against.
# Measured on this H100: layer_norm 0.283 ms, rms_norm 0.731 ms.
ln = F.layer_norm(x, (cols,), g, b, eps=1e-5)
rms = F.rms_norm(x, (cols,), g, eps=1e-6)

# Written out, so the two-statistic vs one-statistic difference is visible:
mu = x.mean(dim=-1, keepdim=True)                       # (rows, 1)
var = x.var(dim=-1, keepdim=True, unbiased=False)       # (rows, 1)
ln_ref = (x - mu) * torch.rsqrt(var + 1e-5) * g + b
assert torch.allclose(ln, ln_ref, atol=1e-4)

rms_ref = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + 1e-6) * g
assert torch.allclose(rms, rms_ref, atol=1e-4)

# The unstable one-pass formula, for contrast. Try it with x += 1e4 and
# watch the variance go to zero or negative in fp32; Welford does not.
var_bad = (x * x).mean(-1, keepdim=True) - mu * mu

Fused attention, the Triton kernel, and the library calls

The Triton tab is the measured flash-forward kernel (340.2 TFLOPS fp16 at batch \(\times\) heads \(= 64\), \(L = 4096\), \(d = 64\), max abs error 1.2e-4 against SDPA), non-causal and simplified to its load-bearing 40 lines. One program owns a block of query rows, streams all key/value blocks through on-chip memory, and carries the online-softmax state \((m, \ell, o)\) in registers. The PyTorch and JAX tabs are what production code should almost always call instead, with the kernel choice pinned so a silent fallback to the materializing path cannot reintroduce the \(O(L^2)\) memory.

# Minimal FlashAttention forward (non-causal). Measured: 0.808 ms,
# 340.2 TFLOPS at bh=64, L=4096, d=64 (SDPA on same shape: 0.926 ms).
import torch
import triton
import triton.language as tl

@triton.jit
def flash_fwd(q_ptr, k_ptr, v_ptr, o_ptr, L, scale,
              BM: tl.constexpr, BN: tl.constexpr, D: tl.constexpr):
    pid_m = tl.program_id(0)               # which query block
    pid_bh = tl.program_id(1)              # which (batch, head)
    offs_m = pid_m * BM + tl.arange(0, BM) # (BM,) query rows
    offs_d = tl.arange(0, D)               # (D,) head dim
    base = pid_bh * L * D
    q = tl.load(q_ptr + base + offs_m[:, None] * D + offs_d[None, :])  # (BM,D)

    m_i = tl.full((BM,), float("-inf"), tl.float32)   # running row max
    l_i = tl.zeros((BM,), tl.float32)                 # running normalizer
    acc = tl.zeros((BM, D), tl.float32)               # unnormalized output

    for start in range(0, L, BN):          # stream K/V blocks: S never in HBM
        offs_n = start + tl.arange(0, BN)
        k = tl.load(k_ptr + base + offs_n[:, None] * D + offs_d[None, :])
        v = tl.load(v_ptr + base + offs_n[:, None] * D + offs_d[None, :])
        s = tl.dot(q, tl.trans(k)) * scale            # (BM, BN) score tile
        m_new = tl.maximum(m_i, tl.max(s, axis=1))    # merge maxima
        p = tl.exp(s - m_new[:, None])                # local exponentials
        alpha = tl.exp(m_i - m_new)                   # rebase factor
        l_i = l_i * alpha + tl.sum(p, axis=1)         # exact merge of sums
        acc = acc * alpha[:, None] + tl.dot(p.to(v.dtype), v)
        m_i = m_new
    acc = acc / l_i[:, None]                          # normalize once
    tl.store(o_ptr + base + offs_m[:, None] * D + offs_d[None, :],
             acc.to(tl.float16))

def flash(q, k, v):                        # q,k,v: (BH, L, D) fp16
    BH, L, D = q.shape
    o = torch.empty_like(q)
    grid = (triton.cdiv(L, 128), BH)
    flash_fwd[grid](q, k, v, o, L, D ** -0.5, BM=128, BN=64, D=D)
    return o
// flash_fwd_minimal.cu: correctness-first fused attention, fp32, no
// tensor cores. One thread owns one query row; K/V tiles are staged
// through shared memory and the online-softmax state lives in registers.
// Measured at bh=32, L=4096, d=64: 14.304 ms / 9.6 TFLOPS, against
// PyTorch fp32 SDPA at 8.358 ms / 16.4 TFLOPS. Max abs err 8.9e-7.
// ptxas: 167 regs, 32768 B smem, 0 spills.  Assumes L % 64 == 0.
#include <cuda_runtime.h>
#include <math_constants.h>

__global__ void flash_fwd(const float* __restrict__ Q,
                          const float* __restrict__ K,
                          const float* __restrict__ V,
                          float* __restrict__ O, int L, float scale) {
    const int D = 64, BN = 64;                  // head dim, key-block size
    int bh   = blockIdx.y;                      // which (batch, head)
    int qrow = blockIdx.x * blockDim.x + threadIdx.x;   // this thread's query
    const float* Qb = Q + (long)bh * L * D;
    const float* Kb = K + (long)bh * L * D;
    const float* Vb = V + (long)bh * L * D;
    float* Ob = O + (long)bh * L * D;

    float q[D], acc[D];                         // 128 registers of state:
    for (int d = 0; d < D; ++d) {               // this is why ptxas says 167
        q[d]   = Qb[(long)qrow * D + d];        // q stays resident forever
        acc[d] = 0.f;                           // unnormalized output o
    }
    float m = -CUDART_INF_F, l = 0.f;           // running max and normalizer

    __shared__ float Ks[BN][D];                 // 16 KB each: 32 KB total
    __shared__ float Vs[BN][D];

    for (int t = 0; t < L; t += BN) {           // stream K/V: S never in HBM
        for (int d = 0; d < D; ++d) {           // cooperative tile load
            Ks[threadIdx.x][d] = Kb[(long)(t + threadIdx.x) * D + d];
            Vs[threadIdx.x][d] = Vb[(long)(t + threadIdx.x) * D + d];
        }
        __syncthreads();                        // tile staged before use
        for (int n = 0; n < BN; ++n) {
            float s = 0.f;
            for (int d = 0; d < D; ++d) s += q[d] * Ks[n][d];
            s *= scale;                         // one score, in a register
            float mn    = fmaxf(m, s);          // m' = max(m, s)
            float alpha = expf(m - mn);         // rebase factor
            float p     = expf(s - mn);
            l = l * alpha + p;                  // exact normalizer merge
            for (int d = 0; d < D; ++d)         // exact output merge
                acc[d] = acc[d] * alpha + p * Vs[n][d];
            m = mn;
        }
        __syncthreads();                        // done before overwrite
    }
    for (int d = 0; d < D; ++d)                 // normalize once, at the end
        Ob[(long)qrow * D + d] = acc[d] / l;
}
// launch: dim3 grid(L / 64, BH); flash_fwd<<<grid, 64>>>(Q, K, V, O, L, 1.f/8.f);
//
// What this kernel gets right: O(L) memory, exact online softmax, coalesced
// tile loads. What it gives up, and why it loses to SDPA by 1.7x: the two
// matmuls are scalar FMA loops instead of tensor-core MMAs, there is no
// double buffering of the K/V tiles, and 167 registers per thread cap
// occupancy hard. Correctness first, then the datapath.
import torch
import torch.nn.functional as F
from torch.nn.attention import SDPBackend, sdpa_kernel

# (batch, heads, L, d) -- measured on this H100 at bh=128, L=8192, d=64:
# flash path 3.575 ms / 615.1 TFLOPS vs naive 98.946 ms / 35.0 GB peak.
q = torch.randn(8, 16, 8192, 64, device="cuda", dtype=torch.bfloat16)
k = torch.randn_like(q)
v = torch.randn_like(q)

# Pin the fused kernel; a silent fallback to the math path would
# quietly reintroduce the O(L^2) score matrix.
with sdpa_kernel(SDPBackend.FLASH_ATTENTION):
    o = F.scaled_dot_product_attention(q, k, v)       # (8, 16, 8192, 64)

# The naive reference the fused kernel is checked against:
def naive_attention(q, k, v):
    s = q @ k.transpose(-2, -1) / (q.shape[-1] ** 0.5)  # (B,H,L,L) in HBM!
    return torch.softmax(s, dim=-1) @ v

o_ref = naive_attention(q[:1, :1].float(), k[:1, :1].float(),
                        v[:1, :1].float())
assert torch.allclose(o[:1, :1].float(), o_ref, atol=2e-2)
import jax
import jax.numpy as jnp

# jax.nn.dot_product_attention takes (batch, L, heads, d) and dispatches
# to the cuDNN fused flash kernel on Hopper with implementation="cudnn".
key = jax.random.PRNGKey(0)
kq, kk, kv = jax.random.split(key, 3)
q = jax.random.normal(kq, (8, 8192, 16, 64), dtype=jnp.bfloat16)
k = jax.random.normal(kk, (8, 8192, 16, 64), dtype=jnp.bfloat16)
v = jax.random.normal(kv, (8, 8192, 16, 64), dtype=jnp.bfloat16)

@jax.jit
def attn(q, k, v):
    return jax.nn.dot_product_attention(q, k, v, implementation="cudnn")

o = attn(q, k, v)                                  # (8, 8192, 16, 64)

# Same computation, materializing form (the one that OOMs at L=16384):
@jax.jit
def naive(q, k, v):
    s = jnp.einsum("bqhd,bkhd->bhqk", q, k) / jnp.sqrt(64.0)
    return jnp.einsum("bhqk,bkhd->bqhd", jax.nn.softmax(s, axis=-1), v)

Fusion without writing kernels

The measured fusion numbers, 0.911 ms unfused to 0.214 ms fused (\(4.26\times\)) for a four-op elementwise chain on \(2^{26}\) elements, and a hand-written fused CUDA version of the same chain at 0.205 ms (2616.6 GB/s), establish that the compilers now reach hand-kernel bandwidth for elementwise fusion. This is the roofline's memory term collapsed by a decorator.

import torch

def chain(x):                       # 4 elementwise ops: eager launches 4
    y = x * 1.5 + 0.5               # kernels, 8 extra HBM round trips of
    y = torch.nn.functional.gelu(y) # intermediates
    y = y * torch.sigmoid(y)
    return y + x                    # residual

fused = torch.compile(chain)        # Inductor emits ONE Triton kernel:
                                    # read x once, write result once
x = torch.randn(1 << 26, device="cuda")
# measured on this H100: eager 0.911 ms -> compiled 0.214 ms (4.26x);
# hand-written fused CUDA kernel of the same chain: 0.205 ms, 2616.6 GB/s
# (tolerance: Inductor's gelu lowering differs by ~1e-6 in fp32)
assert torch.allclose(fused(x), chain(x), atol=1e-5)
import jax
import jax.numpy as jnp

def chain(x):
    y = x * 1.5 + 0.5
    y = jax.nn.gelu(y)
    y = y * jax.nn.sigmoid(y)
    return y + x

fused = jax.jit(chain)              # XLA fuses the whole chain into one
                                    # kernel; same 2-pass traffic as the
x = jnp.ones((1 << 26,))            # hand-written CUDA version
print(fused(x)[0])                  # first call compiles, later calls run
# Inspect the fusion decision directly:
print(jax.jit(chain).lower(x).compile().as_text()[:400])

Profiling, and what the counters actually say

Every number on this page was produced by timing a kernel and dividing first-principles bytes or flops by the result, which is the measurement a profiler cannot get wrong. A profiler earns its place at the next question, namely, given that a kernel is slower than its roof, which resource is responsible. Two tools divide the work. Nsight Systems is a timeline. It records the whole process, showing CPU threads, CUDA API calls, kernel executions, memory copies, and NCCL collectives on a common time axis. It answers the questions that live between kernels, which are usually the important ones, and it is where launch gaps, missing overlap, a data loader starving the GPU, and a collective that failed to hide under the backward pass all become visible. Nsight Compute is a microscope on one kernel. It replays the kernel while sampling hardware performance counters and produces a per-kernel report with the roofline position, the memory-hierarchy traffic at every level, and the reasons warps were not issuing.

Four families of counter carry nearly all the diagnostic value. Achieved occupancy is the time-averaged ratio of resident warps to the hardware maximum, and its use is comparative. If it is far below the theoretical occupancy computed from registers and shared memory, the cause is usually load imbalance or a tail effect where most blocks have finished, not a resource limit. Memory throughput is reported per level, and the informative quantity is the ratio between the levels, bytes requested by the kernel, bytes actually transferred from L2, and bytes actually transferred from DRAM. A large gap between requested and transferred at the DRAM level is the sector overfetch that the strided-read table measured. The profiler names it as a low number of sectors per request. Warp stall reasons attribute each cycle a warp spent not issuing. Long scoreboard means waiting on a global-memory load, the signature of a latency-bound or bandwidth-bound kernel. Short scoreboard means waiting on shared memory, which is what the tiled matmul on this page would show. Barrier means waiting at __syncthreads, indicating imbalance within a block. MIO throttle and LG throttle mean the load/store issue queues are full, which is a sign that the kernel is issuing too many small accesses rather than fewer wide ones. Instruction mix closes the loop by saying what fraction of issued instructions were the arithmetic the kernel exists to perform, as opposed to address computation, predication, and index arithmetic. A matmul in which fewer than half the instructions are FFMA or MMA is spending its issue slots on bookkeeping, which is exactly the problem TMA was built to remove.

Reading a profile is a fixed sequence, and following it prevents the most common failure, which is optimizing the wrong thing convincingly. Start on the timeline, not in the kernel, and confirm the kernel in question is actually a meaningful fraction of the step, since a kernel that is 3 percent of the time cannot repay more than 3 percent. Compute the kernel's intensity by hand and place it on the roofline. The profiler's own roofline chart is a cross-check, not the source of truth, because it counts bytes at the DRAM boundary and therefore credits cache hits in a way that flatters the kernel. If the kernel is near its roof, stop, because the only remaining move is a different algorithm with a different intensity, which is precisely what fusion and FlashAttention are. If it is far below, read the stall reasons to name the resource, then check the corresponding traffic ratio to confirm. Only then change code, and re-measure with the same first-principles arithmetic rather than the profiler's summary, because a profiler's replay perturbs clocks and caches in ways that make its absolute times unreliable. The measured progressions on this page each took exactly this path. reduce0's stall reasons are barrier and short scoreboard, reduce4's are long scoreboard, and a kernel whose dominant stall is long scoreboard while running at 99 percent of streaming bandwidth is finished.

How it is done in practice

The division of labor in a production stack

No production system hand-writes most of its kernels. The division of labor is stable. Matmuls go to cuBLAS or CUTLASS-generated kernels (the measured gap on this page, 25.1 custom versus 51.2 cuBLAS fp32 and 702.4 fp16, is why). Attention goes to purpose-built fused kernels (FlashAttention, cuDNN's fused path, or a Triton equivalent). Collectives go to NCCL. The long tail of elementwise, normalization, and reshaping work goes to compilers (torch.compile's Inductor emitting Triton, XLA for JAX), whose fusion decisions recover the measured \(4.26\times\) automatically. Hand-written kernels are reserved for the residue, new primitives the libraries do not know (a new attention variant, a new quantization scheme), and the highest-value inner loops where the last 20 percent pays for the engineering. The skill that stays scarce is not writing CUDA. It is the roofline diagnosis that decides which bucket a slow operator belongs in, using exactly the arithmetic this page practices.

Scaling out, the same laws, one level up

Multi-GPU training re-runs every idea on this page with the network as the memory system. Data parallelism is a grid-stride loop over devices with an all-reduce as the tree reduction. Its bucketing (tens of megabytes, per the measured \(m_{1/2} \approx 10\) MB) and its overlap of gradient communication with backward compute are the fusion and latency-hiding arguments again. Tensor parallelism partitions the matmul tiles across devices, and its all-gathers are coalescing at rack scale. Pipeline parallelism is the classical instruction pipeline with microbatches as instructions and bubble fraction as the new Amdahl term. Its efficiency has a closed form worth deriving once. With \(p\) stages and \(m\) microbatches, a stage is busy for \(m\) microbatch-times, while the pipeline occupies \(m + p - 1\) microbatch-times end to end (\(p-1\) to fill, \(m\) to stream, and the drain overlaps the fill of the backward direction), so the fraction of stage-time wasted is \((p-1)/(m+p-1)\), the same hyperbola as \(S(p)\). At \(p = 8\) and \(m = 8\) that is 47 percent idle. At \(m = 64\) it is 10 percent, which is why interleaved and zero-bubble schedules spend their complexity budget on raising the effective \(m\). The measured NVLink numbers set the intuition for why this matters. The 328.3 GB/s between two H100s is a tenth of on-package HBM bandwidth, and cross-node InfiniBand is another order of magnitude down, so the parallelism hierarchy (tensor parallel within a node, data and pipeline parallel across) is the memory hierarchy argument with new constants.

Determinism, testing, and the operational realities

Three practical habits separate working kernel engineering from folklore. First, correctness harnesses compare against a slow reference at fp64 or fp32 with explicit tolerances (the measured kernels on this page carry max-abs-error columns, 2.9e-6 for the fused layernorm, 8.9e-7 for the fp32 flash kernel), because a fast wrong kernel is worse than a slow right one and tensor-core accumulation order makes bitwise comparison meaningless. Second, timing discipline means warmup runs to absorb JIT and clock ramp, medians over many runs, cudaEvent or synchronized wall clocks, and bandwidth or TFLOPS derived from first-principles byte and flop counts, never from a profiler's guess. Third, reproducibility has a price tag. Atomics-based reductions (including the histogram and reduction kernels above) are nondeterministic in float, and frameworks expose deterministic modes that swap in tree-ordered alternatives at measurable cost. Training-infrastructure teams choose per workload, and debugging a loss-curve divergence starts by ruling this in or out.

The current research frontier

Hardware-software co-design past the SIMT model

Hopper cracked the classical CUDA abstraction open, and the frontier is the race to program what leaked out. The Tensor Memory Accelerator (TMA) makes bulk shared-memory copies asynchronous and address-generation-free, warpgroup matmul instructions (WGMMA) issue from a quarter-SM at a time, and thread-block clusters let blocks on different SMs share distributed shared memory. Exploiting all three requires producer-consumer kernel structures (some warps feed tiles, others compute) that look more like dataflow pipelines than SPMD, and the systems that do it, FlashAttention-3 from Colfax, Meta, NVIDIA, and Princeton collaborators, NVIDIA's CUTLASS 3 with its CuTe layout algebra, and ThunderKittens from Hazy Research at Stanford, define the current state of the art. ThunderKittens's bet is that a small C++ embedded library of register and shared-memory tile types recovers 90 percent of hand-tuned performance with a tenth of the code. CUTLASS's bet is that a complete layout algebra composes to 100 percent, and Triton's bet is that a compiler should own the problem. Blackwell's fifth-generation tensor cores and fp4/fp6 microscaling formats raise the same questions again with a wider gap between peak and portable.

Compilation and the end of the kernel boundary

A second line of work attacks kernel boundaries themselves. Persistent "megakernel" designs keep one grid resident and stream work through it, removing launch gaps that are pure Amdahl serial fraction at small batch. Hazy Research's megakernel prototypes for low-latency LLM inference and the grouped/persistent GEMMs inside DeepSeek's DeepGEMM are the visible examples, and DeepSeek's FlashMLA shows the attention variant co-designed with the KV-cache compression it serves. On the compiler side, Triton (now the backend of torch.compile) is being met by JAX's Pallas, which exposes the same tile-level model portably across GPU and TPU, and by MLIR-based efforts industry-wide. The open contest is whether scheduling decisions (pipelining depth, warp specialization) stay human-written, become autotuned search as in the long line of work descending from Halide and TVM, or are learned outright, with AlphaTensor-style algorithm discovery and LLM-driven kernel generation (NVIDIA and academic groups have both reported Triton kernels written and verified by models) as the aggressive end. Communication is being fused too. Overlapping tensor-parallel all-gathers with the matmuls that consume them, in-network reduction in NVLink switches (SHARP), and compiler-scheduled communication in DeepSeek's DualPipe indicate that the roofline's two roofs are becoming three, with the network as a first-class term.

Open source to read

Each entry names the file to open first, and together they cover every layer this page discussed.

  • NVIDIA/cutlass. GPU matmul at full performance, in the open. Read media/docs/efficient_gemm.md first, the best short document on hierarchical tiling anywhere, then the CuTe layout tutorial under media/docs/cute/ to see the Hopper-era abstraction.
  • NVIDIA/cccl. CUB and Thrust live here now. Open cub/cub/block/block_reduce.cuh to see the reduction progression of this page productionized, then cub/cub/agent/agent_scan.cuh for decoupled look-back, the single-pass scan.
  • triton-lang/triton. The tile-level kernel language. Open python/tutorials/03-matrix-multiplication.py and 06-fused-attention.py. The tutorials are the canonical readable implementations of this page's two capstone kernels.
  • Dao-AILab/flash-attention. The production fused attention kernel. Open csrc/flash_attn/src/flash_fwd_kernel.h and find the online softmax rebase from this page's derivation in the inner loop, then the Hopper-specific code under hopper/ for the FlashAttention-3 structures.
  • HazyResearch/ThunderKittens. Tile-type kernel programming from the Hazy Research group at Stanford. Open kernels/attn/h100/h100.cu for a fused attention kernel in the tile idiom, shorter than seems plausible.
  • NVIDIA/nccl. The collectives measured in the alpha-beta section. Open src/device/all_reduce.h to see ring and tree all-reduce as actual device code, then src/graph/search.cc for how topology becomes rings.
  • rapidsai/cudf. Data-parallel primitives applied to dataframes, scans, scatters, and hash joins at GPU bandwidth. Open cpp/src/join/hash_join.cu to watch the histogram and scan patterns of this page do database work.
  • pytorch/pytorch. The reference implementations every kernel on this page was checked against. Open aten/src/ATen/native/cuda/SoftMax.cu to compare a production softmax with this page's version, then aten/src/ATen/native/cuda/Reduce.cuh for the templated reduction machinery and torch/_inductor/ for the compiler that produced the measured 4.26x fusion.
  • vllm-project/vllm. Attention as an inference system rather than a kernel. Open csrc/attention/attention_kernels.cuh, the paged-attention kernel that reads keys and values through a block table instead of a contiguous tensor. The online-softmax rebase of this page's derivation is in its inner loop, with the indirection that makes a fragmented KV cache possible layered on top.
  • open-mpi/ompi. The collectives of the message-passing section as portable CPU code. Open ompi/mca/coll/base/coll_base_allreduce.c and read the ring, recursive-doubling, and segmented-ring variants side by side, each with the cost regime it is selected for.
  • ROCm/composable_kernel. The same tiling theory expressed for AMD's CDNA architecture, which is the best way to see which parts of this page are physics and which are NVIDIA-specific. Open include/ck/tensor_operation/gpu/device/device_gemm.hpp and compare its tile descriptors with CUTLASS's.
  • ispc/ispc. The SPMD-on-SIMD compiler for CPUs. Open examples/cpu/mandelbrot/mandelbrot.ispc next to its serial C++ twin. The diff is exactly the SPMD contract, and the masking machinery the compiler generates is the CPU's version of warp divergence.

Common misconceptions

"Amdahl's law says parallel computing has a low ceiling." Amdahl's law fixes the problem size, and problems refuse to stay fixed. Gustafson's accounting shows linear scaled speedup when the parallel work grows with the machine, which is the regime of essentially all large-scale training and simulation. The correct use of Amdahl is diagnostic and local. Profile the serial residue (launch gaps, data loading, the optimizer step) and compute the ceiling it imposes at your actual scale, as in Problem 1.

"Maximize occupancy." Occupancy is sufficient warps to hide latency, and past that point it competes with the resources that create per-thread performance. The measured blocked128 kernel runs at 25 percent occupancy and beats the 100-percent-occupancy tiled32 by \(3.1\times\), because 112 registers per thread buy an \(8\times 8\) register tile worth more than 48 extra resident warps. Volkov documented this in 2010 and it remains the most durable piece of GPU performance lore.

"Coalescing means aligned access." Coalescing is about which 32-byte sectors a warp's 32 addresses touch, not about alignment per se. The measured strided reads degrade smoothly (2380.5 to 1856.3 to 1218.3 GB/s at strides 1, 2, 4) long before any alignment rule is violated, and the floor at large stride (325.2 GB/s, \(1/7\) of contiguous) is set by the 4-useful-of-32-fetched-bytes ratio. The layout question to ask is always what sectors one warp's instruction touches.

"Shared memory is a cache." Shared memory is an explicitly managed scratchpad with a bank structure, and treating it as a transparent cache forfeits both of its advantages. It does nothing unless the kernel stages data into it (the tiling transformations above), and it penalizes layouts a cache would forgive. The measured transpose gains \(1.58\times\) (1702.9 to 2694.0 GB/s) from one column of padding whose only purpose is to spread accesses across banks.

"Atomics are slow, so avoid them." Atomics are slow under contention and nearly free without it. The measured histogram improved \(144.6\times\) not by removing atomics but by privatizing them. Per-block shared-memory histograms take the contended traffic, and the global atomics that remain (256 per block) are rare. The reduction kernels end with a single atomicAdd per block for the same reason. Count expected conflicts, not atomic instructions.

"FlashAttention is an approximation, or saves flops." It is exact to within rounding (measured max abs error 1.2e-4 in fp16 against the materializing reference) and performs the same \(4L^2d\) flops, slightly more with recomputation in the backward pass. Everything it saves is memory traffic. The measured \(27.68\times\) at \(L = 8192\) comes from never writing the 17.18 GB score matrix. Confusing IO savings with flop savings leads to wrong predictions everywhere else too. Recomputation strategies win because flops are the cheap resource.

"Work-efficient algorithms are always the right choice." Problem 5's arithmetic shows the crossover. Oversubscribed machines favor work-efficient algorithms (Blelloch scan, \(6\times\) at \(n/p = 32\)), while saturated ones favor step-efficient algorithms (Hillis-Steele wins at \(p = n\) despite \(9.5\times\) the work). Production scans are hybrids for exactly this reason, and the same logic governs choices like bitonic-in-block versus radix-across-device sorting.

"A ring all-reduce is always the right collective." Ring is bandwidth-optimal, moving \(2(N-1)/N \le 2\) times the buffer through each link regardless of \(N\), and latency-pessimal, paying \(2(N-1)\alpha\) that grows linearly in the rank count. Problem 4's arithmetic puts the crossover at 12.6 MB for a plausible interconnect. Below it a tree is \(6.9\times\) faster, above it a ring is \(6.0\times\) faster. NCCL switches automatically, which is the actual reason to call a library rather than implement the ring that every tutorial shows.

"Fusing more operations is always better." Fusion removes intermediate memory traffic, which is free money only while the kernel is memory-bound. Fusing into a kernel that is already compute-bound adds register pressure and can cut occupancy below what latency hiding needs, and fusing across a reduction forces either a grid-wide synchronization or a recomputation. The measured elementwise chain gained \(4.26\times\) because it lived at \(I \approx 1\). The fused attention block at 406.6 TFLOPS would gain nothing from further fusion because its bottleneck moved to the tensor cores. Fusion is a roofline move, so check which roof first.

"Speedup numbers speak for themselves." A speedup has three free parameters, the baseline, the scale, and the denominator of the serial fraction. Against a parallel code on one thread instead of the best serial code, against a naive kernel instead of the library call (this page's flash kernel is \(27.7\times\) over naive but \(1.15\times\) over SDPA on the same shape), or with Gustafson's \(s\) quoted where Amdahl's \(f\) is implied (Problem 2's factor of 120), the same measurement supports opposite conclusions. Demand all three parameters before believing any factor.

Self-check

References

  1. Hennessy, J. and Patterson, D. Computer Architecture: A Quantitative Approach, 6th ed., Morgan Kaufmann, 2017.
  2. Hwu, W.-m., Kirk, D., and El Hajj, I. Programming Massively Parallel Processors: A Hands-on Approach, 4th ed., Morgan Kaufmann, 2022.
  3. Herlihy, M., Shavit, N., Luchangco, V., and Spear, M. The Art of Multiprocessor Programming, 2nd ed., Morgan Kaufmann, 2020.
  4. Amdahl, G. "Validity of the Single Processor Approach to Achieving Large Scale Computing Capabilities," AFIPS Spring Joint Computer Conference, 1967. doi:10.1145/1465482.1465560
  5. Gustafson, J. "Reevaluating Amdahl's Law," CACM 31(5), 1988. doi:10.1145/42411.42415
  6. Brent, R. "The Parallel Evaluation of General Arithmetic Expressions," JACM 21(2), 1974. doi:10.1145/321812.321815
  7. Hillis, W. D. and Steele, G. L. "Data Parallel Algorithms," CACM 29(12), 1986. doi:10.1145/7902.7903
  8. Blelloch, G. "Prefix Sums and Their Applications," Technical Report CMU-CS-90-190, Carnegie Mellon University, 1990. cs.cmu.edu/~guyb/papers/Ble93.pdf
  9. Valiant, L. "A Bridging Model for Parallel Computation," CACM 33(8), 1990. doi:10.1145/79173.79181
  10. Culler, D., Karp, R., Patterson, D., Sahay, A., Schauser, K., Santos, E., Subramonian, R., and von Eicken, T. "LogP: Towards a Realistic Model of Parallel Computation," PPoPP, 1993. doi:10.1145/155332.155333
  11. Nickolls, J., Buck, I., Garland, M., and Skadron, K. "Scalable Parallel Programming with CUDA," ACM Queue 6(2), 2008. doi:10.1145/1365490.1365500
  12. Harris, M. "Optimizing Parallel Reduction in CUDA," NVIDIA Developer Technology, 2007. developer.download.nvidia.com/.../reduction.pdf
  13. Volkov, V. "Better Performance at Lower Occupancy," GPU Technology Conference, 2010, and Understanding Latency Hiding on GPUs, PhD thesis, UC Berkeley, 2016.
  14. Williams, S., Waterman, A., and Patterson, D. "Roofline: An Insightful Visual Performance Model for Multicore Architectures," CACM 52(4), 2009. doi:10.1145/1498765.1498785
  15. Pharr, M. and Mark, W. R. "ispc: A SPMD Compiler for High-Performance CPU Programming," InPar, 2012. doi:10.1109/InPar.2012.6339601
  16. Merrill, D. and Garland, M. "Single-pass Parallel Prefix Scan with Decoupled Look-back," NVIDIA Technical Report NVR-2016-002, 2016.
  17. Milakov, M. and Gimelshein, N. "Online Normalizer Calculation for Softmax," 2018. arXiv:1805.02867
  18. Rabe, M. and Staats, C. "Self-attention Does Not Need O(n^2) Memory," 2021. arXiv:2112.05682
  19. Dao, T., Fu, D., Ermon, S., Rudra, A., and Ré, C. "FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness," NeurIPS, 2022. arXiv:2205.14135
  20. Dao, T. "FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning," 2023. arXiv:2307.08691
  21. Shah, J., Bikshandi, G., Zhang, Y., Thakkar, V., Ramani, P., and Dao, T. "FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision," 2024. arXiv:2407.08608
  22. Tillet, P., Kung, H. T., and Cox, D. "Triton: An Intermediate Language and Compiler for Tiled Neural Network Computations," MAPL, 2019. doi:10.1145/3315508.3329973
  23. Spector, B., Arora, S., Singhal, A., Fu, D., and Ré, C. "ThunderKittens: Simple, Fast, and Adorable AI Kernels," 2024. arXiv:2410.20399
  24. NVIDIA. NVIDIA H100 Tensor Core GPU Architecture (Hopper whitepaper), 2022.
  25. NVIDIA. CUDA C++ Programming Guide, version 12.8, 2025. docs.nvidia.com/cuda/cuda-c-programming-guide
  26. Blumofe, R. and Leiserson, C. "Scheduling Multithreaded Computations by Work Stealing," FOCS, 1994, and JACM 46(5), 1999. doi:10.1145/324133.324234
  27. Patarasuk, P. and Yuan, X. "Bandwidth Optimal All-reduce Algorithms for Clusters of Workstations," Journal of Parallel and Distributed Computing 69(2), 2009. doi:10.1016/j.jpdc.2008.09.002
  28. Thakur, R., Rabenseifner, R., and Gropp, W. "Optimization of Collective Communication Operations in MPICH," IJHPCA 19(1), 2005. doi:10.1177/1094342005051521
  29. Chan, T., Golub, G., and LeVeque, R. "Algorithms for Computing the Sample Variance: Analysis and Recommendations," The American Statistician 37(3), 1983. doi:10.1080/00031305.1983.10483115
  30. Batcher, K. "Sorting Networks and Their Applications," AFIPS Spring Joint Computer Conference, 1968. doi:10.1145/1468075.1468121
  31. Dennard, R., Gaensslen, F., Yu, H.-N., Rideout, V., Bassous, E., and LeBlanc, A. "Design of Ion-implanted MOSFETs with Very Small Physical Dimensions," IEEE Journal of Solid-State Circuits 9(5), 1974. doi:10.1109/JSSC.1974.1050511
  32. Kwon, W., Li, Z., Zhuang, S., Sheng, Y., Zheng, L., Yu, C., Gonzalez, J., Zhang, H., and Stoica, I. "Efficient Memory Management for Large Language Model Serving with PagedAttention," SOSP, 2023. arXiv:2309.06180
  33. Zhu, X., Chen, W., Zheng, W., and Ma, X. "Gemini: A Computation-Centric Distributed Graph Processing System," OSDI, 2016. usenix.org/conference/osdi16
  34. Ben-Nun, T. and Hoefler, T. "Demystifying Parallel and Distributed Deep Learning: An In-depth Concurrency Analysis," ACM Computing Surveys 52(4), 2019. arXiv:1802.09941
  35. Chen, T., Moreau, T., Jiang, Z., Zheng, L., Yan, E., Cowan, M., Shen, H., Wang, L., Hu, Y., Ceze, L., Guestrin, C., and Krishnamurthy, A. "TVM: An Automated End-to-End Optimizing Compiler for Deep Learning," OSDI, 2018. arXiv:1802.04799
Key takeaway. Parallel computing is two ledgers kept honestly, the algorithm's ledger of work and depth, and the machine's ledger of bytes moved and flops available. Amdahl and Gustafson bound what parallelism can buy. Brent's bound, recovered decentrally by work stealing, says a greedy schedule collects nearly all of it if the parallelism exists. The ring all-reduce's \(2(N-1)/N\) says what it costs to spread across machines, and everything after that is memory. On this page's H100, one measured lesson repeated at every level. The reduction gained \(8.5\times\) purely from access patterns, the transpose gained \(1.58\times\) from one column of shared-memory padding, the matmul gained \(3.1\times\) by moving reuse from shared memory into registers at a quarter of the occupancy, a full causal multi-head attention block gained \(17.8\times\) at 13.7 times less memory, and attention alone gained \(27.7\times\) at \(L = 8192\) by refusing to write a 17 GB matrix that was never the answer to anything. The roofline, calibrated here at 3063.5 GB/s and 744.6 bf16 TFLOPS with a ridge at 243 flops per byte, is the one-line summary of all of it. Know a kernel's intensity, know which roof it lives under, and spend effort only on transformations that move the binding resource, because on modern machines the flops were almost never the problem.