fix(gpu): recover device-only declines at the remaining cliff sites (R4 DEEP, comp-tree, R3 barycentric)

crypto/math-cuda/src/barycentric.rs

@@ -143,6 +143,8 @@ pub fn barycentric_base_on_device(
     inv_denoms_ext3: &[u64],
     n: usize,
 ) -> Result<Vec<u64>> {
+    #[cfg(feature = "test-faults")]
+    crate::faults::check_sticky(&crate::faults::FAULT_BARYCENTRIC_STICKY)?;
     assert_eq!(coset_points.len(), n);
     assert_eq!(inv_denoms_ext3.len(), 3 * n);
     let num_cols = main_handle.m;
@@ -204,6 +206,8 @@ pub fn barycentric_base_on_device_with_dev_inv_denoms(
     inv_offset_u64: usize,
     n: usize,
 ) -> Result<Vec<u64>> {
+    #[cfg(feature = "test-faults")]
+    crate::faults::check_sticky(&crate::faults::FAULT_BARYCENTRIC_STICKY)?;
     main_handle.wait_ready_on(stream)?;
     assert!(coset_points_dev.len() >= n);
     let inv_end = inv_offset_u64
@@ -255,6 +259,8 @@ pub fn barycentric_ext3_on_device(
     inv_denoms_ext3: &[u64],
     n: usize,
 ) -> Result<Vec<u64>> {
+    #[cfg(feature = "test-faults")]
+    crate::faults::check_sticky(&crate::faults::FAULT_BARYCENTRIC_STICKY)?;
     assert_eq!(coset_points.len(), n);
     assert_eq!(inv_denoms_ext3.len(), 3 * n);
     let num_cols = aux_handle.m;
@@ -308,6 +314,8 @@ pub fn barycentric_ext3_on_device_with_dev_inv_denoms(
     inv_offset_u64: usize,
     n: usize,
 ) -> Result<Vec<u64>> {
+    #[cfg(feature = "test-faults")]
+    crate::faults::check_sticky(&crate::faults::FAULT_BARYCENTRIC_STICKY)?;
     aux_handle.wait_ready_on(stream)?;
     assert!(coset_points_dev.len() >= n);
     let inv_end = inv_offset_u64

crypto/math-cuda/src/deep.rs

@@ -41,6 +41,8 @@ pub fn deep_composition_ext3(
     row_stride: usize,
     domain_size: usize,
 ) -> Result<Vec<u64>> {
+    #[cfg(feature = "test-faults")]
+    crate::faults::check_sticky(&crate::faults::FAULT_DEEP_STICKY)?;
     let be = backend()?;
     let stream = be.next_stream();
     deep_composition_ext3_impl(
@@ -86,6 +88,8 @@ pub fn deep_composition_ext3_with_dev_parts(
     row_stride: usize,
     domain_size: usize,
 ) -> Result<Vec<u64>> {
+    #[cfg(feature = "test-faults")]
+    crate::faults::check_sticky(&crate::faults::FAULT_DEEP_STICKY)?;
     let be = backend()?;
     let stream = be.next_stream();
     deep_composition_ext3_impl(
@@ -262,6 +266,8 @@ pub fn deep_composition_ext3_with_dev_parts_and_inv_denoms(
     row_stride: usize,
     domain_size: usize,
 ) -> Result<Vec<u64>> {
+    #[cfg(feature = "test-faults")]
+    crate::faults::check_sticky(&crate::faults::FAULT_DEEP_STICKY)?;
     let deep_out = deep_fully_resident_launch(
         stream,
         main_lde,
@@ -324,6 +330,8 @@ pub fn deep_composition_ext3_fully_resident_keep(
     row_stride: usize,
     domain_size: usize,
 ) -> Result<GpuDeepCodeword> {
+    #[cfg(feature = "test-faults")]
+    crate::faults::check_sticky(&crate::faults::FAULT_DEEP_STICKY)?;
     assert!(
         domain_size.is_power_of_two() && domain_size >= 2,
         "bit-reverse needs a power-of-two codeword"

crypto/math-cuda/src/faults.rs

@@ -0,0 +1,51 @@
+//! Sticky fault-injection hooks for the GPU error-path tests.
+//!
+//! Unlike the one-shot hooks in `fri` and `inverse` (which disarm after
+//! firing, so a drain-and-retry absorbs the injected error before it can
+//! surface), a sticky hook keeps failing once its armed call count is
+//! reached, until explicitly disarmed. The device-decline recovery tests
+//! need that: a stage falls through to its host path only when every device
+//! arm of that stage declines in the same prove.
+
+use std::sync::atomic::{AtomicI64, Ordering};
+
+use crate::Result;
+
+/// R3 barycentric entries (`barycentric_{base,ext3}_on_device{,_with_dev_inv_denoms}`).
+pub static FAULT_BARYCENTRIC_STICKY: AtomicI64 = AtomicI64::new(-1);
+/// R4 DEEP composition entries (`deep_composition_ext3*`).
+pub static FAULT_DEEP_STICKY: AtomicI64 = AtomicI64::new(-1);
+/// R2 comp-poly tree entries (`build_comp_poly_tree_from_{evals_ext3_keep,slabs_dev}`).
+pub static FAULT_COMP_TREE_STICKY: AtomicI64 = AtomicI64::new(-1);
+
+/// Countdown check shared by the sticky hooks: negative = disarmed (the
+/// production state); N > 0 counts down across calls and the Nth call — and
+/// every call after it — returns Err (the counter parks at 0); 0 therefore
+/// doubles as the "fired" marker. Disarm by storing -1.
+///
+/// The transition is a single `fetch_update`, so concurrent table dispatches
+/// (the prover runs a rayon task per table) cannot race the load against the
+/// decrement: each caller walks the counter one step (the closure returns
+/// `None` at `<= 0`, so it parks at 0 and never underflows), which keeps both
+/// the sticky guarantee and the `== 0` fired check sound. The fire decision
+/// reads `fetch_update`'s own result — `Ok(prev)` for the call that
+/// decremented, `Err(cur)` for a no-op — so no second load is needed.
+pub fn check_sticky(counter: &AtomicI64) -> Result<()> {
+    // One atomic transition, so concurrent dispatches saturate at 0 rather
+    // than underflowing: a decrement only happens from a positive value.
+    let fired = counter
+        .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |v| match v {
+            n if n < 0 => None, // disarmed: never fires
+            0 => None,          // already parked: stay fired (sticky)
+            _ => Some(v - 1),   // count down toward the parked 0
+        })
+        // Ok(prev): this call decremented — the 1 → 0 step fires.
+        // Err(cur): no-op — fires only if already parked at 0.
+        .map_or_else(|cur| cur == 0, |prev| prev <= 1);
+    if fired {
+        return Err(cudarc::driver::DriverError(
+            cudarc::driver::sys::CUresult::CUDA_ERROR_UNKNOWN,
+        ));
+    }
+    Ok(())
+}

crypto/math-cuda/src/lib.rs

@@ -9,6 +9,8 @@ pub mod barycentric;
 pub mod constraint_interp;
 pub mod deep;
 pub mod device;
+#[cfg(feature = "test-faults")]
+pub mod faults;
 pub mod fri;
 pub mod inverse;
 pub mod lde;

crypto/math-cuda/src/merkle.rs

@@ -497,6 +497,8 @@ pub fn build_comp_poly_tree_from_slabs_dev(
     m: usize,
     lde_size: usize,
 ) -> Result<crate::lde::GpuMerkleTree> {
+    #[cfg(feature = "test-faults")]
+    crate::faults::check_sticky(&crate::faults::FAULT_COMP_TREE_STICKY)?;
     assert!(m > 0);
     assert!(lde_size.is_power_of_two() && lde_size >= 2);
     assert_eq!(buf.len(), 3 * m * lde_size, "slab buffer shape");
@@ -544,6 +546,8 @@ pub fn build_comp_poly_tree_from_slabs_dev(
 pub fn build_comp_poly_tree_from_evals_ext3_keep(
     parts_interleaved: &[&[u64]],
 ) -> Result<crate::lde::GpuMerkleTree> {
+    #[cfg(feature = "test-faults")]
+    crate::faults::check_sticky(&crate::faults::FAULT_COMP_TREE_STICKY)?;
     let (nodes_dev, num_leaves, stream) = build_comp_poly_tree_nodes_dev(parts_interleaved)?;
     let mut root = [0u8; 32];
     stream.memcpy_dtoh(&nodes_dev.slice(0..32), &mut root)?;

crypto/stark/src/gpu_lde.rs

@@ -117,6 +117,7 @@ pub fn reset_all_gpu_call_counters() {
     GPU_DEVICE_ONLY_DOWNGRADES.store(0, Ordering::Relaxed);
     GPU_RESIDENT_AUX_RETRIES.store(0, Ordering::Relaxed);
     GPU_RESIDENT_AUX_DOWNGRADES.store(0, Ordering::Relaxed);
+    GPU_COMPOSITION_PARTS_DOWNLOADS.store(0, Ordering::Relaxed);
 }
 
 pub(crate) static GPU_EXTEND_HALVES_CALLS: AtomicU64 = AtomicU64::new(0);
@@ -1464,16 +1465,19 @@ 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);
-/// R2 downgrades, and only those: times a device-only table fell back to the
-/// host evaluator and had its resident LDEs downloaded into the host buffers
-/// first ([`materialize_lde_trace_host`], the sole site that bumps this).
-/// Nonzero means the device-only gate cleared a table whose R2 dispatch then
-/// declined at runtime — the table continued host-backed, correct but slower —
-/// so every count is a gate miss, and the fix is to mirror the missing
-/// condition into the gate. The R1 resident-aux downgrade is counted by
-/// [`GPU_RESIDENT_AUX_DOWNGRADES`] instead: it fires on tables the gate never
-/// marked device-only, so summing the two would blame the gate for declines it
-/// never made.
+/// Device-only trace downgrades: times a device-only table fell back to a
+/// host arm and had its resident LDEs downloaded into the host buffers first
+/// ([`materialize_lde_trace_host`], the sole function that bumps this —
+/// entered from the R2 host evaluator, the R3 barycentric arms and the R4
+/// DEEP host loop). Nonzero means the device-only gate cleared a table whose
+/// downstream dispatch then declined at runtime — the table continued
+/// host-backed, correct but slower. A count is either a gate miss (a static
+/// condition worth mirroring into the gate) or a transient device decline
+/// (VRAM pressure), which by definition cannot be gated out — see
+/// [`materialize_lde_trace_host`]'s own note. The R1 resident-aux downgrade
+/// is counted by [`GPU_RESIDENT_AUX_DOWNGRADES`] instead: it fires on tables
+/// the gate never marked device-only, so summing the two would blame the gate
+/// for declines it never made.
 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)
@@ -1493,6 +1497,18 @@ pub fn gpu_resident_aux_downgrades() -> u64 {
     GPU_RESIDENT_AUX_DOWNGRADES.load(Ordering::Relaxed)
 }
 
+/// Times the composition-poly parts of a device-only table were downloaded
+/// from the resident R2 handle so a host consumer could run
+/// ([`download_composition_parts_host`], the sole site that bumps this). The
+/// parts-side counterpart of [`GPU_DEVICE_ONLY_DOWNGRADES`]: that one covers
+/// the trace LDEs, this one the H part evaluations whose R2 host drain was
+/// skipped, when the R2 commit, the R3 parts OOD or the R4 DEEP H terms later
+/// fall back to the host path.
+pub(crate) static GPU_COMPOSITION_PARTS_DOWNLOADS: AtomicU64 = AtomicU64::new(0);
+pub fn gpu_composition_parts_downloads() -> u64 {
+    GPU_COMPOSITION_PARTS_DOWNLOADS.load(Ordering::Relaxed)
+}
+
 /// Times the R1 resident-aux LDE declined and the prover drained the device to
 /// retry it (prover.rs). Nonzero means the device hit transient VRAM pressure —
 /// the retry is what keeps a decline from becoming a
@@ -1725,6 +1741,109 @@ where
     true
 }
 
+/// Parts counterpart of [`materialize_lde_trace_host`]: download the resident
+/// composition-poly parts (de-interleaved ext3 slabs, natural evaluation
+/// order) into per-part host Vecs. Serves the host consumers of the part
+/// evaluations — the R2 Merkle commit, the R3 parts OOD and the R4 DEEP H
+/// terms — when a device dispatch declines on a table whose R2 host drain was
+/// skipped (device-only). Returns `None` when the handle cannot serve the
+/// data: a non-ext3 field, a failed download or sync.
+pub(crate) fn download_composition_parts_host<E>(
+    h: &math_cuda::lde::GpuLdeExt3,
+    stream: &Arc<math_cuda::CudaStream>,
+) -> Option<Vec<Vec<FieldElement<E>>>>
+where
+    E: IsField + 'static,
+{
+    if TypeId::of::<E>() != TypeId::of::<Degree3GoldilocksExtensionField>() {
+        return None;
+    }
+    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;
+    }
+    // Per part: de-interleave the 3 slabs into row-major ext3 and reinterpret
+    // the u64 buffer in place — mirroring `materialize_lde_trace_host` rather
+    // than copying again through `u64_to_ext3_vec`. The row fill is parallel;
+    // this path fires often under VRAM pressure and otherwise dominates the
+    // D2H it follows.
+    let parts = (0..m)
+        .map(|p| {
+            let mut interleaved = vec![0u64; lde * 3];
+            #[cfg(feature = "parallel")]
+            interleaved
+                .par_chunks_exact_mut(3)
+                .enumerate()
+                .for_each(|(r, dst)| {
+                    for (k, d) in dst.iter_mut().enumerate() {
+                        *d = slabs[(p * 3 + k) * lde + r];
+                    }
+                });
+            #[cfg(not(feature = "parallel"))]
+            for (r, dst) in interleaved.chunks_exact_mut(3).enumerate() {
+                for (k, d) in dst.iter_mut().enumerate() {
+                    *d = slabs[(p * 3 + k) * lde + r];
+                }
+            }
+            // SAFETY: E == Ext3 per the tower check above; FieldElement<Ext3>
+            // is [u64; 3]. `vec![0u64; lde*3]` has len == capacity == lde*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<E>,
+                    v.len() / 3,
+                    v.capacity() / 3,
+                )
+            }
+        })
+        .collect();
+    GPU_COMPOSITION_PARTS_DOWNLOADS.fetch_add(1, Ordering::Relaxed);
+    Some(parts)
+}
+
+/// Repopulate empty host part evaluations from the resident R2 parts handle
+/// held by `lde_trace`. Already-populated evaluations are left untouched (the
+/// R2 host drain ran, nothing is missing). Returns false only when the parts
+/// are empty and the handle cannot serve them — a missing handle or bound
+/// stream, a handle whose part count disagrees with the evaluations, or a
+/// failed download — so the caller's abort carries the device-only contract's
+/// message.
+pub(crate) fn materialize_composition_parts_host<F, E>(
+    lde_trace: &crate::trace::LDETraceTable<F, E>,
+    evals: &mut [Vec<FieldElement<E>>],
+) -> bool
+where
+    F: IsField + IsSubFieldOf<E> + 'static,
+    E: IsField + 'static,
+{
+    if evals.first().is_none_or(|p| !p.is_empty()) {
+        return true;
+    }
+    let Some(h) = lde_trace.gpu_composition_parts() else {
+        return false;
+    };
+    let Some(stream) = lde_trace.bound_stream() else {
+        return false;
+    };
+    if h.m != evals.len() {
+        return false;
+    }
+    let Some(parts) = download_composition_parts_host::<E>(h, &stream) else {
+        return false;
+    };
+    for (dst, src) in evals.iter_mut().zip(parts) {
+        *dst = src;
+    }
+    true
+}
+
 pub fn gpu_batch_invert_calls() -> u64 {
     GPU_BATCH_INVERT_CALLS.load(Ordering::Relaxed)
 }
@@ -1764,6 +1883,49 @@ pub fn inverse_fault_fired() -> bool {
     math_cuda::inverse::FAULT_INVERSE_REMAINING_UNTIL_ERR.load(Ordering::Relaxed) < 0
 }
 
+/// Test-only: make the Nth upcoming math-cuda barycentric dispatch — and
+/// every one after it — return Err. Sticky, unlike the one-shot hooks above:
+/// the retry arms would absorb a single-shot fault before the fall-through
+/// could reach a device-only cliff site. Pass -1 to disarm (the production
+/// state). Only available with the `test-cuda-faults` feature.
+#[cfg(feature = "test-cuda-faults")]
+pub fn schedule_barycentric_fault_sticky(n_calls_until_err: i64) {
+    math_cuda::faults::FAULT_BARYCENTRIC_STICKY.store(n_calls_until_err, Ordering::Relaxed);
+}
+
+/// Test-only: whether the sticky barycentric fault reached its firing point
+/// (the countdown parks at 0 once it fires and stays there until disarmed).
+#[cfg(feature = "test-cuda-faults")]
+pub fn barycentric_fault_fired() -> bool {
+    math_cuda::faults::FAULT_BARYCENTRIC_STICKY.load(Ordering::Relaxed) == 0
+}
+
+/// Sticky counterpart of [`schedule_barycentric_fault_sticky`] for the R4
+/// DEEP composition dispatches (`deep_composition_ext3*`).
+#[cfg(feature = "test-cuda-faults")]
+pub fn schedule_deep_fault_sticky(n_calls_until_err: i64) {
+    math_cuda::faults::FAULT_DEEP_STICKY.store(n_calls_until_err, Ordering::Relaxed);
+}
+
+/// Test-only counterpart of [`barycentric_fault_fired`] for the DEEP hook.
+#[cfg(feature = "test-cuda-faults")]
+pub fn deep_fault_fired() -> bool {
+    math_cuda::faults::FAULT_DEEP_STICKY.load(Ordering::Relaxed) == 0
+}
+
+/// Sticky counterpart of [`schedule_barycentric_fault_sticky`] for the R2
+/// comp-poly tree builds (`build_comp_poly_tree_from_*`).
+#[cfg(feature = "test-cuda-faults")]
+pub fn schedule_comp_tree_fault_sticky(n_calls_until_err: i64) {
+    math_cuda::faults::FAULT_COMP_TREE_STICKY.store(n_calls_until_err, Ordering::Relaxed);
+}
+
+/// Test-only counterpart of [`barycentric_fault_fired`] for the comp-tree hook.
+#[cfg(feature = "test-cuda-faults")]
+pub fn comp_tree_fault_fired() -> bool {
+    math_cuda::faults::FAULT_COMP_TREE_STICKY.load(Ordering::Relaxed) == 0
+}
+
 /// R2 GPU dispatch: batched ext3 LDE over `parts_coefs` (composition-poly
 /// coefficient parts). Returns both the host LDE eval Vecs (needed for the
 /// R2 Merkle commit and R3 OOD path) and a device-resident `GpuLdeExt3`