fix(core): rewind the pool in KVCache::trim for pool-backed caches

[rapsealk]

Summary

KVCache::trim returned 0 for a pool-backed cache, so the pad rows a tile-padded prefill writes on M5 stayed in the pool and offset stayed at the padded length. Decode attended over the pad rows and rotated the next token at the padded position, which made the default (paged) server return wrong greedy output for any prompt whose length is not a multiple of 32. trim now rewinds the pool block table and moves offset with it, and the decode lookahead teardown routes through the same call.

Related issues

Closes #2096. Sibling of #1755 (the model-owned variant of the same pad trim). Refs #1760 for the replay abort described under "Not fixed here".

Type of change

• feat — new user-visible feature
• fix — bug fix
• perf — performance improvement (include before/after numbers in the PR body)
• refactor — internal restructuring without behavior change
• chore — build, CI, dependencies, release infrastructure
• docs — documentation only
• test — tests only

What changed

• src/lib/mlxcel-core/src/cache.rs: for a pool-backed cache, KVCache::trim calls PagedBlockPool::rewind_tokens through the cache's own backing and subtracts the removed count from offset. The four pad trims in scheduler/prefill.rs and every speculative rollback already call trim and expect a real one; none of them pairs it with a pool rewind.
• src/server/batch/scheduler/decode_tick.rs: apply_lookahead_trim drops its pool-only branch and trims every cache through KVCache::trim. The branch rewound the block table with rewind_paged_tokens and left cache.offset, which is where BatchedAttentionMetadata takes the RoPE offset from, one ahead.
• src/lib/mlxcel-core/src/cache/paged_batch_decode_tests.rs: trim_drops_padded_prefill_rows_from_a_pool_backed_cache, checkpoint-free. It writes 64 rows, trims 27, and asserts the returned count, offset, the pool len, and that the next update_and_fetch returns the 37 real rows plus the new token at row 37. On the pre-fix code it fails at the first assertion (left: 0, right: 27).

Real-checkpoint validation

Apple M5 Pro, macOS 27.0, temperature 0, max_tokens 24, five prompts per run, fresh server per run. "default" is mlxcel-server -m <model> with no other flags; the reference is --decode-storage-backend dense.

Checkpoint	Before, default vs dense	After, default vs dense	After, 5 concurrent vs dense
mlx-community/Llama-3.2-1B-Instruct-4bit	0 of 5 identical	5 of 5	5 of 5
HuggingFaceTB/SmolLM2-135M-Instruct	0 of 5	5 of 5	5 of 5
mlx-community/Qwen3-0.6B-4bit	1 of 5	5 of 5	5 of 5

• "Before" is the v0.7.0 release binary. Unpatched main at 9a0be04 is wrong in the same way on the prompts I repeated there: SmolLM2 returns the same text as below, and Llama-3.2-1B answers Repeat: one two three with One!Two!Three! (dense: One two three.).
• The concurrent column was run twice per checkpoint.
• With only the cache.rs change, concurrent was 4 of 5 on Llama-3.2-1B and SmolLM2. On Llama-3.2-1B it was 5 of 5 with MLXCEL_FORCE_SYNC=1, and 5 of 5 once the lookahead teardown used the same trim.
• Qwen3 is a weak data point: at 24 tokens its dense content is mostly empty because it is still reasoning.

SmolLM2-135M, What is the capital of France?, default flags:

before: 'The concept of "The Greatest Thing That\'s Happened to me\n\nI\'m an AI assistant, and I'
        paged decode v2: gather: 65 visible KV tokens across 1 request(s) ...   (37-token prompt)
after:  'The capital of France is Paris.'     (three repeats, one streamed, same text)
        paged decode v2: gather: 38 visible KV tokens across 1 request(s) ...

Test plan

• cargo fmt --all -- --check
• cargo clippy --workspace --all-targets --features metal,accelerate -- -D warnings
• cargo test --workspace --profile test-fast --features metal,accelerate --no-fail-fast -- --test-threads=1: 11818 passed, 4 failed. The same 4 fail on unpatched main in this environment, see below.
• cargo deny check (not installed on this host; not run)
• Validated with a real checkpoint: the three above
• If the change moves the numbers: not traced. The change touches no kernel or arithmetic; it removes rows that should not have been in the attention window.

The 4 failures, identical before and after this change:

• fp8_block_requantize_round_trip_stays_within_half_an_e4m3_step, mxfp4_gather_qmm_matches_host_reference, mxfp8_gather_qmm_matches_host_reference. This host has Command Line Tools only and no Metal shader compiler, so the main build linked the v0.7.0 release mlx.metallib instead of compiling the kernels at main's MLX pin. These three exercise kernels that changed between the two pins.
• speculative::prompt_lookup::tests::warmup_runs_every_verify_width_and_rolls_each_back records a first forward width of 32 instead of 5 on M5. It uses a dense KVCache::new(), which this change does not touch.

Because of that build caveat I also applied the same change to a v0.7.0 source build, where the release metallib matches exactly, and it fixes the same reproduction there.

Notes for reviewers

• Not fixed here: replaying a prompt whose length is a multiple of 16 still aborts the server on M5 ([broadcast_shapes] Shapes (32,63) and (1,9,32,64) cannot be broadcast). That is the whole-prompt hit of fix(prompt-cache): a whole-prompt cache hit re-runs the last token on top of a cache that already holds it #1760; the evidence is in fix(server): pool-backed KVCache::trim is a no-op, so tile-padded prefill leaves pad rows in the cache and default M5 serving returns wrong greedy output #2096 under "Related".
• Not measured: hardware older than M5, CUDA, and the model-owned families. The batched-prefill pad trim (prefill.rs:700) pads to the longest row in the window on any hardware, so this should also change output there for concurrent prompts of different lengths under the paged backend; I have no such host.
• CachePool::rewind_paged_tokens has no caller left outside tests. I left it in place.
• No scheduler-level test for the lookahead teardown: the dense-natural scheduler fixtures need a checkpoint under models/. The unit test pins the trim contract both call sites now share.
• CI: the first run failed only in WebUI installed artifact. Its activity-performance gate reported two-visible as investigate because the baseline CV was 5.10% against a 5% budget; the median decode degradation was -0.25%. The same gate failed on main in runs 36961726625 (baseline CV 6.2%) and 36735057944 (degradation 2.7%). I could not re-run it from a fork, so I merged main (c0b7134) into the branch; the run on that head is green.

Checklist

• PR title uses a Conventional Commits prefix (feat:, fix:, etc.)
• One logical change per PR (split unrelated changes)
• Updated docs/ if user-facing behavior or supported models changed (no flag or documented behavior changes)
• Updated // Used by: ... comments on any shared function I modified (see docs/code-guidelines.md)
• No secrets, credentials, or .env files committed

🤖 Generated with Claude Code