Skip to content

fix(vision): stop VLM wrappers re-enabling a padded prefill their backbone refuses - #1202

Merged
inureyes merged 1 commit into
mainfrom
fix/issue-1201-vlm-padded-prefill-delegation
Aug 17, 2026
Merged

fix(vision): stop VLM wrappers re-enabling a padded prefill their backbone refuses#1202
inureyes merged 1 commit into
mainfrom
fix/issue-1201-vlm-padded-prefill-delegation

Conversation

@inureyes

Copy link
Copy Markdown
Member

Summary

LanguageModel::supports_padded_prefill defaults to true. Every hybrid and recurrent text model in this tree overrides it to false, with the reason in the comment:

nemotron_h  false // Hybrid SSM model: padding tokens corrupt Mamba recurrent state
mamba2      false // Padding tokens corrupt Mamba2 recurrent state
jamba       false // Padding tokens corrupt Mamba recurrent state in hybrid architecture
kimi_linear false // Padding tokens corrupt DeltaCache recurrent state
rwkv7       false // Padding tokens corrupt RWKV recurrent state

Qwen35Model is one of them. Qwen35VLModel and MiniCPMV46VLModel each hold one in a text_model field and each already forward four other capability predicates (supports_batching, supports_batched_prefill, supports_paged_decode_backend, supports_snapshot_reuse) while leaving these two defaulted.

The result is a wrapper that answers true where its own backbone answers false, re-enabling an optimization the backbone disabled for correctness. The causal mask and trim_caches_to_actual_len undo the pad positions' effect on the KV caches; a GatedDeltaNet state that has already absorbed them cannot be rewound.

Impact

Silent. It compiles, it runs, and greedy output changes, on:

  • Neural Accelerator hardware only (should_align_prefill() requires has_neural_accelerator && macos_supports_na), so M5-class and not M1 Ultra, and
  • prompts whose token count is not already a multiple of 32, since align_to_na_tile is otherwise a no-op.

qwen3.8-27b-4bit is affected because its architectures field is Qwen3_5ForConditionalGeneration, so detect_text_or_vlm routes even a text-only run through Qwen35VLModel.

Measured on M5 Max, 75-token prompt, --temp 0 -n 120:

arm generated verdict
classic, before 630 chars padded to 96, recurrent state corrupted
classic, MLXCEL_NO_PADDED_PREFILL=1 611 chars correct
classic, after this change 611 chars correct
MTP speculative, any config 611 chars never padded, correct all along

With a 1312-token prompt (41 x 32) every arm agrees before and after, which is the control that isolates tile alignment as the variable.

How it was found

#1201, from a disagreement rather than from a crash. MTP's prefill_and_seed forwards the raw prompt length, classic padded it, and their temperature-0 outputs diverged at one token. The first four hypotheses were all wrong (the #1199 qmv_wide split, block width, the GDN chain-parity kernel, nondeterminism); each was ruled out by a control before this one was found. The instrumented build settled it:

[PROBE] should_align_prefill: true (na=true macos=true)
[PROBE] PADDED branch #1: actual=75 padded=96
[PROBE] UNPADDED branch: actual=75 supports_padded=true   <- with the env override

supports_padded=true on a model whose own impl returns false is the whole bug in one line.

The guard

tests/vlm_wrapper_capability_delegation.rs fails when a wrapper in src/vision/ holds a backbone that refuses padded prefill and does not forward the predicate.

It is source-level rather than runtime because constructing every wrapper needs weights, and the property is about which methods an impl block carries, which the source states directly. Two things keep it from rotting:

  • The refusing-backbone set is derived by scanning src/models for impls that answer a bare false, not hard-coded, so a hybrid family added next month is covered the day it lands.
  • "Wraps" means a field type, not a mention, so a doc comment naming a backbone the wrapper does not embed will not trip it.

Verified to fail: removing the Qwen35VLModel override reproduces the exact diagnostic naming the file, the wrapper, and the backbone.

Test plan

  • -p mlxcel --lib fails the same 13 tests over two runs, matching the baseline; -p mlxcel-core --lib is 1472 passed / 4 failed. All are the pre-existing M5 reduced-precision failures of fix(core): f32 matmul and single-query SDPA lose ~fp16 precision on M5, breaking four numeric tests #1065.
  • The new guard passes, and fails when the override is removed.
  • cargo fmt --check and cargo clippy --all-targets --features metal,accelerate -- -D warnings clean.
  • Behavior confirmed end to end: after the change classic decode and MTP produce byte-identical text with no environment variables set.

Not covered

Only the two wrappers the scan flags are changed. Every other wrapper in src/vision/ holds a plain transformer, where the true default is correct, and the guard now says so mechanically rather than by inspection.

Whether the server path has the same gap is not checked here: src/server/batch/scheduler.rs has its own should_align_prefill() and four align_to_na_tile sites, and whether those consult the model predicate at all is a separate question worth its own look.

Fixes #1201

…kbone refuses

LanguageModel::supports_padded_prefill defaults to true. Every hybrid and
recurrent text model overrides it to false, and the comments say why: a
tile-aligned prefill appends up to 31 pad positions, and while the causal
mask and trim_caches_to_actual_len undo their effect on the KV caches, a
Mamba / GatedDeltaNet / RWKV / DeltaCache state that has already absorbed
them cannot be rewound.

Qwen35VLModel and MiniCPMV46VLModel both hold a Qwen35Model, which answers
false, and both forwarded four other capability predicates while leaving
these two defaulted. So on Neural Accelerator hardware a text-only run
through either wrapper padded the prompt to a 32-token tile and corrupted
the backbone's recurrent state. Nothing failed. Greedy output just changed,
and only when the prompt length was not already a multiple of 32.

Measured on M5 Max with qwen3.8-27b-4bit, whose architectures field routes
it through the VLM wrapper. A 75-token prompt padded to 96 and produced a
different 120-token completion than the same prompt with padding disabled.
After this change the padded and unpadded outputs agree, and so does
speculative MTP decode, which never padded and was therefore right all
along. That disagreement is how this was found (#1201); the classic path was
the wrong one.

The guard is a source-level test rather than a runtime one, because
constructing every wrapper needs weights and the property is about which
methods each impl block carries. It derives the refusing-backbone set by
scanning src/models rather than hard-coding it, so a new hybrid family is
covered on the day it lands. Verified to fail by removing one override.

Fixes #1201
@inureyes inureyes added type:bug Bug fixes, error corrections, or issue resolutions priority:critical Requires immediate attention area:models Model architectures, weights, loading, metadata 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 1e7b1d1 into main Aug 17, 2026
8 checks passed
@inureyes
inureyes deleted the fix/issue-1201-vlm-padded-prefill-delegation branch August 17, 2026 11:57
inureyes added a commit that referenced this pull request Aug 17, 2026
…are alone (#1205)

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.
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 area:models Model architectures, weights, loading, metadata 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.

fix(speculative): Qwen 3.5 MTP output is not byte-identical to classic decode, and the exactness probe does not catch it

1 participant