perf(speculative): the default draft block width of 16 loses on GB10 #1797

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.

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 fix(cuda): attribute and remove the fixed cost of multi-row speculative verify #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. Correction, 2026-09-20: PR #1771 merged on 2026-09-11, a few hours after this issue was filed, so the Laguna pairing is measurable now. Both checkpoints are on the GB10 host and Laguna has a server-side DFlash target (LagunaWrapper in src/server/batch/dflash_target.rs), not only the offline arm from #1351.

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 fix(cuda): attribute and remove the fixed cost of multi-row speculative verify #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 perf(cuda): the multi-row speculative verify forward is launch-bound on GB10 #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 fix(cuda): attribute and remove the fixed cost of multi-row speculative verify #1795 established at every width. Not met. Greedy output differs from classic decode at every width measured, including the 16 that main resolves today, so this is the existing behavior rather than a consequence of the new default. A MLXCEL_QMV_MULTIROW=0 control rules the multirow quantized-matmul kernel out as the cause. Tracked in fix(speculative): Qwen 3.5 DFlash verify diverges from classic decode #1935; see docs/benchmark_results/draft-block-width-default-gb10-2026-09-20.md.
• 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. Since PR feat(speculative): add the Laguna DFlash drafter #1771 merged, the statement-of-omission alternative no longer applies and the Laguna NVFP4 pairing is measured on the same host in the same session. A width only earns a policy entry if it beats classic decode with disjoint n = 3 ranges; if the NVFP4 optimum lands on the same width as affine, the record says the 5-to-7 accumulator mechanism does not explain it rather than smoothing it over.
• 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.