Write Cuda Softmax Kernel

Skill: Write CUDA Softmax Kernel

Purpose

Guide the agent through designing and implementing a correct, numerically stable CUDA softmax kernel, covering online (single-pass) computation, row-parallel decomposition, warp-level reductions, fp16/bf16 precision pitfalls, masked softmax variants, and when to fuse with attention versus implementing standalone.

Use this when

• You need softmax along the last dimension of a 2D or 3D tensor and need a custom kernel for fusion or layout reasons
• You are implementing masked softmax (e.g., causal attention mask, padding mask) where the mask pattern is not supported by existing library routines
• You need to fuse softmax with the subsequent matrix multiply in an attention kernel (flash attention pattern) to avoid materializing the full attention score matrix
• You are targeting a specific hardware or latency budget where you need to control the decomposition precisely
• The input shape (sequence length, number of heads) does not match the assumptions of available library softmax implementations

Do not use this when

• Standard softmax on well-shaped inputs with no custom masking: cuDNN cudnnSoftmaxForward and cudnnSoftmaxBackward are highly optimized for common attention shapes
• The softmax is part of a standard multi-head attention block: use FlashAttention-2 (or equivalent) which fuses QK^T, softmax, and AV into a single tiled kernel with O(seqlen) memory instead of O(seqlen^2)
• The sequence dimension is very small ( 1024: multiple thread blocks per row, requiring a multi-block reduction (two-pass or online with global memory synchronization). This is complex; consider restructuring or blocking at a higher level if possible.

1. Choose online (single-pass) vs two-pass formulation.

• Two-pass: first pass computes the row maximum, second pass computes sum(exp(x - max)), third pass normalizes. Requires two or three reads of the input row. Simple to implement correctly.
• Online (single-pass): computes max and sum in a single scan using the online softmax update rule (Milakov & Gimelshein, 2018). As each new element is seen, if it exceeds the current max, the running sum is rescaled: new_sum = old_sum * exp(old_max - new_max) + exp(x - new_max). Requires one read of the input row (for forward pass; backward still needs two passes). Use online for memory-bandwidth-limited situations.

For most attention-scale softmax (seqlen 0; offset >>= 1) localmax = fmaxf(localmax, _shfldownsync(0xffffffff, localmax, offset)); `` Broadcast the warp max to all lanes: rowmax = __shflsync(0xffffffff, localmax, 0);`. If the block has multiple warps, store warp maxes to shared memory, sync, reduce the smem array in the first warp.

1. Compute the shifted exponentials and sum. Each thread computes exp(x_i - row_max) for each element it owns. Sum these locally, then perform a warp+block reduction for the total sum (same pattern as step 4 but with addition).

1. Normalize. Each thread divides its exp(x_i - row_max) by the total sum and writes to the output. This is a second pass over the elements. If using online softmax, this normalization is implicit in the update rule.

1. Handle masked softmax. For additive masking (adding -inf or a large negative value to masked positions): apply the mask before the max computation. After subtracting max and computing exp, the masked positions contribute exp(-inf - max) = 0 to the sum. For boolean masks: convert to additive mask by mask ? 0.0f : -INFINITY before the exp computation. Special case: if all positions in a row are masked, the sum is 0 and division produces NaN. Decide how to handle this — common choices are: output 0, output 1/D (uniform), or output NaN (let the caller handle it). Document the choice.

1. Handle fp16/bf16 inputs. For fp16 inputs: perform max reduction and exp-sum accumulation in fp32. Convert each fp16 element to fp32 before the subtraction and exp. The final output can be written as fp16 after division. Never compute hexp in fp16 with a direct exp(__half) in the warp reduction — fp16 has very limited dynamic range and exponents in [-87, 88] in fp32 translate to [-10, 10] in fp16 before saturation.

1. Implement the backward pass if needed. The softmax backward computes dL/dx_i = softmax_i * (dL/dy_i - sum_j(softmax_j * dL/dy_j)). This requires the forward softmax output y and gradient dL/dy. The inner sum is a dot product reduction over the row, identical in structure to the forward pass. Reuse the same warp+block reduction pattern.

Kernel design rules

• The subtraction of the row maximum before exp() is non-negotiable. Never compute raw exp(x) without this normalization.
• For inputs in fp16 or bf16: accumulate max, sum, and the dot product (for backward) in fp32. Convert fp16 to fp32 on load, store fp16 on output write.
• Use __expf() (fast math intrinsic) instead of expf() when the application can accept ~2 ULP error in the exp computation. For attention weights, this is almost always acceptable and saves ~20% on the exp computation.
• The warp reduction for max must use fmaxf (float max, propagates NaN to the left — i.e., fmaxf(NaN, x) = x, fmaxf(x, NaN) = NaN). Be aware: if the input contains NaN, the max will be NaN, and the entire row output will be NaN. If NaN inputs are possible, add a NaN guard.
• Thread block size should be a multiple of 32. For rows that fit in one block, 128 or 256 threads is typical. Assign elements to threads using a stride of blockDim.x so that loads are coalesced when threads in the same warp access consecutive elements.
• For the shared memory inter-warp reduction: size the smem array as blockDim.x / 32 elements. Do not over-allocate.
• For masked softmax: if using additive masking with -INFINITY, ensure the identity for the max reduction is initialized to -INFINITY (not 0 or FLT_MIN), so that fully-masked rows produce 0 in the output after exp(−∞ − max) = 0 / sum = 0, with sum handled carefully.

Correctness requirements

• Numerical stability: every code path must subtract the row maximum before computing exp(). This is a mandatory correctness requirement, not just a performance hint.
• Full-row masked outputs: when all elements of a row are masked (sum = 0 after exp), division by zero produces NaN or inf. Decide and implement an explicit fallback (e.g., output zero row, or output 1/row_length uniform). Never silently produce NaN unless the API contract documents it.
• Warp mask correctness: __shfl_down_sync masks must include only active lanes. At the tail of the input (when row_length is not a multiple of 32), threads beyond the row length must use the identity value (−∞ for max, 0 for sum) and the shfl mask must exclude invalid lanes, or be initialized to identity so they do not affect the result.
• Synchronization for inter-warp smem reduction: the write to smem (warp leaders) must be followed by __syncthreads() before any thread reads from smem. The read must be followed by another __syncthreads() before any subsequent smem write.
• Two-pass consistency: in the two-pass approach, the max used in the normalization pass must be the same max computed in the first pass. If multiple blocks are used, the global max must be fully determined (via a barrier or two-kernel approach) before the second pass begins.
• Output bounds: the output write must be predicated to the valid row range. Threads assigned to positions beyond the row length must not write to the output array.
• dtype correctness: the final division and output write must cast back to the original output dtype. Writing fp32 results to an fp16 output buffer requires explicit conversion; an implicit truncation may work but is non-obvious and should be done with __float2half_rn.

Performance requirements

• Softmax is memory-bandwidth-bound for long rows. Target near peak memory bandwidth for the target device.
• For rows that fit in shared memory: minimize the number of global memory reads. The two-pass approach reads each row twice (once for max, once for exp+normalize). The online approach reads once. For large seqlen, this matters; for seqlen = row_length) excluded from global reads and contributing identity values to the reduction?
• [ ] Is there a __syncthreads() after all warp leaders write to smem and before any thread reads from smem?
• [ ] Is the case of a fully-masked row (sum = 0) handled explicitly and documented?
• [ ] For masked softmax: is the mask applied additively before the max computation, not after?
• [ ] Is the output written back in the correct dtype, with explicit conversion if input and output dtypes differ?
• [ ] Is the warp shfl mask correct for all row lengths, including those not divisible by 32?
• [ ] Has the kernel been tested on: rowlength=1, rowlength=32, rowlength=33, rowlength=1024, rowlength with all elements equal, rowlength with a large spread of values, fully-masked row?
• [ ] For the online softmax variant: is the rescaling factor exp(old_max - new_max) applied correctly each time the max is updated?
• [ ] Is CUB, cuDNN, or FlashAttention considered and explicitly noted as the preferred choice for standard cases?

Source & license

This open-source skill is cataloged on AgentStack and links to its original source — we do not rehost the code.

• Author: tensormux
• Source: tensormux/kernel-skills
• License: MIT

Install and usage instructions live in the source repository linked above.