Repository navigation
fix(speculative): prefill the Qwen 3.5 DFlash burst like classic decode - #1939
Merged
Merged
Conversation
…Qwen 3.5 Three `#[ignore]`-gated diagnostics on `models/mlx/qwen3.5-4b-4bit`, written to separate the candidates issue #1935 lists rather than to assert a fixed expectation. They are what located the divergence, and they stay so the next reader can re-run the separation instead of re-deriving it. `block_chain_exactness_verdicts_on_the_real_checkpoint` reports `probe_block_chain_exactness` per width. On GB10 it returns Equal at widths 2 and 4 and Diverges at 8 and 16, which is the result that rules the issue's proposed gate out as a complete fix: at the shipped default of 4 the probe passes and the served text still differs from classic decode. `per_layer_block_versus_chain_bisect` compares a verify block against the single-token chain layer by layer at block position 0, both arms inside `forward_speculative`. Every layer is byte-identical at width 4, so the verify forward itself is exact there. `classic_forward_versus_speculative_forward_at_one_row` compares the pair the block-versus-chain probe holds constant: `forward_internal`, which is what classic decode reaches, against `forward_speculative`, which is what the DFlash burst runs including its prompt prefill. They diverge in the prefill, first in the layer-4 cache state, which is the layer after the first full-attention layer. Refs #1935
…ecode does The DFlash burst's first call is the prompt prefill, and for Qwen 3.5 it went through `forward_speculative`, which sends every full-attention layer through `attend_per_position`. That branch exists so a verify block's rows match what single-token decode would produce; it is not the computation classic decode runs over a prompt, which is one batched causal attention. The burst therefore started from different KV and gated-delta state than classic decode, and its greedy output drifted away within a few tokens at every verify width alike, which is what issue #1935 reported. The MTP path already carries this fix: `forward_prefill_with_last_hidden` exists precisely because `forward_speculative` "does not match the classic prefill numerics", and `Qwen35MtpTargetAdapter::prefill_and_seed` uses it. The DFlash path never got it, because `SpeculativeTarget::prefill_forward_with_capture_layers` defaults to the verify forward and Qwen 3.5 did not override it. New `Qwen35Model::forward_prefill_with_capture_layers` runs `forward_internal`'s loop, batched causal mask and all, while capturing the per-layer hidden states the drafter is seeded from, and returns no gated-delta rollback snapshots, which a prefill never needs and which are prompt-sized rather than block-sized. Both `SpeculativeTarget` impls for the family, `Qwen35Model` and `Qwen35VLModel`, override the hook to call it, so the B = 1 and B > 1 burst arms are both covered. Measured on GB10 with `models/mlx/qwen3.5-4b-4bit` and `models/mlx/qwen3.5-4b-dflash`, greedy, 200 tokens, one server per arm: the two prefills' last-row logits differed in 173206 of 496640 bytes before this, and the served completion parted from classic decode at generated token 9 at every width from 2 to 16. After it, widths 2 and 4 agree with classic decode through token 104. Widths 8 and 16 still part at token 9, from a different cause: the verify block itself is not byte-identical to the single-token chain there, which the block-versus-chain probe reports directly. Refs #1935
…xactness The Qwen 3.5 DFlash target took `DFlashTargetModel::exactness_allows`'s permissive default while LFM2, Muse Glimmer and Laguna all override it with a measured block-versus-chain probe. On GB10 that default is wrong above width 4: `Qwen35Model::probe_block_chain_exactness` reports the verify block byte-identical to the single-token chain at widths 2 and 4 and divergent at 8 and 16, in 229971 of 496640 logit bytes at block position 0. The boundary is MLX's `M * B < 8` switch from the `qmv` family to `qmm_sm80`, which is a real numerical difference between two kernels rather than a defect, so the burst has to decline there instead of serving text that differs from classic decode. New `Qwen35Model::dflash_exactness_allows` runs that probe through the shared `mtp_exactness_gate`, which owns the per-(model, width) memoization, the decline log line and the `MLXCEL_MTP_ALLOW_INEXACT` override, and both Qwen 3.5 `DFlashTargetModel` impls call it. Deliberately not `mtp_exactness_allows`: its first precondition is `metal_is_available()`, so on CUDA it would decline every burst without measuring one. Verified served on GB10, one server per arm, greedy, 200 tokens: widths 8 and 16 now decline and return text byte-identical to the drafter-less baseline, widths 2 and 4 pass the probe and engage, and `MLXCEL_MTP_ALLOW_INEXACT=1` at width 16 still engages the burst with its forfeit warning. `greedy_parity_dflash_qwen35_4b` now covers widths 2, 4, 8 and 16 through a new multi-width helper rather than the single width 16 it used to run, which the gate now declines. Each width must either run a burst and match the baseline byte for byte or be declined by the gate, and the helper names which of the two happened instead of letting a decline read as a pass. The record under `docs/benchmark_results/` carries the measurement, the controls that excluded the gated-delta scan, the multirow kernel and TF32, one control that was void rather than negative, and the residual divergence at widths 2 and 4 that this change does not fix. Refs #1935
inureyes
added
status:review
5 tasks done
A served burst's round structure is the one thing an in-process replay of issue #1935's divergence could not reproduce from the outside: the generator collects per-round accept lengths and nothing logged them, so a replay had to guess at a synthetic accept cycle and could not claim to have walked the same cache history. One `debug` line per burst in the DFlash target's finalize path carries the vector. Adding one to each entry (the log reports accepted draft tokens, the replay takes kept rows) and passing them as `MLXCEL_Q35_PROBE_ACCEPTS` reproduces the served run's 79 rounds and 56 rewinds exactly, and the target then agrees with classic decode at 0 of 201 greedy positions. That is what moves the residual divergence at widths 2 and 4 from "not reproduced in process" to "excluded from the target's prefill, verify and rewind under the served run's own round structure". The record gains that arm, the recipe for capturing an accept sequence, a note that the prefill cache comparison has to precede the chain loop or it reports the chain's own offset gap, and a correction: the claim that classic decode reaches `forward_internal` now rests on the scheduler's call graph through `forward_with_sequence_id`, not on the offline and served endpoints agreeing, which is a consistency check rather than a proof of the path. Refs #1935
inureyes
added a commit
that referenced
this pull request
Sep 21, 2026
The 2026-09-11 record states that greedy text is byte-identical to classic decode on every after arm including block 16. Issue #1935 and PR #1939 established that this never held: the DFlash burst prefilled the prompt through `forward_speculative`, which drives full-attention layers one query position at a time instead of the single batched causal attention classic decode runs, so the burst entered its first round from different KV and gated-delta state at every width. `attend_per_position` predates the measurement, so the property had not held when the record was written rather than regressing afterwards. The record is corrected in place with a banner rather than edited, matching how the superseded #1820 record is handled, because #1797's issue body already cites this table and a reader arriving at the file directly needs the correction there. The banner also withdraws the attribution of block 16's divergence to the ignored chain-parity flag, which PR #1939 showed changes acceptance and burst time without changing the completion, and flags that every acceptance-dependent number needs re-checking against the fixed prefill. Refs #1935
2 tasks done
inureyes
added a commit
that referenced
this pull request
Sep 21, 2026
#1945) * docs(benchmarks): re-measure the draft block width on main after the #1939 prefill fix Issue #1797 seeded a measured default of 4 for `(12, 1, Affine)` on a DFlash drafter from a sweep taken before PR #1939 fixed the burst's prompt prefill. That fix changes the state every round starts from, which changes acceptance, which changes throughput per width, and nobody had looked since. The same harness, the same pairing, widths 2, 3, 4 and 6 at n = 3 with the classic brackets at both ends, on a binary built from `origin/main` and verified to carry none of the issue #1935 branch's gate. The ordering moved. Width 3 runs at 1.19x classic and its slowest run clears width 4's fastest by 0.80 tok/s, twice the session's full bracket spread, with the drift working against width 3 rather than for it since width 4 ran later. Width 2 and width 4 overlap and this sweep does not order them. Width 6 is below every classic run. Acceptance falls monotonically as the block widens (0.754, 0.599, 0.511, 0.432) against a per-round device-sync cost that rises (19.9, 24.8, 30.3, 47.4 ms), and the product of those peaks at 3 rather than at 4. The record states the drift rather than printing a clean table over it: the two classic brackets do not overlap, by about half a percent, and that is the resolution floor for everything between them. Nothing is differenced against the #1797 table, which ran on a different binary and a different tree. This record does not change the default; it files the evidence so the number can be re-examined against 3. Refs #1797, #1935 * docs(benchmarks): separate the two drift numbers and say the width still matters The reading section quoted the classic brackets' full spread (0.40 tok/s) in one sentence and the shift between their means (0.27) in the next without naming either, which invites a reader to add them. Both now say which they are. And the closing paragraph said the width was academic for this pairing because PR #1944's gate declines the burst on CUDA. It is not: `MLXCEL_MTP_ALLOW_INEXACT=1` is a documented escape and an operator who sets it gets exactly this curve, so the default still decides what they run at. The paragraph now says plainly that on this evidence 4 is not the right number for this pairing on this host and 3 is. Refs #1797, #1935
inureyes
added a commit
that referenced
this pull request
Sep 21, 2026
The `draft-block-width-gb10-2026-09-20/harness/` directory is a live tool despite its dated path: the post-#1939 width sweep and the #1943 architecture A/B both re-used it, and `draft-block-width-post-1939-gb10-2026-09-21/` ships no `harness/` of its own for exactly that reason. Its four files told a reader that `build.rs` auto-detects `121a` here, so a row from an unpinned build is not comparable with the pinned records. After this branch that is backwards: auto-detection resolves to plain `121`, the same list those records pin. Each now keeps its pin and its reason, states that the pin no longer corrects a shape but still names what was compiled, and identifies the rows that are genuinely suspect, those recorded between #1934 and #1943 on an unpinned build, rather than implying every auto-detected build is. The record prose under `docs/benchmark_results/` and the `TECHNICAL_REPORTS` entry are left as written; they document what was true when they were measured. Refs #1943.
5 of 6 tasks
inureyes
added a commit
that referenced
this pull request
Sep 21, 2026
#1946) * chore(build): stop auto-detect at the shipped Blackwell spelling With MLX_CUDA_ARCHITECTURES unset, `sm_arch_with_suffix` appended CUDA's architecture-specific `a` suffix to every SM from 90 up, so a default `cargo build --features cuda` on a GB10 resolved the whole MLX target to `121a`. That is the whole-target form PR #1938 decided against: every translation unit compiles under `__CUDA_ARCH_SPECIFIC__`, CUTLASS keys `CUTLASS_ARCH_MMA_SM121A_ENABLED` on the same macro, and the per-source injection in `src/lib/mlx-cpp/CMakeLists.txt` never fires at all, because an entry already carrying `a` is recorded as architecture-specific and skipped rather than given a duplicate `--generate-code`. A local build and a shipped build then differ in the machine code their decode kernels run, not merely in architecture coverage. Both functions that spell a list from a detected capability now stop suffixing at a shared `FIRST_PLAIN_SM` boundary: `sm_arch_with_suffix` in the build script, and `CudaArchMismatch::suggested_architecture`, whose startup message was telling an operator on a Blackwell card to rebuild with `MLX_CUDA_ARCHITECTURES=121a`, the one value this repository rejects everywhere else. Hopper keeps `90a` because `release.yml` ships `90a`. The stale rationale for it is replaced rather than repeated: upstream MLX commit `44540d12` moved `qmm_sm80`, `qmm_sm90` and `gather_gemm` to runtime NVRTC compilation and deleted the `MLX_CUDA_SM90A_ENABLED` definition the old comments cited, and `jit_module.cpp` now derives the NVRTC `--gpu-architecture` from the running device, appending `a` itself from compute capability 9 up. `check_cuda_arch_lists.py` grows rule 6 for the auto-detect path, which rule 3 cannot see because no workflow writes that list down, and `check_cuda_arch_lists_test.sh` grows three cases for it. Refs #1943, #1934, PR #1938. * docs(installation): describe one CUDA architecture rule, not two The section stated the auto-detect rule as suffixing every SM from 90 up, citing a Hopper quantized-kernel gate, and then told operators not to set `a` on Blackwell and that building `121a` is a step backwards. Both passages described the same code path in opposite terms, and `check_docs` in the guard script only verifies that the release lists appear verbatim, so nothing caught it. It now states the rule once, the way the build script and the release workflow both spell it, and explains what `121a` costs and why the per-source injection makes it unnecessary. The Hopper paragraph records what is actually true at the current MLX pin: upstream `44540d12` moved the Hopper quantized kernel to runtime NVRTC compilation and removed the `MLX_CUDA_SM90A_ENABLED` definition this section used to cite, `90a` is kept because both release lists ship it, and `CUTLASS_ARCH_MMA_SM90A_ENABLED` still keys on `__CUDA_ARCH_FEAT_SM90_ALL` so a later pin can make it matter again. Refs #1943. * chore: state the Hopper cross-compile result precisely The first pass said the quantized translation units produce identical SASS at `90` and `90a`. Measured across all five, that was too strong in one direction and too weak in another. `qmm_sm90.cu`, `qmm.cu` and `qmm_sm80.cu` emit no device function at either spelling, because upstream `44540d12` moved those kernels to runtime NVRTC compilation; their objects carry 21 lines of fatbin header and zero kernel symbols. `qmv.cu` (2898 kernels) and `fp_qmv.cu` (54) do emit device code, and it is instruction-identical, but their cubins are not byte-identical: the `90a` build adds `EF_CUDA_ACCELERATORS` to the `.headerflags` line, which is the ELF flag marking a cubin architecture-specific. So the suffix changes no instruction at this pin, which is what the decision rests on, and the text now says that rather than implying the objects match byte for byte. Also drops the two remaining statements of the pre-#1943 rule that the new guard does not read: the inline comment in `detect_cuda_arch`, and the `90a is load-bearing` rationale on the x86_64 release job, which cited the same deleted macro. The architecture lists themselves are untouched. Refs #1943. * docs(build): say why the no-detection fallback stays 90a The issue tied the `90a` fallback literal to the Hopper decision, and the cross-compile shows the suffix buys no instruction at this pin, so leaving the literal in place needed a reason rather than silence. The reason is the same one that keeps Hopper suffixed at all: `release.yml` ships `90a`, and a fallback spelled `90` would hand a host without `nvidia-smi` a build shape no published archive uses, which is the divergence this issue closed on Blackwell. The doc comment on `resolve_cuda_architectures` said only that `90a` is the last resort. Refs #1943. * docs(core): stop the compiled-list example showing a rejected shape The doc comment on `compiled_cuda_architectures` illustrated the string with `"80;86;89;90a;100a-real;100;120a-real;120"`, which names Blackwell architecture-specific twice. A doc comment on a parser is where a reader copies a list from, and after this branch the project's own guard rejects that shape, so the file's example and the file's rule disagreed. It now shows the real x86_64 release list, which exercises the same parser features that mattered, pre-Blackwell plain entries and a Hopper `a` entry, and cannot drift into recommending something the build refuses. The `-real` and `-virtual` forms are still described on `parse_cuda_arch_entry` and round-tripped in `entry_display_round_trips_through_the_parser`, which is where an illustration of the grammar belongs. The same pass corrects the header on the release-list constants in the test module, which still justified Hopper's `90a` by the quantized-kernel gate that upstream deleted. Refs #1943. * chore(ci): correct the reasons that named the pre-#1943 auto-detect shape Four comments justified themselves by what auto-detection used to produce. Three CUDA jobs in `ci.yml` said their pin exists because auto-detection "yields `121a` here"; after this branch it yields `121`. Each now states the reason that survives, that a gate should name the list it compiles rather than infer it from whichever GPU the runner has, so the pin keeps a rationale instead of becoming a bare value someone deletes as redundant. `scripts/bench_block_width.sh` is the one that inverted rather than going stale. It told a reader that a row measured on an auto-detected build is not comparable with records pinned at plain `121`, which after this branch is backwards, since an auto-detected GB10 build now resolves to exactly `121`. A wrong comparability rule in a benchmark harness produces discarded runs or published numbers someone trusts, so it now says which rows are actually suspect: those recorded between #1934 and #1943 on an unpinned build. The aarch64 release job carried the same deleted-macro rationale the x86_64 one did, missed in the earlier pass because only one of the two was grepped. Refs #1943. * chore(bench): correct the reused width harness on auto-detected builds The `draft-block-width-gb10-2026-09-20/harness/` directory is a live tool despite its dated path: the post-#1939 width sweep and the #1943 architecture A/B both re-used it, and `draft-block-width-post-1939-gb10-2026-09-21/` ships no `harness/` of its own for exactly that reason. Its four files told a reader that `build.rs` auto-detects `121a` here, so a row from an unpinned build is not comparable with the pinned records. After this branch that is backwards: auto-detection resolves to plain `121`, the same list those records pin. Each now keeps its pin and its reason, states that the pin no longer corrects a shape but still names what was compiled, and identifies the rows that are genuinely suspect, those recorded between #1934 and #1943 on an unpinned build, rather than implying every auto-detected build is. The record prose under `docs/benchmark_results/` and the `TECHNICAL_REPORTS` entry are left as written; they document what was true when they were measured. Refs #1943.
6 of 7 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
1 participant
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
Greedy DFlash output on
qwen3.5-4b-4bitwithqwen3.5-4b-dflashwas not byte-identical to classic decode at any verify width. The cause is not on the candidate list issue #1935 gave: the burst prefilled the prompt throughforward_speculative, which sends every full-attention layer throughattend_per_position. That branch is right for a verify block, whose rows must match what single-token decode would produce, and wrong for a prompt, which classic decode runs as one batched causal attention. The burst therefore entered its first round holding different KV and gated-delta state than classic decode, at every width alike.The MTP path already carries this fix.
forward_prefill_with_last_hiddenexists becauseforward_speculative"does not match the classic prefill numerics", in its own words since #1182. The DFlash path never got it, becauseSpeculativeTarget::prefill_forward_with_capture_layersdefaults to the verify forward and Qwen 3.5 did not override it. At PR #1795's merge commite391ae9cthe same default was already in place, so the byte-identity that record reports at widths 2, 4 and 8 never held rather than having broken since.A second, independent defect remains at widths 8 and 16 and is gated rather than fixed: the verify block there is genuinely not byte-identical to the single-token chain (229971 of 496640 logit bytes at block position 0), because
M * B >= 8leaves MLX'sqmvfamily forqmm_sm80. That is a real difference between two kernels, and the exactness probe the other three DFlash families already run is the right answer to it.What changed
Qwen35Model::forward_prefill_with_capture_layersrunsforward_internal's loop while capturing the per-layer hidden states the drafter is seeded from, and captures no gated-delta rollback snapshots, which a prefill never needs and which are prompt-sized. BothSpeculativeTargetimpls for the family override the prefill hook to call it, covering the B = 1 and B > 1 burst arms.Qwen35Model::dflash_exactness_allowsruns the block-versus-chain probe through the sharedmtp_exactness_gate, and both Qwen 3.5DFlashTargetModelimpls call it. Deliberately notmtp_exactness_allows, whosemetal_is_available()precondition would decline every CUDA burst without measuring one.greedy_parity_dflash_qwen35_4bcovers widths 2, 4, 8 and 16 instead of only 16, which the gate now declines. Each width must either run a burst and match the drafter-less baseline byte for byte, or be declined by the gate; the helper names which happened rather than letting a decline read as a pass.#[ignore]-gated real-checkpoint diagnostics insrc/models/qwen3_5_dflash_probe_tests.rs. They are what located the divergence and they stay so the separation can be re-run.docs/benchmark_results/, and theMLXCEL_MTP_ALLOW_INEXACTrow now says it governs this path too.Measured on GB10
In process against a real 200-token classic transcript, after the fix: the burst's prefill and classic decode's prefill agree in all 496640 logit bytes of the last row and in every layer's cache state, and the burst shape reproduces classic decode's greedy argmax at 0 of 201 positions disagreeing at widths 2 and 4, including 100 rewinds with the rejected rows carrying a wrong token id and across six accept patterns that mix full-accept rounds with partial ones.
Served, one server per arm: widths 8 and 16 decline and return text byte-identical to the baseline, widths 2 and 4 engage, and
MLXCEL_MTP_ALLOW_INEXACT=1at width 16 still engages with its forfeit warning.What this does not fix
Widths 2 and 4 still part from classic decode at generated token 105, down from token 9. That is its own finding, not a remnant of the prefill bug, which was masking it. It is invariant to
MLXCEL_QMV_MULTIROW=0,MLX_USE_CUDA_GRAPHS=0,--no-warmupand omittinglogprobs, and widths 2 and 4 produce identical text despite entirely different drafting patterns.The exclusion is stronger than "not reproduced in process". This branch adds a
debugline carrying the burst's per-round accept lengths, and replaying the served width-4 run's own 79 rounds and 56 rewinds in process, at the same depths in the same order and with the served capture list, 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 target's prefill, verify or rewind under the served run's own cache history, and one hypothesis is left standing: the drafter executing MLX work between the target's forwards, which every arm here omits. The next step is to driveDFlashGenerator::runin process with the real drafter bound.The issue's last acceptance criterion is therefore not met, which is why this does not close it. Concretely,
greedy_parity_dflash_qwen35_4bfails on this branch at widths 2 and 4 and passes at 8 and 16 by decline. That red is the open finding above, not a regression this PR introduces: onmainthe same test fails at all four widths, and every width's failure there is the prefill bug.Test plan
cargo fmt --all -- --checkcargo clippy -p mlxcel --lib --tests -- -D warningson the 1.97.1 toolchainrust-toolchain.tomlpins: cleancargo test --release --features cuda -p mlxcel --lib -- --test-threads=1scoped toqwen3_5,dflashandspeculative_exactness: 30, 22 and 11 passed, 0 failedcargo test --release --features cuda -p mlxcel --test speculative_parity --no-run: compilesNV_ERR_NO_MEMORYfor the sessionmake verify-test-cudawas not run locally:--profile test-fastgivesmlxcel-coreits own build-script unit and recompiles MLX from source rather than reusing the warm release build.cargo check --features cuda,xla-iree --all-targetscannot run on this host at all, becausemlxcel-xla's build script aborts withIREE_DIST is not setand CI provisions that as a job step; CI is the check for it.Refs #1935