Reduction — The Heartbeat of Matrix Math

Reduction is GPU compression

A reduction collapses many parallel values into one (or a small handful). Sum, max, mean, norm, dot product, softmax denominator — all reductions. The reason it deserves its own track lesson: every neural network is full of them, and getting reduction wrong wastes 50%+ of your hot-path time.

The winning recipe is hierarchical:

Registers — each thread accumulates its own partial sum in a register (no memory traffic).
Warp-level shuffle — combine 32 partials within a warp using __shfl_down_sync. Zero shared-memory traffic.
Shared-memory tree — each warp writes its sum to one shared-memory slot, then a tree reduction across slots gives the block sum.
Block result — either an atomic add to a global accumulator, or a second kernel launch that reduces block sums.

Where LLMs spend their reduction budget:

Attention softmax — find row max + sum of exp(x - max) per row. Two reductions per attention head.
LayerNorm / RMSNorm — mean and (root-)mean-square per token, every layer.
Loss computation — billions of logits → one scalar.
Gradient accumulation — micro-batch gradients summed before optimizer step.

Code

Warp-level reduction with shuffle (CUDA)·cuda

// Sum across 32 threads in a warp without shared memory.
__device__ float warp_reduce_sum(float v) {
    for (int off = warpSize / 2; off > 0; off >>= 1) {
        v += __shfl_down_sync(0xffffffff, v, off);
    }
    return v;   // thread 0 of the warp now holds the warp's sum
}

__global__ void block_sum(const float* in, float* out, int N) {
    extern __shared__ float warp_sums[];   // one slot per warp in the block
    int tid = threadIdx.x;
    int gid = blockIdx.x * blockDim.x + tid;
    int lane = tid % warpSize;
    int warp = tid / warpSize;

    float v = (gid < N) ? in[gid] : 0.f;
    v = warp_reduce_sum(v);                // warp sums in registers

    if (lane == 0) warp_sums[warp] = v;    // first lane writes shared
    __syncthreads();

    // First warp reduces the warp sums
    if (warp == 0) {
        v = (lane < blockDim.x / warpSize) ? warp_sums[lane] : 0.f;
        v = warp_reduce_sum(v);
        if (lane == 0) atomicAdd(out, v);  // contribute to global
    }
}

SIMD-group reduction in Metal (the Apple equivalent)·metal

#include <metal_stdlib>
using namespace metal;

// simd_sum is the built-in equivalent of warp shuffle reduction.
// 32 threads in a SIMD-group share their value, return the sum.
kernel void block_sum(const device float *in       [[buffer(0)]],
                      device       float *out      [[buffer(1)]],
                      constant     uint  &N        [[buffer(2)]],
                      threadgroup  float *warp_sums [[threadgroup(0)]],
                      uint gid                     [[thread_position_in_grid]],
                      uint tid                     [[thread_index_in_threadgroup]],
                      uint simd_lane               [[thread_index_in_simdgroup]],
                      uint simd_id                 [[simdgroup_index_in_threadgroup]])
{
    float v = (gid < N) ? in[gid] : 0.f;
    v = simd_sum(v);                       // SIMD-group sum
    if (simd_lane == 0) warp_sums[simd_id] = v;
    threadgroup_barrier(mem_flags::mem_threadgroup);

    if (simd_id == 0) {
        v = (simd_lane < /* num warps */ 8) ? warp_sums[simd_lane] : 0.f;
        v = simd_sum(v);
        if (simd_lane == 0) atomic_fetch_add_explicit(
            (device atomic_float *)out, v, memory_order_relaxed);
    }
}

Exercise

Implement block_sum as above on whichever GPU you have, and benchmark vs a naive 'one thread reads all N elements' implementation. The hierarchical version should be 50–100× faster on large inputs because it actually uses parallelism. Then time it vs the equivalent in NumPy / PyTorch / MLX — you'll usually be within 2× of the framework, which means your hierarchical reduction is competitive.

Reduction — The Heartbeat of Matrix Math

Reduction is GPU compression

Code

External links

Exercise

Progress

Comments 0