diff --git a/docs/benchmark_results/dflash-verify-divergence-qwen35-gb10-2026-09-20.md b/docs/benchmark_results/dflash-verify-divergence-qwen35-gb10-2026-09-20.md new file mode 100644 index 000000000..c56ddb65c --- /dev/null +++ b/docs/benchmark_results/dflash-verify-divergence-qwen35-gb10-2026-09-20.md @@ -0,0 +1,141 @@ +# Qwen 3.5 DFlash verify divergence from classic decode (GB10, 2026-09-20) + +Issue #1935, from the sweep recorded in `draft-block-width-default-gb10-2026-09-20.md` (issue #1797, PR #1937). That sweep found greedy speculative output on `models/mlx/qwen3.5-4b-4bit` with `models/mlx/qwen3.5-4b-dflash` not byte-identical to classic decode at any verify width, identically at 2, 3, 4, 5, 6, 7, 8 and 16, and could not separate the candidates. This record separates them. + +Nothing here is a throughput measurement. Every arm compares greedy token streams, which do not depend on host load, so the arms were deliberately run while the host was busy with other sessions' builds and with mlxcel CI. Timing numbers appear only where they are evidence that a kill switch was live, and are labelled as such. + +## Host + +``` +host: spark-101 +kernel: 7.0.0-1019-nvidia +arch: aarch64 +nvidia_driver: 580.178.04 +gpu: NVIDIA GB10, compute capability 12.1 +cuda_toolkit: 13.0 +mlx_pin: 81ba1c6a0e50a9268b931579c2d4f1158b9aab5a +rustc: 1.97.1 (8bab26f4f 2026-07-14), the version rust-toolchain.toml pins +mem_total_gib: 121.7 +``` + +Driver budget: the boot's cumulative kernel `NV_ERR_NO_MEMORY` count was 0 before this session and 0 after it, across roughly thirty server starts and a dozen in-process model loads. That matches the zero delta the #1797 sweep measured on the same driver and contradicts the 38-per-run figure from `580.173.02`, which is why the next long sweep should still calibrate from its own first run rather than inherit either number. + +## Method + +Served arms: one `mlxcel-server` per arm, `--ignore-eos --max-batch-size 1`, a fixed 158-token Python prompt (the #1797 harness's `prompt_retry.txt`), one non-streaming `POST /v1/completions` at `temperature 0` with `max_tokens 200`, compared by the sha256 of the completion text and by the per-token list the `logprobs` field returns. In-process arms: `#[ignore]`-gated tests in `src/models/qwen3_5_dflash_probe_tests.rs`, which load the real checkpoint and drive the model directly. + +The classic reference is anchored twice. Two classic served arms at the start and end of the session are byte-identical to each other (`2c76b0a181`), and `MLXCEL_PRINT_TOKEN_IDS=1 mlxcel generate --no-chat-template --temp 0` on the same prompt produces the same 200 tokens byte for byte. It also supplies the prompt and reference token ids the replay arms take. + +Every in-process arm below uses `Qwen35Model::forward_internal` as the stand-in for served classic decode, and that stands on the call graph rather than on the equality above. The scheduler's B = 1 decode calls `LanguageModel::forward_with_sequence_id` (`src/server/batch/scheduler/decode_tick.rs`), whose Qwen 3.5 impl is `forward_with_sequence_caches(input_ids, None, seq_id)`, which resolves the MRoPE entry and calls `forward_internal`. The offline CLI reaches the same function through `LanguageModel::forward`, which is the same call with `seq_id = None`. The two endpoints agreeing on 200 argmaxes is then a consistency check on that reading, not the reading itself. + +## The root cause: the burst prefilled the prompt as if it were a verify block + +`Qwen35Model::forward_speculative` sends every full-attention layer through `Qwen3NextAttention::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 see. That is correct for a verify block and is not the computation 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 prefilled the prompt the per-position way. It therefore entered its first round holding different KV and gated-delta state than classic decode, and its greedy output drifted away within a few tokens at every width alike. + +Measured, before the fix: the two prefills' last-row logits differ in 173206 of 496640 bytes, and the first cache state to differ is layer 4's, the layer after the first full-attention layer (`full_attention_interval` is 4, so layers 0 to 2 are gated-delta and layer 3 is the first attention layer). Layers 0 to 3 agree because their inputs are identical; layer 3's attention output differs without its own K and V differing, which is what puts the difference in the attention computation rather than in a projection. + +The MTP path already carries this fix. `Qwen35Model::forward_prefill_with_last_hidden` exists for exactly this reason, and its doc comment has said so since #1182: `forward_speculative` "does not match the classic *prefill* numerics", and "the classic path's first token is sampled from the batched causal prefill logits, so MTP's first bonus must come from the same computation to stay byte-identical". The DFlash path never got it. + +### How long this has been true + +`attend_per_position` landed on 2026-05-24 (`4038da96`, issue #78), eleven days after the B = 1 speculative burst was wired (`ca938b67`, 2026-05-13). At PR #1795's merge commit `e391ae9c` (2026-09-11), `Qwen35DecoderLayer::forward_with_capture` already called `forward_with_position_ids_verify(..., true)` and `Qwen35Model` still did not override `prefill_forward_with_capture_layers`, so the burst prefilled the same way then. + +This answers issue #1935's first question without a rebuild at that commit, and more conclusively than a rebuild would: the property did not break after PR #1795, it never held as PR #1795's record stated. `dflash-verify-fixed-cost-gb10-2026-09-11.md` reports byte-identity at widths 2, 4 and 8 on this pairing, and that claim should be read as not established. + +## What the fix changes + +`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 (a prefill is never rolled back, and they are prompt-sized rather than block-sized). Both `SpeculativeTarget` impls for the family override the hook to call it, so the B = 1 and B > 1 burst arms are both covered. + +In process, against the real 200-token classic transcript: + +| arm | prefill last-row logit bytes differing | greedy argmax positions disagreeing, of 201 | +|---|---:|---:| +| before the fix | 173206 of 496640 | not measured; the served arms parted at token 9 | +| after, block 4 | 0 | 0 | +| after, block 2 | 0 | 0 | +| after, block 8 | 0 | 2, first at 93 | + +After the fix every layer's cache state is byte-identical between the two prefills, not merely the logits. The width-8 disagreement is the block-versus-chain divergence described below and is not something a prefill fix could reach. + +One thing about that cache comparison is worth stating because the first version of the test got it wrong: it has to run immediately after both prefills, before the classic arm's single-token loop advances its caches. Run afterwards it reports the offset gap the chain loop itself opened (358 against 158) and reads as a layer-0 divergence that is not there. + +Served, same session, same binary per row: + +| arm | completion sha256 | first token differing from classic | +|---|---|---:| +| classic (twice, opening and closing) | `2c76b0a181` | | +| before the fix, widths 2 and 4 | `57d291aca2` | 9 | +| before the fix, widths 8 and 16 | `57d291aca2` | 9 | +| after the fix, widths 2 and 4 | `3e60b1574c` | 105 | +| after the fix, widths 8 and 16 | `57d291aca2` | 9 | + +Widths 8 and 16 landing back on the pre-fix hash is not the fix failing to apply there. Their verify block is independently divergent, the flip at token 9 is a near tie (the classic choice carries logprob -0.875 against -0.750 for the speculative one), and once a greedy trajectory flips at a near tie the rest of it follows deterministically, so two different causes reach the same completion. + +## The second, separate defect: widths 8 and 16 are genuinely not exact + +`Qwen35Model::probe_block_chain_exactness`, run on the real checkpoint, reports: + +| width | verdict | +|---|---| +| 2 | byte-identical to the single-token chain | +| 4 | byte-identical to the single-token chain | +| 8 | differs at block position 0 in 229971 of 496640 logit bytes | +| 16 | differs at block position 0 in 229971 of 496640 logit bytes | + +The boundary is `if (can_use_qmv && (M * B < 8))` in `mlx/backend/cuda/quantized/quantized.cpp`: a verify block of 8 or more rows leaves the `qmv` family for `qmm_sm80`, and the two are not bit-equal. That is a real numerical difference between two kernels rather than a defect, and it is the failure class the exactness probe exists for. It is also what the #1797 record's forced-`MLX_ENABLE_TF32=0` pass was seeing when its speculative arms split into 2 through 7 against 8 and 16; that grouping was a correct observation of a second effect, described there as if it were the only one. + +`Qwen35Model::dflash_exactness_allows` now runs that probe through the shared `mtp_exactness_gate`, and both Qwen 3.5 `DFlashTargetModel` impls call it. Deliberately not `mtp_exactness_allows`, whose first precondition is `metal_is_available()` and which would therefore decline every CUDA burst without measuring one. + +Served, with the gate in place: + +| arm | completion sha256 | bursts completed | what happened | +|---|---|---:|---| +| classic | `2c76b0a181` | 0 | | +| width 2 | `3e60b1574c` | 2 | probe passed, burst engaged | +| width 4 | `3e60b1574c` | 2 | probe passed, burst engaged | +| width 8 | `2c76b0a181` | 0 | probe declined, served by classic decode, byte-identical to classic | +| width 16 | `2c76b0a181` | 0 | probe declined, served by classic decode, byte-identical to classic | +| width 16, `MLXCEL_MTP_ALLOW_INEXACT=1` | `57d291aca2` | 2 | override engaged the burst and logged the forfeit | + +## What remains, and what it is not + +Widths 2 and 4 still part from classic decode at generated token 105. This is its own finding and it does not share a cause with the prefill: the prefill is now byte-exact, and the prefill bug was masking this one rather than producing it. + +It is reproducible and deterministic, and it is invariant to everything cheap. Widths 2 and 4 produce identical completions despite drafting 1 and 3 proposals per round over 120 and 79 rounds, so it does not track the drafter's proposal pattern, its round boundaries, or its allocation pattern. It is unchanged by `MLXCEL_QMV_MULTIROW=0`, by `MLX_USE_CUDA_GRAPHS=0`, by `--no-warmup`, and by omitting `logprobs` from the request. + +Against that, the in-process target path is exhaustively exact over the same transcript at widths 2 and 4: + +- prefill: every layer's cache state byte-identical to `forward_internal`'s; +- verify block: byte-identical to the single-token chain at every one of the 32 layers at block position 0, and 0 of 201 greedy argmax positions disagreeing over the full transcript; +- rollback: 0 disagreements across accept patterns `1,2,3`, `1,2,3,4`, `4,1`, `4,4,1`, `1,4,2,4,3` and `2,4,1,3,4,4,1,1`, which mix full-accept rounds that never rewind with partial ones that do, and with the rejected rows carrying a wrong token id (9999) the way a real round's rejected rows carry the drafter's wrong proposals; +- rollback again, against the served run's **own** accept sequence rather than a synthetic cycle: the width-4 server logs its per-round accept lengths (a `debug` line this change adds), and replaying those 79 rounds in process reproduces the served burst's exact cache history, 56 rewinds at the same depths in the same order. It disagrees with the chain at 0 of 201 positions, with the rejected rows wrong and with them correct alike; +- with the served capture list `[1, 8, 15, 22, 29]` rather than no capture, and with the round loop's own whole-block `argmax_last_axis` rather than one call per sliced row: 0 disagreements either way. + +The last of those is the strongest exclusion available short of running the real drafter: the served burst's round structure is not an approximation of it any more, it is the same 79 rounds with the same rewind depths, and the target reproduces classic decode through all of them. Rejected-row content is excluded as well, since feeding the correct tokens and feeding 9999 give the same answer. + +So the residual is in the served burst wrapper rather than in the target's forward, verify or rewind, 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 drive `DFlashGenerator::run` in process with the real drafter bound and bisect from there, which is the one configuration these arms do not reproduce. That is a test worth building and is not in this change. + +## Controls, including one that was void + +**Gated-delta chunked scan: excluded.** `MLXCEL_GDN_CHAIN_PARITY=0` swaps the verify block's bit-exact sequential loop for the 64-row chunked scan. Before the prefix fix the switch was live (burst 5511 ms against 7826 ms, acceptance 0.5325 against 0.5281) and changed the completion not at all, which is what excluded the scan as the cause of the token-9 divergence. After the fix the same switch does change the completion (`3e60b1574c` against `08f8cfb1fe`), which is the expected direction: it forfeits an exactness the fix had just made load-bearing. Both readings belong in the record, because the first one on its own would suggest the scan can never matter. + +**Multirow quantized matmul: excluded, twice.** `MLXCEL_QMV_MULTIROW=0` routes a 4-row verify through the same `qmv_kernel` classic decode uses. The completion is unchanged before the fix and unchanged after it. + +**TF32: not the discriminator.** `MLX_ENABLE_TF32=0` moves both the classic arm (`92e90219fd`) and the width-4 arm (`f89987e15e`), and does not bring them together. + +**`MLXCEL_SDPA_FALLBACK_MAX_QUERIES=0`: void, not negative.** This switch diverts a 2-to-32-query masked SDPA call from cuDNN to MLX's ops fallback. On this checkpoint it can do nothing: `head_dim` is 256 and `cuda_sdpa_materializes_scores` records that cuDNN's flash SDPA requires `head_dim <= 128`, so these calls never reach cuDNN and there is nothing to divert. The arm changed neither the completion nor the burst time (5511 ms against 5480 ms), and the unchanged time is how it was identified as inert. Recording it as "no effect" would have read as evidence that the verify's attention dispatch does not matter, which it is not. + +**`MLXCEL_SDPA_VECTOR_LARGE_D=0`: live, and informative about the shape of the problem rather than the cause.** It moves classic decode's single-row attention off the fused `sdpa_vector` kernel, and classic's completion changes (`4c37650547`). So this pairing's greedy output is not robust to the attention kernel it decodes with, which is worth knowing before reading any byte-identity claim about it. + +## Files + +Diagnostics: `src/models/qwen3_5_dflash_probe_tests.rs`, six `#[ignore]`-gated tests. The three that take a recorded transcript read it from `MLXCEL_Q35_PROBE_PROMPT` and `MLXCEL_Q35_PROBE_REFERENCE`; `MLXCEL_PRINT_TOKEN_IDS=1 mlxcel generate` prints both lines. To replay a served burst's own round structure, read its accept lengths from the server at `RUST_LOG=mlxcel::server::batch::dflash_target=debug`, add one to each (the log reports accepted draft tokens, the test takes kept rows) and pass them as `MLXCEL_Q35_PROBE_ACCEPTS`. + +``` +cargo test --release --features cuda -p mlxcel --lib -- --ignored --test-threads=1 \ + --nocapture qwen3_5_dflash_probe_tests +``` + +`--test-threads=1` is required for every CUDA test run on this host. diff --git a/docs/environment-variables.md b/docs/environment-variables.md index 48fcdea93..82ba324d4 100644 --- a/docs/environment-variables.md +++ b/docs/environment-variables.md @@ -325,7 +325,7 @@ The OpenAI audio endpoints (`/v1/audio/speech`, `/v1/audio/transcriptions`, `/v1 | `MLXCEL_ENABLE_MTP_BATCH` | truthy value | off | **Advanced.** Forces the batched Gemma 4 MTP burst path for parity/debug testing. Not governed by the adaptive policy (issue #333), which scopes to the validated, byte-identical B=1 path. | | `MLXCEL_METAL4_ATTENTION` | `0`/`false`/`no`/`off` to disable; unset or any other value to enable | on (where the hardware has it) | **Advanced, diagnostic kill switch.** Forces the M5 neural-accelerator fused attention route off on hardware that has it, so `layers::metal4_causal_attention` is skipped and the ordinary SDPA path runs instead. Off-switch only: it cannot turn the route on where `has_neural_accelerator && macos_supports_na` is false, so setting it on an M1 is inert. This route is the first suspect whenever an M5 disagrees numerically with an earlier Apple GPU generation (issue #1065), and before this switch the only way to A/B the hypothesis was to patch `should_use_metal4_attention` and rebuild, which is what the #1182 M5 investigation had to do. Read once per process, so set it before starting `mlxcel` or `mlxcel-server`. Inert on non-Metal builds. | | `MLXCEL_GDN_CHAIN_PARITY` | `0` to disable, any other value (or unset) to enable | on | **Advanced, diagnostic escape hatch.** Gates the chain-parity gated-delta Metal kernel used by Qwen 3.5 MTP's speculative verify and rollback-replay paths (issue #1165). The standard gated-delta kernel carries float32 recurrent state across a `T = K` verify block and rounds it to the storage dtype only once at the end, while the classic single-token decode chain rounds after every token; a `T = K` verify block is therefore NOT bit-identical to `K` consecutive single-token decode steps unless the state is rounded per in-block step. The chain-parity kernel (`gated_delta_step_seqpar`) does that rounding, which is what makes Qwen 3.5 MTP's temperature-0 output byte-identical to classic decode. **Setting this to `0` forfeits that exactness contract**, restoring the pre-#1165 block numerics for A/B attribution of the parity kernel's own cost and acceptance effect; do not set it to `0` in a deployment that needs byte-identical speculative output. On the non-Metal ops fallback (CUDA, CPU) the flag selects, for an unmasked scalar-gated verify block of at most 32 rows (`CHAIN_PARITY_SEQUENTIAL_MAX_T`), the sequential loop of fused single steps instead of the 64-row chunked scan (issue #1782). That loop is the same arithmetic as the classic decode chain, so it is bit-identical to `T` single-token steps, and it launches far fewer kernels: the chunked scan pays a fixed 63-step Horner loop per linear-attention layer whatever the row count (about 4400 launches and a fifth of the device time per 2-row verify round on GB10 with `qwen3.5-4b-4bit`, twice that when a partial accept replays through it). `0` restores the chunked scan there for A/B attribution. **The kernel is necessary but not sufficient**: byte-identity also requires every quantized projection to dispatch to the same MLX kernel at `M = block_size` as at `M = 1`, which is not true on every GPU generation or at every block width, so the runtime probe behind `MLXCEL_MTP_ALLOW_INEXACT` is what actually decides whether MTP engages. See `docs/benchmark_results/qwen38-mtp-m1ultra-2026-08-16.md` for the measured kernel cost (inside the dispatch-noise band). | -| `MLXCEL_MTP_ALLOW_INEXACT` | `1`/`true`/`yes`/`on` to enable; unset or anything else to disable | **off** | Engage Qwen 3.5 MTP speculative decoding even when the startup exactness probe reports that the multi-token verify block is **not** byte-identical to the single-token decode chain. Before enabling MTP the runtime now measures the property instead of predicting it: one synthetic verify block and the equivalent single-token chain are run from the same state on the loaded checkpoint at the configured `--draft-block-size`, and their logits are compared byte for byte (three independent synthetic inputs, each two short prefills plus `K + 1` forwards; measured 4.9 s for the first call and 1.3 s for a later one per input on a Qwen3.8-27B 4-bit target on an M1 Ultra, the difference being MLX's one-time kernel compilation; more than one input because a kernel pair can disagree by only a byte or two out of ten thousand, at which amplitude a single draw can read as equal; memoized per (model, block width) and warmed at worker startup so it never lands on the request path). A divergence means temperature-0 speculative output would silently differ from `mlxcel generate` without `--draft-model`, so the default is to decline and run classic decode. The static conditions (Metal backend, `supports_metal_gated_delta_kernel` geometry) still apply and are checked first; this probe covers what they cannot, namely which MLX kernel each **quantized projection** dispatches to at `M = K` versus `M = 1`. That choice depends on the GPU generation, the quantization mode, the operand sizes and the block width: `use_qmv_wide` in [`mlx/backend/metal/quantized.cpp`](https://github.com/ml-explore/mlx/blob/main/mlx/backend/metal/quantized.cpp) sends `M >= 2` to a different reduction whenever `mode != "affine" || arch_gen >= 15`, and `get_qmv_batch_limit` sends `M` above 10, 12, 18 or 32 (by architecture size and generation) to the matrix-matrix kernel. Measured: an affine 4-bit Qwen3.8-27B target on an M1 Ultra is byte-identical at block widths 1 through 11 and diverges at 12 (the `arch_size == 'd'` branch); the same checkpoint on an M5 Max diverges from block width 2, where the `M >= 2` split fires before any batch limit can be observed. Within one checkpoint the limit is per projection, not per model: Gemma 4 12B's attention shapes hold to 17 on an M1 Ultra while its MLP shapes break at 12, so a model's own cliff is the minimum over its shapes, which is why this is measured rather than tabulated. Note the ordering that #1199 introduced: on a failing probe the gate first retries with `qmv_wide` disabled and keeps the narrow kernel when that restores exactness, and only a probe that fails **both** ways consults this flag. On Apple GPU generation 15+ the narrow retry passes, so this flag alone is inert there: the process is pinned narrow, output stays byte-identical, and the log shows the retry's INFO line rather than the override warning (verified live on M3 Ultra, 2026-08-22, byte-identical output with and without the flag; see `benchmark_results/qmv-wide-pin-tax-m3ultra-2026-08-22.md`). To research the fast kernel there, set `MLXCEL_QMV_WIDE=1` together with this flag: the pin skips the retry and this flag then engages MTP on the wide kernel, forfeiting byte-identity with the loud WARN. This flag alone is load-bearing only where no exact kernel selection exists at the configured block width. Read once per process. | +| `MLXCEL_MTP_ALLOW_INEXACT` | `1`/`true`/`yes`/`on` to enable; unset or anything else to disable | **off** | Engage Qwen 3.5 MTP speculative decoding even when the startup exactness probe reports that the multi-token verify block is **not** byte-identical to the single-token decode chain. Before enabling MTP the runtime now measures the property instead of predicting it: one synthetic verify block and the equivalent single-token chain are run from the same state on the loaded checkpoint at the configured `--draft-block-size`, and their logits are compared byte for byte (three independent synthetic inputs, each two short prefills plus `K + 1` forwards; measured 4.9 s for the first call and 1.3 s for a later one per input on a Qwen3.8-27B 4-bit target on an M1 Ultra, the difference being MLX's one-time kernel compilation; more than one input because a kernel pair can disagree by only a byte or two out of ten thousand, at which amplitude a single draw can read as equal; memoized per (model, block width) and warmed at worker startup so it never lands on the request path). A divergence means temperature-0 speculative output would silently differ from `mlxcel generate` without `--draft-model`, so the default is to decline and run classic decode. The static conditions (Metal backend, `supports_metal_gated_delta_kernel` geometry) still apply and are checked first; this probe covers what they cannot, namely which MLX kernel each **quantized projection** dispatches to at `M = K` versus `M = 1`. That choice depends on the GPU generation, the quantization mode, the operand sizes and the block width: `use_qmv_wide` in [`mlx/backend/metal/quantized.cpp`](https://github.com/ml-explore/mlx/blob/main/mlx/backend/metal/quantized.cpp) sends `M >= 2` to a different reduction whenever `mode != "affine" || arch_gen >= 15`, and `get_qmv_batch_limit` sends `M` above 10, 12, 18 or 32 (by architecture size and generation) to the matrix-matrix kernel. Measured: an affine 4-bit Qwen3.8-27B target on an M1 Ultra is byte-identical at block widths 1 through 11 and diverges at 12 (the `arch_size == 'd'` branch); the same checkpoint on an M5 Max diverges from block width 2, where the `M >= 2` split fires before any batch limit can be observed. Within one checkpoint the limit is per projection, not per model: Gemma 4 12B's attention shapes hold to 17 on an M1 Ultra while its MLP shapes break at 12, so a model's own cliff is the minimum over its shapes, which is why this is measured rather than tabulated. Note the ordering that #1199 introduced: on a failing probe the gate first retries with `qmv_wide` disabled and keeps the narrow kernel when that restores exactness, and only a probe that fails **both** ways consults this flag. On Apple GPU generation 15+ the narrow retry passes, so this flag alone is inert there: the process is pinned narrow, output stays byte-identical, and the log shows the retry's INFO line rather than the override warning (verified live on M3 Ultra, 2026-08-22, byte-identical output with and without the flag; see `benchmark_results/qmv-wide-pin-tax-m3ultra-2026-08-22.md`). To research the fast kernel there, set `MLXCEL_QMV_WIDE=1` together with this flag: the pin skips the retry and this flag then engages MTP on the wide kernel, forfeiting byte-identity with the loud WARN. This flag alone is load-bearing only where no exact kernel selection exists at the configured block width. Read once per process. Since issue #1935 the same probe and the same override also govern the **Qwen 3.5 DFlash** burst through `Qwen35Model::dflash_exactness_allows`, which is the measured half alone (it does not require Metal, because the DFlash burst runs on CUDA). Measured on GB10 with an affine 4-bit Qwen 3.5 4B target: the verify block is byte-identical to the chain at widths 2 and 4 and is not at 8 and 16, where `M * B >= 8` leaves MLX's `qmv` family for `qmm_sm80`, so an unmodified server engages the burst at the measured default of 4 and declines to classic decode at 16. | | `MLXCEL_QMV_WIDE` | `0`/`false`/`no`/`off` to disable; `1` (or any other value) to pin wide | unset (wide, until the MTP gate's retry turns it off) | Operator pin for MLX's `qmv_wide` kernel, the faster reduction for `M >= 2` quantized matmuls on Apple GPU generation 15+ (overlay in `src/lib/mlx-cpp/patches/mlx/backend/metal/quantized.cpp`, added by #1199). **Setting the variable at all, to any value, counts as an operator pin**: the MTP exactness gate's retry (`retry_without_qmv_wide`) is skipped in both directions, so `MLXCEL_QMV_WIDE=1` keeps the wide kernel and makes a failing probe decline MTP instead of buying exactness back, and `MLXCEL_QMV_WIDE=0` runs the whole process narrow from the start. Unset, the kernel is wide until a failing MTP probe's retry finds the narrow kernel exact and pins the process narrow for good. The pin is process-wide and sits on the dispatch path of every quantized matmul; what non-MTP work pays for the narrow state is measured in `benchmark_results/qmv-wide-pin-tax-m3ultra-2026-08-22.md` (nothing measurable on batched decode, about 15 ms per prompt-cache-hit request's suffix prefill). Read once per process at first dispatch; `mlxcel_core::set_qmv_wide` can move it at runtime and the gate is its only caller. | | `MLXCEL_MTP_BLOCK_CONTROLLER` | `proxy` to pin the acceptance-proxy gate; `requested` to honour the requested width immediately (measurement mode); any other value (or unset) selects the throughput comparator | throughput | Which controller decides the B=1 MTP verify width when `--draft-block-size` exceeds the drafter's configured depth (issue #1207). The default measures the decision: after a short warm-up the round loop alternates measurement windows (32 rounds) between the configured depth and the requested ceiling, compares emitted tokens per millisecond of round time, holds whichever measures faster (a challenger needs a 2% lead; ties go to the configured depth), and re-challenges the loser on a growing backoff (4, 16, then every 64 windows), with a collapsing challenger aborted after 4 rounds once it trails by more than 35% so a harmful ceiling (the Qwen 3.8 pairing measures 5.80 against 21.30 tok/s at width 12) costs rounds rather than windows. Evidence lives in the generator, so a server process keeps it across requests; drafters that set `prefer_requested_block_size` (Qwen 3.5 MTP) bypass both controllers and always honour the request, exactly as before. Set to `proxy` to restore upstream's fully-accepted-prefix gate, which issue #1207 measured holding the Gemma 4 12B pairing about 5% below its optimum (93.54 against 98.16 tok/s at requested width 5 on an M5 Max) because at 0.585 acceptance the configured prefix is rarely fully accepted no matter how profitable widening is. The batched (B>1) loop stays on the proxy gate regardless: the row-averaged accept length is the only per-round signal it measures today. Set to `requested` to bypass both controllers and draft at the requested width from the first round: this is the width-sweep measurement mode (#1207's own sweep needed a temporary code patch to hold widths; this value is that patch as a switch), not a deployment setting, and the harmful-ceiling protection is off under it. Read once per generator. | | `MLXCEL_MTP_TICK_SLICE` | `0`/`false`/`no`/`off` to disable, any other value (or unset) to enable | on | Tick-cooperative B=1 MTP serving (issue #734). When on (the default), a B=1 MTP request on the Gemma 4 family is served one speculative round per scheduler tick, alternating with the classic decode/prefill actions, so concurrent classic-decode rows advance between rounds and the head-of-line stall a speculative request imposes drops from the whole burst to about one round (`burst_wall_ms` in the finalize log reports the max single-tick wall). Tokens stream per round instead of in one end-of-burst lump. Set to an off value to restore the legacy run-to-completion burst (the whole request served inside one tick). The interleaving trades roughly 27% of the speculative request's own aggregate decode throughput (cross-tick round gaps) for that bounded stall, so a deployment serving speculative requests without concurrent classic traffic can turn it off to keep the full-throughput burst. Greedy output, acceptance accounting, and every other env gate are unchanged in both modes; DFlash and the batched B>1 paths always run to completion regardless of this flag. | diff --git a/src/models/qwen3_5.rs b/src/models/qwen3_5.rs index 8bc0e99d0..62b059ee7 100644 --- a/src/models/qwen3_5.rs +++ b/src/models/qwen3_5.rs @@ -1527,6 +1527,40 @@ impl Qwen35Model { mtp_exactness_gate(key, || self.probe_block_chain_exactness(block_size)) } + /// Whether the **DFlash** burst may engage at `block_size` verify rows on + /// this host (issue #1935). + /// + /// Deliberately not [`Self::mtp_exactness_allows`]: that gate's first + /// precondition is `metal_is_available()`, because Qwen 3.5 MTP's + /// byte-identity rests on the Metal chain-parity gated-delta kernel. The + /// DFlash burst runs on CUDA too, where that precondition would decline + /// every burst without measuring one, so this entry point is the measured + /// half alone: [`Self::probe_block_chain_exactness`], memoized per + /// (model, width) per process by + /// [`mtp_exactness_gate`](crate::models::speculative_exactness::mtp_exactness_gate), + /// which also owns the decline log line, the `qmv_wide` retry and the + /// `MLXCEL_MTP_ALLOW_INEXACT` override. + /// + /// What this catches on GB10 with `qwen3.5-4b-4bit`: at widths 2 and 4 the + /// verify block is byte-identical to the chain and the burst engages, and + /// at 8 and 16 it is not (229971 of 496640 logit bytes differ at block + /// position 0), which is the `M * B >= 8` switch from `qmv` to + /// `qmm_sm80` in MLX's CUDA quantized dispatch. That is a real numerical + /// difference between two kernels rather than a defect, so the burst + /// declines there instead of serving text that differs from classic + /// decode. + /// + /// Used by: `DFlashTargetModel::exactness_allows` for `Qwen35Model` and + /// `Qwen35VLModel`, which the server burst gate reads. + pub fn dflash_exactness_allows(&self, block_size: usize) -> bool { + let key = ProbeKey { + block_size: block_size as u32, + hidden_size: self.config.hidden_size as u32, + num_hidden_layers: self.config.num_hidden_layers as u32, + }; + mtp_exactness_gate(key, || self.probe_block_chain_exactness(block_size)) + } + /// Measure, on this loaded checkpoint, whether a `T = block_size` /// verify block produces byte-identical logits to `block_size` /// consecutive single-token decode steps. @@ -2165,6 +2199,105 @@ impl Qwen35Model { } } + /// Prefill `input_ids` through the **standard batched-attention** forward + /// while capturing the per-layer hidden states a DFlash drafter is seeded + /// from (issue #1935). + /// + /// This is deliberately NOT [`Self::forward_speculative`], for the reason + /// [`Self::forward_prefill_with_last_hidden`] already states for the MTP + /// path: `forward_speculative` sends every full-attention layer through + /// `attend_per_position`, which is byte-aligned with single-token *decode* + /// and is not the *prefill* computation classic decode runs. The classic + /// path samples its first token from the batched causal prefill logits and + /// carries that prefill's KV and gated-delta state into every token after + /// it, so a burst that prefills the per-position way starts from a + /// different state and its greedy output drifts away from classic decode + /// within a few tokens, at every verify width alike. Measured on GB10 with + /// `qwen3.5-4b-4bit`: the two prefills' last-row logits differ in 173206 of + /// 496640 bytes and the served completions part company at generated token + /// 9. The MTP path took this fix; the DFlash path did not, which is what + /// issue #1935 reported. + /// + /// GDN rollback snapshots are deliberately not captured: a prefill is never + /// rolled back, and they are prompt-sized rather than block-sized, which is + /// the memory argument `SpeculativeTarget::prefill_forward_with_capture_layers` + /// exists for. The returned `gdn_states` is therefore empty, and + /// `rollback_partial` must never be called against this output. + /// + /// Keep the loop in lockstep with [`Self::forward_internal`] (mask + /// construction, layer dispatch, final norm, LM head); the capture is the + /// only addition. + /// + /// Used by: `SpeculativeTarget::prefill_forward_with_capture_layers` for + /// `Qwen35Model` and `Qwen35VLModel`, which is the DFlash burst's first + /// call on both the B = 1 and B > 1 arms. + pub fn forward_prefill_with_capture_layers( + &self, + input_ids: &MlxArray, + caches: &mut [Qwen3NextCache], + capture_layer_ids: &[usize], + ) -> VerifyOutput { + let mut h = self.embed_tokens.forward(input_ids); + + let shape = mlxcel_core::array_shape(&h); + let seq_len = shape[1]; + + let fa_idx = self.config.full_attention_interval - 1; + let fa_mask = if seq_len > 1 { + let offset = if fa_idx < caches.len() { + caches[fa_idx].offset() + } else { + 0 + }; + Some(create_causal_mask(seq_len, offset)) + } else { + None + }; + + let mut hidden_slots: Vec>> = + (0..capture_layer_ids.len()).map(|_| None).collect(); + + for (i, (layer, cache)) in self.layers.iter().zip(caches.iter_mut()).enumerate() { + let mask = if layer.is_linear { + None + } else { + fa_mask.as_deref() + }; + h = layer.forward(&h, mask, cache, None); + + for (slot_idx, &want_idx) in capture_layer_ids.iter().enumerate() { + if want_idx == i { + hidden_slots[slot_idx] = Some(mlxcel_core::copy(&h)); + } + } + } + + let hidden_states: Vec> = hidden_slots + .into_iter() + .map(|opt| { + opt.unwrap_or_else(|| { + mlxcel_core::zeros( + &[shape[0], seq_len, self.config.hidden_size as i32], + mlxcel_core::array_dtype(&h), + ) + }) + }) + .collect(); + + let h = self.norm.forward(&h); + let logits = if let Some(ref lm_head) = self.lm_head { + lm_head.forward(&h) + } else { + self.embed_tokens.as_linear(&h) + }; + + VerifyOutput { + logits, + hidden_states, + gdn_states: Vec::new(), + } + } + /// Rewind both KV (attention) and GDN (linear-attention) caches to the /// position of the last accepted token after a DFlash verify-pass block. /// @@ -2589,6 +2722,32 @@ impl mlxcel_core::drafter::dflash::SpeculativeTarget for Qwen35Model { self.forward_speculative(verify_input, caches, capture_layer_ids) } + /// The burst's prompt prefill, which is NOT a verify block and must not + /// be computed like one (issue #1935). + /// + /// The trait default routes here to `verify_forward_with_capture_layers`, + /// which sends every full-attention layer through `attend_per_position`. + /// That is right for a verify block, whose rows have to match what + /// single-token decode would produce, and wrong for the prompt, which + /// classic decode runs as one batched causal attention. Prefilling the + /// per-position way leaves the burst holding different KV and gated-delta + /// state than classic decode, and its greedy output then drifts at every + /// verify width alike. See + /// [`Qwen35Model::forward_prefill_with_capture_layers`]. + fn prefill_forward_with_capture_layers( + &self, + verify_input: &MlxArray, + caches: &mut [Self::Cache], + capture_layer_ids: &[usize], + ) -> Self::VerifyOut { + let capture_layer_ids = if capture_layer_ids.is_empty() { + mlxcel_core::drafter::dflash::config::DEFAULT_TARGET_LAYER_IDS + } else { + capture_layer_ids + }; + self.forward_prefill_with_capture_layers(verify_input, caches, capture_layer_ids) + } + fn rollback_partial( &self, caches: &mut [Self::Cache], @@ -3909,3 +4068,7 @@ mod chain_parity_gate_tests { assert!(!qwen35_config_supports_chain_parity_kernel(&config)); } } + +#[cfg(test)] +#[path = "qwen3_5_dflash_probe_tests.rs"] +mod qwen3_5_dflash_probe_tests; diff --git a/src/models/qwen3_5_dflash_probe_tests.rs b/src/models/qwen3_5_dflash_probe_tests.rs new file mode 100644 index 000000000..1ced612b9 --- /dev/null +++ b/src/models/qwen3_5_dflash_probe_tests.rs @@ -0,0 +1,727 @@ +// Copyright 2025-2026 Lablup Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Real-checkpoint block-versus-chain diagnostics for the Qwen 3.5 DFlash +//! verify path (issue #1935). +//! +//! Served greedy output on `qwen3.5-4b-4bit` + `qwen3.5-4b-dflash` is not +//! byte-identical to classic decode at any DFlash verify width, and the +//! difference is the same for every width from 2 to 16. These tests locate +//! where in the forward the two arms part company, on the real checkpoint, +//! without a server. +//! +//! `#[ignore]`-gated: they need the checkpoint on disk and a GPU. +//! +//! ```text +//! cargo test --release --features cuda -p mlxcel --lib -- --ignored \ +//! --test-threads=1 --nocapture models::qwen3_5::qwen3_5_dflash_probe_tests +//! ``` + +use super::{Qwen3NextCache, Qwen35Model}; +use crate::models::speculative_exactness::BlockChainExactness; +use mlxcel_core::{MlxArray, UniquePtr}; + +/// Default target checkpoint; override with `MLXCEL_Q35_PROBE_TARGET`. +const DEFAULT_TARGET: &str = "models/mlx/qwen3.5-4b-4bit"; + +/// Prompt length for the synthetic-id arms. Long enough that the full +/// attention layers have a real KV history behind the verify block. +const PROMPT_LEN: usize = 64; + +fn target_dir() -> String { + std::env::var("MLXCEL_Q35_PROBE_TARGET").unwrap_or_else(|_| DEFAULT_TARGET.to_string()) +} + +fn widths() -> Vec { + match std::env::var("MLXCEL_Q35_PROBE_WIDTHS") { + Ok(v) => v + .split(',') + .filter(|s| !s.trim().is_empty()) + .map(|s| s.trim().parse::().expect("width")) + .collect(), + Err(_) => vec![2, 4, 8, 16], + } +} + +/// Load the target and hand back its Qwen 3.5 text backbone. +/// +/// Returns `None` (and says so) when the checkpoint is not on disk, so a +/// host without it reports a skip instead of a failure. +fn load_text_model() -> Option<(crate::LoadedModel, String)> { + let dir = target_dir(); + if !std::path::Path::new(&dir).exists() { + eprintln!("[1935] skipping: {dir} not on disk"); + return None; + } + let (model, _tokenizer) = crate::backend::select_backend() + .load_model(std::path::Path::new(&dir)) + .expect("target model must load"); + Some((model, dir)) +} + +fn text_model(model: &crate::LoadedModel) -> &Qwen35Model { + match model { + crate::LoadedModel::Qwen35(m) => m, + crate::LoadedModel::Qwen35VLM(vlm) => &vlm.text_model, + _ => panic!("target is not a Qwen 3.5 text backbone"), + } +} + +fn ids_of(tokens: &[i32]) -> UniquePtr { + mlxcel_core::from_slice_i32(tokens, &[1, tokens.len() as i32]) +} + +/// Deterministic synthetic ids, the same shape the in-tree probe draws. +fn synthetic(len: usize, vocab: usize, stride: usize, offset: usize) -> Vec { + (0..len) + .map(|i| ((i * stride + offset) % vocab) as i32) + .collect() +} + +/// Raw bytes of one `[1, T, W]` row. +fn row_bytes(t: &MlxArray, index: i32) -> Vec { + let shape = mlxcel_core::array_shape(t); + let row = mlxcel_core::slice(t, &[0, index, 0], &[shape[0], index + 1, shape[2]]); + mlxcel_core::array_to_raw_bytes(&row) +} + +fn differing(a: &[u8], b: &[u8]) -> usize { + a.iter().zip(b).filter(|(x, y)| x != y).count() +} + +/// Which arm each verdict came from, for the printed table. +fn verdict_line(label: &str, verdict: &BlockChainExactness) -> String { + format!("[1935] {label}: {}", verdict.reason()) +} + +/// **Diagnostic 1.** Does `probe_block_chain_exactness` (the gate the issue +/// proposes to wire up for this family) actually fire on this checkpoint, +/// and at which widths? +/// +/// A `Diverges` verdict means the proposed gate would decline the burst; an +/// `Equal` verdict at a width the served arms diverge at would mean the +/// probe measures the wrong pair and the gate would be a no-op. +#[test] +#[ignore = "needs the real Qwen 3.5 4B checkpoint and a GPU"] +fn block_chain_exactness_verdicts_on_the_real_checkpoint() { + let Some((model, dir)) = load_text_model() else { + return; + }; + let text = text_model(&model); + eprintln!("[1935] target {dir}"); + for width in widths() { + let verdict = text.probe_block_chain_exactness(width); + eprintln!("{}", verdict_line(&format!("width {width}"), &verdict)); + } +} + +/// **Diagnostic 2.** Per-layer bisect of one verify block against the +/// single-token chain, on the real checkpoint. +/// +/// Both arms run `forward_speculative` from the same prefill, so the only +/// difference is `T = width` against `T = 1`. Position 0 of the block sees +/// exactly the context the chain's first step sees, so a difference there is +/// a pure row-count effect with identical mathematics behind it, and the +/// first layer that shows one names the operator responsible. +/// +/// Layers 0, 1, 2 are gated-delta (linear attention) and layer 3 is the +/// first full-attention layer on a `full_attention_interval = 4` checkpoint, +/// so the first differing layer index separates the two candidate families +/// the issue lists. +#[test] +#[ignore = "needs the real Qwen 3.5 4B checkpoint and a GPU"] +fn per_layer_block_versus_chain_bisect() { + let Some((model, dir)) = load_text_model() else { + return; + }; + let text = text_model(&model); + let vocab = text.config.vocab_size; + let layers = text.num_layers(); + let capture: Vec = (0..layers).collect(); + let block_size: usize = std::env::var("MLXCEL_Q35_PROBE_BLOCK") + .ok() + .and_then(|v| v.parse().ok()) + .unwrap_or(4); + eprintln!("[1935] target {dir}, {layers} layers, block {block_size}"); + + let prompt = synthetic(PROMPT_LEN, vocab, 7, 1); + let block = synthetic(block_size, vocab, 13, 3); + + // Chain arm: prefill, then one token per forward. + let mut chain_caches: Vec = text.make_speculative_caches(); + let _ = text.forward_speculative(&ids_of(&prompt), &mut chain_caches, &capture); + let mut chain_hidden: Vec>> = Vec::with_capacity(block_size); + let mut chain_logits: Vec> = Vec::with_capacity(block_size); + for token in &block { + let out = text.forward_speculative(&ids_of(&[*token]), &mut chain_caches, &capture); + chain_hidden.push( + out.hidden_states + .iter() + .map(|h| row_bytes(h, 0)) + .collect::>(), + ); + chain_logits.push(row_bytes(&out.logits, 0)); + } + + // Block arm: fresh caches, same prefill, the whole block at once. + let mut block_caches: Vec = text.make_speculative_caches(); + let _ = text.forward_speculative(&ids_of(&prompt), &mut block_caches, &capture); + let out = text.forward_speculative(&ids_of(&block), &mut block_caches, &capture); + + let mut first_layer: Option = None; + for (layer, chain_row) in chain_hidden[0].iter().enumerate().take(layers) { + let block_row = row_bytes(&out.hidden_states[layer], 0); + let d = differing(&block_row, chain_row); + if d > 0 { + if first_layer.is_none() { + first_layer = Some(layer); + } + if layer < 8 || first_layer == Some(layer) { + eprintln!( + "[1935] position 0, layer {layer} ({}): {d} of {} hidden bytes differ", + if text.config.is_linear_layer(layer) { + "gated-delta" + } else { + "full attention" + }, + block_row.len() + ); + } + } + } + match first_layer { + None => eprintln!( + "[1935] position 0: every layer's hidden state is byte-identical between the \ + {block_size}-row block and the single-token chain" + ), + Some(layer) => eprintln!( + "[1935] position 0: FIRST divergence at layer {layer} ({})", + if text.config.is_linear_layer(layer) { + "gated-delta" + } else { + "full attention" + } + ), + } + let logit_diff = differing(&row_bytes(&out.logits, 0), &chain_logits[0]); + eprintln!( + "[1935] position 0: {logit_diff} of {} logit bytes differ", + chain_logits[0].len() + ); +} + +/// Bytes of a whole tensor, or an empty vector when absent. +fn all_bytes(t: Option<&UniquePtr>) -> Vec { + match t { + Some(t) => mlxcel_core::array_to_raw_bytes(t.as_ref().expect("live tensor")), + None => Vec::new(), + } +} + +/// One line describing how far two cache vectors agree. +fn compare_caches(a: &[Qwen3NextCache], b: &[Qwen3NextCache]) -> Option<(usize, String)> { + for (i, (x, y)) in a.iter().zip(b).enumerate() { + match (x, y) { + (Qwen3NextCache::Attention(kx), Qwen3NextCache::Attention(ky)) => { + let dk = differing(&all_bytes(kx.keys.as_ref()), &all_bytes(ky.keys.as_ref())); + let dv = differing( + &all_bytes(kx.values.as_ref()), + &all_bytes(ky.values.as_ref()), + ); + if dk > 0 || dv > 0 || kx.offset != ky.offset { + return Some(( + i, + format!( + "attention layer: {dk} key bytes, {dv} value bytes differ, \ + offsets {} vs {}", + kx.offset, ky.offset + ), + )); + } + } + (Qwen3NextCache::Linear(gx), Qwen3NextCache::Linear(gy)) => { + let dc = differing( + &all_bytes(gx.conv_state.as_ref()), + &all_bytes(gy.conv_state.as_ref()), + ); + let ds = differing( + &all_bytes(gx.state_cache.as_ref()), + &all_bytes(gy.state_cache.as_ref()), + ); + if dc > 0 || ds > 0 || gx.offset != gy.offset { + return Some(( + i, + format!( + "gated-delta layer: {dc} conv-state bytes, {ds} recurrent-state \ + bytes differ, offsets {} vs {}", + gx.offset, gy.offset + ), + )); + } + } + _ => return Some((i, "cache variants differ".to_string())), + } + } + None +} + +/// **Diagnostic 3.** The pair the block-versus-chain probe does not compare: +/// the classic decode forward (`forward_internal`, which is what +/// `LanguageModel::forward` reaches through the model-owned sequence state) +/// against the speculative verify forward (`forward_speculative`), both at +/// `T = 1` and both from their own prefill of the same prompt. +/// +/// The DFlash burst runs every one of its forwards through +/// `forward_speculative` on caches of its own, so this is the pair a served +/// speculative completion is actually compared against when an operator +/// diffs it with classic decode. `probe_block_chain_exactness` holds both of +/// its arms inside `forward_speculative` and therefore cannot see a +/// difference here. +#[test] +#[ignore = "needs the real Qwen 3.5 4B checkpoint and a GPU"] +fn classic_forward_versus_speculative_forward_at_one_row() { + let Some((model, dir)) = load_text_model() else { + return; + }; + let text = text_model(&model); + let vocab = text.config.vocab_size; + let steps: usize = std::env::var("MLXCEL_Q35_PROBE_STEPS") + .ok() + .and_then(|v| v.parse().ok()) + .unwrap_or(8); + eprintln!("[1935] target {dir}, comparing classic and speculative forwards at T=1"); + + let prompt = synthetic(PROMPT_LEN, vocab, 7, 1); + let block = synthetic(steps, vocab, 13, 3); + + let mut classic_caches = text.make_internal_caches(); + let classic_prefill = text.forward_internal(&ids_of(&prompt), None, &mut classic_caches, None); + let mut spec_caches = text.make_speculative_caches(); + let spec_prefill = text.forward_speculative(&ids_of(&prompt), &mut spec_caches, &[]); + + let last = PROMPT_LEN as i32 - 1; + let prefill_diff = differing( + &row_bytes(&classic_prefill, last), + &row_bytes(&spec_prefill.logits, last), + ); + eprintln!( + "[1935] prefill last row: {prefill_diff} of {} logit bytes differ", + row_bytes(&classic_prefill, last).len() + ); + match compare_caches(&classic_caches, &spec_caches) { + None => eprintln!("[1935] prefill: every layer's cache state is byte-identical"), + Some((layer, why)) => { + eprintln!("[1935] prefill: FIRST cache divergence at layer {layer}: {why}") + } + } + + for (i, token) in block.iter().enumerate() { + let c = text.forward_internal(&ids_of(&[*token]), None, &mut classic_caches, None); + let s = text.forward_speculative(&ids_of(&[*token]), &mut spec_caches, &[]); + let d = differing(&row_bytes(&c, 0), &row_bytes(&s.logits, 0)); + eprintln!( + "[1935] step {i}: {d} of {} logit bytes differ", + row_bytes(&c, 0).len() + ); + if d > 0 { + if let Some((layer, why)) = compare_caches(&classic_caches, &spec_caches) { + eprintln!("[1935] step {i}: FIRST cache divergence at layer {layer}: {why}"); + } + break; + } + } +} + +/// **Diagnostic 4.** Rollback exactness: after a partial accept, is the +/// rewound cache state byte-identical to the state a decode chain of the +/// accepted length would hold? +/// +/// The round loop verifies a `block_size`-row block, accepts a prefix of it +/// and rolls the caches back to the accepted length. If that rewind is not +/// exact, every later token inherits the drift even when the verify block +/// itself is byte-identical to the chain, which on this checkpoint it is at +/// widths 2 and 4. +/// +/// `MLXCEL_Q35_PROBE_BLOCK` sets the block width, `MLXCEL_Q35_PROBE_KEEP` the +/// number of rows accepted (the round loop's `accepted + 1`). +#[test] +#[ignore = "needs the real Qwen 3.5 4B checkpoint and a GPU"] +fn rollback_after_partial_accept_matches_the_decode_chain() { + let Some((model, dir)) = load_text_model() else { + return; + }; + let text = text_model(&model); + let vocab = text.config.vocab_size; + let block_size: usize = std::env::var("MLXCEL_Q35_PROBE_BLOCK") + .ok() + .and_then(|v| v.parse().ok()) + .unwrap_or(4); + let keep: usize = std::env::var("MLXCEL_Q35_PROBE_KEEP") + .ok() + .and_then(|v| v.parse().ok()) + .unwrap_or(2); + assert!(keep >= 1 && keep < block_size, "keep must be in 1..block"); + eprintln!("[1935] target {dir}, block {block_size}, keeping {keep} rows"); + + let prompt = synthetic(PROMPT_LEN, vocab, 7, 1); + let block = synthetic(block_size, vocab, 13, 3); + + // Chain arm: the state after `keep` single-token steps. + let mut chain_caches = text.make_speculative_caches(); + let _ = text.forward_speculative(&ids_of(&prompt), &mut chain_caches, &[]); + for token in block.iter().take(keep) { + let _ = text.forward_speculative(&ids_of(&[*token]), &mut chain_caches, &[]); + } + + // Block arm: verify the whole block, then roll back to `keep` rows. + let mut block_caches = text.make_speculative_caches(); + let _ = text.forward_speculative(&ids_of(&prompt), &mut block_caches, &[]); + let out = text.forward_speculative(&ids_of(&block), &mut block_caches, &[]); + let kept = text.rollback_speculative_cache( + &mut block_caches, + &out.gdn_states, + &[keep as i32 - 1], + block_size as i32, + ); + eprintln!("[1935] rollback returned max accepted index {kept}"); + + match compare_caches(&chain_caches, &block_caches) { + None => eprintln!("[1935] rollback: every layer's cache state is byte-identical"), + Some((layer, why)) => { + eprintln!("[1935] rollback: FIRST cache divergence at layer {layer}: {why}") + } + } + + // The token that follows is what the next round would verify from. + let next = synthetic(1, vocab, 29, 11); + let c = text.forward_speculative(&ids_of(&next), &mut chain_caches, &[]); + let b = text.forward_speculative(&ids_of(&next), &mut block_caches, &[]); + let d = differing(&row_bytes(&c.logits, 0), &row_bytes(&b.logits, 0)); + eprintln!( + "[1935] next step after rollback: {d} of {} logit bytes differ", + row_bytes(&c.logits, 0).len() + ); +} + +fn env_ids(name: &str) -> Option> { + let raw = std::env::var(name).ok()?; + Some( + raw.trim() + .trim_matches(|c| c == '[' || c == ']') + .split(',') + .filter(|s| !s.trim().is_empty()) + .map(|s| s.trim().parse::().expect("integer token id")) + .collect(), + ) +} + +/// The capture list the served burst actually uses. Passing an empty slice +/// through `forward_speculative` captures nothing, and a capture inserts a +/// `copy` of the layer's output into the graph, which can change what MLX +/// fuses. `MLXCEL_Q35_PROBE_CAPTURE=0` turns it off for the A/B. +fn probe_capture_layer_ids() -> Vec { + match std::env::var("MLXCEL_Q35_PROBE_CAPTURE").as_deref() { + Ok("0") => Vec::new(), + Ok(v) if !v.trim().is_empty() => v + .split(',') + .filter(|s| !s.trim().is_empty()) + .map(|s| s.trim().parse::().expect("layer id")) + .collect(), + _ => mlxcel_core::drafter::dflash::config::DEFAULT_TARGET_LAYER_IDS.to_vec(), + } +} + +/// `true` when the burst arm should take its per-position argmax the way the +/// round loop does, one `argmax_last_axis` over the whole `[1, T, vocab]` +/// block, rather than one call per sliced row. A multi-row reduction is not +/// obliged to break an exact tie the same way a single-row one does. +fn probe_block_argmax() -> bool { + std::env::var("MLXCEL_Q35_PROBE_BLOCK_ARGMAX") + .map(|v| v != "0") + .unwrap_or(false) +} + +/// Per-position argmax over a whole `[1, T, vocab]` block, which is what +/// `round_loop::argmax_logits_to_array` computes. +fn argmax_block(logits: &MlxArray) -> Vec { + let arg = mlxcel_core::argmax_last_axis(logits); + mlxcel_core::eval(&arg); + let itemsize = mlxcel_core::array_itemsize(&arg); + let bytes = mlxcel_core::array_to_raw_bytes(&arg); + match itemsize { + 4 => bytes + .chunks_exact(4) + .map(|c| i32::from_ne_bytes([c[0], c[1], c[2], c[3]])) + .collect(), + 8 => bytes + .chunks_exact(8) + .map(|c| i64::from_ne_bytes([c[0], c[1], c[2], c[3], c[4], c[5], c[6], c[7]]) as i32) + .collect(), + other => panic!("unexpected argmax itemsize {other}"), + } +} + +fn argmax_row(logits: &MlxArray, index: i32) -> i32 { + let shape = mlxcel_core::array_shape(logits); + let row = mlxcel_core::slice(logits, &[0, index, 0], &[shape[0], index + 1, shape[2]]); + let arg = mlxcel_core::argmax_last_axis(&row); + let arg = mlxcel_core::reshape(&arg, &[]); + mlxcel_core::eval(&arg); + mlxcel_core::item_i32(&arg) +} + +/// **Diagnostic 5.** Replay a real classic-decode transcript through both +/// shapes and report where their greedy choices part. +/// +/// The classic arm is `forward_internal`, prefill and one token per step, +/// which is what `mlxcel generate` and the server's classic decode both run +/// (verified byte-identical on GB10). The burst arm prefills through +/// `forward_prefill_with_capture_layers` and then feeds the same tokens in +/// `block_size`-row verify blocks through `forward_speculative`, which is +/// the shape a DFlash round takes when every proposal is correct. No +/// rollback is exercised, so a divergence here is in the prefill or the +/// verify block rather than in the rewind. +/// +/// Supply the transcript from a real run: +/// +/// ```text +/// MLXCEL_Q35_PROBE_PROMPT="[...]" MLXCEL_Q35_PROBE_REFERENCE="[...]" +/// ``` +/// +/// (`MLXCEL_PRINT_TOKEN_IDS=1 mlxcel generate -m -p --no-chat-template --temp 0` +/// prints both lines.) +#[test] +#[ignore = "needs the real Qwen 3.5 4B checkpoint, a GPU and a recorded transcript"] +fn reference_replay_classic_versus_burst_shape() { + let Some((model, dir)) = load_text_model() else { + return; + }; + let text = text_model(&model); + let Some(prompt) = env_ids("MLXCEL_Q35_PROBE_PROMPT") else { + eprintln!("[1935] skipping: set MLXCEL_Q35_PROBE_PROMPT and MLXCEL_Q35_PROBE_REFERENCE"); + return; + }; + let reference = env_ids("MLXCEL_Q35_PROBE_REFERENCE").expect("MLXCEL_Q35_PROBE_REFERENCE"); + let block_size: usize = std::env::var("MLXCEL_Q35_PROBE_BLOCK") + .ok() + .and_then(|v| v.parse().ok()) + .unwrap_or(4); + eprintln!( + "[1935] target {dir}, prompt {} ids, reference {} ids, block {block_size}", + prompt.len(), + reference.len() + ); + + // Both prefills first, so the cache comparison below sees two caches that + // have consumed the same tokens. Comparing after the classic chain has run + // would only report the offset gap the chain itself opened. + let mut classic_caches = text.make_internal_caches(); + let classic_prefill = text.forward_internal(&ids_of(&prompt), None, &mut classic_caches, None); + let capture = probe_capture_layer_ids(); + eprintln!( + "[1935] capture layers {capture:?}, block argmax {}", + probe_block_argmax() + ); + let mut burst_caches = text.make_speculative_caches(); + let burst_prefill = + text.forward_prefill_with_capture_layers(&ids_of(&prompt), &mut burst_caches, &capture); + + let prefill_bytes = differing( + &row_bytes(&classic_prefill, prompt.len() as i32 - 1), + &row_bytes(&burst_prefill.logits, prompt.len() as i32 - 1), + ); + eprintln!("[1935] prefill last row: {prefill_bytes} logit bytes differ"); + match compare_caches(&classic_caches, &burst_caches) { + None => eprintln!("[1935] prefill: every layer's cache state is byte-identical"), + Some((layer, why)) => { + eprintln!("[1935] prefill: FIRST cache divergence at layer {layer}: {why}") + } + } + + let mut classic: Vec = vec![argmax_row(&classic_prefill, prompt.len() as i32 - 1)]; + for token in &reference { + let out = text.forward_internal(&ids_of(&[*token]), None, &mut classic_caches, None); + classic.push(argmax_row(&out, 0)); + } + + let mut burst: Vec = vec![argmax_row(&burst_prefill.logits, prompt.len() as i32 - 1)]; + for chunk in reference.chunks(block_size) { + let out = text.forward_speculative(&ids_of(chunk), &mut burst_caches, &capture); + if probe_block_argmax() { + burst.extend(argmax_block(&out.logits)); + } else { + for pos in 0..chunk.len() { + burst.push(argmax_row(&out.logits, pos as i32)); + } + } + } + + let first = classic + .iter() + .zip(&burst) + .position(|(a, b)| a != b) + .map(|i| i as i64) + .unwrap_or(-1); + let disagreements = classic.iter().zip(&burst).filter(|(a, b)| a != b).count(); + eprintln!( + "[1935] greedy argmax over {} positions: {disagreements} disagree, first at {first} \ + (-1 means none)", + classic.len() + ); + if first >= 0 { + let i = first as usize; + eprintln!( + "[1935] at position {i}: classic {} against burst {}, reference has {:?}", + classic[i], + burst[i], + reference.get(i) + ); + } +} + +/// **Diagnostic 6.** The round loop's cache dynamics, without a drafter. +/// +/// Diagnostic 5 feeds the transcript in aligned blocks with every proposal +/// correct, so it never rolls back. A served burst rolls back on most rounds +/// (measured acceptance on this pairing is about 0.53), and a rollback trims +/// the attention caches and replays the gated-delta state. This arm walks the +/// same transcript but accepts a varying prefix of each block and rewinds the +/// rest, which is the cache history a served burst actually builds, and +/// compares its greedy argmax against the single-token chain position by +/// position. +/// +/// `MLXCEL_Q35_PROBE_ACCEPTS` is a comma-separated cycle of accepted row +/// counts (default `1,2,3`, clamped to the block width). +#[test] +#[ignore = "needs the real Qwen 3.5 4B checkpoint, a GPU and a recorded transcript"] +fn round_loop_cache_dynamics_with_rollback_match_the_chain() { + let Some((model, dir)) = load_text_model() else { + return; + }; + let text = text_model(&model); + let Some(prompt) = env_ids("MLXCEL_Q35_PROBE_PROMPT") else { + eprintln!("[1935] skipping: set MLXCEL_Q35_PROBE_PROMPT and MLXCEL_Q35_PROBE_REFERENCE"); + return; + }; + let reference = env_ids("MLXCEL_Q35_PROBE_REFERENCE").expect("MLXCEL_Q35_PROBE_REFERENCE"); + let block_size: usize = std::env::var("MLXCEL_Q35_PROBE_BLOCK") + .ok() + .and_then(|v| v.parse().ok()) + .unwrap_or(4); + let accepts: Vec = std::env::var("MLXCEL_Q35_PROBE_ACCEPTS") + .ok() + .map(|v| { + v.split(',') + .filter(|s| !s.trim().is_empty()) + .map(|s| s.trim().parse::().expect("accept count")) + .collect() + }) + .unwrap_or_else(|| vec![1, 2, 3]); + eprintln!( + "[1935] target {dir}, block {block_size}, accept cycle {accepts:?}, \ + reference {} ids", + reference.len() + ); + + let mut classic_caches = text.make_internal_caches(); + let classic_prefill = text.forward_internal(&ids_of(&prompt), None, &mut classic_caches, None); + let mut classic: Vec = vec![argmax_row(&classic_prefill, prompt.len() as i32 - 1)]; + for token in &reference { + let out = text.forward_internal(&ids_of(&[*token]), None, &mut classic_caches, None); + classic.push(argmax_row(&out, 0)); + } + + let capture = probe_capture_layer_ids(); + eprintln!( + "[1935] capture layers {capture:?}, block argmax {}", + probe_block_argmax() + ); + let mut burst_caches = text.make_speculative_caches(); + let burst_prefill = + text.forward_prefill_with_capture_layers(&ids_of(&prompt), &mut burst_caches, &capture); + let mut burst: Vec = vec![argmax_row(&burst_prefill.logits, prompt.len() as i32 - 1)]; + + // A served round's rejected rows hold the drafter's WRONG proposals, not + // the tokens that end up emitted. Their keys and values are written into + // the attention caches and their projections into the gated-delta + // snapshot before the rewind trims them, so a rewind that reads one row + // too far shows up here and nowhere else. `MLXCEL_Q35_PROBE_WRONG` turns + // that on; with it off the rejected rows carry the correct tokens, which + // is the weaker test. + let wrong: Option = std::env::var("MLXCEL_Q35_PROBE_WRONG") + .ok() + .filter(|v| v != "0") + .map(|v| v.parse::().unwrap_or(9999)); + if let Some(w) = wrong { + eprintln!("[1935] rejected rows carry the wrong token id {w}"); + } + + let mut i = 0usize; + let mut round = 0usize; + let mut rollbacks = 0usize; + while i < reference.len() { + let end = (i + block_size).min(reference.len()); + let rows = end - i; + let keep_now = accepts[round % accepts.len()].clamp(1, rows); + let mut fed: Vec = reference[i..end].to_vec(); + if let Some(w) = wrong { + for slot in fed.iter_mut().skip(keep_now) { + *slot = w; + } + } + let out = text.forward_speculative(&ids_of(&fed), &mut burst_caches, &capture); + let keep = keep_now; + if probe_block_argmax() { + let all = argmax_block(&out.logits); + burst.extend(all.into_iter().take(keep)); + } else { + for pos in 0..keep { + burst.push(argmax_row(&out.logits, pos as i32)); + } + } + if keep < rows { + // The round loop's own condition: rewind whenever the block was + // not fully accepted. `accepted` is the index of the last kept row. + text.rollback_speculative_cache( + &mut burst_caches, + &out.gdn_states, + &[keep as i32 - 1], + rows as i32, + ); + rollbacks += 1; + } + i += keep; + round += 1; + } + + let first = classic + .iter() + .zip(&burst) + .position(|(a, b)| a != b) + .map(|i| i as i64) + .unwrap_or(-1); + let disagreements = classic.iter().zip(&burst).filter(|(a, b)| a != b).count(); + eprintln!( + "[1935] {round} rounds, {rollbacks} rollbacks; greedy argmax over {} compared \ + positions: {disagreements} disagree, first at {first} (-1 means none)", + classic.len().min(burst.len()) + ); + if first >= 0 { + let i = first as usize; + eprintln!( + "[1935] at position {i}: classic {} against burst {}", + classic[i], burst[i] + ); + } +} diff --git a/src/server/batch/dflash_target.rs b/src/server/batch/dflash_target.rs index d0eafbbf2..54c16026d 100644 --- a/src/server/batch/dflash_target.rs +++ b/src/server/batch/dflash_target.rs @@ -207,8 +207,10 @@ pub(crate) trait DFlashTargetModel: /// Whether a `block_size`-row verify block is byte-identical to the /// single-token decode chain on this host, and the burst may therefore - /// promise the temperature-0 contract. Qwen 3.5 DFlash has never gated on - /// this and keeps the permissive default; LFM2 runs the measured probe. + /// promise the temperature-0 contract. Every family with a server-side + /// DFlash target now overrides this with the measured block-versus-chain + /// probe; the permissive default remains only for a target that has not + /// been measured yet, and issue #1935 is what a permissive default costs. fn exactness_allows(&self, _block_size: usize) -> bool { true } @@ -252,12 +254,18 @@ impl DFlashTargetModel for crate::models::Qwen35Model { fn make_dflash_caches(&self) -> Vec { self.make_speculative_caches() } + fn exactness_allows(&self, block_size: usize) -> bool { + self.dflash_exactness_allows(block_size) + } } impl DFlashTargetModel for crate::vision::Qwen35VLModel { fn make_dflash_caches(&self) -> Vec { self.text_model.make_speculative_caches() } + fn exactness_allows(&self, block_size: usize) -> bool { + self.text_model.dflash_exactness_allows(block_size) + } } impl DFlashTargetModel for crate::models::Lfm2Model { @@ -717,6 +725,13 @@ where decode_ms = diagnostics.total_decode_time_ms, "DFlash diagnostics" ); + // The per-round accept lengths, at `debug` because they are one + // line per burst and only a parity investigation reads them. They + // are what lets an in-process replay reproduce a served burst's + // exact cache history (rewind depth per round) rather than a + // synthetic accept cycle, which is the gap issue #1935's residual + // divergence is being chased through. + tracing::debug!(accept_lens = ?output.accept_lens, "DFlash accept lengths"); let mut tokens = Vec::with_capacity(output.tokens.len() + 1); tokens.push(first_bonus); tokens.extend(output.tokens); diff --git a/src/vision/qwen3_5_vl.rs b/src/vision/qwen3_5_vl.rs index 855cd337c..89f7a475f 100644 --- a/src/vision/qwen3_5_vl.rs +++ b/src/vision/qwen3_5_vl.rs @@ -182,6 +182,24 @@ impl mlxcel_core::drafter::dflash::SpeculativeTarget for Qwen35VLModel { .forward_speculative(verify_input, caches, capture_layer_ids) } + /// Prompt prefill through the batched causal attention path rather than + /// the per-position verify path, for the reason the text model's own + /// override states (issue #1935). + fn prefill_forward_with_capture_layers( + &self, + verify_input: &MlxArray, + caches: &mut [Self::Cache], + capture_layer_ids: &[usize], + ) -> Self::VerifyOut { + let capture_layer_ids = if capture_layer_ids.is_empty() { + mlxcel_core::drafter::dflash::config::DEFAULT_TARGET_LAYER_IDS + } else { + capture_layer_ids + }; + self.text_model + .forward_prefill_with_capture_layers(verify_input, caches, capture_layer_ids) + } + fn rollback_partial( &self, caches: &mut [Self::Cache], diff --git a/tests/speculative_parity.rs b/tests/speculative_parity.rs index 3ada77531..a9d70cbe3 100644 --- a/tests/speculative_parity.rs +++ b/tests/speculative_parity.rs @@ -282,6 +282,109 @@ async fn server_round( } } +/// Multi-width byte-equality for a DFlash pairing (issue #1935). +/// +/// Runs the drafter-less baseline once and then one speculative server per +/// width, and requires each width to land in one of exactly two acceptable +/// states, naming which: +/// +/// * the burst ran (`Speculative burst completed`, no decline line) and its +/// text is byte-identical to the baseline, or +/// * the exactness gate declined the burst, in which case the request was +/// served by classic decode and byte-equality is not evidence about the +/// verify path. The decline is reported as a decline. +/// +/// A width that runs the burst and differs from the baseline fails, which is +/// the condition issue #1935 was filed for. A width that declines does not +/// fail, because declining is the gate working: on GB10 widths 8 and 16 are +/// genuinely not byte-identical to the single-token chain, and no arrangement +/// of this code makes them so. +/// +/// The single-width [`assert_server_byte_equality`] stays for the MTP +/// pairings, which have one measured width each. +async fn assert_server_byte_equality_at_widths( + pairing: &Pairing, + target_path: &std::path::Path, + draft_path: &std::path::Path, + widths: &[usize], +) { + let draft_str = draft_path.to_string_lossy().to_string(); + + // Baseline first and once: every width compares against the same text, + // and a drafter-less server is the cheapest arm to hold the GPU. + let baseline = server_round(pairing.name, "baseline", target_path, &[]).await; + + let mut ran = Vec::new(); + let mut declined = Vec::new(); + let mut mismatched = Vec::new(); + + for &width in widths { + let width_str = width.to_string(); + let label = format!("speculative b={width}"); + let spec = server_round( + pairing.name, + &label, + target_path, + &[ + "--model-draft", + &draft_str, + "--draft-kind", + pairing.kind, + "--draft-block-size", + &width_str, + ], + ) + .await; + + let was_declined = spec.logs.contains("DFlash speculative dispatch declined"); + let burst_ran = spec.logs.contains("Speculative burst completed"); + assert!( + was_declined || burst_ran, + "[{}] b={width}: the server neither completed a speculative burst nor logged a \ + decline, so this arm proves nothing either way. Captured logs:\n{}", + pairing.name, + spec.logs, + ); + + let identical = spec.content == baseline.content + && spec.completion_tokens == baseline.completion_tokens; + if was_declined { + declined.push(width); + assert!( + identical, + "[{}] b={width}: the burst declined to classic decode, so the response must \ + equal the drafter-less baseline byte for byte; it does not, which means the \ + decline path itself changed the output.\nspeculative: {:?}\nbaseline: {:?}", + pairing.name, spec.content, baseline.content, + ); + } else { + ran.push(width); + if !identical { + mismatched.push((width, spec.content.clone())); + } + } + } + + eprintln!( + "[{}] widths that ran the burst: {ran:?}; widths the gate declined: {declined:?}", + pairing.name + ); + assert!( + mismatched.is_empty(), + "[{}] these widths ran a speculative burst and produced text that is NOT \ + byte-identical to the drafter-less baseline at temperature 0: {:?}.\nbaseline: {:?}", + pairing.name, + mismatched.iter().map(|(w, _)| *w).collect::>(), + baseline.content, + ); + assert!( + !ran.is_empty(), + "[{}] every width declined, so this run says nothing about verify-path parity. \ + Widths tried: {widths:?}", + pairing.name, + ); +} + /// byte-equality phase: spawn `mlxcel-server` twice against the /// same `target_path` — once speculative (drafter attached via /// `--model-draft --draft-kind --draft-block-size`), once drafter-less — @@ -699,7 +802,12 @@ async fn greedy_parity_dflash_qwen35_4b() { } // ---- Phase 2: end-to-end byte-equality (subprocess) ---- - assert_server_byte_equality(pairing, &target_path, &draft_path).await; + // Every width issue #1935 measured, not just the one the pairing names. + // The divergence it reported was identical at 2, 4, 8 and 16, so a + // single-width arm could not tell a width-specific kernel boundary from a + // path difference shared by all of them, and the single width this test + // used to run was 16, which the gate now declines. + assert_server_byte_equality_at_widths(pairing, &target_path, &draft_path, &[2, 4, 8, 16]).await; } /// Greedy parity for the Gemma 4 31B + MTP assistant pairing.