Softmax Implementation

Softmax appears in every attention computation and at the output of every LLM. A naive implementation requires three passes over the data (max, sum, normalize) — reading from HBM three times. The online algorithm does it in one pass, and this same trick is what makes FlashAttention possible.

---

1. The Softmax Operation

$$\text{softmax}(x_i) = \frac{e^{x_i}}{\sum_{j} e^{x_j}}$$

Numerical Stability Problem

Direct computation of $e^{x_i}$ overflows for $x > 88$ (FP32) or $x > 11$ (FP16). Fix: subtract the maximum first.

$$\text{softmax}(x_i) = \frac{e^{x_i - m}}{\sum_{j} e^{x_j - m}}, \qquad m = \max_j x_j$$

The Three-Pass Algorithm

Pass 1: Find max    → m = max(x)        (read all N elements)
Pass 2: Compute sum → s = Σ exp(x_i - m) (read all N elements)
Pass 3: Normalize   → y_i = exp(x_i - m) / s  (read all N, write all N)

Three reads of the entire array from HBM. For attention with $S=4096$ and 32 heads: that’s 3 × 32 × 4096² × 2 bytes = 3 GB of HBM traffic just for softmax.

---

2. V0: Naive (One Thread Per Row)

__global__ void softmaxV0(float* input, float* output, int N) {
    int row = blockIdx.x;
    float* x = input + row * N;
    float* y = output + row * N;

    // Pass 1: max
    float m = -INFINITY;
    for (int i = 0; i < N; i++) m = fmaxf(m, x[i]);

    // Pass 2: sum of exp
    float s = 0.0f;
    for (int i = 0; i < N; i++) s += expf(x[i] - m);

    // Pass 3: normalize
    for (int i = 0; i < N; i++) y[i] = expf(x[i] - m) / s;
}

Problems: One thread does all work for a row (no parallelism). Three sequential passes over HBM.

---

3. V1: Parallel Three-Pass with Block Reduction

Use a full block to process one row in parallel:

__global__ void softmaxV1(float* input, float* output, int N) {
    __shared__ float smax, ssum;
    int row = blockIdx.x;
    int tid = threadIdx.x;
    float* x = input + row * N;

    // Pass 1: Parallel max reduction
    float local_max = -INFINITY;
    for (int i = tid; i < N; i += blockDim.x)
        local_max = fmaxf(local_max, x[i]);

    // Warp shuffle + shared memory reduction for max
    local_max = warp_reduce_max(local_max);
    if (tid % 32 == 0) shared_warp_max[tid/32] = local_max;
    __syncthreads();
    if (tid == 0) {
        float m = -INFINITY;
        for (int i = 0; i < blockDim.x/32; i++) m = fmaxf(m, shared_warp_max[i]);
        smax = m;
    }
    __syncthreads();

    // Pass 2: Parallel exp-sum
    float local_sum = 0.0f;
    for (int i = tid; i < N; i += blockDim.x)
        local_sum += expf(x[i] - smax);

    // Reduce sum (similar pattern)
    // ...
    __syncthreads();

    // Pass 3: Normalize
    for (int i = tid; i < N; i += blockDim.x)
        output[row * N + i] = expf(x[i] - smax) / ssum;
}

Better: Parallelized within each pass. But still three passes over HBM.

---

4. Online Softmax: One Pass

The Key Insight

We can compute max and sum simultaneously by maintaining a running correction:

When we see a new element $x_{n+1}$ that’s larger than our current max $m_n$:

• Update max: $m_{n+1} = x_{n+1}$
• Correct the running sum: $s_{n+1} = s_n \cdot e^{m_n - m_{n+1}} + e^{x_{n+1} - m_{n+1}}$

The correction factor $e^{m_n - m_{n+1}}$ rescales all previous exponentials to the new max.

The Online Algorithm

// Single pass: compute max and sum simultaneously
float m = -INFINITY;  // running max
float s = 0.0f;       // running sum (corrected)

for (int i = 0; i < N; i++) {
    float new_m = fmaxf(m, x[i]);
    s = s * expf(m - new_m) + expf(x[i] - new_m);  // correct and add
    m = new_m;
}

// Final normalization pass (still needed)
for (int i = 0; i < N; i++) {
    y[i] = expf(x[i] - m) / s;
}

Two passes instead of three — and the first pass does max + sum simultaneously. 33% fewer HBM reads.

Why This Matters for FlashAttention

FlashAttention processes attention in tiles. Within each tile, it needs softmax statistics (max and sum) that are updated as new tiles are processed. The online algorithm enables this:

1. Process tile 0: compute partial max $m_0$ and partial sum $s_0$
2. Process tile 1: update $m_1 = \max(m_0, \text{tile\_max}_1)$, correct $s_1 = s_0 \cdot e^{m_0 - m_1} + \text{tile\_sum}_1$
3. Continue for all tiles…

Without online softmax, FlashAttention would need to read the entire attention matrix to find the global max first — defeating the purpose of tiling.

---

5. V2: Parallel Online Softmax

Parallelize the online algorithm across a block:

__global__ void softmaxOnline(float* input, float* output, int N) {
    int row = blockIdx.x;
    int tid = threadIdx.x;
    float* x = input + row * N;

    // Each thread computes local online max and sum
    float m = -INFINITY;
    float s = 0.0f;
    for (int i = tid; i < N; i += blockDim.x) {
        float new_m = fmaxf(m, x[i]);
        s = s * expf(m - new_m) + expf(x[i] - new_m);
        m = new_m;
    }

    // Combine across threads: need to merge (m, s) pairs
    // Use warp shuffle to merge pairs
    for (int offset = 16; offset > 0; offset >>= 1) {
        float other_m = __shfl_down_sync(0xFFFFFFFF, m, offset);
        float other_s = __shfl_down_sync(0xFFFFFFFF, s, offset);
        float new_m = fmaxf(m, other_m);
        s = s * expf(m - new_m) + other_s * expf(other_m - new_m);
        m = new_m;
    }

    // Block-level merge via shared memory (same pattern)
    // ...

    // Final normalize pass
    for (int i = tid; i < N; i += blockDim.x) {
        output[row * N + i] = expf(x[i] - m) / s;
    }
}

The merge operation: when combining two partial results $(m_1, s_1)$ and $(m_2, s_2)$:

$$m = \max(m_1, m_2), \qquad s = s_1 \cdot e^{m_1 - m} + s_2 \cdot e^{m_2 - m}$$

This is associative — so it works with tree reduction, warp shuffle, etc.

---

6. Fused Softmax + Cross-Entropy

In training, softmax is almost always followed by cross-entropy loss. Fusing them:

1. Avoids writing softmax output to HBM (just compute it inline)
2. Uses the numerically stable log_softmax form:

$$\log \text{softmax}(x_i) = x_i - m - \log(s)$$ $$\text{Cross-entropy} = -\log \text{softmax}(x_{\text{target}}) = -(x_{\text{target}} - m - \log(s))$$

And the gradient has the beautiful simple form:

$$\frac{\partial L}{\partial x_i} = \text{softmax}(x_i) - \mathbb{1}[i = \text{target}]$$

PyTorch’s F.cross_entropy uses this fused implementation — you should too.

---

7. Performance Summary

Version	Passes	Key optimization	Relative speed
V0 Naive (single thread)	3	None	1×
V1 Parallel three-pass	3	Block reduction	~10×
V2 Parallel online	2	Online max+sum	~15×
+ Fused kernel	1	Avoid intermediate write	~20×
PyTorch F.cross_entropy	1	All of the above	Production speed

---

Self-Check

• Explain why naive $e^{x}$ overflows and how subtracting max fixes it
• Explain the three-pass algorithm and why three HBM reads is expensive
• Derive the online softmax correction: $s_{n+1} = s_n \cdot e^{m_n - m_{n+1}} + e^{x_{n+1} - m_{n+1}}$
• Explain why online softmax enables FlashAttention’s tiled computation
• Implement the merge operation for combining two (max, sum) pairs
• Explain why fused softmax + cross-entropy is both faster and more numerically stable

---

References

• Online Normalizer Calculation for Softmax (Milakov & Gimelshein, 2018)
• FlashAttention (Dao et al., 2022) — uses online softmax
• PyTorch — CrossEntropyLoss