chapter four

4 Reduction Patterns

 

This chapter covers

  • How reduction trees map onto the GPU communication hierarchy
  • Numerical issues in floating-point reductions and how to mitigate them
  • Triton’s reduction primitives and the limits of their synchronization scope
  • Reduce-then-map kernels like RMSNorm and numerically stable softmax

We now turn to one of the most fundamental operations in parallel computing: reduction. Reductions appear throughout machine learning: summing gradients during backpropagation, computing the denominator in softmax, or accumulating statistics for normalization layers. Operations like normalization and probability distributions rely on global statistics of this kind. The denominator of a softmax calculation is a single scalar value, but it is derived from every element in a tensor row. When that row is distributed across thousands of executing threads, no individual thread has enough information to compute its final output using only the data in its local registers. The arithmetic itself forces the threads to cooperate.

4.1 The mechanics of reduction

4.1.1 Floating-point associativity and numerical concerns

4.1.2 Triton’s primitives for reduction

4.1.3 Atomic operations

4.1.4 RMS normalization in Triton

4.2 Numerical stability and chained reductions (Softmax)

4.2.1 The "exploding exponential" problem

4.2.2 Case study: softmax in Triton

4.2.3 Fused softmax kernel

4.3 Summary