Inference and KV-Caching

Autoregressive Generation

After training, the model's purpose shifts to generation, or inference. An LLM generates text autoregressively, producing one token at a time. A naive implementation of this process is computationally expensive. For each new token, the model would re-process the entire sequence of previously generated tokens from scratch. This leads to a quadratic time complexity, $O(n^2)$ , where $n$ is the sequence length. As the sequence grows, generation slows to a crawl.

The bottleneck is the self-attention mechanism, which recomputes Key ( $K$ ) and Value ( $V$ ) tensors for all tokens at every step. However, for causal attention, the $K$ and $V$ tensors for past tokens remain unchanged when generating a new token. We can exploit this property by caching these tensors instead of recomputing them. This is the core principle of the Key-Value (KV) cache.

To reduce computational overhead, Key-Value (KV) caching is used, allowing the model to reuse previously computed attention vectors, speeding up autoregressive decoding.

The KV cache transforms per-token generation complexity from $O(n)$ to $O(1)$ in the attention layer, as the number of matrix multiplications becomes constant with respect to sequence length. However, this efficiency comes at the cost of memory. The cache size grows linearly with the sequence length and can become the primary memory bottleneck during inference.

S_{cache} = 2 \times N_{layers} \times N_{heads} \times L_{seq} \times d_{head} \times P_{bytes}

Implementation Details

Implementing the KV cache requires modifying the forward pass of your Transformer block. Instead of processing a full sequence, the inference loop will pass a single token and the cache itself. The cache can be implemented as a list of tensors (one for K, one for V) for each layer.

# Inside the TransformerBlock's forward method
def forward(self, x, kv_cache=None):
    # x is of shape (batch_size, 1, d_model)
    if kv_cache is None:
        # Initialize cache for the first token
        # Tuple of (key_cache, value_cache)
        # Shape: (batch, num_heads, 0, head_dim)
        past_key_values = ( 
            torch.zeros((x.shape[0], self.num_heads, 0, self.head_dim), device=x.device), 
            torch.zeros((x.shape[0], self.num_heads, 0, self.head_dim), device=x.device) 
        )
    else:
        past_key_values = kv_cache

    # Self-attention with cache
    attn_output, present_key_values = self.attention(x, past_key_values)
    
    # ... (residual connection, FFN, etc.) ...
    
    return output, present_key_values

# Inside the Attention module's forward method
def attention(self, x, past_key_values):
    # Project to Q, K, V for the current token
    q, k, v = self.wq(x), self.wk(x), self.wv(x)
    
    # Reshape and apply RoPE as before
    # ... q, k, v are now (batch, num_heads, 1, head_dim)
    
    # Retrieve and concatenate with cache
    past_k, past_v = past_key_values
    keys = torch.cat([past_k, k], dim=2)
    values = torch.cat([past_v, v], dim=2)
    
    # The new cache for this layer
    present_key_values = (keys, values)

    # Perform scaled dot-product attention with the full key/value sets
    # attn_weights = (q @ keys.transpose(-2, -1)) * (1.0 / math.sqrt(keys.size(-1)))
    # ... apply softmax and compute output
    
    return attn_output, present_key_values

The autoregressive loop manages this process. It starts with a prompt, generates a token, appends it to the input, and feeds it back with the updated cache to generate the next token. This continues until a stopping condition is met, such as generating an end-of-sequence (EOS) token or reaching a maximum length.

Controlling Generation

The raw output of the model is a tensor of logits, representing unnormalized scores for every token in the vocabulary. To convert this into a probability distribution, we apply a softmax function. The temperature parameter is used to control the randomness of the output by scaling the logits before the softmax.

P(x_i) = \frac{e^{z_i / T}}{\sum_j e^{z_j / T}}

Simply sampling from this distribution can produce incoherent text. We need more sophisticated strategies. Top-K sampling restricts the sampling pool to the $K$ tokens with the highest probabilities. A more adaptive method is Nucleus (Top-P) sampling, which selects the smallest set of tokens whose cumulative probability exceeds a threshold $P$ . This allows the sampling pool size to vary based on the model's confidence.

import torch
import torch.nn.functional as F

def sample(logits, temperature=1.0, top_k=None, top_p=None):
    # logits shape: (batch_size, vocab_size)

    # Apply temperature
    if temperature > 0:
        logits = logits / temperature

    # Apply Top-K
    if top_k is not None:
        v, _ = torch.topk(logits, top_k)
        # Set all other logits to negative infinity
        logits[logits < v[:, [-1]]] = -float('Inf')

    probs = F.softmax(logits, dim=-1)

    # Apply Top-P (Nucleus Sampling)
    if top_p is not None:
        sorted_probs, sorted_indices = torch.sort(probs, descending=True)
        cumulative_probs = torch.cumsum(sorted_probs, dim=-1)
        
        # Find indices to remove
        sorted_indices_to_remove = cumulative_probs > top_p
        # Shift right to keep the first token that exceeds the threshold
        sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
        sorted_indices_to_remove[..., 0] = 0

        # Scatter the boolean mask to the original indices
        indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove)
        probs[indices_to_remove] = 0
        # Renormalize the probabilities
        probs = probs / probs.sum(dim=-1, keepdim=True)

    # Sample from the final distribution
    next_token = torch.multinomial(probs, num_samples=1)
    return next_token

By combining the KV cache for efficiency with controlled sampling techniques like Top-P, we can build a robust and performant inference engine for our LLM. The final step is to orchestrate these components within a generation loop that handles tokenization, state management, and stopping criteria.

Architecting Transformers from Scratch in Python

Architecting Transformers from Scratch in Python

Inference and KV-Caching

Autoregressive Generation

Implementation Details

Controlling Generation