diff --git a/.github/workflows/benchmark-pr.yml b/.github/workflows/benchmark-pr.yml index 91f5b02ac..956852588 100644 --- a/.github/workflows/benchmark-pr.yml +++ b/.github/workflows/benchmark-pr.yml @@ -273,7 +273,8 @@ jobs: # Optional table parallelism for the HEADLINE benchmark only (the memory # growth sweep always runs at default parallelism). `/bench k=N` overrides; - # otherwise default (cores/3). /bench-growth no longer forces k=1. + # otherwise the build's default (num_airs on cuda, cores/3 on CPU). + # /bench-growth no longer forces k=1. TABLE_K="" if [ "$EVENT_NAME" = "issue_comment" ]; then TABLE_K=$(echo "$COMMENT_BODY" | grep -o 'k=[0-9]*' | head -1 | cut -d= -f2) diff --git a/crypto/stark/src/instruments.rs b/crypto/stark/src/instruments.rs index 796aaf46f..0f68059f4 100644 --- a/crypto/stark/src/instruments.rs +++ b/crypto/stark/src/instruments.rs @@ -22,7 +22,7 @@ use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH}; // siblings overlap in wall time. Read them as per instance wall time. // - `scripts/profiling/phase_table.py` SUMS spans that share a label, so a // label used once per table reports the sum over all tables, which can -// exceed the enclosing phase's wall clock by up to `table_parallelism()`. +// exceed the enclosing phase's wall clock by up to the scheduler's `k`. // Give a per instance span its own label; never reuse a phase label for it. // // let _s = instruments::span("trace_build"); // RAII, stops on drop @@ -278,8 +278,9 @@ pub struct MultiProveTiming { /// root must be absorbed before the shared LogUp challenges are sampled. pub main_commits: Duration, /// Wall clock of the fused per-table region: aux build, aux commit and - /// rounds 2-4, which run as one task per table across `table_parallelism()` - /// drivers. There is no phase-level wall for the aux stages on their own + /// rounds 2-4, which run as one task per table across + /// `table_parallelism(num_airs)` drivers. There is no phase-level wall for + /// the aux stages on their own /// any more; their CPU time shows up in `round1_sub`. pub rounds_2_4: Duration, /// Sub-op breakdown for Round 1 (main + aux LDE vs Merkle). diff --git a/crypto/stark/src/prover.rs b/crypto/stark/src/prover.rs index 4047458bc..0fb561ab3 100644 --- a/crypto/stark/src/prover.rs +++ b/crypto/stark/src/prover.rs @@ -267,8 +267,9 @@ where /// aux commit and rounds 2-4 into one task: /// - main: produced by the Round 1 main commit, which is a phase-wide barrier, /// so all N tables' main LDEs are live at once (O(N × main_cols × lde_size)). -/// - aux: produced and consumed inside the same fused task, so at most -/// `table_parallelism()` of them coexist (O(k × aux_cols × lde_size)). +/// - aux: produced and consumed inside the same fused task, so at most the +/// scheduler's `k` coexist (O(k × aux_cols × lde_size)) — which under `cuda` +/// is `num_airs`, so there they are all-N-live like the main ones. /// /// Under `debug-checks` the fused task is split around the cross-table bus /// balance check, so there the aux LDEs are all-N-live like the main ones. @@ -573,41 +574,93 @@ where (d, t) } -/// Number of tables to process concurrently in `multi_prove`. +/// Explicit `TABLE_PARALLELISM` override, honoured by both `k` values below so +/// setting it pins the scheduler and the storage estimate to the same number. +#[cfg(feature = "parallel")] +fn parallelism_override() -> Option { + std::env::var("TABLE_PARALLELISM") + .ok() + .and_then(|s| s.parse().ok()) +} + +#[cfg(feature = "parallel")] +fn host_cores() -> usize { + std::thread::available_parallelism() + .map(|n| n.get()) + .unwrap_or(4) +} + +/// Number of tables `multi_prove` proves concurrently, out of `num_airs` of +/// them. +/// +/// Defaults: **every table** under `cuda`, `num_cores / 3` on CPU builds +/// (benchmarked optimal on both M3 Pro and EPYC 9454P — every table there is +/// pure host work, so `k` genuinely competes for cores). Both arms are +/// overridden by the `TABLE_PARALLELISM` env var, and the result is clamped to +/// `1..=num_airs`. Without the `parallel` feature this is 1 and the env var is +/// ignored. /// -/// Defaults: `num_cores / 3` on CPU builds (benchmarked optimal on both M3 Pro -/// and EPYC 9454P — every table there is pure host work), `num_cores * 2 / 3` -/// under `cuda`, where most in-flight tables sit in GPU waits so more of them -/// pay (swept flat at ~2/3 of the cores on a 16-core/RTX 5090 box). Both arms -/// are overridden by the `TABLE_PARALLELISM` env var. Without the `parallel` -/// feature this is hardcoded to 1 and the env var is ignored. +/// # Why the `cuda` arm has no core term /// -/// Not only the prover's `k`: `auto_storage::decide` feeds this into the -/// RAM-vs-Disk storage estimate, so the `cuda` arm also doubles that transient -/// term (see `peak_bytes`). -pub fn table_parallelism() -> usize { +/// Measured over 881 runs on two RTX 5090 boxes (sweep record linked from +/// PR #911): the work `k` divides is device- and workload-bound — invariant to +/// host core count over an 8× range — so `available_parallelism()` is the +/// wrong quantity to scale `k` by. `k` is not a thread count; it counts +/// concurrent drivers whose per-table work all runs on the one global rayon +/// pool. Worst case against the best measured `k`: `num_airs` +1.6 % (inside +/// noise), the old `cores*2/3` +13.0 %. Bounding concurrency is memory +/// admission's job (`VramGate`), not this count's. +pub fn table_parallelism(num_airs: usize) -> usize { #[cfg(feature = "parallel")] { - std::env::var("TABLE_PARALLELISM") - .ok() - .and_then(|s| s.parse().ok()) - .unwrap_or_else(|| { - let cores = std::thread::available_parallelism() - .map(|n| n.get()) - .unwrap_or(4); - // GPU builds: with the admission scheduler most in-flight - // tables sit in GPU waits, so more of them pay (swept flat at - // ~2/3 of the cores on a 16-core/RTX 5090 box). CPU builds - // stay at cores/3 — every table is pure host work there. - #[cfg(feature = "cuda")] - { - (cores * 2 / 3).max(1) - } - #[cfg(not(feature = "cuda"))] - { - (cores / 3).max(1) - } - }) + // GPU builds: run every table. The work `k` divides is device- and + // workload-bound, not core-bound — see the doc comment. + #[cfg(feature = "cuda")] + let k = parallelism_override().unwrap_or(num_airs); + // CPU builds: every table is pure host work, so `k` competes for + // the same cores the rayon pool wants. + #[cfg(not(feature = "cuda"))] + let k = parallelism_override().unwrap_or_else(|| (host_cores() / 3).max(1)); + k.clamp(1, num_airs.max(1)) + } + #[cfg(not(feature = "parallel"))] + { + let _ = num_airs; + 1 + } +} + +/// How many tables' rounds 2-4 transients the *RAM* estimate assumes are alive +/// at once (`auto_storage::peak_bytes` sums the transient bytes of the top-k +/// tables, and `decide` turns that into RAM vs Disk). +/// +/// Deliberately not `table_parallelism(num_airs)`. That is a ceiling, not a +/// bound: on a `cuda` build what actually limits how many tables are in flight +/// is `VramGate`'s byte budget, which this host-side estimate cannot see. +/// Feeding an unbounded count in here would sum *every* table's transients — +/// on many-PAGE shapes that inflates the estimate by up to +44 % (512 PAGE +/// tables at blowup 4) and would spill proofs to disk that fit in RAM. On the +/// shapes that reach this path today (~21 tables, one PAGE table) the top-k sum +/// has all but saturated, so this value and `num_airs` agree to well under 1 %. +/// +/// Kept at exactly the value it had when the scheduler shared it, so splitting +/// the two does not move any storage decision. +/// +/// TODO: derive this from a byte budget rather than a table count, so it +/// tracks what `VramGate` admits instead of standing in for it. +pub fn storage_estimate_parallelism() -> usize { + #[cfg(feature = "parallel")] + { + parallelism_override().unwrap_or_else(|| { + #[cfg(feature = "cuda")] + { + (host_cores() * 2 / 3).max(1) + } + #[cfg(not(feature = "cuda"))] + { + (host_cores() / 3).max(1) + } + }) } #[cfg(not(feature = "parallel"))] { @@ -3075,7 +3128,7 @@ pub trait IsStarkProver< twiddle_caches.push(twiddles); } - let k = table_parallelism().min(num_airs).max(1); + let k = table_parallelism(num_airs); // VRAM budgeted admission. The budget caps the summed device working set // of the tables proved concurrently so large blocks don't exhaust VRAM. diff --git a/crypto/stark/src/tests/prover_tests.rs b/crypto/stark/src/tests/prover_tests.rs index ff4a0313c..480969a84 100644 --- a/crypto/stark/src/tests/prover_tests.rs +++ b/crypto/stark/src/tests/prover_tests.rs @@ -609,3 +609,43 @@ fn commit_rows_bit_reversed_matches_commit_bit_reversed() { } } } + +/// `k` is a count of concurrent table drivers — `run_admitted` spawns exactly +/// this many OS threads and indexes `order` with them — so it has to stay +/// inside `1..=num_airs` in every arm, including under a `TABLE_PARALLELISM` +/// override (CI's prover shard 1 sets one). +#[test] +fn table_parallelism_stays_within_one_and_num_airs() { + use crate::prover::table_parallelism; + + assert_eq!(table_parallelism(0), 1, "no tables still needs one driver"); + for n in [1usize, 2, 7, 31, 64, 1024] { + let k = table_parallelism(n); + assert!(k >= 1 && k <= n, "k={k} outside 1..={n}"); + } + + // Monotone in `num_airs` in every arm: cuda `n`, CPU `min(cores/3, n)`, + // override `min(override, n)`. + let mut prev = 0; + for n in 1..=64 { + let k = table_parallelism(n); + assert!(k >= prev, "k fell from {prev} to {k} at num_airs={n}"); + prev = k; + } +} + +/// The cuda default is every table: the sweep in `thoughts/k-sweep-877b/` found +/// no core count at which a smaller `k` wins, and `T(k) = S + max(Tmax, W/k)` +/// has no term that ever favours one. Skipped when the env var pins `k`. +#[cfg(all(feature = "cuda", feature = "parallel"))] +#[test] +fn cuda_table_parallelism_defaults_to_num_airs() { + use crate::prover::table_parallelism; + + if std::env::var("TABLE_PARALLELISM").is_ok() { + return; + } + for n in [1usize, 7, 31, 1024] { + assert_eq!(table_parallelism(n), n, "cuda k must be num_airs"); + } +} diff --git a/prover/src/auto_storage.rs b/prover/src/auto_storage.rs index 6b5ed8a5d..b4718974c 100644 --- a/prover/src/auto_storage.rs +++ b/prover/src/auto_storage.rs @@ -30,7 +30,7 @@ use crate::tables::register::{ }; use crate::tables::shift::{bus_interactions as shift_buses, cols::NUM_COLUMNS as SHIFT_COLS}; use crate::tables::trace_builder::TableLengths; -use stark::prover::table_parallelism; +use stark::prover::storage_estimate_parallelism; use stark::storage_mode::StorageMode; use sysinfo::System; @@ -222,7 +222,7 @@ pub fn decide(lengths: &TableLengths, blowup_factor: u8) -> StorageMode { log::info!("storage_mode: Disk (forced via FORCE_DISK_SPILL)"); return StorageMode::Disk; } - let estimated = peak_bytes(lengths, blowup_factor, table_parallelism()); + let estimated = peak_bytes(lengths, blowup_factor, storage_estimate_parallelism()); let mode = select_storage_mode(estimated, available_ram_bytes()); log::info!("estimated_peak_bytes: {estimated}, storage_mode: {mode:?}"); mode @@ -230,30 +230,33 @@ pub fn decide(lengths: &TableLengths, blowup_factor: u8) -> StorageMode { /// Peak RAM estimate in bytes for a proof whose trace shape matches `lengths`. /// -/// `table_parallelism` is the prover's `k` (`stark::prover::table_parallelism`), -/// and it is not only a prover knob: `decide` feeds it in here, so the `cuda` -/// arm's `cores * 2 / 3` doubles the transient term below versus the CPU arm's -/// `cores / 3` and makes `Disk` more likely. That direction is safe (it -/// over-estimates), but it means a change to `k` changes the storage decision. +/// `table_parallelism` is how many tables' rounds 2-4 transients this assumes +/// are alive at once. `decide` passes `storage_estimate_parallelism()`, which +/// is deliberately *not* the scheduler's `k` — that one is `num_airs` under +/// `cuda`, and summing every table's transients here inflates the estimate on +/// many-PAGE shapes (up to +44 %) and makes `Disk` more likely than the real +/// heap warrants. See that function for why the honest bound is a byte budget +/// rather than a count. pub fn peak_bytes(lengths: &TableLengths, blowup_factor: u8, table_parallelism: usize) -> u64 { let blowup = blowup_factor as u64; let k = table_parallelism.max(1); let specs = table_specs(lengths); // Persistent: every table's main LDE + Merkle really is alive at once (the - // Round 1 main commit is a phase-wide barrier). The aux LDE no longer is — - // it is produced and consumed inside one table's fused task, so at most k - // coexist — but it is still counted for every table here, which keeps this - // an over-estimate rather than making the bound unsound. + // Round 1 main commit is a phase-wide barrier). The aux LDE is produced and + // consumed inside one table's fused task, so only the scheduler's k coexist + // — exactly all of them on `cuda`, fewer on CPU builds. Counted for every + // table either way, which is exact on `cuda` and an over-estimate on CPU + // rather than an unsound bound. let persistent_total: u64 = specs .iter() .map(|s| persistent_per_table(*s, blowup)) .fold(0u64, u64::saturating_add); - // Transient: only k tables run the fused aux+rounds task at a time. The - // top-k tables by transient bytes bound it; with the scheduler's - // heaviest-first admission that top-k is also the set actually admitted - // first, so this is the realistic peak, not a worst case. + // Transient: k tables' fused aux+rounds tasks assumed in flight at once. + // The top-k tables by transient bytes bound that; with the scheduler's + // heaviest-first admission that top-k is also the set admitted first, so + // this is the realistic peak, not a worst case. let mut transient_per: Vec = specs .iter() .map(|s| transient_per_table(*s, blowup)) diff --git a/prover/src/tests/auto_storage_tests.rs b/prover/src/tests/auto_storage_tests.rs index 5d976f81b..e26674d27 100644 --- a/prover/src/tests/auto_storage_tests.rs +++ b/prover/src/tests/auto_storage_tests.rs @@ -95,3 +95,43 @@ fn unknown_available_defaults_to_disk() { let mode = select_storage_mode(peak_bytes(&empty_lengths(), 2, ALL_TABLES), None); assert_eq!(mode, StorageMode::Disk); } + +/// A shape with one PAGE table — everything the monolithic path proves today. +/// The top-k sum has saturated well before the table count, so the estimate is +/// insensitive to `k` in that range: this is why raising the *scheduler's* `k` +/// to `num_airs` does not move the storage decision on a normal workload. +#[test] +fn peak_bytes_is_k_saturated_on_single_page_shapes() { + let mut lengths = empty_lengths(); + lengths.cpu_padded_rows = 1 << 20; + lengths.memw_padded_rows = 1 << 20; + lengths.decode_rows = 1 << 16; + lengths.unique_page_count = 1; + + let bounded = peak_bytes(&lengths, 2, 12); + let unbounded = peak_bytes(&lengths, 2, ALL_TABLES); + assert!( + unbounded * 100 <= bounded * 101, + "estimate moved {bounded} -> {unbounded} on a one-page shape" + ); +} + +/// …and why `decide` must not simply be handed the scheduler's `k`. PAGE tables +/// are all the same size, so once there are many of them the top-k truncation +/// is doing real work: summing every table's transients inflates the estimate +/// by >20 % here, which spills proofs to disk that fit in RAM. +#[test] +fn unbounded_k_inflates_peak_bytes_on_many_page_shapes() { + let mut lengths = empty_lengths(); + lengths.cpu_padded_rows = 1 << 20; + lengths.memw_padded_rows = 1 << 20; + lengths.decode_rows = 1 << 16; + lengths.unique_page_count = 128; + + let bounded = peak_bytes(&lengths, 2, 21); + let unbounded = peak_bytes(&lengths, 2, ALL_TABLES); + assert!( + unbounded * 10 > bounded * 12, + "expected >20 % inflation, got {bounded} -> {unbounded}" + ); +} diff --git a/prover/tests/calibration.rs b/prover/tests/calibration.rs index ff11bcf4b..c7d4d66f5 100644 --- a/prover/tests/calibration.rs +++ b/prover/tests/calibration.rs @@ -11,7 +11,7 @@ use lambda_vm_prover::tables::MaxRowsConfig; use lambda_vm_prover::tables::trace_builder::count_table_lengths; use lambda_vm_prover::test_utils::{asm_elf_bytes, run_asm_elf}; use stark::proof::options::GoldilocksCubicProofOptions; -use stark::prover::table_parallelism; +use stark::prover::storage_estimate_parallelism; use std::sync::Arc; use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; use std::thread; @@ -36,7 +36,8 @@ fn peak_bytes_does_not_underestimate_measured_heap() { count_table_lengths(&elf, &logs, &max_rows, &[]).expect("count_table_lengths succeeds"); let opts = GoldilocksCubicProofOptions::with_blowup(2).expect("blowup=2 is valid"); - let predicted = peak_bytes(&lengths, opts.blowup_factor, table_parallelism()) as usize; + let predicted = + peak_bytes(&lengths, opts.blowup_factor, storage_estimate_parallelism()) as usize; drop(logs);