feat(scheduler): sequence-aware cache trim hook so model-owned families can use padded prefill #1755

Problem

LanguageModel::supports_padded_prefill() defaults to true with a doc comment asserting that "padding tokens only affect the external KV cache which is trimmed afterwards" (src/lib/mlxcel-core/src/generate.rs:644-654). That assumption is false for any family whose layout is SequenceStateLayout::model_owned.

The scheduler's post-pad trim iterates CachePool::get_caches_mut (src/server/batch/scheduler/prefill.rs:824, :992, :1190), which returns the SequenceCacheSet::caches vector, and that vector is empty for a model-owned sequence (src/lib/mlxcel-core/src/cache.rs:6322). The trim therefore reaches nothing, the pad positions stay in the model's own caches, and offset runs ahead of the real token count by the pad width. The path is gated by should_align_prefill() (M5 or newer with neural-accelerator support, src/server/batch/scheduler/mod.rs:104-107), so it is live on M5 hardware.

PR #1752 works around it by having Gemma 3, AFMoE and Llama 4 answer false (src/models/gemma3.rs:1515, src/models/afmoe.rs:1832, src/models/llama4.rs:1866), joining Gemma 4 (src/models/gemma4.rs:5915) and Muse Glimmer (src/models/muse_glimmer.rs:412), which already did. Three model-owned families still leave the default true and carry the latent defect: src/models/deepseek_v4.rs:955, src/models/bailing_moe_linear.rs:2648, and src/models/qwen3_next.rs:1736.

Proposed solution

Add a sequence-aware cache trim hook the scheduler can call for model-owned families, so padded prefill becomes usable rather than opt-out. The existing trim_internal_caches(&self, excess: i32) cannot serve: it takes no SequenceId and is wired only into the CLI generate paths, so it cannot address one sequence's state inside a batched server. The new hook takes the sequence id and the pad width, and the scheduler calls it at the three trim sites above whenever the model's layout is ModelOwned.

Also audit the three families that still default to true and either fix them through the new hook or opt them out explicitly in the meantime, and correct the trait doc comment so the false assumption is not inherited by the next family that implements the trait.

Acceptance criteria

• A model-owned family running a padded prefill on M5 hardware ends prefill with offset equal to the real token count, not the padded width. (PR feat(scheduler): sequence-aware pad trim for model-owned families #2140: verified on CUDA with forced alignment, MLXCEL_FORCE_PADDED_PREFILL; M5 manual check pending)
• deepseek_v4, bailing_moe_linear and qwen3_next are each either fixed through the hook or carry an explicit supports_padded_prefill() -> false with a comment naming the reason. (PR feat(scheduler): sequence-aware pad trim for model-owned families #2140, plus PipelineServerModel)
• At least one family that PR feat(core): share KV snapshot serialization across attention cache families #1752 opted out (Gemma 3, AFMoE or Llama 4) is restored to padded prefill through the hook, proving the hook works end to end rather than only existing. (PR feat(scheduler): sequence-aware pad trim for model-owned families #2140: Gemma 3)
• The generate.rs doc comment states the model-owned exception. (PR feat(scheduler): sequence-aware pad trim for model-owned families #2140)
• A test pins that a model-owned sequence's offset after a padded prefill equals the unpadded length. (PR feat(scheduler): sequence-aware pad trim for model-owned families #2140: scheduler_model_owned_pad_trim_tests)

Verification

cargo test --workspace --profile test-fast --features metal,accelerate
cargo clippy --workspace --all-targets -- -D warnings

Manual, on M5 or newer: run a prompt whose length is not a multiple of 32 through mlxcel-server for one model-owned family and confirm the generated text matches the same prompt run with padding disabled.

Context: PR #1752 and issue #1335. This is not a fix for #1335 itself.