From 2c1ae05c021c9e2957876dabec68d28f08312ac7 Mon Sep 17 00:00:00 2001 From: Joaquin Carletti Date: Wed, 5 Aug 2026 16:44:21 -0300 Subject: [PATCH 1/5] 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 2/5] 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 3/5] 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 4/5] 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 5/5] 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 =