Inference From Scratch

This book is a work in progress. Chapters are being released incrementally.

What is this book?

Most engineers who work with language models treat them as black boxes: send a prompt, get a response. This book opens every box.

You will build a complete LLM inference engine from scratch — in Python, in plain PyTorch, one concept at a time. By the end, you will have written your own autoregressive loop, your own attention kernel, your own KV cache, your own scheduler, your own batching system, and your own speculative decoder. Every optimization will be measured before and after against the same HTTP benchmark harness you build in Chapter 1.

The engine you end up with is a simplified vLLM, built entirely by you, line by line.

Who is this for?

ML engineers who use inference systems but want to understand what is happening inside them
Systems engineers curious about what makes LLM serving different from regular serving
Researchers who want to implement and experiment with inference optimizations
Anyone who has read “Attention is All You Need” and wants to see how that attention becomes a production serving system

Prerequisites: Python, basic PyTorch, some familiarity with neural networks. No CUDA or GPU programming experience required — that is covered from scratch.

What you will build

A single serving engine that evolves chapter by chapter. Each chapter adds one optimization, measures the speedup, and swaps it into the running server in one line. The benchmark numbers are real — run on the same hardware, the same model, the same prompts.

By the end of Chapter 13, your engine handles:

✅ Autoregressive generation with your own loop
✅ Full transformer model built from scratch with real weights loaded
✅ KV cache eliminating redundant computation
✅ Static and continuous batching for multi-request throughput
✅ Flash attention reducing memory from O(n²) to O(n)
✅ Paged attention eliminating memory fragmentation
✅ Prefix caching reusing KV blocks across requests
✅ INT4 quantization halving memory footprint
✅ Speculative decoding generating multiple tokens per step

Chapters

#	Question	Key concept	Status
1	What does an inference engine look like?	Baseline HF server, TTFT/TPOT/throughput	🚧 Coming soon
2	How does generation actually work?	Autoregressive loop, prefill vs decode	🚧 Coming soon
3	What is the model doing?	Transformer from scratch, load real weights	🚧 Coming soon
4	Why is decode slow, and how do we fix it?	KV cache	🚧 Coming soon
5	How do we serve multiple requests at once?	Static batching, scheduler	🚧 Coming soon
6	Profiling — where is time actually spent?	Roofline model, memory vs compute bound	🚧 Coming soon
7	Writing GPU kernels with Triton	GEMM, tiling, shared memory	🚧 Coming soon
8	How do we make attention faster?	Online softmax, Flash Attention	🚧 Coming soon
9	Continuous batching	Iteration-level scheduling, no padding waste	🚧 Coming soon
10	How do we manage memory efficiently?	Paged attention, block tables	🚧 Coming soon
11	How do we avoid recomputing common prefixes?	Prefix caching, hash-based block reuse	🚧 Coming soon
12	How do we run bigger models on less memory?	INT4/GPTQ quantization	🚧 Coming soon
13	How do we generate faster without changing the model?	Speculative decoding, draft-target	🚧 Coming soon

The code

Every chapter has a companion directory of numbered Python files. Each file is self-contained and runnable. The files build on each other within the chapter — run them in order.

chapters/
├── ch01_baseline/
│   ├── 00_start.py        ← run this first
│   ├── 01_load.py
│   ├── 02_sampling.py
│   ├── 03_server.py
│   └── 04_benchmark.py
├── ch02_generate_loop/
│   └── ...

Clone the repo and run any chapter independently:

git clone https://github.com/arunpatala/inference-from-scratch-book.git
cd inference-from-scratch-book/CODE/chapters/ch01_baseline
pip install torch transformers fastapi uvicorn requests
python 00_start.py

A note on the model

All examples use small, freely available models that run on CPU so you can follow along without a GPU. The optimizations are GPU-native — the book explains what each one does on GPU hardware and why it matters at scale — but the code runs and produces correct output on any machine.

Built in public. Feedback and corrections welcome via GitHub issues.

What does an inference engine look like?

An inference engine is a surprisingly small thing. Strip away the HTTP routing, the batching logic, the memory management, and the GPU kernels, and what remains is a program that does three things: it takes text in, runs it through a neural network, and returns text out. That loop, at its core, is what this entire book is about.

This chapter builds the simplest possible version of that loop in one sitting. By the end you will have a working server that accepts prompts over HTTP and returns generated text, along with a benchmark harness to measure how fast it runs. Every subsequent chapter takes one component from this baseline and replaces it with something faster, more capable, or more principled — but the overall shape stays the same.

The chapter is organised as a sequence of self-contained files, each one adding a single layer of understanding. Read them in order, run each one, and by the time you reach the benchmark you will have a clear mental model of what an inference engine actually does — and exactly where the interesting engineering problems live.

File	What it covers
`00_start.py`	The complete engine in 35 lines
`01_load.py`	What model loading is actually doing
`02_sampling.py`	How temperature, top-k, and top-p change the output
`03_contracts.py`	The shared types every chapter uses
`04_engine.py`	The engine abstraction
`05_server.py`	Wrapping the engine in an HTTP API
`06_benchmark_offline.py`	Measuring throughput directly
`06_benchmark_online.py`	Measuring throughput over HTTP

Start here

This is the entire inference engine.

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from config import MODEL_NAME, PROMPT

tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model     = AutoModelForCausalLM.from_pretrained(MODEL_NAME)

def generate_text(prompt, model, tokenizer, max_length=512, temperature=1, top_k=50, top_p=0.95):
    inputs  = tokenizer.encode(prompt, return_tensors="pt")
    outputs = model.generate(
        inputs,
        max_length=max_length,
        temperature=temperature,
        top_k=top_k,
        top_p=top_p,
        do_sample=True,
    )
    return tokenizer.decode(outputs[0], skip_special_tokens=True)

print(generate_text(PROMPT, model, tokenizer))

Run it:

python 00_start.py

You will see a pause, then some text. That pause is the entire cost of being an AI: loading hundreds of millions of numbers from disk into memory, arranging them into matrices, and preparing to multiply them together at high speed. The text that follows is the result of doing exactly that, once per token, until the model decides it is done.

Every line of the program above is currently a black box. That is intentional. The goal of this first pass is to see the whole thing working before we start taking it apart. The rest of this chapter, and the rest of this book, opens each box in turn.

What just happened

There are three steps inside generate_text, and they are the same three steps inside every language model inference system ever built, from the smallest research prototype to the largest production cluster.

Step 1: Tokenization

inputs = tokenizer.encode(prompt, return_tensors="pt")

A language model cannot read text. What it reads is integers — specifically, indices into a fixed vocabulary of tens of thousands of subword units called tokens. The tokenizer’s job is to convert the raw string into that sequence of integers.

The word “tokenizer” hides a surprising amount of machinery. Before a single integer is produced, the input text is normalised (Unicode is standardised, whitespace is cleaned), then split into preliminary chunks on word boundaries and punctuation, then passed through a learned compression algorithm — typically Byte Pair Encoding or a variant of it — that merges common character sequences into single tokens. A word like “unbelievable” might become three tokens: “un”, “believ”, “able”. A rare technical term might be split into individual characters. The choice of how to split determines how efficiently the model uses its fixed context window, and different models make different choices.

The result of tokenizer.encode is a PyTorch tensor of integers. The sentence “According to all known laws of aviation, there is no way a bee should be able to fly” becomes something like [12902, 2311, 477, 3245, 4141, ...] — the exact values depend on the model’s vocabulary. This is what the model sees. It has never seen the characters ‘b’, ‘e’, ‘e’ in sequence; it has seen the integer that represents the token “bee”. The argument return_tensors="pt" wraps those integers in a PyTorch tensor so they can be passed directly to the model.

Step 2: Generation

outputs = model.generate(
    inputs,
    max_length=max_length,
    temperature=temperature,   # controls randomness (0 = deterministic)
    top_k=top_k,               # only consider top-k most likely tokens
    top_p=top_p,               # only consider tokens in top-p probability mass
    do_sample=True,
)

This single call is where the interesting work happens, and also where the most is hidden. Inside model.generate, the model runs a two-phase process that will occupy several later chapters.

The first phase is called prefill. The model takes the full sequence of input token IDs and runs them through all of its layers in a single parallel forward pass. Every token attends to every earlier token simultaneously, building up rich contextual representations. The output of prefill is a probability distribution over the vocabulary: given this prompt, how likely is each possible next token? This phase is compute-intensive because it processes all tokens at once, and its cost grows with the length of the input.

The second phase is called decode. The model samples one token from that distribution, appends it to the sequence, and then runs another forward pass — but this time processing only the new token, because the representations of all earlier tokens were computed in prefill and can be reused. The model outputs another distribution, another token is sampled, and the loop continues. Each decode step is cheap compared to prefill, but the steps must happen sequentially: you cannot start computing token five until token four has been chosen. This sequential constraint is the fundamental reason why inference is hard to scale, and it is the thread that runs through most of the optimisations in later chapters.

The parameters temperature, top_k, and top_p all control the sampling step — how the model chooses which token to pick from the distribution at each decode step. At temperature=0 the choice is deterministic: always take the most likely token, which is called greedy decoding. At higher temperatures the distribution is flattened, making lower-probability tokens more likely and the output more varied. do_sample=True tells HuggingFace to use sampling rather than greedy search. Section 1.2 explores all three parameters in detail.

Step 3: Detokenization

return tokenizer.decode(outputs[0], skip_special_tokens=True)

The third step is the inverse of the first. model.generate returns a tensor of token IDs that includes both the original prompt and the newly generated tokens. outputs[0] selects the first (and only) sequence in the batch. tokenizer.decode maps each integer back to its string representation and joins them into a single string. The argument skip_special_tokens=True strips out control tokens like the model’s end-of-sequence marker, which should not appear in the final output.

The full pipeline, then, is: text → integers → neural network loop → integers → text. Everything else in this book is engineering that makes the middle step faster, cheaper, or more capable of handling many requests at once.

What to notice when you run it

There are two pauses. The first, before any output appears, is the model loading: the weights being read from disk and arranged in memory. The second, shorter pause is generation itself. On a CPU with the 10M parameter model used in this chapter, generation takes a fraction of a second per response. On a GPU with a 70B parameter model, it takes several seconds — which is exactly why throughput optimisation matters and why this book exists.

Run the program a second time without changing anything. The output will be different. That is sampling: the model is not computing a single deterministic answer, it is drawing from a probability distribution, and that distribution has many plausible continuations. Whether that nondeterminism is a feature or a bug depends on the application, and controlling it precisely is what the sampling parameters are for.

What comes next

Each subsequent section in this chapter takes one part of the program above and opens it. 01_load.py shows what AutoTokenizer.from_pretrained and AutoModelForCausalLM.from_pretrained actually do — what is downloaded, how the weights are laid out in memory, and how long loading takes as model size grows. 02_sampling.py shows what happens when you change temperature, top_k, and top_p, running the same prompt through each configuration so you can see the difference directly. 03_contracts.py introduces the shared data types that every chapter in this book uses, so that when we replace the engine in later chapters the surrounding code stays the same.

By the time you reach the server and the benchmark, you will have seen every component of this system working in isolation. Then, when Chapter 2 starts changing things, you will know exactly which box is being opened and why.

Loading a model

The first two lines of 00_start.py are the most expensive lines in the entire program.

tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model     = AutoModelForCausalLM.from_pretrained(MODEL_NAME, torch_dtype=torch.float32)

On a cold start, these two calls download several files from the HuggingFace Hub, parse a configuration file, construct a neural network architecture in memory, and then copy hundreds of millions of numbers into that architecture one tensor at a time. On subsequent runs the files are cached locally, so only the last two steps repeat — but even then, loading a large model can take several seconds. Understanding exactly what is happening in this time is the subject of this section.

Run:

python 01_load.py
python 01_load.py --model HuggingFaceTB/SmolLM2-135M

What `AutoTokenizer.from_pretrained` downloads

tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)

A tokenizer is not a neural network — it is a deterministic lookup table combined with a set of rules. from_pretrained downloads two things. The first is a vocabulary file: a mapping from every token the model knows (typically 32,000 to 128,000 subword strings) to an integer index. The second is a configuration file that specifies which tokenization algorithm to use — Byte Pair Encoding (BPE), WordPiece, Unigram, or one of several others — along with any special tokens the model was trained with, such as the beginning-of-sequence marker or the padding token.

The Auto prefix means that HuggingFace inspects the configuration file and automatically selects the correct tokenizer class for the model you asked for. If the model was trained with a BPE tokenizer, you get a BPETokenizer. If it uses a SentencePiece model, you get that instead. You do not need to know which type is being used; the encode and decode methods work identically regardless.

Once loaded, the tokenizer is fast. It runs entirely on CPU, requires no GPU, and processes most inputs in under a millisecond. Its output — a tensor of integer indices — is small: a 200-word prompt might produce 250 token IDs, each one a 32-bit integer.

What `AutoModelForCausalLM.from_pretrained` does

model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, torch_dtype=torch.float32)

This call does three things in sequence: it reads the model’s configuration, constructs the neural network architecture in memory, and loads the trained weights into it.

The configuration file — typically config.json in the model repository — describes the shape of the network: how many layers it has, how wide those layers are, how many attention heads each layer contains, what its vocabulary size is, and which activation function to use. AutoModelForCausalLM reads this file, looks up the right model class in an internal registry (mapping “llama” in the config to LlamaForCausalLM, for instance), and calls its constructor. At this point the network exists in memory but its parameters are random or uninitialized — it has the right shape, but no knowledge.

The second step copies the trained parameters into that skeleton. The weights are stored on disk as a safetensors file — a format designed to be faster and safer than the legacy .pkl format. Each tensor in the file is matched by name to a parameter in the model and copied in. For a 10M parameter model this takes milliseconds. For a 70B parameter model it takes minutes — each parameter is a 32-bit float requiring 4 bytes of memory, so 70 billion parameters occupy 280 gigabytes in float32. In practice, production models are almost always loaded in float16 or bfloat16 (2 bytes per parameter), halving that to 140 GB. Quantisation to 8-bit or 4-bit reduces it further still, and Chapters 12 and 13 address this directly.

The torch_dtype=torch.float32 argument we pass in the chapter code forces all weights to be stored as 32-bit floats. This makes the arithmetic easier to reason about — the numbers are exactly representable, without the rounding that lower-precision formats introduce — at the cost of using twice the memory of float16. For the small model used in this chapter, the difference is negligible.

Placing the model on a device

model = model.to(DEVICE).eval()

PyTorch models and tensors always live on a specific device — either the CPU or a particular GPU. When a model is first loaded with from_pretrained, it lands on CPU. The call .to(DEVICE) moves every parameter tensor to the target device. If DEVICE is "cuda", each tensor is copied from system RAM into GPU VRAM over the PCIe bus. If DEVICE is "cpu", nothing moves — the model stays where it is.

The call .eval() switches the model into inference mode. During training, certain layers behave differently than they do at inference time. Dropout layers randomly zero out activations to prevent overfitting; at inference time, they should pass their input through unchanged. Batch normalisation layers use per-batch statistics during training and running statistics during inference. Calling .eval() activates the inference behaviour in all such layers. Omitting this call produces subtly wrong results that can be difficult to diagnose, because the model will still run — it just will not run correctly.

The two calls are chained because .to() returns the model itself, making it idiomatic to write them on a single line.

Measuring what you loaded

params = sum(p.numel() for p in model.parameters())
print(f"  parameters : {params / 1e6:.1f}M")

Every nn.Parameter in the model is a tensor. p.numel() returns the total number of elements in that tensor — for a weight matrix of shape [4096, 4096] that is 16,777,216 elements. Summing this across all parameters gives the total parameter count. Dividing by one million and formatting to one decimal place gives the familiar shorthand: “10.5M”, “135M”, “7B”.

Parameter count is a useful proxy for model capability and cost, but it is not the whole story. Memory consumption also depends on the data type: a 7B parameter model in float16 occupies 14 GB, but the same model in float32 occupies 28 GB. GPU memory also includes space for the KV cache during generation — the amount of memory needed to store the intermediate attention keys and values for every token in every layer — which can rival the model weights in size for long contexts. Chapter 4 covers the KV cache in detail.

Verifying it works: a tokenizer smoke test

sample = "Hello, world!"
ids    = tokenizer.encode(sample, return_tensors="pt")
print(f"  text       : '{sample}'")
print(f"  token ids  : {ids.tolist()}")
print(f"  decoded    : '{tokenizer.decode(ids[0])}'")

Before running any generation, it is worth confirming that the tokenizer is working correctly. This snippet encodes a short string to token IDs and then decodes it back. If the decoded string matches the input — modulo any normalisation the tokenizer applies — everything is wired up correctly. It is a cheap sanity check that catches common mistakes like loading the wrong tokenizer for a given model, which produces plausible-looking but incorrect token IDs.

The round-trip encode → decode should be lossless for any string the tokenizer supports. Punctuation and whitespace can sometimes look different after the round-trip because the tokenizer uses a byte-level representation internally, but the semantic content should be identical.

What comes next

At this point, the model and tokenizer are loaded, placed on the correct device, and verified to be working. The next section, 02_sampling.py, keeps the model exactly as it is and instead explores how changing the sampling parameters — temperature, top_k, and top_p — changes the character of the generated text. The model is not the variable; the sampling strategy is.

Sampling

Every time a language model generates a token it produces a vector of raw scores — one per word in its vocabulary — and then has to choose which word to output next. How it makes that choice has an enormous impact on the character of the text it produces. The same model, given the same prompt, can output a careful and predictable answer, a varied and creative one, or something approaching nonsense, depending entirely on the sampling strategy applied to those scores. This section makes that choice concrete and observable.

def generate(model, tokenizer, prompt: str, max_new_tokens: int = BENCHMARK_TOKENS, **kwargs):
    input_ids = tokenizer.encode(prompt, return_tensors="pt").to(DEVICE)
    with torch.no_grad():
        output_ids = model.generate(input_ids, max_new_tokens=max_new_tokens, **kwargs)
    new_ids = output_ids[0][input_ids.shape[1]:]
    return tokenizer.decode(new_ids, skip_special_tokens=True)

Run it:

python 02_sampling.py
python 02_sampling.py --prompt "Once upon a time"

The generate helper above strips the prompt tokens from the output before decoding — output_ids[0][input_ids.shape[1]:] selects only the newly generated tokens — so what you see printed is the model’s continuation, not a repetition of the input. The **kwargs pattern passes whatever sampling arguments the caller provides directly to model.generate, which is what allows the demo below to explore every strategy through the same function.

From logits to a probability distribution

Before any sampling strategy makes sense, it helps to understand what the model actually produces at each decode step. The final layer of the network outputs a vector of floating-point numbers called logits — one number for each token in the vocabulary, which for most modern models means somewhere between 32,000 and 128,000 values. A logit is a raw score: a higher number means the model considers that token a more plausible continuation of the sequence so far. The numbers themselves are not probabilities; they can be negative, arbitrarily large, and they do not sum to one.

To turn logits into probabilities the model applies the softmax function, which exponentiates each logit and then divides by the sum of all exponentiated logits. The result is a proper probability distribution: every value is positive and they sum to exactly one. A token with a logit of 3.5 and one with a logit of 3.0 will not have probabilities of 3.5% and 3.0% — the exponential function amplifies differences, so the gap between them in probability space is larger than the gap in logit space.

Once the probabilities are computed, some form of sampling algorithm picks the next token. The sampling parameters — temperature, top_k, top_p — all operate on the logits before or alongside the softmax step, reshaping the distribution before a token is drawn from it. Chapter 3 will implement this from scratch; for now the goal is to observe the effects.

Greedy decoding

{"do_sample": False}

The simplest strategy is to always pick the token with the highest probability. This is called greedy decoding, and it is what do_sample=False selects. The demo runs the same prompt twice with identical settings.

{"label": "Greedy / temperature=0  [run 1]", "kwargs": {"do_sample": False}},
{"label": "Greedy / temperature=0  [run 2]", "kwargs": {"do_sample": False}},

Both runs produce identical output. That is the defining property of greedy decoding: given the same prompt and the same model weights, the output is completely deterministic. There is no randomness anywhere in the computation. This makes greedy decoding reproducible and easy to debug, which is why it is the default for tasks like translation or summarisation where there is typically one correct answer. Its weakness is that it tends toward repetitive, formulaic text — the most probable token at each step is often also the most expected one, and the model can get trapped following its own highest-probability paths into loops.

Temperature

{"do_sample": True, "temperature": 0.1},
{"do_sample": True, "temperature": 1.0},
{"do_sample": True, "temperature": 2.0},

Temperature is a scalar that divides every logit before the softmax is applied. The modified formula is:

\[ p_i = \frac{e^{x_i / T}}{\sum_{j} e^{x_j / T}} \]

where \(T\) is the temperature and \(x_i\) is the raw logit for token \(i\). When \(T < 1\), dividing by a number less than one makes each logit larger in magnitude. Because the softmax uses exponentials, this amplifies the differences between tokens: already-probable tokens become even more dominant, and unlikely tokens become nearly impossible. The output concentrates around the model’s top choices and becomes more predictable — and more prone to repetition. When \(T > 1\), the logits shrink toward zero, softmax differences flatten, and the probability mass spreads more evenly across the vocabulary. Lower-ranked tokens get a larger share of the probability, and the output becomes more varied — at the cost of sometimes choosing tokens that are implausible in context.

At temperature=0 the difference between logits becomes infinite in the limit, and the distribution collapses to a point mass on the single highest-probability token — equivalent to greedy decoding. At very high temperatures the distribution approaches uniform, meaning every token in the vocabulary has roughly equal probability regardless of context.

Running the demo shows this directly. At temperature=0.1 the two runs with the same prompt produce nearly identical output — the distribution is so concentrated that the same token wins almost every time. At temperature=1.0 the two runs diverge. At temperature=2.0 the output may be surprising, non sequitur, or grammatically strange, because the model is now sampling from a flattened distribution that gives real probability to tokens it would ordinarily never choose.

Top-k sampling

{"do_sample": True, "temperature": 1.0, "top_k": 1},
{"do_sample": True, "temperature": 1.0, "top_k": 5},
{"do_sample": True, "temperature": 1.0, "top_k": 200},

Top-k sampling addresses a specific problem with pure temperature sampling: even at moderate temperatures, the long tail of the vocabulary — tokens with very low probability — retains some nonzero probability mass, and occasionally one of them gets sampled. Most of those tokens are genuinely bad choices that no amount of creativity justifies.

Top-k solves this by discarding all but the \(k\) most probable tokens before sampling. The logits of the remaining vocabulary are set to negative infinity, which makes their softmax probability exactly zero. Sampling then proceeds over just the \(k\) survivors. When \(k = 1\), only the single best token survives and the result is identical to greedy decoding. When \(k\) equals the full vocabulary size, nothing is filtered and the result is identical to unconstrained sampling. Values between those extremes offer a tunable trade-off.

The weakness of top-k is that it uses a fixed count regardless of the shape of the distribution. At a position where the model is highly confident — the next token is almost certainly “the” — a top-k of 50 still allows 49 other tokens to compete, including plausible but wrong ones. At a position where the model is genuinely uncertain — many continuations are reasonable — a top-k of 5 may be too restrictive. The right value of \(k\) is context-dependent in a way the parameter cannot capture.

Top-p (nucleus) sampling

{"do_sample": True, "temperature": 1.0, "top_p": 0.1},
{"do_sample": True, "temperature": 1.0, "top_p": 0.95},

Top-p sampling, introduced by Holtzman et al. in 2019, solves the rigidity of top-k by filtering based on cumulative probability mass rather than a fixed count. Tokens are ranked from most to least probable, and the model keeps adding tokens to the candidate set until their cumulative probability exceeds the threshold \(p\). All remaining tokens are set to zero probability. Sampling then proceeds over the nucleus — the smallest set of tokens that together account for at least \(p\) of the probability mass.

The result is a candidate set that adapts to the situation. When the model is confident, probability mass is concentrated in a few tokens and the nucleus is small — perhaps two or three tokens — even if \(p = 0.95\). When the model is uncertain, the mass is spread across many tokens and the nucleus expands accordingly. The number of candidates is not fixed; it is determined by the shape of the distribution at each step.

In the demo, top_p=0.1 produces very conservative output: only the tokens that together account for the first 10% of probability mass are considered, which typically means just one or two tokens. top_p=0.95 allows a much richer set of candidates. Most production systems use values between 0.9 and 0.95 as a reasonable default.

What to observe when you run it

The most informative thing to look for is the greedy runs first. Their identical outputs confirm that when sampling is disabled the computation is purely deterministic. Then compare the two temperature=1.0 runs — they should differ, sometimes substantially. That difference comes entirely from the randomness introduced by sampling, not from anything in the model itself.

The comparison between top_k=1 and do_sample=False is also instructive: they produce identical output, which confirms that restricting the candidate pool to a single token is the same as always taking the argmax. And the comparison between top_p=0.1 and top_p=0.95 shows nucleus sampling adapting to different thresholds — the conservative setting staying close to the model’s top choices, the permissive setting occasionally venturing further.

None of the math underlying these strategies has been implemented here — model.generate handles all of it. Section 1.4 introduces the shared SamplingParams dataclass that every subsequent chapter uses to carry these parameters through the system. Chapter 3 implements the sampling logic from scratch, making the formulas above executable and inspectable.

Contracts

Before wiring an engine into a server, or a server into a benchmark, it helps to agree on what data flows between them. In software, that agreement is called a contract — a definition of the types that cross a boundary, making it possible for each side to be written, tested, and replaced independently.

The three types introduced in this section appear in every single chapter of this book. They are defined once in shared/ and never change. What changes from chapter to chapter is the engine that implements them. Chapter 2 replaces the HuggingFace call inside generate() with our own loop. Chapter 4 adds a KV cache. Chapter 9 adds continuous batching. In every case, the surrounding code — the server, the benchmark harness, the metrics display — stays identical because the contracts stay identical.

from shared.metrics import RequestMetrics
from shared.result import GenerateResult
from shared.sampling_params import SamplingParams

Run it:

python 03_contracts.py

SamplingParams — how should tokens be chosen?

@dataclass
class SamplingParams:
    temperature: float = 0.0
    top_k:       int   = -1
    top_p:       float = 1.0
    ignore_eos:  bool  = False
    max_tokens:  int   = 1024

    @property
    def is_greedy(self) -> bool:
        return (self.temperature <= 0.0 or self.top_k == 1) and self.top_p == 1.0

SamplingParams is a dataclass that carries the sampling configuration for a single request. It travels from the caller of generate() down through the engine to the sampler that picks each output token.

The defaults are deliberately conservative. temperature=0.0 means greedy decoding — always take the most probable token, no randomness. top_k=-1 means top-k filtering is disabled. top_p=1.0 means all tokens remain in the candidate pool. Together, these defaults produce deterministic, reproducible output with a single line:

params = SamplingParams()   # greedy

To switch to nucleus sampling with moderate creativity, the caller constructs:

params = SamplingParams(temperature=0.8, top_p=0.95)

The is_greedy property gives every engine a single place to check whether randomness is needed before calling the sampler, avoiding the cost of multinomial sampling when the result is predetermined. If both top_p is 1.0 and temperature is zero (or top_k is 1), the distribution has already collapsed to a point mass and sampling adds nothing.

The reason to put these parameters in a dataclass rather than passing them as individual keyword arguments is consistency. As the number of sampling strategies grows — minimum-p sampling, repetition penalty, beam search — the function signature stays the same. The caller constructs a SamplingParams object; the engine receives it; the sampler reads from it. Adding a new field to the dataclass does not break any existing call site.

RequestMetrics — how do we measure performance?

@dataclass
class RequestMetrics:
    throughput:    float
    input_tokens:  int
    output_tokens: int
    ttft_ms:       float = 0.0   # 0.0 = not measured
    tpot_ms:       float = 0.0   # 0.0 = not measured

RequestMetrics records what happened during a single generate() call. It lives inside every GenerateResult and is serialised into the server’s HTTP response under a metrics key.

The three required fields are always present. throughput is output tokens per second — the most direct measure of how fast the engine is running. input_tokens and output_tokens are the prompt and completion lengths, which together determine how much work was done and provide the numbers that go into the OpenAI-compatible usage field of the response.

The two optional fields — ttft_ms and tpot_ms — default to 0.0 and are only included in the serialised output when non-zero. TTFT (time to first token) measures how long the prefill phase took: the delay between receiving the request and producing the first output token, which is what the user experiences as initial latency. TPOT (time per output token) measures the average duration of each decode step, which determines how fast text streams to the user once generation has begun.

In Chapter 1, both fields stay at zero. HuggingFace runs prefill and decode as a single opaque call, so there is no way to separate them from the outside. Chapter 2 builds the loop step by step, at which point timing each phase becomes straightforward and both fields get populated. The 0.0 default is not a placeholder — it is a signal that the measurement was not available, and as_dict() is designed to omit it cleanly:

def as_dict(self) -> dict:
    d = {
        "throughput":    round(self.throughput, 1),
        "input_tokens":  self.input_tokens,
        "output_tokens": self.output_tokens,
    }
    if self.ttft_ms:
        d["ttft_ms"] = round(self.ttft_ms, 2)
    if self.tpot_ms:
        d["tpot_ms"] = round(self.tpot_ms, 2)
    return d

A Ch01 engine produces {"throughput": 12.5, "input_tokens": 20, "output_tokens": 40}. A Ch02 engine produces the same dict extended with "ttft_ms" and "tpot_ms". The server code that formats the response never needs to know which chapter it is running; it calls result.metrics.as_dict() and gets back whatever the engine was able to measure.

GenerateResult — the envelope

@dataclass
class GenerateResult:
    text:       str
    output_ids: list[int]
    metrics:    RequestMetrics

GenerateResult is the return type of every generate() call in the book. It is deliberately minimal: three fields, no methods, no logic.

text is the decoded output string — what gets shown to the user or returned in the HTTP response body. output_ids is the raw list of generated token IDs, not including the prompt tokens. Keeping the raw IDs alongside the text matters because some later chapters — speculative decoding in Chapter 13, prefix caching in Chapter 11 — need to reason about the token sequence directly, not just the string it decodes to. metrics holds a RequestMetrics instance, or any subclass of it.

That last point — “or any subclass” — is what makes the design extensible. Chapter 13 introduces SpecDecodeMetrics, which inherits from RequestMetrics and adds two additional fields:

@dataclass
class SpecDecodeMetrics(RequestMetrics):
    tokens_proposed: int = 0
    tokens_accepted: int = 0

    @property
    def acceptance_rate(self) -> float:
        return self.tokens_accepted / max(self.tokens_proposed, 1)

A Chapter 13 engine returns a GenerateResult with a SpecDecodeMetrics inside it. The server still calls result.metrics.as_dict(), and because SpecDecodeMetrics overrides as_dict() to include the extra fields, the acceptance rate appears in the response automatically. No code outside the engine changes.

The data flow

Putting the three types together:

prompt (str)
    │
    ▼
engine.generate(prompt, max_new_tokens, SamplingParams)
    │
    ▼
GenerateResult
    ├── text        → returned to caller / shown in HTTP response
    ├── output_ids  → token IDs, kept for chapters that need them
    └── metrics     → RequestMetrics.as_dict() → "metrics" in JSON

Every chapter from here to the end of the book follows this shape. The prompt goes in, a GenerateResult comes out. What varies is everything that happens in between — and the contracts above are what make it possible to swap that middle out without touching anything else.

The Engine

Every chapter of this book centres on one class: the engine. Its job is to accept a prompt string and return a GenerateResult. What changes from chapter to chapter is how it produces that result — which data structures it uses internally, how it schedules work, whether it caches past computation. What never changes is the interface the engine exposes.

This section introduces BaselineEngine, the Chapter 1 implementation, and the abstract base class it fulfils.

Run it:

python 04_engine.py

The interface — `BaseOfflineEngine`

class BaseOfflineEngine(ABC):

    @abstractmethod
    def generate(
        self,
        prompt:         str,
        max_new_tokens: int,
        params:         SamplingParams | None = None,
    ) -> GenerateResult: ...

    @abstractmethod
    def generate_batch(
        self,
        prompts:        list[str],
        max_new_tokens: int,
        params:         SamplingParams | None = None,
    ) -> list[GenerateResult]: ...

BaseOfflineEngine is an abstract base class — a class that can never be instantiated on its own but enforces that any subclass provides specific methods. In Python, this guarantee is expressed with the @abstractmethod decorator: attempting to instantiate a class that inherits from BaseOfflineEngine without implementing both generate and generate_batch raises a TypeError at the moment of construction, before any inference code runs. The error is caught early, close to where the contract was broken, rather than failing silently during a benchmark run.

The design choice of an abstract base class over a plain protocol or duck-typed interface matters for this book in particular. Each chapter defines its own engine class. The FastAPI server in 05_server.py is built by a shared factory function (make_app) that accepts any BaseOfflineEngine and returns a working application. The benchmark scripts in 06_benchmark_offline.py call engine.generate() directly. None of that code needs to know which chapter’s engine it is talking to — the abstract base class is the contract that makes that substitution safe. This is the same pattern production inference systems use: vLLM’s LLMEngine, SGLang’s LLM class, and TGI’s TextGenerationInterface are all variations of the same idea.

Loading the model

def __init__(self, model_name: str = MODEL_NAME):
    self.tokenizer = AutoTokenizer.from_pretrained(model_name)
    self.model     = (
        AutoModelForCausalLM
        .from_pretrained(model_name, torch_dtype=torch.float32)
        .to(DEVICE)
        .eval()
    )

__init__ loads the model exactly once, when the engine is constructed. This matters because loading a multi-gigabyte checkpoint from disk and moving its weights to the GPU can take several seconds; doing it once and keeping the model resident in memory for the lifetime of the engine is the pattern every production system follows.

torch_dtype=torch.float32 asks HuggingFace to load weights in 32-bit floating point. For a small teaching model this is fine. In production, you would typically pass torch.float16 or torch.bfloat16 to halve the memory footprint. Chapter 12 covers quantisation, which goes further still.

.to(DEVICE).eval() moves all weight tensors from CPU RAM to the target device and switches the model from training mode to inference mode. In training mode, certain layers — BatchNorm, Dropout — behave differently: dropout randomly zeros activations during the forward pass, and batch norm tracks running statistics. Calling .eval() disables both. For a transformer used purely for inference, .eval() is always correct. Omitting it is a common source of non-deterministic outputs that are difficult to debug.

The generate method

def generate(
    self,
    prompt:         str,
    max_new_tokens: int,
    params:         SamplingParams | None = None,
) -> GenerateResult:
    params    = params or SamplingParams()
    input_ids = self.tokenizer.encode(prompt, return_tensors="pt").to(DEVICE)

    t0 = time.perf_counter()
    with torch.no_grad():
        # ← Chapter 2 replaces this with our own token-by-token loop
        output_ids = self.model.generate(
            input_ids,
            max_new_tokens=max_new_tokens,
            do_sample=not params.is_greedy,
            temperature=params.temperature if not params.is_greedy else 1.0,
            top_p=params.top_p,
        )
    elapsed_s = time.perf_counter() - t0

    new_ids = output_ids[0][input_ids.shape[1]:]
    text    = self.tokenizer.decode(new_ids, skip_special_tokens=True)
    n_tok   = len(new_ids)

The generate method is five steps in one place: parameter defaulting, tokenisation, inference, output slicing, and metrics construction.

If the caller passes None for params, the line params = params or SamplingParams() substitutes the default SamplingParams object, which uses greedy decoding. This means callers that do not care about sampling get deterministic output without any extra code.

The prompt string is then encoded to a tensor of token IDs and sent to the same device as the model weights. return_tensors="pt" produces a PyTorch tensor with a batch dimension of one, so input_ids has shape [1, prompt_length].

The torch.no_grad() context manager tells PyTorch to skip building the computation graph — the record of operations needed to compute gradients during backpropagation. During inference no gradient is ever needed, so building the graph would waste both memory and time. On a long prompt the difference is measurable. Forgetting torch.no_grad() is one of the most common sources of out-of-memory errors in inference code.

Inside the context, self.model.generate() runs HuggingFace’s autoregressive loop. The loop feeds the model the growing sequence of tokens, one step at a time, extracts the logit vector for the last position, applies the requested sampling strategy — greedy argmax, temperature scaling, nucleus filtering — and appends the chosen token back onto the sequence. It repeats until max_new_tokens tokens have been generated or an end-of-sequence token appears. The entire loop is hidden inside a single function call. This is what makes Chapter 1 a black box: the loop exists, but we cannot see inside it. Chapter 2 rewrites it from scratch so that each step is visible.

The do_sample flag translates SamplingParams.is_greedy into HuggingFace’s vocabulary: False means greedy, True enables multinomial sampling. When do_sample=True, the temperature and top_p parameters are forwarded from SamplingParams. When do_sample=False, temperature is forced to 1.0 because it has no effect in greedy mode and HuggingFace raises a warning if it is set to anything else.

After the call, output_ids[0] strips the batch dimension, giving a one-dimensional tensor of the complete sequence — prompt tokens followed by generated tokens. Slicing from input_ids.shape[1] forward isolates only the generated part. The prompt tokens were already shown to the user; they should not appear in the output string.

Returning a GenerateResult

    return GenerateResult(
        text       = text,
        output_ids = new_ids.tolist(),
        metrics    = RequestMetrics(
            throughput    = n_tok / max(elapsed_s, 1e-9),
            input_tokens  = input_ids.shape[1],
            output_tokens = n_tok,
        ),
    )

The GenerateResult envelope is constructed with the three fields introduced in 03_contracts.py. output_ids calls .tolist() to convert the PyTorch tensor into a plain Python list, because GenerateResult is a plain dataclass that knows nothing about PyTorch and should not.

RequestMetrics is constructed with only the three required fields. throughput is n_tok / elapsed_s — output tokens divided by wall-clock seconds. The max(elapsed_s, 1e-9) guard prevents a division by zero on extremely fast machines or when the token count is very small. ttft_ms and tpot_ms are intentionally absent: HuggingFace runs prefill and decode as a single fused call, so the moment the first token was produced is invisible from outside model.generate(). Chapter 2 builds the loop step by step, which makes both timings accessible.

generate_batch

def generate_batch(
    self,
    prompts:        list[str],
    max_new_tokens: int,
    params:         SamplingParams | None = None,
) -> list[GenerateResult]:
    return [self.generate(p, max_new_tokens, params) for p in prompts]

generate_batch is a sequential loop over generate. Each prompt is processed one at a time, with no parallelism. This implementation satisfies the interface contract — the method exists and returns the right type — but it does not exploit any of the efficiency opportunities that batching makes possible. On a GPU, the compute units sit mostly idle while the model processes a single short sequence; running multiple sequences in parallel would keep them busy and amortise the fixed cost of a forward pass across many tokens simultaneously. Chapter 5 replaces this loop with true static batching. Chapter 9 extends that to continuous batching, which handles requests of different lengths arriving at different times.

Running the engine

When run as a script, 04_engine.py constructs the engine, runs one greedy generation call, and prints the result:

[engine] loading arnir0/Tiny-LLM on cpu ...
[engine] ready

[Prompt]
  According to all known laws of aviation, ...

[Output]
  The bee flies because it does not know it cannot.

[Metrics]
  throughput      12.3
  input_tokens    18
  output_tokens   40

The engine is now a fully functional unit. The next step is to wrap it in an HTTP server so that any client — a benchmark script, a web UI, another service — can call it over the network using a standard API.

The Server

The engine from the previous section is a Python object. Anything that wants to use it must be Python code running in the same process. Wrapping the engine in an HTTP server removes that constraint: any client anywhere on the network — a curl command, a Python script, a browser, another microservice — can call the engine by making an ordinary HTTP POST request. This is how every production LLM deployment works. OpenAI’s API, Anthropic’s API, and self-hosted systems like vLLM and SGLang all expose the same basic shape: a POST endpoint that accepts a JSON body describing the request and returns a JSON body containing the generated text.

This section builds that server in roughly twenty lines of application code, reusing a factory from shared/server.py that every chapter server will use unchanged.

Run it:

uvicorn 05_server:app --port 8001

Test it (in a second terminal):

curl http://localhost:8001/v1/chat/completions \
     -H "Content-Type: application/json" \
     -d '{"messages":[{"role":"user","content":"Hello"}],"max_tokens":50}'

FastAPI and uvicorn

from fastapi import FastAPI
from pydantic import BaseModel

FastAPI is a Python web framework built on two foundations. The networking layer is handled by uvicorn, an ASGI server. ASGI — Asynchronous Server Gateway Interface — is the protocol that connects a Python application to the outside world: uvicorn opens a TCP socket, waits for HTTP connections, reads each request from the wire, and hands control to FastAPI as a structured Python object. FastAPI then inspects the request, routes it to the right handler function, validates the request body with Pydantic, calls the handler, serialises the return value back to JSON, and instructs uvicorn to write the response bytes to the socket.

For our purposes the important practical consequence is that uvicorn and FastAPI are separate processes of concern. Running uvicorn 05_server:app --port 8001 tells uvicorn to import the app object from 05_server.py and start listening on port 8001. The application code never calls accept() or manages sockets; it only defines handlers.

Request model

class Message(BaseModel):
    role:    str
    content: str

class ChatRequest(BaseModel):
    model:       str   = "default"
    messages:    list[Message]
    max_tokens:  int   = 100
    temperature: float = 0.0
    top_p:       float = 1.0

ChatRequest is a Pydantic BaseModel. Pydantic models declare their fields as class-level annotations with optional defaults. When FastAPI receives an incoming POST body, it deserialises the raw JSON bytes and attempts to construct a ChatRequest from them. If any required field is missing, if a field’s value cannot be coerced to the declared type, or if a list element does not match the inner type, Pydantic raises a ValidationError and FastAPI returns a 422 Unprocessable Entity response to the client automatically. The handler function is never called. This means validation is essentially free from the handler’s perspective: by the time the function body runs, every field is guaranteed to be present and correctly typed.

The field names follow the OpenAI chat completions API convention. messages is a list of Message objects, each with a role ("user", "assistant", "system") and a content string. The convention exists because the ecosystem around LLM APIs — client libraries, proxies, benchmarking tools, evaluation harnesses — speaks this language. Matching it means our server is compatible with any tool that targets OpenAI’s endpoint, which in Chapter 6 lets the benchmark script treat our server identically to a production one.

The make_app factory

def make_app(engine: BaseOfflineEngine) -> FastAPI:
    app  = FastAPI()
    _uid = [0]

    @app.post("/v1/chat/completions")
    def chat(req: ChatRequest) -> dict:
        prompt = req.messages[-1].content
        params = SamplingParams(
            temperature=req.temperature,
            top_p=req.top_p,
            max_tokens=req.max_tokens,
        )
        result: GenerateResult = engine.generate(prompt, req.max_tokens, params)
        _uid[0] += 1
        return _format_response(result, req, uid=_uid[0])

    @app.get("/health")
    def health() -> dict:
        return {"status": "ok"}

    return app

make_app is a factory function: it takes an engine, builds a FastAPI application wired to that engine, and returns the application object. The engine is captured in the closure over the chat handler. This is the same pattern used by shared/server.py for every chapter in the book — the factory is written once and never changes; only the engine passed to it changes.

The chat handler extracts the user’s message from the last element of the messages list (req.messages[-1].content). Multi-turn conversation history — earlier messages in the list — is not passed to the engine in Chapter 1. Building a proper multi-turn engine requires concatenating or summarising context, which is a topic for later chapters. For now, the server behaves as a stateless single-turn responder.

SamplingParams is constructed from the request fields, translating the HTTP API’s vocabulary into the internal contract. The engine.generate() call is synchronous and blocking: the handler thread calls into the engine, waits until the model produces all max_tokens output tokens, and only then returns the result. This works correctly for a single concurrent request. It becomes a bottleneck under load — a second request arriving while the first is generating must wait in a queue — and Chapter 9 addresses this by separating the inference loop from the HTTP handler using asynchronous scheduling.

_uid is a module-level list containing a single integer used as a request counter. A list rather than a bare integer is used because Python closures over a name, not a value: reassigning an integer with _uid = _uid + 1 inside the closure would create a new local binding, leaving the outer name unchanged. Mutating the first element of a list (_uid[0] += 1) modifies the existing object, which the closure sees correctly. In production code a threading lock or an atomic counter would be appropriate; here a simple integer is sufficient because the server handles one request at a time.

The /health endpoint returns a trivial JSON object. 06_benchmark_online.py polls this endpoint after starting the server process to confirm it is ready before sending benchmark requests. Any HTTP 200 response from /health signals readiness.

The response format

def _format_response(result: GenerateResult, req: ChatRequest, uid: int) -> dict:
    m = result.metrics
    return {
        "id":     f"cmpl-{uid}",
        "object": "chat.completion",
        "choices": [{
            "index":         0,
            "message":       {"role": "assistant", "content": result.text},
            "finish_reason": "stop",
        }],
        "usage": {
            "prompt_tokens":     m.input_tokens,
            "completion_tokens": m.output_tokens,
            "total_tokens":      m.input_tokens + m.output_tokens,
        },
        "metrics": m.as_dict(),
    }

The response follows the OpenAI chat completions object shape. id is a unique identifier per completion — benchmark scripts and clients use it to match responses to requests. object is a fixed string identifying the response type. choices is a list of completions; we always return exactly one, at index: 0. Each choice wraps the generated text in a message object with role: "assistant" and sets finish_reason: "stop" to indicate the model reached max_new_tokens without an error or a content filter triggering.

usage reports token counts using OpenAI’s field names: prompt_tokens, completion_tokens, and their sum as total_tokens. These values come directly from RequestMetrics.input_tokens and RequestMetrics.output_tokens, which the engine measured during generation. Tools that consume the response — cost calculators, rate limiters, evaluation harnesses — depend on this field being accurate.

metrics is a non-standard extension, absent from the OpenAI specification. It surfaces the full RequestMetrics.as_dict() output — throughput and token counts for Chapter 1 — so that clients and benchmark scripts can read performance data from the HTTP response without needing a separate metrics endpoint. Future chapters add ttft_ms and tpot_ms to RequestMetrics; when they do, those fields appear in metrics automatically because as_dict() includes them conditionally.

Wiring the engine to the server

engine = BaselineEngine()
app    = make_app(engine)

These two lines are the entire chapter-specific portion of 05_server.py. The engine is constructed at module import time, meaning the model is loaded when uvicorn imports the module, before the first request arrives. make_app wraps it. The resulting app object is what uvicorn holds a reference to and calls for every incoming HTTP request.

BaselineEngine appears in full in 05_server.py rather than being imported from 04_engine.py because this file is designed to be read as a standalone introduction to the engine-plus-server pattern: everything a reader needs to understand how the server works is visible in one file. The thin engine.py wrapper in the same directory uses importlib to import BaselineEngine from 04_engine.py for the benchmark scripts, which need a clean import path. Both point to the same implementation; neither is a copy of different code.

A curl response from this server looks like:

{
  "id": "cmpl-1",
  "object": "chat.completion",
  "choices": [{
    "index": 0,
    "message": {"role": "assistant", "content": "The bee flies anyway."},
    "finish_reason": "stop"
  }],
  "usage": {
    "prompt_tokens": 18,
    "completion_tokens": 40,
    "total_tokens": 58
  },
  "metrics": {
    "throughput": 12.3,
    "input_tokens": 18,
    "output_tokens": 40
  }
}

The next step is to measure how fast this server actually runs — both directly, by calling the engine in-process, and over the network, by sending HTTP requests to the running server.

The Benchmark

Knowing that an engine produces correct output is not enough. Every chapter of this book introduces a new engine implementation, and the only way to know whether a change is actually faster is to measure it systematically, with the same methodology, against the same numbers. This section builds the measurement tooling and produces the baseline numbers that every subsequent chapter will try to beat.

There are two benchmark scripts. 06_benchmark_offline.py calls the engine directly in-process. 06_benchmark_online.py sends HTTP requests to a running server. Both produce the same table of numbers; the difference between them is the overhead of the network stack.

Run offline:

python 06_benchmark_offline.py

Run online (server is launched and stopped automatically):

python 06_benchmark_online.py

What to measure

A benchmark number is only useful if you understand what it is measuring and what it is not. Three quantities matter for a single-user, single-request system like Chapter 1:

Throughput (tokens per second) is the rate at which the model produces output tokens. For a generative model, this is output_tokens / elapsed_seconds. It is the primary number that every chapter optimisation targets. Higher is better.

Requests per second is 1 / avg_wall_seconds. For a single-threaded server, this is simply the reciprocal of average latency. It becomes independently interesting in Chapter 9, where continuous batching allows multiple requests to overlap, decoupling throughput from single-request latency.

Wall time per request is the total time from calling engine.generate() (or sending the HTTP request) to receiving the complete result. It is the number the end user experiences as latency. Because Chapter 1 generates all tokens before returning, wall time grows linearly with max_new_tokens.

Time-to-first-token (TTFT) and time-per-output-token (TPOT) are absent from Chapter 1. HuggingFace’s model.generate() runs prefill and decode as a single fused call and returns only after all tokens are produced; the moment the first token was generated is not observable from outside the call. Chapter 2 builds the loop step by step and populates both. The benchmark infrastructure tracks them in BenchStats with 0.0 defaults and omits them from the printed table until they are non-zero.

Warmup runs

for _ in range(n_warmup):
    engine.generate(prompt, max_tokens, params)

The first few calls to engine.generate() are always slower than subsequent ones. Several independent effects cause this. PyTorch compiles CUDA kernels lazily: the first time a particular operation is encountered with a particular tensor shape, PyTorch compiles and caches a CUDA kernel for it. That compilation can take tens or hundreds of milliseconds. On CPU, the OS page-fault mechanism brings model weight pages into RAM only as they are accessed; the first forward pass touches every weight page for the first time. Python’s import machinery and interpreter warm up. On subsequent calls, compiled kernels are cached, weights are resident in memory, and the interpreter is in a steady state.

Warmup runs are discarded. They are not included in any computed average. The number of warmup runs is configured by WARMUP_RUNS in shared/config.py (default: 3), which is enough to get past both kernel compilation and page faults on the models used in this book. Production benchmarks like those published by NVIDIA NIM use a dedicated warmup phase for the same reason, though at larger scale they may require dozens of iterations.

BenchStats

@dataclass
class BenchStats:
    req_per_s:     float
    tok_per_s:     float
    wall_ms:       float
    output_tokens: float
    n_runs:        int
    ttft_ms:       float       = 0.0
    tpot_ms:       float       = 0.0
    wall_times:    list[float] = field(default_factory=list)

BenchStats aggregates all measurements from a benchmark run into one object. req_per_s, tok_per_s, and wall_ms are averages over n_runs measured requests. wall_times keeps the raw per-request timings; later chapters that analyse latency distributions — p50, p95, p99 — can compute percentiles from this list without re-running the benchmark. ttft_ms and tpot_ms default to 0.0 and remain at zero when not populated, following the same optional-field convention used in RequestMetrics.

Offline benchmark

for _ in range(n_runs):
    t0      = time.perf_counter()
    result  = engine.generate(prompt, max_tokens, params)
    wall_ms = (time.perf_counter() - t0) * 1000

    wall_times.append(wall_ms)
    tok_counts.append(result.metrics.output_tokens)
    if result.metrics.ttft_ms:
        ttfts.append(result.metrics.ttft_ms)
    if result.metrics.tpot_ms:
        tpots.append(result.metrics.tpot_ms)

offline_benchmark calls engine.generate() directly. time.perf_counter() measures wall-clock elapsed time with nanosecond resolution — it is the correct timer for short durations. time.time() uses the system clock and can jump backward or forward if the OS adjusts it; time.perf_counter() is monotonic and immune to that. The multiplication by 1000 converts seconds to milliseconds, which are the natural unit for per-request latency on the models used in this book.

The output token count comes from result.metrics.output_tokens rather than re-counting result.output_ids. The engine already measured this; there is no reason to recompute it. ttft_ms and tpot_ms are collected only when non-zero, preserving the same 0.0-signals-unmeasured convention used throughout the contracts.

Token counts matter because max_new_tokens is a ceiling, not a guarantee. HuggingFace’s model.generate() stops early if it produces an end-of-sequence token before reaching max_new_tokens. If the average output is 35 tokens when you asked for 40, computing tokens-per-second using 40 would give a number that is 14% optimistic. Reading output_tokens from the result gives the correct denominator.

Aggregation divides total output tokens by total wall time:

avg_wall   = _avg(wall_times)       # ms
avg_tokens = _avg(tok_counts)
req_per_s  = 1000 / avg_wall        # 1 / avg_seconds
tok_per_s  = avg_tokens / (avg_wall / 1000)

Online benchmark

t0   = time.perf_counter()
r    = requests.post(endpoint, json=payload)
r.raise_for_status()
wall_ms = (time.perf_counter() - t0) * 1000

data = r.json()
wall_times.append(wall_ms)
tok_counts.append(data["usage"]["completion_tokens"])

online_benchmark sends HTTP POST requests to the running server. The wall time is measured around the full requests.post() call — from the moment the client sends the request to the moment the response body is fully received. This includes TCP connection time (negligible for localhost), HTTP framing, JSON serialisation on the server, and JSON deserialisation on the client. The difference between this number and the offline wall time is the overhead the network stack and server layer add. For a localhost server this overhead is typically 0.5–5ms per request.

The token count is read from data["usage"]["completion_tokens"] — the OpenAI-standard field in the response. The non-standard metrics field is also read for ttft_ms and tpot_ms when present, using .get() with a fallback to an empty dict so that the benchmark works correctly against servers that do not include a metrics field.

r.raise_for_status() converts any HTTP 4xx or 5xx response into a Python exception immediately. Without it, a server error would silently produce None JSON and the benchmark would either crash with an unhelpful KeyError or, worse, record a zero-token measurement and silently corrupt the averages.

Server lifecycle

proc = launch_server(
    server_module = "05_server:app",
    port          = SERVER_PORT,
    health_url    = f"{SERVER_URL}/health",
    cwd           = CHAPTER_DIR,
)
try:
    stats = online_benchmark(SERVER_URL, ...)
    print_results("Ch01 baseline — online (HTTP)", stats)
finally:
    stop_server(proc)

launch_server starts uvicorn as a child process using subprocess.Popen. The parent process then polls GET /health in a loop, sleeping one second between attempts, until the server returns HTTP 200 or a 60-second timeout expires. The poll loop exists because uvicorn takes a variable amount of time to start: it must import the server module, which triggers model loading, which on CPU can take several seconds. Polling /health is more reliable than sleeping a fixed duration.

stop_server calls proc.terminate(), which sends SIGTERM to the uvicorn process, followed by proc.wait(), which blocks until the process exits. This ensures the port is fully released before the benchmark script itself exits. Without proc.wait(), the port can remain bound briefly after the script ends, causing the next run to fail with “address already in use”.

The try / finally block around the benchmark is critical. If online_benchmark raises an exception — a server error, a network timeout, a keyboard interrupt — stop_server still runs in the finally branch. Without it, a failed benchmark would leave a uvicorn process running in the background, holding the port and the GPU memory.

The printed table

  ──────────────────────────────────────────────────────
  Ch01 baseline — offline (no HTTP)
  ──────────────────────────────────────────────────────
  Req/s        : 0.82
  Tok/s        : 32.6
  Wall time    : 1224ms / request
  Tokens out   : 39 (avg over 20 runs)
  ──────────────────────────────────────────────────────

print_results omits TTFT and TPOT lines when they are zero, keeping the Chapter 1 output clean. Future chapters print a longer table. The label argument identifies which chapter and mode produced the table, making it straightforward to compare results across chapters or across offline and online runs side by side.

These numbers are the baseline. Chapter 2 prints the same table immediately after building its own token-by-token loop. If the Chapter 2 numbers match Chapter 1’s to within a few percent, the new loop is correct — it is doing the same computation. If they differ significantly, there is a bug. The benchmark is the test.