Inference

chat.py at a glance

Inference is a separate script from training, but it shares the model code. chat.py:

1. Loads the tokenizer from a checkpoint folder.
2. Loads model weights and config from the same folder.
3. Runs an interactive REPL: tokenize the user's input, call the model to produce the next tokens, detokenize, loop.

No inference server, no batching — the goal is to smoke-test a trained model, not to serve production traffic.

Why a KV cache is needed

During training, the model runs the full sequence through once per step — one forward pass produces predictions for every position in parallel. That is fine because the sequence is already known.

During autoregressive generation, the sequence grows one token at a time. A naive implementation would re-run the whole sequence through the model for every new token: to generate T tokens would cost O(T²) forward-pass work.

The fix exploits a property of causal attention: at position t, the attention layer only reads K and V from positions ≤ t. Those K and V values do not change once computed. So we can:

step 1: feed the prompt through the model
        → cache K, V for every prompt position, every layer
        → sample token_1 from the logits at the last position

step 2: feed only token_1 through the model
        → the attention layer reads the cached K, V, and appends
          K[t+1], V[t+1] for this new position
        → sample token_2 from the logits at this position

step 3: feed only token_2, ...

Per-step cost becomes O(T) instead of O(T²). This is the single most important inference optimization in any decoder-only LM.

Memory footprint

The KV cache lives per attention layer, and its size grows linearly with the generated length:

cache_size  =  2 · n_layers · n_kv_heads · d_head · T

The 2 is for K and V. Note n_kv_heads, not n_heads — this is precisely why GQA matters for deployment: cutting n_kv_heads 5× cuts per-token cache memory 5×.

Reference writeup:https://huggingface.co/docs/transformers/main/kv_cache.

One code path for training and inference

The same attention module handles both regimes. At training time the module does not receive a cache argument and recomputes everything from the input sequence. At inference time it accepts a cache object, appends to it, and reads from it:

def attention(self, x, kv_cache=None):
    q = self.q_proj(x)
    k = self.k_proj(x)
    v = self.v_proj(x)
    q, k = apply_rope(q, k, positions)
    if kv_cache is not None:
        k = torch.cat([kv_cache.k, k], dim=2)
        v = torch.cat([kv_cache.v, v], dim=2)
        kv_cache.update(k, v)
    return scaled_dot_product_attention(q, k, v, is_causal=(kv_cache is None))

There is no separate inference implementation — which prevents the two from drifting. Any architecture change (a new norm, a new positional encoding, a new attention variant) automatically applies to chat.py without a second port.

Running it

python chat.py

Expects a checkpoint folder written by training (checkpoints/latest/ by default) with the tokenizer saved alongside the weights.

End of pipeline

This is where the v0 pipeline ends: text in, token IDs out, model forward, token IDs back, text out. Every component along the way has its own page in this section; the landing page at ./ has the full stage list.