Reduction 5: Stable Softmax, Two Row Reductions

The tree you built in stages 3 and 4 was never really about addition. Softmax needs the same tree twice per row, once with fmaxf and once with +, and the max one is there to keep exp from overflowing rather than to make anything faster. This is the reduction pattern arriving at a real Transformer kernel.

Languages: CUDA C++, Rust, HIP C++

Lesson

**Start with the max, and start with why.**

The textbook formula is `softmax(x)_i = exp(x_i) / sum(exp(x_j))`. Written that way it does not survive contact with real data. `expf` overflows float32 once its argument passes about 88, and this harness feeds you a row spanning [-50, +99.5]. `expf(99.5)` is `inf`, `inf / inf` is `NaN`, and your output is ruined.

Subtracting the row maximum fixes it exactly:

```
exp(x_i - m) / sum(exp(x_j - m))
```

The m cancels top and bottom, so this is the same function, evaluated safely. Every exponent is now at most 0, the largest term is exactly `1.0`, and 1024 of them cannot overflow. This is not a rounding nicety you can skip when in a hurry: the correctness check grades it, and a kernel without the subtraction returns `NaN` and fails. **Numerical stability is part of correctness here, not an optimization.**

That requirement is what makes softmax two reductions rather than one:

```
row max -> keeps exp() safe (operator: fmaxf)
row exp sum -> builds the denominator (operator: +)
normalize -> writes probabilities
```

And this is the stage's real lesson. The tree you wrote in stage 3 was never specifically about addition. It is a shape: pair up, combine, halve, repeat. Swap `+` for `fmaxf` and the identical structure computes a maximum. The operator only has to be associative and have an identity element, which for max is negative infinity, the same role `0.0f` played for sum.

The starter runs all three passes serially in thread 0 while the other 255 threads idle. Keep the block-per-row layout, which is standard for row kernels because rows are independent, and give each of the 256 threads 4 of the 1024 columns. Then run the tree twice.

One synchronization subtlety, and it is the kind that produces intermittent wrong answers rather than a crash: the same shared scratch array serves both trees. A barrier must separate *reading* the max out of `red[0]` from *overwriting* `red[]` with partial sums, or a fast thread clobbers the max while a slow one is still reading it.

The gate is a speedup against a baseline the harness runs alongside your kernel: the same serial one-thread-per-row code as the starter, timed on the same GPU in the same run. You need 4x.

Checkpoint the maximum, shifted exponential sum, and normalized output in that order. Do not reuse shared scratch until every participating thread has finished its preceding read.

What you will learn

Explain why softmax must subtract the row max before exponentiating, and why that is correctness rather than polish.
Recognize the reduction tree as a shape parameterized by an associative operator and its identity, not as a sum.
Map one block to one row and cover the row with a block-stride loop.
Run two shared-memory tree reductions with different operators (fmaxf, +) in one kernel.
Place the barrier that protects reused shared scratch between the two trees.

Starter code

    // Correct and stable, but serial: thread 0 owns the whole row while the
    // other 255 threads of this block do nothing.
    int row = blockIdx.x;
    if (row < rows && threadIdx.x == 0) {
        const float* src = X + row * cols;
        float m = src[0];
        for (int c = 1; c < cols; ++c) m = fmaxf(m, src[c]);
        float sum = 0.0f;
        for (int c = 0; c < cols; ++c) sum += expf(src[c] - m);
        for (int c = 0; c < cols; ++c) Y[row * cols + c] = expf(src[c] - m) / sum;
    }