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

Summary

On M5 class hardware, mlxcel-server returns wrong greedy output under its default flags for every dense family tested. Tile-padded prefill writes pad rows into the pool-backed KV cache and then calls KVCache::trim to remove them, but trim is a no-op for pool-backed caches. The pad rows stay visible to decode and offset stays at the padded length, so the next token is rotated at the wrong RoPE position and attends over up to 31 pad rows. Setting --decode-storage-backend dense or --parallel 1 (both of which use dense caches, where trim is a real rewind) produces correct output, which localizes the defect to the pool-backed trim path.

Background

Verified statically against main; the project cannot be built or run in this triage environment (no Metal/MLX), so the runtime reproduction in the report below is taken as given and only the code paths it cites were checked.

• KVCache::trim (src/lib/mlxcel-core/src/cache.rs:2341) returns 0 immediately when self.paged_backing.is_some() (lines 2357-2359), without rewinding the pool block table or decrementing self.offset. The comment states the dense-side slicing is skipped for pool-backed caches, but no pool-side rewind is substituted, so the call removes nothing.
• should_align_prefill() (src/server/batch/scheduler/mod.rs:104) is true when the hardware has a Neural Accelerator and macOS supports it (M5 and newer), which turns on tile-aligned prefill via align_to_na_tile and a padded causal mask (prefill.rs:584, :790, :960, :1195).
• Each prefill path pads to a 32-token tile and then tries to drop the pad positions with c.trim(excess) (src/server/batch/scheduler/prefill.rs:700, :871, :1037, :1249). On a pool-backed cache this trim does nothing, so the pad rows and the inflated offset persist into decode.
• apply_lookahead_trim (src/server/batch/scheduler/decode_tick.rs:435) special-cases pool-backed sequences and rewinds the pool with rewind_paged_tokens (line 453) while leaving cache.offset one ahead, which is the same offset-versus-pool split in the lookahead teardown.
• The batched-prefill site pads every row to the longest prompt in the window regardless of hardware, so concurrent prompts of different lengths can leave pad rows behind even without an M5. This widens the blast radius beyond the M5 default case.

Proposed Solution

Make KVCache::trim perform a real trim for pool-backed caches instead of returning early: rewind the pool block table through the cache's paged_backing (the pool API the report calls PagedBlockPool::rewind_tokens) and decrement self.offset by the number of tokens actually removed, matching what every existing caller already expects. Route apply_lookahead_trim through the same code path and remove its pool-only branch so pool and offset stay consistent in both the prefill pad-drop and the lookahead teardown.

Implementation Notes

• src/lib/mlxcel-core/src/cache.rs:2341-2359: replace the if self.paged_backing.is_some() { return 0; } early return with a pool rewind plus self.offset -= n, keeping the existing live_len() clamp (line 2348) so a prior trim_front cannot roll offset below live_start.
• src/server/batch/scheduler/prefill.rs:700, :871, :1037, :1249: no change needed if trim becomes authoritative; these callers already pass the pad excess and expect it to be removed.
• src/server/batch/scheduler/decode_tick.rs:435-487: drop the pool-backed special case and call the unified trim so the lookahead rewind also corrects cache.offset.
• Confirm no other caller depends on the current pool-backed no-op behavior (the report states nothing else pairs a trim with a separate pool rewind).

Acceptance Criteria

• After a padded prefill, a pool-backed cache has offset and pool len both equal to the real prompt length.
• On M5 class hardware, default flags and --decode-storage-backend dense return the same greedy continuation for prompts whose length is not a multiple of 32.
• The same equivalence holds for concurrent requests of differing prompt lengths.
• A checkpoint-free unit test pins the trim behavior on a pool-backed cache (offset and pool length after trim).

Original Suggestion

Title: fix(server): padded prefill leaves its pad rows in pool-backed KV caches, so default serving on M5 returns wrong greedy output

Summary

On M5 hardware mlxcel-server returns wrong greedy output under its default flags, for every dense-natural family tried. A tile-padded prefill writes its pad rows into the pool-backed KV cache and then calls KVCache::trim to drop them, and KVCache::trim is a no-op for a pool-backed cache (src/lib/mlxcel-core/src/cache.rs:2352-2359). The pad rows stay visible to decode and offset stays at the padded length, so the next token is rotated at the wrong position and attends over up to 31 pad rows.

Found while evaluating mlxcel-server as a serving runtime. Reproduces on the v0.7.0 release binary and on main at 9a0be04.

Reproduction

Any prompt whose token count is not a multiple of 32, on an M5 Mac, default flags (--parallel auto = 4, decode_storage=auto resolves to paged):

mlxcel-server -m mlx-community/Llama-3.2-1B-Instruct-4bit --port 8080
curl -s http://127.0.0.1:8080/v1/chat/completions -H 'Content-Type: application/json' \
  -d '{"model":"m","messages":[{"role":"user","content":"What is the capital of France?"}],"max_tokens":24,"temperature":0}'

# Reference: the same request against the dense backend.
mlxcel-server -m mlx-community/Llama-3.2-1B-Instruct-4bit --port 8080 --decode-storage-backend dense

Expected behavior

The default (paged) backend returns the same greedy continuation as --decode-storage-backend dense, --parallel 1, and mlxcel generate.

Actual behavior

v0.7.0 release binary, temperature 0, max_tokens 24, five prompts sent one after another to a fresh server. Identical across restarts.

On the larger checkpoint the output stays plausible, which is why it is easy to miss. On SmolLM2-135M it is unrelated text: What is the capital of France? returns The concept of "The Greatest Thing That's Happened to me … (dense, and mlx_lm: The capital of France is Paris.).

Flag matrix, SmolLM2-135M, same request:

Root cause

The first decode step logs the visible KV length, and it is the padded length plus one, not the prompt length plus one:

prefill{seq_id=seq-1 prompt_len=37 cached=0 start=0}
decode_step{batch_size=1}: paged decode v2: gather: 65 visible KV tokens across 1 request(s) is below the 4096-token dispatch floor

A 37-token prompt gives 65, a 41-token prompt 65, a 17-token prompt 33, a 31-token prompt 33.

1. should_align_prefill() (src/server/batch/scheduler/mod.rs:104) is true on M5 or newer, so the prefill paths pad the prompt to a 32-token tile with align_to_na_tile and a padded causal mask.
2. After the forward, each path drops the pad positions with c.trim(excess) (src/server/batch/scheduler/prefill.rs:700, :871, :1037, :1235).
3. KVCache::trim returns 0 for a pool-backed cache: "Treat a dense-side trim as a no-op for them" (cache.rs:2352-2359). The pool's len and the cache's offset both stay at the padded length.
4. Decode gathers all 64 rows and writes the next token at row 64. BatchedAttentionMetadata takes its RoPE offsets from cache.offset, so that token is also rotated at position 64 instead of 37.

--parallel 1 and --decode-storage-backend dense use dense caches, where trim is a logical rewind, so they are unaffected.

The same split shows up in the decode lookahead teardown. apply_lookahead_trim (src/server/batch/scheduler/decode_tick.rs:434-487) special-cases pool-backed sequences and calls rewind_paged_tokens, which rewinds the pool but leaves cache.offset one ahead. With the pad trim fixed, five concurrent requests still gave 4 of 5 continuations identical to dense; with MLXCEL_FORCE_SYNC=1 (no lookahead) it was 5 of 5.

By reading the code, not measured: the batched-prefill site (prefill.rs:700) pads every row to the longest prompt in the window on any hardware, so concurrent prompts of different lengths would leave pad rows behind without an M5 as well.

Proposed solution

Make KVCache::trim rewind the pool block table for a pool-backed cache (PagedBlockPool::rewind_tokens through the cache's own backing) and move offset by the count it removed. Every caller already expects a real trim, and nothing else pairs a trim call with a pool rewind. Route apply_lookahead_trim through the same call and drop its pool-only branch.

I have this change with a regression test and will open a PR.

Acceptance criteria

• After a padded prefill, a pool-backed cache has offset and pool len equal to the real prompt length.
• On an M5, default flags and --decode-storage-backend dense return the same greedy continuation for prompts whose length is not a multiple of 32.
• The same holds for concurrent requests.
• A checkpoint-free test pins the trim on a pool-backed cache.

Environment

• mlxcel version: v0.7.0 (release asset mlxcel-macos-aarch64.zip), and main @ 9a0be04 built from source
• Install method: pre-built release binary (GitHub Releases); built from source for main
• Platform: macOS (Apple Silicon), macOS 27.0, Apple M5 Pro, 48 GB
• Model checkpoints: mlx-community/Llama-3.2-1B-Instruct-4bit, HuggingFaceTB/SmolLM2-135M-Instruct, mlx-community/Qwen3-0.6B-4bit

Related

• feat(scheduler): sequence-aware cache trim hook so model-owned families can use padded prefill #1755 is the model-owned variant of the same class: the pad trim reaches nothing because the cache vector is empty. This issue is the pool-backed variant, where the trim reaches the cache and the cache ignores it.

• fix(prompt-cache): a whole-prompt cache hit re-runs the last token on top of a cache that already holds it #1760: the process abort that led me here is that issue's whole-prompt hit, and on M5 it does not only diverge. Replaying a prompt whose length is a multiple of the APC block size (16) aborts the server on the K/V path. A 32-token replay aborts on both backends, a 48-token replay under --parallel 1:

  prompt-cache hit: adopted 32/32 tokens seq_id=seq-2 matched=32 total=32
  prompt-cache hit covered the entire prompt; re-running the last token through prefill to produce a sampling logit seq_id=seq-2
  libc++abi: terminating due to uncaught exception of type std::invalid_argument: [broadcast_shapes] Shapes (32,63) and (1,9,32,64) cannot be broadcast.
  

  The padded mask is built for start = 31 cached rows and the cache holds 32, which settles that issue's open question as (a) for the K/V path. Clamping matched_len to tokens.len() - 1 in try_adopt_cached_prefix after lookup_longest_prefix removes the abort in those three cases. Not part of this issue or its PR.

Pre-submission checks

• I searched existing issues and did not find a duplicate
• This is not a security vulnerability