Skip to content

fix(server): gate tile-aligned prefill on the model, not on the hardware alone - #1205

Merged
inureyes merged 1 commit into
mainfrom
fix/server-padded-prefill-predicate
Aug 17, 2026
Merged

fix(server): gate tile-aligned prefill on the model, not on the hardware alone#1205
inureyes merged 1 commit into
mainfrom
fix/server-padded-prefill-predicate

Conversation

@inureyes

Copy link
Copy Markdown
Member

Summary

execute_full_prefill, start_chunked_prefill and continue_chunked_prefill padded the prompt to a 32-token Neural Accelerator tile whenever the host had one:

if should_align_prefill() && seq.vlm_embeddings.is_none() {   // hardware only

should_align_prefill() in the scheduler is hw.has_neural_accelerator && hw.macos_supports_na. It never asks the model. Every hybrid and recurrent family answers supports_padded_prefill() == false, because pad positions are absorbed by a conv / SSM / GatedDeltaNet state and trimming the KV caches afterwards does not rewind it.

So on an M5-class host, every such model served through these paths had its recurrent state corrupted whenever a prompt or chunk was not already tile-aligned. That is Mamba, Mamba2, Jamba, RWKV, Nemotron-H, Falcon-H1, Kimi Linear, LFM2, Plamo2, GraniteMoeHybrid, Qwen 3.5 and the rest of the false list, VLM-wrapped or not.

The file already knew

Two other places in the same scheduler get it right, which is what makes this a slip rather than a disagreement:

// No NA-tile padding here: several snapshot-only families report
// `supports_padded_prefill() == false` because padding tokens corrupt
// their conv / SSM recurrent state, and a padded segment would make the
// captured state describe more tokens than the key claims.

and the batched path reads it into a local:

let can_pad_prefill = self.model.supports_padded_prefill();

Measured, not inferred

M5 Max, qwen3.8-27b-4bit (its backbone answers false), the same /v1/completions request at temperature: 0, before and after this change:

before: ...broken down into its key components.\n\n### 1. The Core: ...
after:  ...illustrated with a series of images.\n\n### 1. The Core: ...
                                   ^ diverges at char 74 of the completion

The after side is the correct one: no padding on a model that says it cannot take it.

A CLI-to-server comparison is not the discriminator here, because the server parses the reasoning channel differently, so the before/after A/B on the same endpoint is what isolates the variable.

Relationship to #1201

This is the server-side twin, and wider. #1201 needed a VLM wrapper to lose the predicate by inheriting the trait default instead of delegating; the offline path itself asked correctly. Here the call sites never ask, so it reaches every hybrid model directly.

Same corruption, same hardware gate, different way of losing the guard.

The guard test grows a second case

tests/vlm_wrapper_capability_delegation.rs already fails when a wrapper does not forward the predicate. It now also fails when an align_to_na_tile call site is not guarded by it.

failure mode check
a type answers the trait default instead of forwarding wrapper scan (#1202)
a call site never asks call-site scan (this PR)

The call-site rule is deliberately shallow: the guard must appear within a dozen lines before the call, either as the predicate itself or as a local bound to it, since the batched path hoists it fourteen lines up. A guard further away is one a reader cannot see either. The definition, its doc examples and its own unit tests are skipped.

Verified to fail by removing one of the three guards: it names the file, the line and the call.

Test plan

  • -p mlxcel --lib fails the same 13 tests as baseline, by set comparison rather than by count. The count reads 12 to 14 across runs, which is the pre-existing parallel-execution flakiness in the "prefill is causal" family, not this change.
  • Both guard tests pass, and the new one fails when a guard is removed.
  • cargo fmt --check and cargo clippy --all-targets --features metal,accelerate -- -D warnings clean.
  • Behavior confirmed on a running server, before and after, same request.

Not covered

Whether padding actually buys anything is a separate question this PR does not settle. Measured on the CLI earlier today at 68, 75, 237 and 1324 tokens, tile alignment made no difference to prefill throughput outside noise, because where the padding is a large fraction the prefill is fixed-cost dominated, and where the prefill is compute-bound the padding is at most 31 tokens. If that holds on the server too, the optimization is costing correctness risk for nothing and should be reconsidered rather than guarded.

Refs #1201

…are alone

execute_full_prefill and both chunked-prefill paths padded the prompt to a
32-token Neural Accelerator tile whenever the hardware had one, without
asking whether the model tolerates it. Every hybrid and recurrent family
answers supports_padded_prefill() == false, because pad positions are
absorbed by a conv / SSM / GatedDeltaNet state that trimming the KV caches
afterwards does not rewind. So every such model served on an M5-class host
had its recurrent state corrupted whenever a prompt or chunk was not already
tile-aligned.

The same file knows the hazard. The boundary-snapshot path declines to pad
and says why in a comment, and the batched path reads the predicate into
can_pad_prefill. These three sites simply never asked.

Measured on M5 Max with qwen3.8-27b-4bit, whose backbone answers false: the
same /v1/completions request at temperature 0 returns different text before
and after this change, diverging at character 74 of the completion. The
after side is the correct one.

This is the server-side twin of #1201, and wider: that one needed a VLM
wrapper to lose the predicate by not delegating it, while this one reaches
every hybrid model directly, VLM-wrapped or not.

The guard test added with #1202 grows a second case for the second way to
lose a guard. The first covers a type that answers the trait default instead
of forwarding; this covers a call site that never asks. It accepts a local
bound to the predicate as the guard, since the batched path hoists it
fourteen lines up, and skips the definition and its own unit tests. Verified
to fail by removing one of the three guards.
@inureyes inureyes added type:bug Bug fixes, error corrections, or issue resolutions priority:critical Requires immediate attention area:core mlxcel-core: MLX FFI, primitives, KV cache, layers platform:macos macOS (Apple Silicon) specific status:review Under review labels Aug 17, 2026
@inureyes
inureyes merged commit fa260a5 into main Aug 17, 2026
8 checks passed
@inureyes
inureyes deleted the fix/server-padded-prefill-predicate branch August 17, 2026 15:08
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area:core mlxcel-core: MLX FFI, primitives, KV cache, layers platform:macos macOS (Apple Silicon) specific priority:critical Requires immediate attention status:review Under review type:bug Bug fixes, error corrections, or issue resolutions

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant