write-triton-softmax-kernel

Skill: Write a Triton Softmax Kernel

Purpose

Guide the agent through implementing a numerically stable, performant row-wise softmax kernel in Triton. This covers single-program-per-row assignment, online max+sum reduction with tl.max/tl.sum, masking for rows wider than BLOCK_SIZE, fp32 accumulation to avoid overflow and precision loss, and the masked softmax variant for attention.

---

Use this when

• You need a fused softmax that avoids a separate max-reduction kernel pass and a separate division pass — i.e., you want a single kernel that reads each row once.
• You need a masked softmax (attention mask applied before the exp) and PyTorch's built-in path is not fusing the mask application correctly.
• You are fusing softmax with a preceding or following elementwise operation (e.g., scale by 1/sqrt(d_k) before softmax, or multiply output by values V immediately after).
• The row dimension is large enough that a per-row kernel is worthwhile (row_size >= 256 is a reasonable floor). Below this, torch.nn.functional.softmax is likely faster.

---

Do not use this when

• The input is a standard 2D or 3D tensor with no mask and no fusion requirement. torch.nn.functional.softmax backed by cuDNN or torch.compile will handle this efficiently.
• The softmax dimension is across rows rather than within rows (i.e., column-wise softmax). The per-row strategy does not apply without transposing the problem.
• The row size is very small (< 64). A warp-level reduction in CUDA or a fused torch.compile graph is more efficient.
• You need a stable online-softmax for arbitrarily long sequences in a streaming fashion — this requires a more complex multi-block reduction strategy beyond a single-program-per-row approach.

---

Inputs the agent should gather first

Before writing any code, confirm:

1. Input shape — (N, D) or (B, H, N, D) or similar. Which dimension is the softmax applied over?
2. Row dimension size D — is it fixed (compile-time constant), or dynamic? Is it always a power of 2?
3. Input dtype — fp16, bf16, or fp32. Accumulation dtype must be fp32 regardless.
4. Whether a mask is applied — additive mask (large negative values added before softmax) or boolean mask (invalid positions should be treated as -inf)?
5. Whether to fuse a downstream operation — e.g., multiply by a V matrix tile, apply dropout, or write to a specific output layout.
6. Whether the row fits in one BLOCK_SIZE — or whether a loop over multiple blocks per row is needed.
7. Hardware target — relevant for choosing BLOCK_SIZE and deciding whether to use triton.autotune.