Reduction 6: Online Softmax, One Read Instead of Three

Lesson

**This stage optimizes something different from every stage before it.**

Stages 3 and 4 attacked *work*: contention on one address, then idle lanes. Your stage 5 kernel has neither problem. Every thread is busy, nothing is serialized, and there is no arithmetic left to remove. What is left is *traffic*. The row crosses global memory four times, and three of those crossings read the same 4 KB.

The blocker to fusing them looks fundamental: you cannot compute `sum(exp(x - m))` before you know `m`. The online normalizer dissolves it with one identity. Maintain a running pair `(m, d)` over a stream of values, and when a new value `x` arrives:

```
mNew = max(m, x)
d = d * exp(m - mNew) + exp(x - mNew)
m = mNew
```

The `d * exp(m - mNew)` factor retroactively corrects every term you already accumulated, exactly as if you had known the true max all along. It costs one multiply per update, and it buys you the ability to finish the sum in the same pass that discovers the max.

Two properties make it work in a real kernel:

- **It merges.** The same rule combines two partial pairs, so the block's tree combine folds `(mA, dA)` and `(mB, dB)` the same way it folded plain sums. This is the tree from stage 3 again, with a pair instead of a scalar.
- **It composes at any granularity.** Threads, blocks, tiles: the rule does not care how the stream was split.

That leaves the normalize pass, which still needs the original values. The row is 1024 floats, 4 KB, so stage it in shared memory on the way past and normalize from that copy. Global traffic goes from 3 reads + 1 write to 1 read + 1 write.

One caution about how you judge it. The wall clock may barely move, because the L2 cache absorbs some of the redundant reads at this size, and a faster kernel that still moves the same bytes has not actually saved anything on a bigger problem. So this lab grades the SASS sector count, which counts the traffic the hardware really issued. **Reducing arithmetic and removing traversals are different optimizations and need different instruments.** That is the last idea in the path, and it is the one that transfers furthest.

Metric contract: algorithmic reads/writes count source-level sweeps. GLOBAL_MEMORY_SECTORS counts SASS-derived issued/requested sectors in the selected profiling scope, including repeated launches. It is cache-blind: it is neither L2 traffic nor device-memory bytes. Trace elapsed time is a separate measurement; a global load may hit cache.

One input read here requires a row cache that fits the supported row size. Online normalization alone does not imply arbitrary-length softmax with constant storage and one read.

What you will learn

• Distinguish removing arithmetic work from removing memory traversals, and pick the metric that measures each.
• Derive and apply the online-normalizer update: rescale the running sum when the running max changes.
• Merge two (max, sum) partial pairs with the same rescaling rule inside a tree combine.
• Use a shared-memory row cache so the normalize pass re-reads staged data, not global memory.
• Read the improvement in GLOBAL_MEMORY_SECTORS rather than the wall clock, and know why.

Starter code

    // The 3-pass parallel softmax you built in the previous lab: stable and
    // parallel, but the row crosses global memory three times before the write.
    __shared__ float red[256];

    int row = blockIdx.x;
    if (row >= rows) return;
    int tid = threadIdx.x;
    const float* src = X + row * cols;
    float* dst = Y + row * cols;

    float localMax = -3.402823466e+38f;
    for (int c = tid; c < cols; c += blockDim.x) localMax = fmaxf(localMax, src[c]);
    red[tid] = localMax;
    __syncthreads();
    for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {
        if (tid < stride) red[tid] = fmaxf(red[tid], red[tid + stride]);
        __syncthreads();
    }
    float m = red[0];
    __syncthreads();

    float localSum = 0.0f;
    for (int c = tid; c < cols; c += blockDim.x) localSum += expf(src[c] - m);
    red[tid] = localSum;
    __syncthreads();
    for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {
        if (tid < stride) red[tid] += red[tid + stride];
        __syncthreads();
    }
    float denom = red[0];

    for (int c = tid; c < cols; c += blockDim.x) {
        dst[c] = expf(src[c] - m) / denom;
    }