[MHA] Add a delta race, resume, finalist rounds and evidence to the MHA forward tuner by amd-bartgips · Pull Request #5849 · ROCm/aiter

amd-bartgips · 2026-09-25T13:36:21Z

Motivation

The MHA forward tuner in #5762 times every candidate once and publishes a winner only if it beats what the shape runs today by at least 3%. That is enough for a first tuned table, but it has three gaps:

A full field is slow: every Triton and Gluon tile config is timed to full depth, even the ones that are clearly losing after a few calls.
An interrupted run starts over.
One timing per candidate cannot tell a real 4% win from run-to-run noise, and re-tuning a shape gates against auto-select rather than against the row already published.

This PR closes those gaps, on top of the shared tuning infrastructure in #5846.

What this PR adds

--strategy race: candidates are timed in randomized short blocks and dropped once they are proven more than delta slower than the leader. The race core is family-independent (aiter/utility/block_race.py).
Resume: a journal of finished measurements, so --resume continues an interrupted run. Its name digests the shape, the candidate field and the race thresholds, so a changed run never replays another run's blocks.
Finalist rounds: under the exhaustive strategy, the fastest candidates are timed again over fresh workers. A challenger must then beat the incumbent by the 3% minimum and by twice the combined standard error of the round means, whichever is larger.
Evidence manifest per run (<tune_file>.evidence.json): per-shape outcome (published, retained, failed, unmeasured), every promotion decision with its margin, bar and noise floor, the selection proofs, and the software and measurement settings.
Re-tuning defends the published row: with --all, a shape that already has a row is gated against that row, not against auto-select. Without --all, already-tuned shapes are still skipped.

Reading this diff/PR

Only the last two commits are this PR's: Add the race parameters to the tuning policy and Add the delta race, resume journal, finalist rounds and evidence to the MHA tuner, +2925 / −209 across 11 files. Below them are #5846's seven commits, cherry-picked so this branch builds on #5762; they drop out when #5846 merges.

aiter/utility/tuning_policy.py (+92): FinalistPolicy (finalists, rounds, significance sigma), RacePolicy (error budget and block counts), SAMPLE_SEED, standard_error_bar in percent, and a noise_pct floor on gate_against_incumbent, which then promotes only at max(min_improvement_pct, noise_pct) and reports that bar as bar_pct.
op_tests/tuners/tune_mha_fwd.py: the race, journal, finalist rounds, evidence and published-row incumbent. Every gate decision goes through gate_against_incumbent. The minimum improvement comes from --min_improvement_pct, the same flag --compare uses, and --delta defaults to it. The finalists, rounds, sigma, race parameters and seed come from the shared policy; mha_fwd_policy.py only aliases them as MHA_FWD_PROMOTION, MHA_FWD_FINALISTS and MHA_FWD_RACE.
aiter/utility/block_race.py (674) and test_block_race.py (536): the race core and its tests, from synthetic latencies.

Nothing here runs in a serving process: the dispatch path is #5762's and unchanged.

Technical Details

One bar, two noise floors. A challenger replaces the incumbent only if it is faster by at least max(min_improvement_pct, noise floor). The minimum, 3% by default, is the bar --compare --update_improved applies, so a row the search publishes is one --compare would also keep. The noise floor is what the search can resolve: twice the combined standard error of the finalist round means under the exhaustive strategy, and the race's delta under --strategy race. More rounds can therefore reject a win that is within noise, but never publish one below the minimum. --min_improvement_pct sets the minimum for both the search and --compare; --delta only changes the race's indifference zone.

The race. Each block times every surviving candidate for a few calls in a random order. After a minimum number of blocks, a candidate is eliminated once a paired t-test shows it is more than delta slower than the leader, with a Bonferroni-corrected error budget. The incumbent races as a protected entrant, so the final gate compares two estimates from the same race and a process-wide offset cancels. A race that ends without certifying a winner publishes nothing rather than falling back to the fastest point estimate.

Correctness. race() only sees timings, never kernel output. Every candidate is correctness-checked in the pass that compiles and warms it, before it enters the race.

Faults. With mp_tuner's typed statuses from #5846, a candidate that faults the worker is reported as crash and costs only itself, not the measurements already taken for its shape. That holds when the worker dies outright and when the group aborts before launch. Candidates behind the fault are not_run, never passed to result_callback, so they have no journal entry and --resume retries them.

When the published row no longer resolves. A re-tune then gates against auto-select instead. If auto-select holds, no row is written, and writing no row cannot remove the old one, so the tuner warns that the row still overrides auto-select and should be deleted.

Test Plan

python3 -m unittest over the level 0 and 1 modules in .github/workflows/tuning-tests.yaml, plus test_compare_logic and op_tests.test_jit_cache_transaction (CPU): race elimination, resume and journal binding, race-versus-exhaustive agreement, finalist rounds and the noise floor on the gate, the published-row incumbent, skipping already-tuned shapes, and the policy module.
On one GPU: test_tune_pipeline.TestTunePipeline.test_mha_fwd_mp1 (measure, gate, publish, prove and write evidence for one shape) and test_mp_tuner_fault.
A real gfx950 run of both strategies on two shapes.

Test Result

On gfx950 / MI355X, one GPU (HIP_VISIBLE_DEVICES=1):

320 CPU tests pass; ruff check and black are clean on the changed files.
test_mha_fwd_mp1 and both test_mp_tuner_fault tests pass.
Two hd128 bf16 varlen shapes (batch 2 / 512 tokens non-causal, batch 4 / 8192 tokens causal), --candidate-sample 8 (10 candidates per shape with the shipped defaults and the incumbent): the exhaustive strategy took 64 s and the race 17 s. Both reached the same decision on both shapes: asm_v3 auto-select (35 us and 87 us) was not beaten by 3%, so no row was written.
On a 200-candidate gfx950 field, measured with the earlier 2% bar, the race took 104 s against 315 s for the exhaustive strategy.

Not validated here: the machine has only gfx950, so the gfx942 paths are untested on this revision.

Related PRs

[MHA] Add tuned-config dispatch and an exhaustive tuner for packed varlen #5762: the MHA forward tuned-config dispatch and exhaustive tuner this PR builds on. This PR targets its branch.
[Tuning] Share mp_tuner typed statuses and add a central tuning policy #5846: shared mp_tuner typed statuses, MpTunerTask, measurement_kwargs and the central tuning policy with its single min_improvement_pct, which this PR extends with the finalist and race policies and the noise floor.
[FlyDSL] Add gfx950 FHMoE tuner on gemm_moe_tune.py #5751: the FHMoE tuner. Its follow-up can reuse PromotionPolicy, gate_against_incumbent and, if it wants a race, block_race with RacePolicy.

Submission Checklist

Look over the contributing guidelines at https://github.com/ROCm/TheRock/blob/main/GOVERNANCE.md#pull-requests.

--compare --update_improved and gate_against_incumbent both decide whether a faster configuration is worth replacing what a shape runs today, but read two different numbers (3% and 2%). They now read one: PromotionPolicy.min_improvement_pct, 3.0, named after the --min_improvement_pct flag whose default it already was. - The gate works in percent and promotes at "at least" the bar, as --compare does, and reports margin_pct. - COMPARE_MIN_IMPROVEMENT_PCT and indifference_delta are gone; neither was on main. - finalists and finalist_rounds move to the follow-up that times finalists (#5849), since nothing on main uses them. - batch comes from a new RunPolicy, so no TunerCommon default is a bare literal any more. It sits apart from MeasurementPolicy because it changes how often the tuned CSV is written, not which candidate wins. No default a tuner sees changes. Co-authored-by: Cursor <cursoragent@cursor.com>

--compare --update_improved and gate_against_incumbent both decide whether a faster configuration is worth replacing what a shape runs today, but read two different numbers (3% and 2%). They now read one: PromotionPolicy.min_improvement_pct, 3.0, named after the --min_improvement_pct flag whose default it already was. - The gate works in percent and promotes at "at least" the bar, as --compare does, and reports margin_pct. - COMPARE_MIN_IMPROVEMENT_PCT and indifference_delta are gone; neither was on main. - finalists and finalist_rounds move to the follow-up that times finalists (#5849), since nothing on main uses them. - batch comes from a new RunPolicy, so no TunerCommon default is a bare literal any more. It sits apart from MeasurementPolicy because it changes how often the tuned CSV is written, not which candidate wins. No default a tuner sees changes.

Moved unchanged from the MHA forward tuner branch (#5762) so that other tuners can use them without waiting for it: - mp_tuner: optional per-candidate status and detail (return_status), a per-candidate progress queue with result_callback, a faulted group keeping the candidates it had already measured, MpTunerTask, broader accelerator-fault detection, and a pool restart on a stale PID map. - TunerCommon.measurement_kwargs, so --warmup and --iters reach the timed call instead of run_perftest's defaults. - _read_csv builds a new frame rather than assigning into a slice. - FileBaton treats a zombie lock holder as dead, so a builder that exited without releasing its lock no longer wedges the module build.

A caller that does not pass return_status picks the fastest finite latency and cannot tell a partial group from a complete one. Keeping the candidates measured before a fault would let it publish the winner of a search that stopped early, so those callers get every candidate of the group back as failed, as before. Callers with return_status keep the measured candidates and see the crash or timeout status on the rest. test_mp_tuner_fault runs one shape group with a faulting candidate through the real pool on one GPU, in both call forms.

aiter/utility/tuning_policy.py holds the thresholds whose justification does not mention a kernel family: - MeasurementPolicy: warmup, iters, errRatio and timeout, which TunerCommon.ARG_DEFAULTS now reads (same values as before). - COMPARE_MIN_IMPROVEMENT_PCT: the --compare --update_improved bar. - PromotionPolicy: the indifference delta a challenger has to clear over the incumbent, and the finalist count and rounds for re-timing. - gate_against_incumbent: the promote-or-retain decision on those numbers, for tuners that measure the incumbent themselves. Standard library only, so the runtime and CPU tests can import it. A family that needs a different threshold overrides it with dataclasses.replace in its own module.

--compare --update_improved and gate_against_incumbent both decide whether a faster configuration is worth replacing what a shape runs today, but read two different numbers (3% and 2%). They now read one: PromotionPolicy.min_improvement_pct, 3.0, named after the --min_improvement_pct flag whose default it already was. - The gate works in percent and promotes at "at least" the bar, as --compare does, and reports margin_pct. - COMPARE_MIN_IMPROVEMENT_PCT and indifference_delta are gone; neither was on main. - finalists and finalist_rounds move to the follow-up that times finalists (#5849), since nothing on main uses them. - batch comes from a new RunPolicy, so no TunerCommon default is a bare literal any more. It sits apart from MeasurementPolicy because it changes how often the tuned CSV is written, not which candidate wins. No default a tuner sees changes.

test_mp_tuner_logic and test_mp_tuner_fault imported Triton before torch, and the start-time test switched its pool from spawn to fork, both for a dynamic-loader abort in spawned workers. Main runs the same test, and the worker-PID test from #5841, under spawn with no Triton import, and the abort does not reproduce on MI355X. The tests now use spawn, the start method mp_tuner forces, and a CPU test no longer needs Triton installed.

Worker death is now caught by #5841's check, which this branch rebases onto. Without it, the immediate pool restart on a stale PID map resubmitted the dead group with a fresh start time, so a candidate that kills its worker restarted the pool forever. - The worker-exited branch drains the progress queue before failing the group, as the other failure branches do, so a return_status caller keeps the candidates the worker measured before it died. The detail names the exited worker. - With return_status, the candidates behind a fault come back as not_run instead of carrying the fault's status. The callback stream is unchanged: it still receives only the candidate that failed. - The mp_tuner docstring describes both result forms, every status and result_callback. - test_mp_tuner_fault adds a candidate that exits its worker with another shape queued behind it, in both call forms. With the dead-worker check disabled it fails at its deadline after repeated pool restarts. The file now runs in the per-PR shards (67 s on one MI355X).

work_group catches an error raised outside worker(), such as a later candidate's gen_data or reference failing, and returns a result list itself. That list failed every candidate of the group, including those already measured, and put all of them on the progress queue. The parent saw a normal return, so its partial-result handling never ran: a return_status caller lost the measured candidates, a checkpointing caller received a second, crash result for each of them, and the candidates behind the abort were checkpointed as tried, so a resume skipped them. work_group now builds those results with _failed_group_results, the same rules the parent applies to a group the pool lost: with return_status, measured candidates keep their results, the one that aborted gets the failure, and the rest are not_run; without it, the whole group fails as before. Only the aborting candidate is published. test_mp_tuner_fault adds a group whose second candidate fails in gen_data (reported by Araceli); the typed case failed before this change.

RacePolicy holds an elimination race's error budget and block counts, and SAMPLE_SEED seeds candidate sampling and block order, so a run is repeatable. The race's indifference zone defaults to PromotionPolicy.min_improvement_pct, so the race and the gate after it agree on what a tie is. Their first consumer is the MHA forward tuner's race.

…he MHA tuner On top of the exhaustive tuner: - --strategy race: candidates are timed in randomized short blocks and dropped once proven more than delta slower than the leader (aiter/utility/block_race.py), about 3x faster than timing the whole field to full depth on a 200-candidate field. - A journal of finished measurements, so --resume continues an interrupted run; its name digests the shape, the field and the race thresholds, so a changed run never replays another's blocks. - Finalist rounds: the fastest candidates are timed again over fresh workers. FinalistPolicy (finalists, rounds, significance_sigma) sits next to RacePolicy in the shared tuning policy. - An evidence manifest per run: per-shape outcome, every promotion decision with its margin_pct, bar_pct and noise_floor_pct, selection proofs, software and measurement settings, and the promotion bar. - With --all, the published row is the incumbent a challenger has to beat, instead of auto-select. Gate decisions go through tuning_policy.gate_against_incumbent, which takes a noise_pct floor and promotes only at max(min_improvement_pct, noise_pct): the minimum improvement says what is worth a row, the noise says what the search resolved, so enough finalist rounds never publish a real but tiny win that --compare would then refuse. The floor is two combined standard errors of the round means under the exhaustive strategy (standard_error_bar, in percent) and the race's delta under --strategy race. The tuner reads the bar from --min_improvement_pct, so one flag sets it for both the search and --compare; --delta defaults to it and is only needed to make the race's indifference zone differ. mp_tuner's typed statuses keep a faulting candidate from costing its shape the measurements already taken; the tuner checkpoints through result_callback, so candidates not reached before a fault are retried on resume. A re-tune whose published row no longer resolves now says it gated against auto-select, and warns that the row still overrides it.

This was referenced Sep 25, 2026

[MHA] Add tuned-config dispatch and an exhaustive tuner for packed varlen #5762

Draft

[Tuning] Share mp_tuner typed statuses and add a central tuning policy #5846

Draft

amd-bartgips force-pushed the users/bartgips/fmha-single-store branch from 8776074 to 8a11428 Compare September 28, 2026 14:17

amd-bartgips force-pushed the users/bartgips/fmha-tuning-advanced branch from ae0d49f to aa3b347 Compare September 28, 2026 14:17

amd-bartgips force-pushed the users/bartgips/fmha-tuning-advanced branch from aa3b347 to 1752978 Compare October 1, 2026 12:38

amd-bartgips added 9 commits October 1, 2026 12:52

amd-bartgips force-pushed the users/bartgips/fmha-tuning-advanced branch from 1752978 to 6647aaa Compare October 1, 2026 12:55

amd-bartgips force-pushed the users/bartgips/fmha-single-store branch from 8a11428 to ab166eb Compare October 1, 2026 12:55

Navigation Menu

Repository navigation

[MHA] Add a delta race, resume, finalist rounds and evidence to the MHA forward tuner - #5849

[MHA] Add a delta race, resume, finalist rounds and evidence to the MHA forward tuner#5849
amd-bartgips wants to merge 9 commits into
users/bartgips/fmha-single-storefrom
users/bartgips/fmha-tuning-advanced

amd-bartgips commented Sep 25, 2026 •

edited

Loading

Uh oh!

This branch has not been deployed

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Conversation

amd-bartgips commented Sep 25, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Motivation

What this PR adds

Reading this diff/PR

Technical Details

Test Plan

Test Result

Related PRs

Submission Checklist

Uh oh!

This branch has not been deployed

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

amd-bartgips commented Sep 25, 2026 •

edited

Loading