Using a Transformer Model: From Training to Inference

Using a Transformer Model: From Training to Inference - MachineLearningMastery.com Using a Transformer Model: From Training to Inference - MachineLearningMastery.com 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. This 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. In this chapter, you will learn about: The autoregressive generation loop The difference between prefill and decode Why key-value caching is necessary How to implement a simple KV cache How to reason about the memory used by the cache Let’s get started. Using a Transformer Model: From Training to Inference Photo by Jacob Smith. Some rights reserved. Overview This chapter is divided into four parts; they are: Autoregressive Generation Prefill and Decode A Simple KV Cache Memory Usage of the KV Cache Autoregressive Generation A 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: 1 The cat sat on the the 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. The generation loop is therefore simple: Tokenize the prompt. Run the model to obtain logits for the next token. Choose a token from the logits. Append that token to the input. Repeat until a stopping rule is reached. This 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. A very small greedy decoding loop can be written as follows: 1 2 3 4 5 6 7 8 9 10 11 12 13 import 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 In the code above, model is a PyTorch model, max_new_tokens is a positive integer, and all other variables are PyTorch tensors. The for-loop iterates max_new_tokens times, and at each iteration, it feeds the entire sequence back into the model to get the logits for the next token. The argmax() function selects the highest-scoring token. The cat() function 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. This 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$. The 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)$.) Inference systems mitigate this by splitting generation into two phases: prefill and decode. Prefill and Decode Generation 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. During 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. After 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. This 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. This distinction is important enough that serving systems usually measure prefill and decode separately: 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. A 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. The 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. The attention-score matrix grows during generation. Prefill computes the prompt block once (in blue). Each decode iteration appends one row (due to expanded $Q$) and one column (due to expanded $K$) for the newly generated token. A Simple KV Cache The 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. This model is not intended to produce useful text. Its purpose is to show how the cache is created during prefill and extended during decode. 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 import math import torch import torch.nn as nn import 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 The cache is a list with one element per transformer layer. Each element is a pair (k, v). The shape of each tensor is: 1 [batch_size, num_heads, sequence_length, head_dim] During prefill, sequence_length is the prompt length. During decode, the model receives one token at a time and appends one position to the cache. You may notice that only keys and values are stored in the cache but not the query tensor. Note that the forward() method 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. Here is a minimal generation loop using the cache: 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 torch.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) The 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(): when past_kv is 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 rather than the entire sequence. This is the [truncated for AI cost control]