🚧 Week 3 Day 2: Chunked Prefill

🚧 This chapter is under review and may change.

A long prompt can monopolize the device while active decode requests wait for their next token. Chunked prefill gives each scheduler iteration a prompt-token budget, limiting how long decode work can be delayed.

The scheduler policy becomes:

admit at most prefill_max_step prompt tokens
decode one token for every active request
repeat until the queue and active batch are empty

Task 1: Bound Prefill Work

Update Request.try_prefill in src/tiny_llm/batch.py to select one prompt slice, call the model with the slice’s absolute offset, and mark the request ready only after the full prompt has been processed.

for start in range(0, len(prompt_tokens), prefill_max_step):
    chunk = prompt_tokens[start : start + prefill_max_step]
    model(chunk, offset=start, cache=cache)

The final chunk may be smaller than the configured budget. Test prompts shorter than one chunk, exactly one chunk, and one token longer than a chunk.

Day 2’s supplied tests are named by behavior rather than task number. This one checks chunk boundaries against one-shot prefill, together with Task 3’s materialization:

pdm run test --week 3 --day 2 -- -k chunk_boundaries

With chunked prefill in place, the Day 1 loop already interleaves: a long prompt advances one chunk per iteration while active requests keep decoding. The test for that order, -k active_decode_advances, also checks that every request cache is released, so it turns green with Task 4.

Task 2: Build Rectangular Causal Masks

When a cache already holds S - L tokens and a chunk contributes L new tokens, the mask is L x S. Every query can attend to the old prefix and to earlier positions in its own chunk.

For a five-token prefix and a three-token chunk, the mask is 3 x 8:

0  0  0  0  0  0  -inf  -inf
0  0  0  0  0  0     0  -inf
0  0  0  0  0  0     0     0

Use the absolute cache offset for RoPE and S - L as the causal diagonal offset. Chunked prefill should then produce the same logits as one-shot prefill. The focused test below does not run the model; it checks the mask itself. causal_mask(L=2, S=5) must keep the earlier positions visible, and attending through that mask must give the final query the same result as attending over the whole sequence at once:

pdm run test --week 3 --day 2 -- -k rectangular_mask

Task 3: Materialize Between Chunks

MLX is lazy. Extending an unevaluated cache repeatedly creates a long graph and can grow memory usage. Call each layer cache’s materialize() hook after every chunk so the next scheduler iteration starts from materialized state. For the dense cache, implement TinyKvFullCache.materialize to evaluate its key/value tuple. The paged cache that arrives on Day 3 will evaluate its page pool storage without first gathering it into a dense tensor.

The hook is part of the cache lifecycle rather than the scheduler’s storage logic. This lets the scheduler use dense and paged caches without inspecting their internal representation.

Task 4: Release Every Request Cache

Each request owns one cache per layer from the moment Request is constructed. The scheduler must release each of those caches exactly once, however the request ends:

• A request that completes during prefill, from an immediate EOS or by reaching max_seq_len, never enters a slot: release its caches directly.
• A request that completes in a slot is released by remove_request(i), so your BatchingKvCache.remove_request must call the request cache’s release() before clearing the slot.
• A prompt longer than max_seq_len never gets caches: Day 1 already rejects it before create_kv_cache() is called, and its Task 4 test still requires that.
• If prefill, materialization, decoding or detokenization raises, the error should still reach the caller, but first release every cache that is still live, whether it belongs to the pending request or to a slot. Wrap the scheduler loop in try/finally for this. A cache can be referenced from both places while a request is being admitted, so release each cache object only once.

These tests count live caches through normal completion, an oversized prompt and an injected failure at each of those four points. Then rerun the interleaving test from Task 1, which also checks release:

pdm run test --week 3 --day 2 -- -k "release or oversized"
pdm run test --week 3 --day 2 -- -k active_decode_advances

Task 5: Measure the Fairness Tradeoff

Run the same request trace with several prefill_max_step values. Report total throughput and the longest interval between consecutive decode steps. Smaller chunks usually improve fairness but add scheduler and launch overhead; choose a default from the measured tradeoff rather than treating one chunk size as universal.

week3_result_root="$HOME/tiny-llm-week3-results"
mkdir -p "$week3_result_root"
week3_result_dir="$(mktemp -d "$week3_result_root/run-XXXXXX")"

pdm run test --week 3 --day 2
pdm run batch-main --solution tiny_llm --loader week2

pdm run bench-chunked-prefill --solution tiny_llm --offline --model qwen3-0.6b \
  --prefill-steps 32 128 512 --num-seqs 8 --batch-size 4 \
  --min-input-len 64 --max-input-len 512 \
  --min-output-len 32 --max-output-len 32 \
  --warmup 1 --repeats 4 --cooldown-seconds 1 \
  --json-output "$week3_result_dir/chunked-prefill-tiny-llm.json"

This command measures your tiny_llm solution and writes a new JSON file outside the tracked historical corpus. For a current reference control, rerun the same flags with --solution ref and a different filename in the new directory; compare those new measurements with your own run. The table below is historical task #367 evidence from source 18aec8503929d80c986324578068ecac2463c2ac, before the five-day Week 2 revamp. Its trace used seed 0 and the same 32-token output budget for every request. Each chunk size ran twice in forward order and twice in reverse order in fresh processes. Every row used the same canonical Week 3 MLX quantized-projection seam and the same course-owned scheduler, dense cache, and attention code; only the prefill budget changed. The saved JSON stores every prompt token id, the per-request output budget, and their canonical SHA-256 checksum. Running today’s reference can produce different values; it does not reproduce the predecessor table.

A decode-completion gap is the wall-clock interval between two consecutive synchronized decode calls while at least one decode request remains active. It therefore includes intervening prefill and scheduler work; idle time with no decode request is excluded. On the measured M4 Pro, the four-process medians were:

The 512-token row was the full-prompt Day 1 control for that trace. Relative to that row, the 128-token budget gives up 9.8% output throughput while reducing the p95 completion gap by 38.3% and the maximum gap by 55.2%. The 32-token budget reduces the p95 gap further but gives up substantially more throughput. For that predecessor workload, 128 was a measured compromise, not a universal optimum for your current run. The ledger at benchmark_results/task367-final-main/task367-final-main-benchmark-ledger.md keeps this predecessor absolute result separate from task #360’s causal projection-seam ablation.

Your feedback is greatly appreciated. Join our Discord community.
Found an issue? Open an issue or pull request at github.com/skyzh/tiny-llm.
tiny-llm-book © 2025 by Alex Chi Z is licensed under CC BY-NC-SA 4.0.