kernels.py

"""
Hand-written Triton kernel: numerically-stable row-wise softmax.

This is the operation used inside self-attention (softmax over the last
dimension of the attention-score matrix), so comparing it here is directly
relevant to LLM inference rather than a toy example. We compare three
implementations on identically-shaped attention-score tensors:

    1. PyTorch eager        torch.softmax(x, dim=-1)
    2. torch.compile        torch.compile(softmax)
    3. Hand-written Triton  triton_softmax(x)  (this file)

`triton_softmax` requires a CUDA GPU to actually launch (Triton JIT-compiles
to a GPU kernel). Importing this module does not require a GPU.
"""
from __future__ import annotations

import time
from typing import List, TypedDict

import numpy as np
import torch
import triton
import triton.language as tl


@triton.jit
def _softmax_fwd_kernel(
    output_ptr,
    input_ptr,
    input_row_stride,
    output_row_stride,
    n_cols,
    BLOCK_SIZE: tl.constexpr,
):
    row_idx = tl.program_id(0)
    row_start_ptr = input_ptr + row_idx * input_row_stride
    col_offsets = tl.arange(0, BLOCK_SIZE)
    input_ptrs = row_start_ptr + col_offsets
    mask = col_offsets < n_cols

    row = tl.load(input_ptrs, mask=mask, other=-float("inf"))

    # numerically stable softmax
    row_minus_max = row - tl.max(row, axis=0)
    numerator = tl.exp(row_minus_max)
    denominator = tl.sum(numerator, axis=0)
    softmax_out = numerator / denominator

    output_row_start_ptr = output_ptr + row_idx * output_row_stride
    output_ptrs = output_row_start_ptr + col_offsets
    tl.store(output_ptrs, softmax_out, mask=mask)


def triton_softmax(x: torch.Tensor) -> torch.Tensor:
    """Row-wise softmax over the last dimension. Any leading dims are flattened."""
    if not x.is_cuda:
        raise RuntimeError("triton_softmax requires a CUDA tensor")

    orig_shape = x.shape
    x2d = x.reshape(-1, orig_shape[-1]).contiguous()
    n_rows, n_cols = x2d.shape

    block_size = triton.next_power_of_2(n_cols)
    num_warps = 4
    if block_size >= 2048:
        num_warps = 8
    if block_size >= 4096:
        num_warps = 16

    out = torch.empty_like(x2d)
    _softmax_fwd_kernel[(n_rows,)](
        out,
        x2d,
        x2d.stride(0),
        out.stride(0),
        n_cols,
        BLOCK_SIZE=block_size,
        num_warps=num_warps,
    )
    return out.reshape(orig_shape)


class SoftmaxVariantResult(TypedDict, total=False):
    variant: str
    mean_latency_ms: float
    std_latency_ms: float
    max_abs_error: float
    error: str


def benchmark_softmax_variants(
    device: str,
    batch: int = 8,
    heads: int = 12,
    seq: int = 512,
    n_warmup: int = 5,
    n_iters: int = 30,
    seed: int = 0,
) -> List[SoftmaxVariantResult]:
    """Compare eager / torch.compile / hand-written Triton softmax on an
    attention-score-shaped tensor (batch, heads, seq, seq)."""
    if device != "cuda":
        return [{"variant": "all", "error": "skipped: requires a CUDA GPU"}]

    torch.manual_seed(seed)
    x = torch.randn(batch, heads, seq, seq, device=device, dtype=torch.float16)
    reference = torch.softmax(x.float(), dim=-1).cpu()

    compiled_softmax = torch.compile(lambda t: torch.softmax(t, dim=-1))
    variants = {
        "eager": lambda t: torch.softmax(t, dim=-1),
        "torch.compile": compiled_softmax,
        "triton": triton_softmax,
    }

    results: List[SoftmaxVariantResult] = []
    for name, fn in variants.items():
        try:
            out = None
            for _ in range(n_warmup):
                out = fn(x)
            torch.cuda.synchronize()

            times_ms = []
            for _ in range(n_iters):
                torch.cuda.synchronize()
                t0 = time.perf_counter()
                out = fn(x)
                torch.cuda.synchronize()
                times_ms.append((time.perf_counter() - t0) * 1000)

            max_abs_error = (out.float().cpu() - reference).abs().max().item()
            results.append(
                {
                    "variant": name,
                    "mean_latency_ms": float(np.mean(times_ms)),
                    "std_latency_ms": float(np.std(times_ms)),
                    "max_abs_error": max_abs_error,
                }
            )
        except Exception as e:  # noqa: BLE001 - want to keep benchmarking other variants
            results.append({"variant": name, "error": str(e)})

    return results