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

[inureyes]

Problem

MTP's temperature-0 contract is that speculative output is byte-identical to classic decode. It is not, on qwen3.8-27b-4bit on an M5 Max, even with the #1189 exactness probe passing and #1199's qmv_wide switch engaged.

--temp 0 -n 120 --show-reasoning, main at 0199be00, generated span compared directly:

classic:  ... normalization, and training." Need produce final detailed explanation. Need include architect...
mtp:      ... normalization, and training." Need provide detailed explanation. Need final only. Need likely...
                                                 ^ diverges at char 159 of 630

One token flips ("produce" / "provide") and the continuation is a different generation from there.

What it is not

control	result
classic run twice	identical
MTP run twice	identical
classic with and without MLXCEL_QMV_WIDE=0	identical
MTP with #1199's switch vs qmv_wide forced on	identical
MTP at --draft-block-size 2 vs 3	identical

• Not nondeterminism. Both paths reproduce exactly.
• Not the qmv_wide kernel split (perf(speculative): recover MTP throughput on Apple GPU generation 15+ without giving up temperature-0 byte-identity #1187). Turning the switch off changes MTP's output not at all, so the thing feat(speculative): let the MTP gate buy back exactness by dropping qmv_wide #1199 fixed is not the thing breaking this.
• Not block-width dependent. Widths 2 and 3 produce byte-identical output and diverge from classic at the same character. Any verify-block-versus-chain effect would scale with width. This does not, which points away from the block mechanism entirely.
• Not the GDN chain-parity kernel. MLXCEL_GDN_CHAIN_PARITY defaults on, and the shape-fallback warning in models::gated_delta does not fire on this checkpoint, so the parity kernel is being taken.

Why the probe misses it

models::speculative_exactness's module docs already say a passing probe is evidence rather than proof, and this is that caveat cashing out. The probe compares one verify block against one single-token chain from one synthetic state, three draws deep. The contract is a claim about a whole generation from the real sequence of states the decode actually visits. The first is passing here; the second is not.

That gap is the real defect. A gate whose green light does not imply the property it gates is worse than one that is merely conservative, because #1199 now ships MTP on generation 15+ on the strength of it.

Suspects not yet examined

Width independence is the strongest clue: the difference survives when the block is as narrow as it can be, so look at what the MTP path does that classic does not do at all rather than what it does per block.

• The seed prefill (prefill_seed_ms is ~160 ms in the round-loop diagnostics), which runs a target forward classic decode never runs.
• Cache state after a partial-accept rollback, versus a cache that only ever grew by one.
• The accept hook's target-hidden read, if it forces an evaluation ordering the classic path does not.

Acceptance criteria

• The mechanism is identified with evidence, not inferred from where the divergence lands.
• Either MTP output is byte-identical to classic decode at temperature 0 on this pairing, or the contract is restated to what actually holds and the gate is changed to test that instead.
• The probe is extended, or paired with an end-to-end check, so a passing gate implies the contract it is named after. A generation-level diff on a small checkpoint is cheap and would have caught this.
• Checked on Gemma 4 MTP as well as Qwen 3.5, and on an M1 Ultra, so the answer is not scoped to one family on one generation.

Refs #1189, #1186, #1199

[inureyes]

Root cause: the tile-aligned padded prefill, which only the classic path does

Not an MTP defect. mlxcel generate's classic path pads the prompt to a 32-token Neural Accelerator tile before prefilling and then trims the KV caches back; Qwen35MtpTargetAdapter::prefill_and_seed forwards the raw prompt length. Different shapes, different kernels, different KV values, and a near-tie argmax flip 35 tokens later.

MLXCEL_NO_PADDED_PREFILL=1 turns the padding off and the two agree exactly:

arms	result
classic vs MTP, default	diverges at char 159 of 630
classic vs MTP, MLXCEL_NO_PADDED_PREFILL=1	identical
unpadded classic vs default MTP	identical
classic padded vs classic unpadded	diverges at char 159

The last two rows are the whole story. MTP's output never changes, because MTP never pads. Classic's output is what moves, and it moves to exactly MTP's output when the padding is removed. So MTP has been byte-identical to classic decode all along, against an unpadded reference, and the reference is what shifted.

This explains everything the earlier controls showed. Block widths 2 and 3 agree because prefill is upstream of the block. The #1199 qmv_wide switch changes nothing here because this is not the verify path. The probe passes honestly, because it compares a block against a chain that both ran through the same MTP prefill.

The gates involved:

fn should_align_prefill() -> bool {
    if std::env::var("MLXCEL_NO_PADDED_PREFILL").is_ok() { return false; }
    let hw = hardware::get_hardware();
    hw.has_neural_accelerator && hw.macos_supports_na
}

generate.rs:1267 takes the padded branch when that holds and the model supports it, then align_to_na_tile(75) gives 96. So this is generation-17-only: on hardware without a neural accelerator neither path pads, both agree, and MTP's parity claim has been true there the whole time.

Two separate things to fix

1. MTP's prefill should pad the same way the classic path does. That is the parity fix, and it is also free throughput MTP is currently leaving on the table, since the padding exists because the NA prefers tile-aligned sequences. The masking and cache trimming are already written and tested (pad_tokens_for_prefill, trim_caches_to_actual_len); prefill_and_seed needs to route through them instead of calling forward_prefill_with_last_hidden_for_sequence with the raw length.

2. Padded prefill changes greedy output on M5, and that is worth stating out loud. The padding is mathematically neutral: positions past actual_len are masked and trimmed. The difference is pure floating-point, from feeding a 96-wide sequence to kernels instead of a 75-wide one. But it means mlxcel generate --temp 0 output on an M5 depends on prompt length modulo 32, and that the same command on an M1 Ultra produces a third answer, since that host never pads. Nothing in the docs says so. Whether to document it or to make it uniform is a policy call bigger than this issue.

What this retires

The suspects listed above are ruled out: the seed prefill's extra target forward, rollback cache state, and the accept hook are all innocent. So is the reframing on #1199: its contract claim was wrong for a different reason than I gave. The switch does restore block-vs-chain exactness, and the end-to-end gap was never about that.

[inureyes]

The same bug exists on the server, and it is wider there

Reopening the topic rather than the issue, because the fix is separate: PR #1205.

This issue was about a VLM wrapper losing supports_padded_prefill by inheriting the trait default instead of forwarding it. The server loses the same guard a different way: three call sites never ask.

src/server/batch/scheduler.rs, in execute_full_prefill, start_chunked_prefill and continue_chunked_prefill:

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

The scheduler's should_align_prefill() is hw.has_neural_accelerator && hw.macos_supports_na, with no model term. So a hybrid model served on an M5-class host was padded regardless of what it reports, and unlike this issue that needs no wrapper: it reaches Mamba, Jamba, RWKV, Nemotron-H, Falcon-H1, Kimi Linear, Qwen 3.5 and every other family on the false list directly.

Confirmed the same way, by A/B on one endpoint rather than by reading: the same /v1/completions request at temperature 0 returns different text before and after the guard, diverging at character 74.

Two things worth recording from this.

The file already knew. Its boundary-snapshot path declines to pad and explains why in a comment naming the exact hazard, and its batched path reads the predicate into can_pad_prefill. Two of five sites got it right, three did not, and nothing failed.

The guard test added here only covered half the failure space. It checked that a type forwards the predicate; it did not check that a call site consults it. #1205 adds the second half. The general shape is that a capability predicate can be lost either by a type that answers a permissive default or by a caller that never asks, and a guard for one is not a guard for the other.