diff --git a/Cargo.lock b/Cargo.lock index fd763f24b..4567c3e1e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -227,7 +227,7 @@ checksum = "4ea181bf566f71cb9a5d17a59e1871af638180a18fb0035c92ae62b705207123" dependencies = [ "bitflags 1.3.2", "clap_lex 0.2.4", - "indexmap", + "indexmap 1.9.3", "textwrap", ] @@ -289,10 +289,12 @@ dependencies = [ "executor", "lambda-vm-prover", "rkyv", + "serde", "stark", "tempfile", "tikv-jemalloc-ctl", "tikv-jemallocator", + "toml", ] [[package]] @@ -569,6 +571,12 @@ dependencies = [ "log", ] +[[package]] +name = "equivalent" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "877a4ace8713b0bcf2a4e7eec82529c029f1d0619886d18145fea96c3ffe5c0f" + [[package]] name = "errno" version = "0.3.14" @@ -726,6 +734,16 @@ dependencies = [ "hashbrown 0.12.3", ] +[[package]] +name = "indexmap" +version = "2.14.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d466e9454f08e4a911e14806c24e16fba1b4c121d1ea474396f396069cf949d9" +dependencies = [ + "equivalent", + "hashbrown 0.17.1", +] + [[package]] name = "is-terminal" version = "0.4.17" @@ -1518,6 +1536,15 @@ dependencies = [ "serde_core", ] +[[package]] +name = "serde_spanned" +version = "0.6.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bf41e0cfaf7226dca15e8197172c295a782857fcb97fad1808a166870dee75a3" +dependencies = [ + "serde", +] + [[package]] name = "sha2" version = "0.10.9" @@ -1784,6 +1811,47 @@ version = "0.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1f3ccbac311fea05f86f61904b462b55fb3df8837a366dfc601a0161d0532f20" +[[package]] +name = "toml" +version = "0.8.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dc1beb996b9d83529a9e75c17a1686767d148d70663143c7854d8b4a09ced362" +dependencies = [ + "serde", + "serde_spanned", + "toml_datetime", + "toml_edit", +] + +[[package]] +name = "toml_datetime" +version = "0.6.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "22cddaf88f4fbc13c51aebbf5f8eceb5c7c5a9da2ac40a13519eb5b0a0e8f11c" +dependencies = [ + "serde", +] + +[[package]] +name = "toml_edit" +version = "0.22.27" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "41fe8c660ae4257887cf66394862d21dbca4a6ddd26f04a3560410406a2f819a" +dependencies = [ + "indexmap 2.14.0", + "serde", + "serde_spanned", + "toml_datetime", + "toml_write", + "winnow", +] + +[[package]] +name = "toml_write" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5d99f8c9a7727884afe522e9bd5edbfc91a3312b36a77b5fb8926e4c31a41801" + [[package]] name = "tracing" version = "0.1.44" @@ -2126,6 +2194,15 @@ version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec" +[[package]] +name = "winnow" +version = "0.7.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "df79d97927682d2fd8adb29682d1140b343be4ac0f08fd68b7765d9c059d3945" +dependencies = [ + "memchr", +] + [[package]] name = "wit-bindgen" version = "0.46.0" diff --git a/bin/cli/Cargo.toml b/bin/cli/Cargo.toml index b9140e34c..f5fd8a465 100644 --- a/bin/cli/Cargo.toml +++ b/bin/cli/Cargo.toml @@ -15,6 +15,8 @@ tempfile = "3" tikv-jemallocator = "0.6" tikv-jemalloc-ctl = { version = "0.6", features = ["stats"], optional = true } env_logger = "0.11" +serde = { version = "1", features = ["derive"] } +toml = "0.8" [features] jemalloc-stats = ["dep:tikv-jemalloc-ctl"] diff --git a/bin/cli/src/config.rs b/bin/cli/src/config.rs new file mode 100644 index 000000000..6c0a612d1 --- /dev/null +++ b/bin/cli/src/config.rs @@ -0,0 +1,314 @@ +//! `--config` file support: a TOML profile that overrides the prover's code +//! defaults. Every key is optional — the file is a diff over the defaults, +//! never a replacement (several defaults are computed per machine, e.g. the +//! VRAM budget). Precedence, highest first: explicit CLI flag > env var > +//! config file > code default. See `docs/prover.example.toml`. + +use prover::tables::MaxRowsConfig; +use serde::Deserialize; +use std::path::Path; + +#[derive(Debug, Default, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct FileConfig { + #[serde(default)] + pub prove: ProveSection, + #[serde(default)] + pub tables: TablesSection, + #[serde(default)] + pub scheduler: SchedulerSection, + #[serde(default)] + pub gpu: GpuSection, +} + +#[derive(Debug, Default, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct ProveSection { + /// Blowup factor (power of 2), like `--blowup`. + pub blowup: Option, + /// Continuation epoch size as log2(cycles), like `--epoch-size-log2`. + pub epoch_size_log2: Option, +} + +/// Per-table row caps as log2(rows). Defaults equalize memory per instance +/// (`effective_width x rows ~ constant`); raising one cap makes that table's +/// instances proportionally heavier than everyone else's. +#[derive(Debug, Default, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct TablesSection { + pub cpu: Option, + pub memw: Option, + pub memw_aligned: Option, + pub dvrm: Option, + pub mul: Option, + pub lt: Option, + pub shift: Option, + pub load: Option, + pub branch: Option, + pub memw_register: Option, + pub eq: Option, + pub bytewise: Option, + pub store: Option, + pub cpu32: Option, +} + +#[derive(Debug, Default, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct SchedulerSection { + /// Tables proven concurrently (`TABLE_PARALLELISM`). + pub table_parallelism: Option, + /// Device VRAM budget in MiB (`LAMBDA_VM_VRAM_BUDGET_MB`). + pub vram_budget_mb: Option, +} + +#[derive(Debug, Default, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct GpuSection { + /// Minimum LDE size (log2) for the GPU commit paths + /// (`LAMBDA_VM_GPU_LDE_THRESHOLD`). + pub lde_threshold_log2: Option, + /// Minimum trace size (log2) for the GPU barycentric paths + /// (`LAMBDA_VM_GPU_BARY_THRESHOLD`). + pub bary_threshold_log2: Option, + /// CUDA mempool release threshold in MiB + /// (`LAMBDA_VM_MEMPOOL_RELEASE_MB`). + pub mempool_release_mb: Option, + /// Kill-switch: force the CPU composition path + /// (`LAMBDA_VM_DISABLE_GPU_COMPOSITION`). + pub disable_composition: Option, + /// Kill-switch: force the CPU LogUp aux build (`LAMBDA_VM_NO_GPU_LOGUP`). + pub disable_logup: Option, + /// Kill-switch: keep host copies of device-resident buffers + /// (`LAMBDA_VM_DISABLE_DEVICE_ONLY`). + pub disable_device_only: Option, +} + +/// Table caps must stay provable and addressable: 2^5 (the test floor) +/// through 2^25 (u32 LDE domain headroom). +const TABLE_LOG2_RANGE: std::ops::RangeInclusive = 5..=25; + +impl FileConfig { + pub fn load(path: &Path) -> Result { + let text = + std::fs::read_to_string(path).map_err(|e| format!("config {}: {e}", path.display()))?; + let cfg: FileConfig = + toml::from_str(&text).map_err(|e| format!("config {}: {e}", path.display()))?; + cfg.validate() + .map_err(|e| format!("config {}: {e}", path.display()))?; + Ok(cfg) + } + + fn validate(&self) -> Result<(), String> { + let t = &self.tables; + for (name, v) in [ + ("cpu", t.cpu), + ("memw", t.memw), + ("memw_aligned", t.memw_aligned), + ("dvrm", t.dvrm), + ("mul", t.mul), + ("lt", t.lt), + ("shift", t.shift), + ("load", t.load), + ("branch", t.branch), + ("memw_register", t.memw_register), + ("eq", t.eq), + ("bytewise", t.bytewise), + ("store", t.store), + ("cpu32", t.cpu32), + ] { + if let Some(v) = v + && !TABLE_LOG2_RANGE.contains(&v) + { + return Err(format!( + "[tables] {name} = {v}: log2 caps must be within {:?}", + TABLE_LOG2_RANGE + )); + } + } + if let Some(b) = self.prove.blowup + && !b.is_power_of_two() + { + return Err(format!("[prove] blowup = {b}: must be a power of 2")); + } + Ok(()) + } + + /// Code defaults overridden by the set `[tables]` keys. + pub fn max_rows(&self) -> MaxRowsConfig { + let mut m = MaxRowsConfig::default(); + let t = &self.tables; + let set = |dst: &mut usize, v: Option| { + if let Some(log2) = v { + *dst = 1usize << log2; + } + }; + set(&mut m.cpu, t.cpu); + set(&mut m.memw, t.memw); + set(&mut m.memw_aligned, t.memw_aligned); + set(&mut m.dvrm, t.dvrm); + set(&mut m.mul, t.mul); + set(&mut m.lt, t.lt); + set(&mut m.shift, t.shift); + set(&mut m.load, t.load); + set(&mut m.branch, t.branch); + set(&mut m.memw_register, t.memw_register); + set(&mut m.eq, t.eq); + set(&mut m.bytewise, t.bytewise); + set(&mut m.store, t.store); + set(&mut m.cpu32, t.cpu32); + m + } + + /// Warn (stderr) for every file key that a set env var is shadowing, so + /// "I changed the file and nothing happened" is diagnosable at a glance. + pub fn warn_env_shadowing(&self) { + let pairs: [(&str, bool); 8] = [ + ( + "TABLE_PARALLELISM", + self.scheduler.table_parallelism.is_some(), + ), + ( + "LAMBDA_VM_VRAM_BUDGET_MB", + self.scheduler.vram_budget_mb.is_some(), + ), + ( + "LAMBDA_VM_GPU_LDE_THRESHOLD", + self.gpu.lde_threshold_log2.is_some(), + ), + ( + "LAMBDA_VM_GPU_BARY_THRESHOLD", + self.gpu.bary_threshold_log2.is_some(), + ), + ( + "LAMBDA_VM_MEMPOOL_RELEASE_MB", + self.gpu.mempool_release_mb.is_some(), + ), + ( + "LAMBDA_VM_DISABLE_GPU_COMPOSITION", + self.gpu.disable_composition.is_some(), + ), + ("LAMBDA_VM_NO_GPU_LOGUP", self.gpu.disable_logup.is_some()), + ( + "LAMBDA_VM_DISABLE_DEVICE_ONLY", + self.gpu.disable_device_only.is_some(), + ), + ]; + for (var, in_file) in pairs { + if in_file && std::env::var_os(var).is_some() { + eprintln!( + "warning: config file value ignored: env var {var} is set and takes priority" + ); + } + } + } + + /// Install the stark/math-cuda knobs (env vars keep priority at each read + /// site). Call once, before the first prove. + pub fn install_runtime_overrides(&self) { + stark::runtime_overrides::install(stark::runtime_overrides::RuntimeOverrides { + table_parallelism: self.scheduler.table_parallelism, + gpu_lde_threshold: self.gpu.lde_threshold_log2.map(|l| 1usize << l), + gpu_bary_threshold: self.gpu.bary_threshold_log2.map(|l| 1usize << l), + disable_gpu_composition: self.gpu.disable_composition, + no_gpu_logup: self.gpu.disable_logup, + disable_device_only: self.gpu.disable_device_only, + vram_budget_mb: self.scheduler.vram_budget_mb, + mempool_release_mb: self.gpu.mempool_release_mb, + }); + } +} + +#[cfg(test)] +mod tests { + use super::*; + + /// Every key the file format accepts, exercised end to end. A field + /// added to the structs without updating this fixture (or the example + /// file) fails here. + const FULL: &str = r#" + [prove] + blowup = 4 + epoch_size_log2 = 22 + + [tables] + cpu = 20 + memw = 20 + memw_aligned = 20 + dvrm = 20 + mul = 21 + lt = 21 + shift = 21 + load = 21 + branch = 21 + memw_register = 21 + eq = 21 + bytewise = 21 + store = 21 + cpu32 = 20 + + [scheduler] + table_parallelism = 8 + vram_budget_mb = 24000 + + [gpu] + lde_threshold_log2 = 18 + bary_threshold_log2 = 13 + mempool_release_mb = 4096 + disable_composition = true + disable_logup = true + disable_device_only = true + "#; + + #[test] + fn full_fixture_round_trips() { + let cfg: FileConfig = toml::from_str(FULL).unwrap(); + cfg.validate().unwrap(); + assert_eq!(cfg.prove.blowup, Some(4)); + assert_eq!(cfg.prove.epoch_size_log2, Some(22)); + let m = cfg.max_rows(); + assert_eq!(m.cpu, 1 << 20); + assert_eq!(m.lt, 1 << 21); + assert_eq!(m.cpu32, 1 << 20); + assert_eq!(cfg.scheduler.table_parallelism, Some(8)); + assert_eq!(cfg.gpu.lde_threshold_log2, Some(18)); + assert_eq!(cfg.gpu.disable_device_only, Some(true)); + } + + #[test] + fn empty_and_partial_files_keep_defaults() { + let cfg: FileConfig = toml::from_str("").unwrap(); + let d = MaxRowsConfig::default(); + let m = cfg.max_rows(); + assert_eq!(m.cpu, d.cpu); + assert_eq!(m.store, d.store); + + let cfg: FileConfig = toml::from_str("[tables]\ncpu = 20\n").unwrap(); + let m = cfg.max_rows(); + assert_eq!(m.cpu, 1 << 20); + assert_eq!(m.memw, d.memw); // untouched keys keep code defaults + } + + #[test] + fn unknown_keys_and_bad_values_are_rejected() { + assert!(toml::from_str::("[tables]\ncppu = 20\n").is_err()); + assert!(toml::from_str::("[typo_section]\nx = 1\n").is_err()); + let cfg: FileConfig = toml::from_str("[tables]\ncpu = 30\n").unwrap(); + assert!(cfg.validate().is_err()); // out of range + let cfg: FileConfig = toml::from_str("[prove]\nblowup = 3\n").unwrap(); + assert!(cfg.validate().is_err()); // not a power of two + } + + /// The example file must parse (it ships fully commented, so it must + /// stay valid TOML with no unknown keys when uncommented sections drift + /// is checked by hand — this pins at least the syntactic contract). + #[test] + fn example_file_parses() { + let text = std::fs::read_to_string(concat!( + env!("CARGO_MANIFEST_DIR"), + "/../../docs/prover.example.toml" + )) + .expect("docs/prover.example.toml exists"); + let cfg: FileConfig = toml::from_str(&text).expect("example parses"); + cfg.validate().expect("example validates"); + } +} diff --git a/bin/cli/src/main.rs b/bin/cli/src/main.rs index a04e920db..3ec2c432c 100644 --- a/bin/cli/src/main.rs +++ b/bin/cli/src/main.rs @@ -6,6 +6,8 @@ use std::path::{Path, PathBuf}; use std::process::ExitCode; use std::time::Instant; +mod config; + use clap::{Parser, Subcommand, ValueHint}; #[global_allocator] @@ -164,8 +166,15 @@ enum Commands { private_input: Option, /// Blowup factor (power of 2). Higher = fewer queries, smaller proof, slower proving. - #[arg(long, default_value = "2")] - blowup: u8, + /// Default: 2, or the config file's [prove].blowup. + #[arg(long)] + blowup: Option, + + /// TOML profile overriding the prover's tuning defaults (table row + /// caps, scheduler, GPU knobs). Explicit flags and env vars win over + /// the file; see docs/prover.example.toml. + #[arg(long, value_hint = ValueHint::FilePath)] + config: Option, /// Print proving time #[arg(long)] @@ -205,9 +214,14 @@ enum Commands { #[arg(value_parser, value_hint = ValueHint::FilePath)] elf: PathBuf, - /// Blowup factor used during proving (must match) - #[arg(long, default_value = "2")] - blowup: u8, + /// Blowup factor used during proving (must match). + /// Default: 2, or the config file's [prove].blowup. + #[arg(long)] + blowup: Option, + + /// TOML profile; only [prove].blowup applies to verification. + #[arg(long, value_hint = ValueHint::FilePath)] + config: Option, /// Print verification time #[arg(long)] @@ -259,12 +273,40 @@ fn main() -> ExitCode { output, private_input, blowup, + config, time, cycles, elements, continuations, epoch_size_log2, } => { + // Precedence: explicit CLI flag > env var > config file > default. + // The file's runtime knobs are installed before any prover work; + // env vars keep priority at each read site. + let file_cfg = match load_file_config(config.as_deref()) { + Ok(c) => c, + Err(e) => { + eprintln!("{e}"); + return ExitCode::FAILURE; + } + }; + file_cfg.warn_env_shadowing(); + file_cfg.install_runtime_overrides(); + let blowup = blowup.or(file_cfg.prove.blowup).unwrap_or(2); + let epoch_size_log2 = match epoch_size_log2 { + Some(v) => Some(v), + None => match file_cfg.prove.epoch_size_log2 { + Some(v) => match parse_epoch_size_log2(&v.to_string()) { + Ok(parsed) => Some(parsed), + Err(e) => { + eprintln!("config [prove].epoch_size_log2: {e}"); + return ExitCode::FAILURE; + } + }, + None => None, + }, + }; + let max_rows = file_cfg.max_rows(); if continuations { cmd_prove_continuation( elf, @@ -272,20 +314,39 @@ fn main() -> ExitCode { private_input, epoch_size_log2, blowup, + max_rows, time, cycles, ) } else { - cmd_prove(elf, output, private_input, blowup, time, cycles, elements) + cmd_prove( + elf, + output, + private_input, + blowup, + max_rows, + time, + cycles, + elements, + ) } } Commands::Verify { proof, elf, blowup, + config, time, continuations, } => { + let file_cfg = match load_file_config(config.as_deref()) { + Ok(c) => c, + Err(e) => { + eprintln!("{e}"); + return ExitCode::FAILURE; + } + }; + let blowup = blowup.or(file_cfg.prove.blowup).unwrap_or(2); if continuations { cmd_verify_continuation(proof, elf, blowup, time) } else { @@ -296,6 +357,13 @@ fn main() -> ExitCode { } } +fn load_file_config(path: Option<&std::path::Path>) -> Result { + match path { + Some(p) => config::FileConfig::load(p), + None => Ok(config::FileConfig::default()), + } +} + fn read_private_input(path: Option<&PathBuf>) -> Result, String> { match path { Some(path) => { @@ -542,11 +610,13 @@ fn cmd_execute( ExitCode::SUCCESS } +#[allow(clippy::too_many_arguments)] fn cmd_prove( elf_path: PathBuf, output_path: PathBuf, private_input_path: Option, blowup: u8, + max_rows: prover::tables::MaxRowsConfig, time: bool, cycles: bool, elements: bool, @@ -617,12 +687,7 @@ fn cmd_prove( "Generating proof (blowup={blowup}, queries={})...", opts.fri_number_of_queries ); - let proof = prover::prove_with_options_and_inputs( - &elf_data, - &private_inputs, - &opts, - &Default::default(), - ); + let proof = prover::prove_with_options_and_inputs(&elf_data, &private_inputs, &opts, &max_rows); let prove_elapsed = start.elapsed(); let proof = match proof { Ok(proof) => proof, @@ -732,12 +797,14 @@ fn cmd_verify(proof_path: PathBuf, elf_path: PathBuf, blowup: u8, time: bool) -> } } +#[allow(clippy::too_many_arguments)] fn cmd_prove_continuation( elf_path: PathBuf, output_path: PathBuf, private_input_path: Option, epoch_size_log2: Option, blowup: u8, + max_rows: prover::tables::MaxRowsConfig, time: bool, cycles: bool, ) -> ExitCode { @@ -797,11 +864,12 @@ fn cmd_prove_continuation( #[cfg(feature = "jemalloc-stats")] let tracker = heap_tracker::HeapTracker::start(); let start = Instant::now(); - let bundle = match prover::continuation::prove_continuation( + let bundle = match prover::continuation::prove_continuation_with_max_rows( &elf_data, &private_inputs, epoch_size_log2, &opts, + &max_rows, ) { Ok(b) => b, Err(e) => { diff --git a/crypto/math-cuda/src/device.rs b/crypto/math-cuda/src/device.rs index 8fd7f13de..d661ef828 100644 --- a/crypto/math-cuda/src/device.rs +++ b/crypto/math-cuda/src/device.rs @@ -136,6 +136,21 @@ const LOGUP_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/logup.cubin const CONSTRAINT_INTERP_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/constraint_interp.cubin")); +/// Runtime overrides for the device knobs, installed once at startup (before +/// backend init) by the config-file layer in the stark crate. Env vars keep +/// priority at each read site. +static RUNTIME_OVERRIDES: std::sync::OnceLock<(Option, Option)> = + std::sync::OnceLock::new(); + +/// `(vram_budget_mb, mempool_release_mb)`; first install wins. +pub fn set_runtime_overrides(vram_budget_mb: Option, mempool_release_mb: Option) { + let _ = RUNTIME_OVERRIDES.set((vram_budget_mb, mempool_release_mb)); +} + +fn runtime_overrides() -> (Option, Option) { + RUNTIME_OVERRIDES.get().copied().unwrap_or((None, None)) +} + /// Number of CUDA streams in the pool. Larger pools let many rayon-parallel /// callers overlap on the GPU without serializing on stream ownership. The /// default stream is deliberately excluded because it synchronises with all @@ -266,6 +281,7 @@ fn retain_default_mempool(ctx: &CudaContext) { let threshold: u64 = std::env::var("LAMBDA_VM_MEMPOOL_RELEASE_MB") .ok() .and_then(|s| s.parse::().ok()) + .or(runtime_overrides().1) .map(|mb| mb.saturating_mul(1024 * 1024)) .unwrap_or(u64::MAX); let _ = sys::cuMemPoolSetAttribute( @@ -290,6 +306,9 @@ fn detect_vram_budget_bytes(ctx: &CudaContext) -> u64 { { return mb.saturating_mul(1024 * 1024); } + if let (Some(mb), _) = runtime_overrides() { + return mb.saturating_mul(1024 * 1024); + } use cudarc::driver::sys; // SAFETY: raw driver query writing into two stack slots. The caller's // context is already current (it was just created in `init`). Any error diff --git a/crypto/stark/src/gpu_lde.rs b/crypto/stark/src/gpu_lde.rs index 2167fcb94..16454d269 100644 --- a/crypto/stark/src/gpu_lde.rs +++ b/crypto/stark/src/gpu_lde.rs @@ -50,6 +50,7 @@ fn gpu_lde_threshold() -> usize { std::env::var("LAMBDA_VM_GPU_LDE_THRESHOLD") .ok() .and_then(|s| s.parse().ok()) + .or(crate::runtime_overrides::get().gpu_lde_threshold) .unwrap_or(DEFAULT_GPU_LDE_THRESHOLD) }) } @@ -144,6 +145,8 @@ pub(crate) fn gpu_composition_disabled() -> bool { *ENV_DISABLED.get_or_init(|| { std::env::var("LAMBDA_VM_DISABLE_GPU_COMPOSITION") .map(|v| v == "1") + .ok() + .or(crate::runtime_overrides::get().disable_gpu_composition) .unwrap_or(false) }) } @@ -165,6 +168,8 @@ pub(crate) fn device_only_disabled() -> bool { *ENV_DISABLED.get_or_init(|| { std::env::var("LAMBDA_VM_DISABLE_DEVICE_ONLY") .map(|v| v == "1") + .ok() + .or(crate::runtime_overrides::get().disable_device_only) .unwrap_or(false) }) } @@ -998,6 +1003,7 @@ fn gpu_bary_threshold() -> usize { std::env::var("LAMBDA_VM_GPU_BARY_THRESHOLD") .ok() .and_then(|s| s.parse().ok()) + .or(crate::runtime_overrides::get().gpu_bary_threshold) .unwrap_or(DEFAULT_GPU_BARY_THRESHOLD) }) } diff --git a/crypto/stark/src/lib.rs b/crypto/stark/src/lib.rs index 6f8e7c82e..1a96c7377 100644 --- a/crypto/stark/src/lib.rs +++ b/crypto/stark/src/lib.rs @@ -29,6 +29,7 @@ pub mod profile_markers; pub mod proof; pub mod prover; pub mod r4_denoms; +pub mod runtime_overrides; #[cfg(feature = "disk-spill")] pub mod storage_mode; pub mod table; diff --git a/crypto/stark/src/logup_gpu.rs b/crypto/stark/src/logup_gpu.rs index 3fd49134d..a68b179f1 100644 --- a/crypto/stark/src/logup_gpu.rs +++ b/crypto/stark/src/logup_gpu.rs @@ -364,7 +364,11 @@ where return None; } // Escape hatch for A/B measurement: force the CPU aux build. - if std::env::var_os("LAMBDA_VM_NO_GPU_LOGUP").is_some() { + if std::env::var_os("LAMBDA_VM_NO_GPU_LOGUP").is_some() + || crate::runtime_overrides::get() + .no_gpu_logup + .unwrap_or(false) + { return None; } @@ -434,7 +438,11 @@ where if trace_len < GPU_LOGUP_MIN_ROWS || main_cols.is_empty() || interactions.is_empty() { return None; } - if std::env::var_os("LAMBDA_VM_NO_GPU_LOGUP").is_some() { + if std::env::var_os("LAMBDA_VM_NO_GPU_LOGUP").is_some() + || crate::runtime_overrides::get() + .no_gpu_logup + .unwrap_or(false) + { return None; } let desc = build_fingerprint_descriptor(interactions); diff --git a/crypto/stark/src/prover.rs b/crypto/stark/src/prover.rs index 9a369b042..38c14c119 100644 --- a/crypto/stark/src/prover.rs +++ b/crypto/stark/src/prover.rs @@ -577,6 +577,7 @@ pub fn table_parallelism() -> usize { std::env::var("TABLE_PARALLELISM") .ok() .and_then(|s| s.parse().ok()) + .or(crate::runtime_overrides::get().table_parallelism) .unwrap_or_else(|| { let cores = std::thread::available_parallelism() .map(|n| n.get()) diff --git a/crypto/stark/src/runtime_overrides.rs b/crypto/stark/src/runtime_overrides.rs new file mode 100644 index 000000000..86df72c0c --- /dev/null +++ b/crypto/stark/src/runtime_overrides.rs @@ -0,0 +1,66 @@ +//! Process-wide runtime overrides for the prover's tuning knobs, installed +//! once at startup (the CLI's `--config` file) and consulted by the knob +//! readers with the precedence `env var > override > code default` — a set +//! env var always wins, so existing scripts keep working unchanged. +//! +//! Install before the first prove: the readers cache their resolved value on +//! first use, so a later install would be silently ignored (`install` returns +//! whether it won the race). + +use std::sync::OnceLock; + +/// Values a config file may override. `None` = not set → code default. +#[derive(Debug, Clone, Default)] +pub struct RuntimeOverrides { + /// Tables proven concurrently (`TABLE_PARALLELISM`). + pub table_parallelism: Option, + /// Minimum LDE size for the GPU commit paths + /// (`LAMBDA_VM_GPU_LDE_THRESHOLD`). + pub gpu_lde_threshold: Option, + /// Minimum trace size for the GPU barycentric paths + /// (`LAMBDA_VM_GPU_BARY_THRESHOLD`). + pub gpu_bary_threshold: Option, + /// Kill-switch: skip the GPU composition path + /// (`LAMBDA_VM_DISABLE_GPU_COMPOSITION`). + pub disable_gpu_composition: Option, + /// Kill-switch: skip the GPU LogUp paths (`LAMBDA_VM_NO_GPU_LOGUP`). + pub no_gpu_logup: Option, + /// Kill-switch: keep host copies of device-resident buffers + /// (`LAMBDA_VM_DISABLE_DEVICE_ONLY`). + pub disable_device_only: Option, + /// Device VRAM budget in MiB (`LAMBDA_VM_VRAM_BUDGET_MB`). + pub vram_budget_mb: Option, + /// CUDA mempool release threshold in MiB + /// (`LAMBDA_VM_MEMPOOL_RELEASE_MB`). + pub mempool_release_mb: Option, +} + +static OVERRIDES: OnceLock = OnceLock::new(); + +/// Install the overrides. Returns `false` if a set was already installed (the +/// first install wins; callers should treat `false` as a startup-order bug). +pub fn install(overrides: RuntimeOverrides) -> bool { + #[cfg(feature = "cuda")] + let device = (overrides.vram_budget_mb, overrides.mempool_release_mb); + let won = OVERRIDES.set(overrides).is_ok(); + // The device-side knobs are read inside math-cuda; forward them. + #[cfg(feature = "cuda")] + if won { + math_cuda::device::set_runtime_overrides(device.0, device.1); + } + won +} + +pub(crate) fn get() -> &'static RuntimeOverrides { + static EMPTY: RuntimeOverrides = RuntimeOverrides { + table_parallelism: None, + gpu_lde_threshold: None, + gpu_bary_threshold: None, + disable_gpu_composition: None, + no_gpu_logup: None, + disable_device_only: None, + vram_budget_mb: None, + mempool_release_mb: None, + }; + OVERRIDES.get().unwrap_or(&EMPTY) +} diff --git a/docs/prover.example.toml b/docs/prover.example.toml new file mode 100644 index 000000000..a8a05882e --- /dev/null +++ b/docs/prover.example.toml @@ -0,0 +1,53 @@ +# Prover configuration profile — pass with `cli prove --config `. +# +# Every key is OPTIONAL: the file is a diff over the code defaults, never a +# replacement. Anything not set here keeps the default noted in its comment +# (several are computed per machine and have no fixed number). Precedence, +# highest first: explicit CLI flag > env var > this file > code default. +# +# Copy this file, uncomment only what you want to change. + +[prove] +# blowup = 2 # power of 2; must match at verification +# epoch_size_log2 = 20 # continuation epoch cycles (log2); min 18. + # Higher = fewer epochs and less fixed overhead, + # more peak RAM/VRAM per epoch. + +[tables] +# Per-table row caps, log2(rows). A table whose ops exceed its cap splits +# into multiple instances. Defaults equalize memory per instance +# (effective_width x rows ~ constant): wide tables cap low, narrow ones high. +# Raising one cap alone makes that table's instances heavier than the rest. +# cpu = 19 +# memw = 19 +# memw_aligned = 19 +# dvrm = 19 +# cpu32 = 19 +# mul = 20 +# lt = 20 +# shift = 20 +# load = 20 +# branch = 20 +# memw_register = 20 +# eq = 20 +# bytewise = 20 +# store = 20 + +[scheduler] +# table_parallelism = 5 # tables proven concurrently. + # Default: computed from the machine's cores. + # Env: TABLE_PARALLELISM +# vram_budget_mb = 26000 # device admission budget. + # Default: 80% of the detected VRAM. + # Env: LAMBDA_VM_VRAM_BUDGET_MB + +[gpu] +# lde_threshold_log2 = 19 # min LDE size for the GPU commit paths. + # Env: LAMBDA_VM_GPU_LDE_THRESHOLD (in rows) +# bary_threshold_log2 = 14 # min trace size for GPU barycentric. + # Env: LAMBDA_VM_GPU_BARY_THRESHOLD (in rows) +# mempool_release_mb = 0 # CUDA mempool release threshold; unset = retain. + # Env: LAMBDA_VM_MEMPOOL_RELEASE_MB +# disable_composition = false # kill-switch (LAMBDA_VM_DISABLE_GPU_COMPOSITION) +# disable_logup = false # kill-switch (LAMBDA_VM_NO_GPU_LOGUP) +# disable_device_only = false # kill-switch (LAMBDA_VM_DISABLE_DEVICE_ONLY) diff --git a/prover/src/continuation.rs b/prover/src/continuation.rs index 73e4c877e..496bc7ab1 100644 --- a/prover/src/continuation.rs +++ b/prover/src/continuation.rs @@ -1036,6 +1036,24 @@ pub fn prove_continuation( private_inputs: &[u8], epoch_size_log2: u32, opts: &ProofOptions, +) -> Result { + prove_continuation_with_max_rows( + elf_bytes, + private_inputs, + epoch_size_log2, + opts, + &MaxRowsConfig::default(), + ) +} + +/// [`prove_continuation`] with explicit per-table row caps (the config-file +/// path); the plain entry point uses [`MaxRowsConfig::default`]. +pub fn prove_continuation_with_max_rows( + elf_bytes: &[u8], + private_inputs: &[u8], + epoch_size_log2: u32, + opts: &ProofOptions, + max_rows: &MaxRowsConfig, ) -> Result { if epoch_size_log2 < 2 { return Err(Error::InvalidContinuationEpochSize( @@ -1210,7 +1228,7 @@ pub fn prove_continuation( // only image consumers in the build) are skipped. None::<&std::collections::HashMap>, &job.register_init, - &MaxRowsConfig::default(), + max_rows, private_inputs, job.is_final, true,