Greedy decoding / generation

The problem, from first principles

Everything built so far in this Part produces logits GIVEN a full sequence of tokens; actually GENERATING new text means running the model repeatedly, each time feeding it everything produced SO FAR and asking it what should come next. The simplest possible decision rule: always pick whichever token the model's own logits rank HIGHEST, append it, and repeat, "greedy" decoding, since it greedily takes the single best-looking option at every step with no lookahead, no consideration of how that choice affects future steps.

Every autoregressive generation step must respect causality: [02-modern-transformer-architecture/01-encoder-decoder-arrangements]'s decoder framing applies directly here, since the model must never be allowed to "see" a token that hasn't been generated (and appended) yet, exactly [04-seq-modeling/04-attention/02-causal-mask]'s causal mask, rebuilt fresh at every step as the sequence GROWS.

From theory to code

Implement greedy_decode(token_ids, token_embedding_table, blocks_params, num_heads, tied, output_weight, num_new_tokens): repeatedly call [04-full-forward-pass]'s full_lm_forward (with a causal mask sized to the CURRENT sequence length) on the sequence so far, take argmax of the LAST position's logits, append that token, and repeat num_new_tokens times.

Constraints

• The causal mask is rebuilt at EVERY step, sized to the current (growing) sequence length, never a fixed size computed once upfront.
• Only the LAST position's logits determine the next token (logits[..., -1, :]); earlier positions' logits were already used to predict THEIR own next tokens in previous steps, and are not reconsidered.
• argmax breaks ties deterministically (NumPy's own convention: the FIRST maximal index), so running the same inputs twice produces IDENTICAL output, no randomness anywhere in greedy decoding.
• The final result's first token_ids.shape[-1] positions are exactly the ORIGINAL input, unchanged; only new tokens are appended after it.

Hints

seq_len = token_ids.shape[-1]
mask = build_causal_mask(seq_len)
logits = full_lm_forward(token_ids, token_embedding_table, blocks_params, num_heads, tied, output_weight, mask=mask)
next_token_logits = logits[..., -1, :]
next_token = np.argmax(next_token_logits, axis=-1)