GPT-2 XL: one region per block

This example compiles each GPT-2 XL transformer block into one ATen region and measures it against PyTorch eager and torch.compile. A transformer block is seventeen weights and one straight line of tensor operations, which is what an ATen region is. The block is written once and @ppy.opt(3) compiles it into a single C++ function. Forty-eight calls to it are GPT-2 XL, and the loop between them is all the Python that is left.

Run it

gpt2.ppy is the same model with the same region. It is GPT-2 XL where there is a device for it and a six-layer stack where there is not, so the runs below are the CPU shape; the comparison table is what it does on the 4090. ppy run compiles the driver through LLVM as well, which has nothing to gain on a program that is a loop and forty-eight calls. The three paths print the same number and the same answer.

python  gpt2.ppy
ppy     gpt2.ppy
ppy run gpt2.ppy

What it prints

python gpt2.ppy

torch 2.14.0+cpu | cuda False
# region active: False
# 6 layers, 12 heads, width 768, batch 1, length 128
mean=0.004 absmean=0.439 meansq=0.301
# forward: 22.246 ms
cached decode matches the whole forward: True

ppy gpt2.ppy

torch 2.14.0+cpu | cuda False
# region active: True
# 6 layers, 12 heads, width 768, batch 1, length 128
mean=0.004 absmean=0.439 meansq=0.301
# forward: 22.646 ms
cached decode matches the whole forward: True

ppy run gpt2.ppy

torch 2.14.0+cpu | cuda False
# region active: True
# 6 layers, 12 heads, width 768, batch 1, length 128
mean=0.004 absmean=0.439 meansq=0.301
# forward: 24.435 ms
cached decode matches the whole forward: True

The block

@ppy.opt(3)
def block(x, ln1_w, ln1_b, q_w, q_b, ..., batch, length, heads, head_dim, width, eps):
    h = torch.layer_norm(x, (width,), ln1_w, ln1_b, eps)
    q = F.linear(h, q_w, q_b).reshape((batch, length, heads, head_dim)).transpose(1, 2)
    k = F.linear(h, k_w, k_b).reshape((batch, length, heads, head_dim)).transpose(1, 2)
    v = F.linear(h, v_w, v_b).reshape((batch, length, heads, head_dim)).transpose(1, 2)
    a = F.scaled_dot_product_attention(q, k, v, is_causal=True)
    merged = a.transpose(1, 2).reshape((batch, length, width))
    attended = torch.add(x, F.linear(merged, o_w, o_b))
    n = torch.layer_norm(attended, (width,), ln2_w, ln2_b, eps)
    f = F.gelu(F.linear(n, fc_w, fc_b), approximate="tanh")
    return torch.add(attended, F.linear(f, proj_w, proj_b))

What it becomes

The dimensions are parameters of the region rather than constants folded into it, so one compiled function serves every batch and every length:

at::Tensor ppy_region_gpt2_block(const at::Tensor& x, const at::Tensor& ln1_w, ...,
                                 int64_t batch, int64_t length, int64_t heads,
                                 int64_t head_dim, int64_t width, double eps) {
    auto h = at::layer_norm(x, {width}, ln1_w, ln1_b, eps);
    auto q = ((at::linear(h, q_w, q_b)).reshape({batch, length, heads, head_dim})).transpose(1, 2);
    auto k = ((at::linear(h, k_w, k_b)).reshape({batch, length, heads, head_dim})).transpose(1, 2);
    auto v = ((at::linear(h, v_w, v_b)).reshape({batch, length, heads, head_dim})).transpose(1, 2);
    auto a = at::scaled_dot_product_attention(q, k, v, {}, 0.0, true);
    auto merged = ((a).transpose(1, 2)).reshape({batch, length, width});
    auto attended = at::add(x, at::linear(merged, o_w, o_b));
    auto n = at::layer_norm(attended, {width}, ln2_w, ln2_b, eps);
    auto f = at::gelu(at::linear(n, fc_w, fc_b), "tanh");
    return at::add(attended, at::linear(f, proj_w, proj_b));
}

The region does not reimplement any of it:

• at::scaled_dot_product_attention is the same call F.scaled_dot_product_attention makes, so it reaches the same flash-attention kernel.
• at::linear reaches the same cuBLAS GEMM.
• Every at:: call goes through the dispatcher. That is why autograd is unchanged, and why the training step below differentiates through the region as it does through the Python.

What the region removes is the ten Python round trips per block: four hundred and eighty per forward pass.

The cache the region keeps

A region may hand back several tensors, which is what a KV cache needs. The block returns its output beside the key and value it just grew, so the cache never leaves the region to be reassembled in Python.

@ppy.opt(3)
def step(x, k_cache, v_cache, ..., batch, length, ..., causal: bool
         ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    ...
    k_all = torch.cat((k_cache, k), 2)
    v_all = torch.cat((v_cache, v), 2)
    a = F.scaled_dot_product_attention(q, k_all, v_all, is_causal=causal)
    ...
    return torch.add(attended, F.linear(f, proj_w, proj_b)), k_all, v_all

It becomes a std::tuple<at::Tensor, at::Tensor, at::Tensor>, which pybind11 hands Python as an ordinary tuple. causal is a bool parameter rather than a constant folded in, and length an int64_t, so one compiled function prefills a prompt under a causal mask and then decodes a token at a time out of the cache it built.

gpt2.ppy prints whether the two agree, because that is the check that matters: the last position of a whole forward pass, and the same position reached one token at a time, are the same logits.

A region may also write into a buffer the caller owns, which is the other way to keep a cache: torch.narrow(k_cache, 2, position, length).copy_(k) fills a slot of a preallocated tensor instead of growing a new one. That is an in-place write through a parameter, so it carries WriteMemory, and a function doing it cannot be @ppy.pure; the checker says so by name. The example keeps the cat form because at a 256-token context the copy is a small part of a step and every column pays it equally. At a context of thousands it would not be.

Compared with PyTorch eager and torch.compile at model scale

	PPy ATen regions	PyTorch eager	torch.compile	torch.compile reduce-overhead
forward b1 s32	8.78 ± 0.12	10.77 ± 0.08	7.46 ± 0.28	13.36 ± 0.02
forward b1 s128	9.24 ± 0.18	11.06 ± 0.08	7.88 ± 0.28	14.11 ± 0.01
forward b1 s512	14.87 ± 0.01	14.90 ± 0.01	14.74 ± 0.07	22.86 ± 0.02
forward b8 s512	95.33 ± 0.30	95.50 ± 0.18	95.55 ± 0.26	102.96 ± 0.66
decode 128 tokens	1204.24 ± 28.34	1484.26 ± 17.63	1057.01 ± 41.66	1658.16 ± 1.84
training step b4 s512	136.62 ± 0.23	136.73 ± 0.21	135.58 ± 0.25	143.53 ± 0.67

Milliseconds per pass of the whole model, bfloat16, best of ten in each of five processes. The decode row is a 128-token prompt and then 128 tokens one at a time out of the cache, timed without the prefill.

Every column is the same forty-eight blocks on the same weights, and all four agree on the answer to three decimals. That is what the agreement is checked on, because the greedy token is not stable under any of this: on weights drawn from a generator the top two logits of a position are usually a tie, and fusing a reduction differently flips it. For the same reason the decode loop feeds the next token of a fixed sequence rather than the model's own argmax. Sampling would send the four columns down four different sequences.

Where the time goes

The rows vary one thing: how much of the wall clock is the Python between operators.

• At b1 s32 the GEMMs are small and reading the weights is about three milliseconds, so the rest is dispatch. Removing it is worth 18% to the region and 31% to torch.compile.
• At b8 s512 the same forty-eight blocks are 95 ms of cuBLAS and the Python is a rounding error. The first three columns are the same number, and they should be.
• The training step is the same story with a backward pass on it.

Decode is the far end of that axis. One token through forty-eight blocks is 6144 region calls for 128 tokens, with almost no arithmetic per call. The region takes 19% off eager, torch.compile 29%, and the two are within 14% of each other. It is the row where removing the interpreter is worth the most, and it is still the row where a fusing compiler wins.

What a region leaves unfused

A region removes the interpreter between the operators and leaves the kernels where they were. Where Inductor's own win is also interpreter overhead, the two land close. Inductor could fuse the two layer_norms, the gelu, and the residual adds, but there is not enough of that in a model that is 98% GEMM for the fusion to show at this size.

The columns do not show what each cost to get. The region is one C++ translation unit, 19 s to build once and cached on disk for every process after. torch.compile spends about 18 s on the first shape's forward graph in every new process, and several minutes to reach all five shapes and the backward with a cold Inductor cache.

The reduce-overhead column

reduce-overhead (the CUDA-graphs mode) is slower here in every row than plain torch.compile, and slower than eager too. It reproduces: the same numbers in its own process with a single shape rather than five, and no warning from Inductor that it declined to capture.

It also does not run a KV cache as written. A cache is a tensor the caller keeps across invocations, which is the memory CUDA graphs reuse. The decode loop fails outright (accessing tensor output of CUDAGraphs that has been overwritten by a subsequent run) until the counterpart announces each invocation with cudagraph_mark_step_begin() and sets cudagraph_trees_generation_cloning = "user_visible", the two things its own error message asks for. The cloning that second switch turns on is part of what the last column costs.

The training step needed the same care for the same reason: gradients are set to None before each backward rather than accumulated, which is what an optimizer does anyway.

Setup

The PPy column is ppy <file> (the staged Python that loads the compiled regions), not ppy run. Compiling the driver natively has nothing to do here: outside the regions the program is forty-eight calls and a loop.

NVIDIA GeForce RTX 4090 24 GB, driver 580.126.20, CUDA 13.0; PyTorch 2.14.0+cu130 on CPython 3.13.15, on a rented Pod. This is not the machine the other tables come from: scripts/compare_docs.py skips this comparison unless PPY_TORCH_CUDA_PYTHON names a PyTorch with a device, so the bench runner leaves it alone, and scripts/cloud/runpod_bench.py is what re-measures it.

Limitations

Everything a region reaches is straight-line: no loop, no branch, and no tensor it did not receive as a parameter.

Read on: PyTorch ATen regions · Plugins: PyTorch · Training with torch

gpt2.ppy and the programs under compare/ are hand-written; there is no .py source and no conversion step.

46_gpt2/gpt2.ppy

"""GPT-2: one ATen region per transformer block.

Every block is seventeen weights and one straight line of tensor operations,
which is exactly what a region is. The Python loop over blocks stays in
Python; what leaves it is the ten dispatcher round trips inside each block.
"""

import time

import ppy
import torch
import torch.nn.functional as F

EPS: float = 1e-5


@ppy.opt(3)
def block(
    x: torch.Tensor,
    ln1_w: torch.Tensor,
    ln1_b: torch.Tensor,
    q_w: torch.Tensor,
    q_b: torch.Tensor,
    k_w: torch.Tensor,
    k_b: torch.Tensor,
    v_w: torch.Tensor,
    v_b: torch.Tensor,
    o_w: torch.Tensor,
    o_b: torch.Tensor,
    ln2_w: torch.Tensor,
    ln2_b: torch.Tensor,
    fc_w: torch.Tensor,
    fc_b: torch.Tensor,
    proj_w: torch.Tensor,
    proj_b: torch.Tensor,
    batch: int,
    length: int,
    heads: int,
    head_dim: int,
    width: int,
    eps: float,
) -> torch.Tensor:
    """One pre-norm block: attention with a causal mask, then the MLP."""
    # A region takes every weight as a parameter of its own: that is what a
    # region is, and seventeen tensors beside six dimensions is the shape of
    # one transformer block.
    # pylint: disable=too-many-arguments,too-many-positional-arguments,too-many-locals
    h = torch.layer_norm(x, (width,), ln1_w, ln1_b, eps)
    q = F.linear(h, q_w, q_b).reshape((batch, length, heads, head_dim)).transpose(1, 2)
    k = F.linear(h, k_w, k_b).reshape((batch, length, heads, head_dim)).transpose(1, 2)
    v = F.linear(h, v_w, v_b).reshape((batch, length, heads, head_dim)).transpose(1, 2)
    a = F.scaled_dot_product_attention(q, k, v, is_causal=True)
    merged = a.transpose(1, 2).reshape((batch, length, width))
    attended = torch.add(x, F.linear(merged, o_w, o_b))
    n = torch.layer_norm(attended, (width,), ln2_w, ln2_b, eps)
    f = F.gelu(F.linear(n, fc_w, fc_b), approximate="tanh")
    return torch.add(attended, F.linear(f, proj_w, proj_b))


@ppy.opt(3)
def step(
    x: torch.Tensor,
    k_cache: torch.Tensor,
    v_cache: torch.Tensor,
    ln1_w: torch.Tensor,
    ln1_b: torch.Tensor,
    q_w: torch.Tensor,
    q_b: torch.Tensor,
    k_w: torch.Tensor,
    k_b: torch.Tensor,
    v_w: torch.Tensor,
    v_b: torch.Tensor,
    o_w: torch.Tensor,
    o_b: torch.Tensor,
    ln2_w: torch.Tensor,
    ln2_b: torch.Tensor,
    fc_w: torch.Tensor,
    fc_b: torch.Tensor,
    proj_w: torch.Tensor,
    proj_b: torch.Tensor,
    batch: int,
    length: int,
    heads: int,
    head_dim: int,
    width: int,
    eps: float,
    causal: bool,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """The same block, keeping its own key and value.

    It hands back three tensors, so the cache never leaves the region to be
    reassembled in Python; `causal` is a parameter rather than a constant, so
    one compiled function prefills a prompt and then decodes a token at a
    time out of what it built.
    """
    # A region takes every weight as a parameter of its own: that is what a
    # region is, and seventeen tensors beside six dimensions is the shape of
    # one transformer block.
    # pylint: disable=too-many-arguments,too-many-positional-arguments,too-many-locals
    h = torch.layer_norm(x, (width,), ln1_w, ln1_b, eps)
    q = F.linear(h, q_w, q_b).reshape((batch, length, heads, head_dim)).transpose(1, 2)
    k = F.linear(h, k_w, k_b).reshape((batch, length, heads, head_dim)).transpose(1, 2)
    v = F.linear(h, v_w, v_b).reshape((batch, length, heads, head_dim)).transpose(1, 2)
    k_all = torch.cat((k_cache, k), 2)
    v_all = torch.cat((v_cache, v), 2)
    a = F.scaled_dot_product_attention(q, k_all, v_all, is_causal=causal)
    merged = a.transpose(1, 2).reshape((batch, length, width))
    attended = torch.add(x, F.linear(merged, o_w, o_b))
    n = torch.layer_norm(attended, (width,), ln2_w, ln2_b, eps)
    f = F.gelu(F.linear(n, fc_w, fc_b), approximate="tanh")
    return torch.add(attended, F.linear(f, proj_w, proj_b)), k_all, v_all


def scaled(rows: int, columns: int, spread: float, device: str) -> torch.Tensor:
    """A weight of the given shape, small enough that 48 blocks stay finite."""
    return torch.mul(torch.randn(rows, columns), spread).to(device)


class Block:  # pylint: disable=too-many-instance-attributes
    """The seventeen weights one block's region takes, in the order it takes them."""

    def __init__(self, width: int, device: str) -> None:
        inner: int = 4 * width
        spread: float = 0.02
        self.ln1_w: torch.Tensor = torch.ones(width).to(device)
        self.ln1_b: torch.Tensor = torch.zeros(width).to(device)
        self.q_w: torch.Tensor = scaled(width, width, spread, device)
        self.q_b: torch.Tensor = torch.zeros(width).to(device)
        self.k_w: torch.Tensor = scaled(width, width, spread, device)
        self.k_b: torch.Tensor = torch.zeros(width).to(device)
        self.v_w: torch.Tensor = scaled(width, width, spread, device)
        self.v_b: torch.Tensor = torch.zeros(width).to(device)
        self.o_w: torch.Tensor = scaled(width, width, spread, device)
        self.o_b: torch.Tensor = torch.zeros(width).to(device)
        self.ln2_w: torch.Tensor = torch.ones(width).to(device)
        self.ln2_b: torch.Tensor = torch.zeros(width).to(device)
        self.fc_w: torch.Tensor = scaled(inner, width, spread, device)
        self.fc_b: torch.Tensor = torch.zeros(inner).to(device)
        self.proj_w: torch.Tensor = scaled(width, inner, spread, device)
        self.proj_b: torch.Tensor = torch.zeros(width).to(device)


class Model:  # pylint: disable=too-many-instance-attributes
    """GPT-2's shape: embeddings, a stack of blocks, and a tied output head."""

    def __init__(self, layers: int, heads: int, head_dim: int, vocab: int, device: str) -> None:
        width: int = heads * head_dim
        self.layers: int = layers
        self.heads: int = heads
        self.head_dim: int = head_dim
        self.width: int = width
        self.vocab: int = vocab
        self.wte: torch.Tensor = scaled(vocab, width, 0.02, device)
        self.wpe: torch.Tensor = scaled(1024, width, 0.01, device)
        self.lnf_w: torch.Tensor = torch.ones(width).to(device)
        self.lnf_b: torch.Tensor = torch.zeros(width).to(device)
        self.blocks: list[Block] = [Block(width, device) for _i in range(layers)]
        #: The key and value of every layer so far, and how many tokens that is.
        self.position: int = 0
        self.ks: list[torch.Tensor] = []
        self.vs: list[torch.Tensor] = []

    def reset(self, batch: int, device: str) -> None:
        """An empty cache per layer, and the position the next token sits at."""
        self.position = 0
        empty = torch.zeros(batch, self.heads, 0, self.head_dim).to(device)
        self.ks = [empty for _i in range(self.layers)]
        self.vs = [empty for _i in range(self.layers)]

    def run(self, ids: torch.Tensor, batch: int, length: int, causal: bool) -> torch.Tensor:
        """Prefill when `length` is the prompt, one decode step when it is 1."""
        start: int = self.position
        h: torch.Tensor = torch.add(self.wte[ids], self.wpe[start : start + length])
        for index in range(self.layers):
            layer = self.blocks[index]
            h, k_all, v_all = step(
                h,
                self.ks[index],
                self.vs[index],
                layer.ln1_w,
                layer.ln1_b,
                layer.q_w,
                layer.q_b,
                layer.k_w,
                layer.k_b,
                layer.v_w,
                layer.v_b,
                layer.o_w,
                layer.o_b,
                layer.ln2_w,
                layer.ln2_b,
                layer.fc_w,
                layer.fc_b,
                layer.proj_w,
                layer.proj_b,
                batch,
                length,
                self.heads,
                self.head_dim,
                self.width,
                EPS,
                causal,
            )
            self.ks[index] = k_all
            self.vs[index] = v_all
        self.position = start + length
        normed: torch.Tensor = torch.layer_norm(h, (self.width,), self.lnf_w, self.lnf_b, EPS)
        return torch.matmul(normed, self.wte.t())

    def forward(self, ids: torch.Tensor, batch: int, length: int) -> torch.Tensor:
        """Logits for every position, the way the reference implementation computes them."""
        h: torch.Tensor = torch.add(self.wte[ids], self.wpe[0:length])
        for layer in self.blocks:
            h = block(
                h,
                layer.ln1_w,
                layer.ln1_b,
                layer.q_w,
                layer.q_b,
                layer.k_w,
                layer.k_b,
                layer.v_w,
                layer.v_b,
                layer.o_w,
                layer.o_b,
                layer.ln2_w,
                layer.ln2_b,
                layer.fc_w,
                layer.fc_b,
                layer.proj_w,
                layer.proj_b,
                batch,
                length,
                self.heads,
                self.head_dim,
                self.width,
                EPS,
            )
        normed: torch.Tensor = torch.layer_norm(h, (self.width,), self.lnf_w, self.lnf_b, EPS)
        return torch.matmul(normed, self.wte.t())


def measure(model: Model, ids: torch.Tensor, batch: int, length: int, rounds: int) -> float:
    """The best of `rounds` forward passes, in milliseconds."""
    best: float = 1e9
    for _round in range(rounds):
        if ids.is_cuda:
            torch.cuda.synchronize()
        start: float = time.perf_counter()
        model.forward(ids, batch, length)
        if ids.is_cuda:
            torch.cuda.synchronize()
        elapsed: float = (time.perf_counter() - start) * 1000.0
        best = elapsed if elapsed < best else best
    return best


def answer(logits: torch.Tensor) -> str:
    """Three aggregates: what the paths in `compare/` agree on to three decimals."""
    wide = logits.float()
    return (
        f"mean={float(wide.mean()):.3f} "
        f"absmean={float(wide.abs().mean()):.3f} "
        f"meansq={float(torch.mul(wide, wide).mean()):.3f}"
    )


def main() -> None:
    on_cuda: bool = torch.cuda.is_available()
    device: str = "cuda" if on_cuda else "cpu"
    # GPT-2 XL where there is a device for it, and a stack short enough to run
    # on a CPU in a few seconds where there is not.
    layers: int = 48 if on_cuda else 6
    heads: int = 25 if on_cuda else 12
    vocab: int = 50257 if on_cuda else 8192
    length: int = 512 if on_cuda else 128
    batch: int = 1

    print("torch", torch.__version__, "| cuda", on_cuda)
    print("# region active:", getattr(block, "__ppy_region__", False))
    print(f"# {layers} layers, {heads} heads, width {heads * 64}, batch {batch}, length {length}")

    torch.manual_seed(0)
    model = Model(layers, heads, 64, vocab, device)
    ids = torch.remainder(torch.arange(batch * length).reshape((batch, length)), vocab).to(device)

    warm = model.forward(ids, batch, length)
    print(answer(warm))
    print(f"# forward: {measure(model, ids, batch, length, 5):.3f} ms")

    # The same weights through the cache: prefill all but the last token, then
    # take that one step incrementally. The two paths compute the same logits,
    # which is the check that the cache the region keeps is the right one.
    model.reset(batch, device)
    model.run(ids[:, 0 : length - 1], batch, length - 1, True)
    stepped = model.run(ids[:, length - 1 : length], batch, 1, False)
    whole = warm[:, length - 1 : length]
    # `allclose` takes `rtol` before `atol`, and a logit near zero fails a
    # relative test that a difference of 1e-6 has no business failing.
    matches: bool = bool(torch.allclose(stepped, whole, 1e-4, 1e-4))
    print("cached decode matches the whole forward:", matches)


if __name__ == "__main__":
    main()

Counterpart programs

The programs the comparison above measured, each written the way its tool expects. The PPy one is first.

gpt2_bench.ppy (PPy)

"""GPT-2 XL forward and training step, with each block compiled to one ATen region."""

import time

import ppy
import torch
import torch.nn.functional as F

EPS: float = 1e-5
LAYERS: int = 48
HEADS: int = 25
HEAD_DIM: int = 64
VOCAB: int = 50257
#: A CPU run of this bench is a smoke test; only a device needs synchronizing.
ON_CUDA: bool = torch.cuda.is_available()


@ppy.opt(3)
def block(
    x: torch.Tensor,
    ln1_w: torch.Tensor,
    ln1_b: torch.Tensor,
    q_w: torch.Tensor,
    q_b: torch.Tensor,
    k_w: torch.Tensor,
    k_b: torch.Tensor,
    v_w: torch.Tensor,
    v_b: torch.Tensor,
    o_w: torch.Tensor,
    o_b: torch.Tensor,
    ln2_w: torch.Tensor,
    ln2_b: torch.Tensor,
    fc_w: torch.Tensor,
    fc_b: torch.Tensor,
    proj_w: torch.Tensor,
    proj_b: torch.Tensor,
    batch: int,
    length: int,
    heads: int,
    head_dim: int,
    width: int,
    eps: float,
) -> torch.Tensor:
    # A region takes every weight as a parameter of its own: that is what a
    # region is, and seventeen tensors beside six dimensions is the shape of
    # one transformer block.
    # pylint: disable=too-many-arguments,too-many-positional-arguments,too-many-locals
    h = torch.layer_norm(x, (width,), ln1_w, ln1_b, eps)
    q = F.linear(h, q_w, q_b).reshape((batch, length, heads, head_dim)).transpose(1, 2)
    k = F.linear(h, k_w, k_b).reshape((batch, length, heads, head_dim)).transpose(1, 2)
    v = F.linear(h, v_w, v_b).reshape((batch, length, heads, head_dim)).transpose(1, 2)
    a = F.scaled_dot_product_attention(q, k, v, is_causal=True)
    merged = a.transpose(1, 2).reshape((batch, length, width))
    attended = torch.add(x, F.linear(merged, o_w, o_b))
    n = torch.layer_norm(attended, (width,), ln2_w, ln2_b, eps)
    f = F.gelu(F.linear(n, fc_w, fc_b), approximate="tanh")
    return torch.add(attended, F.linear(f, proj_w, proj_b))


@ppy.opt(3)
def step(
    x: torch.Tensor,
    k_cache: torch.Tensor,
    v_cache: torch.Tensor,
    ln1_w: torch.Tensor,
    ln1_b: torch.Tensor,
    q_w: torch.Tensor,
    q_b: torch.Tensor,
    k_w: torch.Tensor,
    k_b: torch.Tensor,
    v_w: torch.Tensor,
    v_b: torch.Tensor,
    o_w: torch.Tensor,
    o_b: torch.Tensor,
    ln2_w: torch.Tensor,
    ln2_b: torch.Tensor,
    fc_w: torch.Tensor,
    fc_b: torch.Tensor,
    proj_w: torch.Tensor,
    proj_b: torch.Tensor,
    batch: int,
    length: int,
    heads: int,
    head_dim: int,
    width: int,
    eps: float,
    causal: bool,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """The same block, keeping its own key and value.

    It hands back three tensors, so the cache never leaves the region to be
    reassembled in Python; `causal` is a parameter rather than a constant, so
    one compiled function prefills a prompt and then decodes a token at a
    time out of what it built.
    """
    # A region takes every weight as a parameter of its own: that is what a
    # region is, and seventeen tensors beside six dimensions is the shape of
    # one transformer block.
    # pylint: disable=too-many-arguments,too-many-positional-arguments,too-many-locals
    h = torch.layer_norm(x, (width,), ln1_w, ln1_b, eps)
    q = F.linear(h, q_w, q_b).reshape((batch, length, heads, head_dim)).transpose(1, 2)
    k = F.linear(h, k_w, k_b).reshape((batch, length, heads, head_dim)).transpose(1, 2)
    v = F.linear(h, v_w, v_b).reshape((batch, length, heads, head_dim)).transpose(1, 2)
    k_all = torch.cat((k_cache, k), 2)
    v_all = torch.cat((v_cache, v), 2)
    a = F.scaled_dot_product_attention(q, k_all, v_all, is_causal=causal)
    merged = a.transpose(1, 2).reshape((batch, length, width))
    attended = torch.add(x, F.linear(merged, o_w, o_b))
    n = torch.layer_norm(attended, (width,), ln2_w, ln2_b, eps)
    f = F.gelu(F.linear(n, fc_w, fc_b), approximate="tanh")
    return torch.add(attended, F.linear(f, proj_w, proj_b)), k_all, v_all


def scaled(rows: int, columns: int, spread: float, device: str) -> torch.Tensor:
    """One weight, drawn in the order every counterpart draws it."""
    return torch.mul(torch.randn(rows, columns), spread).to(device).bfloat16()


def flat(width: int, value: float, device: str) -> torch.Tensor:
    return torch.mul(torch.ones(width), value).to(device).bfloat16()


def empty_cache(batch: int, device: str) -> torch.Tensor:
    """A key or value cache with no tokens in it yet."""
    return torch.zeros(batch, HEADS, 0, HEAD_DIM).to(device).bfloat16()


class Block:  # pylint: disable=too-many-instance-attributes
    """The seventeen weights of one block, in the order its region takes them."""

    def __init__(self, width: int, device: str) -> None:
        inner: int = 4 * width
        spread: float = 0.02
        self.ln1_w: torch.Tensor = flat(width, 1.0, device)
        self.ln1_b: torch.Tensor = flat(width, 0.0, device)
        self.q_w: torch.Tensor = scaled(width, width, spread, device)
        self.q_b: torch.Tensor = flat(width, 0.0, device)
        self.k_w: torch.Tensor = scaled(width, width, spread, device)
        self.k_b: torch.Tensor = flat(width, 0.0, device)
        self.v_w: torch.Tensor = scaled(width, width, spread, device)
        self.v_b: torch.Tensor = flat(width, 0.0, device)
        self.o_w: torch.Tensor = scaled(width, width, spread, device)
        self.o_b: torch.Tensor = flat(width, 0.0, device)
        self.ln2_w: torch.Tensor = flat(width, 1.0, device)
        self.ln2_b: torch.Tensor = flat(width, 0.0, device)
        self.fc_w: torch.Tensor = scaled(inner, width, spread, device)
        self.fc_b: torch.Tensor = flat(inner, 0.0, device)
        self.proj_w: torch.Tensor = scaled(width, inner, spread, device)
        self.proj_b: torch.Tensor = flat(width, 0.0, device)

    def trainable(self) -> None:
        """Autograd on the two projections a training step differentiates."""
        self.fc_w = self.fc_w.requires_grad_(True)
        self.proj_w = self.proj_w.requires_grad_(True)

    def zero_grads(self) -> None:
        """What an optimizer's `zero_grad(set_to_none=True)` does."""
        self.fc_w.grad = None
        self.proj_w.grad = None


class Model:  # pylint: disable=too-many-instance-attributes
    """GPT-2 XL's shape: embeddings, 48 blocks, and a tied output head."""

    def __init__(self, device: str) -> None:
        width: int = HEADS * HEAD_DIM
        self.width: int = width
        self.wte: torch.Tensor = scaled(VOCAB, width, 0.02, device)
        self.wpe: torch.Tensor = scaled(1024, width, 0.01, device)
        self.lnf_w: torch.Tensor = flat(width, 1.0, device)
        self.lnf_b: torch.Tensor = flat(width, 0.0, device)
        self.blocks: list[Block] = [Block(width, device) for _i in range(LAYERS)]
        #: The key and value of every layer so far, and how many tokens that is.
        self.position: int = 0
        self.ks: list[torch.Tensor] = []
        self.vs: list[torch.Tensor] = []

    def trainable(self) -> None:
        for layer in self.blocks:
            layer.trainable()

    def zero_grads(self) -> None:
        for layer in self.blocks:
            layer.zero_grads()

    def forward(self, ids: torch.Tensor, batch: int, length: int) -> torch.Tensor:
        h: torch.Tensor = torch.add(self.wte[ids], self.wpe[0:length])
        for layer in self.blocks:
            h = block(
                h,
                layer.ln1_w,
                layer.ln1_b,
                layer.q_w,
                layer.q_b,
                layer.k_w,
                layer.k_b,
                layer.v_w,
                layer.v_b,
                layer.o_w,
                layer.o_b,
                layer.ln2_w,
                layer.ln2_b,
                layer.fc_w,
                layer.fc_b,
                layer.proj_w,
                layer.proj_b,
                batch,
                length,
                HEADS,
                HEAD_DIM,
                self.width,
                EPS,
            )
        normed: torch.Tensor = torch.layer_norm(h, (self.width,), self.lnf_w, self.lnf_b, EPS)
        return torch.matmul(normed, self.wte.t())

    def reset(self, batch: int, device: str) -> None:
        """An empty cache per layer, and the position the next token sits at."""
        self.position = 0
        self.ks = [empty_cache(batch, device) for _i in range(LAYERS)]
        self.vs = [empty_cache(batch, device) for _i in range(LAYERS)]

    def run(self, ids: torch.Tensor, batch: int, length: int, causal: bool) -> torch.Tensor:
        """Prefill when `length` is the prompt, one decode step when it is 1."""
        start: int = self.position
        h: torch.Tensor = torch.add(self.wte[ids], self.wpe[start : start + length])
        for index in range(LAYERS):
            layer = self.blocks[index]
            h, k_all, v_all = step(
                h,
                self.ks[index],
                self.vs[index],
                layer.ln1_w,
                layer.ln1_b,
                layer.q_w,
                layer.q_b,
                layer.k_w,
                layer.k_b,
                layer.v_w,
                layer.v_b,
                layer.o_w,
                layer.o_b,
                layer.ln2_w,
                layer.ln2_b,
                layer.fc_w,
                layer.fc_b,
                layer.proj_w,
                layer.proj_b,
                batch,
                length,
                HEADS,
                HEAD_DIM,
                self.width,
                EPS,
                causal,
            )
            self.ks[index] = k_all
            self.vs[index] = v_all
        self.position = start + length
        normed: torch.Tensor = torch.layer_norm(h, (self.width,), self.lnf_w, self.lnf_b, EPS)
        return torch.matmul(normed, self.wte.t())


def ids_for(batch: int, length: int, device: str) -> torch.Tensor:
    return torch.remainder(torch.arange(batch * length).reshape((batch, length)), VOCAB).to(device)


def forward_ms(model: Model, ids: torch.Tensor, batch: int, length: int, rounds: int) -> float:
    """The best of `rounds` forward passes, in milliseconds."""
    best: float = 1e9
    for _round in range(rounds):
        if ON_CUDA:
            torch.cuda.synchronize()
        start: float = time.perf_counter()
        model.forward(ids, batch, length)
        if ON_CUDA:
            torch.cuda.synchronize()
        elapsed: float = (time.perf_counter() - start) * 1000.0
        best = elapsed if elapsed < best else best
    return best


def step_ms(model: Model, ids: torch.Tensor, batch: int, length: int, rounds: int) -> float:
    """The best forward-and-backward, which is what a training step costs."""
    best: float = 1e9
    for _round in range(rounds):
        if ON_CUDA:
            torch.cuda.synchronize()
        start: float = time.perf_counter()
        # A step that accumulated instead would leave `reduce-overhead` with
        # a gradient buffer its CUDA graph had already overwritten.
        model.zero_grads()
        model.forward(ids, batch, length).sum().backward()
        if ON_CUDA:
            torch.cuda.synchronize()
        elapsed: float = (time.perf_counter() - start) * 1000.0
        best = elapsed if elapsed < best else best
    return best


def decode_ms(model: Model, ids: torch.Tensor, batch: int, prompt: int, steps: int) -> float:
    """One prefill, then `steps` tokens out of the cache; only the decode is timed.

    The token fed at each step is the next one of `ids`, not the model's own
    argmax: on weights drawn from a generator the top two logits are usually
    a tie, so sampling would send the four paths down different sequences and
    measure different work.
    """
    model.reset(batch, "cuda" if ON_CUDA else "cpu")
    model.run(ids[:, 0:prompt], batch, prompt, True)
    if ON_CUDA:
        torch.cuda.synchronize()
    start: float = time.perf_counter()
    for offset in range(steps):
        at: int = prompt + offset
        model.run(ids[:, at : at + 1], batch, 1, False)
    if ON_CUDA:
        torch.cuda.synchronize()
    return (time.perf_counter() - start) * 1000.0


def best_decode(model: Model, ids: torch.Tensor, batch: int, prompt: int, steps: int) -> float:
    best: float = 1e9
    for _round in range(3):
        elapsed: float = decode_ms(model, ids, batch, prompt, steps)
        best = elapsed if elapsed < best else best
    return best


def answer(logits: torch.Tensor) -> str:
    """What every path must agree on.

    Not the greedy token, and not an extreme: on weights drawn from a
    generator the top two logits of a position are often a tie, and both
    fusing the reductions differently and ordering an argmax differently
    flip it. Three aggregates over the whole tensor do hold, to three
    decimals, across every path here.
    """
    wide = logits.float()
    return (
        f"mean={float(wide.mean()):.3f} "
        f"absmean={float(wide.abs().mean()):.3f} "
        f"meansq={float(torch.mul(wide, wide).mean()):.3f}"
    )


def main() -> None:
    device: str = "cuda" if ON_CUDA else "cpu"
    torch.manual_seed(0)
    model = Model(device)

    for batch, length in [(1, 32), (1, 128), (1, 512), (8, 512)]:
        ids = ids_for(batch, length, device)
        for _warm in range(8):
            model.forward(ids, batch, length)
        print(answer(model.forward(ids, batch, length)))
        print(f"# forward b{batch} s{length}: {forward_ms(model, ids, batch, length, 10):.3f} ms")

    prompt: int = 128
    steps: int = 128
    ids = ids_for(1, prompt + steps, device)
    model.reset(1, device)
    model.run(ids[:, 0:prompt], 1, prompt, True)
    last = model.run(ids[:, prompt : prompt + 1], 1, 1, False)
    print(answer(last))
    print(f"# decode {steps} tokens: {best_decode(model, ids, 1, prompt, steps):.3f} ms")

    model.trainable()
    ids = ids_for(4, 512, device)
    for _warm in range(8):
        model.zero_grads()
        model.forward(ids, 4, 512).sum().backward()
    print(f"# training step b4 s512: {step_ms(model, ids, 4, 512, 10):.3f} ms")


if __name__ == "__main__":
    main()

gpt2_torch.py (Python)

"""The same GPT-2 XL in plain PyTorch: eager, or through `torch.compile`.

The weights are drawn in the same order as `gpt2_bench.ppy` draws them, so
all four columns are the same model on the same numbers. The mode is an
argument because eager and the two compiled modes differ by one line.

    python gpt2_torch.py eager
    python gpt2_torch.py compile
    python gpt2_torch.py reduce-overhead
"""

from __future__ import annotations

import sys
import time

import torch
import torch.nn.functional as F

EPS = 1e-5
LAYERS = 48
HEADS = 25
HEAD_DIM = 64
VOCAB = 50257
ON_CUDA = torch.cuda.is_available()


def block(
    x,
    ln1_w,
    ln1_b,
    q_w,
    q_b,
    k_w,
    k_b,
    v_w,
    v_b,
    o_w,
    o_b,
    ln2_w,
    ln2_b,
    fc_w,
    fc_b,
    proj_w,
    proj_b,
    batch,
    length,
    heads,
    head_dim,
    width,
    eps,
):
    h = torch.layer_norm(x, (width,), ln1_w, ln1_b, eps)
    q = F.linear(h, q_w, q_b).reshape((batch, length, heads, head_dim)).transpose(1, 2)
    k = F.linear(h, k_w, k_b).reshape((batch, length, heads, head_dim)).transpose(1, 2)
    v = F.linear(h, v_w, v_b).reshape((batch, length, heads, head_dim)).transpose(1, 2)
    a = F.scaled_dot_product_attention(q, k, v, is_causal=True)
    merged = a.transpose(1, 2).reshape((batch, length, width))
    attended = torch.add(x, F.linear(merged, o_w, o_b))
    n = torch.layer_norm(attended, (width,), ln2_w, ln2_b, eps)
    f = F.gelu(F.linear(n, fc_w, fc_b), approximate="tanh")
    return torch.add(attended, F.linear(f, proj_w, proj_b))


def step(
    x,
    k_cache,
    v_cache,
    ln1_w,
    ln1_b,
    q_w,
    q_b,
    k_w,
    k_b,
    v_w,
    v_b,
    o_w,
    o_b,
    ln2_w,
    ln2_b,
    fc_w,
    fc_b,
    proj_w,
    proj_b,
    batch,
    length,
    heads,
    head_dim,
    width,
    eps,
    causal,
):
    h = torch.layer_norm(x, (width,), ln1_w, ln1_b, eps)
    q = F.linear(h, q_w, q_b).reshape((batch, length, heads, head_dim)).transpose(1, 2)
    k = F.linear(h, k_w, k_b).reshape((batch, length, heads, head_dim)).transpose(1, 2)
    v = F.linear(h, v_w, v_b).reshape((batch, length, heads, head_dim)).transpose(1, 2)
    k_all = torch.cat((k_cache, k), 2)
    v_all = torch.cat((v_cache, v), 2)
    a = F.scaled_dot_product_attention(q, k_all, v_all, is_causal=causal)
    merged = a.transpose(1, 2).reshape((batch, length, width))
    attended = torch.add(x, F.linear(merged, o_w, o_b))
    n = torch.layer_norm(attended, (width,), ln2_w, ln2_b, eps)
    f = F.gelu(F.linear(n, fc_w, fc_b), approximate="tanh")
    return torch.add(attended, F.linear(f, proj_w, proj_b)), k_all, v_all


def scaled(rows, columns, spread, device):
    return torch.mul(torch.randn(rows, columns), spread).to(device).bfloat16()


def flat(width, value, device):
    return torch.mul(torch.ones(width), value).to(device).bfloat16()


def empty_cache(batch, device):
    return torch.zeros(batch, HEADS, 0, HEAD_DIM).to(device).bfloat16()


class Block:
    def __init__(self, width, device):
        inner = 4 * width
        spread = 0.02
        self.ln1_w = flat(width, 1.0, device)
        self.ln1_b = flat(width, 0.0, device)
        self.q_w = scaled(width, width, spread, device)
        self.q_b = flat(width, 0.0, device)
        self.k_w = scaled(width, width, spread, device)
        self.k_b = flat(width, 0.0, device)
        self.v_w = scaled(width, width, spread, device)
        self.v_b = flat(width, 0.0, device)
        self.o_w = scaled(width, width, spread, device)
        self.o_b = flat(width, 0.0, device)
        self.ln2_w = flat(width, 1.0, device)
        self.ln2_b = flat(width, 0.0, device)
        self.fc_w = scaled(inner, width, spread, device)
        self.fc_b = flat(inner, 0.0, device)
        self.proj_w = scaled(width, inner, spread, device)
        self.proj_b = flat(width, 0.0, device)

    def trainable(self):
        self.fc_w = self.fc_w.requires_grad_(True)
        self.proj_w = self.proj_w.requires_grad_(True)

    def zero_grads(self):
        self.fc_w.grad = None
        self.proj_w.grad = None


class Model:
    def __init__(self, device):
        width = HEADS * HEAD_DIM
        self.width = width
        self.wte = scaled(VOCAB, width, 0.02, device)
        self.wpe = scaled(1024, width, 0.01, device)
        self.lnf_w = flat(width, 1.0, device)
        self.lnf_b = flat(width, 0.0, device)
        self.blocks = [Block(width, device) for _i in range(LAYERS)]
        self.position = 0
        self.ks = []
        self.vs = []

    def trainable(self):
        for layer in self.blocks:
            layer.trainable()

    def zero_grads(self):
        for layer in self.blocks:
            layer.zero_grads()

    def forward(self, ids, batch, length):
        h = torch.add(self.wte[ids], self.wpe[0:length])
        for layer in self.blocks:
            h = block(
                h,
                layer.ln1_w,
                layer.ln1_b,
                layer.q_w,
                layer.q_b,
                layer.k_w,
                layer.k_b,
                layer.v_w,
                layer.v_b,
                layer.o_w,
                layer.o_b,
                layer.ln2_w,
                layer.ln2_b,
                layer.fc_w,
                layer.fc_b,
                layer.proj_w,
                layer.proj_b,
                batch,
                length,
                HEADS,
                HEAD_DIM,
                self.width,
                EPS,
            )
        normed = torch.layer_norm(h, (self.width,), self.lnf_w, self.lnf_b, EPS)
        return torch.matmul(normed, self.wte.t())

    def reset(self, batch, device):
        self.position = 0
        self.ks = [empty_cache(batch, device) for _i in range(LAYERS)]
        self.vs = [empty_cache(batch, device) for _i in range(LAYERS)]

    def run(self, ids, batch, length, causal):
        start = self.position
        h = torch.add(self.wte[ids], self.wpe[start : start + length])
        for index in range(LAYERS):
            layer = self.blocks[index]
            h, k_all, v_all = step(
                h,
                self.ks[index],
                self.vs[index],
                layer.ln1_w,
                layer.ln1_b,
                layer.q_w,
                layer.q_b,
                layer.k_w,
                layer.k_b,
                layer.v_w,
                layer.v_b,
                layer.o_w,
                layer.o_b,
                layer.ln2_w,
                layer.ln2_b,
                layer.fc_w,
                layer.fc_b,
                layer.proj_w,
                layer.proj_b,
                batch,
                length,
                HEADS,
                HEAD_DIM,
                self.width,
                EPS,
                causal,
            )
            self.ks[index] = k_all
            self.vs[index] = v_all
        self.position = start + length
        normed = torch.layer_norm(h, (self.width,), self.lnf_w, self.lnf_b, EPS)
        return torch.matmul(normed, self.wte.t())


def ids_for(batch, length, device):
    return torch.remainder(torch.arange(batch * length).reshape((batch, length)), VOCAB).to(device)


def best_ms(work, rounds):
    best = 1e9
    for _round in range(rounds):
        if ON_CUDA:
            torch.cuda.synchronize()
        start = time.perf_counter()
        work()
        if ON_CUDA:
            torch.cuda.synchronize()
        best = min(best, (time.perf_counter() - start) * 1000.0)
    return best


def answer(logits):
    wide = logits.float()
    return (
        f"mean={float(wide.mean()):.3f} "
        f"absmean={float(wide.abs().mean()):.3f} "
        f"meansq={float(torch.mul(wide, wide).mean()):.3f}"
    )


def main():
    mode = sys.argv[1] if len(sys.argv) > 1 else "eager"
    device = "cuda" if ON_CUDA else "cpu"
    torch.manual_seed(0)
    model = Model(device)

    forward = model.forward
    runner = model.run
    # A KV cache is a tensor the caller keeps across invocations, which is
    # exactly what CUDA graphs reuse the memory of. `reduce-overhead` refuses
    # the decode loop outright -- "accessing tensor output of CUDAGraphs that
    # has been overwritten by a subsequent run" -- until each invocation is
    # announced, which is what its own error message asks for.
    mark = None
    if mode == "compile":
        forward = torch.compile(model.forward)
        runner = torch.compile(model.run)
    elif mode == "reduce-overhead":
        forward = torch.compile(model.forward, mode="reduce-overhead")
        runner = torch.compile(model.run, mode="reduce-overhead")
        mark = torch.compiler.cudagraph_mark_step_begin
        # Announcing the step is not enough on its own: the cache is a
        # user-visible output that has to stay live across generations, and
        # this is the switch the error message names for that case.
        torch._inductor.config.triton.cudagraph_trees_generation_cloning = "user_visible"
    elif mode != "eager":
        raise SystemExit(f"unknown mode {mode!r}")

    def advance(*arguments):
        if mark is not None:
            mark()
        return runner(*arguments)

    for batch, length in [(1, 32), (1, 128), (1, 512), (8, 512)]:
        ids = ids_for(batch, length, device)
        for _warm in range(8):
            forward(ids, batch, length)
        print(answer(forward(ids, batch, length)))
        elapsed = best_ms(lambda: forward(ids, batch, length), 10)  # noqa: B023
        print(f"# forward b{batch} s{length}: {elapsed:.3f} ms")

    prompt, steps = 128, 128
    ids = ids_for(1, prompt + steps, device)

    def decode_ms():
        # The token fed at each step is the next one of `ids`, not the model's
        # own argmax: sampling would send the four paths down different
        # sequences and measure different work.
        model.reset(1, device)
        advance(ids[:, 0:prompt], 1, prompt, True)
        if ON_CUDA:
            torch.cuda.synchronize()
        started = time.perf_counter()
        for offset in range(steps):
            at = prompt + offset
            advance(ids[:, at : at + 1], 1, 1, False)
        if ON_CUDA:
            torch.cuda.synchronize()
        return (time.perf_counter() - started) * 1000.0

    for _warm in range(2):
        decode_ms()
    model.reset(1, device)
    advance(ids[:, 0:prompt], 1, prompt, True)
    print(answer(advance(ids[:, prompt : prompt + 1], 1, 1, False)))
    print(f"# decode {steps} tokens: {min(decode_ms() for _ in range(3)):.3f} ms")

    model.trainable()
    ids = ids_for(4, 512, device)
    for _warm in range(8):
        model.zero_grads()
        forward(ids, 4, 512).sum().backward()

    def training_step():
        # Accumulating instead would leave `reduce-overhead` with a gradient
        # buffer its CUDA graph had already overwritten.
        model.zero_grads()
        forward(ids, 4, 512).sum().backward()

    print(f"# training step b4 s512: {best_ms(training_step, 10):.3f} ms")


if __name__ == "__main__":
    main()

Source: examples/46_gpt2.