refactor: one prefill planner and one pad-trim hook for CLI and server #2170

Part of #2166. Phase 3. Depends on #2167 (Phase 0 baseline, threshold, and the prefill chunk-size decision). Can run in parallel with Phases 1 and 2. Phase 4a builds on this.

Problem / Background

Prefill chunking, tile padding and pad trim are implemented separately on the CLI and the server, with different defaults. The CLI chunks at 2048 (DEFAULT_PREFILL_CHUNK, src/lib/mlxcel-core/src/generate.rs:282, overridable by MLXCEL_PREFILL_CHUNK). The server chunks at 512 (src/server/config.rs:1172). There are two trim hooks, so #2140 had to touch both. Separately, a server request served from the prompt prefix cache gives different greedy output than its first, uncached run. This was observed with Qwen3-1.7B during #2162, and the result is the same with and without MLXCEL_SDPA_DETERMINISTIC=1. A closely related effect is measured: with the prompt cache on, classic serving prefills a chat prompt in two forwards split at the history boundary, and #2185 found that this split alone made Gemma 4 MTP output differ from classic decode until the MTP prefill mirrored it (parity held with --no-prompt-cache).

Current Behavior

CLI prefill sites in generate.rs:

• prefill_prompt_last_logits (:377), using pad_tokens_for_prefill :200, trim_caches_to_actual_len :228, effective_prefill_chunk :316 and chunked_prefill_last_logits :333.
• The inline copy in generate_with_stats (around :2733-2783).
• The two embedding variants, generate_streaming_with_embeddings (:2233, prefill around :2263-2308) and generate_with_stats_and_embeddings (:2416, prefill around :2453-2498), using pad_embeddings :242.

Server prefill sites in src/server/batch/scheduler/:

• execute_full_prefill (prefill.rs:726), start_chunked_prefill (:915), continue_chunked_prefill (:1136) and run_padded_batched_prefill (:543).
• Chunk ranges come from next_chunked_prefill_range (mod.rs:289), and the pad trim from pad_trim.rs:81.
• execute_full_prefill (prefill.rs:732) and start_chunked_prefill (prefill.rs:920) call capture_history_boundary_snapshot (prompt_cache.rs:1088, feat: Donate a history-boundary snapshot during prefill for snapshot-only families #1143), which forwards the history segment separately when the prompt cache is active. Its doc comment (prompt_cache.rs:1076-1087) states that the split is not bit-exact against an unsplit prefill and names MLXCEL_DISABLE_BOUNDARY_SNAPSHOT=1 as the way to get the unsplit shape.
• The B=1 Gemma 4 MTP burst prefills outside these sites, row-wise in mtp_prefill_ranges (src/models/gemma4_mtp_target.rs:207). Since fix(gemma4): make MTP verify byte-identical to decode on CUDA #2185 it takes the classic split point from history_boundary_split (prompt_cache.rs:1212, same gates as the capture, in the same order) through with_prefill_boundary, set in src/server/batch/scheduler.rs:241 and scheduler/speculative_finalize.rs:327. mtp_prefill_ranges_mirror_classic_history_boundary_split (src/models/gemma4_mtp_target_tests.rs:1501) pins the mirror.

Trim hooks: LanguageModel::trim_internal_caches (generate.rs:652, CLI fallback state) and LanguageModel::trim_sequence_state (generate.rs:675, per SequenceId, defaulting to the former for non-batching models). Seven families in src/models/ (lfm2, plamo2, nemotron_h, granitemoehybrid, gemma3, falcon_h1, inkling/runtime) implement one or both; only gemma3 (src/models/gemma3.rs:1711, :1720) implements both. VisionLanguageModel (src/vision/mod.rs:375-385) forwards both, and four VLM wrappers forward trim_internal_caches (inkling_vl.rs:388, lfm2_vl.rs:240, granite4_vision.rs:357, nemotron_h_nano_omni_vl.rs:611).

A third per-sequence hook sits next to them since #2182: LanguageModel::rewind_decode_appends(seq_id, n) with supports_decode_lookahead_rewind() (generate.rs:698-717), delegated in src/loaded_model.rs:561-571 and src/vision/mod.rs:387-397, implemented only by Gemma 3 (gemma3.rs:1732-1750). It is not a pad trim: it unwinds the scheduler's speculative single-token decode appends exactly, restoring overwritten sliding-window rows from the RotatingKVCache undo log (src/lib/mlxcel-core/src/cache/decode_undo.rs), and the scheduler fails the request when it returns Err. It is not one of the two hooks this phase merges. Gemma 3 also builds its CLI fallback slot (make_fallback_caches, gemma3.rs:1569) without the undo log that its scheduler sequences carry (make_configured_caches, :1555), so the None and Some(seq) targets of trim_state address differently built state.

Proposed Solution

1. One PrefillPlan in src/lib/mlxcel-core/src/prefill_plan.rs: PrefillPlan::new(prompt_len, chunk, caps: PrefillCaps) -> PrefillPlan. PrefillCaps carries supports_chunked_prefill, supports_padded_prefill, supports_maskless_padded_prefill, is_embedding_input and the tile alignment from should_align_prefill. The plan yields an ordered list of PrefillPiece { range, padded_len } plus a trim_after: Option<usize>. Encode the VLM rule in it: embedding input is never chunked and is padded only where the model opts in. Every CLI and server site listed above executes a plan and decides nothing itself. The chunk size comes from one policy function, with the value that ADR 0007 (test: ADR 0007 plus CLI/server parity and B=1 decode benchmark harness #2167) picks.
2. One trim hook. Replace the two hooks with fn trim_state(&self, seq: Option<SequenceId>, excess: i32) -> Result<(), String>. None addresses the fallback single-sequence state the CLI uses. Each of the seven families implements one method, and the VLM wrappers forward one method.
3. Prompt-cache divergence. Find why a prefix-cache hit (adopting cached KV, then prefilling only the suffix) differs from a full prefill of the same tokens. Candidates to check first: the history-boundary split (which prefill partition each run used, given that the uncached run may itself be split by capture_history_boundary_snapshot and the hit forwards only the suffix), a different chunk split of the suffix, tile padding applied to the suffix but not the full prompt, and dense vs paged storage on the adopted blocks. Either make cached and uncached prefill bitwise equivalent under MLXCEL_SDPA_DETERMINISTIC=1, or document in docs/CONTINUOUS_BATCHING.md exactly why they cannot be and bound the difference (first divergent token index, logit max-abs-diff) in a test.

Scope

In scope: the planner, all CLI and server prefill sites, the trim hook across generate.rs, src/loaded_model.rs:546-555, src/vision/ and the seven families, and the prompt-cache investigation and fix.

Out of scope: the KV storage abstraction (Phase 4a), and changing which families support padded or chunked prefill.

Implementation Notes

• Reuse: keep pad_tokens_for_prefill, trim_caches_to_actual_len, pad_embeddings and chunked_prefill_last_logits as the executors the plan drives. Fold effective_prefill_chunk and next_chunked_prefill_range into the plan.
• Constraints: the server keeps interleaving chunked prefill with decode ticks. The plan is a data structure the scheduler steps through, not a blocking loop. The MTP row-wise prefill must keep matching the classic partition: if the plan changes where or whether classic prefill splits at the history boundary, history_boundary_split and mtp_prefill_ranges change with it, and mtp_prefill_ranges_mirror_classic_history_boundary_split stays green. trim_state must leave rewind_decode_appends and its LoadedModel and VLM delegations working (feat(scheduler): sequence-aware cache trim hook so model-owned families can use padded prefill #1755 is the dropped-delegation precedent).
• Edge cases: a prompt shorter than one chunk is a single piece. A 1-token prompt gets no padding. A model with supports_chunked_prefill() == false always gets one piece. A trim error aborts the request, as pad_trim.rs does today.

Acceptance Criteria

• One planner and one hook: rg -n 'fn trim_internal_caches|fn trim_sequence_state' src returns nothing, and no prefill site computes chunk boundaries or padding outside prefill_plan.rs
• A test asserts that chunked and unchunked prefill give the same last-position logits wherever the plan says they must (extending chunked_prefill_matches_single_pass at generate.rs:3362 to the server path)
• A cache-hit vs cache-miss greedy equality test on the server path (in src/server/batch/scheduler_prompt_cache_tests.rs) passes, or the documented bound is asserted by a test and explained in docs/CONTINUOUS_BATCHING.md
• The model-owned pad-trim tests (scheduler_model_owned_pad_trim_tests.rs) and family tests are green. A real-checkpoint run of one hybrid family (for example nemotron_h or lfm2) gives the same greedy output before and after on the CLI and the server
• The planner is used by every real prefill path, with no unused module left behind

Verification

make verify                                   # macOS gate
make verify-fmt verify-test-cuda              # Linux/CUDA gate
cargo test --profile test-fast --features cuda -p mlxcel-core prefill
cargo test --profile test-fast --features cuda --lib server::batch
MLXCEL_SDPA_DETERMINISTIC=1 make <Phase 0 parity target> MODEL=<Qwen3-1.7B-4bit path>

References verified against main at fc21784 (2026-10-07).