diff --git a/crypto/stark/src/gpu_lde.rs b/crypto/stark/src/gpu_lde.rs index 98830fcc7..ced0bc530 100644 --- a/crypto/stark/src/gpu_lde.rs +++ b/crypto/stark/src/gpu_lde.rs @@ -26,6 +26,8 @@ use math::field::extensions_goldilocks::Degree3GoldilocksExtensionField; use math::field::goldilocks::GoldilocksField; use math::field::traits::{IsFFTField, IsField, IsSubFieldOf}; use math::traits::AsBytes; +#[cfg(feature = "parallel")] +use rayon::prelude::{IndexedParallelIterator, ParallelIterator, ParallelSliceMut}; use crate::config::{Commitment, FriLayerMerkleTreeBackend}; use crate::domain::Domain; @@ -54,6 +56,38 @@ fn gpu_lde_threshold() -> usize { }) } +/// Serialize the SUBMISSION of the device R2 window (constraint eval + +/// decompose) across tables. Concurrent R2 windows under VRAM pressure can +/// transiently corrupt a whole H buffer (root mechanism unidentified; reruns +/// on the same resident inputs come out correct), yielding a proof that fails +/// verification. Holding this lock empirically suppresses that at negligible +/// cost — the windows rarely overlap. +/// +/// How much it enforces depends on the table. One that keeps its host trace +/// ends the window in a blocking D2H (the `want_host` arm of +/// [`try_decompose_extend_d2_dev`]), so the guard is held until that table's +/// kernels have completed — a real execution barrier. A device-only table's +/// window is enqueue-only, so two tables' R2 kernels can still overlap on +/// device; what the lock orders there is submission and allocation, which is +/// enough to suppress the corruption in practice but is not a guarantee that +/// R2 kernels never run concurrently. +/// +/// `LAMBDA_VM_GPU_SERIALIZE_R2=0` disables the lock (e.g. to bisect or once +/// the underlying race is fixed). +pub(crate) fn r2_serialize_guard() -> Option> { + static ENABLED: OnceLock = OnceLock::new(); + static LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(()); + if *ENABLED.get_or_init(|| !std::env::var("LAMBDA_VM_GPU_SERIALIZE_R2").is_ok_and(|v| v == "0")) + { + // The guarded state is (), so a panic while holding the lock carries + // no information — recover instead of burying the original panic + // under a cascade of PoisonErrors from every other table. + Some(LOCK.lock().unwrap_or_else(|e| e.into_inner())) + } else { + None + } +} + /// Incremented by the `try_expand_*` functions per base-field column handed to /// the GPU dispatch (an ext3 column counts as 3, one per base component), /// before the GPU call. A failed call returns without decrementing it, so it @@ -82,6 +116,7 @@ pub fn reset_all_gpu_call_counters() { GPU_COMPOSITION_CALLS.store(0, Ordering::Relaxed); GPU_OPENING_GATHER_CALLS.store(0, Ordering::Relaxed); GPU_DEVICE_ONLY_CALLS.store(0, Ordering::Relaxed); + GPU_DEVICE_ONLY_DOWNGRADES.store(0, Ordering::Relaxed); } pub(crate) static GPU_EXTEND_HALVES_CALLS: AtomicU64 = AtomicU64::new(0); @@ -171,12 +206,23 @@ pub(crate) fn device_only_disabled() -> bool { /// Stage-3 device-only gate: `true` when a table's round-1 LDE can be left /// device-resident (host D2H skipped) because every downstream round is -/// guaranteed to take its GPU path. A strict AND of all preconditions that -/// imply the R2 composition, R3 barycentric, R4 DEEP, and R4 opening GPU paths -/// all fire and read the device LDE. The per-round `host_trace_empty` -/// hard-abort guards are the safety net: if any precondition is nonetheless -/// violated at runtime (mis-gate or transient GPU error), the prove aborts -/// loudly rather than reading the empty host trace. +/// guaranteed to take its GPU path. A strict AND of the numeric and shape +/// preconditions that imply the R2 composition, R3 barycentric, R4 DEEP, and +/// R4 opening GPU paths all fire and read the device LDE — but not the whole +/// predicate on its own: the caller `IsStarkProver::device_only_for` +/// (prover.rs) adds the AIR-level preconditions this signature does not +/// carry, notably the d=2 quotient part count the device-resident R2 path +/// 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 +/// [`GPU_DEVICE_ONLY_DOWNGRADES`] 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. /// /// `zerofier_uniform` must be the R1-derived conservative form (all constraints /// share `end_exemptions == 0`), which implies `ZerofierEvaluations::is_uniform` @@ -184,8 +230,11 @@ pub(crate) fn device_only_disabled() -> bool { /// /// LOCKSTEP: this gate must IMPLY the runtime dispatch checks in /// `ConstraintEvaluator::try_evaluate_composition_gpu` (plus the R3/R4 device -/// arms). A fallback condition added to a dispatch without a mirror here turns -/// every gate-true table into a hard-abort — loud, but an avoidable crash. +/// arms). A fallback condition added to a dispatch without a mirror here +/// costs every gate-true table either a hard-abort at R3/R4 — loud, but an +/// avoidable crash — or, at R2 and the R1 resident-aux commit, a silent +/// downgrade to the host path, which is what [`GPU_DEVICE_ONLY_DOWNGRADES`] +/// exists to surface. pub(crate) fn device_only_gate( lde_size: usize, n: usize, @@ -1413,6 +1462,242 @@ pub fn gpu_fri_calls() -> u64 { /// are counted here, so a single failed dispatch does not necessarily lower /// the total; R3's fallbacks are CPU-only, so a failure there does. pub(crate) static GPU_BATCH_INVERT_CALLS: AtomicU64 = AtomicU64::new(0); +/// Times a table had to fall back to a host trace whose data first had to be +/// downloaded off the device, because a device path declined at runtime (see +/// [`materialize_lde_trace_host`] and [`materialize_aux_trace_host`]). +/// Nonzero means a device dispatch declined and the table continued +/// host-backed — correct but slower. Not every one is a gate miss: the R1 +/// resident-aux site is entered whenever `aux_resident()` is set, whatever +/// the device-only gate said, so it also counts declines on tables that were +/// never device-only. Mirroring the missing condition into the gate is the +/// fix for the device-only case; a resident-aux decline is usually transient +/// VRAM pressure instead. +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) +} + +/// Recover a device-only table for the host path: download the resident main +/// and aux LDEs from their device handles into the host buffers and clear the +/// device-only flag. A side whose host buffer is already populated (a mixed +/// state: one commit fell back to CPU while the other stayed device-only) is +/// kept as is — only the missing side is downloaded. 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. +pub(crate) fn materialize_lde_trace_host( + lde_trace: &mut crate::trace::LDETraceTable, +) -> bool +where + F: IsField + IsSubFieldOf + 'static, + E: IsField + 'static, +{ + if !lde_trace.host_trace_empty() { + return true; + } + if !is_goldilocks_ext3_tower::() { + return false; + } + let Some(stream) = lde_trace.bound_stream() else { + return false; + }; + + // Main: column-major device buf -> row-major host Vec. An empty Vec tells + // `set_host_data` to keep the buffer that is already there. + let main_data: Vec> = + if lde_trace.num_main_cols() == 0 || !lde_trace.main_data.is_empty() { + Vec::new() + } else { + let Some(h) = lde_trace.gpu_main() else { + return false; + }; + if h.m != lde_trace.num_main_cols() || h.lde_size != lde_trace.num_rows() { + return false; + } + let Some(data) = download_main_lde_row_major::(h, &stream) else { + return false; + }; + data + }; + + // Aux: de-interleaved ext3 slabs -> row-major interleaved host Vec. + let aux_data: Vec> = + if lde_trace.num_aux_cols() == 0 || !lde_trace.aux_data.is_empty() { + Vec::new() + } else { + let Some(h) = lde_trace.gpu_aux() else { + return false; + }; + if h.m != lde_trace.num_aux_cols() || h.lde_size != lde_trace.num_rows() { + return false; + } + if h.wait_ready_on(&stream).is_err() { + return false; + } + let Ok(slabs) = stream.clone_dtoh(h.buf.as_ref()) else { + return false; + }; + if stream.synchronize().is_err() { + return false; + } + let (m, lde) = (h.m, h.lde_size); + // Short download: degrade like the sibling paths + // (`download_main_lde_row_major`, `materialize_aux_trace_host`) + // rather than panic on the slab slicing below. + if slabs.len() != m * lde * 3 { + return false; + } + // Parallel de-interleaved slabs → row-major interleaved: each row + // chunk gathers from the source slabs independently. + let mut interleaved = vec![0u64; m * lde * 3]; + if m > 0 { + #[cfg(feature = "parallel")] + { + interleaved + .par_chunks_exact_mut(m * 3) + .enumerate() + .for_each(|(r, dst)| { + for (c, dst_col) in dst.chunks_exact_mut(3).enumerate() { + for (k, d) in dst_col.iter_mut().enumerate() { + *d = slabs[(c * 3 + k) * lde + r]; + } + } + }); + } + #[cfg(not(feature = "parallel"))] + { + for (r, dst) in interleaved.chunks_exact_mut(m * 3).enumerate() { + for (c, dst_col) in dst.chunks_exact_mut(3).enumerate() { + for (k, d) in dst_col.iter_mut().enumerate() { + *d = slabs[(c * 3 + k) * lde + r]; + } + } + } + } + } + // SAFETY: E == Ext3 per the tower check; FieldElement backing + // is [u64; 3]. + unsafe { + let mut v = std::mem::ManuallyDrop::new(interleaved); + debug_assert!( + v.len().is_multiple_of(3) && v.capacity().is_multiple_of(3), + "interleaved len/capacity must be a multiple of 3 for Fp3 reinterpret" + ); + Vec::from_raw_parts( + v.as_mut_ptr() as *mut FieldElement, + v.len() / 3, + v.capacity() / 3, + ) + } + }; + + lde_trace.set_host_data(main_data, aux_data); + GPU_DEVICE_ONLY_DOWNGRADES.fetch_add(1, Ordering::Relaxed); + true +} + +/// Download a resident main LDE (column-major device buf) into the row-major +/// host Vec the CPU rounds read. Shared by the R1 and R2 downgrade paths. +pub(crate) fn download_main_lde_row_major( + h: &math_cuda::lde::GpuLdeBase, + stream: &std::sync::Arc, +) -> Option>> +where + F: IsField + 'static, +{ + if TypeId::of::() != TypeId::of::() { + return None; + } + h.wait_ready_on(stream).ok()?; + let col_major = stream.clone_dtoh(h.buf.as_ref()).ok()?; + stream.synchronize().ok()?; + let (m, lde) = (h.m, h.lde_size); + if col_major.len() != m * lde { + return None; + } + // Parallel col-major → row-major transpose: each row chunk gathers from + // the source columns independently. + let mut row_major = vec![0u64; m * lde]; + if m > 0 { + #[cfg(feature = "parallel")] + { + row_major + .par_chunks_exact_mut(m) + .enumerate() + .for_each(|(r, dst)| { + for (c, d) in dst.iter_mut().enumerate() { + *d = col_major[c * lde + r]; + } + }); + } + #[cfg(not(feature = "parallel"))] + { + for (r, dst) in row_major.chunks_exact_mut(m).enumerate() { + for (c, d) in dst.iter_mut().enumerate() { + *d = col_major[c * lde + r]; + } + } + } + } + // SAFETY: F == Goldilocks (gated above); FieldElement is + // #[repr(transparent)] over u64. + Some(unsafe { + let mut v = std::mem::ManuallyDrop::new(row_major); + Vec::from_raw_parts( + v.as_mut_ptr() as *mut FieldElement, + v.len(), + v.capacity(), + ) + }) +} + +/// 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. +pub(crate) fn materialize_aux_trace_host(trace: &mut crate::trace::TraceTable) -> bool +where + F: IsField + IsSubFieldOf + 'static, + E: IsField + 'static, +{ + if !is_goldilocks_ext3_tower::() { + return false; + } + let (buf, rows, cols) = match trace.aux_resident.as_ref() { + Some(ra) => (ra.buf.clone(), ra.num_rows, ra.num_aux_cols), + None => return false, + }; + let Ok(be) = math_cuda::device::backend() else { + return false; + }; + let stream = be.next_stream(); + let Ok(raw) = stream.clone_dtoh(buf.as_ref()) else { + return false; + }; + if stream.synchronize().is_err() || raw.len() != rows * cols * 3 { + return false; + } + let data = u64_to_ext3_vec::(&raw); + trace.aux_table = crate::table::Table::new(data, cols); + trace.num_aux_columns = cols; + // The declined device LDE attempt can leave kernels enqueued on another + // stream still reading this buffer; its owning stream is long idle, so + // dropping here would complete the stream-ordered free immediately and + // the pool could hand the memory to a concurrent table's allocation + // while those kernels run. Drain the device before the drop — this is a + // rare recovery path. + if be.ctx.synchronize().is_err() { + return false; + } + trace.aux_resident = None; + GPU_DEVICE_ONLY_DOWNGRADES.fetch_add(1, Ordering::Relaxed); + true +} + pub fn gpu_batch_invert_calls() -> u64 { GPU_BATCH_INVERT_CALLS.load(Ordering::Relaxed) } @@ -1586,8 +1871,8 @@ where retain_host_lde, ) .inspect_err(|e| { - // This path has no CPU fallback (the host aux trace is empty), so the - // caller hard-aborts; surface the swallowed driver error (e.g. OOM). + // Surface the swallowed driver error (e.g. OOM): the caller drains + // the device and retries, then downgrades the table to the host path. eprintln!( "[gpu] resident aux LDE failed (rows={} cols={} blowup={}): {e:?}", ra.num_rows, ra.num_aux_cols, blowup_factor diff --git a/crypto/stark/src/prover.rs b/crypto/stark/src/prover.rs index 4047458bc..8659cf730 100644 --- a/crypto/stark/src/prover.rs +++ b/crypto/stark/src/prover.rs @@ -310,8 +310,15 @@ where // safety property — if the `device_only` gate held but the GPU keep path // fell back to CPU, the buffer is populated and this stays false, so the // proof runs on the host trace as normal. A mixed state (one buffer - // empty, the other full) is treated as device-only so any host read - // hard-aborts rather than indexing an empty buffer. + // empty, the other full) still sets the flag, and is legal rather than + // an error: the aux commit may be more conservative than the main one + // (never less), so an aux side that kept its host copy can sit next to + // a device-only main. The R3 barycentric arms therefore guard on the + // individual buffer — the side that still holds host data stays + // readable — while the flag keeps the R4 and host-evaluator guards + // armed. Reading the real state also picks up an R1 resident-aux + // downgrade: it repopulates the host buffers before this point, so the + // flag simply comes out false. #[cfg(feature = "cuda")] let main_empty = num_main_cols > 0 && main_data.is_empty(); #[cfg(feature = "cuda")] @@ -1010,29 +1017,40 @@ pub trait IsStarkProver< } /// Stage-3 device-only gate for one table (see - /// [`crate::gpu_lde::device_only_gate`]). Derived purely from the AIR + - /// domain so the round-1 main-commit and aux-commit closures compute the - /// identical value and skip both host D2Hs consistently — the per-table - /// `host_trace_empty` flag covers both the main and aux buffers, so they - /// must be left empty together. + /// [`crate::gpu_lde::device_only_gate`]). Derived from the AIR + domain; + /// the main commit uses it as is, while the aux commit additionally + /// requires the main commit to have produced a device handle — the aux + /// side may be more conservative than the main side (never less), which + /// keeps a mixed GPU-aux/CPU-main state out. #[cfg(feature = "cuda")] fn device_only_for( air: &dyn AIR, domain: &Domain, ) -> bool { // Preconditions the downstream GPU paths require that the numeric gate - // below does not capture. A table missing either would pass the gate, - // skip its host D2H, then hard-abort in round 2: + // below does not capture. A table missing any of them would pass the + // gate and skip its host D2H, leaving round 2 to recover through + // `materialize_lde_trace_host` — correct, but a downgrade, and an + // abort if the resident handles cannot serve the data: // - R2 composition unconditionally needs a device aux handle // (`gpu_aux()?`), so the table must declare an aux trace. // - The composition path needs a uniform zerofier with ≥1 group. An // empty constraint set makes `all(end_exemptions == 0)` vacuously // true here but `is_uniform()` false downstream (0 groups). + // - The device-resident R2 path exists only for the d=2 quotient + // decomposition, checked below once `n` is in hand. if !air.has_aux_trace() || air.constraints_meta().is_empty() { return false; } - let lde_size = domain.interpolation_domain_size * domain.blowup_factor; let n = domain.interpolation_domain_size; + // The device-resident R2 path only exists for the d=2 quotient + // decomposition; any other part count skips it entirely and needs the + // host evaluator, which device-only would leave without data until the + // R2 downgrade recovered it. + if air.composition_poly_degree_bound(n) / n != 2 { + return false; + } + let lde_size = domain.interpolation_domain_size * domain.blowup_factor; let offsets_contiguous = crate::gpu_lde::offsets_are_contiguous(&air.context().transition_offsets); let zerofier_uniform = air.constraints_meta().iter().all(|m| m.end_exemptions == 0); @@ -1588,37 +1606,51 @@ pub trait IsStarkProver< // when the evaluation itself already ran on device). #[cfg(feature = "cuda")] let mut precomputed_parts: Option>>> = None; + // A downloaded `H` awaiting the host decompose: produced under the + // lock below, consumed after it — the host iFFT + LDEs are pure CPU + // work and must not serialize other tables' device windows. + #[cfg(feature = "cuda")] + let mut downloaded_h: Option>> = None; #[cfg(feature = "cuda")] - if number_of_parts == 2 - && let Some(h_dev) = evaluator.evaluate_dev( + if number_of_parts == 2 { + // Serializing this window across tables (device constraint eval + + // decompose, where H is born) empirically eliminates a transient + // whole-buffer H corruption seen under concurrent R2 windows on + // VRAM pressure. What the guard orders is submission: a + // device-only table's window is enqueue-only, so its kernels may + // still overlap another table's on device. The commit, the host + // decompose of a downloaded `H` and every host arm run outside + // the lock. + let _r2_serial_guard = crate::gpu_lde::r2_serialize_guard(); + if let Some(h_dev) = evaluator.evaluate_dev( air, &round_1_result.lde_trace, domain, transition_coefficients, boundary_coefficients, &round_1_result.rap_challenges, - ) - { - match crate::gpu_lde::try_decompose_extend_d2_dev::( - &h_dev, - twiddles.inv_2x(domain), - &twiddles.composition(domain).weights, - !round_1_result.lde_trace.host_trace_empty(), ) { - Some((parts, handle)) => { - gpu_composition_parts = Some(handle); - precomputed_parts = Some(parts); - } - None => { - if let Some(h) = - crate::gpu_lde::download_comp_h_to_field::(&h_dev) - { - precomputed_parts = - Some(Self::decompose_and_extend_d2(&h, domain, twiddles)); + match crate::gpu_lde::try_decompose_extend_d2_dev::( + &h_dev, + twiddles.inv_2x(domain), + &twiddles.composition(domain).weights, + !round_1_result.lde_trace.host_trace_empty(), + ) { + Some((parts, handle)) => { + gpu_composition_parts = Some(handle); + precomputed_parts = Some(parts); + } + None => { + downloaded_h = + crate::gpu_lde::download_comp_h_to_field::(&h_dev); } } } } + #[cfg(feature = "cuda")] + if let Some(h) = downloaded_h.take() { + precomputed_parts = Some(Self::decompose_and_extend_d2(&h, domain, twiddles)); + } #[cfg(not(feature = "cuda"))] let precomputed_parts: Option>>> = None; @@ -1630,14 +1662,28 @@ pub trait IsStarkProver< // Every arm below runs the HOST evaluator, which reads `get_main` / // `get_aux`. Under device-only those buffers are intentionally empty, // so landing here means the device decompose AND the `H` download both - // failed. Abort with the device-only contract's message rather than a - // bare index-out-of-bounds from somewhere inside the evaluator. + // failed. The gate is a static predicate and cannot mirror every + // dynamic decline, so recover rather than abort: download the resident + // LDEs into the host buffers (which also clears the device-only flag) + // and let the host arms run — slower for this table, never wrong. The + // assert is left for the case where the handles themselves cannot + // serve the data, so that failure carries the device-only contract's + // message rather than a bare index-out-of-bounds from somewhere inside + // the evaluator. #[cfg(feature = "cuda")] - if precomputed_parts.is_none() { + if precomputed_parts.is_none() && round_1_result.lde_trace.host_trace_empty() { + let recovered = + crate::gpu_lde::materialize_lde_trace_host(&mut round_1_result.lde_trace); assert!( - !round_1_result.lde_trace.host_trace_empty(), - "R2 composition fell back to the host evaluator, but the trace \ - is device-only (empty)" + recovered, + "R2 composition fell back to the host evaluator on a device-only \ + trace and the resident handles could not be downloaded: \ + table={} n={} num_parts={} main_cols={} aux_cols={}", + air.name(), + trace_length, + number_of_parts, + round_1_result.lde_trace.num_main_cols(), + round_1_result.lde_trace.num_aux_cols(), ); } @@ -3379,23 +3425,29 @@ pub trait IsStarkProver< if air.has_aux_trace() { let lde_size = domain.interpolation_domain_size * domain.blowup_factor; - // Same gate as the Round 1 main commit: skip the aux - // host D2H when device-only, so both buffers are left - // empty together for this table. + // Device-only for the aux commit: the main commit's + // gate AND a produced main device handle. The aux side + // may be MORE conservative than main (never less) — if + // the GPU main commit declined and fell back to CPU, + // skipping the aux D2H here would leave a device-only + // trace with no main handle to serve it. #[cfg(feature = "cuda")] - let device_only = Self::device_only_for(*air, domain); + let mut device_only = Self::device_only_for(*air, domain) + && gpu_main_cells[idx].lock().unwrap().is_some(); // Resident GPU path: aux columns already on device (from // the resident LogUp aux build) — LDE straight from device // memory, no upload, no host column extraction. When the // resident build fired the host aux trace is empty, so a - // device LDE failure is a hard abort, not a fall through to - // the host path below (which would commit a zero aux trace). + // device LDE failure downloads the resident aux trace and + // continues on the host arms below (falling through as-is + // would commit a zero aux trace). #[cfg(feature = "cuda")] - if let Some(ra) = trace.aux_resident() { + if trace.aux_resident().is_some() { #[cfg(feature = "instruments")] let t_sub = Instant::now(); - let (tree, handle, aux_data) = + let num_cols = trace.aux_resident().map_or(0, |ra| ra.num_aux_cols); + let expand = |ra: &math_cuda::logup::ResidentAux| { crate::gpu_lde::try_expand_leaf_and_tree_ext3_row_major_keep_dev::< Field, FieldExtension, @@ -3406,21 +3458,91 @@ pub trait IsStarkProver< &twiddles.coset_weights, !device_only, ) - .ok_or_else(|| { - ProvingError::Fft( - "resident aux LDE failed; host aux trace is empty" - .to_string(), - ) - })?; - let num_cols = ra.num_aux_cols; - #[cfg(feature = "instruments")] - crate::instruments::accum_r1_aux(t_sub.elapsed(), Duration::ZERO); - let root = tree.root; - return Ok(( - Some(TableCommit::plain(tree, root)), - (aux_data, num_cols), - Some(handle), - )); + }; + let mut expanded = expand(trace.aux_resident().expect("checked above")); + if expanded.is_none() + && let Ok(be) = math_cuda::device::backend() + && be.ctx.synchronize().is_ok() + { + // The decline is usually transient VRAM + // pressure from concurrent tables; a device + // drain releases those peaks, so one retry + // tends to keep the table fully resident + // instead of paying the host downgrade. + eprintln!( + "[gpu] resident aux LDE declined: table={} \ + (retrying after device drain)", + air.name(), + ); + expanded = expand(trace.aux_resident().expect("checked above")); + } + if let Some((tree, handle, aux_data)) = expanded { + #[cfg(feature = "instruments")] + crate::instruments::accum_r1_aux(t_sub.elapsed(), Duration::ZERO); + let root = tree.root; + return Ok(( + Some(TableCommit::plain(tree, root)), + (aux_data, num_cols), + Some(handle), + )); + } + // The device aux LDE declined at runtime (transient + // VRAM pressure, usually) and there is no host aux + // trace to fall back to. Same class as the R2 + // downgrade: download the resident aux trace — and + // the main LDE if this table was device-only — and + // continue fully host-backed on the arms below. + let mut recovered = crate::gpu_lde::materialize_aux_trace_host(*trace); + // Once the aux download lands, the host aux trace is + // populated: a later failure is the main-LDE + // download's, and the error has to name that step + // instead of claiming an empty aux trace. + let aux_recovered = recovered; + if recovered && device_only { + let mut cell = main_lde_cells[idx].lock().unwrap(); + if let Some((data, _)) = cell.as_mut() + && data.is_empty() + && trace.num_main_columns > 0 + { + recovered = match ( + gpu_main_cells[idx].lock().unwrap().as_ref(), + math_cuda::device::backend(), + ) { + (Some(h), Ok(be)) => { + match crate::gpu_lde::download_main_lde_row_major::( + h, + &be.next_stream(), + ) { + Some(v) => { + *data = v; + true + } + None => false, + } + } + _ => false, + }; + } + } + if !recovered { + return Err(ProvingError::Fft( + if aux_recovered { + "resident aux LDE declined; the aux trace was recovered \ + but the main-LDE download failed" + } else { + "resident aux LDE declined and the aux-trace download \ + recovery failed" + } + .to_string(), + )); + } + eprintln!( + "[gpu] resident-aux downgrade: table={} rows={} \ + (device aux LDE declined; continuing on host)", + air.name(), + trace.num_rows(), + ); + device_only = false; } // Fused GPU path (cuda only): row-major ext3 NTT — single diff --git a/crypto/stark/src/trace.rs b/crypto/stark/src/trace.rs index b34023ac3..ccf35cca5 100644 --- a/crypto/stark/src/trace.rs +++ b/crypto/stark/src/trace.rs @@ -328,12 +328,18 @@ where pub(crate) lde_step_size: usize, pub(crate) blowup_factor: usize, /// Full-residency (Stage 3): when true the round-1 D2H was intentionally - /// skipped and `main_data`/`aux_data` are empty — every round reads the LDE - /// off the device instead. Any code path that would read the host trace must - /// hard-abort on this flag rather than index an empty buffer, so a mis-gate - /// or an unexpected GPU fallback fails loudly instead of producing a wrong - /// proof. Set by `build_round1` when the device-only gate kept this table's - /// round-1 LDE on the GPU. + /// skipped and at least one of `main_data`/`aux_data` is empty — those + /// columns are read off the device instead. Set by `build_round1` when the + /// device-only gate kept this table's round-1 LDE on the GPU, and cleared + /// again by `set_host_data` once a downgrade has downloaded the resident + /// LDEs back into the host buffers. + /// + /// The R4 and host-evaluator guards hard-abort on this flag rather than + /// index an empty buffer, so a mis-gate or an unexpected GPU fallback + /// fails loudly instead of producing a wrong proof. The R3 barycentric + /// arms instead check the individual buffer they are about to read: mixed + /// states (one side host-backed, the other device-only) are valid, and the + /// populated side stays readable. #[cfg(feature = "cuda")] pub(crate) host_trace_empty: bool, /// Per table GPU residency session: owns this table's device LDE buffers @@ -525,8 +531,11 @@ where } /// Mark this table's host LDE trace as intentionally empty (Stage-3 - /// device-only path): the round-1 D2H was skipped and every host-trace read - /// must hard-abort instead of indexing the empty buffers. + /// device-only path): the round-1 D2H was skipped, so the R4 and + /// host-evaluator reads hard-abort on the flag instead of indexing the + /// empty buffers, while the R3 arms consult the individual buffer. Cleared + /// by [`Self::set_host_data`] once a downgrade has downloaded the resident + /// LDEs back to the host. #[cfg(feature = "cuda")] pub fn set_host_trace_empty(&mut self, empty: bool) { self.host_trace_empty = empty; @@ -541,9 +550,33 @@ where self.num_rows = num_rows; } + /// Install downloaded host buffers on a device-only table and clear the + /// flag: from here every host read is valid again. An empty Vec keeps + /// that side's existing buffer (either the side has no columns or it + /// already held a host copy in a mixed state). Only meaningful from + /// [`crate::gpu_lde::materialize_lde_trace_host`], which guarantees the + /// buffers match the device handles' layout. + #[cfg(feature = "cuda")] + pub(crate) fn set_host_data( + &mut self, + main_data: Vec>, + aux_data: Vec>, + ) { + if !main_data.is_empty() { + self.main_data = main_data; + } + if !aux_data.is_empty() { + self.aux_data = aux_data; + } + self.host_trace_empty = false; + } + /// Whether the host LDE trace was intentionally left empty (see - /// [`Self::set_host_trace_empty`]). Guards on every host-read fallback check - /// this before touching `main_data`/`aux_data`. + /// [`Self::set_host_trace_empty`]). The R4 and host-evaluator fallbacks + /// check this before touching `main_data`/`aux_data`; the R3 barycentric + /// arms check the individual buffer instead, since a mixed state leaves + /// one side readable. False again once a downgrade has repopulated the + /// buffers through [`Self::set_host_data`]. #[cfg(feature = "cuda")] pub fn host_trace_empty(&self) -> bool { self.host_trace_empty @@ -781,10 +814,12 @@ where v } else { // Device-only tables have no host trace; a GPU fall-through here would - // read empty `main_data`. Hard-abort instead of a wrong OOD eval. + // read empty `main_data`. Hard-abort instead of a wrong OOD eval. The + // check is on the buffer itself, not the table-wide flag: a mixed + // state can leave a valid host copy on one side only. #[cfg(feature = "cuda")] assert!( - !lde_trace.host_trace_empty(), + lde_trace.num_main_cols() == 0 || !lde_trace.main_data.is_empty(), "R3 barycentric (main) fell back to the host trace, but it is device-only (empty)" ); let inv_denoms_v = @@ -839,10 +874,11 @@ where v } else { // Device-only tables have no host trace; a GPU fall-through here would - // read empty `aux_data`. Hard-abort instead of a wrong OOD eval. + // read empty `aux_data`. Hard-abort instead of a wrong OOD eval. Same + // buffer-level check as the main arm: mixed states are valid here. #[cfg(feature = "cuda")] assert!( - !lde_trace.host_trace_empty(), + lde_trace.num_aux_cols() == 0 || !lde_trace.aux_data.is_empty(), "R3 barycentric (aux) fell back to the host trace, but it is device-only (empty)" ); let inv_denoms_v = diff --git a/prover/tests/cuda_path_integration.rs b/prover/tests/cuda_path_integration.rs index b60cb3a34..5b5e52fbc 100644 --- a/prover/tests/cuda_path_integration.rs +++ b/prover/tests/cuda_path_integration.rs @@ -5,7 +5,9 @@ //! regressions (GPU path fired but produced output that fails verification). //! //! `#[ignore]`'d so the no-GPU CI path skips it. Run via `make test-cuda-integration` -//! or `cargo test -p lambda-vm-prover --release --features cuda --test cuda_path_integration -- --ignored --nocapture`. +//! or `cargo test -p lambda-vm-prover --release --features cuda --test cuda_path_integration -- --ignored --nocapture --test-threads=1`. +//! The single test thread is not optional: the counters these tests assert on +//! are process-global, so parallel proves in one process cross-contaminate them. #![cfg(feature = "cuda")] use lambda_vm_prover::test_utils::asm_elf_bytes; @@ -183,7 +185,10 @@ fn gpu_opening_gather_fires_and_verifies() { /// the happy path (none may fire) plus the GPU-only R2/R3/R4 paths reading the /// device LDE with no host trace behind them. A regression that silently /// reverts to the host D2H drops the counter to 0 (while the proof would still -/// verify), and a mis-gate that forces a host fallback panics one of the guards. +/// verify). A mis-gate that forces a host fallback shows up one of two ways: +/// at R3/R4 it panics one of the guards, while at R2 and the R1 resident-aux +/// commit it recovers silently and is caught by the downgrade-counter +/// assertion below. #[test] #[ignore = "requires GPU; run with --ignored --nocapture"] fn gpu_device_only_residency_fires_and_verifies() { @@ -194,6 +199,14 @@ fn gpu_device_only_residency_fires_and_verifies() { gpu_device_only_calls() > 0, "device-only residency path did not fire (every table kept its host trace)" ); + assert_eq!( + stark::gpu_lde::gpu_device_only_downgrades(), + 0, + "a table was downgraded back to a host trace on the happy path \ + (a device dispatch declined at runtime: on a device-only table the \ + gate should mirror the missing condition; a resident-aux decline is \ + usually VRAM pressure)" + ); assert!( verify(&proof, &elf).expect("verify"), "GPU-produced proof (device-only residency) failed verification"