lablup / lablup/mlxcel

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

Open
#1,755 0 comments 0 reactions 0 assignees View on GitHub
area:core area:inference priority:medium status:ready type:enhancement
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.