Triton 101 - fused row-wise softmax

Open in Colab ▶️ Run this notebook in Colab

Triton 101 - fused row-wise softmax#

Track 07 - GPU · Notebook 02 · Runtime: ≈5 min on any GPU

Prerequisites: 07_gpu/01 (GPU architecture tour).

Reference: Triton tutorials §02 - Fused Softmax.

What#

A Triton kernel that computes numerically stable softmax row-by-row for a 2-D tensor X ∈ R^{M × N}. One program per row; each program loads the whole row into SRAM, subtracts the max, exps, normalises, and writes the result back.

The naive PyTorch implementation is 3-4 kernel launches (max, sub, exp, sum, div). A fused kernel fires once and reuses the values in registers. On memory-bound rows this is the textbook 2-3× speedup.

This is the “hello world” Triton kernel - simplest case where the fusion actually wins.

from llm_systems_cookbook.nb import bootstrap

import time

import torch

IS_CUDA = torch.cuda.is_available()

s = bootstrap("07_gpu_02_triton_101_softmax")

The kernel#

One Triton program per row. BLOCK_SIZE is chosen at call time and must be at least N; we use triton.next_power_of_2(N) so vectorised loads and reductions work cleanly. Masking handles the slack between N and BLOCK_SIZE.

triton = None
tl = None
fused_softmax = None

if IS_CUDA:
    try:
        import triton
        import triton.language as tl

        @triton.jit
        def _softmax_kernel(
            out_ptr, in_ptr,
            in_row_stride, out_row_stride,
            n_cols,
            BLOCK_SIZE: tl.constexpr,
        ):
            row_idx = tl.program_id(0)
            col = tl.arange(0, BLOCK_SIZE)
            mask = col < n_cols

            x = tl.load(in_ptr + row_idx * in_row_stride + col, mask=mask, other=float("-inf"))
            x_max = tl.max(x, axis=0)
            num = tl.exp(x - x_max)
            denom = tl.sum(num, axis=0)
            y = num / denom
            tl.store(out_ptr + row_idx * out_row_stride + col, y, mask=mask)

        def fused_softmax(x: torch.Tensor) -> torch.Tensor:
            M, N = x.shape
            block = triton.next_power_of_2(N)
            out = torch.empty_like(x)
            num_warps = 4 if block >= 1024 else 2
            _softmax_kernel[(M,)](
                out, x,
                x.stride(0), out.stride(0),
                N,
                BLOCK_SIZE=block,
                num_warps=num_warps,
            )
            return out

        print("Triton softmax kernel compiled")
    except Exception as e:  # noqa: BLE001
        print(f"Triton unavailable: {type(e).__name__}: {e}")
        fused_softmax = None

Correctness - two shapes#

Small (128×64) and large (1823×1024) with the non-power-of-2 rows exercising the masked path. Both should match torch.softmax to within FP32 rounding.

def max_abs_err(shape: tuple[int, int]) -> float:
    torch.manual_seed(0)
    x = torch.randn(shape, device="cuda", dtype=torch.float32)
    y_ref = torch.softmax(x, dim=-1)
    y_ours = fused_softmax(x)
    return (y_ref - y_ours).abs().max().item()


# 1e-5 is the realistic FP32 softmax tolerance across GPU families.
# torch.softmax and Triton don't use identical reduction trees, so a
# stricter 1e-6 can flake on some vendors / driver versions.
TOL = 1e-5
if IS_CUDA and fused_softmax is not None:
    err_small = max_abs_err((128, 64))
    err_large = max_abs_err((1823, 1024))
    print(f"max abs err (128,64)   = {err_small:.2e}")
    print(f"max abs err (1823,1024) = {err_large:.2e}")
    s.check(
        "softmax_correct_power_of_two_cols",
        lambda: err_small < TOL,
        msg=f"err = {err_small:.2e}  tol = {TOL:.0e}",
    )
    s.check(
        "softmax_correct_non_power_of_two_cols",
        lambda: err_large < TOL,
        msg=f"err = {err_large:.2e}  tol = {TOL:.0e}",
    )
else:
    s.skip("softmax_correct_power_of_two_cols",        "no Triton / no CUDA")
    s.skip("softmax_correct_non_power_of_two_cols",    "no Triton / no CUDA")

Workload sanity (CPU-safe)#

Even without a GPU, we can verify the arithmetic intensity of row-wise softmax: one forward reads M × N elements, does O(M × N) FLOPs, and writes M × N. Intensity is ~2 FLOPs/byte - firmly memory-bound. That’s why fusion helps: you pay for the bytes once, not three times.

def softmax_intensity(m: int, n: int, dtype_bytes: int = 4) -> float:
    # Approx: 3 ops per element (exp, sum, div), times M*N, over
    # one read + one write of M*N elements.
    return (3 * m * n) / (dtype_bytes * 2 * m * n)


ai = softmax_intensity(1823, 1024)
print(f"softmax arithmetic intensity = {ai:.3f} FLOPs/byte  (memory-bound)")
s.check(
    "softmax_is_memory_bound",
    lambda: ai < 10,
    msg=f"AI = {ai:.3f}",
)

Performance#

Compare against torch.softmax on a large row-wise matrix (1823 × 1024, FP32). Target: ≥1.3× speedup.

def bench(fn, x, n_iter: int = 50) -> float:
    for _ in range(3):
        fn(x)
    torch.cuda.synchronize()
    t0 = time.perf_counter()
    for _ in range(n_iter):
        fn(x)
    torch.cuda.synchronize()
    return n_iter / (time.perf_counter() - t0)


if IS_CUDA and fused_softmax is not None:
    x = torch.randn((1823, 1024), device="cuda", dtype=torch.float32)
    its_torch = bench(lambda x: torch.softmax(x, dim=-1), x)
    its_ours  = bench(fused_softmax, x)
    print(f"torch.softmax  = {its_torch:7.1f} it/s")
    print(f"triton softmax = {its_ours:7.1f} it/s  ({its_ours / its_torch:.2f}x)")
    s.benchmark(
        "triton_softmax_at_least_1_3x_torch",
        measured=its_ours,
        baseline=its_torch,
        must_beat=1.3,
        unit="it/s",
    )
else:
    s.skip("triton_softmax_at_least_1_3x_torch", "no Triton / no CUDA")

Where does the speedup come from?#

One benchmark point isn’t enough. Sweep N across a range of row widths and plot both kernels’ throughput plus the per-width speedup: narrow rows are launch-overhead limited and the fused kernel dominates; very wide rows saturate memory bandwidth and both converge. The shape of the curve - not a single number - is what tells you when fusion is worth the complexity.

import matplotlib.pyplot as plt

if IS_CUDA and fused_softmax is not None:
    Ns = [64, 128, 256, 512, 1024, 2048, 4096]
    M = 1823
    torch_its, ours_its = [], []
    for N in Ns:
        x = torch.randn((M, N), device="cuda", dtype=torch.float32)
        torch_its.append(bench(lambda z: torch.softmax(z, dim=-1), x))
        ours_its.append(bench(fused_softmax, x))
    speedup = [o / t for o, t in zip(ours_its, torch_its, strict=True)]

    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(9, 3.2))
    ax1.plot(Ns, torch_its, "o-", label="torch.softmax")
    ax1.plot(Ns, ours_its,  "o-", label="triton fused")
    ax1.set_xscale("log")
    ax1.set_yscale("log")
    ax1.set_xlabel("cols N")
    ax1.set_ylabel("it/s")
    ax1.set_title("throughput vs row width")
    ax1.legend()
    ax1.grid(True, which="both", alpha=0.3)

    ax2.bar([str(n) for n in Ns], speedup, color="tab:red")
    ax2.axhline(1.0, color="gray", linestyle="--")
    ax2.set_xlabel("cols N")
    ax2.set_ylabel("triton / torch")
    ax2.set_title("speedup (1.0 = tied)")

    fig.tight_layout()
    plt.show()
else:
    print("skipped - no Triton / no CUDA device.")

Exercises#

  1. Autotune num_warps and num_stages over BLOCK_SIZE ∈ {128, 256, 512, 1024, 2048} and report the best on your GPU.

  2. Add a causal mask before the softmax. Compare against torch.softmax(x.masked_fill(causal_mask, -inf), dim=-1).

  3. Implement the backward pass (softmax gradient): given upstream dy, compute dx = y * (dy - (dy * y).sum(dim=-1, keepdim=True)).

References#

  • Triton tutorials, fused softmax.

  • Dao 2022 §2 for the online-softmax version used inside FA1/FA2 (covered in 07_gpu/04).

s.summary()
s.save()