Code · architecture · study mode
This if/else is the entire distinction between reading the prompt (prefill) and generating new text (decode) — it just chooses where the next token comes from.
Full explanation below the code →
There is no separate prefill() and decode() in this codebase. Both phases run the same forward(); the only difference is where the input token comes from, decided by this one branch (lines 1067–1071):
- While pos < num_prompt_tokens — we are still inside the prompt, so feed the next
prompt token (line 1068). This is prefill: teacher-forcing the model with the
known input to build up the KV cache.
- Once pos reaches the prompt length — the prompt is consumed, so feed the model's own
previously generated token, next (line 1070). This is decode: autoregressive
generation.
So "prefill" here is a *phase*, not a batched matrix operation. The prompt tokens are pushed through one at a time, exactly like decode — each is a single-vector forward(), i.e. a GEMV. The KV cache fills up token by token either way.
The lesson's prefill exists as a phase but is still token-by-token GEMV. Production engines turn that same phase into one big GEMM by packing all prompt tokens into a matrix — the main reason production prefill is so much faster than looping one token at a time.
In this code, is prompt processing (prefill) a GEMV or a GEMM operation, and how does that differ from a production engine?