Autoregressive decoder loop in PyTorch

When you call model.generate() in Hugging Face Transformers, a lot happens behind the scenes: KV cache management, sampling, stopping criteria, and logit processing. That's fine for using models, but if you want to understand inference deeply enough to optimize it, you need to build the loop yourself. Today I'm writing the autoregressive decoder loop from scratch in PyTorch, step by step.

The core idea

Autoregressive generation is embarrassingly simple in concept: run the model forward on the input, get logits for the next token, sample a token, append it to the input, repeat. The loop looks like this:

import torch

@torch.no_grad()
def generate_naive(model, input_ids, max_new_tokens=128):
    """The simplest possible generation loop. No KV cache."""
    generated = input_ids.clone()  # (1, seq_len)

    for _ in range(max_new_tokens):
        logits = model(generated).logits        # (1, seq_len, vocab_size)
        next_logits = logits[:, -1, :]           # (1, vocab_size)
        next_token = torch.argmax(next_logits, dim=-1, keepdim=True)  # greedy
        generated = torch.cat([generated, next_token], dim=1)

        if next_token.item() == tokenizer.eos_token_id:
            break

    return generated

This works but is horrifically slow. On every step, the model recomputes attention over the entire sequence, including all the tokens it already processed. Step 100 recomputes attention for tokens 1 through 99 even though their keys and values haven't changed. This is O(n^2) in sequence length, and it's the reason KV caching exists.

Adding KV cache

The fix is to cache the key and value tensors from previous steps. On step t, instead of passing the full sequence, you pass only the new token and the cached KV pairs. The model computes attention between the new query and all cached keys/values, then appends the new K and V to the cache.

@torch.no_grad()
def generate_with_kv_cache(model, input_ids, max_new_tokens=128):
    """Generation with KV cache. Two phases: prefill and decode."""
    # Phase 1: Prefill - process the full prompt, populate KV cache
    outputs = model(input_ids, use_cache=True)
    next_logits = outputs.logits[:, -1, :]
    past_key_values = outputs.past_key_values  # tuple of (K, V) per layer

    next_token = torch.argmax(next_logits, dim=-1, keepdim=True)
    generated_tokens = [next_token]

    # Phase 2: Decode - one token at a time, reusing KV cache
    for _ in range(max_new_tokens - 1):
        outputs = model(
            next_token,
            past_key_values=past_key_values,
            use_cache=True,
        )
        next_logits = outputs.logits[:, -1, :]
        past_key_values = outputs.past_key_values

        next_token = torch.argmax(next_logits, dim=-1, keepdim=True)
        generated_tokens.append(next_token)

        if next_token.item() == tokenizer.eos_token_id:
            break

    return torch.cat([input_ids] + generated_tokens, dim=1)

Now the two phases of inference are explicit in the code:

• Prefill: the first forward pass processes the full input prompt. It's a large matmul (the full sequence times the model width), and it populates the KV cache. This is compute-bound because the batch of tokens is large.
• Decode: each subsequent step passes a single token. The model does a small vector-matrix multiply (one query against all cached keys), reads the entire cache from memory, and produces one token. This is memory-bound because you're loading gigabytes of cache and weights for a tiny amount of compute.

This is the same prefill/decode split I discussed in the ops:byte ratio post, but now you can see exactly where it happens in code.

Sampling strategies

Greedy decoding (argmax) is deterministic but produces boring, repetitive text. Real generation uses sampling with temperature, top-k, and top-p (nucleus sampling):

def sample_token(logits, temperature=0.8, top_k=50, top_p=0.9):
    """Sample a token with temperature, top-k, and top-p."""
    # Temperature scaling
    logits = logits / temperature

    # Top-k: zero out everything below the k-th highest logit
    if top_k > 0:
        top_k_values, _ = torch.topk(logits, top_k, dim=-1)
        threshold = top_k_values[:, -1].unsqueeze(-1)
        logits = logits.masked_fill(logits < threshold, float("-inf"))

    # Top-p (nucleus): keep the smallest set of tokens whose
    # cumulative probability exceeds p
    if top_p < 1.0:
        sorted_logits, sorted_indices = torch.sort(logits, descending=True, dim=-1)
        cumulative_probs = torch.cumsum(torch.softmax(sorted_logits, dim=-1), dim=-1)
        # Remove tokens with cumulative probability above the threshold
        sorted_mask = cumulative_probs - torch.softmax(sorted_logits, dim=-1) >= top_p
        sorted_logits[sorted_mask] = float("-inf")
        # Scatter back to original ordering
        logits = sorted_logits.scatter(1, sorted_indices, sorted_logits)

    probs = torch.softmax(logits, dim=-1)
    return torch.multinomial(probs, num_samples=1)

Temperature controls randomness: lower values sharpen the distribution (more deterministic), higher values flatten it (more creative). Top-k restricts sampling to the k most likely tokens. Top-p restricts to the smallest set of tokens whose cumulative probability exceeds p. In practice, most deployments use temperature 0.6 to 0.8 with top-p 0.9.

Position IDs and attention masks

Two details that trip people up when building the loop manually:

Position IDs. During prefill, positions are [0, 1, 2, ..., N-1]. During decode, each new token's position is the length of the sequence so far. If you don't pass position IDs explicitly, most models compute them from the input shape, which works during prefill but needs the KV cache length during decode:

# During decode step t, the position of the new token
position_ids = torch.tensor([[prefill_length + t]], device=device)

Attention mask. The causal mask ensures each token can only attend to tokens at earlier positions. During prefill, this is a lower-triangular matrix. During decode with KV cache, the mask is a single row: the new token attends to all previous positions (including itself). Most Hugging Face models handle this automatically when you pass past_key_values, but if you're implementing attention from scratch, you need to get the mask shape right.

What this loop reveals

Building the loop yourself makes several things concrete that are invisible when using model.generate():

• Why KV cache memory grows linearly. Each decode step appends a K and V tensor (of shape [num_heads, 1, head_dim]) per layer to the cache. For a 70B model with 80 layers and 64 heads at head_dim=128 in FP16, each token adds 80 × 2 × 64 × 128 × 2 bytes = 2.6 MB to the cache. Generate 2048 tokens and you've used 5.3 GB of GPU memory just for one request's cache.
• Why prefill is fast and decode is slow. Prefill is one big forward pass. Decode is N small forward passes, each loading the full model weights from memory. You can time both phases separately and see the gap directly.
• Why batching helps. In the naive loop, the GPU does one vector-matrix multiply per decode step. With batching, you do multiple queries against the same weights, turning it into a matrix-matrix multiply and raising the arithmetic intensity.

This loop is the foundation of everything else in the Deep Implementation phase. Next: scaled dot-product attention with masking, where we implement the attention mechanism itself from scratch.