feat(scheduler): sequence-aware cache trim hook so model-owned families can use padded prefill
- Dominant language
- Rust
- Stars
- 467
- Forks
- 54
- Avg merge
- 4h 25m
- Merged PRs (30d)
- 310
Description
## 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.
- [ ] `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.
- [ ] At least one family that PR #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.
- [ ] The `generate.rs` doc comment states the model-owned exception.
- [ ] A test pins that a model-owned sequence's offset after a padded prefill equals the unpadded length.
## Verification
```bash
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.
Contributor guide
Research direction
Read the cache ownership and trim paths in src/lib/mlxcel-core/src/cache.rs and src/server/batch/scheduler/prefill.rs, then inspect the model trait in src/lib/mlxcel-core/src/generate.rs and the three named model files. Run the workspace tests with the listed Metal and Accelerate features before changing behavior. Done means model-owned padded prefill keeps offset at the unpadded length, the named families are handled or explicitly opted out, and a regression test covers the result.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- rust
- Domain
- ai, backend, performance
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Active
- Clarity
- Clearly specified
- Newbie friendliness
- 38/100