Chunking, causality, and parallel prefill

Reconcile recurrent state and GPU parallelism through causal block computation.

Applied AI · advanced · Session 15

Mechanism map

PROMPT: 8 tokens, chunks of 4            PREFILL (parallel)
┌──────── CHUNK 1 : t1..t4 ────────┐   ┌──────── CHUNK 2 : t5..t8 ────────┐
│  M = tril(Q Kᵀ)   4×4            │   │  M = tril(Q Kᵀ)   4×4            │
│    ┌1 . . .┐   “ . ” = masked    │   │    ┌1 . . .┐                     │
│    │0 1 . .│     (future)        │   │    │0 1 . .│                     │
│    │1 0 1 .│   “ 0 ” = zero q·k  │   │    │1 0 1 .│                     │
│    └0 1 0 1┘     (allowed)       │   │    └0 1 0 1┘                     │
│  O = M·V + K·S₀                  │   │  O = M·V + K·S₄                  │
└───────────────┬──────────────────┘   └───────────────┬──────────────────┘
   S₀ = 0 ──────┤                            S₄ ───────┤
                ▼  S₄ = S₀ + K₁ᵀV₁                     ▼  S₈ = S₄ + K₂ᵀV₂
          S₄ = [[3,4],          ──SEQUENTIAL──▶  S₈ = [[4,5],
                [5,5]]           (one hop only)        [7,7]]

DECODE: 1 token → M is 1×1, all that remains is O = qᵀS  (no triangle)

The problem — Prefill and decode

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 — Prefill and decode

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 — Prefill and decode

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.

Check: Prefill over the 8 tokens builds a 4×4 triangle per chunk. When decoding token 9, what is the size of that triangle, and which part of the computation disappears entirely?

The problem — Naive recurrence

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 — Naive recurrence

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 — Naive recurrence

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.

Check: The recurrence requires S₈ after S₄. Why can the 8 updates not run in parallel, and why is the GPU nonetheless underused if you do them strictly one at a time?

Visual support — Naive recurrence

naive recurrence: 8 sequential steps
t1 ▶ t2 ▶ t3 ▶ t4 ▶ t5 ▶ t6 ▶ t7 ▶ t8    GPU: ~1 useful token/step

chunking C=4: 2 sequential steps
┌ t1 t2 t3 t4 ┐ ──S₄──▶ ┌ t5 t6 t7 t8 ┐   GPU: 4 tokens/step
└ in parallel ┘         └ in parallel ┘

the dependency remains (S₈ after S₄) —
but it no longer carries every token

The problem — Split into chunks

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 — Split into chunks

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 — Split into chunks

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.

Check: Move from C=4 to C=2 on these 8 tokens: how many intermediate states must be handed off, and do outputs o₅..o₈ change value? Justify with the numbers from the trace.

The problem — Causal triangle

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 — Causal triangle

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 — Causal triangle

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.

Check: In the chunk-1 mask, entry (row 3, column 4) is 0 while k₃·k₄ = 0 as well. Do these two zeros have the same cause? What would happen if k₃·k₄ were 1?

Visual support — Causal triangle

        column j (token read)
          1  2  3  4
    row  ┌1  .  .  .┐    “ . ” above the diagonal:
    i    │0  1  .  .│    forbidden — that is the future
 (reading│1  0  1  .│
  token) └0  1  0  1┘    “ 0 ” below: allowed,
                          just a zero dot product

  two visually similar zeros, two distinct causes

The problem — Incoming and outgoing state

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 — Incoming and outgoing state

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 — Incoming and outgoing state

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.

Check: Chunk 2 receives S₄ = [[3,4],[5,5]] and nothing else from the past. Rebuild o₅ = [3,5] separating the incoming term from the local term, then say what chunk 2 can no longer know about v₁ and v₃.

The problem — Chunk size

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 — Chunk size

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 — Chunk size

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.

Check: You double C from 64 to 128 on a GPU whose SRAM is already full. The scratch triangle scales as C²: by what factor does it grow, and why can throughput fall even though there are fewer transitions?

Visual support — Chunk size

   C      transitions (8/C)     triangle C²
   1             8                    1
   2             4                    4
   4             2                   16
   8             1                   64

 transitions ÷2  ⇔  triangle ×4
 the optimum: the largest C whose triangle still fits in SRAM

Worked case — full trace

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.

8 tokens, d=2, chunk C=4, q_t = k_t, S₀ = [[0,0],[0,0]]
  chunk 1: k=[1,0],[0,1],[1,0],[0,1]   v=[2,3],[5,1],[1,1],[0,4]
  chunk 2: k=[1,0],[0,1],[1,0],[0,1]   v=[0,1],[2,2],[1,0],[0,0]

── CHUNK 1 ───────────────────────────────────────────────────────
  Q Kᵀ masked to the lower triangle (i reads j ≤ i):
      ┌1 0 0 0┐
      │0 1 0 0│      the zeros above the diagonal are
      │1 0 1 0│      causality, not a numerically null value
      └0 1 0 1┘
  O = M·V + K·S₀ , with S₀=0:
    o₁=[2,3]  o₂=[5,1]  o₃=v₁+v₃=[3,4]  o₄=v₂+v₄=[5,5]
  outgoing state  S₄ = K₁ᵀV₁ = [[3,4],[5,5]]

── CHUNK 2 (receives S₄, never sees v₁..v₄ individually) ──────────
  inter-chunk term: q₅ᵀS₄ = [3,4];   q₆ᵀS₄ = [5,5]
    o₅ = [3,4] + v₅        = [3,4]+[0,1] = [3,5]
    o₆ = [5,5] + v₆        = [5,5]+[2,2] = [7,7]
    o₇ = [3,4] + v₅+v₇     = [3,4]+[1,1] = [4,5]
    o₈ = [5,5] + v₆+v₈     = [5,5]+[2,2] = [7,7]
  S₈ = S₄ + K₂ᵀV₂ = [[3+1,4+1],[5+2,5+2]] = [[4,5],[7,7]]
  ✅ o₅..o₈ are EXACTLY those of the token-by-token recurrence

── ❌ BUG 1: dropping the incoming-state term ─────────────────────
  o₅ = v₅ = [0,1] instead of [3,5]: chunk 2 lost the whole prompt

── ❌ BUG 2: full mask instead of the triangle ────────────────────
  o₅ would include v₇: [3,4]+[0,1]+[1,0] = [4,5] ≠ [3,5]
  token 5 reads token 7 → future leakage, invisible during prefill,
  unreproducible at decode: perplexity drops “suspiciously well”

SHAPE CHECK: K(4×2)·Kᵀ(2×4) → M(4×4); M(4×4)·V(4×2) → O(4×2);
K(4×2)·S(2×2) → (4×2). Triangle scratch memory = C² = 16, not 4.

Choosing C: the same computation, two hardware regimes

Criterion (8 tokens) Small chunk, C=1 Large chunk, C=8
Sequential state transitions 8 (no parallelism) 1 (a single state hand-off)
Triangle scratch memory C² = 1 per chunk C² = 64 per chunk
Shape of matmuls sent to the GPU vector × matrix, units idle matrix × matrix, units saturated
Regime where it is the right pick decode: one token at a time prefill: full prompt available

Causal lab

Predict → change one variable → run → explain the delta

/interactives/curriculum/chunk-gate-memory.html?lang=en

Common errors

“Chunking is an approximation: you trade a little accuracy for speed.”

The trace contradicts it digit for digit: o₅..o₈ are [3,5],[7,7],[4,5],[7,7] with or without chunks. It is an exact algebraic reordering. What changes is memory cost (C²) and transition count — never the result.

“A bigger chunk is always faster, since there are fewer sequential steps.”

The scratch triangle grows as C²: from C=64 to C=128 it is multiplied by 4. Once it spills out of SRAM, throughput collapses despite fewer transitions. The optimum is a hardware fact, not a mathematical one.

Boundary, evidence, and sources

Chunking improves execution; it does not automatically change the information capacity of state.

Evidence status: Established mechanisms; numerical simplifications are pedagogical.

  • 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.

Transfer challenge

Redo the full trace with C = 2: chunks {t1,t2}{t3,t4}{t5,t6}{t7,t8}.

  1. Before computing: how many state hand-offs, and of what size each?
  2. Compute S₂, S₄, S₆, then check o₅..o₈ against the C = 4 trace.
  3. Conclude: what changed, and what is not allowed to change?

Synthesis and exit ticket

  • Prefill and decode
  • Naive recurrence
  • Split into chunks
  • Causal triangle
  • Incoming and outgoing state
  • Chunk size

Ticket: mechanism · trace · observation · boundary · evidence · next experiment

Instructor notes: Frame the problem before naming the mechanism. Collect an initial prediction and retain it for the exit ticket.

Instructor notes: Connect every step to the next with a causal verb. Flag any merely decorative arrow.

Instructor notes: Time a real chat by show of hands: “the first word takes two seconds, the next ones thirty milliseconds — why?”. The lived-experience question primes the beat better than the prefill/decode vocabulary.

Instructor notes: Have them mentally time both phases: “how many tokens in parallel at prefill? how many at decode?”. Until a learner states that asymmetry out loud, the rest of the session sounds like a gratuitous optimization.

Instructor notes: Answer: the triangle becomes 1×1, i.e. trivial — only the state read O = qᵀS plus the local term remains. The whole chunkwise machinery is a prefill optimization; decode gains nothing from it. Expected wrong answer: “a 4×4 with padding”.

Instructor notes: Act out the chain: eight learners, each may compute only after receiving the previous one’s slip. Point at the seven who are waiting: “that is the GPU”. The image lasts the whole session.

Instructor notes: Ask someone to perform all 8 updates out loud, one at a time. The felt tedium does the teaching better than any throughput chart.

Instructor notes: Answer: the state updates form a chain — parallelizing them as-is would change the result; but one at a time, each step offers only a vector-matrix product and the matrix units sit idle. Nuance to demand: the GPU is not “slow”, it is idle.

Instructor notes: Have the ▶ arrows counted in each regime: 7 versus 1. Then ask what the count becomes for 8,000 tokens and C = 64 — dividing the critical path is the only thing chunking buys.

Instructor notes: Ask: “what would it take for the right half of the board to work at the same time as the left?”. The spontaneous proposals — copy the past? resend everything? — set up the value of a single state S₄.

Instructor notes: Physically split the board into two zones and forbid the right half from looking left except through a sticky note labeled “S₄”. That spatial constraint is exactly what the code does.

Instructor notes: Answer: chunks {1-2}{3-4}{5-6}{7-8} → hand off S₂, S₄, S₆, three passes instead of one; o₅..o₈ = [3,5],[7,7],[4,5],[7,7], unchanged — chunking is exact, only the execution changes. Expected wrong answer: “finding” imaginary rounding differences.

Instructor notes: Set the trap before the solution: “in a block computed all at once, what stops t5 from reading t7?”. Honest answer: nothing — except the mask. Let the discomfort settle before showing the triangle.

Instructor notes: Erase one zero from the triangle and have the group hunt the bug. The expected answer is not “the number is wrong” but “token 5 read the future, and training will happily reward it”.

Instructor notes: Answer: different causes — (3,4) sits above the diagonal, forbidden by causality; k₃·k₄ = 0 is a COMPUTED zero that was allowed to be nonzero. If k₃·k₄ were 1, cell (4,3) — below the diagonal — would become 1, but (3,4) would stay masked. Expected wrong answer: “it is zero everywhere, same thing”.

Instructor notes: Have the two families of zeros colored differently before stating the rule. The check “do these two zeros share a cause?” returns exactly here — this slide is its preparation.

Instructor notes: Write S₄ = [[3,4],[5,5]] on the board and ask: “v₁ was [2,3] — where is it?”. The silence is the lesson: it is in there, but no longer separable. An explicit callback to session 13’s superposition.

Instructor notes: Have them write S₄ on a slip, flip the chunk-1 board over, then ask for o₅. What they spontaneously demand is exactly the information the state must carry.

Instructor notes: Answer: o₅ = q₅ᵀS₄ + v₅ = [3,4] + [0,1] = [3,5] — inherited term plus local term. Chunk 2 can no longer know HOW [3,4] decomposes into v₁ and v₃: attribution is lost at the boundary (superposition, session 13). Expected wrong answer: believing v₁ and v₃ are recoverable “somewhere” in S₄.

Instructor notes: Poll: “chunk of 4, of 64, of 4,096: any advance?”. Have them vote for a C before exposing the trade-off — voting forces everyone to pick a criterion, and the diverging criteria fuel the beat’s debate.

Instructor notes: Have them compute C² for C = 16, 64, 128 and compare against a stated SRAM budget. Close on “there is no good C, there is a good C for this GPU” — the only honest conclusion.

Instructor notes: Answer: C² goes from 4,096 to 16,384 — ×4 for one doubling. If SRAM overflows, the triangle spills into slow memory and every access costs an order of magnitude more: throughput drops despite half the transitions. Expected wrong answer: reasoning in FLOPs alone, without the memory hierarchy.

Instructor notes: Have the table extended to C = 128 with a fictional SRAM budget (e.g. 8,192 cells): the group must find on its own which C overflows. The final rule then states itself.

Instructor notes: Walk line by line. Locate an inconsistency at the first faulty step, not only on the final line.

Instructor notes: Have learners fill the final row before revealing it: that trade-off is what decides in production.

Instructor notes: Retain initial and final values. Do not allow simultaneous changes that make the delta impossible to attribute.

Instructor notes: For each claim, have the group produce the smallest counterexample before giving the correction.

Instructor notes: Separate verifiable mechanism, reported implementation choice, and experimental result. Evidence precision must match claim precision.

Instructor notes: Answers: three hand-offs (S₂, S₄, S₆), each 2×2 — state size does not depend on C. S₂ = [[2,3],[5,1]]; S₄ = [[3,4],[5,5]] (identical to the trace); S₆ = [[3,5],[7,7]]. o₅..o₈ = [3,5],[7,7],[4,5],[7,7] — unchanged: chunking is exact (common error 1). What changes: transition count and triangle sizes (four 2×2 instead of two 4×4). Expected misconception: hunting for rounding differences. Twelve minutes, pairs.

Instructor notes: Rebuild the chain without looking at the slides, then fill the ticket in at most six lines. Compare with the opening prediction and name what actually changed.