Faster loglikelihood scoring in mlx_lm.evaluate

[markjouh]

Two changes to the loglikelihood path in mlx_lm.evaluate, plus a bug fix found along the way:

1. Skip the vocabulary projection for the final prompt chunk in _process_prompt. The last prefill chunk previously materialized logits for up to step_size positions (a [1, 2048, vocab] tensor) when only the final position's log probabilities are used. When the final chunk is long, the prompt is now processed the way generate_step handles it: everything before the last token is prefilled evaluating only the cache state (the logits are never computed, thanks to lazy evaluation), then the last token is processed by itself so the vocabulary projection runs for one position instead of up to 2048. For prompts whose final chunk is shorter than 512 tokens, the extra model call costs more than the projection it avoids (measured crossover is ~150–250 tokens on M-series), so those keep the previous single-call behavior and are bit-identical to main.

2. Rewind the prompt cache between continuations instead of deep-copying it. When the cache is trimmable, all continuations of a question are scored with the same cache, using trim_prompt_cache to rewind in between. This is bit-exact — scoring reads identical cache contents either way — and saves a cache copy plus a buffer-growth reallocation per continuation. Models with non-trimmable caches (SSM states, sliding-window caches past the window) keep the copy.deepcopy behavior.

3. Fix a latent NameError in the truncated-prompt branch (all_scores / all_is_greedy → scores / is_greedy), which crashed evaluation whenever a completion was longer than the context limit. Same shape as the loglikelihood_rolling fix in Fix NameError in loglikelihood_rolling method #339.

Benchmarks

M5 Max (40-core GPU, 128 GB), mlx 0.32.0, macOS 15. Paired A/B — the two implementations alternate per round in fresh processes, N questions × 4 continuations, median over rounds:

model	context tokens	speedup vs main
Llama-3.2-1B-Instruct-4bit	1500	1.27× (faster in 4/4 rounds)
Llama-3.2-1B-Instruct-4bit	8192	1.07× (4/4)
Qwen3-4B-4bit	1500	1.10× (4/4)
Qwen3-4B-4bit	120 (below split threshold)	1.02× (4/4)

The gain comes from dropping the [chunk, vocab] projection, so it is largest when the vocabulary is a big share of the model (small models, large vocabs) and for prompts just under a chunk boundary. Peak memory is unchanged (the prefill working set dominates it).

End-to-end mlx_lm evaluate --tasks arc_challenge --num-shots 10 --limit 150 (prompts long enough to take the new path): accuracy and normalized accuracy are unchanged for both models, wall time 13.1s → 12.6s (1B) and 27.0s → 26.6s (4B), which includes model load and harness overhead.

Correctness

• Prompts below the split threshold (all the 0-shot configs I ran) take an unchanged code path; verified identical metrics on arc_challenge + hellaswag.
• For long prompts, the last prompt token is now computed exactly the way generation computes it, so scores shift within half-precision chunk-boundary noise: on 10-shot arc_challenge (150 questions × 4 choices), max per-continuation |Δlogprob| was 0.07 for the fp16 1B model and 0.84 for the bf16 Qwen3-4B (multi-token continuations). All aggregate metrics were unchanged in every CLI configuration tested. The only per-question flips I could find (under a different fewshot formatting) were questions whose top-2 choice margins were already below that noise floor (0.04 and 0.19 respectively).
• The cache rewind is bit-exact: scoring the same continuations in forward and reversed order gives identical results (delta exactly 0.0), which also rules out cross-continuation cache contamination.
• Verified both cache strategies on gemma-3-1b-it-4bit: contexts below the sliding window use the rewind path, contexts past the window correctly fall back to deep copies; both agree with a full-sequence forward pass reference.
• Edge cases: length-1 and length-2 prefixes, chunk-boundary prompt lengths, and the previously crashing fully-truncated-prompt branch.
• New end-to-end tests in tests/test_evaluate.py: score agreement with a full-sequence forward pass, and continuation order invariance, over both short and long contexts. Both tests also pass on main.

[zcbenz]

Makes no sense.