GPT, in depth: the decoder-only transformer from forward pass and training loop to KV cache, sampling and hardware limits

What decoder-only means

The original transformer of 2017 had two stacks: an encoder that reads the whole input with bidirectional attention and a decoder that writes the output one token at a time, attending to its own past and to the encoder. GPT keeps only the decoder stack and removes the cross-attention to an encoder, which is why it is called decoder-only. Everything, prompt and answer alike, is one sequence, and one rule governs it: position i may attend to positions 0 through i and never beyond.

That rule, the causal mask, is what makes the training objective work. Because no position can see its own future, the output at every position is a legitimate prediction of the next token, and a single forward pass over a 1,024-token sequence produces 1,024 training examples at once. The same rule is what makes generation cheap to extend: appending a token never changes what earlier positions computed, so their intermediate keys and values can be cached and reused.

GPT-2 moved layer normalisation to the input of each sub-block and added a final normalisation before the output, the pre-norm arrangement that made deep stacks easier to train. Later open model families kept the decoder-only skeleton while swapping parts: rotary position embeddings instead of a learned position table, RMSNorm instead of LayerNorm, gated SwiGLU MLPs, and grouped-query attention to shrink the cache. Architectural details of later proprietary GPT models are not published, so nothing here depends on them.

The forward pass

Follow the shapes for a batch of B sequences of length T, model width d, h heads and vocabulary V. Token ids of shape B x T index an embedding table V x d, giving B x T x d; a learned position table adds a vector per position. Each block then computes x = x + Attn(LN(x)) followed by x = x + MLP(LN(x)). Inside attention, one matrix multiply produces queries, keys and values, each split into h heads of size d/h. Scores are Q K^T / sqrt(d/h), shape B x h x T x T; entries above the diagonal are set to minus infinity so that softmax gives them zero weight; the weights multiply V, heads are concatenated and projected back to d. The MLP expands to 4d, applies GELU and projects back. After L blocks, a final LayerNorm and a multiply by the transposed embedding table give logits of shape B x T x V.

The residual stream is the key idea for understanding behaviour: every block reads the running sum x and adds a correction to it. Nothing overwrites the stream, which is why very deep stacks train stably and why the same d-dimensional vector carries token identity, position and everything the blocks have inferred so far.

A complete model in code

Here is a complete model in PyTorch, close in structure to GPT-2. It uses the fused scaled-dot-product attention function with is_causal=True, which applies the mask without materialising it.

import torch, torch.nn as nn, torch.nn.functional as F

class Block(nn.Module):
    def __init__(self, d, h):
        super().__init__()
        self.h = h
        self.ln1, self.ln2 = nn.LayerNorm(d), nn.LayerNorm(d)
        self.qkv, self.proj = nn.Linear(d, 3 * d), nn.Linear(d, d)
        self.mlp = nn.Sequential(nn.Linear(d, 4 * d), nn.GELU(), nn.Linear(4 * d, d))

    def forward(self, x):
        B, T, d = x.shape
        q, k, v = self.qkv(self.ln1(x)).split(d, dim=2)
        q, k, v = (t.view(B, T, self.h, d // self.h).transpose(1, 2) for t in (q, k, v))
        y = F.scaled_dot_product_attention(q, k, v, is_causal=True)
        x = x + self.proj(y.transpose(1, 2).reshape(B, T, d))
        return x + self.mlp(self.ln2(x))

class GPT(nn.Module):
    def __init__(self, V=50257, T_max=1024, d=768, h=12, L=12):
        super().__init__()
        self.tok, self.pos = nn.Embedding(V, d), nn.Embedding(T_max, d)
        self.blocks = nn.ModuleList(Block(d, h) for _ in range(L))
        self.ln_f = nn.LayerNorm(d)
        self.head = nn.Linear(d, V, bias=False)
        self.head.weight = self.tok.weight          # weight tying

    def forward(self, idx):
        T = idx.shape[1]
        x = self.tok(idx) + self.pos(torch.arange(T, device=idx.device))
        for b in self.blocks:
            x = b(x)
        return self.head(self.ln_f(x))              # B x T x V logits

Weight tying means the matrix that turns a token into a vector is the same matrix, transposed, that turns the final vector into scores. It saves V x d parameters, which for a 50,257-token vocabulary at d = 768 is almost a third of the model.

The training loop

Training is next-token prediction with cross-entropy. Take a window of T + 1 tokens, feed the first T as input and use the last T as targets, so the target at position t is the input at t + 1. That one-position shift is the whole trick, and getting it wrong is the most common bug in hand-written training code.

model = GPT().cuda()
opt = torch.optim.AdamW(model.parameters(), lr=6e-4, betas=(0.9, 0.95), weight_decay=0.1)

for step, chunk in enumerate(loader):              # chunk: B x (T+1) token ids
    x, y = chunk[:, :-1].cuda(), chunk[:, 1:].cuda()
    with torch.autocast("cuda", dtype=torch.bfloat16):
        logits = model(x)
        loss = F.cross_entropy(logits.reshape(-1, logits.size(-1)), y.reshape(-1))
    loss.backward()
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
    opt.step(); opt.zero_grad(set_to_none=True)

Notice what is absent: no loop over positions. Teacher forcing feeds the true previous tokens at every position and the mask guarantees no position cheats, so all T predictions are computed and graded in one pass. The loss is the average negative log-probability assigned to the true next token; its exponential is perplexity. In practice the loop also has learning-rate warm-up and cosine decay, gradient accumulation to reach large effective batch sizes, and documents packed end to end with a separator token so no compute is wasted on padding. The derivation of the loss is in Cross-Entropy Loss for Next-Token Prediction.

Worked example: counting GPT-2 small

Make the abstractions concrete with GPT-2 small: L = 12, d = 768, h = 12, V = 50,257, context 1,024. Each block holds about 12d² weights: 3d² for the query, key and value projection, d² for the output projection and 8d² for the two MLP matrices.

Two rules of thumb follow. Training costs about 6 floating-point operations per parameter per token, two for the forward pass and four for the backward, so training this model on 10 billion tokens costs roughly 6 x 124M x 10B, about 7.4 x 10¹⁸ operations. Generating a token costs about 2 operations per parameter plus the attention term, which grows with context length.

The cache is the other number to know. For each token, every layer stores one key and one value vector of width d. In 16-bit precision that is 2 x 12 x 768 x 2 bytes, or 36,864 bytes per token; a full 1,024-token context is about 37.7 MB per sequence. For large models with long contexts this cache, not the weights, is what limits how many requests fit on a device, which is exactly why grouped-query attention shares keys and values across heads. The full arithmetic is in KV cache math.

Generation: prefill, decode and sampling

Generation has two phases. Prefill runs the whole prompt through the model in one pass, exactly like training, and stores every layer's keys and values. Decode then feeds one new token at a time: its query attends to the cached keys and values plus its own, the model emits logits for one position, a token is sampled and appended, and the loop repeats until an end token or a length limit.

@torch.no_grad()
def generate(model, idx, max_new, temperature=0.8, top_p=0.95):
    for _ in range(max_new):
        logits = model(idx[:, -1024:])[:, -1, :] / temperature   # no cache: simple but O(T^2) work
        probs = F.softmax(logits, dim=-1)
        sorted_p, order = probs.sort(descending=True)
        keep = sorted_p.cumsum(-1) - sorted_p < top_p              # nucleus: smallest set reaching top_p
        sorted_p = sorted_p * keep
        nxt = order.gather(-1, torch.multinomial(sorted_p / sorted_p.sum(-1, keepdim=True), 1))
        idx = torch.cat([idx, nxt], dim=1)
    return idx

This version recomputes the whole prefix each step for clarity; a production loop passes the cache forward and feeds only the newest token. Temperature, top-k and nucleus sampling change which token is chosen, never the model; their effects are compared in Sampling Strategies.

What the hardware is doing

Almost all of GPT's work is matrix multiplication, which is why it runs on GPUs and other accelerators built around dense matrix units fed from high-bandwidth memory. What limits speed depends on the phase. In training and prefill, each weight matrix is loaded once and multiplied against thousands of token vectors, so arithmetic dominates and the job is compute-bound; mixed precision such as bfloat16 substantially raises matrix throughput on hardware with native support and halves memory traffic.

Decode is different. At batch size one, producing a token touches every weight once and does only two operations per weight, so the device spends its time moving bytes, not multiplying them. A model with W bytes of weights on a device with memory bandwidth M bytes per second cannot exceed about M / W tokens per second for a single sequence, whatever its peak compute. That single fact explains most serving engineering: batching many sequences so each weight load serves many tokens, quantising weights to fewer bytes, paging the KV cache so more sequences fit, and speculative decoding to verify several tokens per weight pass. Attention's own cost grows with the square of context in prefill and linearly per token in decode, which is why long contexts need fused attention kernels that never write the T x T score matrix to memory.

Failure modes

• Missing or misaligned shift. Using the same tensor for inputs and targets trains a copy machine; loss collapses to near zero quickly and generation is gibberish.
• Broken mask. A transposed or off-by-one mask leaks the next token; training loss looks excellent and the model fails the moment it must generate. Test by checking that changing token t+1 never changes logits at position t.
• Padding on the wrong side. With learned positions and batched prompts, right-padding generation inputs places the next token after pad tokens. Left-pad for generation and pass an attention mask.
• Running past the context. A learned position table has exactly T_max rows; indexing beyond it crashes or, with clamping, silently degrades output. Truncate the prefix deliberately.
• Precision overflow. float16 attention scores and logits can overflow; prefer bfloat16 or compute softmax in float32.
• Tokenizer mismatch. Serving with a different tokenizer from training produces fluent nonsense. Version the tokenizer with the weights; Byte-Pair Encoding explains why.

Trade-offs

Decoder-only against encoder-decoder: one stack and one objective make GPT simple to scale and let every task be phrased as text continuation, but the prompt gets only causal attention, so an encoder-decoder can be more parameter-efficient for pure input-to-output tasks like translation. Learned positions against rotary embeddings: a learned table is simple but fixes the maximum length; rotary embeddings encode relative position and extend more gracefully. Full multi-head attention against grouped-query attention: separate keys per head are most expressive, sharing them cuts cache size several-fold for a small quality cost. And autoregression itself: one token per step is easy to train and sample, but latency grows with output length, which is the cost speculative and parallel decoding try to recover.

What to do next

1. Implement the model above and confirm it has about 124M parameters with sum(p.numel() for p in model.parameters()).
2. Write the mask test: perturb token t+1 and assert logits at positions up to t are unchanged.
3. Train on a small corpus, such as a few megabytes of text with a character or BPE tokenizer, until samples become coherent.
4. Add a KV cache to generation and measure tokens per second before and after.
5. Compute weight bytes and cache bytes for a model you serve, and estimate its batch-1 decode ceiling from memory bandwidth.
6. Swap learned positions for rotary embeddings and LayerNorm for RMSNorm, and compare loss curves.