Fill In the Sampling and Decode Loop of an LLM Inference Notebook

Context

One candidate account of the Anthropic Fellows final round describes a 55-minute live exercise in Google Colab. The notebook held 20–30 lines of skeleton Python with two incomplete functions that finish an LLM inference step. Using an LLM was explicitly not allowed. Candidates were told to check the notebook environment and GPU beforehand. The skeleton below is a practice reconstruction of that shape; the real notebook's functions were not published.

Notebook Skeleton

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

tok = AutoTokenizer.from_pretrained("gpt2", padding_side="left")
tok.pad_token = tok.eos_token
model = AutoModelForCausalLM.from_pretrained("gpt2").to("cuda").eval()


def sample_next_token(logits, temperature=1.0, top_k=0, top_p=1.0, generator=None):
    """logits: (batch, vocab) for the last position.
    temperature == 0 -> greedy. top_k == 0 / top_p == 1.0 -> disabled.
    Returns (batch,) int64 token ids."""
    # TODO


@torch.no_grad()
def generate(model, input_ids, attention_mask, max_new_tokens, eos_token_id, pad_token_id, **sampling):
    """input_ids / attention_mask: (batch, seq), LEFT-padded.
    Run the prompt once, then decode one token at a time reusing past_key_values.
    Sequences that emit EOS keep emitting pad_token_id. Stop early when all are done.
    Returns (batch, seq + generated) token ids."""
    # TODO


batch = tok(["The capital of France is", "def fibonacci(n):"], return_tensors="pt", padding=True).to("cuda")
ids = generate(model, batch.input_ids, batch.attention_mask, 20, tok.eos_token_id, tok.pad_token_id, temperature=0)
print(tok.batch_decode(ids, skip_special_tokens=True))

Requirements

1. With temperature=0, the output must match model.generate(..., do_sample=False) exactly, for a left-padded batch of different-length prompts.
2. top_k=1 must behave like greedy at any temperature.
3. Sampling must be reproducible with a seeded torch.Generator.
4. The prompt is processed once. Each later step feeds only the newest token plus past_key_values.
5. Be ready to explain what changes for a batch where one sequence finishes early.