fix: use CPU-aware prover scheduler - #897
Conversation
|
/bench |
Benchmark — real block (
|
| Metric | main | PR | Δ |
|---|---|---|---|
| Peak heap | 47344 MB | 52784 MB | +5440 MB (+11.5%) 🔴 |
| Prove time | 136.925s | 121.870s | -15.055s (-11.0%) 🟢 |
🎉 Improvement on the real block — prove time down 11.0%.
Prove-time spread 2.6% (121.870s / 123.141s / 119.938s)
Memory Growth
ethrex distinct-account transfers · default parallelism · 1 sample per point
| Transfers | main (MB) | PR (MB) | Δ |
|---|---|---|---|
| 4 | 12242 | 14631 | +2389 MB (+19.5%) |
| 8 | 14928 | 18288 | +3360 MB (+22.5%) |
| 12 | 19774 | 21514 | +1740 MB (+8.8%) |
| 16 | 21972 | 25555 | +3583 MB (+16.3%) |
| 20 | 23838 | 29814 | +5976 MB (+25.1%) |
Growth rate: 941 MB / transfer (main: 756, Δ: +24.5%)
Fit: R² = 0.9973 (main: 0.9717)
⚠️ Memory scaling regression — growth rate increased by +24.5%
Commit: 5bba8d4 · Baseline: cached · Runner: self-hosted bench
|
/bench |
|
/bench-growth |
|
/ai-review |
Codex Code ReviewNo actionable issues found in the PR diff. |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
| let results: Vec<std::sync::Mutex<Option<T>>> = (0..order.len()) | ||
| .map(|_| std::sync::Mutex::new(None)) | ||
| .collect(); | ||
|
|
||
| for chunk in order.chunks(workers.max(1)) { | ||
| let chunk_results: Vec<(usize, T)> = | ||
| chunk.par_iter().map(|&idx| (idx, task(idx))).collect(); | ||
| for (idx, result) in chunk_results { | ||
| *results[idx].lock().unwrap() = Some(result); | ||
| } | ||
| } | ||
|
|
||
| results | ||
| .into_iter() | ||
| .map(|m| m.into_inner().unwrap()) | ||
| .collect() |
There was a problem hiding this comment.
Low / simplicity — the Mutex wrapper is dead weight in this arm. Unlike the CUDA sibling (where driver threads write concurrently), every write happens in the sequential for (idx, result) loop on the calling thread after collect() has already joined the chunk. A plain Vec<Option<T>> works — exactly as the not(parallel) fallback right below already does:
| let results: Vec<std::sync::Mutex<Option<T>>> = (0..order.len()) | |
| .map(|_| std::sync::Mutex::new(None)) | |
| .collect(); | |
| for chunk in order.chunks(workers.max(1)) { | |
| let chunk_results: Vec<(usize, T)> = | |
| chunk.par_iter().map(|&idx| (idx, task(idx))).collect(); | |
| for (idx, result) in chunk_results { | |
| *results[idx].lock().unwrap() = Some(result); | |
| } | |
| } | |
| results | |
| .into_iter() | |
| .map(|m| m.into_inner().unwrap()) | |
| .collect() | |
| let mut results: Vec<Option<T>> = (0..order.len()).map(|_| None).collect(); | |
| for chunk in order.chunks(workers.max(1)) { | |
| let chunk_results: Vec<(usize, T)> = | |
| chunk.par_iter().map(|&idx| (idx, task(idx))).collect(); | |
| for (idx, result) in chunk_results { | |
| results[idx] = Some(result); | |
| } | |
| } | |
| results |
Separately: results is sized by order.len() but indexed by idx taken from order. That is only sound because every call site passes a full permutation from heaviest_first(estimates). The CUDA version sizes by estimates.len(), which is the invariant that actually holds — worth matching in both new arms.
| .map(|_| std::sync::Mutex::new(None)) | ||
| .collect(); | ||
|
|
||
| for chunk in order.chunks(workers.max(1)) { |
There was a problem hiding this comment.
Medium (performance) — order.chunks(workers) reintroduces exactly the hard barrier the CUDA sibling's doc comment (three lines above line 754) says it was written to remove: "the fixed chunks this replaces made every table wait for the slowest of its chunk". And because order is heaviest-first, chunk 0 is the k heaviest tables including the single long pole — so the other k-1 lanes go idle waiting on it before chunk 1 can even start, and there are ⌈N/k⌉ such stalls (~3 at N=21, k=8).
Keeping the work inside the Rayon pool (the stated goal of this arm) does not require chunking. The same atomic-cursor loop as the CUDA version, driven by rayon::scope instead of std::thread::scope, gives you pool-resident drivers, a k concurrency cap, and no barrier:
let cursor = std::sync::atomic::AtomicUsize::new(0);
rayon::scope(|scope| {
for _ in 0..workers.max(1).min(order.len().max(1)) {
scope.spawn(|_| loop {
let pos = cursor.fetch_add(1, Ordering::Relaxed);
if pos >= order.len() { return; }
let idx = order[pos];
*results[idx].lock().unwrap() = Some(task(idx));
});
}
});Nested Rayon work from a pool thread participates in work-stealing rather than blocking it, so the oversubscription this PR is fixing does not come back. If the chunked form was chosen because it measured faster than this, please say so in the comment — otherwise the barrier looks like collateral damage rather than intent.
| #[cfg(all(not(feature = "cuda"), not(feature = "debug-checks")))] | ||
| let table_results = { | ||
| let aux_outs = run_admitted(&peak_order, &peak_estimates, &vram_gate, k, aux_stage); | ||
| let mut staged = Vec::with_capacity(num_airs); |
There was a problem hiding this comment.
Medium (memory / stale invariant) — staged now holds all N tables' Lde values alive at once on CPU builds, where the fused path held at most k. Lde::aux is ext3 (24 B/element × lde_size), so the added peak is roughly (N − k) × aux_cols × 24 × lde_size — with k = host_cores / 3 that is the large majority of the aux LDE working set.
The RAM estimate still holds, so no storage misprediction: auto_storage::persistent_per_table unconditionally adds aux_lde for every table, so peak_bytes was already over-estimating CPU and now becomes exact. But two load-bearing comments become wrong and should be updated in this PR, since they are precisely what a future reader would consult before touching this:
prover.rsstruct Ldedoc: "aux: produced and consumed inside the same fused task, so at most the scheduler'skcoexist … which undercudaisnum_airs, so there they are all-N-live like the main ones." This is now inverted —cudais the fused/k-bounded case and CPU is all-N-live.auto_storage::peak_bytes: "The aux LDE is produced and consumed inside one table's fused task, so only the scheduler's k coexist — exactly all of them oncuda, fewer on CPU builds. Counted for every table either way, which is exact oncudaand an over-estimate on CPU."
Worth a line in the PR description too — the CPU throughput win is being paid for in peak host RAM.
| #[cfg(all(not(feature = "cuda"), not(feature = "debug-checks")))] | ||
| let table_results = { | ||
| let aux_outs = run_admitted(&peak_order, &peak_estimates, &vram_gate, k, aux_stage); | ||
| let mut staged = Vec::with_capacity(num_airs); | ||
| for out in aux_outs { | ||
| staged.push(std::sync::Mutex::new(Some( | ||
| out.expect("run_admitted fills every slot")?, | ||
| ))); | ||
| } | ||
| run_admitted(&peak_order, &peak_estimates, &vram_gate, k, |idx| { | ||
| let (commitment, lde) = staged[idx] | ||
| .lock() | ||
| .unwrap() | ||
| .take() | ||
| .expect("aux result consumed once per table"); | ||
| rounds_stage(idx, commitment, lde) | ||
| }) | ||
| }; |
There was a problem hiding this comment.
Low / simplicity — this new block is a near-copy of the debug-checks block immediately below it (two admitted passes with a Vec<Mutex<Option<(Round1Commitments, Lde)>>> staging vector in between). The only difference is the run_debug_checks call. Collapsing them removes ~20 duplicated lines and one of the three table_results definitions:
#[cfg(all(feature = "cuda", not(feature = "debug-checks")))]
let table_results = /* fused, unchanged */;
// Split into two admitted passes: CPU wants a phase barrier between the
// aux/commit region and rounds 2-4 (interleaving these host-heavy regions
// regresses through cache and memory-bandwidth contention), and
// debug-checks needs every table's commitments between the two stages for
// the cross-table bus balance check.
#[cfg(any(not(feature = "cuda"), feature = "debug-checks"))]
let table_results = {
let aux_outs = run_admitted(&peak_order, &peak_estimates, &vram_gate, k, aux_stage);
// ... build commitments/ldes ...
#[cfg(feature = "debug-checks")]
Self::run_debug_checks(&pair_cells, &commitments, &domains, &twiddle_caches);
// ... stage + second run_admitted ...
};Not blocking, but three cfg-gated definitions of the same binding is a lot of surface for one behavioural difference.
| /// start order (heaviest table first, so the long pole starts early and small | ||
| /// tables fill around it — the fixed chunks this replaces made every table | ||
| /// wait for the slowest of its chunk). Returns one slot per original index. | ||
| #[cfg(feature = "cuda")] |
There was a problem hiding this comment.
Low (dead code) — with run_admitted now cuda-only, the whole admission mechanism (VramGate::acquire, VramPermit, its Drop) is unreachable on CPU builds; the three #[cfg_attr(not(feature = "cuda"), allow(dead_code))] attributes exist only to silence that. The _gate: &VramGate parameter on both new arms is threaded through purely as a no-op, and VramGate::new is still called unconditionally at the top of multi_prove.
Since the gate is now a device concept only, #[cfg(feature = "cuda")] on the struct/impl plus dropping the parameter from the CPU arms would delete the attributes and the vestigial argument outright. That needs the two call sites cfg-gated on the arg, so it may not be worth it — but if you keep it as-is, please note in the VramGate doc that it is a CUDA-only mechanism, otherwise the allow(dead_code) attributes read as unexplained.
Also, the + Send added to this signature isn't required: the spawned workers capture &task, so Sync alone suffices. Harmless if it's for signature symmetry across the three arms, just flagging it isn't load-bearing.
Review: CPU-aware prover schedulerRead the diff against No Critical or High findings. Everything below is Medium or lower. Medium
Low
Nothing here blocks merge; (1) and (2) are the two I would want addressed, at minimum the stale comments in (2). |
AI ReviewPR #897 · 1 changed files Findings
Status column reflects the verdict from the verifier: deepseek-verifier (openrouter/deepseek/deepseek-v4-pro). AI-005: `impl Drop for VramPermit` missing the `allow(dead_code)` its siblings got on non-cuda
Claim On non-cuda builds Evidence Lines 705/712/718 carry Suggested fix Add AI-009: heaviest_first uses unstable sort — non-deterministic scheduling for equal estimates
Claim
Evidence Lines 839-844: Suggested fix Use 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 (6) — rejected by the verifier
Raw lane outputs, candidates, final issues, and model metrics are uploaded as workflow artifacts. |
|
/bench |
|
/bench-growth |
3 participants
This PR fixes the CPU proving regression introduced by the table scheduler.
The GPU scheduler uses external driver threads to admit and overlap table work. On CPU, that caused table tasks to launch nested Rayon work from outside the Rayon pool, leading to oversubscription, cache contention, and memory-bandwidth contention.
The CPU path now:
The CUDA path and proof format are unchanged.