perf(nrpt): chain-masked probes — one compiled round loop per autotune#52
Merged
Conversation
XLA compile dominates cold tune_chains (~85% of stage-1 wall on GPU; the actual sampling of both probes is <1s) because each probe at a new chain count recompiles the round loop at new shapes. Two changes make the whole autotune share a single compiled loop: 1. Chain masking: nrpt(pad_chains_to=) pads the ladder to a fixed length with copies of the coldest chain and passes the true count into the jitted loop as TRACED data. Swaps beyond the live prefix are forced-rejected (identity permutation — DEO coupling is nearest-neighbour, so the padding is fully decoupled), attempt counters and the round-trip "top" are masked to the live prefix, and all outputs are sliced back to the true count, so optimize_schedule and every downstream consumer see unpadded semantics. tune_chains(pad_probes=) pads every probe to max_chains; autotune(pad_probes=None) auto-enables on accelerator + template mode (padding Gibbs work is ~free on a dispatch-bound GPU, real compute on CPU) and then runs the polish/production stage masked too. 2. Stage-1 probes run at the final n_expl when it is already known (pinned or the deterministic device default): scan length does not change compile time, so stage 3 reuses stage 1's executable instead of recompiling the loop at a new sweep count. The opt-in wall-timed exploration search still probes at n_expl=1. Because JAX's threefry key/uniform streams are prefix-stable, the live prefix of a masked run is BIT-IDENTICAL to the unpadded run — same states, same accept/reject decisions, same discovered N/Lambda/schedule. Regression-tested (TestChainMasking) and verified exactly on the full planted-Ising benchmark suite (L=16/24/32 lps=2, L=24 lps=6 seeds 0-3, lps=10): masked == unmasked on every config, ground-state hit-rate 1.0 preserved, brute-force --verify OK. Measured (RTX 5080, 22x22 grid): cold autotune 7.4s -> 6.0s, warm (persistent cache) 2.9s -> 2.4s; glassy multi-probe search (planted L=24 lps=6 seed=1) 9.0s -> 5.9s (-34%). Full suite 370 passed / 1 skipped on GPU; pre-commit (ruff check/format 0.15.20, pyright 1.1.411) clean. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
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
1 participant
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.
XLA compile dominates cold tune_chains (~85% of stage-1 wall on GPU; the actual sampling of both probes is <1s) because each probe at a new chain count recompiles the round loop at new shapes. Two changes make the whole autotune share a single compiled loop:
Chain masking: nrpt(pad_chains_to=) pads the ladder to a fixed length with copies of the coldest chain and passes the true count into the jitted loop as TRACED data. Swaps beyond the live prefix are forced-rejected (identity permutation — DEO coupling is nearest-neighbour, so the padding is fully decoupled), attempt counters and the round-trip "top" are masked to the live prefix, and all outputs are sliced back to the true count, so optimize_schedule and every downstream consumer see unpadded semantics. tune_chains(pad_probes=) pads every probe to max_chains; autotune(pad_probes=None) auto-enables on accelerator + template mode (padding Gibbs work is ~free on a dispatch-bound GPU, real compute on CPU) and then runs the polish/production stage masked too.
Stage-1 probes run at the final n_expl when it is already known (pinned or the deterministic device default): scan length does not change compile time, so stage 3 reuses stage 1's executable instead of recompiling the loop at a new sweep count. The opt-in wall-timed exploration search still probes at n_expl=1.
Because JAX's threefry key/uniform streams are prefix-stable, the live prefix of a masked run is BIT-IDENTICAL to the unpadded run — same states, same accept/reject decisions, same discovered N/Lambda/schedule. Regression-tested (TestChainMasking) and verified exactly on the full planted-Ising benchmark suite (L=16/24/32 lps=2, L=24 lps=6 seeds 0-3, lps=10): masked == unmasked on every config, ground-state hit-rate 1.0 preserved, brute-force --verify OK.
Measured (RTX 5080, 22x22 grid): cold autotune 7.4s -> 6.0s, warm (persistent cache) 2.9s -> 2.4s; glassy multi-probe search (planted L=24 lps=6 seed=1) 9.0s -> 5.9s (-34%). Full suite 370 passed / 1 skipped on GPU; pre-commit (ruff check/format 0.15.20, pyright 1.1.411) clean.