From 2c1ae05c021c9e2957876dabec68d28f08312ac7 Mon Sep 17 00:00:00 2001 From: Joaquin Carletti Date: Wed, 5 Aug 2026 16:44:21 -0300 Subject: [PATCH 01/10] fix(gpu): recover device-only tables by downloading the resident LDEs on an R2 miss MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The device-only gate is a static predicate over a dynamic dispatch: it cannot mirror every reason the device R2 path might decline (parts count, kernel eligibility, transient errors, shapes a new workload brings), and each miss was a hard abort that deadlocked the epoch pipeline — DECODE on the synthetic workload, then a second table on the real-block bench. Instead of excluding tables one by one, treat the resident handles as the source of truth: on a miss, download the main/aux LDEs back to host, clear the device-only flag, and continue on the host path. Slower for that table, never wrong; the abort remains only when the handles themselves cannot serve the data. gpu_device_only_downgrades() counts recoveries so a persistently-missing condition still gets mirrored into the gate. --- crypto/stark/src/gpu_lde.rs | 119 ++++++++++++++++++++++++++++++++++++ crypto/stark/src/prover.rs | 22 +++++-- crypto/stark/src/trace.rs | 15 +++++ 3 files changed, 152 insertions(+), 4 deletions(-) diff --git a/crypto/stark/src/gpu_lde.rs b/crypto/stark/src/gpu_lde.rs index 98830fcc7..ee357375d 100644 --- a/crypto/stark/src/gpu_lde.rs +++ b/crypto/stark/src/gpu_lde.rs @@ -1413,6 +1413,125 @@ 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 device-only table had to be downgraded back to a host trace +/// because a downstream device path missed at runtime (see +/// [`materialize_lde_trace_host`]). Nonzero values mean the device-only gate +/// admitted a table some dispatch later declined — correct but slower, and +/// worth mirroring the missing condition into the gate. +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. 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) only +/// when a handle is absent or a download fails. +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. + let main_data: Vec> = if lde_trace.num_main_cols() == 0 { + 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; + } + if h.wait_ready_on(&stream).is_err() { + return false; + } + let Ok(col_major) = 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); + let mut row_major = vec![0u64; m * lde]; + for c in 0..m { + for r in 0..lde { + row_major[r * m + c] = col_major[c * lde + r]; + } + } + // SAFETY: F == Goldilocks per the tower check; FieldElement is + // #[repr(transparent)] over u64. + 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(), + ) + } + }; + + // Aux: de-interleaved ext3 slabs -> row-major interleaved host Vec. + let aux_data: Vec> = if lde_trace.num_aux_cols() == 0 { + 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); + let mut interleaved = vec![0u64; m * lde * 3]; + for c in 0..m { + for k in 0..3 { + let slab = &slabs[(c * 3 + k) * lde..(c * 3 + k + 1) * lde]; + for r in 0..lde { + interleaved[(r * m + c) * 3 + k] = slab[r]; + } + } + } + // SAFETY: E == Ext3 per the tower check; FieldElement backing + // is [u64; 3]. + unsafe { + let mut v = std::mem::ManuallyDrop::new(interleaved); + 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 +} + pub fn gpu_batch_invert_calls() -> u64 { GPU_BATCH_INVERT_CALLS.load(Ordering::Relaxed) } diff --git a/crypto/stark/src/prover.rs b/crypto/stark/src/prover.rs index 4047458bc..d93840a93 100644 --- a/crypto/stark/src/prover.rs +++ b/crypto/stark/src/prover.rs @@ -1633,11 +1633,25 @@ pub trait IsStarkProver< // failed. Abort with 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() { + // The device R2 path missed on a device-only table. The gate is a + // static predicate and cannot mirror every dynamic decline, so + // recover instead of aborting: download the resident LDEs from + // the device handles and continue on the host path — slower for + // this table, never wrong. The abort remains only for the case + // where the handles themselves cannot serve the data. + 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(), ); } diff --git a/crypto/stark/src/trace.rs b/crypto/stark/src/trace.rs index b34023ac3..0af547c46 100644 --- a/crypto/stark/src/trace.rs +++ b/crypto/stark/src/trace.rs @@ -537,6 +537,21 @@ where /// `main_data.len()` — the caller supplies it from the device handle's /// `lde_size` instead. #[cfg(feature = "cuda")] + /// Install downloaded host buffers on a device-only table and clear the + /// flag: from here every host read is valid again. 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>, + ) { + self.main_data = main_data; + self.aux_data = aux_data; + self.host_trace_empty = false; + } + pub fn set_num_rows(&mut self, num_rows: usize) { self.num_rows = num_rows; } From c5cb2f483f36dbbca6b4e64b6657718a10a57a94 Mon Sep 17 00:00:00 2001 From: Joaquin Carletti Date: Thu, 6 Aug 2026 17:43:02 -0300 Subject: [PATCH 02/10] fix(gpu): drain-and-retry, then host downgrade, for resident-aux LDE declines A transient CUDA OOM on the resident aux LDE was a hard prove failure: the resident build leaves no host aux trace to fall back to. A device drain releases the concurrent VRAM peaks, so one retry usually keeps the table fully resident; if it still declines, download the resident aux trace (and the main LDE when the table is device-only) and continue host-backed. The drain before dropping the resident buffer also keeps kernels enqueued by the failed attempt from reading pool memory reused by a concurrent table. --- crypto/stark/src/gpu_lde.rs | 141 ++++++++++++++++++++++++++++++------ crypto/stark/src/prover.rs | 101 +++++++++++++++++++++----- 2 files changed, 198 insertions(+), 44 deletions(-) diff --git a/crypto/stark/src/gpu_lde.rs b/crypto/stark/src/gpu_lde.rs index ee357375d..f06609b23 100644 --- a/crypto/stark/src/gpu_lde.rs +++ b/crypto/stark/src/gpu_lde.rs @@ -1458,32 +1458,10 @@ where if h.m != lde_trace.num_main_cols() || h.lde_size != lde_trace.num_rows() { return false; } - if h.wait_ready_on(&stream).is_err() { - return false; - } - let Ok(col_major) = stream.clone_dtoh(h.buf.as_ref()) else { + let Some(data) = download_main_lde_row_major::(h, &stream) else { return false; }; - if stream.synchronize().is_err() { - return false; - } - let (m, lde) = (h.m, h.lde_size); - let mut row_major = vec![0u64; m * lde]; - for c in 0..m { - for r in 0..lde { - row_major[r * m + c] = col_major[c * lde + r]; - } - } - // SAFETY: F == Goldilocks per the tower check; FieldElement is - // #[repr(transparent)] over u64. - 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(), - ) - } + data }; // Aux: de-interleaved ext3 slabs -> row-major interleaved host Vec. @@ -1532,6 +1510,121 @@ where 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; + } + let mut row_major = vec![0u64; m * lde]; + for c in 0..m { + for r in 0..lde { + row_major[r * m + c] = 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 +} + +/// Diagnostic: download a resident ext3 handle (3-slab layout) as per-column +/// host Vecs. Used by the xcheck post-mortem to compare the committed R2 +/// parts against a host recompute. +pub(crate) fn download_ext3_columns( + h: &math_cuda::lde::GpuLdeExt3, +) -> Option>>> +where + E: IsField + 'static, +{ + if TypeId::of::() != TypeId::of::() { + return None; + } + let be = math_cuda::device::backend().ok()?; + let stream = be.next_stream(); + h.wait_ready_on(&stream).ok()?; + let slabs = stream.clone_dtoh(h.buf.as_ref()).ok()?; + stream.synchronize().ok()?; + let (m, lde) = (h.m, h.lde_size); + if slabs.len() != m * lde * 3 { + return None; + } + let mut cols = Vec::with_capacity(m); + for c in 0..m { + let mut interleaved = vec![0u64; lde * 3]; + for k in 0..3 { + let slab = &slabs[(c * 3 + k) * lde..(c * 3 + k + 1) * lde]; + for r in 0..lde { + interleaved[r * 3 + k] = slab[r]; + } + } + cols.push(u64_to_ext3_vec::(&interleaved)); + } + Some(cols) +} + pub fn gpu_batch_invert_calls() -> u64 { GPU_BATCH_INVERT_CALLS.load(Ordering::Relaxed) } diff --git a/crypto/stark/src/prover.rs b/crypto/stark/src/prover.rs index d93840a93..9ad446809 100644 --- a/crypto/stark/src/prover.rs +++ b/crypto/stark/src/prover.rs @@ -3397,19 +3397,22 @@ pub trait IsStarkProver< // host D2H when device-only, so both buffers are left // empty together for this table. #[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, @@ -3420,21 +3423,79 @@ 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); + 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( + "resident aux LDE failed; host aux trace is empty".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 From a68bb358a9e42abe6bda07f586bfa40c121780fe Mon Sep 17 00:00:00 2001 From: Joaquin Carletti Date: Thu, 6 Aug 2026 17:53:40 -0300 Subject: [PATCH 03/10] fix(gpu): serialize the device R2 window to close a transient H corruption race Concurrent device R2 windows under VRAM pressure can transiently produce a fully wrong H for one or two tables while every input stays correct (rerunning the same chain on the same resident inputs matches the host), yielding a proof that fails the composition check. Serializing only the constraint-eval + decompose window across tables eliminates it; commits and host arms stay parallel, and the windows overlap rarely enough that the lock is near-free. LAMBDA_VM_GPU_SERIALIZE_R2=0 lifts the lock to bisect further or once the underlying race is found. --- crypto/stark/src/gpu_lde.rs | 18 ++++++++++++ crypto/stark/src/prover.rs | 57 +++++++++++++++++++++---------------- 2 files changed, 50 insertions(+), 25 deletions(-) diff --git a/crypto/stark/src/gpu_lde.rs b/crypto/stark/src/gpu_lde.rs index f06609b23..c922c0d92 100644 --- a/crypto/stark/src/gpu_lde.rs +++ b/crypto/stark/src/gpu_lde.rs @@ -54,6 +54,24 @@ fn gpu_lde_threshold() -> usize { }) } +/// Serialize 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. +/// Serializing only this window eliminates it at negligible cost — the +/// windows rarely overlap. `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")) + { + Some(LOCK.lock().unwrap()) + } 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 diff --git a/crypto/stark/src/prover.rs b/crypto/stark/src/prover.rs index 9ad446809..fafe223e8 100644 --- a/crypto/stark/src/prover.rs +++ b/crypto/stark/src/prover.rs @@ -1589,32 +1589,39 @@ pub trait IsStarkProver< #[cfg(feature = "cuda")] let mut precomputed_parts: Option>>> = None; #[cfg(feature = "cuda")] - if number_of_parts == 2 - && 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)); + // Serializing this window across tables (device constraint eval + + // decompose, where H is born) eliminates a transient whole-buffer + // H corruption seen under concurrent R2 windows on VRAM pressure. + // The commit and every host arm run outside the lock. + let _r2_serial_guard = crate::gpu_lde::r2_serialize_guard(); + if number_of_parts == 2 + && 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)); + } } } } From 8f62d7b1627c39a887be55df946cb9bd38d26e84 Mon Sep 17 00:00:00 2001 From: Joaquin Carletti Date: Fri, 7 Aug 2026 11:08:33 -0300 Subject: [PATCH 04/10] fix(gpu): device-only requires the d=2 composition path The device R2 path only exists for the d=2 quotient split; a table with any other composition bound (DECODE proves with num_parts == 1) would skip it entirely and hard-abort on its device-only trace. --- crypto/stark/src/prover.rs | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/crypto/stark/src/prover.rs b/crypto/stark/src/prover.rs index fafe223e8..e58e8d355 100644 --- a/crypto/stark/src/prover.rs +++ b/crypto/stark/src/prover.rs @@ -1031,8 +1031,14 @@ pub trait IsStarkProver< 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 leaves without data. + 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); From 49b741a61483c3773f75c95992a4b58dc266f320 Mon Sep 17 00:00:00 2001 From: Joaquin Carletti Date: Fri, 7 Aug 2026 16:41:27 -0300 Subject: [PATCH 05/10] fix(gpu): keep already-present host buffers in the downgrade recovery A mixed state (one commit fell back to CPU while the other stayed device-only) left the recovery refusing to proceed: it treated a missing device handle as fatal even when that side already had a valid host copy. Only the missing side is downloaded now, and the R3 host-arm guards check the buffer they are about to read instead of the table-wide flag. --- crypto/stark/src/gpu_lde.rs | 120 +++++++++++++++++++----------------- crypto/stark/src/trace.rs | 23 ++++--- 2 files changed, 79 insertions(+), 64 deletions(-) diff --git a/crypto/stark/src/gpu_lde.rs b/crypto/stark/src/gpu_lde.rs index c922c0d92..3c9c6be74 100644 --- a/crypto/stark/src/gpu_lde.rs +++ b/crypto/stark/src/gpu_lde.rs @@ -1443,12 +1443,15 @@ pub fn gpu_device_only_downgrades() -> u64 { /// Recover a device-only table for the host path: download the resident main /// and aux LDEs from their device handles into the host buffers and clear the -/// device-only flag. 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) only -/// when a handle is absent or a download fails. +/// 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) only when a missing side has no handle or a +/// download fails. pub(crate) fn materialize_lde_trace_host( lde_trace: &mut crate::trace::LDETraceTable, ) -> bool @@ -1466,62 +1469,65 @@ where return false; }; - // Main: column-major device buf -> row-major host Vec. - let main_data: Vec> = if lde_trace.num_main_cols() == 0 { - 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; + // 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 }; - data - }; // Aux: de-interleaved ext3 slabs -> row-major interleaved host Vec. - let aux_data: Vec> = if lde_trace.num_aux_cols() == 0 { - 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); - let mut interleaved = vec![0u64; m * lde * 3]; - for c in 0..m { - for k in 0..3 { - let slab = &slabs[(c * 3 + k) * lde..(c * 3 + k + 1) * lde]; - for r in 0..lde { - interleaved[(r * m + c) * 3 + k] = slab[r]; + 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); + let mut interleaved = vec![0u64; m * lde * 3]; + for c in 0..m { + for k in 0..3 { + let slab = &slabs[(c * 3 + k) * lde..(c * 3 + k + 1) * lde]; + for r in 0..lde { + interleaved[(r * m + c) * 3 + k] = slab[r]; + } } } - } - // SAFETY: E == Ext3 per the tower check; FieldElement backing - // is [u64; 3]. - unsafe { - let mut v = std::mem::ManuallyDrop::new(interleaved); - Vec::from_raw_parts( - v.as_mut_ptr() as *mut FieldElement, - v.len() / 3, - v.capacity() / 3, - ) - } - }; + // SAFETY: E == Ext3 per the tower check; FieldElement backing + // is [u64; 3]. + unsafe { + let mut v = std::mem::ManuallyDrop::new(interleaved); + 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); diff --git a/crypto/stark/src/trace.rs b/crypto/stark/src/trace.rs index 0af547c46..8e5948f8d 100644 --- a/crypto/stark/src/trace.rs +++ b/crypto/stark/src/trace.rs @@ -538,7 +538,9 @@ where /// `lde_size` instead. #[cfg(feature = "cuda")] /// Install downloaded host buffers on a device-only table and clear the - /// flag: from here every host read is valid again. Only meaningful from + /// 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")] @@ -547,8 +549,12 @@ where main_data: Vec>, aux_data: Vec>, ) { - self.main_data = main_data; - self.aux_data = aux_data; + 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; } @@ -796,10 +802,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 = @@ -854,10 +862,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 = From 37a0ce19bd392a051a544ec083c562d70b151732 Mon Sep 17 00:00:00 2001 From: Joaquin Carletti Date: Mon, 10 Aug 2026 10:57:10 -0300 Subject: [PATCH 06/10] chore(gpu): drop an orphaned diagnostic helper download_ext3_columns came along in a cherry-pick but its only consumer (the cross-check post-mortem) ships separately; dead code under the cuda feature. --- crypto/stark/src/gpu_lde.rs | 35 ----------------------------------- 1 file changed, 35 deletions(-) diff --git a/crypto/stark/src/gpu_lde.rs b/crypto/stark/src/gpu_lde.rs index 3c9c6be74..cb59d6e56 100644 --- a/crypto/stark/src/gpu_lde.rs +++ b/crypto/stark/src/gpu_lde.rs @@ -1614,41 +1614,6 @@ where true } -/// Diagnostic: download a resident ext3 handle (3-slab layout) as per-column -/// host Vecs. Used by the xcheck post-mortem to compare the committed R2 -/// parts against a host recompute. -pub(crate) fn download_ext3_columns( - h: &math_cuda::lde::GpuLdeExt3, -) -> Option>>> -where - E: IsField + 'static, -{ - if TypeId::of::() != TypeId::of::() { - return None; - } - let be = math_cuda::device::backend().ok()?; - let stream = be.next_stream(); - h.wait_ready_on(&stream).ok()?; - let slabs = stream.clone_dtoh(h.buf.as_ref()).ok()?; - stream.synchronize().ok()?; - let (m, lde) = (h.m, h.lde_size); - if slabs.len() != m * lde * 3 { - return None; - } - let mut cols = Vec::with_capacity(m); - for c in 0..m { - let mut interleaved = vec![0u64; lde * 3]; - for k in 0..3 { - let slab = &slabs[(c * 3 + k) * lde..(c * 3 + k + 1) * lde]; - for r in 0..lde { - interleaved[r * 3 + k] = slab[r]; - } - } - cols.push(u64_to_ext3_vec::(&interleaved)); - } - Some(cols) -} - pub fn gpu_batch_invert_calls() -> u64 { GPU_BATCH_INVERT_CALLS.load(Ordering::Relaxed) } From bc9e65794f57dc41adffe431f81dac790c6a263f Mon Sep 17 00:00:00 2001 From: Joaquin Carletti Date: Mon, 10 Aug 2026 11:26:17 -0300 Subject: [PATCH 07/10] fix(gpu): run the host decompose of a downloaded H outside the R2 lock MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The fallback arm (download H, host iFFT + LDEs) executed under the serialization lock, so under VRAM pressure — exactly when that arm runs — it serialized every other table's device window behind pure CPU work. The lock now covers only the device eval + decompose + the H download; the host decompose and every host arm run outside it. The lock is also acquired only for d=2 tables (the others never enter the device path). --- crypto/stark/src/prover.rs | 59 +++++++++++++++++++++----------------- 1 file changed, 33 insertions(+), 26 deletions(-) diff --git a/crypto/stark/src/prover.rs b/crypto/stark/src/prover.rs index e58e8d355..d5cca05e1 100644 --- a/crypto/stark/src/prover.rs +++ b/crypto/stark/src/prover.rs @@ -1010,11 +1010,11 @@ 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, @@ -1594,23 +1594,27 @@ 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 { // Serializing this window across tables (device constraint eval + // decompose, where H is born) eliminates a transient whole-buffer // H corruption seen under concurrent R2 windows on VRAM pressure. - // The commit and every host arm run outside the lock. + // 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 number_of_parts == 2 - && let Some(h_dev) = evaluator.evaluate_dev( - air, - &round_1_result.lde_trace, - domain, - transition_coefficients, - boundary_coefficients, - &round_1_result.rap_challenges, - ) - { + 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), @@ -1622,16 +1626,16 @@ pub trait IsStarkProver< 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)); - } + 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; @@ -3406,9 +3410,12 @@ 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 mut device_only = Self::device_only_for(*air, domain) && gpu_main_cells[idx].lock().unwrap().is_some(); From 5a236765355c943b7e7fd7ac0d32718c7dca8cff Mon Sep 17 00:00:00 2001 From: Joaquin Carletti Date: Mon, 10 Aug 2026 11:26:18 -0300 Subject: [PATCH 08/10] chore(gpu): review follow-ups on the downgrade recovery set_host_data had been inserted between set_num_rows' doc and its signature, stealing its doc comment and un-gating it from the cuda feature; the serialization lock now recovers from poisoning instead of cascading PoisonErrors over the original panic; the downgrade counter joins reset_all_gpu_call_counters and the device-only residency test asserts it stays at zero on the happy path; stale comments about the aux gate mirroring the main gate rewritten with the actual contract. --- crypto/stark/src/gpu_lde.rs | 10 +++++++--- crypto/stark/src/trace.rs | 8 ++++---- prover/tests/cuda_path_integration.rs | 6 ++++++ 3 files changed, 17 insertions(+), 7 deletions(-) diff --git a/crypto/stark/src/gpu_lde.rs b/crypto/stark/src/gpu_lde.rs index cb59d6e56..2a9547153 100644 --- a/crypto/stark/src/gpu_lde.rs +++ b/crypto/stark/src/gpu_lde.rs @@ -66,7 +66,10 @@ pub(crate) fn r2_serialize_guard() -> Option> 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")) { - Some(LOCK.lock().unwrap()) + // 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 } @@ -100,6 +103,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); @@ -1787,8 +1791,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/trace.rs b/crypto/stark/src/trace.rs index 8e5948f8d..c19d2be16 100644 --- a/crypto/stark/src/trace.rs +++ b/crypto/stark/src/trace.rs @@ -537,6 +537,10 @@ where /// `main_data.len()` — the caller supplies it from the device handle's /// `lde_size` instead. #[cfg(feature = "cuda")] + pub fn set_num_rows(&mut self, num_rows: usize) { + 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 @@ -558,10 +562,6 @@ where self.host_trace_empty = false; } - pub fn set_num_rows(&mut self, num_rows: usize) { - self.num_rows = num_rows; - } - /// 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`. diff --git a/prover/tests/cuda_path_integration.rs b/prover/tests/cuda_path_integration.rs index b60cb3a34..859f361ec 100644 --- a/prover/tests/cuda_path_integration.rs +++ b/prover/tests/cuda_path_integration.rs @@ -194,6 +194,12 @@ 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 device-only table was downgraded back to host on the happy path \ + (a device dispatch declined that the gate should mirror)" + ); assert!( verify(&proof, &elf).expect("verify"), "GPU-produced proof (device-only residency) failed verification" From 671d814dfb3b8ec7f1e9f32b295959d53880f797 Mon Sep 17 00:00:00 2001 From: Mauro Toscano <12560266+MauroToscano@users.noreply.github.com> Date: Mon, 10 Aug 2026 15:48:44 -0300 Subject: [PATCH 09/10] fix(gpu): harden the downgrade downloads, parallelize the recovery transposes (#921) * fix(gpu): harden the downgrade download recovery Three fixes on the device-only downgrade path, all in the graceful degradation function whose whole point is to avoid a hard abort. - The aux branch of `materialize_lde_trace_host` sliced the downloaded slabs without checking their length, so a short download would panic inside the recovery instead of degrading. Both sibling download paths already validate (`download_main_lde_row_major` checks `col_major.len() != m * lde`, `materialize_aux_trace_host` checks `raw.len() != rows * cols * 3`); this adds the matching check. - Restore the `len/capacity % 3` guard the other two ext3 `from_raw_parts` sites carry, spelled `is_multiple_of` because clippy's `manual_is_multiple_of` rejects the older form here. - The failure error claimed "host aux trace is empty" on a path where that is false: when the aux download succeeded and the follow-up main-LDE download failed, the host aux trace had just been populated. Track which recovery step failed and name it. Control flow unchanged. * perf(gpu): parallelize the downgrade recovery transposes Both conversions in the recovery path were single-threaded nested loops over the full LDE: the col-major -> row-major main transpose in `download_main_lde_row_major`, and the de-interleaved-slabs -> row-major interleaved aux conversion in `materialize_lde_trace_host`. For MEMW at LDE 2^20 those are a 411 MB and a 327 MB buffer respectively, walked with a strided access on one core. Both now follow the existing idiom in `trace.rs` ("Parallel col-major -> row-major transpose"): parallelize over OUTPUT row chunks with `par_chunks_exact_mut`, so every element is still written exactly once and no unsafe is involved. The index math is unchanged -- chunk `r` of width `m` is `row_major[r * m + c]`, and chunk `r` of width `m * 3` sub-chunked by 3 is `interleaved[(r * m + c) * 3 + k]` -- because the layout was verified against the kernels. Gated on the `parallel` feature with the sequential loop kept for builds without it, and skipped when `m == 0` since `chunks_exact_mut(0)` panics. These loops run on a scheduler driver thread holding no locks, so rayon is safe here, unlike the pinned-staging unpack in math-cuda. --- crypto/stark/src/gpu_lde.rs | 65 ++++++++++++++++++++++++++++++++----- crypto/stark/src/prover.rs | 14 +++++++- 2 files changed, 70 insertions(+), 9 deletions(-) diff --git a/crypto/stark/src/gpu_lde.rs b/crypto/stark/src/gpu_lde.rs index 2a9547153..e4e85792e 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; @@ -1512,12 +1514,37 @@ where 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]; - for c in 0..m { - for k in 0..3 { - let slab = &slabs[(c * 3 + k) * lde..(c * 3 + k + 1) * lde]; - for r in 0..lde { - interleaved[(r * m + c) * 3 + k] = slab[r]; + 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]; + } + } } } } @@ -1525,6 +1552,10 @@ where // 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, @@ -1557,10 +1588,28 @@ where 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]; - for c in 0..m { - for r in 0..lde { - row_major[r * m + c] = col_major[c * lde + r]; + 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 diff --git a/crypto/stark/src/prover.rs b/crypto/stark/src/prover.rs index d5cca05e1..b609786e7 100644 --- a/crypto/stark/src/prover.rs +++ b/crypto/stark/src/prover.rs @@ -3478,6 +3478,11 @@ pub trait IsStarkProver< // 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() @@ -3506,7 +3511,14 @@ pub trait IsStarkProver< } if !recovered { return Err(ProvingError::Fft( - "resident aux LDE failed; host aux trace is empty".to_string(), + 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!( From 1e11b70fcb75880ac9ff39f0c7e3b1f7ac69938d Mon Sep 17 00:00:00 2001 From: Mauro Toscano <12560266+MauroToscano@users.noreply.github.com> Date: Mon, 10 Aug 2026 15:48:57 -0300 Subject: [PATCH 10/10] docs(gpu): align device-only and downgrade docs with the recovery semantics (#920) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This branch turned two of the device-only hard-aborts into downloads that recover and continue host-backed, but the surrounding docs still describe the old contract: "every host read hard-aborts", "the prove aborts loudly", "a mis-gate panics one of the guards". Rewrite those to say what the code now does — R2 and the R1 resident-aux commit recover and bump GPU_DEVICE_ONLY_DOWNGRADES, R3/R4 still abort, and the R3 guards check the individual buffer so mixed states are legal. Also correct the R2 lock comment (it serializes submission, not execution, for device-only tables), note that the numeric gate is not the complete predicate on its own, broaden the downgrade counter's doc to cover resident-aux declines on tables that were never device-only, and drop the false "only" from materialize_lde_trace_host's failure list. Comments, doc comments, two assertion message strings and one doc-comment run command (--test-threads=1, matching the Makefile target). No behavior changes. --- crypto/stark/src/gpu_lde.rs | 75 +++++++++++++++++++-------- crypto/stark/src/prover.rs | 49 +++++++++++------ crypto/stark/src/trace.rs | 32 ++++++++---- prover/tests/cuda_path_integration.rs | 15 ++++-- 4 files changed, 118 insertions(+), 53 deletions(-) diff --git a/crypto/stark/src/gpu_lde.rs b/crypto/stark/src/gpu_lde.rs index e4e85792e..ced0bc530 100644 --- a/crypto/stark/src/gpu_lde.rs +++ b/crypto/stark/src/gpu_lde.rs @@ -56,13 +56,24 @@ fn gpu_lde_threshold() -> usize { }) } -/// Serialize 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. -/// Serializing only this window eliminates it at negligible cost — the -/// windows rarely overlap. `LAMBDA_VM_GPU_SERIALIZE_R2=0` disables the lock -/// (e.g. to bisect or once the underlying race is fixed). +/// 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(()); @@ -195,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` @@ -208,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, @@ -1437,11 +1462,16 @@ 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 device-only table had to be downgraded back to a host trace -/// because a downstream device path missed at runtime (see -/// [`materialize_lde_trace_host`]). Nonzero values mean the device-only gate -/// admitted a table some dispatch later declined — correct but slower, and -/// worth mirroring the missing condition into the gate. +/// 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) @@ -1456,8 +1486,9 @@ pub fn gpu_device_only_downgrades() -> u64 { /// 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) only when a missing side has no handle or a -/// download fails. +/// (→ 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 diff --git a/crypto/stark/src/prover.rs b/crypto/stark/src/prover.rs index b609786e7..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")] @@ -1021,20 +1028,25 @@ pub trait IsStarkProver< 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 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 leaves without data. + // 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; } @@ -1602,10 +1614,13 @@ pub trait IsStarkProver< #[cfg(feature = "cuda")] if number_of_parts == 2 { // Serializing this window across tables (device constraint eval + - // decompose, where H is born) eliminates a transient whole-buffer - // H corruption seen under concurrent R2 windows on VRAM pressure. - // The commit, the host decompose of a downloaded `H` and every - // host arm run outside the lock. + // 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, @@ -1647,16 +1662,16 @@ 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() && round_1_result.lde_trace.host_trace_empty() { - // The device R2 path missed on a device-only table. The gate is a - // static predicate and cannot mirror every dynamic decline, so - // recover instead of aborting: download the resident LDEs from - // the device handles and continue on the host path — slower for - // this table, never wrong. The abort remains only for the case - // where the handles themselves cannot serve the data. let recovered = crate::gpu_lde::materialize_lde_trace_host(&mut round_1_result.lde_trace); assert!( diff --git a/crypto/stark/src/trace.rs b/crypto/stark/src/trace.rs index c19d2be16..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; @@ -563,8 +572,11 @@ where } /// 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 diff --git a/prover/tests/cuda_path_integration.rs b/prover/tests/cuda_path_integration.rs index 859f361ec..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() { @@ -197,8 +202,10 @@ fn gpu_device_only_residency_fires_and_verifies() { assert_eq!( stark::gpu_lde::gpu_device_only_downgrades(), 0, - "a device-only table was downgraded back to host on the happy path \ - (a device dispatch declined that the gate should mirror)" + "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"),