Teacher guide — Chunking, causality, and parallel prefill
Duration: 120 minutes
Positioning: Reconcile recurrent state and GPU parallelism through causal block computation.
Expected evidence: Established mechanisms; numerical simplifications are pedagogical.
Observable outcomes and preparation
- Distinguish prefill and decoding.
- Trace state across chunks.
- Explain causal triangular structure.
Before class, the instructor runs the worked case and lab, prints the exercise packet, prepares a four-column board—assumption, prediction, observation, delta—and checks that every mathematical notation is paired with dimensions. The demonstration must not become slide reading.
Diagnostic
- Explain in one sentence: Prefill and decode. What observation would falsify your explanation?
- Explain in one sentence: Naive recurrence. What observation would falsify your explanation?
- Explain in one sentence: Split into chunks. What observation would falsify your explanation?
Teaching decision: if two of three answers remain nominal or lack a validity condition, rebuild the vocabulary with a numerical example before any formula. A fluent but unfalsifiable answer does not count as mastery.
Timed plan
| Time | Activity | Observable evidence |
|---|---|---|
| 0–10 min | Individual diagnostic, then pair comparison | Three answers and one named uncertainty |
| 10–25 min | Situation and vocabulary | Annotated input → state → output diagram |
| 25–55 min | Develop the mechanism on the board | Shapes, assumptions, and intermediate calculation visible |
| 55–75 min | Worked case with deliberate errors | Reasoned correction, not only the right number |
| 75–95 min | Causal lab: predict, change one variable, run | Prediction / observation / delta table |
| 95–112 min | Exercises 1 and 2 with peer correction | Retained artifact and applied rubric |
| 112–120 min | Exit ticket and transfer | Mechanism, boundary, next experiment |
Teaching notes
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.
Running the 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.
Do not reveal the result at once. Ask learners to predict the next operation, its shape, and the expected sign. After each line ask: “What changed? What stayed fixed? Which assumption did we use?” A calculation error repaired with a causal chain is worth more than a guessed result.
Lab protocol
- Write a qualitative and, where possible, numerical prediction before touching a control.
- Change one variable only; retain a capture or record initial and final values.
- Explain the delta through the mechanism, not “the tool did that.”
- Test one boundary value and state where the model stops representing a real system.
Misconceptions
| # | Observable misconception | Grounded correction | Probe |
|---|---|---|---|
| 1 | “Prefill and decode guarantees the outcome without assumptions or measurement.” | 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. | Ask for a counterexample, then restate the mechanism with its validity condition. |
| 2 | “Naive recurrence guarantees the outcome without assumptions or measurement.” | 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. | Ask for a counterexample, then restate the mechanism with its validity condition. |
| 3 | “Split into chunks guarantees the outcome without assumptions or measurement.” | 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. | Ask for a counterexample, then restate the mechanism with its validity condition. |
Boundary to maintain: Chunking improves execution; it does not automatically change the information capacity of state.
Probing questions
- If we remove or reverse Prefill and decode, which output changes first, and what observation would show it?
- If we remove or reverse Naive recurrence, which output changes first, and what observation would show it?
- If we remove or reverse Split into chunks, which output changes first, and what observation would show it?
- If we remove or reverse Causal triangle, which output changes first, and what observation would show it?
- If we remove or reverse Incoming and outgoing state, which output changes first, and what observation would show it?
- If we remove or reverse Chunk size, which output changes first, and what observation would show it?
Assessment
| Level | Criterion |
|---|---|
| 0 | Repeats terms without connecting input, transformation, and output. |
| 1 | Describes the chain but checks neither shape nor assumption. |
| 2 | Executes the case, explains the result, and names one limitation. |
| 3 | Transfers to a new case, compares an alternative, and proposes a measurement that could invalidate the choice. |
Exit threshold: level 2 on the worked case and at least one exercise; a memorized formula without interpretation remains level 1.
Observation and remediation protocol
During discussion, the instructor records evidence rather than impressions. Evidence of understanding contains a named object, a justified transformation, and a checkable consequence. If a learner gives the right result without a chain, ask for the preceding line. If the chain is coherent but the result is wrong, preserve the chain and isolate the arithmetic error. If vocabulary from another concept is used, compare both mechanisms in an input, state, output, cost, and boundary table. Remediation targets the first break only: vocabulary, shapes, operation, interpretation, or claim scope. After correction, use a neighboring case with a changed value; success on the same example does not prove transfer. For pair work, assign operator and verifier roles, then swap. The verifier does not supply the answer: they request an assumption, check the shape, and ask what observation could contradict the reasoning. The instructor retains the exit ticket and classifies the dominant break. The next session opens with a three-minute problem aimed at that break instead of repeating the whole lesson.
Differentiation
- Support: provide shapes and the first transformation; let the learner complete interpretation and boundary.
- Core path: worked case without result, lab with one assigned variable, diagnostic exercise.
- Extension: change one assumption, compare two mechanisms, and define the metric that would decide between them.
Post-session follow-up
Within twenty-four hours, return each annotated exit ticket with one priority, the exact resource to reopen, and a mini-case different from the worked case. Revision requires three items: a written prediction, a retained trace, and one sentence explaining the delta. At the next session, sample two submissions: one that repaired the break and one that remains ambiguous. Discuss them anonymously, then state the criterion that separates them. Do not use completion rate as evidence of mastery. Evidence of remediation is a correct chain on a new case with a stated boundary. If the same break appears in more than one third of the group, repair the support or demonstration before blaming learners.
Follow-up closes only when the new artifact shows the causal chain, check, and boundary—not merely when a file has been submitted.
Sources and evidence boundary
- Owner-supplied bilingual course packet, Chapter 10.
- Katharopoulos et al., “Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention”, ICML (2020).
- Yang et al., “Gated Linear Attention Transformers with Hardware-Efficient Training” (chunkwise parallel form), ICML (2024).
- Course source packet supplied by the owner; named-product details remain source-reported until primary verification.
Scope: Established mechanisms; numerical simplifications are pedagogical. These references support the session frame; they do not turn a reported product choice into an independently verified result.
Exit ticket
In no more than six lines: mechanism; calculation or trace; observation; boundary; evidence level; next experiment. The instructor marks one priority causal break for revision.