Training View vs. Inference View

The model you train and the model you serve are, computationally, almost unrecognizably different programs. The weights are the same, the matrix multiplications are the same, the attention mechanism is the same. But the shape of the computation changes completely between training and inference, and so does what the hardware is actually doing, how long you wait, and what the bottleneck turns out to be.

A lot of introductions gloss over this. You see the training loop, you see the model.generate() call, and it's easy to come away thinking the difference is just "training has a loss function." The gap is deeper than that, and it shapes everything we'll cover in Arcs 5 and 6. Why serving LLMs is expensive, why we need a KV cache, why speculative decoding exists, why the industry is obsessed with "tokens per second," it all starts here.

Prerequisites

This post assumes you've read:

• Post 4.1: Self-Attention from First Principles for how a forward pass works (embed, attention + FFN blocks reading and writing a residual stream, unembed).
• Post 4.3: Causal Masking for why each position only sees previous tokens, and how teacher forcing falls out of that.

The training view

Let me start with training, because I think it's the easier picture to hold in your head.

During training, you have a dataset of text. You take a chunk of tokens, say a 2,048-token sequence, and feed the entire thing into the model in one shot. Every token goes in simultaneously. One forward pass through the residual stream, one pass through every attention block and every FFN block.

The causal mask (from Post 4.3) ensures that position $i$ can only attend to positions $0$ through $i$. So even though all 2,048 tokens are physically present in the batch, each position's computation only depends on the tokens before it. This is what makes the training signal valid: position 12 computes a representation using only tokens 0 through 12, and we ask it to predict token 13. Position 2,046 uses almost the entire sequence to predict token 2,047 (the final token). Position 2,047 has no next token to predict, so it drops out of the loss.

We get a loss value at every position. 2,047 predictions from a single forward pass (the last position has no next token to predict). That's an extraordinary amount of learning signal per GPU operation.

This is teacher forcing. During training, the model always conditions on the ground-truth prefix. It never generates a token and then conditions on its own generation. It just sees the real sequence, masked so that each position can only peek backward. We covered this in the causal masking post. Teacher forcing isn't something you opt into. It's the natural consequence of feeding a full sequence through a causally masked transformer.

The model never encounters its own mistakes during training. If it assigns low probability to the correct next token at position 50, that's reflected in the loss, but position 51 still gets the real token 50 as input, not whatever the model would have generated. This gap between training conditions and inference conditions is sometimes called "exposure bias," but in practice it works fine, and the efficiency gain is enormous.

The key computational property: training is compute-bound. The GPU is busy doing enormous matrix multiplications across the full sequence length. Every attention block computes an $n \times n$ score matrix. Every FFN block multiplies a batch of $n$ vectors through its weight matrices. The arithmetic logic units are saturated. The GPU is doing what GPUs are designed to do.

The inference view

Now forget all of that. At inference time, the model generates text one token at a time.

You start with a prompt. You feed it through the model. The model produces a probability distribution over the vocabulary. You sample a token from that distribution. Then you append that token to the input and run the model again, and again, and again, until the model emits an end-of-sequence token or you hit some stop condition.

Each step produces exactly one new token. And each step requires a full read of the model's weights, because every attention block and every FFN block needs to process the new token's interactions with all previous tokens.

This is the autoregressive loop, and it's inherently sequential. You can't generate token 5 until you've generated token 4, because token 5's probability distribution depends on what token 4 actually turned out to be. There's a feedback loop: the model's output becomes its own input on the next step.

This is also why latency matters so much for inference. If each decode step takes 30 ms and you're generating a 500-token response, that's 15 seconds of wall-clock time. The user is waiting. During training, nobody cares about latency per sample because you're processing millions of tokens in bulk. Inference has a user on the other end of the connection, watching tokens stream in one by one.

The key computational property: inference (specifically, the decode phase) is memory-bound. For each new token, the GPU reads all the model's weight matrices from memory but only does a tiny amount of math with them. You're multiplying a single vector through billions of parameters. The arithmetic units are mostly idle, waiting for data to arrive from memory.

Arithmetic intensity: the number that explains everything

There's a clean way to quantify this difference. It's called arithmetic intensity. It's the ratio of compute operations (FLOPs) to bytes of data moved from memory:

During training, you're doing a matrix multiplication like $X \cdot W$ where $X$ is a batch of 2,048 vectors. The weight matrix $W$ gets loaded from memory once and reused across all 2,048 vectors. The ratio of compute to memory access is high. For large batch sizes and long sequences, training can hit arithmetic intensities of hundreds to thousands of FLOPs per byte.

During decode, the same weight matrix $W$ gets loaded from memory, but you only multiply it by a single vector. Same memory traffic, a tiny fraction of the compute. The arithmetic intensity drops to something like 1 to 2 FLOPs per byte.

Here's the easiest way to feel this. Imagine you have a massive reference book (the model's weights). Training is like looking up answers for 2,048 questions in one sitting: open the book once, work through all of them. Decode is like looking up one answer, closing the book, then opening it again for the next question. You spend most of your time opening and closing the book, not reading it.

Modern GPUs are designed to do lots of math fast. An A100 can do about 312 teraFLOPs per second in float16, but it can only read from its high-bandwidth memory at about 2 TB/s. That ratio, roughly 156 FLOPs per byte, is the GPU's "ridge point." If your workload sits above that ridge, you're compute-bound and the GPU is happy. Below that, you're memory-bound and the arithmetic units sit idle waiting for data.

There's a classic visualization for this called the roofline model, from Williams, Waterman, and Patterson (2009). This is the single picture I want you to walk away with from this post, and it's the one I'll keep pointing back to for the rest of the arc.

It's a 2D plot. Arithmetic intensity (FLOPs per byte, log scale) goes on the x-axis, achievable performance (FLOPS) on the y-axis. The hardware gives you two ceilings: a sloped one set by memory bandwidth, and a flat one set by peak compute. They cross at the ridge point. Anything to the left of the ridge is memory-bound, anything to the right is compute-bound.

Watch where the markers land. Decode (batch 1) sits at AI = 1, far down on the orange slope. The achievable performance there is roughly 2 TFLOPS, less than 1% of the GPU's peak. Training at AI = 2,048 saturates the green compute ceiling. Same hardware, same model, two different regimes.

The whole rest of Arc 5 and Arc 6 is the field figuring out how to drag the decode marker rightward, toward the ridge. Batching slides it up the slope by getting more tokens out of one weight read. Quantization lifts the slope itself, because fewer bytes per parameter means more FLOPs per byte moved. Speculative decoding spends the idle compute on guesses. Every one of those techniques is a move on this plot, and the numbers change with the hardware (an H100's ridge sits near 296 FLOPs/byte rather than the A100's 156) but the shape of the picture doesn't.

The minimal math

Let me put some numbers on this. For a transformer with:

• $P$ parameters (total model weights),
• $n$ tokens in a sequence,
• $d$ model dimension.

Training forward pass (one sequence, ignoring attention's quadratic term).

The dominant cost is the matrix multiplications in the FFN and attention projection layers. For a full sequence of $n$ tokens, the total FLOPs scale as:

The factor of 2 comes from the multiply-accumulate operations in matrix multiplication. The key thing is that $n$ is large (512, 2,048, 8,192), so you're doing a lot of compute per byte of weights loaded.

Inference decode (one new token per step).

Each step generates one token. You still need to touch all $P$ parameters, but you're only processing a single token's worth of computation:

For a 7-billion-parameter model, that's about 14 billion FLOPs per token. Sounds like a lot, but the FLOPs aren't the bottleneck. The bottleneck is reading 7 billion parameters (14 GB in float16) from GPU memory. At 2 TB/s memory bandwidth on an A100, that's about 7 ms just to read the weights. The actual multiplication takes far less time than the memory read.

Scale that up to a 70B model and the gap gets absurd. Here's where the 140 GB comes from, using Llama 2 70B's shape:

140 GB at 2 TB/s is roughly 70 ms per token, which caps you at about 14 tokens per second from the bandwidth floor alone. The compute side is 140 billion multiply-adds, which an A100 can crunch through in under a millisecond. The GPU spends more than 99% of each decode step waiting for bytes to arrive.

Worked example

Let me trace through a concrete example. Say we have an 8-token sequence: "The cat sat on the fluffy mat ."

Training view.

One forward pass. All 8 tokens enter the model simultaneously. The attention mechanism computes an $8 \times 8$ score matrix at each head (64 entries, though causal masking zeros out the upper triangle so only 36 matter). The FFN processes all 8 positions in parallel.

Output: 7 predictions. Position 0 predicts token 1, position 1 predicts token 2, and so on through position 6 predicting token 7. We compute cross-entropy loss at each position and average them. One backward pass updates all the weights.

Total cost: 1 forward pass + 1 backward pass. The backward pass is roughly 2x the forward pass, so call it 3 forward-pass equivalents.

Inference view.

Say the prompt is "The" and we want to generate the remaining 7 tokens.

• Step 1: Feed "The" through the model. Sample "cat".
• Step 2: Feed "The cat" through the model. Sample "sat".
• Step 3: Feed "The cat sat". Sample "on".
• Step 4: Feed "The cat sat on". Sample "the".
• Step 5: Feed "The cat sat on the". Sample "fluffy".
• Step 6: Feed "The cat sat on the fluffy". Sample "mat".
• Step 7: Feed "The cat sat on the fluffy mat". Sample ".".

Total cost: 7 forward passes. No backward passes. But each forward pass is sequential, because we can't start step 4 until step 3 produces "on".

Even without the backward pass, inference for those 7 tokens costs 7 forward passes, compared to training's single forward pass for all 8 tokens. Training is embarrassingly parallel. Inference is fundamentally sequential. At each decode step we produce exactly one token but still read every weight in the model.

(This is slightly simplified. As we'll see in Post 5.2, the prompt processing step, called "prefill," can process the entire prompt in parallel, similar to training. It's only the generation of new tokens, "decode," that's sequential. But the basic picture holds.)

One more observation about scale. During training, an 8-token sequence gives us 7 loss terms in one shot. A batch of 32 sequences, each 2,048 tokens long, gives us $32 \times 2{,}047 = 65{,}504$ prediction-and-loss computations from a single forward pass. During inference, we get 1 token per forward pass. The asymmetry in useful work per GPU cycle is hard to overstate.

Implementation sketch

Here's what the training loop and the inference loop look like in code. The contrast is the point.

import torch
import torch.nn.functional as F
 
def train_step(model, tokens, optimizer):
    """
    tokens: (batch, seq_len) — e.g., shape (32, 2048)
 
    One forward pass produces predictions at EVERY position.
    We get seq_len - 1 loss terms from a single pass.
    """
    logits = model(tokens)  # (batch, seq_len, vocab_size)
 
    # Shift: position i predicts token i+1
    predictions = logits[:, :-1, :]    # (batch, seq_len-1, vocab)
    targets = tokens[:, 1:]            # (batch, seq_len-1)
 
    # Cross-entropy loss at every position simultaneously
    loss = F.cross_entropy(
        predictions.reshape(-1, predictions.size(-1)),
        targets.reshape(-1),
    )
 
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()
 
    return loss.item()

@torch.no_grad()
def generate(model, prompt_tokens, max_new_tokens=100, temperature=1.0):
    """
    prompt_tokens: (1, prompt_len)
 
    Each iteration does a FULL forward pass but only uses
    the LAST position's logits to sample one new token.
    """
    tokens = prompt_tokens.clone()
 
    for _ in range(max_new_tokens):
        # Forward pass over the full sequence so far
        logits = model(tokens)            # (1, current_len, vocab)
 
        # Only the LAST position's output matters
        next_logits = logits[:, -1, :]    # (1, vocab)
 
        # Sample from the distribution
        probs = torch.softmax(next_logits / temperature, dim=-1)
        next_token = torch.multinomial(probs, num_samples=1)
 
        # Append and repeat — the feedback loop
        tokens = torch.cat([tokens, next_token], dim=1)
 
        if next_token.item() == eos_token_id:
            break
 
    return tokens

Look at the generate function. On every iteration, we run a forward pass over the entire sequence so far, but we only use the last position's output. We throw away all the intermediate logits. That's wildly wasteful, and it's exactly the problem that the KV cache solves (we'll get there in Post 5.4).

Notice @torch.no_grad(). We disable gradient computation entirely during generation. No computation graph, no backward pass, no optimizer step. The model's weights are frozen. This saves memory (no storing activations for backprop), but it doesn't change the fundamental bottleneck. The weights still have to be read from memory for each forward pass.

And the torch.cat on each iteration: we're literally growing the tensor by one token and re-feeding the whole thing. This is the naive implementation. In practice you'd use a KV cache to avoid recomputing attention for tokens you've already processed. I'm showing the naive version because it makes the algorithmic structure clearest: generate one, append, repeat.

Misconceptions

"Training and inference are the same computation, just with or without gradients." Technically true and deeply misleading. The forward pass math is identical. But the shape of the workload is completely different. Training processes thousands of tokens per forward pass with high arithmetic intensity, and decode processes one token per forward pass with low arithmetic intensity. That changes the hardware bottleneck, the optimization strategies, and the cost structure all at once.

"Inference is cheap because there's no backward pass." The backward pass is roughly 2x the forward pass cost, so sure, you save that. But the real expense of inference is the serial nature of decode, not the per-step cost. You need one forward pass per generated token, and they have to happen sequentially. For a 500-token response, that's 500 sequential forward passes. Training would process those 500 tokens in a single forward pass.

"Bigger batch sizes fix the memory-bound problem." This is partially true, and it's the core idea behind continuous batching in production inference systems. If you batch many user requests together, you amortize the weight-reading cost across more tokens, pushing arithmetic intensity back up. But there's a limit. As batch size grows, the KV cache memory for all those sequences eats into GPU memory, and eventually you run out of space. The tension between batch size and memory is one of the central engineering challenges of LLM serving, and we'll explore it throughout Arc 6.

What's next

This post split the world into training (parallel, compute-bound, embarrassingly easy on the hardware) and inference (sequential, memory-bound, the actual engineering problem). Every optimization in Arcs 5 and 6 lives on one side of that split or tries to bridge it.

The next post covers prefill vs. decode, where we split inference itself into two phases with very different performance profiles. Once you see that the prompt-processing step looks a lot like training, the rest of the arc clicks into place.

Additional reading (and watching)

• Vaswani, A., et al. (2017). Attention Is All You Need. NeurIPS 2017. The architecture that this whole arc takes for granted.

• Pope, R., et al. (2022). Efficiently Scaling Transformer Inference. MLSys 2023. The first paper I'd point at for a careful breakdown of the compute-bound vs. memory-bound regimes in production transformer inference.

• Karpathy, A. (2023). Let's build GPT: from scratch, in code, spelled out. YouTube. Walks through both the training loop and the generate loop side by side. Good source of the "aha" if this post hasn't landed yet.

• Williams, S., Waterman, A., & Patterson, D. (2009). Roofline: An Insightful Visual Performance Model for Multicore Architectures. Communications of the ACM, 52(4), 65-76. The original roofline paper. Surprisingly readable, and the framing has aged well.

• Kwon, W., et al. (2023). Efficient Memory Management for Large Language Model Serving with PagedAttention. SOSP 2023. The vLLM paper. Shows how the memory-bound decode regime drives the need for a completely different memory management strategy at serving time.

• Dao, T. (2023). FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning. Key to understanding how kernel-level optimizations differ between the compute-bound (training) and memory-bound (decode) regimes.

• Leviathan, Y., Kalman, M., & Matias, Y. (2023). Fast Inference from Transformers via Speculative Decoding. ICML 2023. The paper that formalized using a small draft model to cheat the sequential decode bottleneck.