Parallel Reduction, Scan, and Histograms

A recurring training example uses a decoder-only model with approximately 1 billion parameters on one server containing eight GPUs. A representative microbatch holds four sequences of 2048 tokens in BF16, and data-parallel workers exchange gradients after backpropagation. One training step repeatedly needs answers such as the largest score in a row, the sum of squared gradients, the position at which a retained token should be written, and the number of tokens assigned to each expert. PyTorch presents those jobs as operations such as max, sum, nonzero, and bincount. A GPU profiler shows that the apparently small operations launch CUDA kernels and sometimes occupy a surprising fraction of the step because millions of threads must exchange or update information.

Chapter 8 studies three reusable patterns behind those kernels. A reduction turns many input values into one result, such as one sum. A prefix scan produces a running result at every input position, which lets a program assign output addresses without a serial counter. A histogram lets many inputs update a smaller collection of counters. The chapter introduces each name after tracing a concrete array through the operation.

NVIDIA Nsight Compute and Nsight Systems provide the empirical evidence used in the chapter. A reduction often appears as a bandwidth-heavy kernel followed by a much smaller second kernel. A contended histogram can show many atomic operations and a large timing change when the input distribution becomes skewed. A scan used for compaction produces a flag array, an offset array, and a final scatter. The observed arrays and timelines make the synchronization rules concrete.

During those experiments, a central processing unit (CPU) launches work on a graphics processing unit (GPU). Large input arrays reside in GPU high-bandwidth memory (HBM) before CUDA threads create smaller partial results on chip. Chapter 9 applies the same hierarchical reasoning to matrix multiplication (GEMM), whose K dimension contains a reduction for every output element.

The implementation proceeds through GPU scopes in increasing size. One CUDA thread first works with values in its registers. The 32 threads in a warp exchange partial results. Warps in one thread block communicate through shared memory and a block barrier. Separate thread blocks usually finish through another kernel launch or a limited number of atomic updates because an ordinary CUDA block barrier cannot synchronize the whole grid. Each lesson states which value exists at each scope and which event makes the value safe to consume.