diff --git a/crypto/math-cuda/src/blake3.rs b/crypto/math-cuda/src/blake3.rs index b8914bec2..15747baa8 100644 --- a/crypto/math-cuda/src/blake3.rs +++ b/crypto/math-cuda/src/blake3.rs @@ -603,6 +603,10 @@ pub fn build_comp_poly_tree_from_slabs_dev( m: usize, lde_size: usize, ) -> Result { + // Same sticky hook as the keccak twin: the comp-tree cliff test arms one + // counter and must reach it under whichever hash the build pins. + #[cfg(feature = "test-faults")] + crate::faults::check_sticky(&crate::faults::FAULT_COMP_TREE_STICKY)?; assert!(m > 0); assert!(lde_size.is_power_of_two() && lde_size >= 2); assert_eq!(buf.len(), 3 * m * lde_size, "slab buffer shape"); @@ -647,6 +651,8 @@ pub fn build_comp_poly_tree_from_slabs_dev( pub fn build_comp_poly_tree_from_evals_ext3_keep( parts_interleaved: &[&[u64]], ) -> Result { + #[cfg(feature = "test-faults")] + crate::faults::check_sticky(&crate::faults::FAULT_COMP_TREE_STICKY)?; assert!(!parts_interleaved.is_empty()); let m = parts_interleaved.len(); let ext3_elems = parts_interleaved[0].len() / 3; diff --git a/crypto/math-cuda/src/device.rs b/crypto/math-cuda/src/device.rs index 687d2d265..c62bf2855 100644 --- a/crypto/math-cuda/src/device.rs +++ b/crypto/math-cuda/src/device.rs @@ -285,42 +285,102 @@ pub struct Backend { inv_twiddles: Mutex>>>>, } -/// Raise the device default memory pool's release threshold so freed -/// stream-ordered allocations are kept for reuse instead of returned to the OS -/// at each sync. Best-effort: any failure (e.g. a device/driver without -/// stream-ordered allocator support) leaves the default behaviour untouched. -fn retain_default_mempool(ctx: &CudaContext) { +/// The environment knob for the device default memory pool's release +/// threshold, in MiB: the bytes of freed stream-ordered memory the pool keeps +/// before handing memory back to the OS at the next sync. Unset means +/// [`DEFAULT_MEMPOOL_RELEASE_THRESHOLD_BYTES`]. The VRAM sampler runs set it +/// to `0`, so `total - free` reads the live working set and not the retained +/// pool. +pub const MEMPOOL_RELEASE_ENV: &str = "LAMBDA_VM_MEMPOOL_RELEASE_MB"; + +/// Retain every freed block (`u64::MAX`): a same-shape allocation skips the +/// driver and reuses the block, which is what the per-table pipeline's +/// repeated LDE/FRI buffers want. +/// +/// Measured, not guessed (RTX 5090, 2026-09-07, `one_lde_buffer::vram_arm` +/// at 2^21 × 316 @ blowup 2, five commits, in-process 1 kHz peak): retain-all +/// 1176.9 / 1057.9 / 1055.3 / 1056.2 / 1055.6 ms against release-0 +/// 1178.5 / 1057.2 / 1077.5 / 1081.2 / 1079.3 ms — retention ≈2% faster once +/// the first commit has populated the pool — and a peak of 15.67 GiB under +/// both: a same-shape allocation reuses the retained block, so retention adds +/// nothing to the peak. The unequal-shape case is covered at block scale by +/// the multi-table q=41 wrap rung under this default (VRAM peak 28,976 MiB, no +/// device decline): the stream-ordered allocator serves a new request from the +/// physical chunks it retains, and the release threshold governs only what a +/// sync hands back to the OS. The explicit release for a moment reuse cannot +/// serve is [`Backend::trim_mempool_to`]; the sampler runs set the knob to `0` +/// so `total - free` reads the live set rather than the pool. +pub const DEFAULT_MEMPOOL_RELEASE_THRESHOLD_BYTES: u64 = u64::MAX; + +/// The effective release threshold in bytes: the knob when set and parseable, +/// the default otherwise. Read once per process; the prover's diagnostics +/// print it so every box log states the posture its run had. +pub fn mempool_release_threshold_bytes() -> u64 { + static CACHED: OnceLock = OnceLock::new(); + *CACHED.get_or_init(|| { + std::env::var(MEMPOOL_RELEASE_ENV) + .ok() + .and_then(|s| s.parse::().ok()) + .map(|mb| mb.saturating_mul(1024 * 1024)) + .unwrap_or(DEFAULT_MEMPOOL_RELEASE_THRESHOLD_BYTES) + }) +} + +/// The device default memory pool, or `None` on a device/driver without +/// stream-ordered allocator support. +/// +/// # Safety +/// +/// `ctx` must be a live context; its device is queried directly. +unsafe fn default_mempool(ctx: &CudaContext) -> Option { use cudarc::driver::sys; - // SAFETY: raw CUDA driver calls. `ctx.cu_device()` is a valid device for - // the just-created context; the out-pointers are valid stack slots; the - // threshold is read as a u64 by the driver. Errors are swallowed. + let mut pool: sys::CUmemoryPool = std::ptr::null_mut(); + // SAFETY: the out-pointer is a valid stack slot; the device is the + // context's own. unsafe { - let dev = ctx.cu_device(); - let mut pool: sys::CUmemoryPool = std::ptr::null_mut(); - if sys::cuDeviceGetDefaultMemPool(&mut pool as *mut _, dev) + sys::cuDeviceGetDefaultMemPool(&mut pool as *mut _, ctx.cu_device()) .result() - .is_err() - { - return; - } - // Default: retain freed stream-ordered blocks indefinitely (u64::MAX) - // for reuse. `LAMBDA_VM_MEMPOOL_RELEASE_MB` overrides the cap (bytes the - // pool keeps before returning memory to the OS) when retained-pool - // growth needs bounding. - let threshold: u64 = std::env::var("LAMBDA_VM_MEMPOOL_RELEASE_MB") .ok() - .and_then(|s| s.parse::().ok()) - .map(|mb| mb.saturating_mul(1024 * 1024)) - .unwrap_or(u64::MAX); - let _ = sys::cuMemPoolSetAttribute( - pool, - sys::CUmemPool_attribute_enum::CU_MEMPOOL_ATTR_RELEASE_THRESHOLD, - &threshold as *const u64 as *mut core::ffi::c_void, - ) - .result(); + .map(|()| pool) } } +/// Set the device default memory pool's release threshold +/// ([`mempool_release_threshold_bytes`]) so freed stream-ordered allocations +/// are kept for reuse instead of returned to the OS at each sync. Best-effort: +/// any failure leaves the driver default (release everything) untouched, and +/// the one-line report says so. +fn retain_default_mempool(ctx: &CudaContext) { + use cudarc::driver::sys; + let threshold = mempool_release_threshold_bytes(); + // SAFETY: raw CUDA driver calls on the just-created context's device; the + // threshold is read as a u64 by the driver. Errors are swallowed. + let set = unsafe { + default_mempool(ctx).is_some_and(|pool| { + sys::cuMemPoolSetAttribute( + pool, + sys::CUmemPool_attribute_enum::CU_MEMPOOL_ATTR_RELEASE_THRESHOLD, + &threshold as *const u64 as *mut core::ffi::c_void, + ) + .result() + .is_ok() + }) + }; + // One line per process, so the box log states the posture the run had. + eprintln!( + "[gpu] mempool release threshold: {}{}", + match threshold { + u64::MAX => "retain all freed blocks".to_string(), + t => format!("{} MiB", t >> 20), + }, + if set { + "" + } else { + " (driver refused; the release-on-sync default stays)" + } + ); +} + /// Device VRAM budget in bytes for table session admission control. /// /// LAMBDA_VM_VRAM_BUDGET_MB overrides it (used to force the throttle in tests). @@ -557,6 +617,36 @@ impl Backend { self.vram_budget_bytes } + /// Live `(free, total)` device memory in bytes, for diagnostics — the + /// admission gates never read it (they must answer the same at R1 and at + /// R4). `None` when the query fails. + pub fn device_mem_info(&self) -> Option<(u64, u64)> { + self.ctx + .mem_get_info() + .ok() + .map(|(free, total)| (free as u64, total as u64)) + } + + /// Hand the default memory pool's unused reserved memory back to the OS, + /// keeping at most `keep_bytes` (`cuMemPoolTrimTo`). Under the retained + /// posture ([`mempool_release_threshold_bytes`]) a sync never releases; + /// this is the explicit release for the moments reuse cannot serve — a + /// differently-shaped table after a large one, or a sampler that must read + /// the live working set. Best effort: `false` when the pool cannot be + /// queried or the trim fails. + pub fn trim_mempool_to(&self, keep_bytes: u64) -> bool { + use cudarc::driver::sys; + // SAFETY: raw driver calls on this backend's live context; the trim + // takes a plain byte count. + unsafe { + default_mempool(&self.ctx).is_some_and(|pool| { + sys::cuMemPoolTrimTo(pool, keep_bytes as usize) + .result() + .is_ok() + }) + } + } + /// Round-robin over the stream pool. Concurrent callers get different /// streams so their kernel launches overlap on the GPU. pub fn next_stream(&self) -> Arc { diff --git a/crypto/stark/src/gpu_lde.rs b/crypto/stark/src/gpu_lde.rs index b5d5c7904..887a7cc22 100644 --- a/crypto/stark/src/gpu_lde.rs +++ b/crypto/stark/src/gpu_lde.rs @@ -5,6 +5,12 @@ //! launch overhead dominates. Produces the same natural-order, non-canonical //! LDE evaluations as the CPU path. //! +//! Admission is bytes-aware and the device path, once admitted, is the ONLY +//! path: see the "Bytes-aware admission" section. A table the device admits +//! proves on the device or the prove aborts with a diagnostic; host RAM is a +//! cache, not a compute path, and the one way back to a host commit is the +//! test-only switch [`TEST_ONLY_HOST_FALLBACK_ENV`]. +//! //! The tree-building entries here are generic over a Merkle backend `B` that //! they never call: the leaf and parent hashing happens in the `math-cuda` //! kernels, and `B` only types the host `MerkleTree` the root is wrapped in. @@ -154,6 +160,10 @@ fn gpu_lde_threshold() -> usize { /// nothing to fall back to and aborts regardless of the host LDE. Lowering /// the commit threshold therefore widens that one abort site even though it /// leaves this envelope alone. +/// +/// Inside the envelope the contract is stricter still: a device-only table's +/// host recovery is refused in production ([`refuse_host_recovery`]), so a +/// runtime decline there is an abort with a diagnostic, never a slow prove. const DEFAULT_DEVICE_ONLY_MIN_LDE: usize = 1 << 19; fn gpu_device_only_threshold() -> usize { @@ -166,8 +176,394 @@ fn gpu_device_only_threshold() -> usize { }) } +// ============================================================================ +// Bytes-aware admission — the ONE predicate every dispatch site consults +// ============================================================================ +// +// Two rules decide whether a dispatch may ask the device for anything, and +// both live here so no site carries a private copy of either: +// +// - the row FLOOR (`gpu_lde_threshold`): below it the host path is the faster +// one — launch overhead dominates — and taking it is policy, not a fallback; +// - the bytes CEILING (the card's admission budget): what the site is about to +// allocate must fit the card. Width enters here and only here: a `2^21 × 449` +// table clears the floor at any width, and used to sail into an allocation +// failure that quietly became a host commit. +// +// What happens past admission is the other half of the contract. Host RAM is a +// cache, not a compute path: a table the device ADMITS proves on the device or +// the prove stops, loudly, with the shape, the bytes and the live VRAM in the +// message ([`abort_or_test_fallback`]). The only way back to a host commit is +// the test-only switch [`TEST_ONLY_HOST_FALLBACK_ENV`]. +// +// LOCKSTEP: [`admit`] is a pure function of `(lde_size, bytes)` and two +// process constants (the floor, the card's budget). R1 sizes its commit and +// asks; every handle-bearing site re-derives admission later from the same +// `lde_size` and its own, smaller, transient, and gets the same answer. Live +// free memory is deliberately NOT an input: it would let a table admitted at +// R1 be declined at R4, which for a device-only table is an abort with the +// work already done. FRI re-derives at width 1 (one ext3 column); the ceiling +// 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 process predicate: `gpu_lde_threshold()` as the floor and the card's +/// admission budget ([`device_vram_budget_bytes`]: 80% of device memory, or +/// `LAMBDA_VM_VRAM_BUDGET_MB`) as the ceiling. Both are fixed for the life of +/// the process, so the same `(lde_size, bytes)` gets the same answer at R1 and +/// at every later re-derivation. +pub(crate) fn admit(lde_size: usize, bytes: u64) -> Admission { + match device_vram_budget_bytes() { + None => Admission::NoDevice, + Some(budget) => admit_bytes(lde_size, bytes, gpu_lde_threshold(), budget), + } +} + +/// One dispatch's identity for the diagnostics: which device stage, and the +/// shape it was sized for. The table's NAME is not known at this layer; the +/// 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 what: &'static str, + pub n: usize, + pub base_cols: usize, + pub blowup: usize, +} + +/// Why a stage that must run on the device could not. +#[derive(Debug)] +pub(crate) enum DevicePathFailure { + OverBudget { bytes: u64, budget: u64 }, + DeviceError(String), +} + +fn gib(bytes: u64) -> f64 { + bytes as f64 / (1u64 << 30) as f64 +} + +/// Live `free / total` device memory for the diagnostics — the one place the +/// live number is read; admission itself never looks at it (LOCKSTEP). +fn live_vram_line() -> String { + match math_cuda::device::backend() + .ok() + .and_then(|be| be.device_mem_info()) + { + Some((free, total)) => format!( + "live VRAM free {:.2} GiB of {:.2} GiB", + gib(free), + gib(total) + ), + None => "live VRAM unavailable".to_string(), + } +} + +fn mempool_line() -> String { + match math_cuda::device::mempool_release_threshold_bytes() { + u64::MAX => "mempool retains freed blocks (release threshold unset)".to_string(), + t => format!("mempool release threshold {} MiB", t >> 20), + } +} + +/// The diagnostic every abort and every test-only fallback prints: the stage +/// 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, + set: Option<&CommitDeviceSet>, + failure: &DevicePathFailure, +) -> String { + let DispatchShape { + what, + n, + base_cols, + blowup, + } = shape; + let reason = match failure { + DevicePathFailure::OverBudget { bytes, budget } => format!( + "over the VRAM admission budget: needs {bytes} B ({:.2} GiB), budget {budget} B ({:.2} GiB)", + gib(*bytes), + gib(*budget) + ), + DevicePathFailure::DeviceError(e) => format!("device error after admission: {e}"), + }; + let set_line = match set { + Some(s) => format!( + "; device set LDE {:.2} GiB + snapshot {:.2} GiB + tree {:.3} GiB + scratch {:.3} GiB = {:.2} GiB", + gib(s.lde_bytes), + gib(s.snapshot_bytes), + gib(s.tree_bytes), + gib(s.scratch_bytes), + gib(s.total()) + ), + None => String::new(), + }; + format!( + "{what}: rows {n} x {base_cols} base cols @ blowup {blowup} (LDE {}); {reason}{set_line}; {}; {}", + n.saturating_mul(*blowup), + live_vram_line(), + mempool_line() + ) +} + +/// The test-only switch that turns the loud abort back into the old host +/// commit. Named so it cannot be read as a production knob; a banner is +/// printed when it is honoured. Production runs never set it. +pub const TEST_ONLY_HOST_FALLBACK_ENV: &str = "LAMBDA_VM_TEST_ONLY_HOST_FALLBACK"; + +/// Whether a device failure may fall back to host compute in this process. +/// `true` under [`TEST_ONLY_HOST_FALLBACK_ENV`], under the +/// [`gpu_force_downgrade`] test hook (whose whole purpose is to exercise the +/// host recovery), and in a `test-cuda-faults` build (the fault-injection +/// suite asserts on the recoveries). None of the three is a production +/// configuration; the banner says so once. +pub(crate) fn test_only_host_fallback() -> bool { + static CACHED: OnceLock = OnceLock::new(); + *CACHED.get_or_init(|| { + let by_env = std::env::var(TEST_ONLY_HOST_FALLBACK_ENV).is_ok_and(|v| v != "0"); + let by_hook = gpu_force_downgrade(); + let by_feature = cfg!(feature = "test-cuda-faults"); + let on = by_env || by_hook || by_feature; + if on { + let why = if by_env { + TEST_ONLY_HOST_FALLBACK_ENV + } else if by_hook { + "LAMBDA_VM_GPU_FORCE_DOWNGRADE" + } else { + "feature test-cuda-faults" + }; + eprintln!( + "[gpu] TEST-ONLY host fallback ENABLED ({why}): a device failure falls back to \ + host compute instead of aborting the prove. Never a production configuration." + ); + } + on + }) +} + +/// Refuse the host path — or take it under the test-only switch. +/// +/// Returns only when [`test_only_host_fallback`] is set, after printing the +/// diagnostic tagged TEST-ONLY; the caller then returns `None` and its host +/// arm runs. Otherwise panics with the diagnostic: the prover's driver threads +/// 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, + set: Option<&CommitDeviceSet>, + failure: DevicePathFailure, +) { + let msg = device_path_diagnostic(shape, set, &failure); + if test_only_host_fallback() { + eprintln!("[gpu] TEST-ONLY host fallback: {msg}"); + return; + } + panic!("[gpu] ABORT: the device path is the production path and it is unavailable — {msg}"); +} + +/// R1 commit admission: `None` when the device is absent or the table is +/// 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<()> { + match admit(lde_size, set.total()) { + Admission::NoDevice | Admission::BelowFloor { .. } => None, + Admission::Admitted { .. } => Some(()), + Admission::OverBudget { bytes, budget } => { + abort_or_test_fallback( + shape, + Some(set), + DevicePathFailure::OverBudget { bytes, budget }, + ); + None + } + } +} + +/// [`admit_commit`] for an input that is ALREADY resident on device (the +/// 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<()> { + let budget = device_vram_budget_bytes()?; + match admit_bytes(usize::MAX, set.total(), 0, budget) { + Admission::OverBudget { bytes, budget } => { + abort_or_test_fallback( + shape, + Some(set), + DevicePathFailure::OverBudget { bytes, budget }, + ); + None + } + _ => Some(()), + } +} + +/// Admission for the R2–R4 transients (parts LDE, trees, DEEP, FRI, the +/// inverted denominators): `true` when admitted. A decline is the caller's +/// documented host arm; over budget is reported so it is never a silent one. +/// Whether that arm may run at all is decided downstream, where a device-only +/// table's recovery goes through [`refuse_host_recovery`]. +fn admit_transient(lde_size: usize, bytes: u64, what: &str) -> bool { + match admit(lde_size, bytes) { + Admission::Admitted { .. } => true, + Admission::NoDevice | Admission::BelowFloor { .. } => false, + Admission::OverBudget { bytes, budget } => { + eprintln!( + "[gpu] {what} declined at LDE {lde_size}: over the VRAM admission budget \ + (needs {bytes} B = {:.2} GiB, budget {budget} B = {:.2} GiB); the host arm runs", + gib(bytes), + gib(budget) + ); + false + } + } +} + +/// The gate on every "download the resident data and continue on the host" +/// recovery of a device-only table. In production that recovery IS the +/// failure to report: it aborts here with the shape and the live VRAM. Under +/// the test-only switch it prints the same line tagged TEST-ONLY and returns, +/// 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() + ); + if test_only_host_fallback() { + eprintln!("[gpu] TEST-ONLY host recovery: {msg}"); + return; + } + panic!( + "[gpu] ABORT: a device-only table would continue on the HOST (host RAM is a cache, \ + not a compute path) — {msg}" + ); +} + /// Test hook: decline the device R2 path unconditionally so device-only /// tables exercise the [`materialize_lde_trace_host`] recovery end to end. +/// Setting it also enables the test-only host fallback +/// ([`test_only_host_fallback`]) — the recovery it exists to exercise would +/// otherwise abort. pub(crate) fn gpu_force_downgrade() -> bool { static CACHED: OnceLock = OnceLock::new(); *CACHED.get_or_init(|| std::env::var("LAMBDA_VM_GPU_FORCE_DOWNGRADE").is_ok_and(|v| v != "0")) @@ -361,15 +757,15 @@ pub(crate) fn device_only_disabled() -> bool { /// requires. /// /// If a precondition is nonetheless violated at runtime (mis-gate or -/// transient GPU error), what happens depends on the round. R2 and the R1 -/// resident-aux commit recover: they 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 — slower, -/// never wrong — aborting only when the resident handles cannot serve the -/// data. R3 and R4 have no such recovery: the R3 barycentric arms assert on -/// the buffer they are about to read and the R4 guards on `host_trace_empty`, -/// both failing loudly rather than reading an empty host trace. +/// 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. /// /// `zerofier_uniform` must be the R1-derived conservative form (all constraints /// share `end_exemptions == 0`), which implies `ZerofierEvaluations::is_uniform` @@ -454,7 +850,11 @@ where } let n = columns[0].len(); let lde_size = n.saturating_mul(blowup_factor); - if lde_size < gpu_lde_threshold() { + // One fresh LDE output per column on device. + let bytes = (lde_size as u64) + .saturating_mul(columns.len() as u64) + .saturating_mul(BASE_BYTES); + if !admit(lde_size, bytes).is_admitted() { return LayoutDispatch::Skip; } if TypeId::of::() != TypeId::of::() { @@ -482,7 +882,8 @@ where } let n = columns[0].len(); let lde_size = n.saturating_mul(blowup_factor); - if lde_size < gpu_lde_threshold() { + // One fresh ext3 LDE output per column on device. + if !admit(lde_size, ext3_bytes(lde_size as u64, columns.len() as u64)).is_admitted() { return LayoutDispatch::Skip; } if TypeId::of::() != TypeId::of::() { @@ -699,15 +1100,17 @@ where let n = h0.len(); let blowup = 2; // extend_half_to_lde extends N → 2N always let lde_size = n * blowup; - if lde_size < gpu_lde_threshold() { - return None; - } if TypeId::of::() != TypeId::of::() { return None; } if TypeId::of::() != TypeId::of::() { return None; } + // Two ext3 inputs staged and two ext3 LDE outputs allocated on device. + let bytes = ext3_bytes(lde_size as u64, 2).saturating_add(ext3_bytes(n as u64, 2)); + if !admit_transient(lde_size, bytes, "R2 extend-halves") { + return None; + } GPU_EXTEND_HALVES_CALLS.fetch_add(1, Ordering::Relaxed); // Weights are built from `g = domain.coset_offset` directly: the // CPU caller previously passed `g²` redundantly. See the @@ -766,10 +1169,11 @@ where } /// Shared admission gate for the device composition-parts producers: the tower -/// must be the Goldilocks/ext3 pair the kernels are written for, and the LDE must -/// be a power of two at or above the commit threshold. Returns the validated LDE -/// size so callers can derive from it. Kept in one place so a future condition -/// (a VRAM check, a tower widening) cannot land on only one of the d=1/d=2 arms. +/// must be the Goldilocks/ext3 pair the kernels are written for, the LDE must +/// be a power of two, and the two ext3 part slabs the decompose allocates must +/// clear [`admit`]. Returns the validated LDE size so callers can derive from +/// it. Kept in one place so a future condition (a tower widening) cannot land +/// on only one of the d=1/d=2 arms. fn dev_comp_parts_gate(num_rows: usize) -> Option where F: IsField + 'static, @@ -781,7 +1185,10 @@ where if TypeId::of::() != TypeId::of::() { return None; } - if num_rows < gpu_lde_threshold() || !num_rows.is_power_of_two() { + if !num_rows.is_power_of_two() { + return None; + } + if !admit_transient(num_rows, ext3_bytes(num_rows as u64, 2), "R2 decompose") { return None; } Some(num_rows) @@ -954,10 +1361,6 @@ where E: IsField + 'static, B: DeviceTreeBackend, { - let lde_size = n.saturating_mul(blowup_factor); - if lde_size < gpu_lde_threshold() { - return None; - } if TypeId::of::() != TypeId::of::() { return None; } @@ -967,6 +1370,15 @@ where if row_major.len() != n * m || m == 0 || n == 0 { return None; } + let lde_size = n.saturating_mul(blowup_factor); + let shape = DispatchShape { + what: "R1 main commit", + n, + base_cols: m, + blowup: blowup_factor, + }; + let set = commit_device_set(n, m, blowup_factor, true); + admit_commit(lde_size, &shape, &set)?; let raw: &[u64] = unsafe { from_raw_parts(row_major.as_ptr() as *const u64, n * m) }; let weights_u64 = unsafe { weights_to_u64::(weights) }; @@ -977,7 +1389,8 @@ where // The keep path keeps the Merkle tree resident on device (in `handle.tree`). // `retain_host_lde=false` additionally skips the row-major D2H (device-only). - let (handle, lde_u64) = math_cuda::lde::coset_lde_row_major_with_merkle_tree_keep( + // Admitted means the device path is the only path: a failure here aborts. + let (handle, lde_u64) = match math_cuda::lde::coset_lde_row_major_with_merkle_tree_keep( raw, predev, device_hash_of::(), @@ -986,8 +1399,17 @@ where blowup_factor, &weights_u64, retain_host_lde, - ) - .ok()?; + ) { + Ok(v) => v, + Err(e) => { + abort_or_test_fallback( + &shape, + Some(&set), + DevicePathFailure::DeviceError(format!("{e:?}")), + ); + return None; + } + }; // Transmute Vec → Vec> (zero-copy, E == GoldilocksField). let lde_out: Vec> = unsafe { @@ -1062,10 +1484,6 @@ where E: IsField + 'static, B: DeviceTreeBackend, { - let lde_size = n.saturating_mul(blowup_factor); - if lde_size < gpu_lde_threshold() { - return None; - } if TypeId::of::() != TypeId::of::() { return None; } @@ -1078,6 +1496,15 @@ where if split_col == 0 || split_col >= m { return None; } + let lde_size = n.saturating_mul(blowup_factor); + let shape = DispatchShape { + what: "R1 main commit (preprocessed split)", + n, + base_cols: m, + blowup: blowup_factor, + }; + let set = commit_device_set(n, m, blowup_factor, true); + admit_commit(lde_size, &shape, &set)?; let raw: &[u64] = unsafe { from_raw_parts(row_major.as_ptr() as *const u64, n * m) }; let weights_u64 = unsafe { weights_to_u64::(weights) }; @@ -1086,7 +1513,8 @@ where GPU_LEAF_HASH_CALLS.fetch_add(1 + build_precomputed as u64, Ordering::Relaxed); GPU_MERKLE_TREE_CALLS.fetch_add(1 + build_precomputed as u64, Ordering::Relaxed); - let (pre_nodes, handle, lde_u64) = math_cuda::lde::coset_lde_row_major_split_trees( + // Admitted means the device path is the only path: a failure here aborts. + let (pre_nodes, handle, lde_u64) = match math_cuda::lde::coset_lde_row_major_split_trees( raw, predev, device_hash_of::(), @@ -1097,8 +1525,17 @@ where split_col, build_precomputed, want_host, - ) - .ok()?; + ) { + Ok(v) => v, + Err(e) => { + abort_or_test_fallback( + &shape, + Some(&set), + DevicePathFailure::DeviceError(format!("{e:?}")), + ); + return None; + } + }; let pre_tree = match pre_nodes { Some(nodes) => Some(tree_from_node_bytes::(nodes)?), @@ -1147,10 +1584,6 @@ where E: IsField + 'static, B: DeviceTreeBackend, { - let lde_size = n.saturating_mul(blowup_factor); - if lde_size < gpu_lde_threshold() { - return None; - } if TypeId::of::() != TypeId::of::() { return None; } @@ -1160,9 +1593,18 @@ where if row_major.len() != n * m || m == 0 || n == 0 { return None; } - // Fp3 = [u64; 3] in memory — reinterpret as flat u64 slice (m3 = m*3). let m3 = m * 3; + let lde_size = n.saturating_mul(blowup_factor); + let shape = DispatchShape { + what: "R1 aux commit", + n, + base_cols: m3, + blowup: blowup_factor, + }; + let set = commit_device_set(n, m3, blowup_factor, false); + admit_commit(lde_size, &shape, &set)?; + let raw: &[u64] = unsafe { from_raw_parts(row_major.as_ptr() as *const u64, n * m3) }; let weights_u64 = unsafe { weights_to_u64::(weights) }; @@ -1172,7 +1614,8 @@ where // The keep path keeps the Merkle tree resident on device (in `handle.tree`). // `retain_host_lde=false` additionally skips the row-major D2H (device-only). - let (handle, lde_u64) = math_cuda::lde::coset_lde_ext3_row_major_with_merkle_tree_keep( + // Admitted means the device path is the only path: a failure here aborts. + let (handle, lde_u64) = match math_cuda::lde::coset_lde_ext3_row_major_with_merkle_tree_keep( raw, device_hash_of::(), n, @@ -1180,8 +1623,17 @@ where blowup_factor, &weights_u64, retain_host_lde, - ) - .ok()?; + ) { + Ok(v) => v, + Err(e) => { + abort_or_test_fallback( + &shape, + Some(&set), + DevicePathFailure::DeviceError(format!("{e:?}")), + ); + return None; + } + }; // Transmute Vec → Vec> (zero-copy, E == Fp3 = [u64;3]). let lde_out: Vec> = unsafe { @@ -1391,10 +1843,14 @@ where if !lde_size.is_power_of_two() || lde_size < 2 { return None; } - if lde_size < gpu_lde_threshold() { + if TypeId::of::() != TypeId::of::() { return None; } - if TypeId::of::() != TypeId::of::() { + // The parts are re-uploaded (`m` ext3 columns over the LDE) and one full + // row-pair tree is built. + let bytes = ext3_bytes(lde_size as u64, lde_parts.len() as u64) + .saturating_add(full_tree_bytes(lde_size as u64)); + if !admit_transient(lde_size, bytes, "R2 composition tree") { return None; } // All parts must have the same LDE length. @@ -1451,8 +1907,15 @@ where if TypeId::of::() != TypeId::of::() { return None; } - if handle.m == 0 || !handle.lde_size.is_power_of_two() || handle.lde_size < gpu_lde_threshold() - { + if handle.m == 0 || !handle.lde_size.is_power_of_two() { + return None; + } + // Only the tree is fresh: the parts are already resident. + if !admit_transient( + handle.lde_size, + full_tree_bytes(handle.lde_size as u64), + "R2 composition tree (resident parts)", + ) { return None; } let be = math_cuda::device::backend().ok()?; @@ -1856,14 +2319,14 @@ pub(crate) static GPU_BATCH_INVERT_CALLS: AtomicU64 = AtomicU64::new(0); /// ([`materialize_lde_trace_host`], the sole function that bumps this — /// entered from the R2 host evaluator, the R3 barycentric arms and the R4 /// DEEP host loop). Nonzero means the device-only gate cleared a table whose -/// downstream dispatch then declined at runtime — the table continued -/// host-backed, correct but slower. A count is either a gate miss (a static -/// condition worth mirroring into the gate) or a transient device decline -/// (VRAM pressure), which by definition cannot be gated out — see -/// [`materialize_lde_trace_host`]'s own note. The R1 resident-aux downgrade -/// is counted by [`GPU_RESIDENT_AUX_DOWNGRADES`] instead: it fires on tables -/// the gate never marked device-only, so summing the two would blame the gate -/// for declines it never made. +/// downstream dispatch then declined at runtime. Only reachable under the +/// test-only host fallback ([`test_only_host_fallback`]): in production the +/// recovery aborts instead. A count is either a gate miss (a static condition +/// worth mirroring into the gate) or a transient device decline (VRAM +/// pressure), which by definition cannot be gated out. The R1 resident-aux +/// downgrade is counted by [`GPU_RESIDENT_AUX_DOWNGRADES`] instead: it fires +/// on tables the gate never marked device-only, so summing the two would +/// blame the gate for declines it never made. pub(crate) static GPU_DEVICE_ONLY_DOWNGRADES: AtomicU64 = AtomicU64::new(0); pub fn gpu_device_only_downgrades() -> u64 { GPU_DEVICE_ONLY_DOWNGRADES.load(Ordering::Relaxed) @@ -1872,7 +2335,8 @@ pub fn gpu_device_only_downgrades() -> u64 { /// 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). Independent +/// ([`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 @@ -1910,14 +2374,18 @@ pub fn gpu_resident_aux_retries() -> u64 { /// 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. The class-level safety -/// net under the device-only gate — a static predicate can never mirror every -/// reason a dynamic dispatch might decline (kernel eligibility, transient -/// errors, shapes a new workload brings), so any miss lands here and degrades -/// to a slower-but-correct CPU round instead of a hard abort. Returns false -/// (→ the caller's abort) when the resident handles cannot serve the data: a -/// missing handle or bound stream, a handle whose shape disagrees with the -/// trace, a failed download or sync, or a field tower with no CUDA lowering. +/// 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 +/// a failure to report, and [`refuse_host_recovery`] aborts here with the +/// shape and the live VRAM. Under the test-only switch +/// ([`test_only_host_fallback`]) the download proceeds and the table continues +/// on the host arms — the fault-injection suites assert on exactly that. +/// Returns false (→ the caller's abort) when the resident handles cannot serve +/// the data: a missing handle or bound stream, a handle whose shape disagrees +/// with the trace, a failed download or sync, or a field tower with no CUDA +/// lowering. pub(crate) fn materialize_lde_trace_host( lde_trace: &mut crate::trace::LDETraceTable, ) -> bool @@ -1928,6 +2396,12 @@ where if !lde_trace.host_trace_empty() { return true; } + refuse_host_recovery( + "R2/R3/R4 host arm on a device-only trace", + lde_trace.num_rows(), + lde_trace.num_main_cols(), + lde_trace.num_aux_cols(), + ); if !is_goldilocks_ext3_tower::() { return false; } @@ -2087,7 +2561,9 @@ 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 declines at runtime. +/// 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, @@ -2100,6 +2576,13 @@ where 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; }; @@ -2260,6 +2743,12 @@ where let Some(h) = lde_trace.gpu_composition_parts() else { return false; }; + refuse_host_recovery( + "host consumer of the composition parts on a device-only table", + h.lde_size, + lde_trace.num_main_cols(), + lde_trace.num_aux_cols(), + ); let Some(stream) = lde_trace.bound_stream() else { return false; }; @@ -2379,9 +2868,6 @@ where return None; } let lde_size = domain_size.checked_mul(blowup_factor)?; - if lde_size < gpu_lde_threshold() { - return None; - } if TypeId::of::() != TypeId::of::() { return None; } @@ -2389,6 +2875,12 @@ where return None; } let m = parts_coefs.len(); + // `m` ext3 coefficient inputs staged and `m` ext3 LDE parts kept resident. + let bytes = ext3_bytes(lde_size as u64, m as u64) + .saturating_add(ext3_bytes(domain_size as u64, m as u64)); + if !admit_transient(lde_size, bytes, "R2 parts LDE") { + return None; + } let mut weights_u64 = Vec::with_capacity(domain_size); let mut w = FieldElement::::one(); @@ -2476,6 +2968,16 @@ where { return None; } + // 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 { + what: "R1 aux commit (resident)", + n: ra.num_rows, + base_cols: ra.num_aux_cols * 3, + blowup: blowup_factor, + }; + let set = commit_device_set(ra.num_rows, ra.num_aux_cols * 3, blowup_factor, false); + admit_resident_commit(&shape, &set)?; let weights_u64 = unsafe { weights_to_u64::(weights) }; GPU_LDE_CALLS.fetch_add((ra.num_aux_cols * 3) as u64, Ordering::Relaxed); @@ -2492,8 +2994,10 @@ where retain_host_lde, ) .inspect_err(|e| { - // Surface the swallowed driver error (e.g. OOM): the caller drains - // the device and retries, then downgrades the table to the host path. + // 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`. eprintln!( "[gpu] resident aux LDE failed (rows={} cols={} blowup={}): {e:?}", ra.num_rows, ra.num_aux_cols, blowup_factor @@ -2578,9 +3082,6 @@ where } let main = lde_trace.gpu_main()?; let lde_size = main.lde_size; - if lde_size < gpu_lde_threshold() { - return None; - } if !lde_size.is_power_of_two() { return None; } @@ -2592,6 +3093,12 @@ where if h_ood.len() != num_parts { return None; } + // The DEEP codeword, plus — on the host arms — the parts and the inverted + // denominators staged onto the device (an upper bound over the three arms). + let bytes = ext3_bytes(lde_size as u64, (2 + num_parts + num_eval_points) as u64); + if !admit_transient(lde_size, bytes, "R4 DEEP") { + return None; + } if trace_ood_columns.len() != num_total_cols || trace_ood_columns.iter().any(|c| c.len() != num_eval_points) { @@ -2784,7 +3291,15 @@ where } let main = lde_trace.gpu_main()?; let lde_size = main.lde_size; - if lde_size < gpu_lde_threshold() || !lde_size.is_power_of_two() { + if !lde_size.is_power_of_two() { + return None; + } + // Everything is resident but the DEEP codeword itself. + if !admit_transient( + lde_size, + ext3_bytes(lde_size as u64, 1), + "R4 DEEP (resident)", + ) { return None; } let num_main = main.m; @@ -2878,7 +3393,9 @@ where return None; } let total = n.checked_mul(k_scalars)?; - if total < gpu_lde_threshold() { + // The inverted denominators (`total` ext3) plus the coset points. + let bytes = ext3_bytes(total as u64, 1).saturating_add((n as u64).saturating_mul(BASE_BYTES)); + if !admit_transient(total, bytes, "batch invert") { return None; } @@ -3080,7 +3597,9 @@ where return None; } let total = n.checked_mul(k_scalars)?; - if total < gpu_lde_threshold() { + // The inverted denominators (`total` ext3) plus the coset points. + let bytes = ext3_bytes(total as u64, 1).saturating_add((n as u64).saturating_mul(BASE_BYTES)); + if !admit_transient(total, bytes, "R3 context") { return None; } @@ -3160,7 +3679,10 @@ where if n0 != domain_size || !n0.is_power_of_two() || n0 < 2 { return None; } - if n0 < gpu_lde_threshold() { + // The evals upload, the geometric layer chain (bounded by one more + // codeword) and the layer trees (bounded by one full tree). + let bytes = ext3_bytes(n0 as u64, 2).saturating_add(full_tree_bytes(n0 as u64)); + if !admit_transient(n0, bytes, "R4 FRI commit") { return None; } // Mismatched twiddles would panic inside `FriCommitState::new`; gate here @@ -3232,7 +3754,12 @@ where return None; } let n0 = codeword.n; - if !n0.is_power_of_two() || n0 < 2 || n0 < gpu_lde_threshold() { + if !n0.is_power_of_two() || n0 < 2 { + return None; + } + // The layer chain and its trees; the codeword is already resident. + let bytes = ext3_bytes(n0 as u64, 1).saturating_add(full_tree_bytes(n0 as u64)); + if !admit_transient(n0, bytes, "R4 FRI commit (resident)") { return None; } // Mismatched twiddles would panic inside `FriCommitState::new_dev`; @@ -3497,6 +4024,157 @@ 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 +/// input stays at 1.2 GiB: +/// +/// ```text +/// LAMBDA_VM_VRAM_BUDGET_MB=1024 cargo test -p stark --release --features cuda \ +/// --lib gpu_lde::admission_box_tests -- --ignored --nocapture +/// ``` +#[cfg(all(test, feature = "cuda"))] +mod admission_box_tests { + use super::*; + use crate::config::BatchedMerkleTreeBackend; + + type F = GoldilocksField; + type Fp = FieldElement; + + /// 2^18 rows x 612 columns at blowup 2: a 2.4 GiB LDE, 3.7 GiB device set + /// against the 1 GiB budget the command line sets. Must abort with the + /// over-budget diagnostic, never commit on the host. + #[test] + #[ignore = "requires GPU and LAMBDA_VM_VRAM_BUDGET_MB=1024 in its own process"] + #[should_panic(expected = "over the VRAM admission budget")] + fn an_over_budget_commit_aborts_instead_of_committing_on_the_host() { + let n: usize = 1 << 18; + let m: usize = 612; + let blowup: usize = 2; + let budget = device_vram_budget_bytes().expect("cuda backend"); + assert!( + budget <= 2 << 30, + "this test needs LAMBDA_VM_VRAM_BUDGET_MB=1024 set before the backend initialises \ + (budget is {budget} B)" + ); + 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, + ); + panic!( + "the over-budget commit returned {} instead of aborting", + if committed.is_some() { + "a device handle" + } else { + "None (a silent host commit)" + } + ); + } +} + /// GPU↔CPU parity for the preprocessed split-tree commit path. Requires the /// `cuda` feature and a visible GPU (skipped otherwise via the dispatch gate /// returning `None` — asserted here, so a silent skip fails the test). diff --git a/prover/src/lfm/chunking.rs b/prover/src/lfm/chunking.rs index 3b15cf031..d43db1faa 100644 --- a/prover/src/lfm/chunking.rs +++ b/prover/src/lfm/chunking.rs @@ -1,9 +1,11 @@ //! Row chunking — how the machine's splittable tables scale past one instance. //! -//! Two chips are chunked, by the same mechanism and for the same reason: -//! [`KeccakChunking`] splits `KECCAK_RND`, and [`Blake3Chunking`] splits -//! `LFM_BLAKE3`. Everything the next paragraphs say about the first holds for -//! the second; the differences are collected under [`Blake3Chunking`]. +//! Three chips are chunked, by the same mechanism and for the same reason: +//! [`KeccakChunking`] splits `KECCAK_RND`, [`Blake3Chunking`] splits +//! `LFM_BLAKE3`, and [`BaluChunking`] splits `LFM_BALU` — the first two because +//! their matrices are WIDE, the third because its matrix is TALL. Everything +//! the next paragraphs say about the first holds for the others; the +//! differences are collected under [`Blake3Chunking`] and [`BaluChunking`]. //! //! `KECCAK_RND` costs 24 rows per permutation at 1480 columns, so a single //! instance saturates a 2^19-row table at ~21.8k permutations while a real @@ -114,6 +116,21 @@ impl Default for KeccakChunking { } } +/// Chunks `total` one-row records need at `per` records per chunk — never +/// zero, so a chip stays present for an empty program. The one rule +/// [`Blake3Chunking`] and [`BaluChunking`] share. +fn row_chunk_count(per: usize, total: usize) -> usize { + total.div_ceil(per).max(1) +} + +/// The half-open record range chunk `chunk` covers at `per` records per +/// chunk, clamped to `total`. +fn row_chunk_range(per: usize, total: usize, chunk: usize) -> core::ops::Range { + let start = per.saturating_mul(chunk).min(total); + let end = start.saturating_add(per).min(total); + start..end +} + /// The environment knob that turns `LFM_BLAKE3` chunking on, read at program /// EMISSION time by the driver that emits the program. /// @@ -266,23 +283,14 @@ impl Blake3Chunking { /// The chip MASK, not this, is what drops an unused family; see /// [`ChipSet::blake3_chunks`](super::airs::ChipSet::blake3_chunks). pub fn chunk_count(self, num_compressions: usize) -> usize { - num_compressions - .div_ceil(self.compressions_per_chunk) - .max(1) + row_chunk_count(self.compressions_per_chunk, num_compressions) } /// The half-open row range chunk `chunk` covers, clamped to /// `num_compressions`. The single rule the group split, the record split and /// the census heights all read, so they cannot disagree about a boundary. pub fn chunk_range(self, num_compressions: usize, chunk: usize) -> core::ops::Range { - let start = self - .compressions_per_chunk - .saturating_mul(chunk) - .min(num_compressions); - let end = start - .saturating_add(self.compressions_per_chunk) - .min(num_compressions); - start..end + row_chunk_range(self.compressions_per_chunk, num_compressions, chunk) } /// Splits per-compression records into exactly [`Self::chunk_count`] @@ -300,6 +308,180 @@ impl Default for Blake3Chunking { } } +/// The environment knob that turns `LFM_BALU` chunking on, read at program +/// EMISSION time like [`BLAKE3_MAX_CHUNK_ROWS_LOG2_ENV`]. +/// +/// Unset means one table — today's machine, byte for byte. Set to `k` means +/// chunks of at most `2^k` rows, i.e. `2^k` ALU operations. +pub const BALU_MAX_CHUNK_ROWS_LOG2_ENV: &str = "LFM_BALU_MAX_CHUNK_ROWS_LOG2"; + +/// Trace rows one ALU operation occupies in `LFM_BALU` — exactly one +/// (`emit_column_groups` opens one row per `Balu` instruction). +pub const BALU_ROWS_PER_OP: usize = 1; + +/// The chunk height the sizing under [`BaluChunking`] arrives at: `2^22` rows. +pub const BALU_TARGET_CHUNK_ROWS_LOG2: u32 = 22; + +/// How a program's ALU operations are distributed over `LFM_BALU` instances — +/// the row-chunking arm for the machine's TALL-NARROW chips. +/// +/// # Why this chip needs it +/// +/// `LFM_BALU` is one row per operation at 4 value + 10 preprocessed columns: +/// narrow, and on the aggregator very tall — the program pads to `2^27` rows +/// at 110 queries and `2^28` at 219, a census contribution that is trivial in +/// cells and enormous in rows. At blowup 2 the ONE-table device set of the R1 +/// commit alone is +/// +/// | rows | LDE `2n·14·8` | snapshot `n·14·8` | tree `(2n−1)·32` | R1 set | +/// |------|---------------|-------------------|------------------|--------| +/// | 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. +/// +/// # Why row chunking, not column streaming +/// +/// Streaming the commit column group by column group lowers only the commit's +/// own peak; rounds 2–4 read every column of every row from the RESIDENT LDE, +/// so the `2n · cols · 8` buffer has to be on the card for the whole table +/// however the leaves were hashed. For a tall-narrow chip that buffer IS the +/// problem (28 GiB at `2^27`), and nothing short of splitting the rows shrinks +/// it. Column streaming is the shape for SHORT-WIDE chips (`LFM_HASH` at +/// `2^21 × 449`), and since the fused commit transposes its one LDE buffer in +/// place it buys little even there: the LDE stays resident for rounds 2–4 +/// either way. +/// +/// # Why splitting the rows is free +/// +/// The property [`KeccakChunking`] and [`Blake3Chunking`] rest on, checked on +/// this chip: every constraint `BaluConstraints` emits reads `main(0, ..)` — +/// no row-to-row coupling — and every bus interaction is a within-row `LfmMem` +/// token gated by the row's own selector or multiplicity column. Operands and +/// results travel by address matching, and the addresses are PREPROCESSED +/// program data, so which instance a row lives in is invisible to the balance. +/// +/// # What it costs +/// +/// Like `LFM_BLAKE3` and unlike `KECCAK_RND`, this chip carries a preprocessed +/// instruction group, so **each chunk is its own committed matrix with its +/// own root and its own height**: a chunked program is a different program +/// identity by name, and the chunk roots ride the artifacts and fold into +/// `program_id` exactly as the BLAKE3 chunk roots do. Wiring — the program +/// field, the per-chunk group, the AIR instances, the artifact roots and the +/// slot map — follows the BLAKE3 template one for one and is not in this +/// 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. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct BaluChunking { + ops_per_chunk: usize, +} + +impl BaluChunking { + /// One table, whatever the operation count — **the default**, and the + /// machine as it stands before the aggregator is emitted chunked. + /// `usize::MAX` rather than an `Option` for [`Blake3Chunking::unbounded`]'s + /// reason: one code path. + pub const fn unbounded() -> Self { + Self { + ops_per_chunk: usize::MAX, + } + } + + /// The policy that fills chunks to at most `max_rows` trace rows. + pub const fn from_max_rows(max_rows: usize) -> Self { + let ops_per_chunk = max_rows / BALU_ROWS_PER_OP; + assert!( + ops_per_chunk > 0, + "an LFM_BALU chunk must hold at least one operation" + ); + Self { ops_per_chunk } + } + + /// The policy that puts at most `ops_per_chunk` operations in each chunk. + /// The small-limit constructor tests use to force several chunks out of a + /// tiny program. + pub const fn from_ops(ops_per_chunk: usize) -> Self { + assert!( + ops_per_chunk > 0, + "an LFM_BALU chunk must hold at least one operation" + ); + Self { ops_per_chunk } + } + + /// The sizing's target: chunks of `2^22` rows + /// ([`BALU_TARGET_CHUNK_ROWS_LOG2`]). + pub const fn target() -> Self { + Self::from_max_rows(1usize << BALU_TARGET_CHUNK_ROWS_LOG2) + } + + /// The policy [`BALU_MAX_CHUNK_ROWS_LOG2_ENV`] names, or `None` when it is + /// unset — on [`Blake3Chunking::from_env`]'s terms, including the panic on + /// a value that is not a row exponent. + pub fn from_env() -> Option { + Self::from_env_value(std::env::var(BALU_MAX_CHUNK_ROWS_LOG2_ENV).ok().as_deref()) + } + + /// [`Self::from_env`] with the variable's value supplied, so the parse is + /// testable without mutating process-global state. + pub fn from_env_value(raw: Option<&str>) -> Option { + let raw = raw?; + let log2: u32 = raw.parse().unwrap_or_else(|_| { + panic!("{BALU_MAX_CHUNK_ROWS_LOG2_ENV} must be a base-2 row exponent, got {raw:?}") + }); + assert!( + log2 < usize::BITS, + "{BALU_MAX_CHUNK_ROWS_LOG2_ENV}={log2} is not a representable row count" + ); + Some(Self::from_max_rows(1usize << log2)) + } + + pub const fn ops_per_chunk(self) -> usize { + self.ops_per_chunk + } + + /// Number of `LFM_BALU` instances a program with `num_ops` operations gets + /// — never zero, so the chip is present (and its constraints verified) + /// even for a program containing no ALU operation at all. + pub fn chunk_count(self, num_ops: usize) -> usize { + row_chunk_count(self.ops_per_chunk, num_ops) + } + + /// The half-open row range chunk `chunk` covers, clamped to `num_ops` — + /// the single rule a group split, a record split and a census height + /// would all read. + pub fn chunk_range(self, num_ops: usize, chunk: usize) -> core::ops::Range { + row_chunk_range(self.ops_per_chunk, num_ops, chunk) + } + + /// Splits per-operation records into exactly [`Self::chunk_count`] slices. + pub fn split(self, ops: &[T]) -> Vec<&[T]> { + (0..self.chunk_count(ops.len())) + .map(|c| &ops[self.chunk_range(ops.len(), c)]) + .collect() + } +} + +impl Default for BaluChunking { + fn default() -> Self { + Self::unbounded() + } +} + #[cfg(test)] mod tests { use super::*; @@ -459,4 +641,96 @@ mod tests { assert_eq!(c.chunk_range(0, 0), 0..0); } } + + /// The default is ONE `LFM_BALU` table at any scale — the machine as it + /// stands. + #[test] + fn the_balu_default_is_a_single_table() { + let c = BaluChunking::default(); + assert_eq!(c, BaluChunking::unbounded()); + for n in [0usize, 1, 1 << 27, 1 << 28, usize::MAX - 1] { + assert_eq!(c.chunk_count(n), 1, "n={n} must stay one table"); + } + assert_eq!(BaluChunking::from_env_value(None), None); + } + + /// The target policy splits the aggregator's `2^27` (110 q) and `2^28` + /// (219 q) rows into 32 and 64 chunks of `2^22`; the knob names the same + /// policy. + #[test] + fn the_balu_target_sizes_the_aggregator() { + let c = BaluChunking::target(); + assert_eq!(c.ops_per_chunk(), 1 << 22); + assert_eq!(c.chunk_count(1 << 27), 32); + assert_eq!(c.chunk_count(1 << 28), 64); + assert_eq!(c.chunk_count((1 << 27) + 1), 33); + assert_eq!(BaluChunking::from_env_value(Some("22")), Some(c)); + for log2 in [0usize, 3, 18, 22] { + assert_eq!( + BaluChunking::from_env_value(Some(&log2.to_string())), + Some(BaluChunking::from_max_rows(1 << log2)), + "{log2} must name 2^{log2} rows per chunk" + ); + } + } + + /// 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. + #[test] + fn the_balu_chunk_sizing_is_the_doc() { + 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 + }; + 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}"); + assert!( + 8 * chunk <= 32 * GIB / 5 * 4, + "eight chunks must fit the budget" + ); + let big_chunk = whole_prove_set(1 << 24); + assert!(2 * big_chunk <= 32 * GIB / 5 * 4 && 3 * big_chunk > 32 * GIB / 5 * 4); + } + + /// `split`, `chunk_count` and `chunk_range` are one rule seen three times + /// for this chip too. + #[test] + fn balu_split_agrees_with_chunk_count_and_range() { + for per in [1usize, 2, 3, 5, 8] { + let c = BaluChunking::from_ops(per); + for n in 0..40usize { + let ops: Vec = (0..n).collect(); + let split = c.split(&ops); + assert_eq!(split.len(), c.chunk_count(n), "per={per} n={n}"); + assert_eq!(split.iter().map(|s| s.len()).sum::(), n); + assert!(split.iter().all(|s| s.len() <= per)); + for (i, s) in split.iter().enumerate() { + assert_eq!(&ops[c.chunk_range(n, i)], *s, "per={per} n={n} chunk {i}"); + } + } + assert_eq!(c.chunk_count(0), 1); + assert_eq!(c.chunk_range(0, 0), 0..0); + } + } + + /// A typo stops the run rather than silently proving a different shape. + #[test] + #[should_panic(expected = "must be a base-2 row exponent")] + fn a_malformed_balu_knob_panics() { + let _ = BaluChunking::from_env_value(Some("2^22")); + } }