Softmax kernel

Three passes over one row

One block per row, 256 threads, row length N:

__shared__ float red[256];
int row = blockIdx.x, tid = threadIdx.x;
float m = -INFINITY;
for (int j = tid; j < N; j += blockDim.x) m = fmaxf(m, x[row * N + j]);
m = block_max(m, red);                    // gpu-08 tree, max instead of add
float s = 0.f;
for (int j = tid; j < N; j += blockDim.x) s += __expf(x[row * N + j] - m);
s = block_sum(s, red);
for (int j = tid; j < N; j += blockDim.x)
  y[row * N + j] = __expf(x[row * N + j] - m) / s;

Subtract the max first

expf overflows to infinity for arguments above roughly 88 in fp32. A row of logits in the hundreds overflows on the first pass, and the ratio of two infinities is NaN. Subtracting the row maximum makes every argument at most zero, so every exponential lands in (0, 1] and the denominator is a positive number that cannot overflow. The maximum is not an optimisation; without it the kernel is wrong on exactly the inputs that motivated it. __expf accuracy and the effect of the fast-math flag are documented per release; confirm in the Math API for the toolkit you build with.

Three reads of the row

Each pass reads x again, so a row that does not fit in registers costs three global reads. When N is at most blockDim.x, keep the value in a register between passes and read once. When it is larger, the kernel is bandwidth bound and the exponentials are free by comparison.

Things that break

A fully masked row of -inf: the maximum is -inf, x - m is NaN, and the whole row is NaN. Decide the convention and test it. Reusing the reduction buffer without a barrier between the max pass and the sum pass lets one warp overwrite it while another still reads. __expf is the fast intrinsic with reduced accuracy, and -use_fast_math can substitute it without telling you, so compare against a tolerant reference and record the flag.