From ab26970af70f8e394a5eae9844b5821ffb99abb2 Mon Sep 17 00:00:00 2001 From: MauroFab Date: Mon, 7 Sep 2026 18:28:28 -0300 Subject: [PATCH 1/4] feat(prover): one device-set model for the VRAM throttle, table names on device aborts, no resident-aux host downgrade Prover side of the width-aware admission (#961). crypto/stark/src/device_set.rs (new, cfg-free) - The device-set arithmetic moves out of `gpu_lde` (cuda-only) so the scheduler's throttle, which runs on every build, reads the model the dispatch layer admits against: `commit_device_set` (one LDE + snapshot + tree + scratch), the new `table_device_set` for rounds 2-4 (aux LDE and resident aux trace, H and the parts, the R3/R4 inverted denominators, DEEP and the FRI chain, each with its tree), and the pure `admit_bytes`. `gpu_lde` re-exports them. Tests pin LFM_HASH under RPO at 2^21 x 449 (21.4 GiB R1, 23.1 GiB whole table: fits alone, where the retired model said 28.3 GiB) and a 2^22 BALU chunk at 4.9 GiB (five concurrent). crypto/stark/src/prover.rs - `estimate_table_vram_bytes` (two LDE buffers + 256 B per LDE row, stale since #956) is replaced: R1 admits on the commit set, the fused rounds on the table set, both read off the AIR and the domain. - `commit_main_trace` and the aux commit sites pass `air.name()` into the R1 dispatch entry points, so a device abort names its table at the site; `run_admitted` prefixes any string panic payload from a table task with `table : ` before re-raising it. - The resident-aux host downgrade is trimmed: after the drain-and-retry declines, the aux commit returns `ProvingError::DevicePath` with the table, the shape and the live device posture. Host RAM is a cache, not a compute path. `materialize_aux_trace_host` is deleted; `GPU_RESIDENT_AUX_DOWNGRADES` is retired at zero (accessor kept for the integration assertion). prover/src/lfm/chunking.rs - The BALU sizing doc and its test read `stark::device_set::table_device_set` instead of restating the arithmetic: a 2^27 table is 150 GiB, a 2^22 chunk 4.9 GiB, a 2^24 chunk 19 GiB; LFM_LANES at 2^24 is 23 GiB. --- crypto/stark/src/device_set.rs | 400 +++++++++++++++++++++++++++++++++ crypto/stark/src/gpu_lde.rs | 379 ++++++------------------------- crypto/stark/src/lib.rs | 1 + crypto/stark/src/prover.rs | 244 ++++++++++---------- prover/src/lfm/chunking.rs | 84 +++---- 5 files changed, 645 insertions(+), 463 deletions(-) create mode 100644 crypto/stark/src/device_set.rs diff --git a/crypto/stark/src/device_set.rs b/crypto/stark/src/device_set.rs new file mode 100644 index 000000000..0f39591f9 --- /dev/null +++ b/crypto/stark/src/device_set.rs @@ -0,0 +1,400 @@ +//! The device working set of one table, term by term — and the pure admission +//! predicate over it. +//! +//! Two consumers read this model and must agree: the per-table scheduler's VRAM +//! throttle (`prover::VramGate`, which bounds the SUM of the tables proved +//! concurrently and runs on every build, GPU or not) and the GPU dispatch +//! layer's admission (`gpu_lde::admit`, which asks whether ONE table fits the +//! card at all). Keeping the arithmetic here, free of any `cuda` gate, is what +//! lets both read one model instead of each carrying a copy. +//! +//! Every term mirrors an allocation in `math_cuda` after the in-place LDE +//! transpose of #956 (one LDE buffer, never two); `one_lde_buffer::vram_arm` +//! measures the commit's three big terms and the doc on +//! `DEFAULT_MEMPOOL_RELEASE_THRESHOLD_BYTES` records the numbers. The model +//! carries no blanket safety factor: the card's admission budget is 80% of +//! device memory, and the 20% it leaves is what the context, the module code +//! and the retained pool live in. + +/// Bytes per Goldilocks element on device. +pub const BASE_BYTES: u64 = 8; + +/// Bytes per ext3 element on device — three adjacent base columns. +pub const EXT3_BYTES: u64 = 3 * BASE_BYTES; + +/// Bytes of one Merkle node. Every commitment hash the device dispatches on +/// emits a 32-byte digest — a four-felt Goldilocks digest is exactly 32 +/// canonical bytes — so the node buffer costs the same under every hash. +pub const MERKLE_NODE_BYTES: u64 = 32; + +/// Cap on the in-place transpose's device scratch, mirrored from +/// `math_cuda::lde::INPLACE_TRANSPOSE_SCRATCH_BYTES` (private there). The +/// admission wants a bound, not the block geometry. +pub const INPLACE_TRANSPOSE_SCRATCH_CAP_BYTES: u64 = 256 << 20; + +/// `(2 · leaves − 1) · 32` for the row-pair tree over `lde_size` rows. +pub const fn full_tree_bytes(lde_size: u64) -> u64 { + lde_size.saturating_sub(1).saturating_mul(MERKLE_NODE_BYTES) +} + +/// Bytes of `cols` ext3 columns over `rows` rows. +pub const fn ext3_bytes(rows: u64, cols: u64) -> u64 { + rows.saturating_mul(cols).saturating_mul(EXT3_BYTES) +} + +/// Bytes of `cols` base columns over `rows` rows. +pub const fn base_bytes(rows: u64, cols: u64) -> u64 { + rows.saturating_mul(cols).saturating_mul(BASE_BYTES) +} + +/// The device working set one fused row-major commit allocates, term by term +/// (`math_cuda::lde::coset_lde_row_major_inner`): ONE LDE buffer, the optional +/// trace-domain snapshot, the full Merkle node buffer, and the small scratch +/// (coset weights plus the capped transpose scratch). +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct CommitDeviceSet { + /// `lde_size · base_cols · 8`: the row-major LDE, transposed in place. + pub lde_bytes: u64, + /// `n · base_cols · 8`: the pre-NTT column-major snapshot the LogUp + /// fingerprint kernel reads in place (main commits only). + pub snapshot_bytes: u64, + /// `(2 · leaves − 1) · 32` with `leaves = lde_size / 2`: one full row-pair + /// tree. The preprocessed split path builds two, sequentially on one + /// stream — the precomputed tree is downloaded and freed before the + /// multiplicity tree is allocated — so one is the peak there too. + pub tree_bytes: u64, + /// Coset weights (`n · 8`) plus the transpose scratch cap. + pub scratch_bytes: u64, +} + +impl CommitDeviceSet { + pub const fn total(&self) -> u64 { + self.lde_bytes + .saturating_add(self.snapshot_bytes) + .saturating_add(self.tree_bytes) + .saturating_add(self.scratch_bytes) + } +} + +/// Size one fused commit's device set. `base_cols` counts BASE-FIELD columns: +/// `m` for a base table, `3m` for an ext3 one (the ext3 row-major layout is +/// three adjacent base columns per element). `snapshot` is whether the +/// trace-domain column-major snapshot is retained (the main commits do, the +/// aux commits do not). +pub fn commit_device_set( + n: usize, + base_cols: usize, + blowup: usize, + snapshot: bool, +) -> CommitDeviceSet { + let n = n as u64; + let cols = base_cols as u64; + let lde = n.saturating_mul(blowup as u64); + CommitDeviceSet { + lde_bytes: base_bytes(lde, cols), + snapshot_bytes: if snapshot { base_bytes(n, cols) } else { 0 }, + tree_bytes: full_tree_bytes(lde), + scratch_bytes: n + .saturating_mul(BASE_BYTES) + .saturating_add(INPLACE_TRANSPOSE_SCRATCH_CAP_BYTES), + } +} + +/// The shape the rounds-2–4 model takes: what the AIR and the domain fix +/// before any device work starts. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct TableShape { + /// Trace rows (the interpolation domain). + pub n: usize, + pub blowup: usize, + /// Base-field main columns, preprocessed ones included. + pub main_cols: usize, + /// Ext3 aux (LogUp) columns. + pub aux_cols: usize, + /// Composition-polynomial parts: `composition_poly_degree_bound(n) / n`. + pub num_parts: usize, + /// OOD evaluation points per trace column: + /// `transition_offsets.len() · step_size`. + pub num_eval_points: usize, +} + +/// The device set of one table across rounds 2–4, term by term — what the +/// scheduler's throttle admits a table's fused task against. Everything the +/// main commit left resident stays counted (the LDE, its snapshot, its tree), +/// and each later round adds what it allocates on top: +/// +/// - R1 aux: the aux LDE (`lde · aux · 24`), the resident aux trace the LogUp +/// build left behind (`n · (aux + 1) · 24`, one extra column for the running +/// sum), and the aux tree; +/// - R2: `H` (`lde · 24`), the parts (`num_parts · lde · 24`) — `H` is still +/// alive while they are decomposed — and the parts tree; +/// - R3: the inverted denominators on the trace domain (`k · n · 24`); +/// - R4: the inverted denominators on the LDE (`(1 + k) · lde · 24`), the DEEP +/// codeword (`lde · 24`), the FRI layer chain (geometric, bounded by one more +/// codeword) and its trees (bounded by one full tree). +/// +/// Rounds do not overlap inside one table, so this is an upper bound on any +/// instant of the task, not a sum of the rounds' peaks; the counted R3/R4 +/// transients are small next to the resident LDEs. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct TableDeviceSet { + pub main: CommitDeviceSet, + pub aux_bytes: u64, + pub composition_bytes: u64, + pub deep_fri_bytes: u64, +} + +impl TableDeviceSet { + pub const fn total(&self) -> u64 { + self.main + .total() + .saturating_add(self.aux_bytes) + .saturating_add(self.composition_bytes) + .saturating_add(self.deep_fri_bytes) + } +} + +/// Size one table's rounds-2–4 device set for `shape`. +pub fn table_device_set(shape: TableShape) -> TableDeviceSet { + let TableShape { + n, + blowup, + main_cols, + aux_cols, + num_parts, + num_eval_points, + } = shape; + let main = commit_device_set(n, main_cols, blowup, true); + let (n, k, aux, parts) = ( + n as u64, + num_eval_points as u64, + aux_cols as u64, + num_parts as u64, + ); + let lde = n.saturating_mul(blowup as u64); + let aux_bytes = if aux == 0 { + 0 + } else { + ext3_bytes(lde, aux) + .saturating_add(ext3_bytes(n, aux + 1)) + .saturating_add(full_tree_bytes(lde)) + }; + let composition_bytes = if parts == 0 { + 0 + } else { + ext3_bytes(lde, 1 + parts).saturating_add(full_tree_bytes(lde)) + }; + let deep_fri_bytes = ext3_bytes(n, k) + .saturating_add(ext3_bytes(lde, 1 + k)) + .saturating_add(ext3_bytes(lde, 2)) + .saturating_add(full_tree_bytes(lde)); + TableDeviceSet { + main, + aux_bytes, + composition_bytes, + deep_fri_bytes, + } +} + +/// What the admission predicate decided for one dispatch. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum Admission { + /// No CUDA backend — no GPU, or cubins that would not load. The host path + /// is the only one; `math_cuda::device::backend` already warned once. A + /// GPU-less host is not the production pipeline, so this is not an abort. + NoDevice, + /// Below the launch-overhead floor: the host path is the faster one. + BelowFloor { lde_size: usize, floor: usize }, + /// Fits the card's admission budget. + Admitted { bytes: u64, budget: u64 }, + /// Does not fit the card even alone. + OverBudget { bytes: u64, budget: u64 }, +} + +impl Admission { + pub const fn is_admitted(&self) -> bool { + matches!(self, Admission::Admitted { .. }) + } +} + +/// The pure predicate, floor and budget supplied. The row floor is checked +/// first — a table below it never asks the device for anything, whatever its +/// width — then the bytes ceiling. +pub const fn admit_bytes(lde_size: usize, bytes: u64, floor: usize, budget: u64) -> Admission { + if lde_size < floor { + return Admission::BelowFloor { lde_size, floor }; + } + if bytes > budget { + return Admission::OverBudget { bytes, budget }; + } + Admission::Admitted { bytes, budget } +} + +/// The admission arithmetic, with the floor and the budget supplied: no +/// device, no backend, pure numbers. +#[cfg(test)] +mod tests { + use super::*; + + const GIB: u64 = 1 << 30; + /// `detect_vram_budget_bytes` on a 32 GiB card: 80% of the total. + const CARD_32_GIB_BUDGET: u64 = 32 * GIB / 5 * 4; + /// The dispatch layer's row floor (`gpu_lde::DEFAULT_GPU_LDE_THRESHOLD`). + const FLOOR: usize = 1 << 14; + + /// The synthetic over-budget table: 2^22 rows x 612 columns at blowup 2. + /// Its LDE alone is 38.25 GiB; with the snapshot and the tree the commit's + /// device set is 57.9 GiB against a 25.6 GiB budget. + #[test] + fn the_over_budget_shape_is_over_budget() { + let n = 1usize << 22; + let set = commit_device_set(n, 612, 2, true); + assert_eq!(set.lde_bytes, (n as u64) * 2 * 612 * 8); + assert_eq!(set.snapshot_bytes, (n as u64) * 612 * 8); + assert_eq!(set.tree_bytes, ((n as u64) * 2 - 1) * 32); + assert!(set.lde_bytes > 38 * GIB && set.lde_bytes < 39 * GIB); + assert!(set.total() > 57 * GIB && set.total() < 58 * GIB); + assert!(matches!( + admit_bytes(n * 2, set.total(), FLOOR, CARD_32_GIB_BUDGET), + Admission::OverBudget { .. } + )); + } + + /// LFM_HASH under RPO — 2^21 rows x (436 value + 13 preprocessed) columns — + /// fits the card at blowup 2 with one LDE buffer and does not at blowup 4. + #[test] + fn lfm_hash_rpo_fits_at_blowup_2_and_not_at_4() { + let n = 1usize << 21; + let b2 = commit_device_set(n, 449, 2, true); + assert!(b2.total() > 21 * GIB && b2.total() < 22 * GIB, "{b2:?}"); + assert!(admit_bytes(n * 2, b2.total(), FLOOR, CARD_32_GIB_BUDGET).is_admitted()); + let b4 = commit_device_set(n, 449, 4, true); + assert!(b4.total() > 35 * GIB && b4.total() < 36 * GIB, "{b4:?}"); + assert!(matches!( + admit_bytes(n * 4, b4.total(), FLOOR, CARD_32_GIB_BUDGET), + Admission::OverBudget { .. } + )); + } + + /// The whole-table set of LFM_HASH under RPO at blowup 2 — 3 aux columns, + /// two parts, two eval points — is ~23 GiB: it proves alone inside the + /// budget, and nothing else proves beside it. The old throttle model + /// (two LDE buffers plus 256 B per LDE row) put the same table at + /// 28.3 GiB, over the budget it is actually under. + #[test] + fn lfm_hash_rpo_whole_table_set() { + let shape = TableShape { + n: 1 << 21, + blowup: 2, + main_cols: 449, + aux_cols: 3, + num_parts: 2, + num_eval_points: 2, + }; + let set = table_device_set(shape); + assert!(set.total() > 23 * GIB && set.total() < 24 * GIB, "{set:?}"); + assert!(set.total() <= CARD_32_GIB_BUDGET); + assert!(2 * set.total() > CARD_32_GIB_BUDGET); + let old_model = (1u64 << 22) * (449 * 8 + 3 * 24) * 2 + (1u64 << 22) * 256; + assert!(old_model > CARD_32_GIB_BUDGET); + } + + /// One LFM_BALU chunk of 2^22 rows (14 main, 2 aux, 2 parts, 2 points) is + /// ~4.9 GiB under this model — the R1 set is 1.8 GiB; the resident aux + /// trace, `H`, the denominators, DEEP and FRI add the rest — so five prove + /// concurrently inside the budget. The sizing behind + /// `prover::lfm::chunking::BaluChunking`, whose test calls this model. + #[test] + fn a_balu_chunk_is_five_gib() { + let set = table_device_set(TableShape { + n: 1 << 22, + blowup: 2, + main_cols: 14, + aux_cols: 2, + num_parts: 2, + num_eval_points: 2, + }); + assert!( + set.main.total() > GIB + GIB / 2 && set.main.total() < 2 * GIB, + "{set:?}" + ); + assert!( + set.total() > 4 * GIB + GIB / 2 && set.total() < 5 * GIB + GIB / 2, + "{set:?}" + ); + assert!(5 * set.total() <= CARD_32_GIB_BUDGET); + assert!(6 * set.total() > CARD_32_GIB_BUDGET); + } + + /// A table without aux or parts (d=1 with no lookups) counts only what it + /// allocates. + #[test] + fn absent_rounds_cost_nothing() { + let set = table_device_set(TableShape { + n: 1 << 16, + blowup: 2, + main_cols: 8, + aux_cols: 0, + num_parts: 0, + num_eval_points: 1, + }); + assert_eq!(set.aux_bytes, 0); + assert_eq!(set.composition_bytes, 0); + assert!(set.deep_fri_bytes > 0); + } + + /// The aux commit has no snapshot; its ext3 columns count as three base + /// columns each. + #[test] + fn aux_sets_have_no_snapshot() { + let set = commit_device_set(1 << 20, 3 * 3, 2, false); + assert_eq!(set.snapshot_bytes, 0); + assert_eq!(set.lde_bytes, ext3_bytes(1 << 21, 3)); + } + + /// The row floor is checked before the ceiling: a tiny table with an + /// absurd byte count is "too small", never "over budget" — it will not ask + /// the device for anything. + #[test] + fn the_floor_is_checked_before_the_budget() { + assert!(matches!( + admit_bytes(1 << 13, u64::MAX, FLOOR, CARD_32_GIB_BUDGET), + Admission::BelowFloor { .. } + )); + assert!(matches!( + admit_bytes(FLOOR, u64::MAX, FLOOR, CARD_32_GIB_BUDGET), + Admission::OverBudget { .. } + )); + } + + /// FRI re-derives admission at width 1: a narrow transient over a large + /// domain always clears the ceiling. The floor stays a row count, so it + /// does not degenerate there either. + #[test] + fn fri_at_width_one_never_degenerates() { + let n0 = 1usize << 24; + let bytes = ext3_bytes(n0 as u64, 1) + full_tree_bytes(n0 as u64); + assert!(admit_bytes(n0, bytes, FLOOR, CARD_32_GIB_BUDGET).is_admitted()); + } + + /// A budget of `u64::MAX` (query failed) makes the ceiling inert — the + /// floor alone decides, which is the pre-admission behaviour. + #[test] + fn an_unbounded_budget_is_inert() { + assert!(admit_bytes(1 << 20, u64::MAX - 1, FLOOR, u64::MAX).is_admitted()); + } + + /// The table's committed width is what the model takes: the row floor is + /// width-blind on purpose, the ceiling is not. + #[test] + fn width_moves_the_ceiling_not_the_floor() { + let n = 1usize << 21; + let narrow = commit_device_set(n, 4, 2, true); + let wide = commit_device_set(n, 612, 2, true); + assert!(admit_bytes(n * 2, narrow.total(), FLOOR, CARD_32_GIB_BUDGET).is_admitted()); + assert!(matches!( + admit_bytes(n * 2, wide.total(), FLOOR, CARD_32_GIB_BUDGET), + Admission::OverBudget { .. } + )); + } +} diff --git a/crypto/stark/src/gpu_lde.rs b/crypto/stark/src/gpu_lde.rs index 35310b409..2707a96d8 100644 --- a/crypto/stark/src/gpu_lde.rs +++ b/crypto/stark/src/gpu_lde.rs @@ -206,124 +206,14 @@ fn gpu_device_only_threshold() -> usize { // is an upper bound, so a narrow transient always clears it — a cells FLOOR // would degenerate there, which is why the floor stays a row count. -/// Bytes per Goldilocks element on device. -const BASE_BYTES: u64 = 8; - -/// Bytes per ext3 element on device — three adjacent base columns. -const EXT3_BYTES: u64 = 3 * BASE_BYTES; - -/// Bytes of one Merkle node. Every commitment hash the device dispatches on -/// emits a 32-byte digest — a four-felt Goldilocks digest is exactly 32 -/// canonical bytes — so the node buffer costs the same under every hash. -const MERKLE_NODE_BYTES: u64 = 32; - -/// Cap on the in-place transpose's device scratch, mirrored from -/// `math_cuda::lde::INPLACE_TRANSPOSE_SCRATCH_BYTES` (private there). The -/// admission wants a bound, not the block geometry. -const INPLACE_TRANSPOSE_SCRATCH_CAP_BYTES: u64 = 256 << 20; - -/// The device working set one fused row-major commit allocates, term by term -/// (`math_cuda::lde::coset_lde_row_major_inner` after the in-place transpose -/// of #956): ONE LDE buffer, the optional trace-domain snapshot, the full -/// Merkle node buffer, and the small scratch (coset weights plus the capped -/// transpose scratch). `one_lde_buffer::vram_arm` prints the same three big -/// terms; the model here is the model it measures. -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -pub struct CommitDeviceSet { - /// `lde_size · base_cols · 8`: the row-major LDE, transposed in place. - pub lde_bytes: u64, - /// `n · base_cols · 8`: the pre-NTT column-major snapshot the LogUp - /// fingerprint kernel reads in place (main commits only). - pub snapshot_bytes: u64, - /// `(2 · leaves − 1) · 32` with `leaves = lde_size / 2`: one full row-pair - /// tree. The preprocessed split path builds two, sequentially on one - /// stream — the precomputed tree is downloaded and freed before the - /// multiplicity tree is allocated — so one is the peak there too. - pub tree_bytes: u64, - /// Coset weights (`n · 8`) plus the transpose scratch cap. - pub scratch_bytes: u64, -} - -impl CommitDeviceSet { - pub const fn total(&self) -> u64 { - self.lde_bytes - .saturating_add(self.snapshot_bytes) - .saturating_add(self.tree_bytes) - .saturating_add(self.scratch_bytes) - } -} - -/// `(2 · leaves − 1) · 32` for the row-pair tree over `lde_size` rows. -pub const fn full_tree_bytes(lde_size: u64) -> u64 { - lde_size.saturating_sub(1).saturating_mul(MERKLE_NODE_BYTES) -} - -/// Bytes of `cols` ext3 columns over `rows` rows. -pub const fn ext3_bytes(rows: u64, cols: u64) -> u64 { - rows.saturating_mul(cols).saturating_mul(EXT3_BYTES) -} - -/// Size one fused commit's device set. `base_cols` counts BASE-FIELD columns: -/// `m` for a base table, `3m` for an ext3 one (the ext3 row-major layout is -/// three adjacent base columns per element). `snapshot` is whether the -/// trace-domain column-major snapshot is retained (the main commits do, the -/// aux commits do not). -pub fn commit_device_set( - n: usize, - base_cols: usize, - blowup: usize, - snapshot: bool, -) -> CommitDeviceSet { - let n = n as u64; - let cols = base_cols as u64; - let lde = n.saturating_mul(blowup as u64); - CommitDeviceSet { - lde_bytes: lde.saturating_mul(cols).saturating_mul(BASE_BYTES), - snapshot_bytes: if snapshot { - n.saturating_mul(cols).saturating_mul(BASE_BYTES) - } else { - 0 - }, - tree_bytes: full_tree_bytes(lde), - scratch_bytes: n - .saturating_mul(BASE_BYTES) - .saturating_add(INPLACE_TRANSPOSE_SCRATCH_CAP_BYTES), - } -} - -/// What the admission predicate decided for one dispatch. -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -pub enum Admission { - /// No CUDA backend — no GPU, or cubins that would not load. The host path - /// is the only one; `math_cuda::device::backend` already warned once. A - /// GPU-less host is not the production pipeline, so this is not an abort. - NoDevice, - /// Below the launch-overhead floor: the host path is the faster one. - BelowFloor { lde_size: usize, floor: usize }, - /// Fits the card's admission budget. - Admitted { bytes: u64, budget: u64 }, - /// Does not fit the card even alone. - OverBudget { bytes: u64, budget: u64 }, -} - -impl Admission { - pub const fn is_admitted(&self) -> bool { - matches!(self, Admission::Admitted { .. }) - } -} - -/// The pure predicate, floor and budget supplied. The row floor is checked -/// first — a table below it never asks the device for anything, whatever its -/// width — then the bytes ceiling. -pub const fn admit_bytes(lde_size: usize, bytes: u64, floor: usize, budget: u64) -> Admission { - if lde_size < floor { - return Admission::BelowFloor { lde_size, floor }; - } - if bytes > budget { - return Admission::OverBudget { bytes, budget }; - } - Admission::Admitted { bytes, budget } -} +// The arithmetic — the device-set model and the pure predicate — lives in +// `crate::device_set`, free of the `cuda` gate, because the per-table +// scheduler's throttle reads the same model on every build. Re-exported here +// so the dispatch layer's callers keep one path. +use crate::device_set::BASE_BYTES; +pub use crate::device_set::{ + Admission, CommitDeviceSet, admit_bytes, commit_device_set, ext3_bytes, full_tree_bytes, +}; /// The process predicate: `gpu_lde_threshold()` as the floor and the card's /// admission budget ([`device_vram_budget_bytes`]: 80% of device memory, or @@ -342,7 +232,10 @@ pub(crate) fn admit(lde_size: usize, bytes: u64) -> Admission { /// prover's driver re-raises the panic payload with the message intact, and /// its own `[gpu]` lines name the table. #[derive(Clone, Copy, Debug)] -pub(crate) struct DispatchShape { +pub(crate) struct DispatchShape<'a> { + /// The AIR's name — the prover passes it down so the diagnostic names the + /// table at the site, not only in the re-raised payload. + pub table: &'a str, pub what: &'static str, pub n: usize, pub base_cols: usize, @@ -376,6 +269,13 @@ fn live_vram_line() -> String { } } +/// The live device posture every device-path failure reports: free/total VRAM +/// and the mempool release threshold. Shared by the dispatch layer's aborts and +/// the prover's clean errors so the two read the same line. +pub(crate) fn device_path_status() -> String { + format!("{}; {}", live_vram_line(), mempool_line()) +} + fn mempool_line() -> String { match math_cuda::device::mempool_release_threshold_bytes() { u64::MAX => "mempool retains freed blocks (release threshold unset)".to_string(), @@ -387,11 +287,12 @@ fn mempool_line() -> String { /// and shape, the device set term by term, the admission budget, the live /// free/total VRAM and the mempool posture. fn device_path_diagnostic( - shape: &DispatchShape, + shape: &DispatchShape<'_>, set: Option<&CommitDeviceSet>, failure: &DevicePathFailure, ) -> String { let DispatchShape { + table, what, n, base_cols, @@ -417,10 +318,9 @@ fn device_path_diagnostic( None => String::new(), }; format!( - "{what}: rows {n} x {base_cols} base cols @ blowup {blowup} (LDE {}); {reason}{set_line}; {}; {}", + "table {table}: {what}: rows {n} x {base_cols} base cols @ blowup {blowup} (LDE {}); {reason}{set_line}; {}", n.saturating_mul(*blowup), - live_vram_line(), - mempool_line() + device_path_status() ) } @@ -467,7 +367,7 @@ pub(crate) fn test_only_host_fallback() -> bool { /// catch and re-raise it on the calling thread (`run_admitted`), so the prove /// stops with the reason instead of finishing hours later on the CPU. pub(crate) fn abort_or_test_fallback( - shape: &DispatchShape, + shape: &DispatchShape<'_>, set: Option<&CommitDeviceSet>, failure: DevicePathFailure, ) { @@ -483,7 +383,7 @@ pub(crate) fn abort_or_test_fallback( /// below the floor (the host commit is the right one), `Some(())` when /// admitted. Over budget is an abort — or, under the test-only switch, a /// reported host commit. -fn admit_commit(lde_size: usize, shape: &DispatchShape, set: &CommitDeviceSet) -> Option<()> { +fn admit_commit(lde_size: usize, shape: &DispatchShape<'_>, set: &CommitDeviceSet) -> Option<()> { match admit(lde_size, set.total()) { Admission::NoDevice | Admission::BelowFloor { .. } => None, Admission::Admitted { .. } => Some(()), @@ -502,7 +402,7 @@ fn admit_commit(lde_size: usize, shape: &DispatchShape, set: &CommitDeviceSet) - /// LogUp aux build's `ResidentAux`): no row floor — the data is there, and a /// decline would not be "take the faster host path" but "download it to /// commit on the host" — only the bytes ceiling. -fn admit_resident_commit(shape: &DispatchShape, set: &CommitDeviceSet) -> Option<()> { +fn admit_resident_commit(shape: &DispatchShape<'_>, set: &CommitDeviceSet) -> Option<()> { let budget = device_vram_budget_bytes()?; match admit_bytes(usize::MAX, set.total(), 0, budget) { Admission::OverBudget { bytes, budget } => { @@ -545,9 +445,8 @@ fn admit_transient(lde_size: usize, bytes: u64, what: &str) -> bool { /// and the recovery proceeds. fn refuse_host_recovery(what: &str, rows: usize, main_cols: usize, aux_cols: usize) { let msg = format!( - "{what}: rows {rows} main cols {main_cols} aux cols {aux_cols}; {}; {}", - live_vram_line(), - mempool_line() + "{what}: rows {rows} main cols {main_cols} aux cols {aux_cols}; {}", + device_path_status() ); if test_only_host_fallback() { eprintln!("[gpu] TEST-ONLY host recovery: {msg}"); @@ -760,12 +659,11 @@ pub(crate) fn device_only_disabled() -> bool { /// transient GPU error), the table's recovery reaches /// [`refuse_host_recovery`]: in production that is a loud abort with the shape /// and the live VRAM — host RAM is a cache, not a compute path. Under the -/// test-only fallback ([`test_only_host_fallback`]) R2 and the R1 resident-aux -/// commit download what the host arms need (the resident LDEs at R2, the -/// resident aux trace plus the main LDE at R1), bump their site's counter -/// ([`GPU_DEVICE_ONLY_DOWNGRADES`] at R2, [`GPU_RESIDENT_AUX_DOWNGRADES`] at -/// R1) and continue host-backed; R3 and R4 have no host recovery of their own -/// and assert on the buffer they are about to read. +/// test-only fallback ([`test_only_host_fallback`]) R2 downloads what the host +/// arms need (the resident LDEs), bumps [`GPU_DEVICE_ONLY_DOWNGRADES`] and +/// continues host-backed; the R1 resident-aux commit has no host recovery at +/// all (a decline after the drain-and-retry is a `ProvingError::DevicePath`), +/// and R3 and R4 assert on the buffer they are about to read. /// /// `zerofier_uniform` must be the R1-derived conservative form (all constraints /// share `end_exemptions == 0`), which implies `ZerofierEvaluations::is_uniform` @@ -1343,7 +1241,9 @@ pub fn gpu_leaf_hash_calls() -> u64 { /// Merkle → single D2H. Keeps the Merkle tree resident on device (in the /// handle's `.tree`); the returned host `MerkleTree` is root only, so query /// openings gather paths from the device tree via [`gather_proofs_dev`]. +#[allow(clippy::too_many_arguments)] pub(crate) fn try_expand_leaf_and_tree_row_major_keep( + table: &str, row_major: &[FieldElement], predev: Option<&math_cuda::CudaSlice>, n: usize, @@ -1372,6 +1272,7 @@ where } let lde_size = n.saturating_mul(blowup_factor); let shape = DispatchShape { + table, what: "R1 main commit", n, base_cols: m, @@ -1464,6 +1365,7 @@ where #[allow(clippy::type_complexity)] #[allow(clippy::too_many_arguments)] pub(crate) fn try_expand_split_trees_row_major_keep( + table: &str, row_major: &[FieldElement], predev: Option<&math_cuda::CudaSlice>, n: usize, @@ -1498,6 +1400,7 @@ where } let lde_size = n.saturating_mul(blowup_factor); let shape = DispatchShape { + table, what: "R1 main commit (preprocessed split)", n, base_cols: m, @@ -1568,6 +1471,7 @@ where /// row-major Keccak → Merkle → single D2H → transpose to GpuLdeExt3 handle. /// Same optimization as the base-field path: no extract_columns, no CPU transpose. pub(crate) fn try_expand_leaf_and_tree_ext3_row_major_keep( + table: &str, row_major: &[FieldElement], n: usize, m: usize, @@ -1597,6 +1501,7 @@ where let m3 = m * 3; let lde_size = n.saturating_mul(blowup_factor); let shape = DispatchShape { + table, what: "R1 aux commit", n, base_cols: m3, @@ -2337,16 +2242,12 @@ pub fn gpu_device_only_downgrades() -> u64 { GPU_DEVICE_ONLY_DOWNGRADES.load(Ordering::Relaxed) } -/// R1 downgrades, and only those: times the resident aux trace was downloaded -/// so the aux commit could continue on the host arms, after the device aux LDE -/// declined and the drain-and-retry either did not run or declined again -/// ([`materialize_aux_trace_host`], the sole site that bumps this — and, like -/// every host recovery, only under the test-only fallback). Independent -/// of the device-only gate — the site is entered whenever `aux_resident()` is -/// set, whatever the gate said — so a table that was never device-only can land -/// here, and a nonzero value points at sustained VRAM pressure rather than a -/// gate miss. Read it against [`GPU_RESIDENT_AUX_RETRIES`]: retries alone mean -/// the drain absorbed the pressure, retries plus downgrades mean it did not. +/// R1 resident-aux downgrades. Retired: a resident aux LDE that declines after +/// the drain-and-retry is now a clean `ProvingError::DevicePath` in the prover +/// (host RAM is a cache, not a compute path), so nothing bumps this and it +/// reads zero. The accessor stays for the integration suite's assertion that +/// it IS zero. Read [`GPU_RESIDENT_AUX_RETRIES`] for the pressure signal: +/// retries mean the drain absorbed a transient decline. pub(crate) static GPU_RESIDENT_AUX_DOWNGRADES: AtomicU64 = AtomicU64::new(0); pub fn gpu_resident_aux_downgrades() -> u64 { GPU_RESIDENT_AUX_DOWNGRADES.load(Ordering::Relaxed) @@ -2378,8 +2279,9 @@ pub fn gpu_resident_aux_retries() -> u64 { /// Recover a device-only table for the host path: download the resident main /// and aux LDEs from their device handles into the host buffers and clear the /// device-only flag. A side whose host buffer is already populated (a mixed -/// state: one commit fell back to CPU while the other stayed device-only) is -/// kept as is — only the missing side is downloaded. +/// state: one side's dispatch declined below the device-only envelope while +/// the other stayed device-only) is kept as is — only the missing side is +/// downloaded. /// /// In production this recovery does not run: host RAM is a cache, not a /// compute path, so a device-only table whose downstream dispatch declined is @@ -2454,7 +2356,7 @@ where } let (m, lde) = (h.m, h.lde_size); // Short download: degrade like the sibling paths - // (`download_main_lde_row_major`, `materialize_aux_trace_host`) + // (`download_main_lde_row_major`, `download_composition_parts_host`) // rather than panic on the slab slicing below. if slabs.len() != m * lde * 3 { return false; @@ -2563,58 +2465,6 @@ where }) } -/// R1 counterpart of [`materialize_lde_trace_host`]: download the resident -/// aux trace (already row-major ext3, matching the host layout) into the -/// trace's aux table, so the aux commit continues on the host arms when the -/// device aux LDE declined at runtime — twice, the caller having drained the -/// device and retried in between. Refused in production on the same terms as -/// its R2 counterpart ([`refuse_host_recovery`]). -pub(crate) fn materialize_aux_trace_host(trace: &mut crate::trace::TraceTable) -> bool -where - F: IsField + IsSubFieldOf + 'static, - E: IsField + 'static, -{ - if !is_goldilocks_ext3_tower::() { - return false; - } - let (buf, rows, cols) = match trace.aux_resident.as_ref() { - Some(ra) => (ra.buf.clone(), ra.num_rows, ra.num_aux_cols), - None => return false, - }; - refuse_host_recovery( - "R1 aux commit on the host from the resident aux trace (device aux LDE declined after \ - the drain-and-retry)", - rows, - trace.num_main_columns, - cols, - ); - let Ok(be) = math_cuda::device::backend() else { - return false; - }; - let stream = be.next_stream(); - let Ok(raw) = stream.clone_dtoh(buf.as_ref()) else { - return false; - }; - if stream.synchronize().is_err() || raw.len() != rows * cols * 3 { - return false; - } - let data = u64_to_ext3_vec::(&raw); - trace.aux_table = crate::table::Table::new(data, cols); - trace.num_aux_columns = cols; - // The declined device LDE attempt can leave kernels enqueued on another - // stream still reading this buffer; its owning stream is long idle, so - // dropping here would complete the stream-ordered free immediately and - // the pool could hand the memory to a concurrent table's allocation - // while those kernels run. Drain the device before the drop — this is a - // rare recovery path. - if be.ctx.synchronize().is_err() { - return false; - } - trace.aux_resident = None; - GPU_RESIDENT_AUX_DOWNGRADES.fetch_add(1, Ordering::Relaxed); - true -} - /// Diagnostic: download a resident ext3 handle (3-slab layout) as per-column /// host Vecs. Used by the xcheck post-mortem to compare the committed R2 /// parts against a host recompute. @@ -2954,6 +2804,7 @@ unsafe fn ext3_slice_to_u64(col: &[FieldElement]) -> &[u64] { /// The resident buffer is only borrowed: the device-input LDE copies it /// device-to-device into its own scratch, so `ra` stays valid afterwards. pub(crate) fn try_expand_leaf_and_tree_ext3_row_major_keep_dev( + table: &str, ra: &math_cuda::logup::ResidentAux, blowup_factor: usize, weights: &[FieldElement], @@ -2976,6 +2827,7 @@ where // No row floor: the aux trace is already on device. Only the bytes // ceiling — a resident input that will not fit its own LDE is an abort. let shape = DispatchShape { + table, what: "R1 aux commit (resident)", n: ra.num_rows, base_cols: ra.num_aux_cols * 3, @@ -3001,8 +2853,9 @@ where .inspect_err(|e| { // Surface the swallowed driver error (e.g. OOM): the caller drains the // device and retries — a device-side recovery, which is why this is a - // decline and not an abort. If the retry declines too, the caller's - // host downgrade is refused by `materialize_aux_trace_host`. + // decline and not an abort. If the retry declines too, the caller + // reports the failure (`ProvingError::DevicePath`); there is no host + // downgrade any more. eprintln!( "[gpu] resident aux LDE failed (rows={} cols={} blowup={}): {e:?}", ra.num_rows, ra.num_aux_cols, blowup_factor @@ -4029,108 +3882,6 @@ where Some(decommits) } -/// The admission arithmetic, with the floor and the budget supplied: no -/// device, no backend, pure numbers. -#[cfg(test)] -mod admission_tests { - use super::*; - - const GIB: u64 = 1 << 30; - /// `detect_vram_budget_bytes` on a 32 GiB card: 80% of the total. - const CARD_32_GIB_BUDGET: u64 = 32 * GIB / 5 * 4; - const FLOOR: usize = DEFAULT_GPU_LDE_THRESHOLD; - - /// The brief's synthetic over-budget table: 2^22 rows x 612 columns at - /// blowup 2. Its LDE alone is 38.25 GiB; with the snapshot and the tree the - /// commit's device set is 57.6 GiB against a 25.6 GiB budget. - #[test] - fn the_brief_shape_is_over_budget() { - let n = 1usize << 22; - let set = commit_device_set(n, 612, 2, true); - assert_eq!(set.lde_bytes, (n as u64) * 2 * 612 * 8); - assert_eq!(set.snapshot_bytes, (n as u64) * 612 * 8); - assert_eq!(set.tree_bytes, ((n as u64) * 2 - 1) * 32); - assert!(set.lde_bytes > 38 * GIB && set.lde_bytes < 39 * GIB); - assert!(set.total() > 57 * GIB && set.total() < 58 * GIB); - assert!(matches!( - admit_bytes(n * 2, set.total(), FLOOR, CARD_32_GIB_BUDGET), - Admission::OverBudget { .. } - )); - } - - /// LFM_HASH under RPO — 2^21 rows x (436 value + 13 preprocessed) columns — - /// fits the card at blowup 2 with one LDE buffer (the point of #956) and - /// does not at blowup 4. The GPU-SEAMS arithmetic, at the committed width. - #[test] - fn lfm_hash_rpo_fits_at_blowup_2_and_not_at_4() { - let n = 1usize << 21; - let b2 = commit_device_set(n, 449, 2, true); - assert!(b2.total() > 21 * GIB && b2.total() < 22 * GIB, "{b2:?}"); - assert!(admit_bytes(n * 2, b2.total(), FLOOR, CARD_32_GIB_BUDGET).is_admitted()); - let b4 = commit_device_set(n, 449, 4, true); - assert!(b4.total() > 35 * GIB && b4.total() < 36 * GIB, "{b4:?}"); - assert!(matches!( - admit_bytes(n * 4, b4.total(), FLOOR, CARD_32_GIB_BUDGET), - Admission::OverBudget { .. } - )); - } - - /// The aux commit has no snapshot; its ext3 columns count as three base - /// columns each. - #[test] - fn aux_sets_have_no_snapshot() { - let set = commit_device_set(1 << 20, 3 * 3, 2, false); - assert_eq!(set.snapshot_bytes, 0); - assert_eq!(set.lde_bytes, ext3_bytes(1 << 21, 3)); - } - - /// The row floor is checked before the ceiling: a tiny table with an - /// absurd byte count is "too small", never "over budget" — it will not ask - /// the device for anything. - #[test] - fn the_floor_is_checked_before_the_budget() { - assert!(matches!( - admit_bytes(1 << 13, u64::MAX, FLOOR, CARD_32_GIB_BUDGET), - Admission::BelowFloor { .. } - )); - assert!(matches!( - admit_bytes(FLOOR, u64::MAX, FLOOR, CARD_32_GIB_BUDGET), - Admission::OverBudget { .. } - )); - } - - /// FRI re-derives admission at width 1: a narrow transient over a large - /// domain always clears the ceiling. The floor stays a row count, so it - /// does not degenerate there either. - #[test] - fn fri_at_width_one_never_degenerates() { - let n0 = 1usize << 24; - let bytes = ext3_bytes(n0 as u64, 1) + full_tree_bytes(n0 as u64); - assert!(admit_bytes(n0, bytes, FLOOR, CARD_32_GIB_BUDGET).is_admitted()); - } - - /// A budget of `u64::MAX` (query failed) makes the ceiling inert — the - /// floor alone decides, which is the pre-admission behaviour. - #[test] - fn an_unbounded_budget_is_inert() { - assert!(admit_bytes(1 << 20, u64::MAX - 1, FLOOR, u64::MAX).is_admitted()); - } - - /// The table's committed width is what the model takes: the row floor is - /// width-blind on purpose, the ceiling is not. - #[test] - fn width_moves_the_ceiling_not_the_floor() { - let n = 1usize << 21; - let narrow = commit_device_set(n, 4, 2, true); - let wide = commit_device_set(n, 612, 2, true); - assert!(admit_bytes(n * 2, narrow.total(), FLOOR, CARD_32_GIB_BUDGET).is_admitted()); - assert!(matches!( - admit_bytes(n * 2, wide.total(), FLOOR, CARD_32_GIB_BUDGET), - Admission::OverBudget { .. } - )); - } -} - /// The abort itself, on a real device. `LAMBDA_VM_VRAM_BUDGET_MB` is read once /// at backend init, so this test runs in its own process with the budget /// lowered to 1 GiB — the shape is then over budget on any card while its host @@ -4167,7 +3918,14 @@ mod admission_box_tests { let data: Vec = (0..n * m).map(|i| Fp::from(i as u64)).collect(); let weights: Vec = (0..n).map(|i| Fp::from(i as u64 + 1)).collect(); let committed = try_expand_leaf_and_tree_row_major_keep::>( - &data, None, n, m, blowup, &weights, true, + "admission_box_test", + &data, + None, + n, + m, + blowup, + &weights, + true, ); panic!( "the over-budget commit returned {} instead of aborting", @@ -4225,7 +3983,16 @@ mod split_tree_tests { let (pre_tree, mult_tree, handle, lde) = try_expand_split_trees_row_major_keep::>( - &data, None, n, m, blowup, &weights, split, true, true, + "split_tree_test", + &data, + None, + n, + m, + blowup, + &weights, + split, + true, + true, ) .expect("GPU split path must engage above the threshold"); let pre_tree = pre_tree.expect("precomputed tree was requested"); diff --git a/crypto/stark/src/lib.rs b/crypto/stark/src/lib.rs index fd9b393e8..a7edb7996 100644 --- a/crypto/stark/src/lib.rs +++ b/crypto/stark/src/lib.rs @@ -11,6 +11,7 @@ pub mod constraint_ir; pub mod constraints; pub mod context; pub mod debug; +pub mod device_set; pub mod domain; #[cfg(any(test, feature = "test-utils"))] pub mod examples; diff --git a/crypto/stark/src/prover.rs b/crypto/stark/src/prover.rs index daf405280..96c4ac445 100644 --- a/crypto/stark/src/prover.rs +++ b/crypto/stark/src/prover.rs @@ -110,6 +110,12 @@ pub enum ProvingError { /// `WrongParameter` because the cause is internal prover machinery, not a /// caller-supplied parameter. Carries the underlying `FFTError`'s message. Fft(String), + /// The device path is the production path and it was unavailable for a + /// table after its device-side recovery ran: a resident aux LDE that + /// declined again after the drain-and-retry. Host RAM is a cache, not a + /// compute path, so there is no host arm to continue on; the message names + /// the table, the shape and the live device posture. + DevicePath(String), } impl From for ProvingError { @@ -718,23 +724,6 @@ pub fn storage_estimate_parallelism() -> usize { } } -/// Heuristic peak device bytes for one table: co-resident LDE columns plus the -/// resident Merkle trees, with a scratch factor for NTT and leaf transients. A -/// deliberate over estimate for a safety ceiling, not a precise allocator. Pass -/// aux_cols == 0 when the aux LDE is not yet resident (R1 main commit). -fn estimate_table_vram_bytes(main_cols: usize, aux_cols: usize, lde_size: usize) -> u64 { - const BYTES_PER_BASE: u64 = 8; - const EXT3_BYTES: u64 = 24; - const SCRATCH_FACTOR: u64 = 2; - const RESIDENT_TREE_BYTES_PER_LDE: u64 = 256; - let lde = lde_size as u64; - let per_row = (main_cols as u64).saturating_mul(BYTES_PER_BASE) - + (aux_cols as u64).saturating_mul(EXT3_BYTES); - let lde_term = lde.saturating_mul(per_row).saturating_mul(SCRATCH_FACTOR); - let tree_term = lde.saturating_mul(RESIDENT_TREE_BYTES_PER_LDE); - lde_term.saturating_add(tree_term) -} - /// Byte-budget admission gate for concurrently proven tables. `acquire` /// blocks until the requested bytes fit under the budget, releasing on /// permit drop. An oversized request is admitted alone (when nothing else @@ -794,6 +783,7 @@ fn run_admitted( estimates: &[u64], gate: &VramGate, workers: usize, + label: impl Fn(usize) -> String + Sync, task: impl Fn(usize) -> T + Sync, ) -> Vec> { let results: Vec>> = estimates @@ -849,6 +839,9 @@ fn run_admitted( match out { Ok(v) => *results[idx].lock().unwrap() = Some(v), Err(payload) => { + // The worker's message names the stage and the + // shape; the driver knows which table it was. + let payload = name_panic_payload(payload, &label(idx)); let mut slot = first_panic.lock().unwrap_or_else(|e| e.into_inner()); if slot.is_none() { *slot = Some(payload); @@ -875,6 +868,24 @@ fn run_admitted( .collect() } +/// Prefix a string panic payload with the table's name so the re-raised +/// message says which table failed; a payload that is not a string is passed +/// through unchanged. +fn name_panic_payload( + payload: Box, + table: &str, +) -> Box { + let message = payload.downcast_ref::().cloned().or_else(|| { + payload + .downcast_ref::<&'static str>() + .map(|m| (*m).to_string()) + }); + match message { + Some(m) => Box::new(format!("table {table}: {m}")), + None => payload, + } +} + /// Table indices sorted heaviest-first by estimate. fn heaviest_first(estimates: &[u64]) -> Vec { let mut order: Vec = (0..estimates.len()).collect(); @@ -1247,8 +1258,10 @@ pub trait IsStarkProver< /// `precomputed`: if present, the leading `num_cols` columns are committed /// as a separate Merkle tree (the precomputed split for preprocessed /// tables) and the root is checked against the AIR-hardcoded commitment. + /// `table` is the AIR's name, for the device diagnostics. #[allow(clippy::type_complexity)] fn commit_main_trace( + #[cfg_attr(not(feature = "cuda"), allow(unused_variables))] table: &str, trace: &TraceTable, domain: &Domain, twiddles: &LdeTwiddles, @@ -1285,6 +1298,7 @@ pub trait IsStarkProver< Field, H::Batched, >( + table, trace_slice, trace.main_rowmajor_dev(), n, @@ -1350,6 +1364,7 @@ pub trait IsStarkProver< Field, H::Batched, >( + table, trace_slice, trace.main_rowmajor_dev(), n, @@ -2041,14 +2056,17 @@ pub trait IsStarkProver< // Every arm below runs the HOST evaluator, which reads `get_main` / // `get_aux`. Under device-only those buffers are intentionally empty, // so landing here means the device decompose AND the `H` download both - // failed. The gate is a static predicate and cannot mirror every - // dynamic decline, so recover rather than abort: download the resident - // LDEs into the host buffers (which also clears the device-only flag) - // and let the host arms run — slower for this table, never wrong. The - // assert is left for the case where the handles themselves cannot - // serve the data, so that failure carries the device-only contract's - // message rather than a bare index-out-of-bounds from somewhere inside - // the evaluator. + // failed. In production that is the failure to report — host RAM is a + // cache, not a compute path — and `materialize_lde_trace_host` aborts + // with the shape and the live VRAM before returning. Only under the + // test-only host fallback (`LAMBDA_VM_TEST_ONLY_HOST_FALLBACK`, the + // `LAMBDA_VM_GPU_FORCE_DOWNGRADE` hook or the `test-cuda-faults` + // feature) does it download the resident LDEs into the host buffers + // (clearing the device-only flag) and let the host arms run; the + // `gpu_force_downgrade` binary asserts on exactly that. The assert is + // left for the case where the handles themselves cannot serve the data, + // so that failure carries the device-only contract's message rather + // than a bare index-out-of-bounds from somewhere inside the evaluator. #[cfg(feature = "cuda")] if precomputed_parts.is_none() && lde_trace.host_trace_empty() { let recovered = crate::gpu_lde::materialize_lde_trace_host(lde_trace); @@ -3689,17 +3707,31 @@ pub trait IsStarkProver< let vram_gate = VramGate::new(vram_budget); - // R1 main commit: only the main LDE and its Merkle scratch are resident, - // so the aux columns add nothing to this phase's working set. + // R1 main commit: the fused commit's device set — one LDE buffer, the + // trace snapshot, the tree and the scratch — the same model the + // dispatch layer admits the commit against (`crate::device_set`). let main_estimates: Vec = air_trace_pairs .iter() .enumerate() .map(|(idx, (_, trace, _))| { - let lde_size = domains[idx].interpolation_domain_size * domains[idx].blowup_factor; - estimate_table_vram_bytes(trace.num_main_columns, 0, lde_size) + let domain = &domains[idx]; + crate::device_set::commit_device_set( + domain.interpolation_domain_size, + trace.num_main_columns, + domain.blowup_factor, + true, + ) + .total() }) .collect(); + // The AIR names, for the driver threads' panic payloads: a device abort + // names its stage and shape, the driver adds which table. + let table_names: Vec = air_trace_pairs + .iter() + .map(|(air, _, _)| air.name().to_string()) + .collect(); + // Spill main traces to mmap before Round 1 LDE. #[cfg(feature = "disk-spill")] if storage_mode == StorageMode::Disk { @@ -3751,6 +3783,7 @@ pub trait IsStarkProver< &main_estimates, &vram_gate, k, + |idx| table_names[idx].clone(), |idx| { let (air, trace, _) = &air_trace_pairs[idx]; let domain = &domains[idx]; @@ -3766,6 +3799,7 @@ pub trait IsStarkProver< let device_only = Self::device_only_for(*air, domain); Self::commit_main_trace( + air.name(), *trace, domain, twiddles, @@ -3899,9 +3933,18 @@ pub trait IsStarkProver< .iter() .enumerate() .map(|(idx, (air, trace, _))| { - let lde_size = domains[idx].interpolation_domain_size * domains[idx].blowup_factor; + let domain = &domains[idx]; + let n = domain.interpolation_domain_size; let (_, aux_cols) = air.trace_layout(); - estimate_table_vram_bytes(trace.num_main_columns, aux_cols, lde_size) + crate::device_set::table_device_set(crate::device_set::TableShape { + n, + blowup: domain.blowup_factor, + main_cols: trace.num_main_columns, + aux_cols, + num_parts: air.composition_poly_degree_bound(n) / n, + num_eval_points: air.context().transition_offsets.len() * air.step_size(), + }) + .total() }) .collect(); @@ -3989,20 +4032,21 @@ pub trait IsStarkProver< // Device-only for the aux commit: the main commit's // gate AND a produced main device handle. The aux side // may be MORE conservative than main (never less) — if - // the GPU main commit declined and fell back to CPU, - // skipping the aux D2H here would leave a device-only - // trace with no main handle to serve it. + // the GPU main commit declined below the floor and + // committed on the host, skipping the aux D2H here would + // leave a device-only trace with no main handle to serve + // it. #[cfg(feature = "cuda")] - let mut device_only = Self::device_only_for(*air, domain) + let device_only = Self::device_only_for(*air, domain) && gpu_main_cells[idx].lock().unwrap().is_some(); // Resident GPU path: aux columns already on device (from // the resident LogUp aux build) — LDE straight from device // memory, no upload, no host column extraction. When the // resident build fired the host aux trace is empty, so a - // device LDE failure downloads the resident aux trace and - // continues on the host arms below (falling through as-is - // would commit a zero aux trace). + // device LDE failure gets one drain-and-retry and is then + // a clean error (falling through as-is would commit a + // zero aux trace). #[cfg(feature = "cuda")] if trace.aux_resident().is_some() { #[cfg(feature = "instruments")] @@ -4014,6 +4058,7 @@ pub trait IsStarkProver< FieldExtension, H::Batched, >( + air.name(), ra, domain.blowup_factor, &twiddles.coset_weights, @@ -4049,80 +4094,19 @@ pub trait IsStarkProver< Some(handle), )); } - // The device aux LDE declined at runtime (transient - // VRAM pressure, usually) and there is no host aux - // trace to fall back to. Same class as the R2 - // downgrade: download the resident aux trace — and - // the main LDE if this table was device-only — and - // continue fully host-backed on the arms below. - let mut recovered = crate::gpu_lde::materialize_aux_trace_host(*trace); - // Once the aux download lands, the host aux trace is - // populated: a later failure is the main-LDE - // download's, and the error has to name that step - // instead of claiming an empty aux trace. - let aux_recovered = recovered; - if recovered && device_only { - let mut cell = main_lde_cells[idx].lock().unwrap(); - // Matched exhaustively on purpose: `MainLdeSlot` - // exists so a consumer between Round 1 and the - // fused task cannot read an empty buffer as if it - // were an LDE, and this recovery is exactly such a - // consumer. - match cell.as_mut() { - // The retained buffer is the one the fused task - // reads, so under device-only it is empty and - // has to come back off the device handle. - Some(MainLdeSlot::Retained((data, _))) => { - if data.is_empty() && trace.num_main_columns > 0 { - recovered = match ( - gpu_main_cells[idx].lock().unwrap().as_ref(), - math_cuda::device::backend(), - ) { - (Some(h), Ok(be)) => { - match crate::gpu_lde::download_main_lde_row_major::< - Field, - >( - h, &be.next_stream() - ) { - Some(v) => { - *data = v; - true - } - None => false, - } - } - _ => false, - }; - } - } - // `RecomputeLde` dropped the buffer by design: - // the fused task rebuilds the main LDE from the - // host trace, which a device decline never - // touched. There is nothing to download and - // nothing to fail — the aux recovery above is - // the whole job. - Some(MainLdeSlot::Dropped { .. }) | None => {} - } - } - if !recovered { - return Err(ProvingError::Fft( - if aux_recovered { - "resident aux LDE declined; the aux trace was recovered \ - but the main-LDE download failed" - } else { - "resident aux LDE declined and the aux-trace download \ - recovery failed" - } - .to_string(), - )); - } - eprintln!( - "[gpu] resident-aux downgrade: table={} rows={} \ - (device aux LDE declined; continuing on host)", + // The device aux LDE declined twice — before and + // after a device drain. There is no host aux trace, + // and host RAM is a cache, not a compute path: this + // is the failure to report, not a downgrade. + return Err(ProvingError::DevicePath(format!( + "table {}: resident aux LDE declined after the drain-and-retry \ + (rows={} aux_cols={} blowup={}); {}", air.name(), trace.num_rows(), - ); - device_only = false; + num_cols, + domain.blowup_factor, + crate::gpu_lde::device_path_status(), + ))); } // Fused GPU path (cuda only): row-major ext3 NTT — single @@ -4143,6 +4127,7 @@ pub trait IsStarkProver< FieldExtension, H::Batched, >( + air.name(), trace_slice, n, num_cols, @@ -4348,17 +4333,31 @@ pub trait IsStarkProver< // shared transcript is untouched past this point (each fork is // per-table), so any order is sound; proofs are drained in index order. #[cfg(not(feature = "debug-checks"))] - let table_results = run_admitted(&peak_order, &peak_estimates, &vram_gate, k, |idx| { - let (commitment, lde) = aux_stage(idx)?; - rounds_stage(idx, commitment, lde) - }); + let table_results = run_admitted( + &peak_order, + &peak_estimates, + &vram_gate, + k, + |idx| table_names[idx].clone(), + |idx| { + let (commitment, lde) = aux_stage(idx)?; + rounds_stage(idx, commitment, lde) + }, + ); // debug-checks needs every table's commitments and traces between the // aux and rounds stages (cross-table bus balance), so it splits the // fused chain into two admitted passes around the check. #[cfg(feature = "debug-checks")] let table_results = { - let aux_outs = run_admitted(&peak_order, &peak_estimates, &vram_gate, k, aux_stage); + let aux_outs = run_admitted( + &peak_order, + &peak_estimates, + &vram_gate, + k, + |idx| table_names[idx].clone(), + aux_stage, + ); let mut commitments = Vec::with_capacity(num_airs); let mut ldes = Vec::with_capacity(num_airs); for out in aux_outs { @@ -4380,10 +4379,17 @@ pub trait IsStarkProver< .zip(ldes) .map(|p| std::sync::Mutex::new(Some(p))) .collect(); - run_admitted(&peak_order, &peak_estimates, &vram_gate, k, |idx| { - let (c, l) = staged[idx].lock().unwrap().take().unwrap(); - rounds_stage(idx, c, l) - }) + run_admitted( + &peak_order, + &peak_estimates, + &vram_gate, + k, + |idx| table_names[idx].clone(), + |idx| { + let (c, l) = staged[idx].lock().unwrap().take().unwrap(); + rounds_stage(idx, c, l) + }, + ) }; let mut proofs = Vec::with_capacity(num_airs); diff --git a/prover/src/lfm/chunking.rs b/prover/src/lfm/chunking.rs index d43db1faa..ffbb9b507 100644 --- a/prover/src/lfm/chunking.rs +++ b/prover/src/lfm/chunking.rs @@ -338,18 +338,19 @@ pub const BALU_TARGET_CHUNK_ROWS_LOG2: u32 = 22; /// | 2^27 | 28.0 GiB | 14.0 GiB | 8.0 GiB | 50 GiB | /// | 2^22 | 0.875 GiB | 0.44 GiB | 0.25 GiB | 1.6 GiB| /// -/// and rounds 2–4 add the aux LDE (2 ext3 columns, `2n·2·24`) and its tree, -/// the two composition parts (`2·2n·24`) and their tree, and the DEEP -/// codeword (`2n·24`): a whole-prove set of ~96 GiB at `2^27` against a -/// 32 GiB card, ~3 GiB per chunk at `2^22` (`the_balu_chunk_sizing_is_the_doc` -/// pins the arithmetic). `2^22` is the height at which eight chunks prove -/// concurrently inside a 25.6 GiB admission budget, which is why it is the -/// target: `2^24` chunks (~12 GiB each) would hold the concurrency at two, and -/// `2^20` chunks would quadruple the per-chunk FRI and query overhead the -/// verifier pays for nothing. That overhead — one FRI commit and one set of -/// openings PER CHUNK — is the counter-pressure against smaller chunks, and -/// the reason the default stays one table until the aggregator is emitted -/// with the knob set. +/// and rounds 2–4 add the aux LDE and its tree, the resident aux trace, `H` +/// and the two composition parts with their tree, the inverted denominators, +/// the DEEP codeword and the FRI chain — the scheduler's throttle model, +/// `stark::device_set::table_device_set`, which +/// `the_balu_chunk_sizing_is_the_doc` reads rather than restates: ~150 GiB at +/// `2^27` against a 32 GiB card, ~4.9 GiB per chunk at `2^22`, so five chunks +/// prove concurrently inside a 25.6 GiB admission budget. `2^22` is the +/// target because `2^24` chunks (~19 GiB each) would serialise the chunks +/// outright, and `2^20` chunks would quadruple the per-chunk FRI and query +/// overhead the verifier pays for nothing. That overhead — one FRI commit and +/// one set of openings PER CHUNK — is the counter-pressure against smaller +/// chunks, and the reason the default stays one table until the aggregator is +/// emitted with the knob set. /// /// # Why row chunking, not column streaming /// @@ -384,8 +385,8 @@ pub const BALU_TARGET_CHUNK_ROWS_LOG2: u32 = 22; /// module. /// /// `LFM_LANES` (4 value + 12 preprocessed, `2^24` rows in the 110-query wrap) -/// is the next chip of this shape; its whole-prove set at `2^24` is ~14 GiB, -/// one doubling from needing the same arm. +/// is the next chip of this shape; its device set at `2^24` is ~23 GiB under +/// the same model — it proves alone — one doubling from needing the same arm. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub struct BaluChunking { ops_per_chunk: usize, @@ -674,36 +675,43 @@ mod tests { } } - /// The device-set arithmetic the `BaluChunking` doc tabulates: one table at - /// `2^27` does not fit a 32 GiB card; a `2^22` chunk's whole-prove set is - /// ~3 GiB, so eight prove concurrently inside the 25.6 GiB budget. Columns: - /// 14 base (4 value + 10 preprocessed), 2 ext3 aux, 2 ext3 composition - /// parts, one ext3 DEEP codeword, blowup 2. + /// The device-set arithmetic the `BaluChunking` doc tabulates, read from + /// the scheduler's own model (`stark::device_set`) so the doc cannot drift + /// from what the throttle admits: one table at `2^27` does not fit a + /// 32 GiB card; a `2^22` chunk is ~4.9 GiB, so five prove concurrently + /// inside the 25.6 GiB budget and a `2^24` chunk proves alone. Columns: + /// 14 base (4 value + 10 preprocessed), 2 ext3 aux, 2 composition parts, + /// 2 OOD points, blowup 2. #[test] fn the_balu_chunk_sizing_is_the_doc() { + use stark::device_set::{TableShape, table_device_set}; const GIB: u64 = 1 << 30; - let whole_prove_set = |n: u64| -> u64 { - let lde = 2 * n; - let tree = (lde - 1) * 32; - let main_lde = lde * 14 * 8; - let snapshot = n * 14 * 8; - let aux_lde = lde * 2 * 24; - let parts = lde * 2 * 24; - let deep = lde * 24; - main_lde + snapshot + tree + aux_lde + tree + parts + tree + deep + const BUDGET: u64 = 32 * GIB / 5 * 4; + let balu = |n: usize| { + table_device_set(TableShape { + n, + blowup: 2, + main_cols: 14, + aux_cols: 2, + num_parts: 2, + num_eval_points: 2, + }) }; - let one_table = whole_prove_set(1 << 27); - assert!(one_table > 95 * GIB && one_table < 97 * GIB, "{one_table}"); - let r1_only = (1u64 << 28) * 14 * 8 + (1u64 << 27) * 14 * 8 + ((1u64 << 28) - 1) * 32; - assert!(r1_only > 49 * GIB && r1_only < 51 * GIB, "{r1_only}"); - let chunk = whole_prove_set(1 << 22); - assert!(chunk < 3 * GIB + GIB / 16, "{chunk}"); + let one_table = balu(1 << 27); + assert!(one_table.main.total() > 50 * GIB && one_table.main.total() < 52 * GIB); assert!( - 8 * chunk <= 32 * GIB / 5 * 4, - "eight chunks must fit the budget" + one_table.total() > 148 * GIB && one_table.total() < 152 * GIB, + "{one_table:?}" ); - let big_chunk = whole_prove_set(1 << 24); - assert!(2 * big_chunk <= 32 * GIB / 5 * 4 && 3 * big_chunk > 32 * GIB / 5 * 4); + let chunk = balu(1 << 22).total(); + assert!( + chunk > 4 * GIB + GIB / 2 && chunk < 5 * GIB + GIB / 2, + "{chunk}" + ); + assert!(5 * chunk <= BUDGET && 6 * chunk > BUDGET); + assert_eq!(BaluChunking::target().chunk_count(1 << 27), 32); + let big_chunk = balu(1 << 24).total(); + assert!(big_chunk <= BUDGET && 2 * big_chunk > BUDGET, "{big_chunk}"); } /// `split`, `chunk_count` and `chunk_range` are one rule seen three times From 6a2fe713e0d47730256acdc3da9dbbbf0c80270a Mon Sep 17 00:00:00 2001 From: MauroFab Date: Tue, 8 Sep 2026 01:21:48 -0300 Subject: [PATCH 2/4] fix(prover): walk the tables by their host transient, not by the device-set estimate MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit At TABLE_PARALLELISM=1 the heaviest-first walk order moves nothing on the device (the gate never blocks) and everything about the host allocator's layout. Measured on the q=20 wrap (2^22, blowup 4, 15 tables, RTX 5090): walking by the device-set model peaked at 51.0 GiB max RSS, walking by the retired 2·lde·(8·main + 24·aux) + 256·lde key at 45.2 GiB, prove flat (137.8 vs 137.9 s). The device-set model's 256 MiB scratch floor tied six small tables and moved LFM_HINT two places. The gate keeps the device-set estimates. Both walks now sort by `device_set::host_transient_bytes` — the host bytes the fused task allocates under residency recompute: the recomputed main LDE, the aux copy, the parts, the DEEP codeword and the FRI chain, lde · (8·main + 24·aux + 24·(parts + 3)) — the retired key's per-row weights with 120 in place of 128 in the constant and no floor. The prover prints the walk it took once per prove. --- crypto/stark/src/device_set.rs | 54 +++++++++++++++++++++ crypto/stark/src/prover.rs | 86 ++++++++++++++++++++++------------ 2 files changed, 109 insertions(+), 31 deletions(-) diff --git a/crypto/stark/src/device_set.rs b/crypto/stark/src/device_set.rs index 0f39591f9..fc5f050de 100644 --- a/crypto/stark/src/device_set.rs +++ b/crypto/stark/src/device_set.rs @@ -196,6 +196,30 @@ pub fn table_device_set(shape: TableShape) -> TableDeviceSet { } } +/// The HOST bytes a table's fused task allocates — the walk-order key. +/// +/// At `TABLE_PARALLELISM=1` the order the tables are walked in changes nothing +/// on the device (the gate never blocks) and everything about the host +/// allocator's layout, and that layout is a measured 5.9 GiB of peak at the +/// wrap (q=20, 2^22, blowup 4, 15 tables, RTX 5090 box, 2026-09-07): walking +/// by the device-set model peaked at 51.0 GiB max RSS, walking by the retired +/// `2·lde·(8·main + 24·aux) + 256·lde` key at 45.2 GiB, prove 137.8 vs +/// 137.9 s. So the walk is a HOST policy, sorted by the host transient of the +/// fused task under residency recompute: the recomputed main LDE (`8·main` +/// per LDE row), the aux LDE's host copy (`24·aux`), the parts (`24·parts`), +/// the DEEP codeword (24) and the FRI chain (48 — two codewords bound the +/// geometric sum): `lde · (8·main + 24·aux + 24·(parts + 3))`. Per LDE row +/// that is the retired key's weights with 120 in place of 128 in the +/// constant, so it walks the tables the retired key did; both phases sort by +/// it so the two walks agree, and the prover prints the walk it took. +pub fn host_transient_bytes(shape: TableShape) -> u64 { + let lde = (shape.n as u64).saturating_mul(shape.blowup as u64); + let per_row = base_bytes(1, shape.main_cols as u64) + .saturating_add(ext3_bytes(1, shape.aux_cols as u64)) + .saturating_add(ext3_bytes(1, shape.num_parts as u64 + 3)); + lde.saturating_mul(per_row) +} + /// What the admission predicate decided for one dispatch. #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub enum Admission { @@ -326,6 +350,36 @@ mod tests { assert!(6 * set.total() > CARD_32_GIB_BUDGET); } + /// The walk key is `lde · (8·main + 24·aux + 24·(parts + 3))`: the retired + /// key's per-row weights (`16·main + 48·aux + 256`, halved) with 120 in + /// place of 128 in the constant, and no scratch floor to tie the small + /// tables. + #[test] + fn the_walk_key_is_the_host_transient() { + let shape = TableShape { + n: 1 << 20, + blowup: 4, + main_cols: 25, + aux_cols: 3, + num_parts: 2, + num_eval_points: 2, + }; + assert_eq!( + host_transient_bytes(shape), + (1u64 << 22) * (8 * 25 + 24 * 3 + 120) + ); + let tiny = TableShape { + n: 4, + blowup: 4, + main_cols: 1, + aux_cols: 1, + num_parts: 2, + num_eval_points: 2, + }; + assert_eq!(host_transient_bytes(tiny), 16 * (8 + 24 + 120)); + assert!(host_transient_bytes(tiny) < host_transient_bytes(shape)); + } + /// A table without aux or parts (d=1 with no lookups) counts only what it /// allocates. #[test] diff --git a/crypto/stark/src/prover.rs b/crypto/stark/src/prover.rs index 96c4ac445..5e72c2609 100644 --- a/crypto/stark/src/prover.rs +++ b/crypto/stark/src/prover.rs @@ -775,9 +775,10 @@ impl Drop for VramPermit<'_> { /// Run `task` once per table index on `workers` OS driver threads, admitting /// each index through `gate` with its estimated bytes. `order` fixes the -/// 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. +/// start order — the caller's walk: largest fused-phase host transient first +/// (`device_set::host_transient_bytes`), so the long pole starts early and +/// small tables fill around it, and so the host allocator sees the layout the +/// measured-good walk produces. Returns one slot per original index. fn run_admitted( order: &[usize], estimates: &[u64], @@ -3707,24 +3708,35 @@ pub trait IsStarkProver< let vram_gate = VramGate::new(vram_budget); - // R1 main commit: the fused commit's device set — one LDE buffer, the - // trace snapshot, the tree and the scratch — the same model the - // dispatch layer admits the commit against (`crate::device_set`). - let main_estimates: Vec = air_trace_pairs + // The shapes the AIR and the domain fix, read once: the device-set + // estimates the gate admits against and the host-transient key the + // walks are sorted by both derive from them (`crate::device_set`). + let table_shapes: Vec = air_trace_pairs .iter() .enumerate() - .map(|(idx, (_, trace, _))| { + .map(|(idx, (air, trace, _))| { let domain = &domains[idx]; - crate::device_set::commit_device_set( - domain.interpolation_domain_size, - trace.num_main_columns, - domain.blowup_factor, - true, - ) - .total() + let n = domain.interpolation_domain_size; + let (_, aux_cols) = air.trace_layout(); + crate::device_set::TableShape { + n, + blowup: domain.blowup_factor, + main_cols: trace.num_main_columns, + aux_cols, + num_parts: air.composition_poly_degree_bound(n) / n, + num_eval_points: air.context().transition_offsets.len() * air.step_size(), + } }) .collect(); + // R1 main commit: the fused commit's device set — one LDE buffer, the + // trace snapshot, the tree and the scratch — the same model the + // dispatch layer admits the commit against. + let main_estimates: Vec = table_shapes + .iter() + .map(|s| crate::device_set::commit_device_set(s.n, s.main_cols, s.blowup, true).total()) + .collect(); + // The AIR names, for the driver threads' panic payloads: a device abort // names its stage and shape, the driver adds which table. let table_names: Vec = air_trace_pairs @@ -3732,6 +3744,29 @@ pub trait IsStarkProver< .map(|(air, _, _)| air.name().to_string()) .collect(); + // The walk order, for BOTH phases: largest fused-phase host transient + // first (`device_set::host_transient_bytes` — the measurement behind + // the choice is on its doc). Deliberately not the device-set estimate + // the gate uses: at `TABLE_PARALLELISM=1` the order moves nothing on + // the device and 5.9 GiB of host peak. + let walk_keys: Vec = table_shapes + .iter() + .map(|&s| crate::device_set::host_transient_bytes(s)) + .collect(); + let walk_order = heaviest_first(&walk_keys); + eprintln!( + "[prover] table walk (fused-phase host transient, largest first): {}", + walk_order + .iter() + .map(|&i| format!( + "{}={:.2}GiB", + table_names[i], + walk_keys[i] as f64 / (1u64 << 30) as f64 + )) + .collect::>() + .join(" ") + ); + // Spill main traces to mmap before Round 1 LDE. #[cfg(feature = "disk-spill")] if storage_mode == StorageMode::Disk { @@ -3779,7 +3814,7 @@ pub trait IsStarkProver< // sequentially below once every commit completed — the one ordering // Fiat-Shamir requires before sampling the shared challenges. let main_results = run_admitted( - &heaviest_first(&main_estimates), + &walk_order, &main_estimates, &vram_gate, k, @@ -3932,20 +3967,7 @@ pub trait IsStarkProver< let peak_estimates: Vec = air_trace_pairs .iter() .enumerate() - .map(|(idx, (air, trace, _))| { - let domain = &domains[idx]; - let n = domain.interpolation_domain_size; - let (_, aux_cols) = air.trace_layout(); - crate::device_set::table_device_set(crate::device_set::TableShape { - n, - blowup: domain.blowup_factor, - main_cols: trace.num_main_columns, - aux_cols, - num_parts: air.composition_poly_degree_bound(n) / n, - num_eval_points: air.context().transition_offsets.len() * air.step_size(), - }) - .total() - }) + .map(|(idx, _)| crate::device_set::table_device_set(table_shapes[idx]).total()) .collect(); // Per-table slots for the fused chain: each driver takes or locks only @@ -4326,7 +4348,9 @@ pub trait IsStarkProver< #[cfg(feature = "instruments")] let __sp = crate::instruments::span("rounds_2to4"); - let peak_order = heaviest_first(&peak_estimates); + // Same walk as R1 (see `walk_order`): the estimates below feed the gate, + // the order is the host policy. + let peak_order = walk_order.clone(); // One fused task per table: while a heavy table works through a // host-bound stretch, the others' GPU stages fill the device. The From a63599061a1a3f82d286fdc15413844d663a0567 Mon Sep 17 00:00:00 2001 From: MauroFab Date: Tue, 8 Sep 2026 02:12:37 -0300 Subject: [PATCH 3/4] fix(prover): keep the table walk exactly as it was; the device-set model gates, it does not schedule MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The device-set model this branch introduces is a size model: it says what a stage puts on the card, so the gate can decide whether it fits. It is not a schedule, and sorting the table walk by it cost 5.9 GiB of host peak. The walk and the gate are now two different functions, and only the gate reads the model. The walk goes back to exactly the key it used before this branch: `2·lde·(8·main + 24·aux) + 256·lde`, aux width zero for the R1 main-commit walk and the AIR's aux width for the fused rounds walk, two separate walks as there have always been. The arithmetic is restored verbatim, so the order is byte-for-byte the one that measured well; the function is renamed `table_walk_weight` and its constants renamed with it, because the numbers are kept for the order they produce and are no longer a claim about bytes (the factor of two assumed the second LDE buffer #956's in-place transpose removed). Measured on the q=20 wrap (2^22, blowup 4, RTX 5090, TABLE_PARALLELISM=1), `/usr/bin/time -v` max RSS, prove time and proof bytes identical throughout: this weight, before the device-set model 47,307,284 kB = 45.1 GiB this weight, under the device-set gate 47,365,192 kB = 45.2 GiB the device-set model's own order 53,472,980 kB = 51.0 GiB a fused-phase host-transient order 53,453,344 kB = 51.0 GiB The last row is why this is a restoration and not a re-derivation: an order justified by a truer quantity is still the wrong order. At q=41 the effect is absent, so the mechanism — why reordering only the small tables moves host peak when one table is resident at a time — is open and tracked separately. The prover prints the walk each phase took, one line per phase, so a run that moves host peak can say which order it ran. --- crypto/stark/src/device_set.rs | 59 ++-------- crypto/stark/src/prover.rs | 189 +++++++++++++++++++++++++++------ 2 files changed, 162 insertions(+), 86 deletions(-) diff --git a/crypto/stark/src/device_set.rs b/crypto/stark/src/device_set.rs index fc5f050de..81771fbdd 100644 --- a/crypto/stark/src/device_set.rs +++ b/crypto/stark/src/device_set.rs @@ -196,29 +196,12 @@ pub fn table_device_set(shape: TableShape) -> TableDeviceSet { } } -/// The HOST bytes a table's fused task allocates — the walk-order key. -/// -/// At `TABLE_PARALLELISM=1` the order the tables are walked in changes nothing -/// on the device (the gate never blocks) and everything about the host -/// allocator's layout, and that layout is a measured 5.9 GiB of peak at the -/// wrap (q=20, 2^22, blowup 4, 15 tables, RTX 5090 box, 2026-09-07): walking -/// by the device-set model peaked at 51.0 GiB max RSS, walking by the retired -/// `2·lde·(8·main + 24·aux) + 256·lde` key at 45.2 GiB, prove 137.8 vs -/// 137.9 s. So the walk is a HOST policy, sorted by the host transient of the -/// fused task under residency recompute: the recomputed main LDE (`8·main` -/// per LDE row), the aux LDE's host copy (`24·aux`), the parts (`24·parts`), -/// the DEEP codeword (24) and the FRI chain (48 — two codewords bound the -/// geometric sum): `lde · (8·main + 24·aux + 24·(parts + 3))`. Per LDE row -/// that is the retired key's weights with 120 in place of 128 in the -/// constant, so it walks the tables the retired key did; both phases sort by -/// it so the two walks agree, and the prover prints the walk it took. -pub fn host_transient_bytes(shape: TableShape) -> u64 { - let lde = (shape.n as u64).saturating_mul(shape.blowup as u64); - let per_row = base_bytes(1, shape.main_cols as u64) - .saturating_add(ext3_bytes(1, shape.aux_cols as u64)) - .saturating_add(ext3_bytes(1, shape.num_parts as u64 + 3)); - lde.saturating_mul(per_row) -} +// This module is the SIZE model and nothing else: what a stage puts on the +// card, so the gate can decide whether it fits. It is deliberately not the +// prover's table walk. The walk is a scheduling policy, keyed on a weight that +// is kept for the order it produces rather than for any byte it names, and it +// lives with the scheduler in `prover::table_walk_weight`. Sorting the walk by +// this model instead cost a measured 5.9 GiB of host peak at the q=20 wrap. /// What the admission predicate decided for one dispatch. #[derive(Clone, Copy, Debug, PartialEq, Eq)] @@ -350,36 +333,6 @@ mod tests { assert!(6 * set.total() > CARD_32_GIB_BUDGET); } - /// The walk key is `lde · (8·main + 24·aux + 24·(parts + 3))`: the retired - /// key's per-row weights (`16·main + 48·aux + 256`, halved) with 120 in - /// place of 128 in the constant, and no scratch floor to tie the small - /// tables. - #[test] - fn the_walk_key_is_the_host_transient() { - let shape = TableShape { - n: 1 << 20, - blowup: 4, - main_cols: 25, - aux_cols: 3, - num_parts: 2, - num_eval_points: 2, - }; - assert_eq!( - host_transient_bytes(shape), - (1u64 << 22) * (8 * 25 + 24 * 3 + 120) - ); - let tiny = TableShape { - n: 4, - blowup: 4, - main_cols: 1, - aux_cols: 1, - num_parts: 2, - num_eval_points: 2, - }; - assert_eq!(host_transient_bytes(tiny), 16 * (8 + 24 + 120)); - assert!(host_transient_bytes(tiny) < host_transient_bytes(shape)); - } - /// A table without aux or parts (d=1 with no lookups) counts only what it /// allocates. #[test] diff --git a/crypto/stark/src/prover.rs b/crypto/stark/src/prover.rs index 5e72c2609..1560be130 100644 --- a/crypto/stark/src/prover.rs +++ b/crypto/stark/src/prover.rs @@ -774,11 +774,12 @@ impl Drop for VramPermit<'_> { } /// Run `task` once per table index on `workers` OS driver threads, admitting -/// each index through `gate` with its estimated bytes. `order` fixes the -/// start order — the caller's walk: largest fused-phase host transient first -/// (`device_set::host_transient_bytes`), so the long pole starts early and -/// small tables fill around it, and so the host allocator sees the layout the -/// measured-good walk produces. Returns one slot per original index. +/// each index through `gate` with its estimated bytes. The two array arguments +/// are deliberately independent: `estimates` is what the gate spends (the +/// device set, `crate::device_set`), while `order` is the caller's walk — a +/// scheduling policy keyed on [`table_walk_weight`], heaviest first, so the +/// long pole starts early and small tables fill around it. Returns one slot +/// per original index. fn run_admitted( order: &[usize], estimates: &[u64], @@ -887,13 +888,80 @@ fn name_panic_payload( } } -/// Table indices sorted heaviest-first by estimate. -fn heaviest_first(estimates: &[u64]) -> Vec { - let mut order: Vec = (0..estimates.len()).collect(); - order.sort_by_key(|&i| std::cmp::Reverse(estimates[i])); +/// The sort weight for the table walk. A SCHEDULING policy, not a size model — +/// see [`crate::device_set`] for the bytes anything is admitted against. +/// +/// The walk is the order the per-table drivers *start* tables in. At +/// `TABLE_PARALLELISM=1` it cannot change what is resident on the device (one +/// table at a time, the gate never blocks), but it does fix the order the host +/// allocator sees the per-table arenas in, and that is worth 5.9 GiB of host +/// peak. Measured on the q=20 wrap (2^22, blowup 4, RTX 5090 box, 2026-09-07), +/// `/usr/bin/time -v` max RSS, with prove time and proof bytes identical across +/// all four runs: +/// +/// | walk | max RSS | +/// |------|---------| +/// | this weight, before the device-set model landed | 47,307,284 kB = 45.1 GiB | +/// | this weight, restored under the device-set gate | 47,365,192 kB = 45.2 GiB | +/// | the device-set model's own order | 53,472,980 kB = 51.0 GiB | +/// | a fused-phase host-transient order | 53,453,344 kB = 51.0 GiB | +/// +/// At q=41 this weight and the device-set order both measure 98.6 GiB, so the +/// effect is shape-dependent and the mechanism — why reordering only the small +/// tables moves host peak at all when one table is resident at a time — is +/// still open. Until it is understood, this order is kept because it is the one +/// that measures well. That is the whole justification, and it is why the +/// weight below lives here with the scheduler and not in the device-set model. +/// +/// Its arithmetic is inherited verbatim from the VRAM estimate the prover +/// sorted by before the device-set model, and it is deliberately NOT re-read as +/// a byte count: the factor of two assumed the second LDE buffer that #956's +/// in-place transpose removed, and the flat 256 B per LDE row stands in for a +/// tree whose real width depends on the digest. Changing these numbers changes +/// the schedule, so any change needs a wrap measurement, not an argument about +/// bytes. +/// +/// Pass `aux_cols == 0` for the R1 main-commit walk and the AIR's aux width for +/// the fused rounds walk: the two phases weigh the tables differently, and they +/// always have. +fn table_walk_weight(main_cols: usize, aux_cols: usize, lde_size: usize) -> u64 { + const BASE_WEIGHT: u64 = 8; + const EXT3_WEIGHT: u64 = 24; + const WIDTH_WEIGHT: u64 = 2; + const PER_LDE_ROW_WEIGHT: u64 = 256; + let lde = lde_size as u64; + let per_row = (main_cols as u64).saturating_mul(BASE_WEIGHT) + + (aux_cols as u64).saturating_mul(EXT3_WEIGHT); + let width_term = lde.saturating_mul(per_row).saturating_mul(WIDTH_WEIGHT); + let row_term = lde.saturating_mul(PER_LDE_ROW_WEIGHT); + width_term.saturating_add(row_term) +} + +/// Table indices sorted heaviest-first by weight. `sort_by_key` is stable, so +/// tables that weigh the same keep their registry order. +fn heaviest_first(weights: &[u64]) -> Vec { + let mut order: Vec = (0..weights.len()).collect(); + order.sort_by_key(|&i| std::cmp::Reverse(weights[i])); order } +/// One line naming the walk a phase took, heaviest first. Printed once per +/// phase per prove: the walk is a measured choice (see [`table_walk_weight`]), +/// so a run that moves host peak has to be able to say which order it took. +fn describe_walk(order: &[usize], weights: &[u64], names: &[String]) -> String { + order + .iter() + .map(|&i| { + format!( + "{}={:.2}GiB", + names[i], + weights[i] as f64 / (1u64 << 30) as f64 + ) + }) + .collect::>() + .join(" ") +} + /// A container for the results of the second round of the STARK Prove protocol. pub(crate) struct Round2 where @@ -3709,8 +3777,9 @@ pub trait IsStarkProver< let vram_gate = VramGate::new(vram_budget); // The shapes the AIR and the domain fix, read once: the device-set - // estimates the gate admits against and the host-transient key the - // walks are sorted by both derive from them (`crate::device_set`). + // estimates the gate admits against derive from them + // (`crate::device_set`). The walk order does NOT — see + // `table_walk_weight`. let table_shapes: Vec = air_trace_pairs .iter() .enumerate() @@ -3744,27 +3813,20 @@ pub trait IsStarkProver< .map(|(air, _, _)| air.name().to_string()) .collect(); - // The walk order, for BOTH phases: largest fused-phase host transient - // first (`device_set::host_transient_bytes` — the measurement behind - // the choice is on its doc). Deliberately not the device-set estimate - // the gate uses: at `TABLE_PARALLELISM=1` the order moves nothing on - // the device and 5.9 GiB of host peak. - let walk_keys: Vec = table_shapes + // The R1 walk: the main commit's weight, aux width zero because the aux + // columns are not resident yet in this phase. Keyed on + // `table_walk_weight`, NOT on `main_estimates` — the estimates above + // are what the gate spends, this is the schedule, and the two are no + // longer the same function. The weight's doc carries the measurement + // that makes this the order rather than any other. + let main_walk_weights: Vec = table_shapes .iter() - .map(|&s| crate::device_set::host_transient_bytes(s)) + .map(|s| table_walk_weight(s.main_cols, 0, s.n * s.blowup)) .collect(); - let walk_order = heaviest_first(&walk_keys); + let main_walk_order = heaviest_first(&main_walk_weights); eprintln!( - "[prover] table walk (fused-phase host transient, largest first): {}", - walk_order - .iter() - .map(|&i| format!( - "{}={:.2}GiB", - table_names[i], - walk_keys[i] as f64 / (1u64 << 30) as f64 - )) - .collect::>() - .join(" ") + "[prover] table walk R1 (walk weight, largest first): {}", + describe_walk(&main_walk_order, &main_walk_weights, &table_names) ); // Spill main traces to mmap before Round 1 LDE. @@ -3814,7 +3876,7 @@ pub trait IsStarkProver< // sequentially below once every commit completed — the one ordering // Fiat-Shamir requires before sampling the shared challenges. let main_results = run_admitted( - &walk_order, + &main_walk_order, &main_estimates, &vram_gate, k, @@ -3970,6 +4032,15 @@ pub trait IsStarkProver< .map(|(idx, _)| crate::device_set::table_device_set(table_shapes[idx]).total()) .collect(); + // The fused phase's own walk, separate from R1's because the aux + // columns are resident here and so carry weight. Keyed on + // `table_walk_weight`, not on `peak_estimates`: the estimates feed the + // gate, the weight fixes the schedule. + let peak_walk_weights: Vec = table_shapes + .iter() + .map(|s| table_walk_weight(s.main_cols, s.aux_cols, s.n * s.blowup)) + .collect(); + // Per-table slots for the fused chain: each driver takes or locks only // its own index, so every mutex is uncontended by construction. let pair_cells: Vec>> = @@ -4348,9 +4419,11 @@ pub trait IsStarkProver< #[cfg(feature = "instruments")] let __sp = crate::instruments::span("rounds_2to4"); - // Same walk as R1 (see `walk_order`): the estimates below feed the gate, - // the order is the host policy. - let peak_order = walk_order.clone(); + let peak_order = heaviest_first(&peak_walk_weights); + eprintln!( + "[prover] table walk rounds 2-4 (walk weight, largest first): {}", + describe_walk(&peak_order, &peak_walk_weights, &table_names) + ); // One fused task per table: while a heavy table works through a // host-bound stretch, the others' GPU stages fill the device. The @@ -5148,3 +5221,53 @@ fn print_bus_balance_report( } } } + +#[cfg(test)] +mod walk_tests { + use super::{heaviest_first, table_walk_weight}; + + /// The weight is `2·lde·(8·main + 24·aux) + 256·lde`. These constants are a + /// schedule, not a size, so this test exists to make a change to them a + /// deliberate act that comes with a wrap measurement. + #[test] + fn the_walk_weight_is_pinned() { + let lde = 1usize << 22; + assert_eq!( + table_walk_weight(25, 3, lde), + (lde as u64) * (2 * (8 * 25 + 24 * 3) + 256) + ); + assert_eq!(table_walk_weight(1, 1, 16), 16 * (2 * (8 + 24) + 256)); + // No floor: the per-LDE-row term scales with height, so two narrow + // tables of different heights never tie. The device-set model's 256 MiB + // scratch floor tied six small tables at the wrap, and that tie is + // where the two walks first diverge. + assert!(table_walk_weight(4, 1, 1 << 16) > table_walk_weight(4, 1, 1 << 15)); + } + + /// The two phases weigh differently: R1 passes `aux_cols == 0` because the + /// aux columns are not resident yet, the fused phase passes the AIR's aux + /// width. A table that is narrow in main and wide in aux therefore moves + /// between the two walks. + #[test] + fn the_two_phases_can_walk_differently() { + let lde = 1usize << 20; + // (main, aux) per table: the second is aux-heavy, the first main-heavy. + let main_walk: Vec = [(60usize, 1usize), (10, 30)] + .iter() + .map(|&(m, _)| table_walk_weight(m, 0, lde)) + .collect(); + let fused_walk: Vec = [(60usize, 1usize), (10, 30)] + .iter() + .map(|&(m, a)| table_walk_weight(m, a, lde)) + .collect(); + assert_eq!(heaviest_first(&main_walk), vec![0, 1]); + assert_eq!(heaviest_first(&fused_walk), vec![1, 0]); + } + + /// Ties keep registry order, so equal-weight tables walk in the order the + /// registry lists them and the walk is reproducible run to run. + #[test] + fn ties_keep_registry_order() { + assert_eq!(heaviest_first(&[5, 9, 5, 9]), vec![1, 3, 0, 2]); + } +} From f1aff94f65697f048e508d1fa3e67fe45baf5835 Mon Sep 17 00:00:00 2001 From: MauroFab Date: Tue, 8 Sep 2026 03:10:04 -0300 Subject: [PATCH 4/4] =?UTF-8?q?docs(prover):=20the=20walk=20order=20is=20a?= =?UTF-8?q?=20correlate,=20not=20a=20cause=20=E2=80=94=20name=20the=20mech?= =?UTF-8?q?anism?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The weight's doc said the mechanism was still open. It is not, and the measurement that closed it is this branch's own. Same walk order, two builds: 45.2 GiB in the diagnostic build, 51.0 GiB here. So the order cannot be the cause. The cause is glibc arena retention — the harness installs no global allocator, and a freed multi-gibibyte buffer goes back to the OS only when the arena top can be trimmed, which depends on what sits above it. The walk order is one input to that layout and the allocations around it are another. Forcing every allocation of a megabyte or more to be mapped and unmapped directly collapses the 5.9 GiB spread to 1.6 MB (46,053,164 kB here against 46,054,788 kB at the pre-#964 tip), so none of it was working set. That makes the allocator the durable fix rather than the walk, and the doc now says so, with the caveat that leaving the retentive state costs ~2.5% prove time. Comments only; no code changed. --- crypto/stark/src/prover.rs | 50 +++++++++++++++++++++++++------------- 1 file changed, 33 insertions(+), 17 deletions(-) diff --git a/crypto/stark/src/prover.rs b/crypto/stark/src/prover.rs index 1560be130..ad6ed4d64 100644 --- a/crypto/stark/src/prover.rs +++ b/crypto/stark/src/prover.rs @@ -893,25 +893,41 @@ fn name_panic_payload( /// /// The walk is the order the per-table drivers *start* tables in. At /// `TABLE_PARALLELISM=1` it cannot change what is resident on the device (one -/// table at a time, the gate never blocks), but it does fix the order the host -/// allocator sees the per-table arenas in, and that is worth 5.9 GiB of host -/// peak. Measured on the q=20 wrap (2^22, blowup 4, RTX 5090 box, 2026-09-07), -/// `/usr/bin/time -v` max RSS, with prove time and proof bytes identical across -/// all four runs: +/// table at a time, the gate never blocks). What it does change is the order +/// the host allocator sees the per-table arenas in. Measured on the q=20 wrap +/// (2^22, blowup 4, RTX 5090 box, 2026-09-07/08), `/usr/bin/time -v` max RSS, +/// proof bytes identical on every row: /// -/// | walk | max RSS | -/// |------|---------| -/// | this weight, before the device-set model landed | 47,307,284 kB = 45.1 GiB | -/// | this weight, restored under the device-set gate | 47,365,192 kB = 45.2 GiB | -/// | the device-set model's own order | 53,472,980 kB = 51.0 GiB | -/// | a fused-phase host-transient order | 53,453,344 kB = 51.0 GiB | +/// | walk | build | `MALLOC_MMAP_THRESHOLD_` | max RSS | +/// |------|-------|--------------------------|---------| +/// | this weight | pre-#964 tip | default | 47,307,284 kB = 45.1 GiB | +/// | this weight | diagnostic, device-set gate | default | 47,365,192 kB = 45.2 GiB | +/// | this weight | this one | default | 53,456,648 kB = 51.0 GiB | +/// | the device-set model's order | this one | default | 53,472,980 kB = 51.0 GiB | +/// | a fused-phase host-transient order | this one | default | 53,453,344 kB = 51.0 GiB | +/// | this weight | this one | 1 MiB | 46,053,164 kB = 43.9 GiB | +/// | this weight | pre-#964 tip | 1 MiB | 46,054,788 kB = 43.9 GiB | /// -/// At q=41 this weight and the device-set order both measure 98.6 GiB, so the -/// effect is shape-dependent and the mechanism — why reordering only the small -/// tables moves host peak at all when one table is resident at a time — is -/// still open. Until it is understood, this order is kept because it is the one -/// that measures well. That is the whole justification, and it is why the -/// weight below lives here with the scheduler and not in the device-set model. +/// Read the first three rows together: the SAME order measures 45.2 GiB in one +/// build and 51.0 GiB in this one. **The order is a correlate, not a cause.** +/// The mechanism is glibc arena retention — the test harness installs no +/// `#[global_allocator]`, and a freed multi-gibibyte buffer is returned to the +/// OS only when the arena top can be trimmed, which depends on what was +/// allocated above it. The walk order is one input to that layout; the +/// allocations made around it are another. The last two rows settle what is +/// being measured: forcing every allocation of a megabyte or more to be mapped +/// and unmapped directly collapses a 5.9 GiB spread to 1.6 MB, so none of it +/// was ever working set. (At q=41 the orders are indistinguishable, 98.5 +/// against 98.6 GiB.) +/// +/// ⇒ **The durable fix is the allocator, not this weight.** Setting +/// `MALLOC_MMAP_THRESHOLD_` for the harness removes the effect outright and +/// takes 1.2 GiB off the good arm too. It is not free: leaving the retentive +/// state costs ~2.5% prove time (136.4 → 140.2 s here, while a build already +/// out of it goes 140.3 → 140.1 s), so the time follows the retention state +/// rather than the order. Until that lands, this order is kept because it is +/// the one the prover had before #964 — not because reordering is a lever, as +/// nothing in this file controls the layout that decides the number. /// /// Its arithmetic is inherited verbatim from the VRAM estimate the prover /// sorted by before the device-set model, and it is deliberately NOT re-read as