From fa0cd4372c8e8e8935b139e26ea80a5db2d7d1df Mon Sep 17 00:00:00 2001 From: MauroFab Date: Tue, 7 Jul 2026 18:07:37 -0300 Subject: [PATCH 1/9] feat(math-cuda): device row-gather kernels for R4 opening values MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit gather_rows_base / gather_rows_ext3: read full rows (all columns) from a device-resident LDE buffer at given row indices, returning them row-major — the building block for gathering R4 query-opening values from the resident LDE instead of the host trace (toward eliminating the round-1 trace D2H). Wrappers gather_rows_base_on_device / gather_rows_ext3_on_device + a standalone parity test vs direct buffer indexing. --- crypto/math-cuda/kernels/barycentric.cu | 42 +++++++++++ crypto/math-cuda/src/barycentric.rs | 78 ++++++++++++++++++++ crypto/math-cuda/src/device.rs | 4 ++ crypto/math-cuda/tests/gather_rows.rs | 95 +++++++++++++++++++++++++ 4 files changed, 219 insertions(+) create mode 100644 crypto/math-cuda/tests/gather_rows.rs diff --git a/crypto/math-cuda/kernels/barycentric.cu b/crypto/math-cuda/kernels/barycentric.cu index 5c18bcb88..f76db471a 100644 --- a/crypto/math-cuda/kernels/barycentric.cu +++ b/crypto/math-cuda/kernels/barycentric.cu @@ -190,3 +190,45 @@ extern "C" __global__ void barycentric_ext3_batched_strided( out_ext3_int[col * 3 + 2] = sum.c; } } + +// Gather full rows from a device-resident base-field LDE (`buf[col*col_stride + +// row]`). One block per gathered row, threads stride over columns. Output is +// row-major `out[q*num_cols + col]` for gathered-row slot `q` — directly the +// concatenation of `gather_main_row(rows[q])` for each q. `rows` are the LDE row +// indices to gather (already the reversed query rows on the host side). +extern "C" __global__ void gather_rows_base( + const uint64_t *__restrict__ columns, + uint64_t col_stride, + uint64_t num_cols, + const uint32_t *__restrict__ rows, + uint64_t num_rows, + uint64_t *__restrict__ out +) { + uint64_t q = blockIdx.x; + if (q >= num_rows) return; + uint64_t row = rows[q]; + for (uint64_t col = threadIdx.x; col < num_cols; col += blockDim.x) { + out[q * num_cols + col] = columns[col * col_stride + row]; + } +} + +// Ext3 variant: M ext3 columns as 3M base slabs, `columns[(col*3+k)*col_stride + +// row]`. Output interleaved ext3: `out[(q*num_cols + col)*3 + k]`. +extern "C" __global__ void gather_rows_ext3( + const uint64_t *__restrict__ columns, + uint64_t col_stride, + uint64_t num_cols, + const uint32_t *__restrict__ rows, + uint64_t num_rows, + uint64_t *__restrict__ out +) { + uint64_t q = blockIdx.x; + if (q >= num_rows) return; + uint64_t row = rows[q]; + for (uint64_t col = threadIdx.x; col < num_cols; col += blockDim.x) { + uint64_t o = (q * num_cols + col) * 3; + out[o + 0] = columns[(col * 3 + 0) * col_stride + row]; + out[o + 1] = columns[(col * 3 + 1) * col_stride + row]; + out[o + 2] = columns[(col * 3 + 2) * col_stride + row]; + } +} diff --git a/crypto/math-cuda/src/barycentric.rs b/crypto/math-cuda/src/barycentric.rs index f299d1839..8fdea14f1 100644 --- a/crypto/math-cuda/src/barycentric.rs +++ b/crypto/math-cuda/src/barycentric.rs @@ -337,3 +337,81 @@ pub fn barycentric_ext3_on_device_with_dev_inv_denoms( stream.synchronize()?; Ok(out) } + +/// Gather full rows from a device-resident base-field LDE handle. `rows` are LDE +/// row indices; returns their column values row-major (`rows.len() * main.m` +/// u64, `out[q*num_cols + col]`) — i.e. the concatenation of +/// `gather_main_row(rows[q])` for each `q`. Runs on the caller's stream. +pub fn gather_rows_base_on_device( + main: &GpuLdeBase, + rows: &[u32], + stream: &Arc, +) -> Result> { + let num_cols = main.m; + if num_cols == 0 || rows.is_empty() { + return Ok(Vec::new()); + } + let be = backend()?; + let rows_dev = stream.clone_htod(rows)?; + let mut out = stream.alloc_zeros::(rows.len() * num_cols)?; + let col_stride = main.lde_size as u64; + let num_cols_u64 = num_cols as u64; + let num_rows_u64 = rows.len() as u64; + let cfg = LaunchConfig { + grid_dim: (rows.len() as u32, 1, 1), + block_dim: (BLOCK_DIM.min(num_cols as u32).max(1), 1, 1), + shared_mem_bytes: 0, + }; + unsafe { + stream + .launch_builder(&be.gather_rows_base) + .arg(main.buf.as_ref()) + .arg(&col_stride) + .arg(&num_cols_u64) + .arg(&rows_dev) + .arg(&num_rows_u64) + .arg(&mut out) + .launch(cfg)?; + } + let host = stream.clone_dtoh(&out)?; + stream.synchronize()?; + Ok(host) +} + +/// Ext3 sibling of [`gather_rows_base_on_device`]: returns `rows.len() * aux.m * +/// 3` u64, interleaved ext3 (`out[(q*num_cols + col)*3 + k]`). +pub fn gather_rows_ext3_on_device( + aux: &GpuLdeExt3, + rows: &[u32], + stream: &Arc, +) -> Result> { + let num_cols = aux.m; + if num_cols == 0 || rows.is_empty() { + return Ok(Vec::new()); + } + let be = backend()?; + let rows_dev = stream.clone_htod(rows)?; + let mut out = stream.alloc_zeros::(rows.len() * num_cols * 3)?; + let col_stride = aux.lde_size as u64; + let num_cols_u64 = num_cols as u64; + let num_rows_u64 = rows.len() as u64; + let cfg = LaunchConfig { + grid_dim: (rows.len() as u32, 1, 1), + block_dim: (BLOCK_DIM.min(num_cols as u32).max(1), 1, 1), + shared_mem_bytes: 0, + }; + unsafe { + stream + .launch_builder(&be.gather_rows_ext3) + .arg(aux.buf.as_ref()) + .arg(&col_stride) + .arg(&num_cols_u64) + .arg(&rows_dev) + .arg(&num_rows_u64) + .arg(&mut out) + .launch(cfg)?; + } + let host = stream.clone_dtoh(&out)?; + stream.synchronize()?; + Ok(host) +} diff --git a/crypto/math-cuda/src/device.rs b/crypto/math-cuda/src/device.rs index be0440d5e..bf75b696e 100644 --- a/crypto/math-cuda/src/device.rs +++ b/crypto/math-cuda/src/device.rs @@ -167,6 +167,8 @@ pub struct Backend { pub barycentric_ext3_batched: CudaFunction, pub barycentric_base_batched_strided: CudaFunction, pub barycentric_ext3_batched_strided: CudaFunction, + pub gather_rows_base: CudaFunction, + pub gather_rows_ext3: CudaFunction, // deep.ptx pub deep_composition_ext3_row: CudaFunction, @@ -367,6 +369,8 @@ impl Backend { .load_function("barycentric_base_batched_strided")?, barycentric_ext3_batched_strided: bary .load_function("barycentric_ext3_batched_strided")?, + gather_rows_base: bary.load_function("gather_rows_base")?, + gather_rows_ext3: bary.load_function("gather_rows_ext3")?, deep_composition_ext3_row: deep.load_function("deep_composition_ext3_row")?, fri_fold_ext3: fri.load_function("fri_fold_ext3")?, fri_update_twiddles: fri.load_function("fri_update_twiddles")?, diff --git a/crypto/math-cuda/tests/gather_rows.rs b/crypto/math-cuda/tests/gather_rows.rs new file mode 100644 index 000000000..8b7f1b4a3 --- /dev/null +++ b/crypto/math-cuda/tests/gather_rows.rs @@ -0,0 +1,95 @@ +//! Parity: the device row-gather (`gather_rows_base` / `gather_rows_ext3`, +//! used by R4 query openings to read opened column values from the resident LDE +//! instead of the host trace) returns exactly the column values obtained by +//! directly indexing the column-major device buffer at each requested row. + +use std::sync::Arc; + +use math_cuda::barycentric::{gather_rows_base_on_device, gather_rows_ext3_on_device}; +use math_cuda::device::backend; +use math_cuda::lde::{GpuLdeBase, GpuLdeExt3}; +use rand::{Rng, SeedableRng}; +use rand_chacha::ChaCha8Rng; + +fn run_base(lde_size: usize, num_cols: usize, seed: u64) { + let mut rng = ChaCha8Rng::seed_from_u64(seed); + // Column-major base LDE: buf[col*lde_size + row]. + let buf: Vec = (0..num_cols * lde_size) + .map(|_| rng.r#gen::()) + .collect(); + + let be = backend().unwrap(); + let stream = be.next_stream(); + let dev = stream.clone_htod(&buf).unwrap(); + stream.synchronize().unwrap(); + let handle = GpuLdeBase { + buf: Arc::new(dev), + m: num_cols, + lde_size, + tree: None, + trace_dev: None, + trace_rows: 0, + }; + + let rows: Vec = (0..9).map(|_| rng.gen_range(0..lde_size) as u32).collect(); + let got = gather_rows_base_on_device(&handle, &rows, &stream).unwrap(); + assert_eq!(got.len(), rows.len() * num_cols, "base gather shape"); + for (q, &row) in rows.iter().enumerate() { + for col in 0..num_cols { + assert_eq!( + got[q * num_cols + col], + buf[col * lde_size + row as usize], + "base gather mismatch: row {row}, col {col}" + ); + } + } +} + +fn run_ext3(lde_size: usize, num_cols: usize, seed: u64) { + let mut rng = ChaCha8Rng::seed_from_u64(seed); + // De-interleaved ext3 LDE: buf[(col*3 + k)*lde_size + row]. + let buf: Vec = (0..num_cols * 3 * lde_size) + .map(|_| rng.r#gen::()) + .collect(); + + let be = backend().unwrap(); + let stream = be.next_stream(); + let dev = stream.clone_htod(&buf).unwrap(); + stream.synchronize().unwrap(); + let handle = GpuLdeExt3 { + buf: Arc::new(dev), + m: num_cols, + lde_size, + tree: None, + }; + + let rows: Vec = (0..9).map(|_| rng.gen_range(0..lde_size) as u32).collect(); + let got = gather_rows_ext3_on_device(&handle, &rows, &stream).unwrap(); + assert_eq!(got.len(), rows.len() * num_cols * 3, "ext3 gather shape"); + for (q, &row) in rows.iter().enumerate() { + for col in 0..num_cols { + let o = (q * num_cols + col) * 3; + for k in 0..3 { + assert_eq!( + got[o + k], + buf[(col * 3 + k) * lde_size + row as usize], + "ext3 gather mismatch: row {row}, col {col}, comp {k}" + ); + } + } + } +} + +#[test] +fn gather_rows_base_matches_direct_indexing() { + for (log_size, cols) in [(6u32, 3usize), (12, 20), (16, 8)] { + run_base(1usize << log_size, cols, 100 + log_size as u64); + } +} + +#[test] +fn gather_rows_ext3_matches_direct_indexing() { + for (log_size, cols) in [(6u32, 2usize), (12, 5), (16, 3)] { + run_ext3(1usize << log_size, cols, 200 + log_size as u64); + } +} From 757dd38392137a0644c244481d2fb656e855196e Mon Sep 17 00:00:00 2001 From: MauroFab Date: Wed, 8 Jul 2026 12:05:20 -0300 Subject: [PATCH 2/9] =?UTF-8?q?spike(stark):=20Phase=204=20Stage=202=20?= =?UTF-8?q?=E2=80=94=20device-gathered=20R4=20opening=20values=20(cross-ch?= =?UTF-8?q?ecked)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Read each query's row-pair straight off the resident device LDE in open_deep_composition_poly instead of indexing the host trace, keeping the round-1 host D2H this stage. For each query, gather LDE rows reverse_index(2c) and reverse_index(2c+1) via the Stage-1 kernels, lift u64 -> FieldElement through TypeId-gated base_u64_to_field / ext3_u64_to_field (inverse of the *_slice_to_u64 seam), and feed them into a new open_polys_from_values. Every device row is cross-checked (release-active assert_eq!) against gather_main_row / gather_aux_row before it is trusted in the proof. Value gathers are gated on the corresponding device-tree proofs so the two device arms stay aligned and no rows are gathered for a non-resident tree. Preprocessed and host-fallback arms are untouched. De-risks Stage 3 (dropping the host D2H). Non-cuda build unaffected (all new code cuda-gated). --- crypto/stark/src/constraint_ir/gpu_interp.rs | 35 +++++ crypto/stark/src/prover.rs | 137 ++++++++++++++++++- 2 files changed, 166 insertions(+), 6 deletions(-) diff --git a/crypto/stark/src/constraint_ir/gpu_interp.rs b/crypto/stark/src/constraint_ir/gpu_interp.rs index 031692675..de6366d74 100644 --- a/crypto/stark/src/constraint_ir/gpu_interp.rs +++ b/crypto/stark/src/constraint_ir/gpu_interp.rs @@ -16,6 +16,8 @@ //! The whole module is `#[cfg(feature = "cuda")]`; without the feature the //! caller only ever has the CPU interpreter. +use std::any::TypeId; + use math::field::element::FieldElement; use math::field::extensions_goldilocks::Degree3GoldilocksExtensionField as GoldilocksExtension; use math::field::goldilocks::GoldilocksField; @@ -72,6 +74,39 @@ unsafe fn base_slice_as_u64(xs: &[FieldElement]) -> &[u64] { unsafe { std::slice::from_raw_parts(xs.as_ptr() as *const u64, xs.len()) } } +/// Lift raw base-field limbs (one canonical-Goldilocks `u64` per element, as +/// produced by the device row-gather kernels) back into owned `FieldElement`s. +/// Returns `None` unless `F == GoldilocksField`. The inverse of +/// [`base_slice_to_u64`]; used to feed device-gathered LDE rows into the generic +/// prover openings. +pub fn base_u64_to_field(raw: &[u64]) -> Option>> { + if TypeId::of::() != TypeId::of::() { + return None; + } + // SAFETY: the gate established `F == GoldilocksField`, whose `FieldElement` + // is `#[repr(transparent)]` over `u64`; `raw` (a `*const u64`) is 8-aligned. + let fe = + unsafe { std::slice::from_raw_parts(raw.as_ptr() as *const FieldElement, raw.len()) }; + Some(fe.to_vec()) +} + +/// Ext3 sibling of [`base_u64_to_field`]: `raw` holds `3` interleaved limbs per +/// element (`[c0, c1, c2]`). Returns `None` unless `E` is the degree-3 +/// Goldilocks extension. Inverse of [`ext3_slice_to_u64`]. +pub fn ext3_u64_to_field(raw: &[u64]) -> Option>> { + if TypeId::of::() != TypeId::of::() { + return None; + } + debug_assert_eq!(raw.len() % 3, 0, "ext3 limbs must come in triples"); + // SAFETY: the gate established the degree-3 Goldilocks extension, whose + // `FieldElement` is `#[repr(transparent)]` over `[u64; 3]`; `raw` (a + // `*const u64`) is 8-aligned, matching `[u64; 3]`'s alignment. + let fe = unsafe { + std::slice::from_raw_parts(raw.as_ptr() as *const FieldElement, raw.len() / 3) + }; + Some(fe.to_vec()) +} + /// Per-proof accumulation inputs (in `FieldElement` form) for /// [`try_eval_composition_gpu`], mirroring the CPU accumulation in /// `constraints::evaluator`. diff --git a/crypto/stark/src/prover.rs b/crypto/stark/src/prover.rs index 5c03292da..b194b759c 100644 --- a/crypto/stark/src/prover.rs +++ b/crypto/stark/src/prover.rs @@ -1809,6 +1809,36 @@ pub trait IsStarkProver< } } + /// Build a [`PolynomialOpenings`] from a device Merkle proof and a pair of + /// row-value vectors already gathered off the resident device LDE — the + /// fully device-sourced counterpart of [`Self::open_polys_with_proofs`], + /// which still reads its evaluations from the host LDE. + #[cfg(feature = "cuda")] + fn open_polys_from_values( + proof: Proof, + evaluations: Vec>, + evaluations_sym: Vec>, + ) -> PolynomialOpenings { + PolynomialOpenings { + proof, + evaluations, + evaluations_sym, + } + } + + /// Slice out query `qi`'s even/odd row (each `ncols` field elements) from the + /// row-major device gather `[even(q0), odd(q0), even(q1), odd(q1), ...]`. + #[cfg(feature = "cuda")] + fn device_row_pair( + vals: &[FieldElement], + qi: usize, + ncols: usize, + ) -> (Vec>, Vec>) { + let even = vals[(2 * qi) * ncols..(2 * qi + 1) * ncols].to_vec(); + let odd = vals[(2 * qi + 1) * ncols..(2 * qi + 2) * ncols].to_vec(); + (even, odd) + } + /// Open the deep composition polynomial on a list of indexes and their symmetric elements. fn open_deep_composition_poly( domain: &Domain, @@ -1828,6 +1858,23 @@ pub trait IsStarkProver< let num_precomputed_cols = main_commit.num_precomputed_cols; let total_cols = lde_trace.num_main_cols(); + // Row-pair LDE positions for every query, `[even(q0), odd(q0), ...]`. + // Each query opens the leaf at `challenge`, which pairs LDE rows + // `reverse_index(2·challenge)` (the queried point) and + // `reverse_index(2·challenge+1)` (its symmetric `-x` point). + #[cfg(feature = "cuda")] + let domain_size = domain.lde_roots_of_unity_coset.len() as u64; + #[cfg(feature = "cuda")] + let query_rows: Vec = indexes_to_open + .iter() + .flat_map(|&c| { + [ + reverse_index(c * 2, domain_size) as u32, + reverse_index(c * 2 + 1, domain_size) as u32, + ] + }) + .collect(); + // R4 trace proofs from the resident device trees, gathered in one batch // over all query positions instead of walking the host trees (byte // identical to the host proofs, guarded by the `merkle_gather` test). @@ -1882,6 +1929,45 @@ pub trait IsStarkProver< ) }); + // Full-residency Stage 2: gather each query's row-pair straight off the + // resident device LDE (a small D2H of only the queried rows), instead of + // indexing the full host LDE trace. The host trace is still resident this + // stage and every device row is cross-checked against it in the loop + // below; Stage 3 drops the host copy once this path is proven. `None` + // when the LDE is not device resident or the tower is not Goldilocks (→ + // host gather). Row-major: `q`-th row's columns at `[q*ncols ..]`. + // Gate the value gathers on the corresponding device-tree proofs so the + // two device arms stay aligned (`*_dev_proofs.is_some() ⇔ + // *_dev_values.is_some()` on the Goldilocks path) and we never gather + // rows for a tree that is not device resident. + #[cfg(feature = "cuda")] + let main_dev_values: Option>> = + main_dev_proofs.as_ref().and_then(|_| { + lde_trace.gpu_main().and_then(|h| { + let stream = lde_trace + .bound_stream() + .expect("bound stream for device-resident main-row gather"); + let raw = + math_cuda::barycentric::gather_rows_base_on_device(h, &query_rows, &stream) + .expect("device main-row gather failed"); + crate::constraint_ir::gpu_interp::base_u64_to_field::(&raw) + }) + }); + + #[cfg(feature = "cuda")] + let aux_dev_values: Option>> = + aux_dev_proofs.as_ref().and_then(|_| { + lde_trace.gpu_aux().and_then(|h| { + let stream = lde_trace + .bound_stream() + .expect("bound stream for device-resident aux-row gather"); + let raw = + math_cuda::barycentric::gather_rows_ext3_on_device(h, &query_rows, &stream) + .expect("device aux-row gather failed"); + crate::constraint_ir::gpu_interp::ext3_u64_to_field::(&raw) + }) + }); + for (qi, index) in indexes_to_open.iter().enumerate() { #[cfg(not(feature = "cuda"))] let _ = qi; @@ -1895,9 +1981,29 @@ pub trait IsStarkProver< #[cfg(feature = "cuda")] { if let Some(proofs) = &main_dev_proofs { - Self::open_polys_with_proofs(domain, proofs[qi].clone(), *index, |row| { - lde_trace.gather_main_row(row) - }) + let proof = proofs[qi].clone(); + if let Some(dev_vals) = &main_dev_values { + let (even, odd) = Self::device_row_pair(dev_vals, qi, total_cols); + // Cross-check the device gather against the (still + // resident) host LDE before trusting it in the proof. + let r_even = reverse_index(*index * 2, domain_size); + let r_odd = reverse_index(*index * 2 + 1, domain_size); + assert_eq!( + even, + lde_trace.gather_main_row(r_even), + "device main-row gather mismatch (even), query {qi}" + ); + assert_eq!( + odd, + lde_trace.gather_main_row(r_odd), + "device main-row gather mismatch (odd), query {qi}" + ); + Self::open_polys_from_values(proof, even, odd) + } else { + Self::open_polys_with_proofs(domain, proof, *index, |row| { + lde_trace.gather_main_row(row) + }) + } } else { Self::open_polys_with(domain, &main_commit.tree, *index, |row| { lde_trace.gather_main_row(row) @@ -1950,9 +2056,28 @@ pub trait IsStarkProver< #[cfg(feature = "cuda")] { if let Some(proofs) = &aux_dev_proofs { - Self::open_polys_with_proofs(domain, proofs[qi].clone(), *index, |row| { - lde_trace.gather_aux_row(row) - }) + let proof = proofs[qi].clone(); + if let Some(dev_vals) = &aux_dev_values { + let ncols = lde_trace.num_aux_cols(); + let (even, odd) = Self::device_row_pair(dev_vals, qi, ncols); + let r_even = reverse_index(*index * 2, domain_size); + let r_odd = reverse_index(*index * 2 + 1, domain_size); + assert_eq!( + even, + lde_trace.gather_aux_row(r_even), + "device aux-row gather mismatch (even), query {qi}" + ); + assert_eq!( + odd, + lde_trace.gather_aux_row(r_odd), + "device aux-row gather mismatch (odd), query {qi}" + ); + Self::open_polys_from_values(proof, even, odd) + } else { + Self::open_polys_with_proofs(domain, proof, *index, |row| { + lde_trace.gather_aux_row(row) + }) + } } else { Self::open_polys_with(domain, &aux.tree, *index, |row| { lde_trace.gather_aux_row(row) From 58577a64c3a7eea9a3cec492c98212cbc7ccc8f3 Mon Sep 17 00:00:00 2001 From: MauroFab Date: Wed, 8 Jul 2026 12:48:03 -0300 Subject: [PATCH 3/9] test(stark,prover): counter-guard the Stage 2 device-opening gather Add GPU_OPENING_GATHER_CALLS, incremented once per main/aux trace whose R4 query rows are read off the device LDE in open_deep_composition_poly, and a gpu_opening_gather_fires_and_verifies e2e test that proves fib_iterative_1M, asserts the counter fired (no silent revert to the host gather), and verifies the proof. Mirrors the existing GPU dispatch counters (e.g. GPU_COMPOSITION_CALLS) so a future regression that quietly drops the residency path fails loudly. --- crypto/stark/src/gpu_lde.rs | 12 ++++++++++++ crypto/stark/src/prover.rs | 10 ++++++++++ prover/tests/cuda_path_integration.rs | 25 ++++++++++++++++++++++++- 3 files changed, 46 insertions(+), 1 deletion(-) diff --git a/crypto/stark/src/gpu_lde.rs b/crypto/stark/src/gpu_lde.rs index b23adb1ac..142144023 100644 --- a/crypto/stark/src/gpu_lde.rs +++ b/crypto/stark/src/gpu_lde.rs @@ -78,6 +78,7 @@ pub fn reset_all_gpu_call_counters() { GPU_BATCH_INVERT_CALLS.store(0, Ordering::Relaxed); GPU_LOGUP_CALLS.store(0, Ordering::Relaxed); GPU_COMPOSITION_CALLS.store(0, Ordering::Relaxed); + GPU_OPENING_GATHER_CALLS.store(0, Ordering::Relaxed); } pub(crate) static GPU_EXTEND_HALVES_CALLS: AtomicU64 = AtomicU64::new(0); @@ -101,6 +102,17 @@ pub fn gpu_composition_calls() -> u64 { GPU_COMPOSITION_CALLS.load(Ordering::Relaxed) } +/// Successful device-resident-LDE opening-value gathers in +/// `open_deep_composition_poly` — one per main/aux trace whose R4 query rows +/// were read straight off the device LDE instead of the host trace (a +/// non-resident tree or non-Goldilocks tower falls back to the host gather and +/// is not counted). Guards against a silent regression where Stage-2 openings +/// quietly revert to the host path. +pub(crate) static GPU_OPENING_GATHER_CALLS: AtomicU64 = AtomicU64::new(0); +pub fn gpu_opening_gather_calls() -> u64 { + GPU_OPENING_GATHER_CALLS.load(Ordering::Relaxed) +} + /// Runtime override to force the GPU composition path off (→ CPU accumulation). /// An escape hatch, and the A/B toggle for benchmarking the path against the CPU /// baseline in one process (no rebuild). Default off (path enabled). diff --git a/crypto/stark/src/prover.rs b/crypto/stark/src/prover.rs index b194b759c..2c17b2129 100644 --- a/crypto/stark/src/prover.rs +++ b/crypto/stark/src/prover.rs @@ -1953,6 +1953,11 @@ pub trait IsStarkProver< crate::constraint_ir::gpu_interp::base_u64_to_field::(&raw) }) }); + #[cfg(feature = "cuda")] + if main_dev_values.is_some() { + crate::gpu_lde::GPU_OPENING_GATHER_CALLS + .fetch_add(1, std::sync::atomic::Ordering::Relaxed); + } #[cfg(feature = "cuda")] let aux_dev_values: Option>> = @@ -1967,6 +1972,11 @@ pub trait IsStarkProver< crate::constraint_ir::gpu_interp::ext3_u64_to_field::(&raw) }) }); + #[cfg(feature = "cuda")] + if aux_dev_values.is_some() { + crate::gpu_lde::GPU_OPENING_GATHER_CALLS + .fetch_add(1, std::sync::atomic::Ordering::Relaxed); + } for (qi, index) in indexes_to_open.iter().enumerate() { #[cfg(not(feature = "cuda"))] diff --git a/prover/tests/cuda_path_integration.rs b/prover/tests/cuda_path_integration.rs index 7081baadf..e45d5d9a2 100644 --- a/prover/tests/cuda_path_integration.rs +++ b/prover/tests/cuda_path_integration.rs @@ -13,7 +13,7 @@ use lambda_vm_prover::{prove, verify}; use stark::gpu_lde::{ gpu_bary_calls, gpu_batch_invert_calls, gpu_comp_poly_tree_calls, gpu_composition_calls, gpu_deep_calls, gpu_extend_halves_calls, gpu_fri_calls, gpu_lde_calls, gpu_logup_calls, - gpu_parts_lde_calls, reset_all_gpu_call_counters, + gpu_opening_gather_calls, gpu_parts_lde_calls, reset_all_gpu_call_counters, }; /// The R2 GPU composition-poly path (fused `H = z·Σβᵢ·Cᵢ + boundary`) fires and @@ -153,3 +153,26 @@ fn gpu_proof_verifies_row_pair_commitment() { "GPU-produced proof (row-pair commitment) failed verification" ); } + +/// The full-residency Stage-2 R4 opening path fires: query row values are +/// gathered straight off the device LDE (not the host trace) and, guarded by the +/// in-prover cross-check against the host gather, still yield a verifying proof. +/// Guards a silent regression where the openings quietly revert to the host +/// path (which would still verify but drop the data-residency win). The +/// cross-check `assert_eq!`s inside `open_deep_composition_poly` are +/// release-active, so a divergent device gather would panic the prove here. +#[test] +#[ignore = "requires GPU; run with --ignored --nocapture"] +fn gpu_opening_gather_fires_and_verifies() { + let elf = asm_elf_bytes("fib_iterative_1M"); + reset_all_gpu_call_counters(); + let proof = prove(&elf).expect("prove"); + assert!( + gpu_opening_gather_calls() > 0, + "device-resident opening gather did not fire (openings fell back to the host LDE)" + ); + assert!( + verify(&proof, &elf).expect("verify"), + "GPU-produced proof (device-gathered openings) failed verification" + ); +} From 248f971759a680b167694d82b598b39d144962c6 Mon Sep 17 00:00:00 2001 From: MauroFab Date: Wed, 8 Jul 2026 14:05:52 -0300 Subject: [PATCH 4/9] =?UTF-8?q?spike(stark):=20Stage=203=20step=201=20?= =?UTF-8?q?=E2=80=94=20host=5Ftrace=5Fempty=20flag=20+=20hard-abort=20guar?= =?UTF-8?q?ds?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Scaffolding for the device-only round-1 residency path (D2H skip lands next). Adds a per-table `host_trace_empty` flag on LDETraceTable (default false, so behavior is identical) and hard-abort `assert!(!host_trace_empty)` guards on every host-LDE-trace fallback that would otherwise index the (future) empty buffers: R2 composition CPU path (boundary + transition), R3 barycentric main/aux host fallbacks, R4 DEEP host loop, and the R4 opening host-fallback arms. The Stage-2 device-gather cross-checks are gated on `!host_trace_empty` so they are skipped when there is no host copy to check against. All guards are inert until step 2 wires the `device_only` gate. --- crypto/stark/src/constraints/evaluator.rs | 11 +++ crypto/stark/src/prover.rs | 87 ++++++++++++++++------- crypto/stark/src/trace.rs | 42 +++++++++++ 3 files changed, 114 insertions(+), 26 deletions(-) diff --git a/crypto/stark/src/constraints/evaluator.rs b/crypto/stark/src/constraints/evaluator.rs index 4b7696bcf..4e82a7a1b 100644 --- a/crypto/stark/src/constraints/evaluator.rs +++ b/crypto/stark/src/constraints/evaluator.rs @@ -233,6 +233,17 @@ where } } + // Reaching here means the GPU composition path fell through to the CPU + // boundary + transition evaluation below, both of which read the host + // trace (`get_main`/`get_aux`, `RowFrame::from_lde`). Under the + // device-only gate that trace is empty, so a fall-through is a mis-gate + // or an unexpected GPU failure: hard-abort rather than read empty buffers. + #[cfg(feature = "cuda")] + assert!( + !lde_trace.host_trace_empty(), + "R2 composition fell back to the host trace, but it is device-only (empty)" + ); + // Fused boundary evaluation: compute (trace[col] - value) on-the-fly // instead of pre-computing all boundary_polys_evaluations. // This eliminates N_constraints × LDE_size intermediate allocations. diff --git a/crypto/stark/src/prover.rs b/crypto/stark/src/prover.rs index 2c17b2129..95d21686b 100644 --- a/crypto/stark/src/prover.rs +++ b/crypto/stark/src/prover.rs @@ -1611,6 +1611,16 @@ pub trait IsStarkProver< } } + // Reaching here means both GPU DEEP arms fell through to the host loop + // below (which reads `get_main`/`get_aux`). Under the device-only gate + // the host trace is empty, so a fall-through is a mis-gate or an + // unexpected GPU failure: hard-abort rather than read empty buffers. + #[cfg(feature = "cuda")] + assert!( + !lde_trace.host_trace_empty(), + "R4 DEEP composition fell back to the host trace, but it is device-only (empty)" + ); + // OOD column compression (Plonky3-style): precompute one value per eval point, // ood_compressed_k = Σ_j gamma[j][k] * ood[j][k]. // The per-LDE-point trace column sums are NOT precomputed — they are fused @@ -1994,27 +2004,41 @@ pub trait IsStarkProver< let proof = proofs[qi].clone(); if let Some(dev_vals) = &main_dev_values { let (even, odd) = Self::device_row_pair(dev_vals, qi, total_cols); - // Cross-check the device gather against the (still - // resident) host LDE before trusting it in the proof. - let r_even = reverse_index(*index * 2, domain_size); - let r_odd = reverse_index(*index * 2 + 1, domain_size); - assert_eq!( - even, - lde_trace.gather_main_row(r_even), - "device main-row gather mismatch (even), query {qi}" - ); - assert_eq!( - odd, - lde_trace.gather_main_row(r_odd), - "device main-row gather mismatch (odd), query {qi}" - ); + // Cross-check the device gather against the host LDE. + // Skipped under device-only (host trace empty): the + // gather was proven bit-identical while the host copy + // was resident, and there is nothing to check against. + if !lde_trace.host_trace_empty() { + let r_even = reverse_index(*index * 2, domain_size); + let r_odd = reverse_index(*index * 2 + 1, domain_size); + assert_eq!( + even, + lde_trace.gather_main_row(r_even), + "device main-row gather mismatch (even), query {qi}" + ); + assert_eq!( + odd, + lde_trace.gather_main_row(r_odd), + "device main-row gather mismatch (odd), query {qi}" + ); + } Self::open_polys_from_values(proof, even, odd) } else { + // Device tree resident but the value gather is absent: + // reading the host gather is invalid under device-only. + assert!( + !lde_trace.host_trace_empty(), + "R4 main opening fell back to the host gather, but it is device-only (empty)" + ); Self::open_polys_with_proofs(domain, proof, *index, |row| { lde_trace.gather_main_row(row) }) } } else { + assert!( + !lde_trace.host_trace_empty(), + "R4 main opening fell back to the host tree, but it is device-only (empty)" + ); Self::open_polys_with(domain, &main_commit.tree, *index, |row| { lde_trace.gather_main_row(row) }) @@ -2070,25 +2094,36 @@ pub trait IsStarkProver< if let Some(dev_vals) = &aux_dev_values { let ncols = lde_trace.num_aux_cols(); let (even, odd) = Self::device_row_pair(dev_vals, qi, ncols); - let r_even = reverse_index(*index * 2, domain_size); - let r_odd = reverse_index(*index * 2 + 1, domain_size); - assert_eq!( - even, - lde_trace.gather_aux_row(r_even), - "device aux-row gather mismatch (even), query {qi}" - ); - assert_eq!( - odd, - lde_trace.gather_aux_row(r_odd), - "device aux-row gather mismatch (odd), query {qi}" - ); + // Cross-check skipped under device-only (host empty). + if !lde_trace.host_trace_empty() { + let r_even = reverse_index(*index * 2, domain_size); + let r_odd = reverse_index(*index * 2 + 1, domain_size); + assert_eq!( + even, + lde_trace.gather_aux_row(r_even), + "device aux-row gather mismatch (even), query {qi}" + ); + assert_eq!( + odd, + lde_trace.gather_aux_row(r_odd), + "device aux-row gather mismatch (odd), query {qi}" + ); + } Self::open_polys_from_values(proof, even, odd) } else { + assert!( + !lde_trace.host_trace_empty(), + "R4 aux opening fell back to the host gather, but it is device-only (empty)" + ); Self::open_polys_with_proofs(domain, proof, *index, |row| { lde_trace.gather_aux_row(row) }) } } else { + assert!( + !lde_trace.host_trace_empty(), + "R4 aux opening fell back to the host tree, but it is device-only (empty)" + ); Self::open_polys_with(domain, &aux.tree, *index, |row| { lde_trace.gather_aux_row(row) }) diff --git a/crypto/stark/src/trace.rs b/crypto/stark/src/trace.rs index 6d40425b7..07e1987e3 100644 --- a/crypto/stark/src/trace.rs +++ b/crypto/stark/src/trace.rs @@ -327,6 +327,14 @@ where pub(crate) num_rows: usize, 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. Always false today until the `device_only` gate is wired. + #[cfg(feature = "cuda")] + pub(crate) host_trace_empty: bool, /// Per table GPU residency session: owns this table's device LDE buffers /// and bound stream. Threaded R1 to R4. Empty on the CPU path. #[cfg(feature = "cuda")] @@ -459,6 +467,8 @@ where lde_step_size, blowup_factor, #[cfg(feature = "cuda")] + host_trace_empty: false, + #[cfg(feature = "cuda")] gpu_session: GpuTableSession::new(), } } @@ -494,6 +504,8 @@ where lde_step_size, blowup_factor, #[cfg(feature = "cuda")] + host_trace_empty: false, + #[cfg(feature = "cuda")] gpu_session: GpuTableSession::new(), } } @@ -511,6 +523,22 @@ where self.gpu_session.aux_lde = Some(h); } + /// 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. + #[cfg(feature = "cuda")] + pub fn set_host_trace_empty(&mut self, empty: bool) { + self.host_trace_empty = empty; + } + + /// 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`. + #[cfg(feature = "cuda")] + pub fn host_trace_empty(&self) -> bool { + self.host_trace_empty + } + #[cfg(feature = "cuda")] pub fn gpu_main(&self) -> Option<&math_cuda::lde::GpuLdeBase> { self.gpu_session.main_lde.as_ref() @@ -742,6 +770,13 @@ where let main_evals: Vec> = if let Some(v) = main_gpu { 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. + #[cfg(feature = "cuda")] + assert!( + !lde_trace.host_trace_empty(), + "R3 barycentric (main) fell back to the host trace, but it is device-only (empty)" + ); let inv_denoms_v = inv_denoms.get_or_insert_with(|| barycentric_inv_denoms(eval_point, &dc.points)); let col_scale = col_scale.get_or_insert_with(|| { @@ -793,6 +828,13 @@ where let aux_evals: Vec> = if let Some(v) = aux_gpu { 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. + #[cfg(feature = "cuda")] + assert!( + !lde_trace.host_trace_empty(), + "R3 barycentric (aux) fell back to the host trace, but it is device-only (empty)" + ); let inv_denoms_v = inv_denoms.get_or_insert_with(|| barycentric_inv_denoms(eval_point, &dc.points)); let col_scale = col_scale.get_or_insert_with(|| { From c7599171076ace89e89dbc10b2a13a59e8790b54 Mon Sep 17 00:00:00 2001 From: MauroFab Date: Wed, 8 Jul 2026 14:18:50 -0300 Subject: [PATCH 5/9] =?UTF-8?q?spike(stark):=20Stage=203=20step=202=20?= =?UTF-8?q?=E2=80=94=20device-only=20gate=20skips=20the=20round-1=20trace?= =?UTF-8?q?=20D2H?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Wire the full-residency win: when a table's `device_only` gate holds, keep its round-1 LDE device-resident and skip the host D2H (the big transfer/alloc the prover paid every round-1). Threads a `retain_host_lde` bool through the math-cuda keep wrappers (`coset_lde_row_major_inner` returns an empty host Vec when false) and the three stark keep entries. `device_only_gate` (gpu_lde.rs) is the strict AND of every precondition that guarantees the R2 composition, R3 barycentric, R4 DEEP and R4 opening GPU paths all fire and read the device LDE: Goldilocks+ext3, lde_size>=gpu_lde_threshold, n>=gpu_bary_threshold, power-of-two, !preprocessed, contiguous offsets, uniform zerofier (all end_exemptions==0), and neither composition nor device-only disabled; forced off under debug-checks so reconstruct_round1 keeps a host trace. `Prover::device_only_for` derives it from the AIR so the main (Phase A) and aux commit closures compute the identical value and skip both host buffers together. build_round1 infers `host_trace_empty` from the ACTUAL buffer state (empty host buffer where data was expected) — a safety property: if the gate held but the GPU keep fell back to CPU, the buffer is populated and the proof runs on the host trace as normal. num_rows is recovered from the device handle's lde_size when the main buffer is empty. Adds GPU_DEVICE_ONLY_CALLS + a LAMBDA_VM_DISABLE_DEVICE_ONLY A/B toggle + a gpu_device_only_residency_fires_and_verifies e2e test. Step-1 host_trace_empty guards are now live. --- crypto/math-cuda/src/lde.rs | 15 ++++- crypto/stark/src/gpu_lde.rs | 79 +++++++++++++++++++++++++ crypto/stark/src/prover.rs | 85 ++++++++++++++++++++++++++- crypto/stark/src/trace.rs | 9 +++ prover/tests/cuda_path_integration.rs | 27 ++++++++- 5 files changed, 211 insertions(+), 4 deletions(-) diff --git a/crypto/math-cuda/src/lde.rs b/crypto/math-cuda/src/lde.rs index 30e524c2c..06c1a02cd 100644 --- a/crypto/math-cuda/src/lde.rs +++ b/crypto/math-cuda/src/lde.rs @@ -400,6 +400,7 @@ enum InnerInput<'a> { /// required by the downstream GPU kernels DEEP/barycentric); callers wrap it in /// the appropriate LDE handle. #[allow(clippy::type_complexity)] +#[allow(clippy::too_many_arguments)] fn coset_lde_row_major_inner( input: InnerInput, n: usize, @@ -408,6 +409,7 @@ fn coset_lde_row_major_inner( weights: &[u64], what: &str, retain_trace_col_major: bool, + retain_host_lde: bool, ) -> Result<( GpuMerkleTree, CudaSlice, @@ -514,7 +516,10 @@ fn coset_lde_row_major_inner( // D2H the row-major LDE first (before the handle transpose). Release the // staging lock before the Merkle nodes transfer to minimise lock contention. - let lde_out = { + // Skipped entirely when `retain_host_lde` is false (the caller keeps the LDE + // device-only): this is the round-1 trace D2H the full-residency path + // eliminates — the big transfer/alloc win — so we return an empty host Vec. + let lde_out = if retain_host_lde { let staging_slot = be.pinned_staging(); let mut staging = staging_slot.lock().unwrap(); staging.ensure_capacity(lde_size * total_cols, &be.ctx)?; @@ -524,6 +529,8 @@ fn coset_lde_row_major_inner( let out = pinned[..lde_size * total_cols].to_vec(); drop(staging); out + } else { + Vec::new() }; // Keep the Merkle tree resident on device; copy only the 32 byte root so the @@ -562,6 +569,7 @@ pub fn coset_lde_row_major_with_merkle_tree_keep( m: usize, blowup_factor: usize, weights: &[u64], + retain_host_lde: bool, ) -> Result<(GpuLdeBase, Vec)> { let (tree, col_major_dev, lde_out, trace_col_major) = coset_lde_row_major_inner( InnerInput::Host(row_major), @@ -571,6 +579,7 @@ pub fn coset_lde_row_major_with_merkle_tree_keep( weights, "coset_lde_row_major lde_size", true, + retain_host_lde, )?; let handle = GpuLdeBase { buf: Arc::new(col_major_dev), @@ -599,6 +608,7 @@ pub fn coset_lde_ext3_row_major_with_merkle_tree_keep( m: usize, blowup_factor: usize, weights: &[u64], + retain_host_lde: bool, ) -> Result<(GpuLdeExt3, Vec)> { let (tree, col_major_dev, lde_out, _) = coset_lde_row_major_inner( InnerInput::Host(row_major), @@ -608,6 +618,7 @@ pub fn coset_lde_ext3_row_major_with_merkle_tree_keep( weights, "coset_lde_ext3_row_major lde_size", false, + retain_host_lde, )?; let handle = GpuLdeExt3 { buf: Arc::new(col_major_dev), @@ -628,6 +639,7 @@ pub fn coset_lde_ext3_row_major_with_merkle_tree_keep_dev( m: usize, blowup_factor: usize, weights: &[u64], + retain_host_lde: bool, ) -> Result<(GpuLdeExt3, Vec)> { let (tree, col_major_dev, lde_out, _) = coset_lde_row_major_inner( InnerInput::Dev(input_dev), @@ -637,6 +649,7 @@ pub fn coset_lde_ext3_row_major_with_merkle_tree_keep_dev( weights, "coset_lde_ext3_row_major_dev lde_size", false, + retain_host_lde, )?; let handle = GpuLdeExt3 { buf: Arc::new(col_major_dev), diff --git a/crypto/stark/src/gpu_lde.rs b/crypto/stark/src/gpu_lde.rs index 142144023..9c21717d2 100644 --- a/crypto/stark/src/gpu_lde.rs +++ b/crypto/stark/src/gpu_lde.rs @@ -79,6 +79,7 @@ pub fn reset_all_gpu_call_counters() { GPU_LOGUP_CALLS.store(0, Ordering::Relaxed); GPU_COMPOSITION_CALLS.store(0, Ordering::Relaxed); GPU_OPENING_GATHER_CALLS.store(0, Ordering::Relaxed); + GPU_DEVICE_ONLY_CALLS.store(0, Ordering::Relaxed); } pub(crate) static GPU_EXTEND_HALVES_CALLS: AtomicU64 = AtomicU64::new(0); @@ -113,6 +114,16 @@ pub fn gpu_opening_gather_calls() -> u64 { GPU_OPENING_GATHER_CALLS.load(Ordering::Relaxed) } +/// Tables whose round-1 LDE was kept device-only (host trace D2H skipped) — the +/// Stage-3 full-residency win. Incremented once per main trace that took the +/// `device_only` path. Zero means every table kept its host copy (gate never +/// engaged), so a residency regression drops this to 0 while proofs still +/// verify. +pub(crate) static GPU_DEVICE_ONLY_CALLS: AtomicU64 = AtomicU64::new(0); +pub fn gpu_device_only_calls() -> u64 { + GPU_DEVICE_ONLY_CALLS.load(Ordering::Relaxed) +} + /// Runtime override to force the GPU composition path off (→ CPU accumulation). /// An escape hatch, and the A/B toggle for benchmarking the path against the CPU /// baseline in one process (no rebuild). Default off (path enabled). @@ -135,6 +146,66 @@ pub(crate) fn gpu_composition_disabled() -> bool { }) } +/// Runtime override to force the Stage-3 device-only path off (keeps the round-1 +/// host D2H). Independent of the composition toggle, so the residency win can be +/// A/B-benched with the GPU composition path left on. Default off (path enabled). +static DEVICE_ONLY_DISABLED: std::sync::atomic::AtomicBool = + std::sync::atomic::AtomicBool::new(false); +pub fn set_device_only_disabled(v: bool) { + DEVICE_ONLY_DISABLED.store(v, Ordering::Relaxed); +} +pub(crate) fn device_only_disabled() -> bool { + if DEVICE_ONLY_DISABLED.load(Ordering::Relaxed) { + return true; + } + // Env fallback (cached): `LAMBDA_VM_DISABLE_DEVICE_ONLY=1`. + static ENV_DISABLED: OnceLock = OnceLock::new(); + *ENV_DISABLED.get_or_init(|| { + std::env::var("LAMBDA_VM_DISABLE_DEVICE_ONLY") + .map(|v| v == "1") + .unwrap_or(false) + }) +} + +/// 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. +/// +/// `zerofier_uniform` must be the R1-derived conservative form (all constraints +/// share `end_exemptions == 0`), which implies `ZerofierEvaluations::is_uniform` +/// (a single cyclic group) — the condition the GPU composition kernel needs. +pub(crate) fn device_only_gate( + lde_size: usize, + n: usize, + is_preprocessed: bool, + offsets_contiguous: bool, + zerofier_uniform: bool, +) -> bool +where + F: 'static, + E: 'static, +{ + // debug-checks reconstruct the LDE from the host trace — keep it resident. + if cfg!(feature = "debug-checks") { + return false; + } + TypeId::of::() == TypeId::of::() + && TypeId::of::() == TypeId::of::() + && !device_only_disabled() + && !gpu_composition_disabled() + && lde_size.is_power_of_two() + && lde_size >= gpu_lde_threshold() + && n >= gpu_bary_threshold() + && !is_preprocessed + && offsets_contiguous + && zerofier_uniform +} + /// `true` when the field tower is concrete Goldilocks + its degree-3 extension — /// the only tower with a CUDA lowering. The one home of this check: every GPU /// dispatch gate calls it, so the tower test cannot drift between sites. @@ -508,6 +579,7 @@ pub(crate) fn try_expand_leaf_and_tree_row_major_keep( m: usize, blowup_factor: usize, weights: &[FieldElement], + retain_host_lde: bool, ) -> Option<( MerkleTree, math_cuda::lde::GpuLdeBase, @@ -540,12 +612,14 @@ where GPU_MERKLE_TREE_CALLS.fetch_add(1, Ordering::Relaxed); // The keep path keeps the Merkle tree resident on device (in `handle.tree`). + // `retain_host_lde=false` additionally skips the row-major D2H (device-only). let (handle, lde_u64) = math_cuda::lde::coset_lde_row_major_with_merkle_tree_keep( raw, n, m, blowup_factor, &weights_u64, + retain_host_lde, ) .ok()?; @@ -575,6 +649,7 @@ pub(crate) fn try_expand_leaf_and_tree_ext3_row_major_keep( m: usize, blowup_factor: usize, weights: &[FieldElement], + retain_host_lde: bool, ) -> Option<( MerkleTree, math_cuda::lde::GpuLdeExt3, @@ -609,12 +684,14 @@ where GPU_MERKLE_TREE_CALLS.fetch_add(1, Ordering::Relaxed); // The keep path keeps the Merkle tree resident on device (in `handle.tree`). + // `retain_host_lde=false` additionally skips the row-major D2H (device-only). let (handle, lde_u64) = math_cuda::lde::coset_lde_ext3_row_major_with_merkle_tree_keep( raw, n, m, blowup_factor, &weights_u64, + retain_host_lde, ) .ok()?; @@ -1176,6 +1253,7 @@ pub(crate) fn try_expand_leaf_and_tree_ext3_row_major_keep_dev( ra: &math_cuda::logup::ResidentAux, blowup_factor: usize, weights: &[FieldElement], + retain_host_lde: bool, ) -> Option<( MerkleTree, math_cuda::lde::GpuLdeExt3, @@ -1203,6 +1281,7 @@ where ra.num_aux_cols, blowup_factor, &weights_u64, + retain_host_lde, ) .ok()?; diff --git a/crypto/stark/src/prover.rs b/crypto/stark/src/prover.rs index 95d21686b..352f5acd4 100644 --- a/crypto/stark/src/prover.rs +++ b/crypto/stark/src/prover.rs @@ -254,6 +254,29 @@ where ) -> Round1 { let (main_data, num_main_cols) = lde.main; let (aux_data, num_aux_cols) = lde.aux; + + // Stage-3 device-only detection, inferred from the ACTUAL buffer state + // (not the gate's intent): a table whose round-1 D2H was skipped has an + // empty host buffer where it should have data. Using the real state is a + // 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. `main_empty` additionally + // drives the `num_rows` recovery, since `from_row_major` reads it from + // `main_data.len()`. 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. + #[cfg(feature = "cuda")] + let main_empty = num_main_cols > 0 && main_data.is_empty(); + #[cfg(feature = "cuda")] + let host_trace_empty = + main_empty || (num_aux_cols > 0 && aux_data.is_empty() && lde.gpu_aux.is_some()); + #[cfg(feature = "cuda")] + let device_num_rows = lde + .gpu_main + .as_ref() + .map(|h| h.lde_size) + .or_else(|| lde.gpu_aux.as_ref().map(|h| h.lde_size)); + #[allow(unused_mut)] let mut lde_trace = LDETraceTable::from_row_major( main_data, @@ -265,6 +288,14 @@ where ); #[cfg(feature = "cuda")] { + if host_trace_empty { + // Recover the LDE row count from the resident handle when the + // (empty) main buffer could not supply it. + if main_empty && let Some(n) = device_num_rows { + lde_trace.set_num_rows(n); + } + lde_trace.set_host_trace_empty(true); + } if let Some(h) = lde.gpu_main { lde_trace.set_gpu_main(h); } @@ -761,11 +792,37 @@ pub trait IsStarkProver< /// as a separate Merkle tree (the precomputed split for preprocessed /// tables) and the root is checked against the AIR-hardcoded commitment. #[allow(clippy::type_complexity)] + /// 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. + #[cfg(feature = "cuda")] + fn device_only_for( + air: &dyn AIR, + domain: &Domain, + ) -> bool { + let lde_size = domain.interpolation_domain_size * domain.blowup_factor; + let n = domain.interpolation_domain_size; + let offsets = &air.context().transition_offsets; + let offsets_contiguous = offsets.iter().enumerate().all(|(i, &o)| o == i); + let zerofier_uniform = air.constraints_meta().iter().all(|m| m.end_exemptions == 0); + crate::gpu_lde::device_only_gate::( + lde_size, + n, + air.is_preprocessed(), + offsets_contiguous, + zerofier_uniform, + ) + } + fn commit_main_trace( trace: &TraceTable, domain: &Domain, twiddles: &LdeTwiddles, precomputed: Option<(Commitment, usize)>, + #[cfg(feature = "cuda")] device_only: bool, #[cfg(feature = "disk-spill")] storage_mode: StorageMode, ) -> Result, ProvingError> where @@ -798,6 +855,7 @@ pub trait IsStarkProver< num_cols, domain.blowup_factor, &twiddles.coset_weights, + !device_only, ) { #[cfg(feature = "instruments")] @@ -805,6 +863,13 @@ pub trait IsStarkProver< let root = tree.root; #[cfg(feature = "instruments")] crate::instruments::accum_r1_main(main_lde_dur, std::time::Duration::ZERO); + // Count a device-only main commit only once the GPU keep path + // actually fired (handle produced + host trace intentionally + // empty), so the counter reflects real residency, not the gate. + if device_only { + crate::gpu_lde::GPU_DEVICE_ONLY_CALLS + .fetch_add(1, std::sync::atomic::Ordering::Relaxed); + } return Ok(( TableCommit::plain(tree, root), (main_data, num_cols), @@ -2325,11 +2390,19 @@ pub trait IsStarkProver< let precomputed = air .is_preprocessed() .then(|| (air.precomputed_commitment(), air.num_precomputed_columns())); + + // Stage-3 device-only gate: when it holds, `commit_main_trace` + // keeps the R1 LDE device-resident and skips the host D2H. + #[cfg(feature = "cuda")] + let device_only = Self::device_only_for(*air, domain); + Self::commit_main_trace( *trace, domain, twiddles, precomputed, + #[cfg(feature = "cuda")] + device_only, #[cfg(feature = "disk-spill")] storage_mode, ) @@ -2533,6 +2606,12 @@ pub trait IsStarkProver< if air.has_aux_trace() { let lde_size = domain.interpolation_domain_size * domain.blowup_factor; + // Same gate as the main commit (Phase A): skip the aux + // 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); + // 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 @@ -2549,7 +2628,10 @@ pub trait IsStarkProver< FieldExtension, BatchedMerkleTreeBackend, >( - ra, domain.blowup_factor, &twiddles.coset_weights + ra, + domain.blowup_factor, + &twiddles.coset_weights, + !device_only, ) .ok_or_else(|| { ProvingError::Fft( @@ -2591,6 +2673,7 @@ pub trait IsStarkProver< num_cols, domain.blowup_factor, &twiddles.coset_weights, + !device_only, ) { #[cfg(feature = "instruments")] diff --git a/crypto/stark/src/trace.rs b/crypto/stark/src/trace.rs index 07e1987e3..280aabb95 100644 --- a/crypto/stark/src/trace.rs +++ b/crypto/stark/src/trace.rs @@ -531,6 +531,15 @@ where self.host_trace_empty = empty; } + /// Override the LDE row count. Needed on the device-only path: the host + /// buffers are empty, so `from_row_major` cannot infer `num_rows` from + /// `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; + } + /// 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 e45d5d9a2..b60cb3a34 100644 --- a/prover/tests/cuda_path_integration.rs +++ b/prover/tests/cuda_path_integration.rs @@ -12,8 +12,8 @@ use lambda_vm_prover::test_utils::asm_elf_bytes; use lambda_vm_prover::{prove, verify}; use stark::gpu_lde::{ gpu_bary_calls, gpu_batch_invert_calls, gpu_comp_poly_tree_calls, gpu_composition_calls, - gpu_deep_calls, gpu_extend_halves_calls, gpu_fri_calls, gpu_lde_calls, gpu_logup_calls, - gpu_opening_gather_calls, gpu_parts_lde_calls, reset_all_gpu_call_counters, + gpu_deep_calls, gpu_device_only_calls, gpu_extend_halves_calls, gpu_fri_calls, gpu_lde_calls, + gpu_logup_calls, gpu_opening_gather_calls, gpu_parts_lde_calls, reset_all_gpu_call_counters, }; /// The R2 GPU composition-poly path (fused `H = z·Σβᵢ·Cᵢ + boundary`) fires and @@ -176,3 +176,26 @@ fn gpu_opening_gather_fires_and_verifies() { "GPU-produced proof (device-gathered openings) failed verification" ); } + +/// The full-residency Stage-3 device-only path fires: at least one table keeps +/// its round-1 LDE device-resident (the host D2H is skipped), and the proof +/// still verifies. This exercises every `host_trace_empty` hard-abort guard on +/// 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. +#[test] +#[ignore = "requires GPU; run with --ignored --nocapture"] +fn gpu_device_only_residency_fires_and_verifies() { + let elf = asm_elf_bytes("fib_iterative_1M"); + reset_all_gpu_call_counters(); + let proof = prove(&elf).expect("prove"); + assert!( + gpu_device_only_calls() > 0, + "device-only residency path did not fire (every table kept its host trace)" + ); + assert!( + verify(&proof, &elf).expect("verify"), + "GPU-produced proof (device-only residency) failed verification" + ); +} From 8edfed1b19932dc506e255f93233aee4604bd040 Mon Sep 17 00:00:00 2001 From: MauroFab Date: Wed, 8 Jul 2026 17:14:40 -0300 Subject: [PATCH 6/9] bench(stark): continuation mode for bench_abba.sh (two-ref GPU residency A/B) Add CONTINUATIONS=1 + EPOCH_SIZE_LOG2 (default 20) + TX_COUNT (default 5) env knobs so the two-ref ABBA harness can prove with --continuations on both sides. Monolithic proving OOMs on large workloads (e.g. ethrex-20tx) and can't surface a per-epoch GPU-residency change at a realistic trace size; continuation mode proves epoch-by-epoch with flat peak memory. Backward compatible (default off = monolithic as before). Also fixes the missing-fixture generator to honor TX_COUNT instead of hardcoding 5. Usage for the device-only residency win vs main: CONTINUATIONS=1 EPOCH_SIZE_LOG2=20 TX_COUNT=20 \ BENCH_FEATURES=jemalloc-stats,prover/cuda \ scripts/bench_abba.sh HEAD origin/main 5 --- scripts/bench_abba.sh | 28 ++++++++++++++++++++++++---- 1 file changed, 24 insertions(+), 4 deletions(-) diff --git a/scripts/bench_abba.sh b/scripts/bench_abba.sh index 9acc7ae86..fd0e7ad3e 100755 --- a/scripts/bench_abba.sh +++ b/scripts/bench_abba.sh @@ -30,6 +30,11 @@ # Env: REBUILD=1 forces a rebuild even if cached binaries exist. # BENCH_FEATURES= cargo features for the cli build (default: jemalloc-stats). # The GPU ABBA workflow passes "jemalloc-stats,prover/cuda" to bench the GPU path. +# CONTINUATIONS=1 proves with --continuations (epochs; flat peak memory) on +# both sides — needed for large workloads that OOM monolithically. +# EPOCH_SIZE_LOG2= continuation epoch size (default 20; min 18). +# TX_COUNT= ethrex transfer fixture to prove (default 5; use 20 for a +# large continuation trace where GPU-residency wins are visible). # # Sizing (ethrex pair-noise sd ~1.2%, 80% power): ~12 pairs for a 1% effect, # ~18 for 0.8%, ~32 for 0.6%. Default 20 -> solid on 0.8-1%, ~60% power at 0.6% @@ -51,9 +56,23 @@ N_PAIRS="${3:-20}" # cli build features. Default matches the CPU bench; the GPU ABBA workflow overrides # with "jemalloc-stats,prover/cuda" to exercise the CUDA prover path. BENCH_FEATURES="${BENCH_FEATURES:-jemalloc-stats}" +# Continuation mode: split execution into epochs (flat peak memory) instead of a +# single monolithic prove. Required for large workloads (e.g. ethrex-20tx) that +# OOM monolithically, and the only way to bench a per-epoch GPU-residency change +# at a realistic trace size. When on, both binaries prove with +# `--continuations --epoch-size-log2 $EPOCH_SIZE_LOG2`. +CONTINUATIONS="${CONTINUATIONS:-0}" +EPOCH_SIZE_LOG2="${EPOCH_SIZE_LOG2:-20}" +# ethrex transfer-count fixture to prove (executor/tests/ethrex_${TX_COUNT}_transfers.bin). +TX_COUNT="${TX_COUNT:-5}" +if [ "$CONTINUATIONS" = "1" ]; then + CONT_ARGS="--continuations --epoch-size-log2 $EPOCH_SIZE_LOG2" +else + CONT_ARGS="" +fi ELF_REL="executor/program_artifacts/rust/ethrex.elf" -INPUT_REL="executor/tests/ethrex_5_transfers.bin" +INPUT_REL="executor/tests/ethrex_${TX_COUNT}_transfers.bin" WORK="/tmp/abba_run" WT="/tmp/abba_wt" PROOF="/tmp/abba_proof.bin" @@ -84,9 +103,9 @@ if [ ! -f "$ELF_REL" ]; then make "$ELF_REL" fi if [ ! -f "$INPUT_REL" ]; then - echo "==> Generating ethrex 5-transfer fixture (missing)" + echo "==> Generating ethrex ${TX_COUNT}-transfer fixture (missing)" ( cd tooling/ethrex-fixtures && cargo build --release ) - tooling/ethrex-fixtures/target/release/ethrex-fixtures 5 "$INPUT_REL" distinct + tooling/ethrex-fixtures/target/release/ethrex-fixtures "$TX_COUNT" "$INPUT_REL" distinct fi ELF="$(cd "$(dirname "$ELF_REL")" && pwd)/$(basename "$ELF_REL")" INPUT="$(cd "$(dirname "$INPUT_REL")" && pwd)/$(basename "$INPUT_REL")" @@ -143,7 +162,8 @@ fi # --- 3. Interleaved A/B/B/A measurement (fresh CSV -- pre-committed batch) --- run_prove() { # $1=binary -> echoes proving time (s) local out t - out="$("$1" prove "$ELF" --private-input "$INPUT" -o "$PROOF" --time 2>&1)" + # shellcheck disable=SC2086 # CONT_ARGS is intentionally word-split (0 or 2 args) + out="$("$1" prove "$ELF" --private-input "$INPUT" -o "$PROOF" --time $CONT_ARGS 2>&1)" rm -f "$PROOF" t="$(printf '%s\n' "$out" | grep -o 'Proving time: [0-9.]*' | awk '{print $3}')" if [ -z "$t" ]; then From 88ac142d3201172381f595f97b7ec652b8db88d9 Mon Sep 17 00:00:00 2001 From: MauroFab Date: Thu, 9 Jul 2026 00:03:05 -0300 Subject: [PATCH 7/9] =?UTF-8?q?fix(stark,math-cuda):=20review=20fixes=20?= =?UTF-8?q?=E2=80=94=20harden=20device=5Fonly=20gate=20+=20inference=20+?= =?UTF-8?q?=20gather=20fallback?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - device_only_for: require has_aux_trace() + a non-empty constraint set (R2 composition needs a device aux handle and a uniform zerofier with >=1 group). - build_round1: recover num_rows from the device handle whenever host_trace_empty, covering the aux-only (num_main_cols==0) case a main_empty-only guard missed. - open_deep_composition_poly: a device row-gather Err falls back to the host gather unless the trace is device-only (was .expect() → panic on a transient error). - merkle_root_parity: pass the retain_host_lde arg (test target no longer compiled). --- crypto/math-cuda/tests/merkle_root_parity.rs | 2 + crypto/stark/src/prover.rs | 71 ++++++++++++++++---- 2 files changed, 59 insertions(+), 14 deletions(-) diff --git a/crypto/math-cuda/tests/merkle_root_parity.rs b/crypto/math-cuda/tests/merkle_root_parity.rs index fcc9d226e..208353d95 100644 --- a/crypto/math-cuda/tests/merkle_root_parity.rs +++ b/crypto/math-cuda/tests/merkle_root_parity.rs @@ -305,6 +305,7 @@ fn new_row_major_pipeline_base_root_matches_cpu() { num_cols, blowup, &weights_u64, + true, ) .expect("new row-major GPU pipeline"); let gpu_root = handle.tree.as_ref().expect("resident merkle tree").root; @@ -368,6 +369,7 @@ fn new_row_major_pipeline_ext3_root_matches_cpu() { num_cols, blowup, &weights_u64, + true, ) .expect("new ext3 row-major GPU pipeline"); let gpu_root = handle.tree.as_ref().expect("resident merkle tree").root; diff --git a/crypto/stark/src/prover.rs b/crypto/stark/src/prover.rs index 352f5acd4..27ff172d2 100644 --- a/crypto/stark/src/prover.rs +++ b/crypto/stark/src/prover.rs @@ -260,11 +260,9 @@ where // empty host buffer where it should have data. Using the real state is a // 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. `main_empty` additionally - // drives the `num_rows` recovery, since `from_row_major` reads it from - // `main_data.len()`. 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. + // 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. #[cfg(feature = "cuda")] let main_empty = num_main_cols > 0 && main_data.is_empty(); #[cfg(feature = "cuda")] @@ -289,9 +287,15 @@ where #[cfg(feature = "cuda")] { if host_trace_empty { - // Recover the LDE row count from the resident handle when the - // (empty) main buffer could not supply it. - if main_empty && let Some(n) = device_num_rows { + // Recover the LDE row count from the resident device handle + // whenever any host buffer is empty. `from_row_major` derives + // `num_rows` from `main_data` (or `aux_data` when there are no + // main columns); if that buffer was skipped it reads 0, so we + // overwrite from the handle's `lde_size` (the true row count). + // Idempotent when `from_row_major` already got it right, and it + // covers the aux-only (`num_main_cols == 0`) device-only case + // that a `main_empty`-only guard missed. + if let Some(n) = device_num_rows { lde_trace.set_num_rows(n); } lde_trace.set_host_trace_empty(true); @@ -803,6 +807,17 @@ pub trait IsStarkProver< air: &dyn AIR, domain: &Domain, ) -> bool { + // Preconditions the downstream GPU paths require that the numeric gate + // below does not capture. A table missing either would pass the gate, + // skip its host D2H, then hard-abort in round 2: + // - 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). + 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; let offsets = &air.context().transition_offsets; @@ -2022,9 +2037,24 @@ pub trait IsStarkProver< let stream = lde_trace .bound_stream() .expect("bound stream for device-resident main-row gather"); - let raw = - math_cuda::barycentric::gather_rows_base_on_device(h, &query_rows, &stream) - .expect("device main-row gather failed"); + let raw = match math_cuda::barycentric::gather_rows_base_on_device( + h, + &query_rows, + &stream, + ) { + Ok(v) => v, + // A gather failure is only fatal under device-only (no host + // trace to fall back to). Otherwise return `None` so the host + // gather arms below serve the openings from the resident LDE. + Err(e) => { + assert!( + !lde_trace.host_trace_empty(), + "device main-row gather failed and the trace is device-only \ + (no host fallback): {e:?}" + ); + return None; + } + }; crate::constraint_ir::gpu_interp::base_u64_to_field::(&raw) }) }); @@ -2041,9 +2071,22 @@ pub trait IsStarkProver< let stream = lde_trace .bound_stream() .expect("bound stream for device-resident aux-row gather"); - let raw = - math_cuda::barycentric::gather_rows_ext3_on_device(h, &query_rows, &stream) - .expect("device aux-row gather failed"); + let raw = match math_cuda::barycentric::gather_rows_ext3_on_device( + h, + &query_rows, + &stream, + ) { + Ok(v) => v, + // Fatal only under device-only; otherwise fall back to host. + Err(e) => { + assert!( + !lde_trace.host_trace_empty(), + "device aux-row gather failed and the trace is device-only \ + (no host fallback): {e:?}" + ); + return None; + } + }; crate::constraint_ir::gpu_interp::ext3_u64_to_field::(&raw) }) }); From 79a2099bc14ba237a2df1459d37b12fc19917846 Mon Sep 17 00:00:00 2001 From: MauroFab Date: Mon, 13 Jul 2026 12:30:40 -0300 Subject: [PATCH 8/9] =?UTF-8?q?fix(stark):=20review=20fixes=20=E2=80=94=20?= =?UTF-8?q?restore=20commit=5Fmain=5Ftrace=20doc=20attachment=20+=20presen?= =?UTF-8?q?t-tense=20host=5Ftrace=5Fempty=20doc?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- crypto/stark/src/prover.rs | 18 +++++++++--------- crypto/stark/src/trace.rs | 3 ++- 2 files changed, 11 insertions(+), 10 deletions(-) diff --git a/crypto/stark/src/prover.rs b/crypto/stark/src/prover.rs index 27ff172d2..45fe0036d 100644 --- a/crypto/stark/src/prover.rs +++ b/crypto/stark/src/prover.rs @@ -787,15 +787,6 @@ pub trait IsStarkProver< }); } - /// Compute the main-trace LDE and commit. Returns a `TableCommit` along - /// with the owned LDE columns (consumed later in Phase D) and (under - /// cuda) the optional device LDE buffer kept alive for downstream rounds - /// when the R1 fused GPU pipeline ran. - /// - /// `precomputed`: if present, the leading `num_cols` columns are committed - /// as a separate Merkle tree (the precomputed split for preprocessed - /// tables) and the root is checked against the AIR-hardcoded commitment. - #[allow(clippy::type_complexity)] /// 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 @@ -832,6 +823,15 @@ pub trait IsStarkProver< ) } + /// Compute the main-trace LDE and commit. Returns a `TableCommit` along + /// with the owned LDE columns (consumed later in Phase D) and (under + /// cuda) the optional device LDE buffer kept alive for downstream rounds + /// when the R1 fused GPU pipeline ran. + /// + /// `precomputed`: if present, the leading `num_cols` columns are committed + /// as a separate Merkle tree (the precomputed split for preprocessed + /// tables) and the root is checked against the AIR-hardcoded commitment. + #[allow(clippy::type_complexity)] fn commit_main_trace( trace: &TraceTable, domain: &Domain, diff --git a/crypto/stark/src/trace.rs b/crypto/stark/src/trace.rs index 280aabb95..b34023ac3 100644 --- a/crypto/stark/src/trace.rs +++ b/crypto/stark/src/trace.rs @@ -332,7 +332,8 @@ where /// 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. Always false today until the `device_only` gate is wired. + /// proof. Set by `build_round1` when the device-only gate kept this table's + /// round-1 LDE on the GPU. #[cfg(feature = "cuda")] pub(crate) host_trace_empty: bool, /// Per table GPU residency session: owns this table's device LDE buffers From 1f94122dd81ae7293d03632242ffbc9913022f3c Mon Sep 17 00:00:00 2001 From: MauroFab Date: Mon, 13 Jul 2026 13:09:39 -0300 Subject: [PATCH 9/9] refactor(stark): one body for the R4 main/aux device openings + shared gate predicates The R4 main and aux arms were twins: the query-row gather closure (x2), the device-vs-host cross-check, and three host-fallback hard-abort guards each existed in a main and an aux copy that had to be edited in lockstep. gather_query_rows_device + open_trace_polys_device hold the one body; the call sites differ only in handle/tree/gather-fn. device_only_gate and device_only_for now call the shared is_goldilocks_ext3_tower / offsets_are_contiguous predicates, and the gate doc states the lockstep requirement against the R2 dispatch checks. --- crypto/stark/src/gpu_lde.rs | 8 +- crypto/stark/src/prover.rs | 292 +++++++++++++++++++----------------- 2 files changed, 162 insertions(+), 138 deletions(-) diff --git a/crypto/stark/src/gpu_lde.rs b/crypto/stark/src/gpu_lde.rs index 9c21717d2..f962dc272 100644 --- a/crypto/stark/src/gpu_lde.rs +++ b/crypto/stark/src/gpu_lde.rs @@ -179,6 +179,11 @@ pub(crate) fn device_only_disabled() -> bool { /// `zerofier_uniform` must be the R1-derived conservative form (all constraints /// share `end_exemptions == 0`), which implies `ZerofierEvaluations::is_uniform` /// (a single cyclic group) — the condition the GPU composition kernel needs. +/// +/// 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. pub(crate) fn device_only_gate( lde_size: usize, n: usize, @@ -194,8 +199,7 @@ where if cfg!(feature = "debug-checks") { return false; } - TypeId::of::() == TypeId::of::() - && TypeId::of::() == TypeId::of::() + is_goldilocks_ext3_tower::() && !device_only_disabled() && !gpu_composition_disabled() && lde_size.is_power_of_two() diff --git a/crypto/stark/src/prover.rs b/crypto/stark/src/prover.rs index 45fe0036d..7fa560e5c 100644 --- a/crypto/stark/src/prover.rs +++ b/crypto/stark/src/prover.rs @@ -811,8 +811,8 @@ pub trait IsStarkProver< } let lde_size = domain.interpolation_domain_size * domain.blowup_factor; let n = domain.interpolation_domain_size; - let offsets = &air.context().transition_offsets; - let offsets_contiguous = offsets.iter().enumerate().all(|(i, &o)| o == i); + 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); crate::gpu_lde::device_only_gate::( lde_size, @@ -1929,6 +1929,108 @@ pub trait IsStarkProver< (even, odd) } + /// Gather every query's row-pair off a device-resident LDE (a small D2H of + /// only the queried rows), lifting the raw limbs to field elements via + /// `convert`. Returns `None` (→ the host arms of the openings) when the + /// gather fails or the tower is not Goldilocks; a gather failure is fatal + /// only under device-only, where no host copy exists to fall back to. Bumps + /// the opening-gather counter exactly when values are produced, so the + /// counter reflects the device path actually serving the openings. One body + /// for the main and aux arms — `what` only labels the messages. + #[cfg(feature = "cuda")] + fn gather_query_rows_device( + lde_trace: &LDETraceTable, + what: &str, + gather: impl FnOnce(&std::sync::Arc) -> math_cuda::Result>, + convert: impl FnOnce(&[u64]) -> Option>>, + ) -> Option>> { + let stream = lde_trace + .bound_stream() + .expect("bound stream for device-resident row gather"); + let raw = match gather(&stream) { + Ok(v) => v, + Err(e) => { + assert!( + !lde_trace.host_trace_empty(), + "device {what}-row gather failed and the trace is device-only \ + (no host fallback): {e:?}" + ); + return None; + } + }; + let vals = convert(&raw); + if vals.is_some() { + crate::gpu_lde::GPU_OPENING_GATHER_CALLS + .fetch_add(1, std::sync::atomic::Ordering::Relaxed); + } + vals + } + + /// One query's trace-poly opening with the device-resident fast paths: + /// device Merkle proof + device-gathered values when both are present, the + /// device proof with a host gather when only the tree is resident, and the + /// full host walk otherwise. One body for the main and aux arms, so the + /// device↔host cross-check and the R4 `host_trace_empty` hard-abort guards + /// exist exactly once. + #[cfg(feature = "cuda")] + #[allow(clippy::too_many_arguments)] + fn open_trace_polys_device( + domain: &Domain, + lde_trace: &LDETraceTable, + dev_proofs: Option<&Vec>>, + dev_values: Option<&Vec>>, + tree: &BatchedMerkleTree, + qi: usize, + challenge: usize, + ncols: usize, + what: &str, + gather: G, + ) -> PolynomialOpenings + where + C: IsField, + FieldElement: AsBytes + Sync + Send, + G: Fn(usize) -> Vec>, + { + let Some(proofs) = dev_proofs else { + assert!( + !lde_trace.host_trace_empty(), + "R4 {what} opening fell back to the host tree, but it is device-only (empty)" + ); + return Self::open_polys_with(domain, tree, challenge, gather); + }; + let proof = proofs[qi].clone(); + let Some(dev_vals) = dev_values else { + // Device tree resident but the value gather is absent: reading the + // host gather is invalid under device-only. + assert!( + !lde_trace.host_trace_empty(), + "R4 {what} opening fell back to the host gather, but it is device-only (empty)" + ); + return Self::open_polys_with_proofs(domain, proof, challenge, gather); + }; + let (even, odd) = Self::device_row_pair(dev_vals, qi, ncols); + // Cross-check the device gather against the host LDE. Skipped under + // device-only (host trace empty): the gather was proven bit-identical + // while the host copy was resident, and there is nothing to check + // against. + if !lde_trace.host_trace_empty() { + let domain_size = domain.lde_roots_of_unity_coset.len() as u64; + let r_even = reverse_index(challenge * 2, domain_size); + let r_odd = reverse_index(challenge * 2 + 1, domain_size); + assert_eq!( + even, + gather(r_even), + "device {what}-row gather mismatch (even), query {qi}" + ); + assert_eq!( + odd, + gather(r_odd), + "device {what}-row gather mismatch (odd), query {qi}" + ); + } + Self::open_polys_from_values(proof, even, odd) + } + /// Open the deep composition polynomial on a list of indexes and their symmetric elements. fn open_deep_composition_poly( domain: &Domain, @@ -2034,67 +2136,43 @@ pub trait IsStarkProver< let main_dev_values: Option>> = main_dev_proofs.as_ref().and_then(|_| { lde_trace.gpu_main().and_then(|h| { - let stream = lde_trace - .bound_stream() - .expect("bound stream for device-resident main-row gather"); - let raw = match math_cuda::barycentric::gather_rows_base_on_device( - h, - &query_rows, - &stream, - ) { - Ok(v) => v, - // A gather failure is only fatal under device-only (no host - // trace to fall back to). Otherwise return `None` so the host - // gather arms below serve the openings from the resident LDE. - Err(e) => { - assert!( - !lde_trace.host_trace_empty(), - "device main-row gather failed and the trace is device-only \ - (no host fallback): {e:?}" - ); - return None; - } - }; - crate::constraint_ir::gpu_interp::base_u64_to_field::(&raw) + Self::gather_query_rows_device( + lde_trace, + "main", + |stream| { + math_cuda::barycentric::gather_rows_base_on_device( + h, + &query_rows, + stream, + ) + }, + |raw| crate::constraint_ir::gpu_interp::base_u64_to_field::(raw), + ) }) }); - #[cfg(feature = "cuda")] - if main_dev_values.is_some() { - crate::gpu_lde::GPU_OPENING_GATHER_CALLS - .fetch_add(1, std::sync::atomic::Ordering::Relaxed); - } #[cfg(feature = "cuda")] let aux_dev_values: Option>> = aux_dev_proofs.as_ref().and_then(|_| { lde_trace.gpu_aux().and_then(|h| { - let stream = lde_trace - .bound_stream() - .expect("bound stream for device-resident aux-row gather"); - let raw = match math_cuda::barycentric::gather_rows_ext3_on_device( - h, - &query_rows, - &stream, - ) { - Ok(v) => v, - // Fatal only under device-only; otherwise fall back to host. - Err(e) => { - assert!( - !lde_trace.host_trace_empty(), - "device aux-row gather failed and the trace is device-only \ - (no host fallback): {e:?}" - ); - return None; - } - }; - crate::constraint_ir::gpu_interp::ext3_u64_to_field::(&raw) + Self::gather_query_rows_device( + lde_trace, + "aux", + |stream| { + math_cuda::barycentric::gather_rows_ext3_on_device( + h, + &query_rows, + stream, + ) + }, + |raw| { + crate::constraint_ir::gpu_interp::ext3_u64_to_field::( + raw, + ) + }, + ) }) }); - #[cfg(feature = "cuda")] - if aux_dev_values.is_some() { - crate::gpu_lde::GPU_OPENING_GATHER_CALLS - .fetch_add(1, std::sync::atomic::Ordering::Relaxed); - } for (qi, index) in indexes_to_open.iter().enumerate() { #[cfg(not(feature = "cuda"))] @@ -2108,49 +2186,18 @@ pub trait IsStarkProver< } else { #[cfg(feature = "cuda")] { - if let Some(proofs) = &main_dev_proofs { - let proof = proofs[qi].clone(); - if let Some(dev_vals) = &main_dev_values { - let (even, odd) = Self::device_row_pair(dev_vals, qi, total_cols); - // Cross-check the device gather against the host LDE. - // Skipped under device-only (host trace empty): the - // gather was proven bit-identical while the host copy - // was resident, and there is nothing to check against. - if !lde_trace.host_trace_empty() { - let r_even = reverse_index(*index * 2, domain_size); - let r_odd = reverse_index(*index * 2 + 1, domain_size); - assert_eq!( - even, - lde_trace.gather_main_row(r_even), - "device main-row gather mismatch (even), query {qi}" - ); - assert_eq!( - odd, - lde_trace.gather_main_row(r_odd), - "device main-row gather mismatch (odd), query {qi}" - ); - } - Self::open_polys_from_values(proof, even, odd) - } else { - // Device tree resident but the value gather is absent: - // reading the host gather is invalid under device-only. - assert!( - !lde_trace.host_trace_empty(), - "R4 main opening fell back to the host gather, but it is device-only (empty)" - ); - Self::open_polys_with_proofs(domain, proof, *index, |row| { - lde_trace.gather_main_row(row) - }) - } - } else { - assert!( - !lde_trace.host_trace_empty(), - "R4 main opening fell back to the host tree, but it is device-only (empty)" - ); - Self::open_polys_with(domain, &main_commit.tree, *index, |row| { - lde_trace.gather_main_row(row) - }) - } + Self::open_trace_polys_device( + domain, + lde_trace, + main_dev_proofs.as_ref(), + main_dev_values.as_ref(), + &main_commit.tree, + qi, + *index, + total_cols, + "main", + |row| lde_trace.gather_main_row(row), + ) } #[cfg(not(feature = "cuda"))] { @@ -2197,45 +2244,18 @@ pub trait IsStarkProver< let aux_trace_polys = round_1_result.aux.as_ref().map(|aux| { #[cfg(feature = "cuda")] { - if let Some(proofs) = &aux_dev_proofs { - let proof = proofs[qi].clone(); - if let Some(dev_vals) = &aux_dev_values { - let ncols = lde_trace.num_aux_cols(); - let (even, odd) = Self::device_row_pair(dev_vals, qi, ncols); - // Cross-check skipped under device-only (host empty). - if !lde_trace.host_trace_empty() { - let r_even = reverse_index(*index * 2, domain_size); - let r_odd = reverse_index(*index * 2 + 1, domain_size); - assert_eq!( - even, - lde_trace.gather_aux_row(r_even), - "device aux-row gather mismatch (even), query {qi}" - ); - assert_eq!( - odd, - lde_trace.gather_aux_row(r_odd), - "device aux-row gather mismatch (odd), query {qi}" - ); - } - Self::open_polys_from_values(proof, even, odd) - } else { - assert!( - !lde_trace.host_trace_empty(), - "R4 aux opening fell back to the host gather, but it is device-only (empty)" - ); - Self::open_polys_with_proofs(domain, proof, *index, |row| { - lde_trace.gather_aux_row(row) - }) - } - } else { - assert!( - !lde_trace.host_trace_empty(), - "R4 aux opening fell back to the host tree, but it is device-only (empty)" - ); - Self::open_polys_with(domain, &aux.tree, *index, |row| { - lde_trace.gather_aux_row(row) - }) - } + Self::open_trace_polys_device( + domain, + lde_trace, + aux_dev_proofs.as_ref(), + aux_dev_values.as_ref(), + &aux.tree, + qi, + *index, + lde_trace.num_aux_cols(), + "aux", + |row| lde_trace.gather_aux_row(row), + ) } #[cfg(not(feature = "cuda"))] {