Build your own KV cache
The fast generate loop·20 min

Greedy decode end-to-end

0 / 1
# grade this stage (from your course directory)
npx theinference test --stage greedy-decode --watch

You have all the pieces: prefill populates every layer's cache and returns the first prediction, and decodeStep advances one cached token. Now you wire them into a loop. Greedy decoding is the simplest possible policy — at each step, take the argmax of the current logits, append that token, and feed it back in. No sampling, no temperature, no randomness. The same prompt always produces the same continuation.

The greedy loop
logits   = prefill(promptIds)          # cache now holds the whole prompt
pos      = promptIds.length
for step in 0..maxTokens-1:
    next   = argmax(logits)            # pick the highest-scoring token id
    emit next
    logits = decodeStep(next, pos)     # one cached step at the next position
    pos   += 1

What to build

Your task

Implement generateTokens and runGreedyGenerate

In src/kvcache.ts, finish argmax(v) (index of the largest entry) and generateTokens(promptIds, maxTokens, w): build a fresh ModelCache, prefill the prompt, then run the greedy loop — argmax the current logits, push that id, decodeStep it at the running position, repeat maxTokens times. The first generated token comes from the prefill logits, so don't decode before your first argmax. Then expose it through the grader adapter, which loads the weights and forwards the call.

tests/adapters.ts
export async function runGreedyGenerate(promptIds: number[], maxTokens: number): Promise<number[]> {
  const w = await weights();
  return generateTokens(promptIds, maxTokens, w);
}
The position keeps counting from the prompt

After prefilling a prompt of length L, the first generated token sits at absolute position L, the next at L + 1, and so on — exactly the rule you enforced in Chapter 3. Start pos at promptIds.length and increment it every step. Reset it to 0 and GPT-2's absolute position embeddings will quietly point the wrong tokens at the wrong slots.

Greedy is deterministic — so we grade the exact token-ID sequence

Because argmax has no randomness, the continuation is fully determined by the prompt and the weights. That lets us grade something stronger than logit-allclose: the exact sequence of token IDs your loop emits must match the reference, integer for integer. A near-tie in the logits that flips the argmax is the only way your cached text could diverge from a recompute run — and the fixture prompt is chosen to avoid ties, so a correct cache produces the identical sequence.

How it's graded

The grader calls runGreedyGenerate(promptIds, maxTokens) and checks the returned array exactly equals the reference continuation — same length, same IDs, same order. There's no tolerance here: it's an integer comparison. If the logits underneath are correct (you proved that in Chapters 3–4), the argmax sequence falls out for free.

Per-step identity over a generationThe generate CLI