{"slug": "using-a-transformer-model-from-training-to-inference", "title": "Using a Transformer Model: From Training to Inference", "summary": "A new chapter on transformer models explains that training and inference use the same PyTorch code but differ fundamentally: training processes fixed-length batches with backward passes, while inference generates tokens one at a time via an autoregressive loop. The chapter introduces the greedy decoding loop, the prefill and decode phases, and the necessity of key-value caching to avoid recomputing hidden states, noting that without caching, generating 100 tokens from a 1,000-token prompt processes O(N^2) tokens.", "body_md": "If you have implemented a transformer model in PyTorch, you can use the same code for both training and inference, but in very different ways. During training, you usually process a batch of fixed-length token sequences and update the model weights. During inference, the weights are fixed and the model generates new tokens one at a time.\n\nThis difference changes almost everything about performance. Training is dominated by large matrix multiplications and the backward pass. Inference is dominated by repeated forward passes, memory movement, and the need to keep previous attention keys and values available for the next token.\n\nIn this chapter, you will learn about:\n\n- The autoregressive generation loop\n- The difference between prefill and decode\n- Why key-value caching is necessary\n- How to implement a simple KV cache\n- How to reason about the memory used by the cache\n\nLet’s get started.\n\n## Overview\n\nThis chapter is divided into four parts; they are:\n\n- Autoregressive Generation\n- Prefill and Decode\n- A Simple KV Cache\n- Memory Usage of the KV Cache\n\n## Autoregressive Generation\n\nA decoder-only transformer model predicts the next token from the tokens that came before it. The strict requirement of using only the previous tokens is enforced by the causal attention mechanism. If the input tokens are:\n\n```\nThe cat sat on the\n\n1\n\nThe cat sat on the\n```\n\nthe model returns a probability distribution over the vocabulary for the next token. A likely next token may be “mat”, but the model does not return a word directly. It returns logits, which are unnormalized scores for every token in the vocabulary.\n\nThe generation loop is therefore simple:\n\n- Tokenize the prompt.\n- Run the model to obtain logits for the next token.\n- Choose a token from the logits.\n- Append that token to the input.\n- Repeat until a stopping rule is reached.\n\nThis is called autoregressive generation because each new token depends on the previous generated tokens. The model cannot generate the tenth output token before it knows the first nine output tokens.\n\nA very small greedy decoding loop can be written as follows:\n\n``` python\nimport torch\n\n@torch.no_grad()\ndef greedy_decode(model, input_ids, max_new_tokens):\n    output_ids = input_ids.clone()\n\n    for _ in range(max_new_tokens):\n        logits = model(output_ids)\n        next_token_logits = logits[:, -1, :]\n        next_token = next_token_logits.argmax(dim=-1, keepdim=True)\n        output_ids = torch.cat([output_ids, next_token], dim=1)\n\n    return output_ids\n\n12345678910111213\n\nimport torch @torch.no_grad()def greedy_decode(model, input_ids, max_new_tokens):    output_ids = input_ids.clone()     for _ in range(max_new_tokens):        logits = model(output_ids)        next_token_logits = logits[:, -1, :]        next_token = next_token_logits.argmax(dim=-1, keepdim=True)        output_ids = torch.cat([output_ids, next_token], dim=1)     return output_ids\n```\n\nIn the code above, `model`\n\nis a PyTorch model, `max_new_tokens`\n\nis a positive integer, and all other variables are PyTorch tensors. The for-loop iterates `max_new_tokens`\n\ntimes, and at each iteration, it feeds the entire sequence back into the model to get the logits for the next token. The `argmax()`\n\nfunction selects the highest-scoring token. The `cat()`\n\nfunction is used to concatenate the new token to the output sequence, which will be used in the next iteration until the stopping rule is reached.\n\nThis code is easy to understand, but it is inefficient. At every iteration, it feeds the entire sequence back into the model. If the prompt has 1,000 tokens and you generate 100 new tokens, the model repeatedly recomputes the hidden states for the same prompt tokens. The model processes $O(N^2)$ tokens in this function, for a prompt of length $N$.\n\nThe actual time complexity of the code is even worse. Without caching, every forward pass recomputes attention for all tokens in the growing sequence. If the sequence length is $N$, self-attention has $O(N^2)$ score computation. For generation, this means you repeat a large amount of work. (Precisely if the output sequence length is $N=P+G$ with prompt length $P$ and number of generated tokens $G$, the computation complexity should be $O(P^2G + PG^2 + G^3)$ naively. With cache, we can reduce it to $O(P^2 + PG)$.)\n\nInference systems mitigate this by splitting generation into two phases: prefill and decode.\n\n## Prefill and Decode\n\nGeneration usually begins with a prompt. The prompt is known before generation starts. The model can process all prompt tokens in one forward pass. This is called the **prefill** phase.\n\nDuring prefill, the model computes hidden states for all prompt tokens and produces logits for the next token. It also computes keys and values for all attention layers. These keys and values can be saved because they will be needed by every future token.\n\nAfter the first new token is selected, generation enters the **decode** phase. In decode, the model receives only the newest token. It computes the query, key, and value for that token, appends the new key and value to the cache, and attends the new query over all cached keys and values.\n\nThis changes the cost of one decode step. Instead of recomputing attention for the whole sequence, the model computes attention for only one new query against all previous keys. The per-token attention cost changes from roughly $O(N^2)$ to $O(N)$ for a sequence of length $N$. The prefill step is still $O(N^2)$, but it is performed only once for the prompt.\n\nThis distinction is important enough that serving systems usually measure prefill and decode separately:\n\n**Prefill** affects time to first token. A slow prefill increases time to the first token.**Decode** affects the speed of streaming output tokens. A slow decode reduces the rate at which output tokens are streamed.\n\nA short prompt with a long answer stresses decode. A long prompt with a short answer stresses prefill. A chat application with a long conversation history stresses both.\n\nThe matrix below illustrates the attention-score matrix $QK^\\top$. Assume the prompt has five tokens. During prefill, the model computes the $5 \\times 5$ block in blue. During decode, one new token is added at a time. Each decode step adds one new row to the matrix, shown in a different shade of red. The elements in black are ignored from calculation due to the causal mask.\n\n## A Simple KV Cache\n\nThe KV cache is where the model stores the attention keys and values produced by previous tokens. To see how it works, you do not need a large model. The following code builds a small transformer-like model with a cache.\n\nThis model is not intended to produce useful text. Its purpose is to show how the cache is created during prefill and extended during decode.\n\n``` python\nimport math\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\nclass SelfAttention(nn.Module):\n    def __init__(self, hidden_size, num_heads):\n        super().__init__()\n        assert hidden_size % num_heads == 0\n        self.num_heads = num_heads\n        self.head_dim = hidden_size // num_heads\n        self.qkv = nn.Linear(hidden_size, 3 * hidden_size)\n        self.out = nn.Linear(hidden_size, hidden_size)\n\n    def forward(self, x, past_kv=None):\n        # Note: Positional encoding and padding masks are not implemented here\n        batch_size, seq_len, hidden_size = x.shape\n\n        qkv = self.qkv(x)\n        qkv = qkv.view(batch_size, seq_len, 3, self.num_heads, self.head_dim)\n        qkv = qkv.permute(2, 0, 3, 1, 4)\n        q, k, v = qkv[0], qkv[1], qkv[2]\n\n        if past_kv is not None:\n            past_k, past_v = past_kv\n            k = torch.cat([past_k, k], dim=2)\n            v = torch.cat([past_v, v], dim=2)\n\n        total_len = k.size(2)\n        past_len = total_len - seq_len\n\n        scores = q @ k.transpose(-2, -1)\n        scores = scores / math.sqrt(self.head_dim)\n\n        # A token may attend to all cached tokens and earlier tokens\n        # in the current chunk, but not future tokens.\n        causal_mask = torch.ones(seq_len, total_len, device=x.device, dtype=torch.bool)\n        causal_mask = torch.tril(causal_mask, diagonal=past_len)\n        scores = scores.masked_fill(~causal_mask, float(\"-inf\"))\n\n        attn = F.softmax(scores, dim=-1)\n        y = attn @ v\n        y = y.transpose(1, 2).contiguous().view(batch_size, seq_len, hidden_size)\n\n        return self.out(y), (k, v)\n\nclass Block(nn.Module):\n    def __init__(self, hidden_size, num_heads):\n        super().__init__()\n        self.attn_norm = nn.LayerNorm(hidden_size)\n        self.attn = SelfAttention(hidden_size, num_heads)\n        self.ffn_norm = nn.LayerNorm(hidden_size)\n        self.ffn = nn.Sequential(\n            nn.Linear(hidden_size, 4 * hidden_size),\n            nn.GELU(),\n            nn.Linear(4 * hidden_size, hidden_size),\n        )\n\n    def forward(self, x, past_kv=None):\n        attn_out, new_kv = self.attn(self.attn_norm(x), past_kv=past_kv)\n        x = x + attn_out\n        x = x + self.ffn(self.ffn_norm(x))\n        return x, new_kv\n\nclass TinyCausalLM(nn.Module):\n    def __init__(self, vocab_size=128, hidden_size=64, num_heads=4, num_layers=2):\n        super().__init__()\n        self.token_emb = nn.Embedding(vocab_size, hidden_size)\n        self.blocks = nn.ModuleList([\n            Block(hidden_size, num_heads) for _ in range(num_layers)\n        ])\n        self.norm = nn.LayerNorm(hidden_size)\n        self.lm_head = nn.Linear(hidden_size, vocab_size, bias=False)\n\n    def forward(self, input_ids, past_kv=None):\n        x = self.token_emb(input_ids)\n        new_cache = []\n\n        if past_kv is None:\n            past_kv = [None] * len(self.blocks)\n\n        for block, layer_past in zip(self.blocks, past_kv):\n            x, layer_cache = block(x, past_kv=layer_past)\n            new_cache.append(layer_cache)\n\n        logits = self.lm_head(self.norm(x))\n        return logits, new_cache\n\n123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990\n\nimport mathimport torchimport torch.nn as nnimport torch.nn.functional as F  class SelfAttention(nn.Module):    def __init__(self, hidden_size, num_heads):        super().__init__()        assert hidden_size % num_heads == 0        self.num_heads = num_heads        self.head_dim = hidden_size // num_heads        self.qkv = nn.Linear(hidden_size, 3 * hidden_size)        self.out = nn.Linear(hidden_size, hidden_size)     def forward(self, x, past_kv=None):        # Note: Positional encoding and padding masks are not implemented here        batch_size, seq_len, hidden_size = x.shape         qkv = self.qkv(x)        qkv = qkv.view(batch_size, seq_len, 3, self.num_heads, self.head_dim)        qkv = qkv.permute(2, 0, 3, 1, 4)        q, k, v = qkv[0], qkv[1], qkv[2]         if past_kv is not None:            past_k, past_v = past_kv            k = torch.cat([past_k, k], dim=2)            v = torch.cat([past_v, v], dim=2)         total_len = k.size(2)        past_len = total_len - seq_len         scores = q @ k.transpose(-2, -1)        scores = scores / math.sqrt(self.head_dim)         # A token may attend to all cached tokens and earlier tokens        # in the current chunk, but not future tokens.        causal_mask = torch.ones(seq_len, total_len, device=x.device, dtype=torch.bool)        causal_mask = torch.tril(causal_mask, diagonal=past_len)        scores = scores.masked_fill(~causal_mask, float(\"-inf\"))         attn = F.softmax(scores, dim=-1)        y = attn @ v        y = y.transpose(1, 2).contiguous().view(batch_size, seq_len, hidden_size)         return self.out(y), (k, v)  class Block(nn.Module):    def __init__(self, hidden_size, num_heads):        super().__init__()        self.attn_norm = nn.LayerNorm(hidden_size)        self.attn = SelfAttention(hidden_size, num_heads)        self.ffn_norm = nn.LayerNorm(hidden_size)        self.ffn = nn.Sequential(            nn.Linear(hidden_size, 4 * hidden_size),            nn.GELU(),            nn.Linear(4 * hidden_size, hidden_size),        )     def forward(self, x, past_kv=None):        attn_out, new_kv = self.attn(self.attn_norm(x), past_kv=past_kv)        x = x + attn_out        x = x + self.ffn(self.ffn_norm(x))        return x, new_kv  class TinyCausalLM(nn.Module):    def __init__(self, vocab_size=128, hidden_size=64, num_heads=4, num_layers=2):        super().__init__()        self.token_emb = nn.Embedding(vocab_size, hidden_size)        self.blocks = nn.ModuleList([            Block(hidden_size, num_heads) for _ in range(num_layers)        ])        self.norm = nn.LayerNorm(hidden_size)        self.lm_head = nn.Linear(hidden_size, vocab_size, bias=False)     def forward(self, input_ids, past_kv=None):        x = self.token_emb(input_ids)        new_cache = []         if past_kv is None:            past_kv = [None] * len(self.blocks)         for block, layer_past in zip(self.blocks, past_kv):            x, layer_cache = block(x, past_kv=layer_past)            new_cache.append(layer_cache)         logits = self.lm_head(self.norm(x))        return logits, new_cache\n```\n\nThe cache is a list with one element per transformer layer. Each element is a pair `(k, v)`\n\n. The shape of each tensor is:\n\n```\n[batch_size, num_heads, sequence_length, head_dim]\n\n1\n\n[batch_size, num_heads, sequence_length, head_dim]\n```\n\nDuring prefill, `sequence_length`\n\nis the prompt length. During decode, the model receives one token at a time and appends one position to the cache.\n\nYou may notice that only keys and values are stored in the cache but not the query tensor. Note that the `forward()`\n\nmethod is to produce the **next** token’s logits. To do so, you only need the last token in the query tensor (which is from the immediate previous token generated) to multiply with every token in the keys to produce attention scores, which are then used to form a weighted sum of the values. That’s why it is only a KV cache while the attention mechanism is a function of query, key, and value.\n\nHere is a minimal generation loop using the cache:\n\n``` python\ntorch.no_grad()\ndef greedy_decode_with_cache(model, input_ids, max_new_tokens):\n    output_ids = input_ids.clone()\n\n    # Prefill: process the whole prompt once.\n    logits, cache = model(input_ids)\n    next_token = logits[:, -1, :].argmax(dim=-1, keepdim=True)\n    output_ids = torch.cat([output_ids, next_token], dim=1)\n\n    # Decode: process only the most recent token.\n    assert max_new_tokens > 0, \"max_new_tokens must be positive\"\n    for _ in range(max_new_tokens - 1):\n        logits, cache = model(next_token, past_kv=cache)\n        next_token = logits[:, -1, :].argmax(dim=-1, keepdim=True)\n        output_ids = torch.cat([output_ids, next_token], dim=1)\n\n    return output_ids\n\nmodel = TinyCausalLM()\nprompt = torch.tensor([[10, 20, 30, 40]])\ngenerated = greedy_decode_with_cache(model, prompt, max_new_tokens=8)\nprint(generated)\n\n1234567891011121314151617181920212223\n\ntorch.no_grad()def greedy_decode_with_cache(model, input_ids, max_new_tokens):    output_ids = input_ids.clone()     # Prefill: process the whole prompt once.    logits, cache = model(input_ids)    next_token = logits[:, -1, :].argmax(dim=-1, keepdim=True)    output_ids = torch.cat([output_ids, next_token], dim=1)     # Decode: process only the most recent token.    assert max_new_tokens > 0, \"max_new_tokens must be positive\"    for _ in range(max_new_tokens - 1):        logits, cache = model(next_token, past_kv=cache)        next_token = logits[:, -1, :].argmax(dim=-1, keepdim=True)        output_ids = torch.cat([output_ids, next_token], dim=1)     return output_ids  model = TinyCausalLM()prompt = torch.tensor([[10, 20, 30, 40]])generated = greedy_decode_with_cache(model, prompt, max_new_tokens=8)print(generated)\n```\n\nThe model still produces one token at a time. The difference is that it no longer recomputes the prompt tokens after prefill. The key logic is in `SelfAttention.forward()`\n\n: when `past_kv`\n\nis provided, the method appends the new key and value to the cached tensors. During decode, the model processes only the most recently generated `next_token`\n\nrather than the entire sequence. This is the basic idea behind the KV cache in production inference engines.\n\n## Memory Usage of the KV Cache\n\nThe KV cache saves compute, but it consumes memory. For each token, each layer stores a key tensor and a value tensor. The approximate memory usage is:\n\n```\nbytes = 2 * num_layers * batch_size * sequence_length\n          * num_kv_heads * head_dim * bytes_per_element\n\n12\n\nbytes = 2 * num_layers * batch_size * sequence_length          * num_kv_heads * head_dim * bytes_per_element\n```\n\nThe factor of `2`\n\nis for keys and values. The `num_kv_heads`\n\nvalue may be smaller than the number of query heads for models that use multi-query attention or grouped-query attention.\n\nFor a model with 32 layers, 32 KV heads, head dimension 128, BF16 cache values, batch size 1, and sequence length 4,096:\n\n```\n2 * 32 * 1 * 4096 * 32 * 128 * 2 bytes\n= 2,147,483,648 bytes\n= 2 GiB\n\n123\n\n2 * 32 * 1 * 4096 * 32 * 128 * 2 bytes= 2,147,483,648 bytes= 2 GiB\n```\n\nThis is only the KV cache for one request. It does not include model weights, temporary activations, tokenization buffers, or framework overhead. If the service handles many users concurrently, KV cache memory quickly becomes a limiting factor.\n\nFor this reason, an inference system must release KV cache memory when a request is finished. A simple script can let Python garbage collection handle this, a production server needs more efficient memory management, typically using cache blocks instead of individual tensors.\n\nThe layout of the cache also matters. In the simple code above, each decode step appends tensors using `torch.cat()`\n\n. This is fine for teaching, but it is inefficient because it repeatedly allocates new tensors and copies old data. Real serving engines pre-allocate cache memory in advance or use a paged layout. Later chapters will revisit this issue in detail.\n\nEfficient KV-cache management is a major differentiator among inference systems.\n\n## Further Reading\n\nBelow are some resources you may find useful:\n\n[Attention Is All You Need](https://arxiv.org/abs/1706.03762), by Vaswani et al.\n\nThis is the original Transformer paper. It introduces scaled dot-product attention, multi-head attention, and the query-key-value formulation used throughout this chapter.[Attention (machine learning)](https://en.wikipedia.org/wiki/Attention_%28machine_learning%29), on Wikipedia.\n\nThis is a useful quick reference for the attention mechanism, including the formula $\\operatorname{softmax}(QK^\\top / \\sqrt{d_k})V$ and the relationship between attention, self-attention, and the Transformer architecture.[Fast Transformer Decoding: One Write-Head is All You Need](https://arxiv.org/abs/1911.02150), by Noam Shazeer.\n\nThis paper introduces multi-query attention. It is directly related to inference because it reduces the amount of key and value data that must be read during incremental decoding.[FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://arxiv.org/abs/2205.14135), by Dao et al.\n\nFlashAttention is not only an inference algorithm; the original paper emphasizes faster Transformer training and memory-efficient exact attention. It is still relevant to inference because prompt prefill and long-context attention also benefit from reducing memory traffic and avoiding materializing the full attention matrix.[Orca: A Distributed Serving System for Transformer-Based Generative Models](https://www.usenix.org/conference/osdi22/presentation/yu), by Yu et al.\n\nThis paper focuses on inference serving. It introduces iteration-level scheduling and selective batching, which are important ideas behind continuous batching for autoregressive generation.[Efficient Memory Management for Large Language Model Serving with PagedAttention](https://arxiv.org/abs/2309.06180), by Kwon et al.\n\nThis paper is directly about LLM inference serving. PagedAttention stores the KV cache in fixed-size blocks instead of requiring each request’s cache to be contiguous, reducing memory fragmentation and allowing larger batches.\n\n## Summary\n\nIn this article, you learned that inference is not just training without the backward pass. The model is used in a different pattern: one prefill step followed by many decode steps. The KV cache avoids recomputing attention keys and values for previous tokens, changing the per-token attention cost during decode from quadratic to linear in the sequence length.\n\nYou also implemented a simple KV cache in a tiny transformer model. This cache is the foundation for many later optimizations, including paged attention, continuous batching, prefix caching, long-context inference, and disaggregated prefill and decode.", "url": "https://wpnews.pro/news/using-a-transformer-model-from-training-to-inference", "canonical_source": "https://machinelearningmastery.com/using-a-transformer-model-from-training-to-inference/", "published_at": "2026-07-31 14:22:34+00:00", "updated_at": "2026-08-03 17:40:59.489799+00:00", "lang": "en", "topics": ["large-language-models", "artificial-intelligence", "machine-learning"], "entities": ["PyTorch"], "alternates": {"html": "https://wpnews.pro/news/using-a-transformer-model-from-training-to-inference", "markdown": "https://wpnews.pro/news/using-a-transformer-model-from-training-to-inference.md", "text": "https://wpnews.pro/news/using-a-transformer-model-from-training-to-inference.txt", "jsonld": "https://wpnews.pro/news/using-a-transformer-model-from-training-to-inference.jsonld"}}