perf(prover): per-table scheduler with VRAM admission for multi_prove - #877
Merged
Conversation
Collaborator
Author
|
/bench-gpu |
GPU Benchmark (ABBA) —
|
Grid-stride the fused row-major NTT past gridDim.y (lde >= 2^24 silently fell back to CPU), assert the device-only contract in the R2 composition commit and preprocessed opening fallbacks, validate htod_via bounds, retain FRI device evals only under device-only, and move the inverse fault-injection hook so every batch-inverse entry is covered.
ColoCarletti
marked this pull request as draft
July 31, 2026 17:48
…ault hook, merkle root-only
…vm into gpu-opt-round2-impl
… zero-total guard - gather_ext3_at asserts positions against the evals buffer host-side (same guard as gather_merkle_paths_dev). - The device-gather cross-checks keep query 0 as a release canary instead of paying every query; debug still checks all of them. - The batch-inverse zero-total guard also compiles under test-faults, so the GPU fallback suite (which runs --release) actually exercises it. - New htod_via round-trip test covering the 64 MB chunk loop and its partial tail.
…892) * fix(gpu): drain htod_via on error, guard the R2 host-evaluator fallback Review follow-ups for the round-2 residency work, rebased onto e75bcbe — only the items that commit did not already cover. htod_via error path. Once a chunk's DMA is in flight, `record_event` / `sync_event` returning `Err` drops the staging `MutexGuard` with the device still reading the pinned slab, so the next locker's `ensure_capacity` can `cuMemFreeHost` it mid-copy. `async_dtoh_via` already guards this exact hazard and the file ships a `DrainOnErr` helper for it; `htod_via` was the one site not using it. R2 host-evaluator fallback. If the device decompose and the `H` download both fail under device-only, control reaches the host evaluator, which reads the intentionally-empty trace and panics with a bare out-of-bounds. Assert the device-only contract instead, matching the other fallback arms. Coverage. `batch_inverse_ext3_dev`'s `n == 1` branch is never exercised — `batch_inverse_n1` goes through the host-only short circuit in `batch_inverse_ext3`, as its own comment says. Add a direct device test. Docs. The preprocessed split-tree comment still claimed both trees come back as full host trees (the multiplicity tree is root-only + device resident), and `FriCommitState`'s doc claimed its input is always Arc-shared with a retained `gpu_evals` (only true on the device-only path). * perf(gpu): set the merkle-tail threshold to the block width TAIL_MAX_PAIRS = 2048 overshoots. The tail grid-strides a single 128-thread block on one SM, so a level of k pairs is k/128 SEQUENTIAL keccak-f1600s where the per-level launches it replaces spread them over k/128 parallel blocks. At 2048 the first four levels alone are 16+8+4+2 = 30 serial permutations against 4 parallel waves — order +100 us per large tree to save 4 launches worth order 10 us, and it sits on the critical path because the caller's 32-byte root memcpy_dtoh host-blocks on everything queued before it. At the block width the entry level is exactly one permutation per thread, so the tail still collapses the top levels into one launch but adds no serialization at all.
…eduler Fiat-Shamir only requires the main roots absorbed in index order before the shared challenges; past that fork every table's chain is independent. Phase A now runs all main commits under a byte-budget admission gate (no chunk barriers), and aux build, aux commit and rounds 2-4 run fused as one task per table, heaviest first — while a big table works through a host-bound stretch, the other tables' GPU stages fill the device. GPU builds default TABLE_PARALLELISM to 2/3 of the cores (swept flat at 10 on a 16-core RTX 5090). ethrex 10tx continuations on RTX 5090: 10.64s -> 8.54s (-19.7%, 8 ABBA pairs).
ColoCarletti
force-pushed
the
gpu-opt-table-scheduler
branch
from
August 3, 2026 19:35
adbe561 to
9d2140b
Compare
ColoCarletti
marked this pull request as ready for review
August 3, 2026 19:35
Collaborator
Author
|
/bench-gpu |
Collaborator
Author
|
/ai-review |
Codex Code Review
|
Contributor
|
/bench-gpu 8 |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
Contributor
Review: per-table scheduler with VRAM admissionThe scheduling argument holds up: Fiat-Shamir only needs the main-commit roots absorbed in index order before the shared challenges, and past the per-table fork every chain is independent. I checked the ordering-sensitive spots and they're preserved — Five findings, none of them proof-correctness: High — Medium — Low — the deleted Details inline. |
AI ReviewPR #877 · 1 changed files Findings
Status column reflects the verdict from the verifier: deepseek-verifier (openrouter/deepseek/deepseek-v4-pro). AI-003: Residual device main-trace snapshots not accounted by the VRAM admission gate
Claim The fused per-table chain defers Evidence trace.rs:188-214 shows Suggested fix Either (a) bulk-clear Reviewer Lanes
Verification Lanes
Native Codex and Claude reviews run separately and post their own comments. They are not included in this structured provenance report. Discarded candidates (5) — rejected by the verifier
Raw lane outputs, candidates, final issues, and model metrics are uploaded as workflow artifacts. |
… the per-table scheduler (#893) * fix(instruments): nest per-table spans under their real parent The per-table scheduler moved `r1_aux_build`, `r1_aux_commit` and `rounds_2to4` inside closures that run on `std::thread::scope` worker threads. `SPAN_DEPTH` is thread-local and a fresh OS thread starts at 0, so all three were stamped `depth = 0` and recorded as root siblings of `prove_total` instead of children of `proving`. That happens even at k = 1. Downstream, `scripts/profiling/phase_table.py` reconstructs the tree with `del stack[d:]`, so a depth-0 span empties the ancestor stack and `prove_total` stops being an ancestor of anything — the "% of total" column documented in `scripts/profiling/README.md` becomes meaningless. `run_admitted` now reads the spawning thread's depth and seeds each driver with it via new `instruments::current_depth` / `enter_depth`. Both call sites are `#[cfg(feature = "instruments")]`, so non-instrumented builds are byte-identical. Also correct the module contract doc: per-table spans genuinely do overlap now — that is inherent to running one driver per in-flight table, not a bug to code around. Only the top-level phase spans remain a strict latency breakdown. * fix(prover): report aux build/commit where they actually accrue `aux_build_elapsed` / `aux_commit_elapsed` were hardcoded to `Duration::ZERO`, but `prover/src/instruments.rs` still computed `round1 = main_commits + aux_build + aux_commit` and still printed the "Aux trace build" / "Aux trace commit" rows. Since `accum_r1_aux` keeps firing, the report showed nonzero LogUp and Aux-LDE/Merkle children under zero parents, and all the aux time silently landed in "Rounds 2-4". Time both stages inside the fused chain and sum them across drivers (`instruments::accum_aux_phases` / `take_aux_phases`), then restructure the report to match what the scheduler actually does: - "Round 1 (main trace commits)" is now exactly the main commits — the last phase-wide barrier, since the main roots must all be in the transcript before the shared LogUp challenges are sampled. - Aux build, aux commit and rounds 2-4 sit under one wall-clock parent, "Rounds 2-4 (aux build+commit fused)", with the aux rows marked as summed across concurrent drivers — they may exceed that wall, the same convention the existing accum_* sub-rows already use. No zero parents over nonzero children remain. Verified on fib_iterative_1M: Round 1 1.68s, Rounds 2-4 6.89s wall, aux build 2.91s and aux commit 3.38s summed over 5 drivers. * fix(bench): drop the heap guards whose snapshots no longer exist The scheduler removed `instruments::snap("After aux build")` and `snap("After aux commit")`. `bench_prover_scaling.sh` still parsed them, printed them and ran heap-growth regressions on them, so two regression guards were comparing nothing and dropping out without complaint. Re-adding a snapshot inside a per-table task would be meaningless — with k tables in flight there is no single moment at which aux build or aux commit has finished — so remove the two rows and their `regress` calls, with a NOTE recording why and pointing at the guards that still cover the fused region ("After main commits" and "Peak heap"). Also repoint the timing regexes at the labels the report actually prints. `Main expand_columns_to_lde` / `Aux expand_columns_to_lde` and `Main commit (Merkle)` / `Aux commit (Merkle)` had not matched since the labels gained their GPU/CPU suffixes, so t_main_lde, t_aux_lde, t_main_merkle and t_aux_merkle silently printed blank. All four populate again — checked by running the script's own awk over a real report. * docs(prover): refresh the comments the per-table scheduler invalidated Nothing functional. All of these described structure the scheduler removed: - `VramGate`'s rustdoc opened with the deleted `plan_table_chunks`'s doc comment ("Plan contiguous table chunks... Returns (start, end) half open ranges"), left behind and contiguous with `VramGate`'s own. - `Lde`'s doc claimed all N tables' LDE columns are live simultaneously. Only the main LDEs still are — the Round 1 main commit is a phase-wide barrier. Each aux LDE is produced and consumed inside one fused task, so at most k coexist. That is a memory improvement the PR made and did not claim; state the real, asymmetric bound. - A "Split into two passes for parallelism: Pass 1 ... Pass 2 ..." block sat two lines above the new comment saying the opposite. - `table_parallelism`'s doc still gave only `num_cores / 3`. Document both arms, that `TABLE_PARALLELISM` overrides both, and that without the `parallel` feature it is hardcoded to 1. - `run_debug_checks` said "called once after Phase C commits"; it now runs between two `run_admitted` passes. Document that, and the "each driver locks only its own index" contract its new `&[Mutex<AirTracePair>]` parameter relies on. - `auto_storage::peak_bytes` described phase D and a "worst possible chunk assignment". With `heaviest_first` the top-k is the set actually admitted first, not a worst case. Also document that `table_parallelism()` is not only the prover's k: `decide` feeds it into the RAM-vs-Disk choice, so the cuda arm's `cores * 2 / 3` doubles that transient term and makes `Disk` likelier. The direction is safe (it over-estimates) but was undocumented. - Remaining "Phase A/B/D" references, plus the "chunks of K" banner and the "Phase D's zip chain" handle comments. * test(prover): cover VramGate, run_admitted and heaviest_first These three had zero direct tests, and PR CI never exercises them concurrently: `ubuntu-latest` has 2-4 vCPU so `cores / 3` floors to k = 1, and `VramGate` is inert on non-cuda builds because `vram_budget = u64::MAX` makes `acquire`'s admit condition always true, so the condvar is never waited on. They are free functions over `&[u64]` with no field, AIR or GPU dependency, so a plain `#[cfg(test)] mod` pins them without a device: - `heaviest_first` returns a permutation of `0..n`, descending by estimate, with ties broken by index (stable sort — so the admission order does not vary run to run). - `run_admitted` fills every slot exactly once, including `order.len() == 0`, `workers > order.len()`, `workers == 1` and `workers == 0`. - `VramGate` admits an over-budget request alone rather than deadlocking, never lets concurrent admissions push `used` past the budget, and wakes waiters on permit drop. - A `u64::MAX` budget never blocks, including when the byte sum saturates. - `run_admitted` seeds its drivers' span depth, which guards the regression fixed earlier in this branch. Reads the depth directly rather than the global span timeline, which other tests in this binary also write to. Deterministic and fast: no sleeps as synchronization: channel rendezvous for ordering, and `recv_timeout` only as a failure deadline so a regression fails instead of hanging.
Per-driver slots were measured: repeated pinned allocation costs more than the shared mutex, whose transfers cross-table overlap already hides.
…timing report honest (#895) * fix(instruments): restore the rounds 2-4 phase span instead of plumbing depth Supersedes the approach in #893. Adversarial review showed the depth field was never the defect. `phase_table.py:121` takes its denominator from `max(s["wall_ns"] for _, s in pathed)` — the longest span, not the root of the ancestor stack — so depth-0 records never broke the "% of total" column, and `scripts/profiling/README.md:77` was accurate all along. `prover/src/continuation.rs` has also recorded spans from worker threads since long before this branch (:1146, :1205, :1299, :1328, :1415), with the comment at :1051-1053 saying so. Seeding worker depth was therefore work that bought nothing, and it would have left overlapping siblings looking like a clean tree — a subtler lie. Removed (`instruments::current_depth` / `enter_depth` / `DepthGuard` and the seeding in `run_admitted`). The real defect is label collision under summing. `phase_table.py:129` does `e["wall_ns"] += s["wall_ns"]`, so spans sharing a label are summed. On origin/main `rounds_2to4` was ONE span around the chunk loop (prover.rs:3503) and measured the phase; this branch made it one span per table, so the row became the sum of N concurrent tables — up to k times the real wall, able to exceed 100% — and no span measured the phase at all. `r1_aux_build` and `r1_aux_commit` were phase spans on main too (:3143, :3225). So: reopen `rounds_2to4` on the calling thread around the whole fused region, and rename the per-table spans `*_table` so a per-instance label can never be summed into a phase row. This also repairs `LAMBDA_VM_NSYS_CAPTURE_SPAN=rounds_2to4` (README.md:115), which with the label on the per-table span had N driver threads calling cuProfilerStart/Stop, the first to finish ending the capture. The report follows, and is compile-coupled to the same change. #893 added per-driver aux timers to fill the zeroed `aux_build` / `aux_commit` buckets; the fused stages have no wall-clock phase of their own any more, so reporting one invites exactly the misreading the label summing caused. Both timers and both `MultiProveTiming` fields are gone. The report now shows only the two phases that remain — "Round 1 (main trace commits)" and "Rounds 2-4 (aux build+commit fused in)" — with the aux CPU-time rows grouped under the fused phase behind headers stating they are summed over tables. That still fixes what #893 set out to fix: no row prints a fabricated 0.00s over live children, and "Round 1" no longer duplicates its own child. Verified on fib_iterative_1M: phase spans sum to their parent (r1_prepass 0.148 + r1_main_commit 2.493 + rounds_2to4 8.943 = 11.584 vs proving 11.585). * revert: trim the bench script back to the minimum #893 also repointed four timing regexes in `scripts/bench_prover_scaling.sh` that had gone stale earlier and independently of this branch. That is unrelated churn in a script with no Makefile target and no workflow referencing it, so it is reverted. What stays removed: the two dead heap rows and their `regress` calls (their `snap()` sources no longer exist and cannot be recreated with k tables in flight) and the two aux timing rows, which follow the report. The NOTE explaining why is kept. Nothing here was failing silently, contrary to the original review note: `regress` prints "(insufficient data)" for a missing key and `print_row` prints "-". * test: drop the scheduler unit tests * ci: force k > 1 on one prover shard so the scheduler runs concurrently Replaces the `VramGate` / `run_admitted` / `heaviest_first` unit tests added in #893 (removed in the previous commit). Every assertion they made was guaranteed by construction, already covered end to end, or unreachable from the call sites: `heaviest_first` is `(0..n).collect()` plus `sort_by_key`; a slot mixup in `run_admitted` is schedule independent, so it trips one of the three `.expect("run_admitted fills every slot")` sites or fails `multi_verify` on every PR today; `order.len() == 0` and `workers > order.len()` cannot happen, since `k` is `.max(1)`'d and `order` is always a full permutation. The one property with teeth — an over-budget table admitted alone — HANGS rather than fails if it regresses, which on an 8-10 minute shard burns to the job timeout unless wrapped in a watchdog. That was ~50 lines of permanent maintenance against approximately zero risk. The actual PR-time gap is that the scheduler never runs concurrently. `table_parallelism()` defaults to `(cores / 3).max(1)` and every job in this workflow is `runs-on: ubuntu-latest` with no larger-runner label, so PR CI proves with exactly one driver thread; `VramGate` is additionally inert on non-cuda builds, where `vram_budget` is `u64::MAX` and `acquire`'s condition always holds. `TABLE_PARALLELISM: 6` on shard 1 only is the smallest change that puts several real table closures in flight at once. The other three shards keep default-k coverage — the expression yields an empty string there, which fails to parse and falls back to the default. `prover/Cargo.toml:8` is `default = ["parallel"]`, so the env arm is the live one. Not a substitute for GPU coverage: `gpu-tests.yml` on merge_group rents a >=16-core RTX 5090, taking the cuda arm (`cores * 2 / 3`) with a finite VRAM budget, so both the concurrent and blocking paths already run before merge. This closes the PR-time gap only. * fix(bench): drop the row killed by the Round 1 relabel "Round 1 (main trace commits)" is lowercase, so `/Main trace commits/` stopped matching, and the row would have printed "-". It is also now redundant: with the aux stages fused out of round 1, `t_main_commits` and `t_round1` are the same number by construction.
MauroToscano
approved these changes
Aug 3, 2026
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
2 participants
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.
Descripción:
Third GPU round on top of #875: ethrex 10tx continuations on RTX 5090 go from 10.64s to 8.54s (-19.7%, 8 ABBA pairs, sd 2.1%). Proofs are unchanged and cross-verify against the previous prover.
Tables inside a prove used to advance in fixed chunks of 5 with a barrier between chunks and between phases (main commit / aux build / aux commit / rounds 2-4): the GPU sat idle whenever the running tables hit a host-bound stretch, because the tables that could have used it were parked at a barrier. Fiat-Shamir only requires the main-commit roots absorbed in index order before the shared challenges; past the per-table transcript fork every table's chain is independent.