Skip to content

perf(speculative): replace the sqrt(K) verify-cost heuristic with measured round cost in the adaptive MTP policy #736

Closed

Description

@inureyes

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_ms and drafter_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, and docs/environment-variables.md if flags change.

Technical Considerations

Activity

  1. inureyes commented on Jul 11, 2026

    @inureyes
    MemberAuthor

    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=1 pins B=1 MTP on for CUDA serving. The measured-round-cost estimator proposed here handles both regimes without hardware heuristics.

  2. self-assigned this
    on Aug 31, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

Labels

Type

No type

Projects

No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions