fix(server): gate tile-aligned prefill on the model, not on the hardware alone - #1205
Merged
Merged
Conversation
…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.
4 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
execute_full_prefill,start_chunked_prefillandcontinue_chunked_prefillpadded the prompt to a 32-token Neural Accelerator tile whenever the host had one:should_align_prefill()in the scheduler ishw.has_neural_accelerator && hw.macos_supports_na. It never asks the model. Every hybrid and recurrent family answerssupports_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
falselist, 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:
and the batched path reads it into a local:
Measured, not inferred
M5 Max,
qwen3.8-27b-4bit(its backbone answersfalse), the same/v1/completionsrequest attemperature: 0, before and after this change:The
afterside 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.rsalready fails when a wrapper does not forward the predicate. It now also fails when analign_to_na_tilecall site is not guarded by it.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 --libfails 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.cargo fmt --checkandcargo clippy --all-targets --features metal,accelerate -- -D warningsclean.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