5.1 Optimizing a naive CUDA Softmax

Softmax sits at the center of Attention in Transformers. The formula looks simple, but floating-point overflow is easy to hit. This article starts from the numerical bugs in a naive implementation, then walks through Safe Softmax, block-level parallelism, warp shuffle, and vectorized loads. Each step has the reasoning and a performance comparison, so you can write a Softmax kernel that is correct and efficient.

📑 Contents

• 1. Softmax basics and numerical issues
• 2. Version V0: naive implementation (one thread per row)
• 3. Version V1: Safe Softmax three-pass scan
• 4. Version V2: warp-shuffle reduction
• 5. Version V3: vectorized float4 loads
• 6. Version V4: two-pass fused kernel
• 7. Performance comparison and engineering tips
• Summary
• Self-check list
• References

---

1. Softmax basics and numerical issues

1.1 What Softmax is

Suppose you have a pile of scores and want each score’s share of the total — high scores get a large share, low scores a small one, and the shares add up to exactly 1. Softmax does that: it squeezes a vector of arbitrary reals into a probability distribution.

Definition: for an input vector $\mathbf{x} = (x_0, x_1, \ldots, x_{N-1})$, Softmax outputs:

$$\text{Softmax}(x_i) = \frac{e^{x_i}}{\sum_{j=0}^{N-1} e^{x_j}}$$

In Transformer self-attention, Softmax runs on each row of $QK^T$ and turns raw attention scores into normalized weights. The input length is usually the sequence length $N$ (2048, 4096, even 128K), so Softmax performance feeds straight into overall model speed.

1.2 Numerical overflow

Softmax uses the exponential $e^{x_i}$, which overflows easily in floating point:

• FP32: max representable value is about $3.4 \times 10^{38}$, which is $e^{88.7}$. When $x_i > 88.7$, $e^{x_i}$ becomes +inf
• FP16: max is about $6.55 \times 10^4$, which is $e^{11.09}$. Overflow starts around $x_i > 11$
• Underflow is just as bad: when $x_i$ is very small (e.g. $-100$), $e^{x_i} \approx 0$, the denominator can underflow to 0, and you get NaN

In real Attention, $QK^T$ values routinely exceed 88 (especially with a large $d_k$ and no scaling). A naive implementation will blow up.

1.3 Fix: subtract the max

You can prove that subtracting the same constant $c$ from every element does not change Softmax:

$$\frac{e^{x_i - c}}{\sum_{j} e^{x_j - c}} = \frac{e^{x_i} \cdot e^{-c}}{\sum_{j} e^{x_j} \cdot e^{-c}} = \frac{e^{x_i}}{\sum_{j} e^{x_j}}$$

Pick $c = \max(\mathbf{x})$. Every exponent argument is then $\leq 0$, results sit in $(0, 1]$, and overflow is gone. That is the core of Safe Softmax (numerically stable Softmax).

1.4 Compute flow and the bottleneck

Safe Softmax naturally splits into a three-pass scan:

Each pass reads the whole row. For length $N$ that is $3N$ global-memory reads. Softmax has low arithmetic intensity (about 1–2 FLOPs per load), so it is a classic memory-bound op. The goal is to cut memory traffic and raise bandwidth utilization.

1.5 Test setup

All code here is CUDA 12.x, measured on an A100 80GB SXM4:

Workload: matrix shape $(M, N) = (4096, 4096)$ — 4096 rows of 4096 elements each (a typical Attention-sized case).

---

2. Version V0: naive implementation (one thread per row)

2.1 Idea

Simplest parallel strategy: each thread owns one row of the input matrix and serially does max, exp-sum, and normalize. Rows are independent, so parallelism is free.

2.2 Kernel

// V0: naive — one thread per row
__global__ void softmax_v0(float* input, float* output, int M, int N) {
    int row = blockIdx.x * blockDim.x + threadIdx.x;
    if (row >= M) return;

    float* x = input  + row * N;
    float* y = output + row * N;

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

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

    // Pass 3: normalize
    float inv_sum = 1.0f / sum;
    for (int i = 0; i < N; i++) {
        y[i] = expf(x[i] - max_val) * inv_sum;
    }
}

2.3 Why it is slow

⚠️ Note: V0 is correct but extremely slow, for three reasons:

1. Too little parallelism: one thread per row; with $N=4096$ elements processed serially, most CUDA cores sit idle
2. Uncoalesced access: 32 threads in a warp work on 32 different rows; addresses in the inner loop are $N \times 4$ bytes apart (a full row stride), so loads are badly uncoalesced
3. Redundant work: expf(x[i] - max_val) is computed twice — once in Pass 2 and again in Pass 3

V0 measured bandwidth utilization: ~5% (~102 GB/s) — it runs, but it is not usable.

---

3. Version V1: Safe Softmax three-pass scan

3.1 Idea

V0’s main problem is insufficient parallelism — $N$ elements of a row go through one thread. Fix: let many threads in one block cooperate on the same row. Each thread owns a slice of the row; shared memory does a parallel reduction for max and sum.

That turns Softmax into two reduces (max-reduce + sum-reduce) plus an element-wise normalize. Each reduce uses the same techniques from the Reduce optimization article.

3.2 Kernel

// V1: one block per row, parallel reduction inside the block
__global__ void softmax_v1(float* input, float* output, int M, int N) {
    extern __shared__ float smem[];  // used for reduction

    int row = blockIdx.x;
    int tid = threadIdx.x;

    float* x = input  + row * N;
    float* y = output + row * N;

    // Pass 1: parallel max
    float max_val = -INFINITY;
    for (int i = tid; i < N; i += blockDim.x) {
        max_val = fmaxf(max_val, x[i]);
    }
    smem[tid] = max_val;
    __syncthreads();

    // Shared-memory reduction for the global max
    for (int s = blockDim.x / 2; s > 0; s >>= 1) {
        if (tid < s) {
            smem[tid] = fmaxf(smem[tid], smem[tid + s]);
        }
        __syncthreads();
    }
    max_val = smem[0];
    __syncthreads();

    // Pass 2: parallel exp sum
    float sum = 0.0f;
    for (int i = tid; i < N; i += blockDim.x) {
        sum += expf(x[i] - max_val);
    }
    smem[tid] = sum;
    __syncthreads();

    // Shared-memory reduction for the total sum
    for (int s = blockDim.x / 2; s > 0; s >>= 1) {
        if (tid < s) {
            smem[tid] += smem[tid + s];
        }
        __syncthreads();
    }
    sum = smem[0];
    __syncthreads();

    // Pass 3: normalize and write
    float inv_sum = 1.0f / sum;
    for (int i = tid; i < N; i += blockDim.x) {
        y[i] = expf(x[i] - max_val) * inv_sum;
    }
}

Launch:

int block_size = 256;
int smem_size = block_size * sizeof(float);
softmax_v1<<<M, block_size, smem_size>>>(d_input, d_output, M, N);

3.3 Performance

Gains over V0:

• Much more parallelism: 256 threads share a row; each thread only handles $N/256 = 16$ elements
• Coalesced access: adjacent threads in a block touch adjacent addresses in the row (stride=1), so loads coalesce cleanly

Still open:

• Shared-memory reduction needs several rounds of __syncthreads()
• Three passes mean the input is read three times

V1 measured bandwidth utilization: ~35% (~714 GB/s) — about 7× V0.

---

4. Version V2: warp-shuffle reduction

4.1 Idea

V1 reduces entirely through shared memory. That has two costs: data shuttles register → SMEM → register, and every reduction round needs __syncthreads(), which hurts more once few threads remain.

From Reduce work we know warp shuffle (__shfl_down_sync) swaps values in registers: lower latency, no bank conflict, and warp-intrinsic sync. We switch to a two-level reduce — shuffle inside each warp, then shared memory across warps: each warp writes one partial to a 32-element SMEM array, and the first warp finishes the reduce.

4.2 Kernel

// Intra-warp reduce: Max
__device__ float warpReduceMax(float val) {
    for (int offset = 16; offset > 0; offset >>= 1) {
        val = fmaxf(val, __shfl_down_sync(0xffffffff, val, offset));
    }
    return val;
}

// Intra-warp reduce: Sum
__device__ float warpReduceSum(float val) {
    for (int offset = 16; offset > 0; offset >>= 1) {
        val += __shfl_down_sync(0xffffffff, val, offset);
    }
    return val;
}

// Two-level reduce: block Max
__device__ float blockReduceMaxShuffle(float val) {
    __shared__ float warp_max[32];

    int lane = threadIdx.x % 32;
    int wid  = threadIdx.x / 32;

    val = warpReduceMax(val);

    if (lane == 0) warp_max[wid] = val;
    __syncthreads();

    int num_warps = blockDim.x / 32;
    val = (lane < num_warps) ? warp_max[lane] : -INFINITY;
    if (wid == 0) val = warpReduceMax(val);

    // Broadcast result to all threads
    __shared__ float block_result;
    if (threadIdx.x == 0) block_result = val;
    __syncthreads();
    return block_result;
}

// Two-level reduce: block Sum
__device__ float blockReduceSumShuffle(float val) {
    __shared__ float warp_sum[32];

    int lane = threadIdx.x % 32;
    int wid  = threadIdx.x / 32;

    val = warpReduceSum(val);

    if (lane == 0) warp_sum[wid] = val;
    __syncthreads();

    int num_warps = blockDim.x / 32;
    val = (lane < num_warps) ? warp_sum[lane] : 0.0f;
    if (wid == 0) val = warpReduceSum(val);

    __shared__ float block_result;
    if (threadIdx.x == 0) block_result = val;
    __syncthreads();
    return block_result;
}

// V2: two-level warp-shuffle reduction
__global__ void softmax_v2(float* input, float* output, int M, int N) {
    int row = blockIdx.x;
    int tid = threadIdx.x;

    float* x = input  + row * N;
    float* y = output + row * N;

    // Pass 1: row max
    float local_max = -INFINITY;
    for (int i = tid; i < N; i += blockDim.x) {
        local_max = fmaxf(local_max, x[i]);
    }
    float max_val = blockReduceMaxShuffle(local_max);

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

    // Pass 3: normalize
    float inv_sum = 1.0f / sum;
    for (int i = tid; i < N; i += blockDim.x) {
        y[i] = expf(x[i] - max_val) * inv_sum;
    }
}

4.3 Why warp shuffle is faster

💡 Tip: Warp-shuffle latency is about 1/4 to 1/6 of shared memory (~5 cycles vs ~20–30 cycles), and you skip __syncthreads(). Late in a reduce, when few threads remain, that latency gap and sync saving matter a lot.

V2 measured bandwidth utilization: ~52% (~1060 GB/s) — about 1.5× V1.

---

5. Version V3: vectorized float4 loads

5.1 Idea

V2’s bottleneck has moved to global loads. Each loop iteration loads one float (4 bytes) per thread, so instruction overhead is high relative to data moved. With float4 each instruction moves 16 bytes (4 floats), which:

• Cuts load-instruction count (4× fewer loop iterations)
• Improves memory-transaction efficiency
• Raises instruction-level parallelism (ILP)

5.2 Kernel

// V3: float4 vectorized loads
__global__ void softmax_v3(float* input, float* output, int M, int N) {
    int row = blockIdx.x;
    int tid = threadIdx.x;

    float* x = input  + row * N;
    float* y = output + row * N;

    // Vectorized pointers
    float4* x4 = reinterpret_cast<float4*>(x);
    float4* y4 = reinterpret_cast<float4*>(y);
    int N4 = N / 4;  // number of float4 elements

    // Pass 1: vectorized max
    float local_max = -INFINITY;
    for (int i = tid; i < N4; i += blockDim.x) {
        float4 data = x4[i];
        local_max = fmaxf(local_max, fmaxf(fmaxf(data.x, data.y),
                                            fmaxf(data.z, data.w)));
    }
    // Tail elements when N is not a multiple of 4
    for (int i = N4 * 4 + tid; i < N; i += blockDim.x) {
        local_max = fmaxf(local_max, x[i]);
    }
    float max_val = blockReduceMaxShuffle(local_max);

    // Pass 2: vectorized exp sum
    float local_sum = 0.0f;
    for (int i = tid; i < N4; i += blockDim.x) {
        float4 data = x4[i];
        local_sum += expf(data.x - max_val) + expf(data.y - max_val)
                   + expf(data.z - max_val) + expf(data.w - max_val);
    }
    for (int i = N4 * 4 + tid; i < N; i += blockDim.x) {
        local_sum += expf(x[i] - max_val);
    }
    float sum = blockReduceSumShuffle(local_sum);

    // Pass 3: vectorized normalize
    float inv_sum = 1.0f / sum;
    for (int i = tid; i < N4; i += blockDim.x) {
        float4 data = x4[i];
        float4 result;
        result.x = expf(data.x - max_val) * inv_sum;
        result.y = expf(data.y - max_val) * inv_sum;
        result.z = expf(data.z - max_val) * inv_sum;
        result.w = expf(data.w - max_val) * inv_sum;
        y4[i] = result;
    }
    for (int i = N4 * 4 + tid; i < N; i += blockDim.x) {
        y[i] = expf(x[i] - max_val) * inv_sum;
    }
}

5.3 Gains

Loop iterations per thread drop from $N/\text{blockDim}$ to $N/(4 \times \text{blockDim})$ — 4× less loop and dispatch overhead. float4 loads also hide memory latency better.

⚠️ Note: float4 needs 16-byte alignment. In practice you allocate with cudaMalloc (256-byte aligned) and keep $N$ a multiple of 4 (otherwise fall back to the tail loop).

V3 measured bandwidth utilization: ~65% (~1325 GB/s) — about 1.25× V2.

---

6. Version V4: two-pass fused kernel

6.1 Redundancy in the three-pass scan

V1–V3 read the input three times. For a $(4096, 4096)$ matrix that is $3 \times 4096 \times 4096 \times 4 = 192$ MB. If Pass 2 (exp-sum) and Pass 3 (normalize) merge into one pass, reads drop to 2× (128 MB) — one third less traffic, theoretical speedup 1.5×.

The catch: Pass 3 needs $d$ (the exp sum), which only exists after Pass 2. If we do not store Pass 2 intermediates, only keep $m$ and $d$, then in a second pass compute $e^{x_i - m}$ and divide by $d$ together, we get the fusion.

A pure two-pass design really wants the Online Softmax algorithm (next article). Here we take a middle path — partially fuse Pass 2 and Pass 3: in the second pass each thread parks its $e^{x_i - m}$ in shared memory, reduces to global $d$, then immediately reads from SMEM and writes the normalized result, skipping the third global read.

6.2 Kernel

// V4: two-pass kernel (Pass 1 max, Pass 2 fused exp-sum + normalize)
// Needs extra shared memory to stash exp values
__global__ void softmax_v4(float* input, float* output, int M, int N) {
    extern __shared__ float smem[];  // size = N (stash exp values)

    int row = blockIdx.x;
    int tid = threadIdx.x;

    float* x = input  + row * N;
    float* y = output + row * N;

    float4* x4 = reinterpret_cast<float4*>(x);
    int N4 = N / 4;

    // Pass 1: row max (vectorized + warp shuffle)
    float local_max = -INFINITY;
    for (int i = tid; i < N4; i += blockDim.x) {
        float4 data = x4[i];
        local_max = fmaxf(local_max, fmaxf(fmaxf(data.x, data.y),
                                            fmaxf(data.z, data.w)));
    }
    for (int i = N4 * 4 + tid; i < N; i += blockDim.x) {
        local_max = fmaxf(local_max, x[i]);
    }
    float max_val = blockReduceMaxShuffle(local_max);

    // Pass 2: compute exp, stash it, and accumulate the sum
    float local_sum = 0.0f;
    for (int i = tid; i < N; i += blockDim.x) {
        float exp_val = expf(x[i] - max_val);
        smem[i] = exp_val;  // stash in shared memory
        local_sum += exp_val;
    }
    float sum = blockReduceSumShuffle(local_sum);

    // Read exp from shared memory and write normalized output
    float inv_sum = 1.0f / sum;
    for (int i = tid; i < N; i += blockDim.x) {
        y[i] = smem[i] * inv_sum;
    }
}

6.3 Limits and when to use it

⚠️ Note: V4 needs $N \times 4$ bytes of shared memory for the exp stash. At $N=4096$ that is 16 KB, fine on A100 (164 KB SMEM per SM). For large $N$ (e.g. 32768+), the full row no longer fits, so this scheme fails.

For large $N$, the real fix is Online Softmax — one scan (or two, without stashing the whole row). That is the next article.

V4 measured bandwidth utilization: ~75% (~1529 GB/s) (at $N=4096$) — about 1.15× V3.

---

7. Performance comparison and engineering tips

7.1 Version summary

📌 Key point: At $N=4096$, V4 reaches ~75% bandwidth. Further gains need Online Softmax, which cuts scan count at the algorithm level.

7.2 What to ship

💡 Tip: In real LLM inference stacks (vLLM, TensorRT-LLM, …), Softmax is usually fused with the rest of Attention (Scale, Mask, MatMul) into one big kernel so intermediates never hit global memory. Standalone Softmax tuning is mainly for understanding the idea and for odd standalone cases (e.g. Softmax on a classification head).

---

📝 Summary

From V0 to V4, Softmax optimization follows three threads:

1. Raise parallelism: one thread for a whole row → many threads in a block cooperating, with reductions for max and sum (V0 → V1)
2. Speed up the reduce: naive shared-memory loops → two-level warp shuffle, exchanging data in registers and cutting sync and memory latency (V1 → V2)
3. Cut memory traffic: scalar loads → float4 vectorization → two-pass fusion that drops a repeated read, pushing toward the memory-bound ceiling (V2 → V3 → V4)

The next leap — from three passes to one — needs a redesign at the algorithm level. That is what Online Softmax is for.

---

🎯 Self-check list

• Explain why naive Softmax overflows in FP32/FP16
• Prove that subtracting the max does not change Softmax
• Say what each of the three Safe Softmax passes does and why you need three
• Write CUDA for a block-level parallel max and sum reduce
• Implement warp max-reduce and sum-reduce with __shfl_down_sync
• Walk through the full two-level reduce (intra-warp + inter-warp)
• Explain float4 alignment requirements and where the speedup comes from
• Analyze the shared-memory capacity limit of the two-pass fusion
• Pick a Softmax implementation from $N$
• Use Nsight Compute to measure Softmax bandwidth utilization and find the bottleneck

---

📚 References

• NVIDIA CUDA C++ Programming Guide - Mathematical Functions
• Online normalizer calculation for softmax - Milakov & Gimelshein, 2018
• FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness
• CUDA C++ Best Practices Guide - Memory Optimizations
• 深入浅出GPU优化系列：softmax优化
• NVIDIA cuDNN Documentation - Softmax
• From Online Softmax to FlashAttention