Training View vs. Inference View

Training scores known sequences in parallel. Generation must wait for a selected token before computing the next prediction. The same transformer operations therefore run on different matrix shapes, with different opportunities to reuse weights. We will compare a dense model’s teacher-forced training forward pass with ordinary cached, batch-one generation; sparse models and large serving batches need their own accounting.

The training loop and a model.generate() call can hide that change in workload. Looking at the matrix shapes explains why a GPU that trains efficiently can spend much of a decode step waiting for data.

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 the mask supports parallel teacher-forced training.

The training view

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.

In teacher-forced next-token training, each position conditions on the ground-truth prefix rather than on sampled predictions. Other training procedures can use generated trajectories; teacher forcing describes this choice of conditioning data, not the causal mask itself.

In this teacher-forced objective, later positions do not receive tokens sampled by the model. 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," and its effect depends on the task and training setup.

Large training matmuls can be compute-bound. The GPU is busy doing enormous matrix multiplications across the full sequence length. Attention has an $n \times n$ logical score matrix; fused implementations need not materialize it. Every FFN block multiplies a batch of $n$ vectors through its weight matrices. With suitable shapes and kernels, the arithmetic units can achieve high utilization.

The inference view

Ordinary autoregressive generation selects one new 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 cached decode step processes one new position through every layer of this dense model. Attention combines its query with cached keys and values; the FFN transforms that position’s vector independently. Both operations use model weights, but only attention mixes information across positions.

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.

For a 500-token response, prefill selects the first output token and 499 decode steps produce the rest. At an assumed 30 ms per decode step, that is about 15 seconds after prefill. These serial dependencies matter to latency even when the total FLOP count is smaller than a training step.

Small-batch decode is often 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: work per byte read

Arithmetic intensity measures compute operations (FLOPs) per byte transferred across a specified memory boundary:

During training, you're doing a matrix multiplication like $X \cdot W$ where $X$ is a batch of 2,048 vectors. In an ideal weights-only estimate, the bytes of $W$ are amortized across all 2,048 vectors. Actual kernels load tiles and may read some bytes more than once. The ratio of compute to memory access is high. For large batch sizes and long sequences, a weights-only estimate gives hundreds to thousands of FLOPs per byte. Actual intensity includes activation traffic, tiling, and backward-pass work.

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.

An A100 80 GB SXM has peak dense FP16 tensor throughput of 312 TFLOP/s and about 2 TB/s of HBM bandwidth (rounded from 2,039 GB/s). That ratio, roughly 156 FLOPs per byte, is the GPU's "ridge point." A workload to the right of that ridge has a compute-limited ceiling; to the left, memory bandwidth sets the lower ceiling. Kernel efficiency determines how close execution gets to either bound.

The roofline model, from Williams, Waterman, and Patterson (2009), relates arithmetic intensity to the hardware's compute and memory-bandwidth limits.

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. To the left, the bandwidth ceiling is lower; to the right, the compute ceiling is lower. Other overhead can keep a workload below either ceiling.

Watch where the markers land. Decode (batch 1) sits at AI = 1, far down on the orange slope. The idealized performance ceiling there is roughly 2 TFLOPS, less than 1% of the GPU's peak. Training at AI = 2,048 reaches the compute-limited region of this idealized roofline; it need not attain the ceiling.

Batching raises weight reuse by processing several requests together. Weight quantization reduces bytes per parameter, while speculative decoding verifies several candidates in one target pass. The resulting hardware limits depend on precision and device: an H100 SXM at dense FP16 rates has a ridge near 295 FLOPs/byte, compared with 156 in this rounded A100 example.

The minimal math

Consider 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 estimate gives about 14 billion FLOPs per token. Reading 14 GB of FP16 weights at an assumed 2 TB/s takes at least 7 ms. In this weights-only roofline calculation, the memory bound is much larger than the compute bound.

For a 70B model, the raw FP16 weights occupy 140 GB. Here's the calculation using Llama 2 70B's shape:

A hypothetical 2 TB/s memory system needs about 70 ms to move 140 GB. The dense projection work is about 140 billion FLOPs (70 billion multiply-adds), whose idealized 312 TFLOP/s lower bound is 0.45 ms. These are separate roofline bounds, not measured sequential phases. A 70B FP16 model does not fit on one 80 GB A100; sharding changes the available aggregate bandwidth and adds communication.

Worked 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 the causal mask sets upper-triangle logits to negative infinity before softmax, leaving 36 allowed entries). 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 computes gradients; the optimizer then updates 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. Each step below names the logical context. With a KV cache, step 1 processes the prompt; later steps process only the newly selected token and reuse the earlier keys and values.

• Step 1: Context: "The". Sample "cat".
• Step 2: Context: "The cat". Sample "sat".
• Step 3: Context: "The cat sat". Sample "on".
• Step 4: Context: "The cat sat on". Sample "the".
• Step 5: Context: "The cat sat on the". Sample "fluffy".
• Step 6: Context: "The cat sat on the fluffy". Sample "mat".
• Step 7: Context: "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".

A forward-pass count measures dependency, not arithmetic: cached decode handles one new row per pass, while training processes all rows and adds a backward pass.

The first pass processes the known prompt, even when it contains many tokens. Post 5.2 separates that prefill work from the subsequent decode steps.

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. A serving batch can process several independent requests together, but each request still waits for its preceding token.

Implementation sketch

The training loop computes losses over known positions; the generation loop appends its own selected tokens:

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.
    """
    model.train()
    optimizer.zero_grad()
    logits = model(tokens)  # application model returns a logits tensor
 
    # 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, eos_token_id, 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.
    """
    if temperature <= 0:
        raise ValueError("this sampling example requires temperature > 0")
    model.eval()
    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

In this naive generate function, each iteration recomputes hidden states for the entire growing sequence while using only the final position’s logits. A KV cache avoids recomputing old positions through the transformer; computing only the final logits is a separate possible optimization. We derive caching in Post 5.4.

@torch.no_grad() prevents autograd from recording the forward pass or retaining activations for a backward pass. Generation here does not update the weights, so it needs neither those records nor an optimizer step. The forward computation still reads the weights.

Misconceptions

"Training and inference are the same computation, just with or without gradients." The same forward-pass operations run on different shapes. Training processes thousands of tokens per forward pass with high arithmetic intensity, and batch-one decode processes one new token per forward pass with low weight reuse. That changes the hardware bottleneck, the optimization strategies, and the cost structure all at once.

"Inference is cheap because there's no backward pass." Inference avoids backward-pass work, but ordinary generation still requires successive model calls. Its latency depends on both that dependency chain and the time spent in each 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 batch must fit both the weights and the requests’ growing caches.

What's next

The next post covers prefill vs. decode, where we split inference itself into two phases with very different performance profiles.

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. Analyzes compute and memory limits in 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.

• 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.

• 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. Verifies candidates from a smaller draft model in a target-model pass.