diff --git a/crypto/stark/src/device_set.rs b/crypto/stark/src/device_set.rs new file mode 100644 index 000000000..81771fbdd --- /dev/null +++ b/crypto/stark/src/device_set.rs @@ -0,0 +1,407 @@ +//! 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, + } +} + +// 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)] +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..ad6ed4d64 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 @@ -785,15 +774,18 @@ 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. +/// 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], gate: &VramGate, workers: usize, + label: impl Fn(usize) -> String + Sync, task: impl Fn(usize) -> T + Sync, ) -> Vec> { let results: Vec>> = estimates @@ -849,6 +841,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,11 +870,112 @@ fn run_admitted( .collect() } -/// 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])); +/// 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, + } +} + +/// 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). 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 | 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 | +/// +/// 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 +/// 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. @@ -1247,8 +1343,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 +1383,7 @@ pub trait IsStarkProver< Field, H::Batched, >( + table, trace_slice, trace.main_rowmajor_dev(), n, @@ -1350,6 +1449,7 @@ pub trait IsStarkProver< Field, H::Batched, >( + table, trace_slice, trace.main_rowmajor_dev(), n, @@ -2041,14 +2141,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 +3792,59 @@ 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. - 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 derive from them + // (`crate::device_set`). The walk order does NOT — see + // `table_walk_weight`. + let table_shapes: 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) + .map(|(idx, (air, trace, _))| { + let domain = &domains[idx]; + 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 + .iter() + .map(|(air, _, _)| air.name().to_string()) + .collect(); + + // 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| table_walk_weight(s.main_cols, 0, s.n * s.blowup)) + .collect(); + let main_walk_order = heaviest_first(&main_walk_weights); + eprintln!( + "[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. #[cfg(feature = "disk-spill")] if storage_mode == StorageMode::Disk { @@ -3747,10 +3892,11 @@ 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), + &main_walk_order, &main_estimates, &vram_gate, k, + |idx| table_names[idx].clone(), |idx| { let (air, trace, _) = &air_trace_pairs[idx]; let domain = &domains[idx]; @@ -3766,6 +3912,7 @@ pub trait IsStarkProver< let device_only = Self::device_only_for(*air, domain); Self::commit_main_trace( + air.name(), *trace, domain, twiddles, @@ -3898,11 +4045,16 @@ pub trait IsStarkProver< let peak_estimates: Vec = air_trace_pairs .iter() .enumerate() - .map(|(idx, (air, trace, _))| { - let lde_size = domains[idx].interpolation_domain_size * domains[idx].blowup_factor; - let (_, aux_cols) = air.trace_layout(); - estimate_table_vram_bytes(trace.num_main_columns, aux_cols, lde_size) - }) + .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 @@ -3989,20 +4141,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 +4167,7 @@ pub trait IsStarkProver< FieldExtension, H::Batched, >( + air.name(), ra, domain.blowup_factor, &twiddles.coset_weights, @@ -4049,80 +4203,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 +4236,7 @@ pub trait IsStarkProver< FieldExtension, H::Batched, >( + air.name(), trace_slice, n, num_cols, @@ -4341,24 +4435,42 @@ pub trait IsStarkProver< #[cfg(feature = "instruments")] let __sp = crate::instruments::span("rounds_2to4"); - let peak_order = heaviest_first(&peak_estimates); + 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 // 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 +4492,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); @@ -5118,3 +5237,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]); + } +} 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