Skip to content

perf(mtp): combine greedy token and confidence reads (+0.65% throughput) - #288

Open
ArthurOstapenko wants to merge 2 commits into
youssofal:mainfrom
ArthurOstapenko:perf/mtp-greedy-confidence-sync
Open

perf(mtp): combine greedy token and confidence reads (+0.65% throughput)#288
ArthurOstapenko wants to merge 2 commits into
youssofal:mainfrom
ArthurOstapenko:perf/mtp-greedy-confidence-sync

Conversation

@ArthurOstapenko

Copy link
Copy Markdown

Summary

ExpectedValue depth selection needs the greedy draft token plus confidence
metrics derived from the row's FP32 top eight. The previous path synchronized
the top-k reduction, then called the regular greedy reader, which launched and
synchronized argmax separately.

This patch queues argmax and top-k before one synchronization and reuses the
selected token. The combined path is limited to greedy draft sampling with the
regular selector. Margin gating, adaptive-width routing, target-prefix routing,
correction caches, adapter ensembles, reranking, and stochastic sampling keep
their existing selection paths.

The new tests cover first-index tie handling, optional one-hot proposals,
confidence values, ExpectedValue policy decisions, selector eligibility, and
key policy and control fields across a complete generation run on a tiny
runtime. No existing tests were changed or removed.

Verification

  • Focused MTP suite: 52 passed.
  • Required CONTRIBUTING.md smoke trio: 268 passed.
  • Full Python suite: 4,148 passed, 20 skipped, 0 failed.
  • Ruff, compileall, repository hygiene, and git diff --check: passed.
  • Wheel and sdist build: passed; both artifacts also passed twine check.
  • scripts/fresh_venv_smoke.sh: passed.

Benchmark Evidence

Measured on 2026-08-18 at commit
5d0040871fd8366ab823cd645a670c28ee571de7.

  • Hardware: Apple M5 Max, 128 GB unified memory.
  • Software: macOS 26.5.2, Python 3.13.13, MLX 0.32.0, mlx-lm 0.31.3.
  • Model: Qwen3.8-27B MTPLX Optimized Quality; affine Q8, group size 64,
    with one BF16 MTP layer.
  • Decode path: ExpectedValue D3, batched stock verification, persistent draft
    cache, cycle MTP history.
  • Sampler: temperature 0, top-p 1, top-k 0; 64 generated tokens per run.

The comparison used one loaded runtime and three prompt classes. Each arm had
one warmup per prompt, followed by 10 measured pairs per prompt. A/B order
alternated by block, prompt order rotated by block, and external wall time was
bracketed by MLX synchronization. The control forced separate token and metric
reads while leaving the model state and remaining decode path unchanged. Every
arm produced the same token, policy, control, and work fingerprints.

  • Paired geometric mean: +0.650% end-to-end throughput.
  • Shared-block bootstrap 95% CI: +0.491% to +0.765%.
  • Per-prompt paired deltas: code +0.691%, JSON/tool +1.062%, creative
    prose +0.199%.
  • First and second halves: +0.745% and +0.555%.
  • Control-first and candidate-first strata: +0.603% and +0.697%.
  • Median draft-phase time fell by 6.36% to 7.81%, depending on the prompt.

Cover first-maximum tie handling, optional one-hot proposals, FP32 top-eight confidence metrics, and the ExpectedValue policy inputs that depend on them.
Schedule greedy argmax and the FP32 top-eight reduction together on the default ExpectedValue path, then reuse the selected token after the shared synchronization. Alternate draft selectors keep their existing readers.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant