The Generation Loop

Notes

The generation loop is the autoregressive inference algorithm that repeatedly computes next-token logits, selects one token, appends it to the context, and feeds the longer sequence back in: given tokens $x_{1:t}$, run $z_t=f_{\theta }(x_{1:t})$, sample or choose $x_{t+1}\sim \mathrm{Decode}(\mathrm{softmax}(z_t/T))$, then set $x_{1:t+1}=[x_{1:t},x_{t+1}]$ and repeat until EOS or $L_{\max }$. In Karpathy’s nanoGPT, this is the core of `GPT.generate`: crop context to `block_size`, run `logits, _ = self(idx_cond)`, take `logits = logits[:, -1, :]`, optionally divide by `temperature`, optionally `top_k` filter, compute `probs = F.softmax(logits, dim=-1)`, draw `idx_next = torch.multinomial(probs, num_samples=1)`, then concatenate `idx = torch.cat((idx, idx_next), dim=1)`.

Example: suppose the prompt is token ids $[10,42,7]$, vocabulary size $5$, and the model’s last-position logits are $z=[1.0,3.0,0.0,2.0,-1.0]$. With temperature $T=1$, $\mathrm{softmax}(z)\approx [0.086,0.636,0.032,0.234,0.012]$, so multinomial sampling might draw token $1$. Append it: $[10,42,7]\to [10,42,7,1]$, run the model again, and suppose the next logits are $[0.5,0.2,2.5,0.0,-0.5]$ with probabilities $\approx [0.104,0.077,0.766,0.063,0.038]$, so token $2$ is drawn and appended. After two generated tokens, the sequence is $[10,42,7,1,2]$; the loop stops if token $2$ is EOS, or continues until `max_new_tokens` is reached.

1. Greedy decoding replaces sampling with $x_{t+1}=\arg {\max }_i{z}_i$; for $z=[1,3,0,2,-1]$, greedy always picks token $1$. 2. Temperature rescales logits: $p_i=\frac{e^{z_i/T}}{{\sum }_j{e}^{z_j/T}}$; with $T=0.5$, $z/T=[2,6,0,4,-2]$ and $p\approx [0.016,0.865,0.002,0.117,0.0003]$, making outputs less random. 3. Top-$k$ keeps only the $k$ largest logits; with $k=2$ for $z=[1,3,0,2,-1]$, keep tokens $1,3$, set others to $-\infty$, and get probabilities $\approx [0,0.731,0,0.269,0]$. 4. Without KV cache, step $t$ recomputes attention over the full prefix, costing roughly ${\sum }_{s=1}^{L}O(s^2)=O(L^3)$ attention work across the generated sequence; with KV cache, each new step attends one query to $s$ cached keys, costing ${\sum }_{s=1}^{L}O(s)=O(L^2)$. 5. Context windows force cropping: if `block_size=4` and current tokens are $[10,42,7,1,2]`, nanoGPT feeds only `[42,7,1,2]` as `idx_cond`, so tokens before the window no longer affect logits.

In real serving systems, this same loop is split into a prefill phase and a decode phase: prefill processes the full prompt, then decode repeatedly runs one-token forward passes using the KV cache. In vLLM, see `LLMEngine.step()` scheduling requests, model execution through worker/model runner paths such as `vllm/worker/model_runner.py`, and sampling logic in `vllm/model_executor/layers/sampler.py`. In HuggingFace TGI, the loop is implemented around batching, prefill, decode, and token selection in the server code, especially `text-generation-inference/server/text_generation_server/models/` and generation/token selection utilities. In TensorRT-LLM, the runtime generation loop appears in executor/session paths such as `tensorrt_llm/runtime/generation.py`, `tensorrt_llm/runtime/model_runner.py`, and C++ executor backends that repeatedly launch optimized decode kernels.

Karpathy’s nanoGPT version is intentionally simple because it stores the whole growing tensor `idx` and calls the model each iteration; for prompt length $3$ and `max_new_tokens=2`, it performs two full forward calls over lengths $3$ and $4$. Production systems instead maintain per-request state: `tokens=[10,42,7,1,2]`, `kv_cache` tensors per layer, stop flags, RNG state, and sampling parameters. The algorithmic shape is unchanged: $\text{while not done: logits}\leftarrow \text{forward(last token, cache)};\text{next}\leftarrow \text{sample(logits)};\text{append(next)}$. The engineering challenge is batching many such loops whose requests have different prompt lengths, generation lengths, and stopping times.