Repository navigation
perf(speculative): make the draft block width a per-device default - #1937
Merged
Merged
Conversation
The DFlash draft block width resolved to a flat 16 with no hardware or quantization input. On GB10 that loses: the Qwen 3.5 DFlash pairing serves 48.63 tok/s at 16 against 58.33 for classic decode, while width 4 serves 76.88 (docs/benchmark_results/dflash-verify-fixed-cost-gb10-2026-09-11.md). The throughput was reachable only by passing --draft-block-size by hand. New src/cli/draft_block_policy.rs owns a pure policy over (drafter kind, CUDA compute capability, target quantization) plus a peek that classifies the target checkpoint's quantization from its config.json, covering MLX's root `quantization` block, the ModelOpt and compressed-tensors `quantization_config` shape that the Laguna NVFP4 release ships, and the `text_config` nesting. The policy is a default only: it resolves after an operator override and after a drafter checkpoint's own declared width, and before the flat per-kind constant, so every host and quantization path with no measurement is unchanged. resolve_draft_block_size now returns the width together with its provenance, which SpeculativeDispatch carries into its startup summary line so an operator can tell a measured default from a flag they forgot they exported. Both call sites thread the target checkpoint path in, because the quantization that picks the CUDA kernel is the target's, not the drafter's. scripts/bench_block_width.sh gains the two DFlash pairings and records the kernel release and NVIDIA driver version in its host line; neither was captured anywhere before, which left cross-session GB10 comparisons with nothing to key on. No policy entry is seeded by this commit; the measurement and the record follow. Refs #1797
The #1782 numbers came from `mlxcel-server`, and `scripts/bench_block_width.sh` drives the offline `mlxcel generate` path, so the served arm needs a driver of its own. A served width cannot be interleaved per sample the way the offline one is: `draft_block_size` is worker-owned in the live settings API, so changing it means restarting the server, and the DFlash generator is built per request and dropped at the end, so a warm-up has nothing to amortize into. The drift control is therefore a bracket, classic first and last with the widths between them, and the summarizer refuses to present a clean table when the two classic ranges do not overlap. The host gate is imported from the #1820 harness rather than copied, so the three predicates it encodes stay in one place: contention is processes that can consume CPU now, a foreign model process is a reason to wait, and the cumulative NV_ERR_NO_MEMORY count is a non-decaying freeze precursor on this host. The sweep records the NVRM delta per arm and refuses to start an arm whose projected cumulative count would pass the budget. The identity script captures what no harness here captured before: the kernel release, the NVIDIA driver version, and the MLX_CUDA_ARCHITECTURES the build was pinned to. build.rs auto-detects 121a on this host while the shipped release and every earlier GB10 record use plain 121, and a draft block width is decided by which CUDA kernel a verify block dispatches to, so the architecture is part of the measured configuration. Refs #1797
The prompt was about twice the 165-token prompt #1782 measured on. Prefill is identical on every arm either way, so it cannot change the width ordering, but a record that cannot state its own prompt length is harder to compare against the one it follows up. The streamed request now asks for the usage block so the measured prompt and completion token counts land in the JSONL instead of being asserted from the source. Refs #1797
"per-kind default" alone cannot be acted on. An unmeasured device and a target whose config.json could not be read take the same fallback and print the same sentence, and they call for different responses: one is a platform nobody has swept, the other is a checkpoint the peek failed on. The two policy arms now carry the compute capability and the quantization classification they were keyed on, so the startup line reads "16 (per-kind default, no measurement for sm_121 with an NVFP4 target)". The sweep driver gains the matching blind spot check. A DFlash burst that declines runs classic decode instead, at classic throughput and with byte-identical text, so a declined arm and a width that genuinely does nothing produce the same table row. The driver now records whether the log carries a decline and whether any DFlash diagnostics line appeared at all, and the summarizer refuses to compute a ratio for a speculative arm that has neither. The diagnostics regex is anchored on the DFlash diagnostics event rather than on a bare block_size field, which also appears in paged-KV and prefill-chunking lines and reported a width for the classic arm. A driver-budget halt is now written into the JSONL rather than raised as a traceback. docs/CONTINUOUS_BATCHING.md no longer states a flat DFlash default of 16, and the identity script cites the binary's own startup line for its architecture list instead of grepping the first numeric string out of a 100 MB file. Refs #1797
Three defects a smoke run found before the sweep, not after it. `CompletionRequest::model` has no serde default, so a body without it is answered 422 by the Json extractor before any generation happens and every arm would have recorded an error. The diagnostics regex assumed the fields preceded the message, but the tracing format emits the message first, so it matched nothing and every speculative arm looked like it had never run a verify block. And the diagnostics fields are per-request totals, so the round device sync the record tabulates is the total divided by the round count; reporting the total gave 2347 ms where the per-round figure is 30.5. The driver now also records the full diagnostics row per burst (rounds, acceptance, emitted per verify, the draft, verify and sync splits) and treats an arm whose every burst reported zero rounds as not a measurement, the way it already treats a declined one. A 32-token request produces exactly that: the first bonus covers the budget and the round loop never runs. Refs #1797
A width seeded into a shipped default has to be measured in the configuration users run. MLX defaults `MLX_ENABLE_TF32` to 1 (`get_var("MLX_ENABLE_TF32", 1)` in mlx/utils.h) and #1782 left it there; the sweep was forcing it to 0, which is neither. The driver now defaults it to 1 and still sets it explicitly, so the record can state it rather than leave it to the reader.
The identity script was reporting the rustc first on PATH, which on this host is 1.97.1 while the build uses the repo's pinned 1.93.1. It now prints both and takes the build toolchain from `BUILD_RUSTC`, which `run_session.sh` sets. It also finds nvcc at the CUDA install root, where this host keeps it.
`speculative_bench` says explicitly that it does not go through the #1797 policy, so its per-pairing block sizes stay comparable across hosts.
Refs #1797
The comment asserted that forcing TF32 off costs the speculative arms about 9% and the classic arm under 2%, reasoning from the linear-attention chunked scan being f32. Measuring it both ways on the same pairing in the same session falsified that: classic loses about 1.2% with TF32 off, the speculative arms lose nothing outside their own spread, and the width ordering does not change. The 9% gap to #1782 that prompted the guess is the different prompt, not the flag. Refs #1797
Adds the record the per-device default is seeded from, with its raw per-arm data and server logs. Affine pairing, MLX's default TF32, n = 3 per arm, classic bracketing the widths at both ends and overlapping: width 4 serves 67.22 to 67.91 tok/s against classic's 57.98 to 58.71, and the current default of 16 serves 41.88 to 41.97. The per-round device sync steps 24.5, 30.0, 40.8 across widths 3, 4, 5, so the fourth row costs 22% and the fifth costs 36%. That is the `dispatch_multirow_width` boundary on a single row; #1782 measured 4 and 6 but not 5 and could only show it as a range. Two findings the sweep was not looking for. Greedy output at every speculative width, including the 16 that `main` resolves today, differs from classic decode, and nothing gates it because the Qwen 3.5 DFlash target takes the permissive default `exactness_allows`. And the Laguna NVFP4 pairing does not speculate at all on an unmodified server: its own probe reports about half the logit bytes differing from the chain and declines to classic at every width, so no NVFP4 policy entry follows. The host line carries the kernel release, the NVIDIA driver version and the pinned `MLX_CUDA_ARCHITECTURES` for the first time in this tree. The whole sweep cost zero `NV_ERR_NO_MEMORY` errors. Refs #1797
The record called "108477 of 200704 logit bytes differ" far too large to be a rounding difference and inferred something structural in the Laguna verify path. That inference is wrong. 200704 bytes is 50176 f32 logits, so every value differing in its low two bytes is already 50%, and the measured 54% is what a last-mantissa-bits difference across the whole vector looks like. Measured under MLXCEL_MTP_ALLOW_INEXACT, the same pairing accepts 0.887 of its drafter's proposals at width 2, which is only possible if the verify's argmax is almost entirely intact. That also fixes an overstated parallel between the two pairings. They share a class, not a consequence: Laguna's difference changes almost no decision and its probe declines anyway because the probe compares bytes, while Qwen 3.5's changes served output about once in thirty tokens and has no probe at all. The gate is what differs usefully, not the defect. Refs #1797
The harness and the record justified pinning MLX_CUDA_ARCHITECTURES=121 by claiming the `a` suffix could move the kernel dispatch boundaries the sweep measures. It cannot. In the pinned MLX tree the only code keyed on `__CUDA_ARCH_SPECIFIC__` is the NVFP4 weight-quantization converter in `nvfp4_quantize.cuh`, reached only from the weight-quantization kernel; `qmv`, `fp_qmv` and `qmm_sm80` compile identically under `121` and `121a`, and `MLX_CUDA_SM90A_ENABLED` keys on the literal `"90a"` so it never fires at either. The pin stays, for the plainer reason: `.github/workflows/release.yml` ships `90a;100;121`, every earlier GB10 record pins `121`, and a measurement that seeds a shipped default should describe the binaries users run rather than a locally auto-detected variant. Refs #1797
The issue predicted that the 4-to-5 accumulator cliff would be absent on the NVFP4 path, because `dispatch_multirow_width` lives only in `qmv` and `fp_qmv` has no accumulator dispatch, while the `M * B < 8` family switch is on both. The Laguna pairing declines to classic at every width on an unmodified server, so the prediction was tested with `MLXCEL_MTP_ALLOW_INEXACT=1`, which makes those rows ineligible to seed anything and useful only as mechanism evidence. Both boundaries behaved as the pinned tree says they should, and the two families came out mirror images. Over width 4 to 5 the affine family loses 14% of its throughput and steps 10.8 ms per round, 1.84x its neighbouring increments; the NVFP4 family gains 1.2% and steps 5.2 ms, 1.47x, just under the threshold fixed before either number was read. Over width 7 to 8 the NVFP4 sync steps 6.4 ms against a 3.2 ms median while the affine one costs 1.9 ms against 6.2 and loses its throughput to acceptance instead. Each family shows exactly the boundary its kernel has. The NVFP4 curve peaks at width 7 at 1.79x classic, four widths away from the affine plateau at 2 to 4, on the same host and the same binary. That is the argument against a single constant, measured rather than asserted. No NVFP4 entry follows, because an unmodified server declines at every one of those widths. Also records the driver budget: 34 arms, zero NV_ERR_NO_MEMORY errors, against the 38 per run a comparable resident set cost on the previous driver. Refs #1797
The previous commit described that pass as having no closing classic bracket. It has one, and the two do not overlap: 29.88 to 30.66 opening against 29.45 to 29.76 closing, so the host lost about 2% across the pass. The ratios are recomputed against both classic arms, which moves the peak from 1.79x to 1.82x at width 7. The record now says which conclusions survive that and which do not. The 4-to-5 comparison the arm exists for is a 14% loss on the affine side against a 1.3% gain here, and 2% of drift cannot turn one into the other; the 7-to-8 drop is 12%; and the per-round device sync is device time for a fixed row count, whose increments reproduced within a few percent across two independent affine passes. What does not survive is any ranking of widths 4, 5 and 6 against each other, whose spread is inside the drift. Nothing is seeded from this pass, so that costs nothing. Also drops the undated claim that the pinned CUDA architecture matches what ships, since the shipped list may move. The pin is for comparability with every earlier GB10 record, which does not expire. Refs #1797
The kernel-dispatch story the forced-TF32-off grouping suggested is wrong, and a control says so rather than an argument. MLXCEL_QMV_MULTIROW=0 routes a 4-row verify through the same single-row qmv_kernel classic decode uses, launched four times. If that kernel were what the verify and the chain disagree about, width 4 would become byte-identical to classic under it. It does not: width 4's completion hashes identically with the kill switch on and off, while classic keeps its own hash. The multirow kernel makes no difference to the result at all. So the divergence lives elsewhere in the verify path. The record now names the candidates it cannot separate, in the order worth trying: the gated-delta chunked scan's bit-equality with the single-token chain, the materialized attention mask the full-attention layers take at T greater than 1, and the first-bonus and rollback handling. Separating them needs a per-layer comparison rather than a throughput sweep. Also records that the host gate held twice on foreign work during this session, once on another session's model process and once on its CUDA compile, neither of which landed inside a timed run. Refs #1797
…ount The control's JSONL and server logs were committed from a snapshot taken before its width-8 arm finished, so the width-8 row was missing from the raw data even though the record cites it. The driver-budget paragraph also counted 34 arms across four passes, which predated the mechanism sweep's closing bracket and the control itself. The session totals 39 arms across five passes, every one of them at an NV_ERR_NO_MEMORY delta of zero, with the cumulative count for the boot still zero afterwards. Refs #1797
inureyes
added
status:review
`SpeculativeDispatch::resolve` gained a target-path argument and the `Mtp` and `DFlash` variants gained `block_size_source`. The unit tests in `src/server/speculative_dispatch_tests.rs` were updated for both; `tests/speculative_dispatch.rs`, which exercises the same API from outside the crate, was not, and it stopped compiling: 21 arity errors and 2 exhaustive patterns missing the new field. A local clippy scoped to `--bin mlxcel-server --bin mlxcel` never builds `tests/`, so nothing local caught it and CI's `--lib --tests` did. Every existing assertion is preserved rather than relaxed. The file pins the dispatch matrix and the flat per-kind block-size defaults, and those defaults are what the policy returns exactly when it has no measured entry for the target, so the fixture is a target directory with no `config.json`: it classifies as `TargetQuantization::Unknown`, which is the no-entry arm on every host, and the tests stay independent of the machine that runs them. One test is added. `resolve` is keyed on the target's quantization rather than the drafter's, which is easy to get backwards and invisible in a block size alone, because both checkpoints can classify the same way. The new case pairs an affine-quantized target with a drafter that declares no quantization and asserts the resolved source carries `Affine`, so reading the wrong path fails it. Refs #1797
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
The DFlash draft block width resolved to a flat 16 with no hardware or quantization input, and on GB10 that loses: served end to end, the Qwen 3.5 DFlash pairing runs 41.88 to 41.97 tok/s at 16 against 57.98 to 58.71 for classic decode, while width 4 runs 67.22 to 67.91. The throughput PR #1795 bought was reachable only by passing
--draft-block-sizeby hand.New
src/cli/draft_block_policy.rsowns a pure policy over (drafter kind, CUDA compute capability, target quantization) plus a peek that classifies the target checkpoint's quantization from itsconfig.json. It is a default only: it resolves after an operator override and after a drafter checkpoint's own declared width, and before the flat per-kind constant, so every host and quantization path with no measurement is bit-for-bit unchanged. One entry is seeded, compute capability 12.1 with an affine-quantized target on a DFlash drafter, at 4.What changed
resolve_draft_block_sizereturns the width with its provenance, andSpeculativeDispatchcarries that into the worker's startup summary, so the log saysblock_size=4, block_size_from=measured default for sm_121 with an affine-quantized targetrather than just a number. The two policy arms name the compute capability and the quantization they were keyed on, because an unmeasured device and an unreadable target config otherwise print the same sentence and call for different responses.scripts/bench_block_width.shgains the two DFlash pairings and records the kernel release, NVIDIA driver version andMLX_CUDA_ARCHITECTURESin its host line; none of the three was captured anywhere in this tree before, which left cross-session GB10 comparisons with nothing to key on. The served sweep harness is new, underdocs/benchmark_results/data/draft-block-width-gb10-2026-09-20/harness/, and imports the perf(cuda): bucket the cuDNN SDPA plan-cache key across verify rounds #1820 host gate rather than copying it.docs/environment-variables.mdanddocs/CONTINUOUS_BATCHING.mdno longer state a flat DFlash default of 16.Acceptance criteria
Every criterion is met except one, and the issue's checklist now records which. The exception:
Not met: greedy byte-identity. Output at the new default is not byte-identical to classic decode. Neither is it at any other width, including the 16 that
mainresolves today, so this is the existing behavior rather than something this PR introduces, and the two classic arms of the same session are byte-identical to each other. Nothing gates it: the Qwen 3.5 DFlash target takes the permissive defaultexactness_allows, while the LFM2, Muse Glimmer and Laguna targets run a block-versus-chain probe. Filed as #1935.The second-family criterion is satisfied by measurement, not by the statement-of-omission escape hatch: PR #1771 merged on 2026-09-11 and the issue's claim that it was still open is stale. The Laguna NVFP4 pairing was measured, its own exactness probe declines at every width on an unmodified server, and no NVFP4 entry follows. Measured again with the contract forfeited by hand it confirms the issue's falsifiable prediction, each family stepping only where its own kernel has a boundary, with the two optima four widths apart; the numbers and the pre-registered decision rule are in the record.
Two things in the record are worth a reviewer's eye beyond the tables: a
MLXCEL_QMV_MULTIROW=0control that rules the multirow quantized-matmul kernel out as the cause of the identity divergence, and a drift check that the NVFP4 mechanism pass failed, with a statement of which of its conclusions survive that and which do not.Test plan
cargo fmt --all -- --checkcargo test --release --features cuda -p mlxcel --lib -- --test-threads=1 draft_block_policy speculative_args speculative_dispatch: 56 passed, 0 failedNV_ERR_NO_MEMORYerrors across the whole sessionmake verify-test-cudawas not run locally:--profile test-fastgivesmlxcel-coreits own build-script unit and so recompiles MLX from source rather than reusing the warm release build. The scoped release-profile run above covers the changed library code. The integration tests, includingtests/cli_help_consistency.rs, were not run locally either; that one asserts the--draft-block-sizeflag name and signature rather than its help text, and neither changed.cargo clippy -p mlxcel --lib --tests -- -D warnings, the gate that caught this, passes locally on the 1.97.1 toolchain thatrust-toolchain.tomlpins. TheOpenXLA feature compilegate could not be reproduced locally:mlxcel-xla's build script aborts before compiling anything becauseIREE_DISTis unset and no IREE distribution is installed on this host, which CI provisions as a job step. That job's failure was the same 23 errors in the same file, so it is expected to clear with this fix, and CI is the check for it rather than a local run.An earlier revision of this body reported six pre-existing clippy lints and claimed CI never lints the cuda feature set. Both were wrong and are withdrawn. The lints came from running clippy under a 1.93.1 toolchain that
rust-toolchain.tomldoes not pin; all six reproduce on unmodifiedmainunder 1.93.1 and none under 1.97.1, which is the comparison that identified them as toolchain noise rather than a repository problem. CI does gate this crate, which is how it caught what the local run could not:tests/speculative_dispatch.rscallsSpeculativeDispatch::resolveand destructures its variants, so the signature change and the newblock_size_sourcefield broke it, and a--bin-scoped clippy never buildstests/.Closes #1797