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