diff --git a/crypto/stark/src/gpu_lde.rs b/crypto/stark/src/gpu_lde.rs index 98830fcc7..3c9c6be74 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 @@ -1413,6 +1431,224 @@ 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. 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 +where + F: IsField + IsSubFieldOf + 'static, + E: IsField + 'static, +{ + if !lde_trace.host_trace_empty() { + return true; + } + if !is_goldilocks_ext3_tower::() { + return false; + } + let Some(stream) = lde_trace.bound_stream() else { + return false; + }; + + // Main: column-major device buf -> row-major host Vec. An empty Vec tells + // `set_host_data` to keep the buffer that is already there. + let main_data: Vec> = + if lde_trace.num_main_cols() == 0 || !lde_trace.main_data.is_empty() { + Vec::new() + } else { + let Some(h) = lde_trace.gpu_main() else { + return false; + }; + if h.m != lde_trace.num_main_cols() || h.lde_size != lde_trace.num_rows() { + return false; + } + let Some(data) = download_main_lde_row_major::(h, &stream) else { + return false; + }; + data + }; + + // Aux: de-interleaved ext3 slabs -> row-major interleaved host Vec. + let aux_data: Vec> = + if lde_trace.num_aux_cols() == 0 || !lde_trace.aux_data.is_empty() { + Vec::new() + } else { + let Some(h) = lde_trace.gpu_aux() else { + return false; + }; + if h.m != lde_trace.num_aux_cols() || h.lde_size != lde_trace.num_rows() { + return false; + } + if h.wait_ready_on(&stream).is_err() { + return false; + } + let Ok(slabs) = stream.clone_dtoh(h.buf.as_ref()) else { + return false; + }; + if stream.synchronize().is_err() { + return false; + } + let (m, lde) = (h.m, h.lde_size); + 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 +} + +/// 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 4047458bc..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); @@ -1589,32 +1595,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)); + } } } } @@ -1633,11 +1646,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(), ); } @@ -3383,19 +3410,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, @@ -3406,21 +3436,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 diff --git a/crypto/stark/src/trace.rs b/crypto/stark/src/trace.rs index b34023ac3..8e5948f8d 100644 --- a/crypto/stark/src/trace.rs +++ b/crypto/stark/src/trace.rs @@ -537,6 +537,27 @@ 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. An empty Vec keeps + /// that side's existing buffer (either the side has no columns or it + /// already held a host copy in a mixed state). Only meaningful from + /// [`crate::gpu_lde::materialize_lde_trace_host`], which guarantees the + /// buffers match the device handles' layout. + #[cfg(feature = "cuda")] + pub(crate) fn set_host_data( + &mut self, + main_data: Vec>, + aux_data: Vec>, + ) { + if !main_data.is_empty() { + self.main_data = main_data; + } + if !aux_data.is_empty() { + self.aux_data = aux_data; + } + self.host_trace_empty = false; + } + pub fn set_num_rows(&mut self, num_rows: usize) { self.num_rows = num_rows; } @@ -781,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 = @@ -839,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 =