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(seq_len) memory instead of O(seq_len^2)
• The sequence dimension is very small (< 32): the warp reduction overhead is not worth it; a simple sequential kernel or even a CPU-side computation may be appropriate

Inputs the agent should gather first

• Input shape: exact dimensions — e.g., [batch, heads, seq_len] for attention scores, or [N, D] for a 2D input. Which axis is the softmax axis (almost always the last dimension)?
• Dtype: fp32, fp16, or bf16 for input and output; whether the accumulator (for exp sum and max) must be fp32 regardless of input dtype
• Masking requirements: is there an additive mask (e.g., -inf for invalid positions), a boolean mask, or no masking? Is the mask shape the same as the input or broadcast?
• Downstream operation: is the softmax output consumed by another matrix multiply (attention pattern), or written to memory for a standalone use?
• Sequence length: is it fixed or variable at runtime? Variable lengths require either padding to a max length or a segmented/jagged dispatch
• Hardware target: SM architecture, to determine warp size, available reduction primitives, and shared memory capacity