Repository navigation
perf(speculative): replace the sqrt(K) verify-cost heuristic with measured round cost in the adaptive MTP policy #736
Description
Context
PR #733 (issue #638) made mtp_policy::estimate_speedup_scaled de-rate the verify cost by sqrt(K) on compute-bound hosts, calibrated against a measured "~1.9x round-cost growth from K=2 to K=8" on GB10. That calibration point was mislabeled: --block-size 8 ran at an effective K=4, because the drafter's configured block size is 4 and the adaptive block controller (effective_mtp_block_size) never expanded at 35% acceptance. The true measurement is 2x verify width for 1.94x time, i.e. verify cost is LINEAR in K on GB10.
The 2026-07-11 investigation explains why it is linear: the MTP verify is a [1, K] sequence-axis forward, and the CUDA quantized dispatch (src/lib/mlx-cpp/patches/mlx/backend/cuda/quantized/quantized.cpp:229-245) routes M*B < 8 to per-row qmv (the sm90 CUTLASS path is Hopper-gated), so a K-wide verify re-reads the weights K times. Measured on GB10: classic forward 69 ms, K=2 round 140 ms, K=4 round ~270 ms; the measured 0.77x/0.52x speedups from PR #733 (docs/benchmark_results/speculative-pairing-gb10-2026-07-10.md) equal tokens-per-round divided by round-cost-in-forwards exactly. Same root cause as the flat batched-decode throughput tracked in #725.
Impact
At K=4 the policy assumes a verify-cost multiple of sqrt(4) = 2 while the measured round cost is ~3.9 classic-forward equivalents, so the estimator under-costs rounds by ~2x on CUDA. The Gemma pairing was declined only because its accepted_len was very low (the estimate landed at ~1.0 against the <=1.0 decline threshold); a moderately better pairing could be wrongly enabled and regress live serving. This is a policy-correctness fix, not a tuning tweak.
What's needed
Stop inferring verify cost from a shape heuristic; use what the profiling window already measures.
- The profiling window already records per-round
verify_msanddrafter_ms(src/server/batch/mtp_policy.rs,ProfileAccumulator). - Add a measured classic single-token step time for the same target, either from the scheduler's own decode timing or a short probe.
- Estimate speedup directly as
(accepted_len + 1) / ((verify_ms + drafter_ms + per-round overhead) / classic_step_ms). - Apple Silicon verdicts must remain correct: either byte-identical decisions, or re-validated against M5 Max measurements.
Acceptance Criteria
- Unit tests pinning the measured-cost estimator.
- The GB10 Gemma pairing declines with clear margin, not borderline.
- A bandwidth-bound favorable pairing still enables.
- The sqrt(K) path removed, or demoted to a fallback used only when no classic-step measurement exists.
- Docs updated: the policy section of
docs/benchmark_results/speculative-pairing-gb10-2026-07-10.md, anddocs/environment-variables.mdif flags change.
Technical Considerations
- Refs perf(speculative): CUDA drafter pairing matrix and continuous-batching integration #638, PR perf(speculative): CUDA pairing matrix and burst HOL observability #733, perf(cuda/quant): amortizing small-M quantized GEMM for batched decode on Blackwell #725, and the docs correction issue docs: correct the compute-bound framing of the GB10 MTP regression to the small-M qmm kernel gap #738.
Activity
docs: correct the GB10 MTP regression framing to the small-M qmv kern…
inureyes commentedon Jul 11, 2026
PR #740 (#725) raises this issue's urgency and flips its direction: the
sqrt(K)de-rate was calibrated against the pre-#725 per-row verify, but with the multirow qmv the verify amortizes toward one classic forward, so the de-rate now under-estimates the speedup on CUDA and can wrongly DECLINE a favourable pairing (measured 1.46x at K=4 on GB10, which the sqrt(4)=2 de-rate estimates at ~1.0, the decline boundary). Until this lands,MLXCEL_ENABLE_MTP_B1=1pins B=1 MTP on for CUDA serving. The measured-round-cost estimator proposed here handles both regimes without hardware heuristics.perf(cuda/quant): amortize small-M quantized matmul with a multirow q…
docs: mark the sqrt(K) MTP policy caveat resolved by #736
perf(server): tick-cooperative speculative decoding to remove the bur…
perf(speculative): settle the adaptive MTP policy from measured round…