Write a fused row softmax in Triton, explain why it is one HBM pass, and say where it stops scaling.

TL;DR: Each Triton program handles one row: it loads the row (masked to the row length, padded to a power-of-two block), computes the row max, subtracts it, exponentiates, sums, divides, and stores. The row lives in registers across the SMs the program occupies, so HBM sees one read and one write: for a 32k × 4k fp32 logits tensor that is 512 MB in and 512 MB out, about 0.35 ms on an H100, against four passes (max, exp, sum, divide) in an unfused version. Numerically, subtracting the max keeps the exponent in range, and the sum accumulates in fp32 even when the input is bf16. The block size is the next power of two above the row width; Triton maps it across the program's warps. The kernel stops scaling when a row no longer fits in registers (tens of thousands of elements): then either a two-pass version (max and sum in one pass over chunks, normalize in a second) or an online softmax that folds the rescaling into a single chunked loop.

How to approach it

Write the kernel and the launch. Explain each line where it matters (masking, the max, fp32 accumulation). Then the byte arithmetic and the expected time. Then the block-size rule and the wide-row limit with the online form. Close with how to validate it.

A strong answer

A typical situation: the interviewer opens a notebook, asks for a softmax over the last dimension of a (M, N) tensor in Triton, then asks why it is faster than torch.softmax in eager mode on some shapes and slower on others.

The kernel:

import torch
import triton
import triton.language as tl

@triton.jit
def softmax_kernel(x_ptr, y_ptr, stride_row, n_cols, BLOCK: tl.constexpr):
    row = tl.program_id(0)                                  # one program per row
    cols = tl.arange(0, BLOCK)                              # BLOCK is a power of two ≥ n_cols
    mask = cols < n_cols
    x = tl.load(x_ptr + row * stride_row + cols, mask=mask, other=-float("inf"))
    x = x.to(tl.float32)                                    # accumulate in fp32 even for bf16 inputs
    row_max = tl.max(x, axis=0)                             # reduction across the block (in registers/shared)
    z = x - row_max                                         # shift for stability; masked lanes stay -inf
    num = tl.exp(z)                                         # exp(-inf) = 0 for the padding
    den = tl.sum(num, axis=0)
    y = num / den
    tl.store(y_ptr + row * stride_row + cols, y.to(y_ptr.dtype.element_ty), mask=mask)

def softmax(x: torch.Tensor) -> torch.Tensor:
    assert x.is_cuda and x.dim() == 2
    M, N = x.shape
    y = torch.empty_like(x)
    BLOCK = triton.next_power_of_2(N)
    num_warps = 4 if BLOCK <= 2048 else (8 if BLOCK <= 8192 else 16)
    softmax_kernel[(M,)](x, y, x.stride(0), N, BLOCK=BLOCK, num_warps=num_warps)
    return y

# check
x = torch.randn(4096, 1000, device="cuda", dtype=torch.bfloat16)
torch.testing.assert_close(softmax(x).float(), torch.softmax(x.float(), dim=-1), atol=1e-2, rtol=1e-2)

Why each line is there: the mask and other=-inf make the padding lanes contribute exp(-inf) = 0 to the sum and never win the max; the cast to fp32 keeps the sum from losing precision over thousands of terms (bf16 has 8 bits of mantissa, so accumulating 4,000 values in it loses most of them); the store casts back to the output dtype. Triton Programming Model explains what a program, a block and tl.arange are and how Triton maps the block across num_warps warps with the reductions done in shared memory.

The byte arithmetic:

tensor: M = 32,768 rows × N = 4,096 columns, fp32 → 512 MB
fused kernel: read 512 MB, write 512 MB → 1 GB; at ~3 TB/s achievable on H100 ≈ 0.35 ms
unfused (max, subtract-exp, sum, divide as separate passes): 4 reads + 3 writes of 512 MB ≈ 3.5 GB → ~1.2 ms
arithmetic: ~5 FLOPs per element × 134M elements ≈ 0.7 GFLOP → microseconds; irrelevant
target: ~90% of a device copy's bandwidth on the same tensor; a copy of 1 GB takes ~0.33 ms, so 0.37 ms is done
bf16 input: half the bytes, half the time; the fp32 accumulation costs nothing extra because it happens in registers
sanity: eager torch.softmax is already a fused kernel for common shapes and hits similar numbers; the Triton version
        wins when it is fused with neighbours (a scale, a mask, a dropout) that eager runs as separate kernels, and
        loses when N is tiny (launch and program overhead per row dominate) or when the eager kernel is better tuned.

Memory-Bound vs Compute-Bound Kernels is why bandwidth is the only number here; Kernel Fusion is the reason the win grows when the softmax absorbs its neighbours.

The block-size rule and the wide-row limit:

BLOCK = next power of two ≥ N; Triton needs power-of-two block shapes for tl.arange; masking handles the rest
num_warps: 4 for rows up to ~2k elements, 8 to 16 for wider; each warp holds BLOCK / (32 × num_warps) elements per
  thread in registers; at BLOCK = 16k and 16 warps that is 32 fp32 values per thread, fine; at BLOCK = 128k it is
  256 per thread, past the register budget (255 per thread on NVIDIA), and the compiler spills
wide rows (N > ~32k): two options
  two-pass: pass 1 loops over chunks of the row computing the running max and, with a rescale, the running sum
    (m_new = max(m, chunk_max); s = s × exp(m − m_new) + Σ exp(chunk − m_new)); pass 2 loops again computing
    exp(x − m) / s and storing → reads the row twice, writes once (1.5× the traffic of the fused version, still far
    better than 7 passes)
  online single pass with a store of unnormalized values and a final scale: same as FlashAttention's trick; the row's
    exponentials are stored with the running max and rescaled at the end, which is a second write; choose by
    measuring
narrow rows (N < 128): one program per row wastes most of a warp; process several rows per program (a 2D block)

FlashAttention Internals is where the online rescaling identity comes from; Occupancy and Register Pressure is the spill limit that sets the wide-row boundary.

Validation: compare against torch.softmax in fp32 with tolerances matched to the dtype; test N at powers of two, one below and one above (masking), N = 1, and a row with a single very large value (the max-subtraction test: without it, exp(1000) overflows to inf and the output is NaN); time with triton.testing.do_bench and report bandwidth as bytes ÷ time against the device copy number.

The reversal condition: if the softmax's input is produced by a GEMM and consumed by another GEMM (attention), the right fusion is not a standalone softmax at all but the attention kernel that never materializes the scores; the standalone fused softmax is for the cases where the logits tensor must exist (a final vocabulary softmax, a loss).

What interviewers probe next

• "Why subtract the max rather than clip?" Softmax is shift-invariant, so subtracting the max changes nothing mathematically and bounds the exponent at 0; clipping changes the result.
• "What if a row is all -inf (fully masked)?" The max is -inf, z is NaN; guard by replacing a -inf max with 0 and emitting zeros or a uniform row, as the product requires.
• "How does Triton do tl.max across a block?" A tree reduction within each warp with shuffles, then across warps through shared memory; the programmer never writes it.
• "How would you fuse a causal mask and a scale?" Load, multiply by the scale, apply tl.where(mask, x, -inf) before the max; no extra passes, which is the point of writing it yourself.

Common mistakes

• Forgetting the max subtraction, and NaNs on large logits.
• Accumulating the sum in bf16.
• A block larger than the register budget for wide rows, and a silent 10× slowdown from spills.
• Padding lanes with 0 instead of -inf and inflating the sum.

Key takeaways

• One program per row, row in registers, max-subtract-exp-sum-divide, one read and one write of HBM.
• fp32 accumulation, -inf padding, power-of-two block with masking, num_warps from row width.
• ~0.35 ms for 512 MB fp32 on an H100; done at ~90% of copy bandwidth.
• Past ~32k-element rows, two-pass or online forms; below ~128, several rows per program.