IIntroduction to AI Systems Performance EngineeringTechnical textbook
Chapter 8: Parallel Reduction, Scan, and Histograms
8.1

Hierarchical reduction

Trace and implement a hierarchical reduction from thread-local accumulation through grid completion

Before you begin

Follow a linked prerequisite when you cannot explain it from memory. Background skills without links are stated as abilities rather than hidden terminology.

  • Warps, blocks, barriers
  • Associative operators

A language-model training program computes one scalar loss from millions of token losses. Later in the same training step, an optimizer computes one gradient norm from millions or billions of parameter-gradient values. In a PyTorch profile, operations such as aten::sum, aten::amax, and vector_norm launch GPU work whose output contains far fewer values than the input. Nsight Systems shows the launch on the timeline, while Nsight Compute can report the input bytes read, the elapsed microseconds, and whether a second launch combines partial answers.

Follow a smaller example before considering millions of values. Eight CUDA threads each load one number. Four pairs can add at the same time, then two pairs can add, and one final pair produces the sum. The shrinking number of partial answers explains the term reduction: the program reduces eight input values to one output value. Real kernels first create partial answers inside each thread because register arithmetic is cheaper than immediate communication.

A reduction turns many values into fewer values under an associative operation: sum, maximum, norm, or a tuple of statistics. The output is small, but every input must contribute. A serial loop expresses the dependency as one long chain. A parallel reduction reshapes it into a tree so independent partial results can form simultaneously and merge through progressively wider scopes.

The reduction tree specifies communication at each hardware scope. Values begin in thread-local registers, combine across lanes, combine across warps in a block, and may require another kernel or atomic operation across blocks. Each wider step costs synchronization and data movement. The implementation should combine partial results before it moves them to a wider scope.

Use addition as the first example, but keep the operator separate from the schedule. The same schedule can reduce maxima, boolean conditions, or a pair such as {sum, count} when the combine operation is associative and has an identity value. Separating operator from schedule explains why libraries can provide general reduction primitives without knowing the application that produced the inputs.

From a serial chain to independent partial results

An operator ⊕ is associative when (a⊕b)⊕c = a⊕(b⊕c), and an identity e satisfies e⊕a = a⊕e = a. Associativity permits the tree to change the grouping of operations. Commutativity, a⊕b = b⊕a, is a separate property that permits reordering. A reduction for an associative but noncommutative operator must keep the original left-to-right order when it assigns and merges partial ranges.

Consider eight values and a sum. A serial loop computes (((((((x0 + x1) + x2) + x3) + x4) + x5) + x6) + x7). Every addition depends on the previous result, so the dependency depth is seven. A balanced tree first computes x0+x1, x2+x3, x4+x5, and x6+x7 at the same time. The tree then combines adjacent pairs and finally combines the last two partial sums. The work is still seven additions, but the dependency depth is three.

The ideal tree assumes one worker per pair and inexpensive communication. A GPU implementation modifies that tree to match hardware scopes. Each thread usually loads and accumulates several elements first. Serial work inside one thread is useful because the work amortizes address calculation and exposes independent memory requests. Only the compact thread partials need to enter the cooperative tree. A grid-stride assignment can reorder elements across thread partials, so use the assignment only when the operator is commutative or when a different range assignment preserves the required order.

Warp and block reduction

Within a warp, a shuffle instruction lets each participating lane read a register value from another lane. In a downward shuffle reduction, lane 0 eventually receives the reduction. The mask is a correctness contract: all non-exited lanes named by the mask must execute the same intrinsic with the same mask. The mask does not make a read from an inactive source lane valid. A simple full-warp path therefore keeps every lane participating and initializes lanes without input data to the identity. A true partial-warp path needs both the exact mask and a bound that prevents any lane from reading a nonparticipant; a tested library primitive is safer than a casual mask substitution.

A block contains one or more warps, so the block needs one additional level. A common implementation writes one partial per warp to shared memory, executes a block barrier, and assigns one warp to reduce those partials. The barrier is required because a warp that consumes shared memory must not run before the producer warps have completed their stores. A block of 256 threads produces only eight warp partials, so the final stage is small.

Completing a reduction across blocks

Ordinary CUDA thread blocks may execute in any order and cannot use a normal block barrier to synchronize the complete grid. A block therefore writes one partial result to global memory. A second kernel reduces those block partials. The two-launch design is simple because the first kernel's completion creates the global ordering required by the second kernel.

An atomic final update can be faster when the first stage produces only a modest number of block partials or when many independent output reductions distribute contention. The atomic path becomes unattractive when every input performs an update to one address. Cooperative grid synchronization and persistent kernels are additional options, but they add launch restrictions and lifecycle complexity. Select among these methods from the number of partials, the surrounding pipeline, and the required ordering guarantees.

Floating-point order and reproducibility

Mathematical addition is associative, but finite-precision floating-point addition is not. For example, adding a very small value to a very large value can round the small contribution away. A different reduction tree groups values differently and can produce a different final bit pattern. An atomic reduction can also vary because block completion order is not fixed.

State the numerical contract before selecting an implementation. Many ML workloads accept an error tolerance and prefer FP32 accumulation for FP16 or BF16 inputs. Some scientific or debugging workloads require deterministic results, which usually means a fixed reduction tree and a fixed launch policy. Compensated summation or wider accumulation can reduce error, but each method adds instructions or storage. Compare error against a higher-precision reference over inputs with large cancellation and varied magnitudes.

Worked example: Reducing a large vector to one sum

A BF16 vector contains 67 million values. The result should have stable FP32-quality accumulation and good bandwidth efficiency.

  1. Have each thread load several contiguous or grid-stride values and accumulate them in FP32 registers. Thread-local accumulation reduces 67 million inputs to one partial per thread without shared communication.
  2. For a 256-thread block, choose a grid-stride loop that lets consecutive lanes read consecutive input elements. Check the final iteration with a bounds predicate and use zero as the identity for missing elements.
  3. Reduce within each warp using shuffle operations. Store one FP32 partial per warp in shared memory, synchronize the block, and have one warp combine those partials.
  4. Write one result per block. Launch a small second reduction over those results rather than issuing millions of atomics to one address.
  5. If the first launch produces only a few hundred partials, compare the second launch with one atomic addition per block. The comparison must include both launches and not only the first kernel.
  6. Compare against a higher-precision reference across adversarial value scales. Measure actual HBM bytes and determine whether time approaches a streaming-read bound.

Conclusion: The hierarchy follows hardware scope: dense reads, private accumulation, warp exchange, block exchange, and a tiny global second pass. Every widening occurs only after data has been compressed.

Common errors

  • Loading one value per thread and immediately synchronizing at every tree level. Local accumulation and warp primitives remove much of that overhead.
  • Using a single global atomic for every element. Correctness can survive while contention serializes the kernel.
  • Demanding bitwise equality with a serial sum without stating why. Parallel order changes rounding; define the numerical requirement explicitly.
  • Using __shfl_down_sync with a mask that includes inactive lanes. The reduction can incorporate undefined values at a partial-warp boundary.

Reference summary

A reduction combines N values with an associative operator and an identity value. Associativity permits regrouping: (a⊕b)⊕c = a⊕(b⊕c). Associativity does not permit arbitrary reordering when the operator is not commutative. Threads accumulate input ranges in registers, warps combine thread partials, blocks combine warp partials through shared memory, and a later kernel or bounded atomic stage combines block outputs. The ownership scheme must preserve any required input order.

cuda
value = has_input ? input : 0.0f;  // additive identity
for (int offset = warpSize / 2; offset > 0; offset /= 2)
  value += __shfl_down_sync(0xffffffff, value, offset);
The full-warp sketch leaves the total in lane 0. Tail lanes still execute and contribute the identity. A partial participating warp needs an exact participation mask and bounds that prevent reads from inactive source lanes; use a proven warp primitive when possible.
  • Load multiple elements per thread to amortize indexing and increase instruction-level parallelism.
  • Use a neutral element for masked tails.
  • Preserve numerical requirements: summation order changes floating-point results.
  • Fuse a cheap map operation into the reduction when it removes an intermediate tensor.
  • For huge reductions, choose multi-kernel staging or a carefully bounded atomic final step.

Knowledge check

Why is one global atomic add per input element usually worse than one per block? When could atomics still be competitive?

Show the answer guide
  • One global atomic add per input element sends every input to the same destination. The updates serialize at that address, and the kernel pays the global atomic cost once for every element.
  • A block reduction first combines values in registers, warp shuffles, and shared memory. One atomic add per block therefore reduces the number of contending global updates from the number of input elements to the number of blocks.
  • Atomics can still be competitive when there are few blocks, when each block performs substantial work before its update, when updates are spread across many destinations, or when a second reduction launch would cost more than the remaining contention.
  • The decision must be measured for the target input size, datatype, destination count, and GPU because atomic throughput and launch overhead vary across devices.
Practice

Measure the crossover between per-element and per-block atomics

Write two CUDA or Triton sum reductions for FP32 arrays of 2^10 through 2^28 elements. In the first kernel, issue one global atomic add per input element. In the second, reduce within 256-thread blocks and issue one atomic add per block. Add a third two-pass reduction if the framework makes that practical. Check every result against a CPU or framework reference, then profile at least six array sizes after warm-up.

Evidence to keep: Submit the kernel code, correctness tests with stated tolerances, a benchmark matrix containing input size and median latency for every implementation, and one profiler report that records atomic transactions or serialization behavior. State the smallest size at which block aggregation wins and give a causal conclusion based on the measurements.

A reduction collapses information. A scan preserves one partial result for every position, making dependency management more subtle.

References

Sources and further reading

The lesson explanations synthesize primary papers and official references. Follow the originals for proofs, exact APIs, architecture details, and version-sensitive behavior.

  1. CUDA C++ Programming Guide↗
  2. CUDA C++ Best Practices Guide↗
  3. Nsight Compute Profiling Guide↗
  4. NVIDIA Nsight Systems User Guide↗

    Official reference for capturing and reading CPU, CUDA API, GPU workload, and communication timelines.

  5. NVIDIA CUTLASS↗