perf(mtp): combine greedy token and confidence reads (+0.65% throughput) - #288
Open
ArthurOstapenko wants to merge 2 commits into
Open
perf(mtp): combine greedy token and confidence reads (+0.65% throughput)#288ArthurOstapenko wants to merge 2 commits into
ArthurOstapenko wants to merge 2 commits into
Conversation
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.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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
CONTRIBUTING.mdsmoke trio: 268 passed.compileall, repository hygiene, andgit diff --check: passed.twine check.scripts/fresh_venv_smoke.sh: passed.Benchmark Evidence
Measured on 2026-08-18 at commit
5d0040871fd8366ab823cd645a670c28ee571de7.with one BF16 MTP layer.
cache, cycle MTP history.
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.
prose +0.199%.