Forward Pass vs Generation

Notes

A training forward pass computes logits for many known target tokens in parallel, while autoregressive generation repeatedly runs the model to sample exactly one next token per sequence and feeds it back into the input. Training: $X\in {\mathbb{Z}}^{B\times T}$, compute $H=f_{\theta }(X)$, logits $L=H{W}_{\text{lm}}\in {\mathbb{R}}^{B\times T\times V}$, loss $\mathcal{L}=-{\sum }_{b,t}\log p_{\theta }(x_{b,t+1}\mid x_{b,\le t})$ in one parallel pass. Generation loop: initialize prompt $x_{1:T}$, then for $i=\mathrm{1..}N$: $l_i=f_{\theta }(x_{\le T+i-1})[-1]$, $p_i=\mathrm{softmax}(l_i/\tau )$, $x_{T+i}\sim p_i$, append token. With KV cache, step $i$ becomes $l_i=f_{\theta }(x_{T+i-1},K_{\le T+i-2},V_{\le T+i-2})[-1]$ instead of recomputing the whole prefix.

Example: let $B=2$, prompt length $T=4$, vocab $V=50,000$, generate $N=3$ tokens. Training forward on two length-4 examples computes logits shape $2\times 4\times 50,000$, i.e. predictions for 8 token positions in one call, then teacher-forces against known labels. Generation prefill computes prompt logits once for shape $2\times 4\times 50,000$ and stores KV cache for 4 positions. Decode step 1 runs only the last token per sequence, produces logits shape $2\times 1\times 50,000$, samples e.g. $[912,44]$; step 2 consumes $[912,44]$, samples $[18,301]$; step 3 consumes $[18,301]$, samples $[502,9]$. The workload changed from one large parallel matrix-heavy pass to many latency-sensitive sequential decode iterations.

1. Parallelism: training parallelizes over $B\times T$ tokens; decoding parallelizes mostly over batch $B$ because time dimension is sequential: $x_{t+1}$ cannot be known before sampling $x_t$. 2. Cost split: generation has prefill cost $O(T^2)$ attention over the prompt and decode cost $O(NT)$ with KV cache, growing one token at a time; without KV cache decode would be $O(N(T+N{)}^2)$ recomputation. 3. Memory: training stores activations for backprop, roughly proportional to layers $\times B\times T\times d$; inference stores KV cache, roughly $2\times L\times B\times T_{\text{current}}\times n_{\text{kv}}\times d_{\text{head}}$. 4. Batching: training batches are static rectangular tensors; generation batches are dynamic because requests have different prompt lengths, output lengths, stop conditions, and sampling parameters. 5. Objective: training computes loss under teacher forcing; generation performs search or sampling using greedy, top-$k$, top-$p$, temperature, beam search, penalties, and stop tokens.

This distinction is explicit in serving systems. In vLLM, prefill/decode scheduling is handled around `vllm/engine/llm_engine.py` and `vllm/core/scheduler.py`, with paged KV cache managed by block allocators such as `vllm/core/block_manager.py`; model execution separates prompt/prefill and decode metadata in worker/model-runner paths such as `vllm/worker/model_runner.py`. HuggingFace TGI implements the continuous generation loop in Rust around `router/src/infer.rs` and batching/scheduling in `router/src/batcher.rs`, while model backends consume prefill/decode batches and return next-token logits through server code under `server/text_generation_server/`. TensorRT-LLM exposes the same split via its inflight batching executor, with generation logic in `cpp/tensorrt_llm/batch_manager/` and Python runtime examples using `tensorrt_llm.runtime.GenerationSession`.

The practical consequence is that inference is not “training without backward.” For a 32-layer model with $n_{\text{kv}}=32$, $d_{\text{head}}=128$, FP16 KV, one sequence of length $4096$ uses $2\times 32\times 4096\times 32\times 128\times 2\approx 2.15$ GB of KV cache, before weights and workspace. If 100 users each decode one token at a time, the server must schedule 100 tiny dependent steps, not one static tensor batch. This is why production engines optimize KV cache paging, continuous batching, prefix caching, speculative decoding, and fast sampling rather than only raw forward throughput.