prefill: configurable chunk size up to 4096, with --prefill-chunk auto

[dwijenpatel]

Before/after on the community long-synthesis case (Gemma 4 26B-A4B,
M5 24 GB, protocol settings): total request wall time 86.1 s -> 60.1 s,
prefill disk reads 205 GB -> 50 GB, byte-identical output, peak process
RSS unchanged. One flag: --prefill-chunk auto. Default behavior is
unchanged: without the flag nothing changes.

Why prefill re-reads so much

Chunked prefill runs two nested loops: for each 128-token chunk, for each of
the 30 layers, read that layer's routed experts and run the chunk through.
The expert reads in the inner loop are the cost. A 3,015-token prompt is 24
chunks, and each chunk touches most of the expert pool per layer, so the
same expert bytes stream from disk ~24 times: 182 GB of prefill reads for a
model whose packed experts total ~12 GB. Prefill time is almost entirely
that I/O (the disk sits saturated for the whole phase).

Read volume is proportional to the number of chunks, so the fix is fewer,
larger chunks. With a chunk that covers the whole prompt, every layer's
experts are read once per prefill: one sweep of the file, the I/O floor.
Prompts up to 4,096 tokens cover today's product context ceiling.

Measured chunk-size sweep (Gemma 4 26B-A4B, this branch, long-synthesis
prompt with --max-new 8 so the numbers are prefill-dominated; disk-active
seconds from 1 Hz iostat):

chunk	prefill disk reads	disk-active time
128 (current)	182 GB	55 s
512	60 GB	33 s
1024	32 GB	28 s
2048	22 GB	26 s
4096	14 GB	26 s

14 GB is roughly the install size (13 GB): one sweep, the floor for chunked
prefill. The time curve flattens once the phase stops being I/O-bound, and
4,096 covers the product context ceiling, hence the new cap.

What changed

• RuntimeConfiguration.allowedPrefillChunkTokens: [32, 64, 128] -> [32 ... 4096]
• PrefillRuntimeConfig.maxChunkTokens: 128 -> 4096 (doc comment records the measurement)
• PrefillChunkScratch: clamp follows maxChunkTokens instead of a hard 128
• KVCacheManager: the sliding-window ring is sized from the configured
  chunk, not the static cap, so memory grows only when a larger chunk is
  opted into
• CLI: --prefill-chunk <n|auto>, validated against the allowed set;
  auto = smallest allowed size covering the prompt; default 128
• docs/RUNTIME_CONTROLS.md: new control documented
• Tests: parameterized coverage widened to all sizes, 4 new CLI-parsing
  tests, a scratch-scaling test; two existing tests updated because they
  froze the old 128 cap as a contract

Evidence

Community benchmark protocol (warmup then fresh measured process, per-case
seeds, --max-new 1024 --max-context 4096, sampling defaults, no other
model processes), plus iostat disk totals and /usr/bin/time -l peak RSS
on the long case. Experiment shape per RUNTIME_CONTROLS: baseline, then one
changed control.

Machine: MacBook Pro (Mac17,2), Apple M5, 10 cores, 24 GB, internal SSD.
macOS 26.5.2 (25F84). Swift 6.3.3 (swift-driver 1.148.6). Model installed
per README; manifest sha256 1cb53c24...; prompt hashes match the frozen
real-generation-v1 set.

Gemma 4 26B-A4B (this branch, f9b9b47 + runtime-controls doc):

case	arm	timing footer	wall	prefill I/O	peak RSS	output
short-explanation	default	prefill=61tok new=497tok decode=22.33s tok/s=22.25	29.7 s	—	1554 MiB	—
short-explanation	auto	prefill=61tok new=497tok decode=21.32s tok/s=23.31	28.6 s	—	1611 MiB	byte-identical
medium-review	default	prefill=430tok new=721tok decode=34.01s tok/s=21.20	46.0 s	—	1617 MiB	—
medium-review	auto	prefill=430tok new=721tok decode=34.10s tok/s=21.14	42.6 s	—	1605 MiB	byte-identical
long-synthesis	default	prefill=3015tok new=610tok decode=30.75s tok/s=19.84	86.1 s	205 GB	1606 MiB	—
long-synthesis	auto	prefill=3015tok new=610tok decode=33.09s tok/s=18.44	60.1 s	50 GB	1578 MiB	byte-identical

Gemma is the interesting correctness case: it has sliding-window layers, so
the KV ring actually grows with the chunk size. Peak RSS stays flat because
the ring is sized from the configured chunk (opt-in memory).

Memory accounting. Peak process RSS does not include device-private
Metal allocations, so to be explicit: prefill scratch
(PrefillChunkScratchLayout.totalPersistentBytes, all MTLBuffer, no model
tensors in Swift heap) grows from ~16 MiB at chunk 128 to ~490 MiB at chunk
4096 on Gemma's dimensions. That cost exists only when a larger chunk is
explicitly requested, and auto never picks a chunk larger than the prompt
needs.

Protocol change: none. Outputs are byte-identical per case at fixed
seed; the bounded-memory model path is preserved (experts still stream via
pread, nothing new touches Swift heap).

Checks: swift build -c release clean, Scripts/test.sh 520/520,
ruby Scripts/check_markdown_links.rb all 22 files resolve.

Limitations

• Prompts longer than 4,096 tokens still pay the multi-chunk re-read above
  the cap.
• auto resolves from the tokenized prompt length at request time; the Mac
  app keeps the default (CLI-only control, as documented).

Happy to split the CLI flag from the ceiling raise if you'd prefer them
separate.