fix(speculative): Qwen 3.5 DFlash verify diverges from classic decode

[inureyes]

Problem / Background

On GB10 (sm_121), models/mlx/qwen3.5-4b-4bit served with models/mlx/qwen3.5-4b-dflash, greedy (temperature 0) completions from mlxcel-server diverge from classic decode at every DFlash verify width measured: 2, 3, 4, 5, 6, 7, 8 and 16. Width 16 is the width main resolves today, so this is the status quo and not a regression introduced by #1797. The two classic arms of the same session, measured from two separate server processes at the start and the end, are byte-identical to each other, so greedy decode itself is reproducible and the divergence is specific to the speculative path. A typical divergence, about 30 tokens in: the classic arm emits for attempt in range(self.max_attempts): where the width-4 arm emits for i in range(self.max_attempts):.

Grouping the arms by the sha256 of their completion text, at MLX's default MLX_ENABLE_TF32=1, gives exactly two groups: the two classic arms, and every speculative width from 2 to 16. All ten speculative arms are byte-identical to each other and none matches classic. So this is not width-specific, not a rollback or first-bonus accident, and not non-determinism: it is the speculative verify path disagreeing with the single-token decode chain, uniformly.

Forcing MLX_ENABLE_TF32=0 splits the speculative arms into two groups instead, 2 to 7 against 8 and 16, which is exactly the M * B < 8 boundary between qmv_multirow_kernel and qmm_sm80_kernel in mlx/backend/cuda/quantized/quantized.cpp. That the grouping tracks the kernel dispatch boundary points at the same failure class the Metal MLXCEL_MTP_ALLOW_INEXACT gate exists for: a quantized projection dispatching to a different kernel at M = K than at M = 1 and not being bit-equal to it. On the CUDA DFlash path that class is unguarded.

A control rules the obvious candidate out. MLXCEL_QMV_MULTIROW=0 (the #725 kill switch) restores the stock per-row path, so a 4-row verify then runs the same qmv_kernel classic decode runs, launched four times instead of once. If the multirow kernel were the cause, width 4 would become byte-identical to classic under it. It does not. Width 4's completion under the kill switch hashes to 57d291a..., which is exactly the hash it has with the multirow path enabled, while classic keeps 2c76b0a.... The multirow kernel makes no difference to the result at all, so it is not what the verify and the chain disagree about, and the grouping by kernel family that the forced-TF32-off pass showed is not the mechanism it looked like.

That leaves the difference somewhere else in the verify path rather than in the quantized projection. The candidates this record cannot separate, in the order a reader should try them: the 24 gated-delta linear-attention layers, where a T-row chunked scan is only bit-equal to T single-token steps when the chain-parity path is taken (MLXCEL_GDN_CHAIN_PARITY, and docs/environment-variables.md:327 says its CUDA arm covers an unmasked scalar-gated block of at most 32 rows, which this shape should satisfy but was not verified here); the materialized attention mask the eight full-attention layers use at T > 1; and the first-bonus and rollback handling around the round loop. Separating them needs a per-layer comparison, not a throughput sweep.

Evidence: docs/benchmark_results/draft-block-width-default-gb10-2026-09-20.md and the per-arm completion text in docs/benchmark_results/data/draft-block-width-gb10-2026-09-20/*.jsonl. Both land with the pull request for issue #1797 and are not on main yet.

Current Behavior

DFlashTargetModel::exactness_allows defaults to true (src/server/batch/dflash_target.rs:212), and neither the Qwen35Model impl (src/server/batch/dflash_target.rs:251) nor the Qwen35VLModel impl (:257) overrides it, so the burst gate at src/server/batch/speculative_burst.rs:1670-1671 always reads permissive for this pairing. Lfm2Model (:270), MuseGlimmerTextWrapper (:322) and LagunaWrapper (:354) do override it and run a measured block-versus-chain probe; Qwen 3.5 on the DFlash path is permissive. So the property #1782 reported for this pairing is asserted by measurement, not enforced by a gate, and nothing re-checks it.

PR #1795's own record states byte-identity at widths 2, 4 and 8 on this pairing on 2026-09-11 (docs/benchmark_results/dflash-verify-fixed-cost-gb10-2026-09-11.md), and marks only width 16 as divergent. If that was right then, the property has been lost since. If it was measured loosely then, the record overstates a contract users may be relying on. Either way an operator turning on speculative decoding today gets different greedy text than without it, silently.

Proposed Solution

1. Determine whether the property held at PR fix(cuda): attribute and remove the fixed cost of multi-row speculative verify #1795's merge commit e391ae9c and broke since, or never held as stated, by re-running the same comparison at that commit on the same GB10 host, with the harness that lands with the perf(speculative): the default draft block width of 16 loses on GB10 #1797 PR (docs/benchmark_results/data/draft-block-width-gb10-2026-09-20/harness/identity.sh and run_session.sh).
2. If it broke, bisect between e391ae9c and main and fix the divergence, trying the candidates in the order listed above.
3. Either way, enforce the contract instead of asserting it. Add Qwen35Model::dflash_exactness_allows(block_size) mirroring Lfm2Model::dflash_exactness_allows (src/models/lfm2_speculative.rs:256): build a ProbeKey from block size, hidden size and layer count, and call mtp_exactness_gate(key, || self.probe_block_chain_exactness(block_size)), reusing the probe already on the model at src/models/qwen3_5.rs:1571. Then override exactness_allows in the Qwen35Model and Qwen35VLModel impls in dflash_target.rs to call it. Do not route this through Qwen35Model::mtp_exactness_allows (src/models/qwen3_5.rs:1518): its first precondition is metal_is_available(), so on CUDA it would decline every DFlash burst instead of measuring one. mtp_exactness_gate (src/models/speculative_exactness.rs:295) already owns the per-(model, width) memoization, the qmv_wide retry, the decline log line and the MLXCEL_MTP_ALLOW_INEXACT override, and it treats any non-Equal verdict (NotRun included) as a decline, so no new policy code is needed.
4. Extend the served-path regression test rather than adding a parallel copy: tests/speculative_parity.rs:610 already holds the #[ignore]-gated greedy_parity_dflash_qwen35_4b, which spawns mlxcel-server with and without --draft-kind dflash at temperature 0 on this exact pairing. Give it the widths this record measured (2, 4, 8, 16), and record in the PR why it did not surface this divergence. A synthetic probe passing is not the same as the served path agreeing (the VLM server-path trap).

Scope

In scope: src/models/qwen3_5.rs (DFlash-facing exactness entry point), src/server/batch/dflash_target.rs (the two Qwen 3.5 impl overrides), tests/speculative_parity.rs (widths and assertions), and a new record under docs/benchmark_results/ for the e391ae9c re-run.

Out of scope: changing the default draft block width (issue #1797), and any change to the Metal MTP gate or to the LFM2, Muse Glimmer and Laguna probes.

Acceptance Criteria

• A record under docs/benchmark_results/ states, per width, whether classic and speculative greedy text were byte-identical at e391ae9c on GB10, and therefore whether this broke after PR fix(cuda): attribute and remove the fixed cost of multi-row speculative verify #1795 or never held as stated. Answered by code archaeology rather than a rebuild, and more conclusively: the defective prefill path was already in place at e391ae9c, so the property never held as stated. See docs/benchmark_results/dflash-verify-divergence-qwen35-gb10-2026-09-20.md (PR fix(speculative): prefill the Qwen 3.5 DFlash burst like classic decode #1939).
• With the gate in place and MLXCEL_MTP_ALLOW_INEXACT unset, a GB10 mlxcel-server start on this pairing logs the probe verdict, and on a divergent verdict declines the burst to classic decode rather than serving speculative text. Verified: widths 8 and 16 decline and return text byte-identical to the drafter-less baseline.
• MLXCEL_MTP_ALLOW_INEXACT=1 still engages the burst and logs the forfeit warning, so the existing escape hatch governs the override. Verified at width 16.
• greedy_parity_dflash_qwen35_4b covers widths 2, 4, 8 and 16, asserts the server logged Speculative burst completed before comparing, and fails on current main on GB10 at every one of those widths. It now requires each width to either run a burst and match the baseline byte for byte or be declined by the gate, and names which happened rather than letting a decline read as a pass.
• After the fix that test passes on GB10 with byte-identical completions at every covered width, or the gate declines the burst and the test records the decline as a decline rather than as a pass. Met by the second branch. The residual at widths 2 and 4 is genuine numerics in CUDA's fused sdpa_vector path, which MLXCEL_SDPA_VECTOR_LARGE_D=0 collapses to byte-identity in two independent arms, and the probe could not observe it because it compares one block immediately after a clean prefill while the difference does not begin until several dozen tokens past the prompt (no prompt length from 8 to 512 changes that verdict). The probe now declines on that configuration instead of reporting a pass it cannot support, so every width declines and every response equals the drafter-less baseline byte for byte. MLXCEL_MTP_ALLOW_INEXACT=1 engages the burst anyway, and MLXCEL_SDPA_VECTOR_LARGE_D=0 buys byte-identity back rather than forfeiting it.

Verification

cargo test --test speculative_parity --release -- --ignored greedy_parity_dflash_qwen35_4b --test-threads=1 --nocapture
docs/benchmark_results/data/draft-block-width-gb10-2026-09-20/harness/identity.sh

--test-threads=1 is required on CUDA. A pass is: the server log shows a completed speculative burst, and the classic and speculative completion sha256s match at every covered width, or the probe declines with its verdict logged and no speculative text is served.

Technical Considerations

Related: issue #1797 (draft block width default; its pull request brings the record and raw data cited above), PR #1795 (merge commit e391ae9c), issue #1782 (the chain-parity CUDA arm) and issue #725 (MLXCEL_QMV_MULTIROW).

[inureyes]

One correction to the body, which was written from a brief prepared before the last control finished.

The body says the forced-MLX_ENABLE_TF32=0 grouping "is exactly the M * B < 8 boundary between qmv_multirow_kernel and qmm_sm80_kernel". That is true as a description of the grouping and misleading as a statement of cause. The MLXCEL_QMV_MULTIROW=0 control reported further down rules the quantized projection kernel out: with the multirow path disabled a 4-row verify runs the same qmv_kernel classic decode runs, and its output is byte-identical to the multirow-enabled run while still differing from classic. So the kernel-family grouping looked like the Metal MLXCEL_MTP_ALLOW_INEXACT failure class and is not it.

The candidates listed at the end of the body are where the difference actually lives.

[inureyes]

Diagnosed and mostly fixed in #1939. The cause is not on this issue's candidate list, so the list should not be worked from further.

Root cause: the burst prefilled the prompt as if it were a verify block. Qwen35Model::forward_speculative sends every full-attention layer through attend_per_position, which computes attention one query position at a time so each row sees exactly the prefix a single-token decode step would. That is right for a verify block and is not what classic decode runs over a prompt, which is one batched causal attention. The DFlash burst's first call is the prompt prefill through SpeculativeTarget::prefill_forward_with_capture_layers, that hook defaults to the verify forward, and Qwen 3.5 did not override it. So the burst entered its first round holding different KV and gated-delta state than classic decode, at every width alike, which is exactly the width-independence this issue reported. Measured: the two prefills' last-row logits differed in 173206 of 496640 bytes, and the first cache state to differ was layer 4's, the layer after the first full-attention layer.

The MTP path already carries this fix. Qwen35Model::forward_prefill_with_last_hidden exists for it and says so in its own doc comment since #1182: forward_speculative "does not match the classic prefill numerics", and the first bonus has to come from the batched causal prefill to stay byte-identical. The DFlash path never got it.

This answers question 1 without a rebuild at e391ae9c, and more conclusively. attend_per_position landed on 2026-05-24 (4038da96, #78). At e391ae9c (2026-09-11) forward_with_capture already called forward_with_position_ids_verify(..., true) and Qwen35Model still did not override the prefill hook, so the burst prefilled the same way then. The property did not break after PR #1795; it never held as stated, and dflash-verify-fixed-cost-gb10-2026-09-11.md's byte-identity claim at widths 2, 4 and 8 should be read as not established.

What killed each listed candidate, by control rather than by argument. Gated-delta chunked scan: a served MLXCEL_GDN_CHAIN_PARITY=0 A/B that was live (burst 5511 against 7826 ms, acceptance 0.5325 against 0.5281) and changed the completion not at all. Materialized attention mask at T > 1 and the verify block generally: the per-layer bisect finds a width-4 verify block byte-identical to the single-token chain at every one of the 32 layers. First bonus and rollback: rollback is byte-exact over 100 rewinds, including with the rejected rows carrying a wrong token id the way a real round's do, across six accept patterns mixing full-accept and partial rounds. One more control was void rather than negative and is recorded that way: MLXCEL_SDPA_FALLBACK_MAX_QUERIES=0 can do nothing on a head_dim 256 checkpoint, which never reaches cuDNN, and its unchanged burst time is how that was identified.

A second, separate defect at widths 8 and 16, which is gated rather than fixed. probe_block_chain_exactness on the real checkpoint returns byte-identical at widths 2 and 4 and divergent at 8 and 16, in 229971 of 496640 logit bytes at block position 0. That is MLX's M * B < 8 switch from qmv to qmm_sm80, a real difference between two kernels rather than a defect, so the gate this issue proposed is the right answer there and bit-identity is not achievable. Note what this means for the gate on its own: at widths 2 and 4 the probe passes, so the gate alone would have admitted the burst at the shipped default of 4 and still served divergent text. It needed the prefill fix underneath it.

What is not fixed. Widths 2 and 4 still part from classic decode at generated token 105, down from token 9. It is a different mechanism that the prefill bug was masking, not a remnant of it, and the last acceptance criterion below stays unchecked because of it. It is invariant to MLXCEL_QMV_MULTIROW=0, MLX_USE_CUDA_GRAPHS=0, --no-warmup and omitting logprobs; widths 2 and 4 produce identical text despite entirely different drafting patterns; and every in-process arm of the target's prefill, verify and rewind is byte-exact over the same 200-token transcript. So it sits in the served burst wrapper rather than in the target's forward, and the next step is to drive DFlashGenerator::run in process with the real drafter bound, which is the one configuration those arms do not reproduce.

Record: docs/benchmark_results/dflash-verify-divergence-qwen35-gb10-2026-09-20.md. Diagnostics: src/models/qwen3_5_dflash_probe_tests.rs.

[inureyes]

Follow-up on the residual at widths 2 and 4, which is now excluded from the target path much more strongly than the comment above could claim.

The one thing an in-process replay could not reproduce from the outside was the served burst's own round structure: the generator collects per-round accept lengths and nothing logged them, so the replay had to walk a synthetic accept cycle. PR #1939 adds a debug line carrying that vector. Replaying the served width-4 run's own 79 rounds and 56 rewinds, at the same depths in the same order and with the served capture list [1, 8, 15, 22, 29], the target disagrees with classic decode at 0 of 201 greedy positions, with the rejected rows carrying a wrong token id and carrying the correct one alike.

So the residual is not in the prefill, not in the verify block, and not in the rewind, under the served run's own cache history rather than an approximation of it. One hypothesis is left standing: the drafter executing MLX work between the target's forwards, which is the only part of a served burst that none of these arms contain. The next step is to drive DFlashGenerator::run in process with the real drafter bound and bisect from there; the accept-length log is what makes that comparison anchorable.

greedy_parity_dflash_qwen35_4b fails on the PR branch at widths 2 and 4 for this reason, and passes at 8 and 16 by decline. That red is this open finding rather than a regression: on main it fails at all four widths and every one of those failures is the prefill bug.