1. Prefill and decode
The problem: An 8,000-token prompt arrives all at once; the answer then leaves token by token. If the engine processes the prompt at generation pace — one token at a time — the user waits whole seconds for the first word.
The idea: Two phases, two regimes: prefill sees all prompt tokens at once (massively parallelizable work); decode adds one token per step (intrinsically sequential work). Same mechanism, opposite execution profiles.
Why / at what price: Separating them lets you optimize each — first-token latency on one side, generation throughput on the other. The price: two code paths for one mechanism, which must produce exactly the same numbers.
Understanding check
Name the input, transformed state, output, and one required assumption. Then compare your chain with the explanation above.
2. Naive recurrence
The problem: Recurrent memory seems to doom prefill: S₅ needs S₄, which needs S₃… Running 8,000 updates one after another leaves a GPU — built for whole matrices — nearly idle at every step.
The idea: The precise diagnosis: the DEPENDENCY is sequential (each S_t depends on S_{t−1}), but most of the per-token COMPUTE — local q·k products, k vᵀ writes — is not. Naive recurrence serializes everything because it never separates the two.
Why / at what price: That observation opens the door to chunking: serialize only what must be. The price of staying naive is directly measurable: matrix units billed by the hour executing vector-matrix products.
Understanding check
Name the input, transformed state, output, and one required assumption. Then compare your chain with the explanation above.
3. Split into chunks
The problem: How do you hand the GPU full matrix blocks without violating order? You need a split where a block’s interior computes in parallel and the distant past arrives compressed — with no double counting and no forgetting.
The idea: A chunk of C tokens computes all its permitted internal interactions at once (a triangular C×C matrix) and reads the earlier past through the incoming state: O = M·V + K·S_in. In the trace, chunk 1 produces S₄, chunk 2 consumes it — and o₅..o₈ are exactly those of the recurrence.
Why / at what price: Algebraic exactness, parallelism recovered. The price: C² scratch memory for the triangle, and real code complexity — two terms to add, hence two ways to get it wrong: precisely the trace’s two bugs.
Understanding check
Name the input, transformed state, output, and one required assumption. Then compare your chain with the explanation above.
4. Causal triangle
The problem: Inside a chunk, all tokens are computed together — including t5 with t7, which is its future. Without a guard, prefill would learn dependencies that decode can never reproduce.
The idea: A lower-triangular matrix embodies the rule “i reads only j ≤ i”: cells above the diagonal are forbidden by construction. Reading trap: a 0 below the diagonal is a null dot product (allowed); a “.” above it is causality.
Why / at what price: The triangle makes the constraint checkable at a glance and free to apply. The price of a wrong mask is vicious: perplexity improves “suspiciously well” in prefill and nothing breaks — until decode, which cannot cheat.
Understanding check
Name the input, transformed state, output, and one required assumption. Then compare your chain with the explanation above.
5. Incoming and outgoing state
The problem: Chunk 2 must never revisit a chunk-1 token — or parallelism collapses — yet o₅ depends on v₁ and v₃. How do you hand over “all the useful past” without handing over the past?
The idea: Through the state: chunk 1 emits S₄ = K₁ᵀV₁, a fixed-size summary; chunk 2 reads it via q_tᵀS₄ and adds its local terms. In the trace: o₅ = [3,4] (inherited) + [0,1] (local) = [3,5]. Boundaries carry order and causality, not tokens.
Why / at what price: One object handed between blocks, fixed size. The price: chunk 2 can no longer decompose [3,4] into v₁ and v₃ — session 13’s compression applies at the boundary. And dropping the incoming term (bug 1) amputates the whole prompt.
Understanding check
Name the input, transformed state, output, and one required assumption. Then compare your chain with the explanation above.
6. Chunk size
The problem: C = 1 brings back the slow recurrence; C = full length blows the triangle up as C². Between the two, who decides? The same code can run several times slower with a C ill-chosen for the GPU.
The idea: C arbitrates two opposing costs: sequential transitions in n/C versus scratch memory in C². Doubling C halves the transitions and quadruples the triangle — the optimum sits where the triangle just saturates fast memory (SRAM).
Why / at what price: A well-chosen C saturates the hardware. The price: the right C does not transfer between GPUs — it is an execution parameter to re-measure, not a model constant. And it never changes the results, only their cost.
Understanding check
Name the input, transformed state, output, and one required assumption. Then compare your chain with the explanation above.
Development from the course source
Chapter 10 — Chunking and parallel prefill
10.1 Why chunks are needed
Goal: combine a recurrent memory rule with fast GPU work. Processing one token after another is easy to describe but fails to use all parallel arithmetic units efficiently during prompt reading, called prefill.
Intuition: Instead of carrying groceries one item at a time, place many items in a box and move the box.
Step by step
-
Split N tokens into chunks, perhaps 64 or 128 tokens in an implementation.
-
Within a chunk, arrange many reads and updates as large matrix operations.
-
Carry the final state S from one chunk to the next.
-
Preserve causal order: token t must not use future tokens.
Worked example: 12 tokens split into chunks of 4 gives chunks 1–4, 5–8 and 9–12. The second chunk receives the state summarizing tokens 1–4. Its four tokens can perform much of their arithmetic together, then produce a state for the third chunk.
Why it matters: Chunking does not change the learning goal. It reorganizes equivalent or carefully derived operations so hardware can execute them efficiently.
Quick check: Why not use one enormous chunk automatically? Answer: larger chunks increase temporary work and storage; the best size depends on hardware and kernels.
10.2 Causal triangular structure
Goal: understand the lower-triangular mask used inside a chunk. A lower-triangular matrix has zeros above its main diagonal.
Intuition: It is a school rule saying each student may read only earlier lines, never answers written later.
Step by step
-
For four positions, allowed links form [[1,0,0,0],[1,1,0,0],[1,1,1,0],[1,1,1,1]].
-
Row 3 may use positions 1, 2 and 3.
-
Row 1 may use only position 1.
Worked example: The triangle protects next-token prediction from cheating during training and prefill.
Why it matters: Causality is a correctness condition, not merely an optimization detail.
Quick check: In a causal block, may token 2 use token 4? Answer: No.
Complete worked case
For 8 tokens in chunks of 4, the first computes a 4×4 triangle then passes S₄. The second receives S₄, computes its local triangle, and produces S₈. No token in the first block can read the second.
Reading method: write the data, state every object shape, perform one transformation, and interpret the result before continuing.