SIGN IN SIGN UP

fix(generate): chunk prefill for prompts shorter than prefill_step_size (#2119)

* fix(generate): chunk prefill for prompts shorter than prefill_step_size

Prefill only chunked when the prompt was longer than prefill_step_size.
The chunk loop discards its output, so MLX never evaluates the
[B, N, vocab] logits; the unchunked path instead hands the whole prompt
to _step, which reads logits[:, -1, :] and materializes every row before
throwing all but one away.

That inverted the memory curve: on Qwen3.5-0.8B a 2047-token prompt
peaked at 2.896 GB while a 2049-token one peaked at 1.965 GB. Two extra
tokens cut peak memory by a third, and most chat turns land on the
expensive side of the boundary. Chunk whenever more than one prompt
token remains.

Measured savings, greedy decode, output unchanged:

  Qwen3.5-0.8B   512 tok  1.491 -> 1.286 GB
  Qwen3.5-0.8B  1024 tok  2.181 -> 1.746 GB
  Qwen3.5-0.8B  2047 tok  2.898 -> 1.973 GB
  Qwen3.8-27B   1229 tok 17.505 -> 17.013 GB

Savings scale with vocab size and prompt length, so large-vocab models
benefit most. Verified across 17 local model families (qwen2/3/3_5,
qwen2_vl, qwen2_5_vl, qwen3_vl, gemma2, gemma3, llama, lfm2, lfm2_vl,
lfm2_moe, ernie4_5, smollm3, phi3, olmoe, llava_next) with identical
generated text in every case.

PaliGemma is excluded: it builds a prefix-LM mask across the whole
prompt, so causal chunking changes its output. It gains a
chunked_prefill_policy that opts out while bidirectional attention is
on, using the hook gemma4 and lfm2 already use. This was a latent bug
for prompts over prefill_step_size, which already chunked.

* test: adapt kv-cache quantization mock to chunked prefill

MockCacheLayer now exposes .state (used by the chunked prefill eval), and the
quantize-call count accounts for one call per prefill chunk in addition to the
per-token calls.
A
Alazer Manakelew committed
34edc23f799bf6e320d119da0de5473c696d5706
Parent: ef8191e
Committed by GitHub <noreply@github.com> on 8/31/2026, 6:49:52 PM