perf(prover): default the cuda table scheduler to K = num_airs

.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)

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).

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,112 @@ 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<usize> {
+    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: `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.
+/// 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.
 ///
-/// 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 {
+/// # Why the `cuda` arm has no core term
+///
+/// Measured over 881 runs on two RTX 5090 boxes
+/// (`scripts/profiling/table-parallelism-sweep/`). All eight core-count curves
+/// fit `T(k) = S + max(Tmax, W/k)` within run-to-run noise, and the divisible
+/// work `W ≈ 5.3–8.0 s` is invariant to host core count over an 8× range, to CPU
+/// model, and to rayon pool width: cutting `RAYON_NUM_THREADS` 32 → 4 moves `W`
+/// by ~0 and merely doubles `S` (the serial continuation producer), with the
+/// best `k` still `num_airs` at every pool width. `available_parallelism()`
+/// sizes precisely that rayon pool, so it is the wrong quantity to scale `k`
+/// by — `k` is not a thread count, it is a count of concurrent drivers whose
+/// per-table work all runs on the one global pool. Worst-case cost against the
+/// best measured `k`, over four core counts on both boxes: `cores/3` +30.2 %,
+/// `cores*2/3` +13.0 %, constant 12 +7.0 %, `num_airs` +1.6 % (both of the
+/// latter's non-zero cells inside noise, p = 0.88 / 0.80).
+///
+/// Taking the ceiling rather than computing an optimum is right in both
+/// regimes of the fit: if `W/num_airs > Tmax` more `k` strictly helps, and if
+/// `W/num_airs < Tmax` the extra drivers are floor-limited and cost nothing —
+/// measured, the one staging slab is held 56 % of wall at `k = 31` and wall
+/// time still improves. What is meant to bound concurrency here is memory
+/// admission rather than a count; that is `VramGate`'s job, and it never binds
+/// at the default budget.
+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)
-                }
-            })
+        let k = parallelism_override().unwrap_or_else(|| {
+            // GPU builds: run every table. The work `k` divides is device- and
+            // workload-bound, not core-bound — see the doc comment.
+            #[cfg(feature = "cuda")]
+            {
+                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"))]
+            {
+                (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 +3147,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.

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");
+    }
+}

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,38 +222,41 @@ 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
 }
 
 /// 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<u64> = specs
         .iter()
         .map(|s| transient_per_table(*s, blowup))

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}"
+    );
+}

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);