Loading editor...

Kernel visualization: online softmax

Softmax needs the maximum of a vector before it can compute any element of it:

Subtracting from each value makes the calculation stable (no overflows) without altering the output.

A naive implementation of a softmax kernel involves three passes across the HBM:

Read the entire sequence to calculate .
Read the entire sequence to calculate the denominator .
Finally, calculate the output values.

Online softmax

Online softmax fuses steps 1 and 2 (shown in the overview), thus requiring only two passes across the HBM. It uses the following formula:

with and . Both and are updated on every element, so a single read of the sequence produces the final and .

Define as the denominator over the first elements, normalized against the maximum known so far:

Split off the last term and rewrite the rest against instead of :

The second line is the only step that does any work, and it is just . The factor does not depend on , so it comes out of the sum, and what is left is exactly .

When the maximum does not move, and the factor is , so the update reduces to an ordinary running sum.

The hardware reads the sequence one tile (a contiguous chunk of numbers) at a time, so each update takes in a whole tile at once instead of a single element.

The tiled version of the same formula is:

Whenever the running maximum moves, every term already accumulated in was computed against the old maximum. Multiplying by fixes all of them at once, because . That single multiply is what removes the need for a separate pass to find first.

Step 3 cannot be fused: no output value can be written until is final, and is not final until the last element has been read.

The naive version: three passes

Each of the three loops is a full trip over . Pass 2 cannot start early, because every exponential it computes needs the final .

Pass 1: calculate the max
- , ,
- , .
Pass 2: calculate the denominator
- , ,
- , .
Pass 3:
- , , ,
- , ,

Online softmax: two passes

Pass 1:
- , .
- ,
Pass 2:
- ,
- ,

// One block per row, blockDim.x == TILE lanes cooperating.
// Launch:  softmax_naive<<<rows, TILE>>>(X, Y, N);
__global__ void softmax_naive(const float* X, float* Y, int N) {
    int row  = blockIdx.x;
    int lane = threadIdx.x;              // 0 .. TILE-1
    const float* x = X + row * N;
    float*       y = Y + row * N;

    __shared__ float red[TILE];         // reduction scratch
    __shared__ float m_s, d_s;          // the row's m and d

    // pass 1: max over the whole row
    float m_local = -INFINITY;
    for (int i = lane; i < N; i += TILE) m_local = fmaxf(m_local, x[i]);
    red[lane] = m_local;
    __syncthreads();
    for (int s = TILE / 2; s > 0; s >>= 1) {
        if (lane < s) red[lane] = fmaxf(red[lane], red[lane + s]);
        __syncthreads();
    }
    if (lane == 0) m_s = red[0];
    __syncthreads();
    float m = m_s;

    // pass 2: sum of exp(x - m)
    float d_local = 0.0f;
    for (int i = lane; i < N; i += TILE) d_local += expf(x[i] - m);
    red[lane] = d_local;
    __syncthreads();
    for (int s = TILE / 2; s > 0; s >>= 1) {
        if (lane < s) red[lane] += red[lane + s];
        __syncthreads();
    }
    if (lane == 0) d_s = red[0];
    __syncthreads();
    float d = d_s;

    // pass 3: normalize and write y
    for (int i = lane; i < N; i += TILE)
        y[i] = expf(x[i] - m) / d;
}

// One block per row, blockDim.x == TILE lanes cooperating.
// Launch:  softmax_online<<<rows, TILE>>>(X, Y, N);
__global__ void softmax_online(const float* X, float* Y, int N) {
    int row  = blockIdx.x;
    int lane = threadIdx.x;              // 0 .. TILE-1
    const float* x = X + row * N;
    float*       y = Y + row * N;

    __shared__ float tile[TILE];        // one tile in SRAM
    __shared__ float red[TILE];         // reduction scratch
    __shared__ float m_s, d_s;          // running m and d for the row

    if (lane == 0) { m_s = -INFINITY; d_s = 0.0f; }
    __syncthreads();

    // pass 1: fold each tile into m and d in one read
    for (int t = 0; t < N; t += TILE) {
        tile[lane] = x[t + lane];
        __syncthreads();

        // tile max via tree reduction
        red[lane] = tile[lane];
        __syncthreads();
        for (int s = TILE / 2; s > 0; s >>= 1) {
            if (lane < s) red[lane] = fmaxf(red[lane], red[lane + s]);
            __syncthreads();
        }
        float m_old = m_s;
        float m_new = fmaxf(m_old, red[0]);

        // tile sum of exp(x - m_new) via tree reduction
        red[lane] = expf(tile[lane] - m_new);
        __syncthreads();
        for (int s = TILE / 2; s > 0; s >>= 1) {
            if (lane < s) red[lane] += red[lane + s];
            __syncthreads();
        }

        if (lane == 0) {
            d_s = d_s * expf(m_old - m_new) + red[0];   // correct, then add
            m_s = m_new;
        }
        __syncthreads();
    }

    // pass 2: write y = exp(x - m) / d with the final m, d
    float m = m_s, d = d_s;
    for (int i = lane; i < N; i += TILE)
        y[i] = expf(x[i] - m) / d;
}