perf(speculative): the default draft block width of 16 loses on GB10
- Dominant language
- Rust
- Stars
- 467
- Forks
- 54
- Avg merge
- 4h 25m
- Merged PRs (30d)
- 310
Description
## Problem / Background
Follow-up to #1782 (PR #1795, merged as `e391ae9c`), deliberately excluded from that PR. With the fixed per-round verify cost removed, speculative decoding on the Qwen 3.5 DFlash pairing is a net win at draft block widths 2 to 4 and a net loss at the effective default of 16, so the throughput the fix bought is unreachable unless the operator passes `--draft-block-size` by hand.
Measured on GB10 (sm_121, DGX Spark), MLX pin `81ba1c6a`, CUDA release build, `models/mlx/qwen3.5-4b-4bit` (affine 4-bit, group 64) with `models/mlx/qwen3.5-4b-dflash`, one `mlxcel-server` per configuration, greedy 200-token streaming completions on a fixed 165-token prompt, n = 3 after a discarded warm-up, GPU held exclusively on an otherwise idle host, same binary on every arm. Full record: `docs/benchmark_results/dflash-verify-fixed-cost-gb10-2026-09-11.md`. Attribution: `TECHNICAL_REPORTS/1782-dflash-verify-fixed-cost-cuda-20260911.en.md`.
| width | e2e tok/s mean (min to max) | vs classic | round device sync |
|---|---|---|---|
| classic | 58.33 (57.27 to 59.38) | | |
| 2 | 72.57 (72.21 to 72.99) | 1.24x | 18.6 ms |
| 3 | 76.42 (76.36 to 76.51) | 1.31x | 24.4 ms |
| 4 | 76.88 (76.62 to 77.28) | 1.32x | 29.9 ms |
| 6 | 59.06 (58.84 to 59.19) | 1.01x | 47.4 ms |
| 7 | 53.45 (53.32 to 53.59) | 0.92x | 53.9 ms |
| 8 | 48.84 (48.48 to 49.32) | 0.84x | 59.2 ms |
| 16 (effective default) | 48.63 (48.48 to 48.83) | 0.83x | 64.1 ms |
Within-configuration spread is under 2% and every arm is greedy byte-identical to classic decode, so the ordering is not the GB10 single-run bimodality of #755. The measured setting for this pairing today is `--draft-block-size 4`.
## Current Behavior
The width is a flat constant resolved at startup with no hardware or quantization input. `resolve_draft_block_size` (`src/cli/speculative_args.rs:222-273`) returns `--draft-block-size` when given, then peeks the checkpoint only for an LFM2 DSpark drafter or a Muse Glimmer assistant, and otherwise falls through to `default_block_size_for_kind` (`src/cli/speculative_args.rs:186-201`), which is `DEFAULT_DFLASH_BLOCK_SIZE = 16` (`src/cli/speculative_args.rs:71`) for every `DrafterKind::Dflash`. The Qwen 3.5 DFlash drafter matches neither peek, so it runs at 16. That call site receives only the override, the resolved kind and the drafter path, so the device and the target's quantization mode are not in scope there at all.
No runtime controller narrows it for this pairing either. `DFlashGenerator::run` builds a `BlockThroughputController` only when the drafter declares a configured depth strictly between 1 and the requested width and does not prefer the requested width (`src/lib/mlxcel-core/src/drafter/dflash/round_loop.rs:699-714`), and `DFlashDrafter::configured_block_size` returns `Some` only for DSpark (`src/lib/mlxcel-core/src/drafter/dflash/drafter.rs:623-626`). The Qwen 3.5 pairing therefore takes the `None` arm and every round runs at `block_size_cfg` (`round_loop.rs:744-747`). Where a controller does exist it is a local inside `run`, and the generator itself is constructed per request (`src/server/batch/dflash_target.rs:607-612`) and dropped when the request finishes, so a measurement window can never amortize across requests. `draft_block_size` is also worker-owned in the live settings API and refused with "owned by the model worker; restart required" (`src/server/runtime_settings.rs:29,513-514`), so a served width cannot be retuned without a restart.
## Why the cliff sits between 4 and 6, as mechanism
Two kernel boundaries in the pinned MLX tree, not one, and they are at different row counts.
1. **Accumulator width, 5 to 7 rows pay for 8.** `dispatch_multirow_width` instantiates exactly three compile-time widths, 2, 4 and `Cap` (`src/lib/mlx-cpp/patches/mlx/backend/cuda/quantized/qmm/qmv.cu:925-939`), and the caller passes `Cap = max_x_rows = 8` (`qmv.cu:975-984`), gating the path at `broadcast_w && x_rows >= 2 && x_rows <= window` (`qmv.cu:978-979`). A 5, 6 or 7 row verify therefore takes the 8-wide instantiation and its register cost without using the extra rows. The in-tree comment measures that cost on a V100 with `cuobjdump -res-usage` at `elems_per_thread` 16: `qmv_kernel` 61 registers against `qmv_multirow_kernel<..., 8>` 168 (`qmv.cu:903-916`). No equivalent register measurement exists for sm_121, so treat the register figure as the V100 datum it is; what is measured on GB10 is the round device sync jumping from 29.9 ms at 4 rows to 47.4 ms at 6. `docs/CONTINUOUS_BATCHING.md:72` already records the multirow window as covering 2 to 7 rows on GB10.
2. **Kernel family switch, at 8 rows.** `if (can_use_qmv && (M * B < 8))` selects `qmv`, else `qmm_sm80` (`src/lib/mlx-cpp/patches/mlx/backend/cuda/quantized/quantized.cpp:283-289`). That is why 8 and 16 cost nearly the same per round (59.2 and 64.1 ms) and why neither wins: the `qmm_sm80` term was not touched by PR #1795.
## Family dependence, which is why this is not one constant
Do not inherit the Qwen 3.5 crossover for other families. Laguna is NVFP4, so `supports_fp_qmv` accepts it and `call_qmv` routes to `fp_qmv` instead of `qmv` (`quantized.cpp:195-196,248-250`). `supports_fp_qmv` carries its own row bound, `if (vec_batch > 8) return false;`, alongside `mode == QuantizationMode::Affine` returning false and compute capability 9 or below returning false (MLX at pin `81ba1c6a`, `mlx/backend/cuda/quantized/qmm/qmm.cu:127-157`). Crucially, `fp_qmv` has **no** multirow accumulator-width dispatch at all: `fp_qmv.cu` instantiates only `fp_qmv_single` and `fp_qmv_batched`, and its `rows_per_block = 8` (`fp_qmv.cu:17`) tiles output rows, not input rows. So boundary 1 above, the 5-to-7 register cliff, does not exist on the NVFP4 path, while boundary 2 does. The two families can therefore have genuinely different optimal widths, and the Laguna DFlash pairing cannot be measured yet because PR #1771 is still open.
## Proposed Solution
Replace the flat `DEFAULT_DFLASH_BLOCK_SIZE` fallback with a default that is a function of the runtime device and the target's quantization path, keeping `--draft-block-size`, `MLXCEL_DRAFT_BLOCK_SIZE` and `LLAMA_ARG_DRAFT_BLOCK_SIZE` as overrides that always win. Concretely:
- Add a pure, unit-testable policy function next to the existing per-kind default in `src/cli/speculative_args.rs`, taking the drafter kind, `Option<(u32, u32)>` compute capability and the target's quantization mode, and returning the width. Keep `default_block_size_for_kind` as its fallback for every host the policy has no measurement for, so an unmeasured platform is unchanged.
- Thread the two new inputs to the call site in `src/server/speculative_dispatch.rs:249` and the offline `generate` path, rather than reading the environment inside the policy.
- Seed it only with widths that have a committed measurement record. On this evidence that is one entry, CUDA compute capability 12.1 with an affine-quantized target on a DFlash drafter, at 4.
- Log the resolved width and the reason at startup so an operator can see whether the default or an override applied, next to the existing `block_size` field on the DFlash diagnostics line (`src/server/batch/dflash_target.rs:640-642`).
Rejected: making the per-request controller do it. `BlockThroughputController` already exists and is deliberately not engaged here, and it cannot help, because the generator is per request (`dflash_target.rs:607`) so the warm-up would be paid on every request, and a short request would pay it without ever reaching the measured width.
Rejected: reusing `MLXCEL_QMV_MULTIROW_MAX_ROWS`. The existing autotuner knob (#906, `src/lib/mlxcel-core/src/autotune/ops/cuda_kernel_knobs.rs:321`, documented at `docs/environment-variables.md:465`) narrows the kernel dispatch window process-wide, so it would tax every batched decode that never asked for it, and its own doc row marks it unvalidated on CUDA. It should be measured against the chosen width, not substituted for it.
## Scope
**In scope:** the width policy in `src/cli/speculative_args.rs` and its two call sites; startup logging of the resolved width; the `MLXCEL_DRAFT_BLOCK_SIZE` row in `docs/environment-variables.md:311`, which currently states the default as "per drafter (`4` for MTP, `16` for DFlash)"; a committed measurement record under `docs/benchmark_results/` for every hardware and quantization arm the policy gains an entry for.
**Out of scope:** the `qmm_sm80` per-round cost at 8 rows and above, which is what leaves 8 and 16 losing and is a kernel question rather than a policy one; the MTP default (`DEFAULT_MTP_BLOCK_SIZE = 4`), which is unmeasured here; the per-request `BlockThroughputController` design; `MLX_MAX_OPS_PER_BUFFER`, filed separately from the same sweep.
## Implementation Notes
- **Reuse**: `mlxcel_core::cuda_arch::cuda_compute_capability()` (`src/lib/mlxcel-core/src/cuda_arch.rs:83-88`) is the existing compute-capability accessor, already consumed by PR #1795's own dtype policy (`src/lib/mlxcel-core/src/drafter/dflash/drafter.rs:202`). Mirror the pure-function-plus-thin-wrapper shape of `drafter_bf16_to_f16_policy` / `drafter_bf16_to_f16_at_load` (`drafter.rs:196-270`), which is unit-tested per arm without a device. For the sweep, extend `scripts/bench_block_width.sh` rather than writing a new harness: it already interleaves widths across rounds so drift spreads across the table instead of pooling at one end, but it hardcodes three MTP pairings and drives `mlxcel generate`, whereas #1782's numbers came from `mlxcel-server`, so both a DFlash pairing and a server arm have to be added.
- **Constraints**: the policy must be a default only, never an override, since the three existing override surfaces all resolve ahead of it. It must return the current 16 on every host and quantization path with no measurement, so nothing changes off GB10. It is resolved once at startup and cannot be revised live, because `draft_block_size` is `WORKER_REASON` in the settings API.
- **Edge cases**: `cuda_compute_capability()` returns `None` on Metal, CPU-only and non-CUDA builds, which must take the unchanged fallback; a width of 1 or 0 must not be reachable, since `round_loop.rs:747-749` breaks out of the loop at `bs <= 1` and would silently disable speculation; the byte-identity probe is keyed on block size (`src/models/muse_glimmer_speculative.rs:275-282` and the LFM2 sibling at `src/models/lfm2_speculative.rs:256-263`), so changing the default changes which width is probed at startup and the new width must be probed, not assumed; a drafter whose checkpoint peek already supplied a width (DSpark, Muse Glimmer) must keep it.
- **Error handling**: an unmeasured hardware or quantization combination is not an error, it takes the fallback silently at `info` level. An out-of-range override keeps today's behavior, which is that concrete generators enforce their own minimums (`src/cli/speculative_args.rs:206-208`).
## Acceptance Criteria
- [ ] On GB10 with `models/mlx/qwen3.5-4b-4bit` plus `models/mlx/qwen3.5-4b-dflash` and no `--draft-block-size`, `mlxcel-server` logs a resolved block size of 4 and the served end-to-end rate beats the classic arm measured in the same session with disjoint n = 3 ranges.
- [ ] On the same host, an explicit `--draft-block-size 16` still yields 16, and `MLXCEL_DRAFT_BLOCK_SIZE=16` with no flag also yields 16.
- [ ] Greedy output at the new default is byte-identical to classic decode on the same prompt, matching the property PR #1795 established at every width.
- [ ] The width is reached through the real serving path, not only the offline `generate` path: verified from `mlxcel-server` startup logs and a served `/v1/completions` request, with the `block_size` field of the DFlash diagnostics line agreeing.
- [ ] A second family is measured before it gets an entry. Either the Laguna NVFP4 pairing once PR #1771 merges, or an explicit statement in the record that it was not measured and therefore takes the fallback.
- [ ] Unit tests cover the policy function per arm (no CUDA, pre-Ampere CUDA, sm_121 affine, sm_121 non-affine, each drafter kind) without requiring a device, in the style of the `drafter_bf16_to_f16_policy` tests.
- [ ] A benchmark record is committed under `docs/benchmark_results/` with host, MLX pin, n, per-arm ranges and the control arm, for every entry the policy gains.
- [ ] `docs/environment-variables.md:311` no longer states a flat DFlash default of 16 when the resolved default is hardware dependent.
## Verification
```
make release-cuda
make verify-fmt
cargo test --profile test-fast --features cuda -p mlxcel speculative_args -- --test-threads=1
make verify-test-cuda
# A/B on GB10, one server per arm, GPU held exclusively, idle host, n = 3 after a discarded warm-up.
# Arm 1 (new default, no flag):
./target/release/mlxcel-server -m models/mlx/qwen3.5-4b-4bit \
--draft-model models/mlx/qwen3.5-4b-dflash --draft-kind dflash \
--ignore-eos --max-batch-size 1
# Arm 2 (control): same command with --draft-block-size 16
# Arm 3 (classic): same command with no --draft-model
```
A pass is: arm 1 logs `block_size = 4`, its three end-to-end rates all sit above the maximum of arm 3's three, and its completion text matches arm 3's byte for byte. `--test-threads=1` is required on this host for every CUDA test run.
## Technical Considerations
`MLX_MAX_OPS_PER_BUFFER=100` was measured in the same sweep and gave about +10% at block 8 but nothing at block 2, so it does not change the ordering above and is filed separately. If a future MLX pin adds a multirow instantiation at 6, or widens `supports_fp_qmv`, boundary 1 moves and the seeded width has to be re-measured against the new pin; the policy should therefore carry the pin it was measured under in a comment, the way `cuda_graph_cache_default` records its removal condition (`src/lib/mlxcel-core/src/hardware.rs:262-267`). Refs #1782, #1795, #1771, #906, #755.
Contributor guide
Research direction
Start with resolve_draft_block_size and default_block_size_for_kind in src/cli/speculative_args.rs, then trace the server and offline generate call sites and the DFlash diagnostics line. Review the existing drafter policy tests and run scripts/bench_block_width.sh for the measured pairing. Done means a tested, measurement-backed default preserves overrides and fallback behavior, logs the reason, updates the documentation, and passes the listed serving and byte-identity checks.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- cpp, rust
- Domain
- backend, performance
- Issue type
- Feature
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Active
- Clarity
- Clearly specified
- Newbie friendliness
- 48/100