Exercises with solutions — Chunking, causality, and parallel prefill

General instruction: every answer must show data, transformation, result, one check, and one limitation. A bare number or copied definition is insufficient.

Starting data: 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.

Boundary to retain: Chunking improves execution; it does not automatically change the information capacity of state.

Exercise 1 — Calculated trace — Prefill and decode

Reproduce and annotate the chain input → state → output. Move from 8 to 10 tokens with chunk size 4. Draw the chunks, local causal masks, and states carried between chunks.

Deliverable: a data → operation → result → interpretation table, plus two sentences about the changed value.

Worked solution

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.

Solved variant: Chunk sizes are [4,4,2]. The first two use 4×4 triangles; the last uses a 2×2 triangle. End states S₄ then S₈ cross boundaries; no position may read a future token.

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. 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. The minimum check covers dimensions, sign, and order of magnitude. If the observed change contradicts the prediction, locate the first operation whose direction changes instead of fixing only the final line.

Rubric Exercise 1 — /10

Criterion	Points
Explicit data and shapes	2
Traceable calculation	3
Prediction before variation	2
Interpretation and check	2
Named limitation	1

Exercise 2 — Diagnose a seductive explanation — Split into chunks

A colleague claims: « Split into chunks proves the system will be accurate, fast, and stable in every context. »

Separate mechanism, assumption, observation, and conclusion.
Name two correct lesson elements and two unsupported extrapolations.
Propose a bounded experiment with controlled variable, metric, and stop threshold.
Rewrite the claim as one defensible sentence.

Reasoned solution

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. 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. Chunking improves execution; it does not automatically change the information capacity of state.

The claim mixes a local relation with a global guarantee. A defensible version states only the observed mechanism, test conditions, and measured metric. Stop the test if shapes become invalid, the metric crosses the declared degradation threshold, or another variable changes.

Rubric Exercise 2 — /10

2 points per element: separation, lesson grounding, extrapolations, protocol, and rewrite.

Exercise 3 — Architecture decision and transfer — Incoming and outgoing state

You must reproduce 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.” under two conditions. Option A uses the full chain through “Incoming and outgoing state.” Option B is a transparent baseline that retains “Prefill and decode,” calculates the expected output directly, and does not use the compression or adjustment mechanism studied. Build a decision record containing:

the workload and dominant constraint;
each option’s mechanism, without slogans;
one quality, memory, or latency prediction;
one case where your preferred procedure loses;
an A/B protocol, metrics, and rollback threshold;
a bounded verdict: choose, defer, or reject.

Elements of a strong solution

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. 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). Established mechanisms; numerical simplifications are pedagogical.

A strong answer does not make the newer mechanism the default winner. It retains a measurable baseline, sets thresholds before testing, and separates component cost from whole-system behavior. The verdict names what remains uncertain and the next evidence that could change it.

Rubric Exercise 3 — /15

Criterion	Points
Framing and baseline	3
Compared causal chains	4
Protocol and metrics	4
Rollback threshold	2
Bounded verdict	2

Extension

Repeat Exercise 3 after reversing the dominant constraint. If you optimized memory, impose a strict quality floor; if you optimized fidelity, halve the memory budget. Identify the first part of the verdict that changes and the evidence required.

Review before submission

Review the packet as if another group had to reproduce it without speaking to you. Are all starting values or assumptions present? Are shapes or roles stated before operations? Does the prediction truly precede the observation? Is the result translated into behavior rather than left as an isolated number? Did you test a boundary value and identify a stop condition? Does the procedure or architecture choice retain a measurable baseline and a rollback threshold set before the test? Finally, highlight one sentence that states what is established, one that remains a hypothesis, and one measurement that could change your verdict. If any element is missing, the work is not reproducible.

Reference appendix for correction

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.

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.