diff --git a/.github/workflows/gpu-tests.yml b/.github/workflows/gpu-tests.yml index ddcce0ee3..c1fc18aa6 100644 --- a/.github/workflows/gpu-tests.yml +++ b/.github/workflows/gpu-tests.yml @@ -3,8 +3,9 @@ name: GPU Tests (merge queue) # Run the GPU test suite (which CPU CI can't, since GitHub runners have no GPU) on a rented # Vast.ai RTX 5090 when a PR is in the merge queue, and block the merge if it fails. # Groups (see scripts/gpu_test.sh): math-cuda kernel parity, cuda_path_integration (GPU proof -# verifies), cuda_fallback (CPU fallback verifies), the prover/stark/crypto/ecsm suite on the -# GPU path, and the comprehensive all-instructions prove. Orchestration runs on a GitHub-hosted +# verifies), cuda_d1_path (the num_parts==1 device DEEP/FRI path), cuda_fallback (CPU fallback +# verifies), the prover/stark/crypto/ecsm suite on the GPU path, and the comprehensive +# all-instructions prove. Orchestration runs on a GitHub-hosted # runner; all GPU work happens on the rented box (provisioned by the template onstart). The box # is ALWAYS destroyed at the end. # @@ -55,7 +56,7 @@ jobs: # Skip on PRs (reports as Skipped = required check satisfied, no GPU rental); run for # real on merge_group and manual dispatch. if: github.event_name != 'pull_request' - # Provisioning + cuda builds + 5 test groups; the prover suite (single-threaded, real + # Provisioning + cuda builds + 6 test groups; the prover suite (single-threaded, real # ELF proves) dominates. Generous ceiling; teardown still always destroys the box. timeout-minutes: 240 steps: diff --git a/Makefile b/Makefile index db12d3722..3e5e3901e 100644 --- a/Makefile +++ b/Makefile @@ -574,7 +574,7 @@ test-disk-spill: GPU_TEST_TIMEOUT := timeout -k 30 2700 # math-cuda kernel tests (requires NVIDIA GPU + nvcc). Group 1 of gpu_test.sh, -# so a hang here also costs Groups 2-5: they run after it, sequentially. +# so a hang here also costs Groups 2-6: they run after it, sequentially. test-math-cuda: $(GPU_TEST_TIMEOUT) cargo test -p math-cuda --release diff --git a/crypto/stark/src/gpu_lde.rs b/crypto/stark/src/gpu_lde.rs index 4de63f35d..8782c6923 100644 --- a/crypto/stark/src/gpu_lde.rs +++ b/crypto/stark/src/gpu_lde.rs @@ -692,10 +692,34 @@ where Some((lde_h0, lde_h1)) } +/// Shared admission gate for the device composition-parts producers: the tower +/// must be the Goldilocks/ext3 pair the kernels are written for, and the LDE must +/// be a power of two at or above the commit threshold. Returns the validated LDE +/// size so callers can derive from it. Kept in one place so a future condition +/// (a VRAM check, a tower widening) cannot land on only one of the d=1/d=2 arms. +fn dev_comp_parts_gate(num_rows: usize) -> Option +where + F: IsField + 'static, + E: IsField + 'static, +{ + if TypeId::of::() != TypeId::of::() { + return None; + } + if TypeId::of::() != TypeId::of::() { + return None; + } + if num_rows < gpu_lde_threshold() || !num_rows.is_power_of_two() { + return None; + } + Some(num_rows) +} + /// Fully device-resident degree-2 decomposition + half extension: takes the /// resident composition evals `H`, decomposes into H0/H1 on device, LDE-extends -/// both and keeps the de-interleaved parts buffer as a `GpuLdeExt3` (commit -/// tree, R3 OOD, R4 DEEP and openings all read the handle). With `want_host` +/// both and keeps the de-interleaved parts buffer as a `GpuLdeExt3` (the commit +/// tree and the R4 openings read `handle.m`; R3 and R4 DEEP read the host part +/// Vec's length and DEEP validates the handle against it — see +/// [`try_comp_h_to_slabs_dev`] for why the two must stay equal). With `want_host` /// the evaluations are also drained to host for the fallback consumers; /// without it (device-only) the returned part Vecs are empty placeholders. /// `None` → the caller downloads `H` and runs the host decompose path. @@ -709,16 +733,7 @@ where F: IsField + 'static, E: IsField + 'static, { - if TypeId::of::() != TypeId::of::() { - return None; - } - if TypeId::of::() != TypeId::of::() { - return None; - } - let lde_size = h.num_rows; - if lde_size < gpu_lde_threshold() || !lde_size.is_power_of_two() { - return None; - } + let lde_size = dev_comp_parts_gate::(h.num_rows)?; let n = lde_size / 2; if weights.len() != n || inv_2x.len() < n { return None; @@ -782,39 +797,46 @@ where } /// Fully device-resident num_parts==1 composition-parts path: `H` itself is the -/// single part, already on the LDE coset, so — unlike [`try_decompose_extend_d2_dev`] -/// — there is no decompose and no re-extension, only a de-interleave into the slab -/// layout the commit / DEEP / FRI consumers read (all of which already read the -/// part count from `handle.m`). The single part is always drained to host — not -/// just because it can be (num_parts==1 tables are never device-only; they keep -/// their preprocessed host trace, see `device_only_for`), but because that host -/// part is what feeds the query-0 composition-opening canary: release-active for -/// `qi == 0` and guarded on a non-empty host part, it is the only end-to-end check -/// that the device m=1 gather / DEEP / FRI layout is correct. Returning empty parts -/// (`vec![Vec::new()]`, as the d=2 device-only arm does) would save the D2H and keep -/// num_parts==1 — but silently disable that canary. +/// single part, already on the LDE coset, so — unlike [`try_comp_h_to_slabs_dev`]'s +/// d=2 sibling [`try_decompose_extend_d2_dev`] — there is no decompose and no +/// re-extension, only a de-interleave into the slab layout the downstream consumers +/// read. No consumer needed changing for `m == 1`, but they do not agree on where +/// the part count comes from, and the difference matters to anyone editing this: +/// +/// - R2 commit and the R4 openings read `handle.m`. +/// - R3's `z^P` exponent and R4 DEEP's gamma count read +/// `lde_composition_poly_evaluations.len()` — the HOST part Vec's length. DEEP only +/// *validates* the handle against it and declines on a mismatch. +/// - FRI never sees the handle at all; it consumes the DEEP codeword. +/// +/// So the invariant to preserve is `handle.m == lde_composition_poly_evaluations.len()` +/// (`materialize_composition_parts_host` also requires it), not "the handle is +/// authoritative". +/// +/// The single part is always drained to host — not just because it can be +/// (num_parts==1 tables are never device-only; `device_only_for`'s degree gate admits +/// only d=2), but because that host part is what feeds the query-0 +/// composition-opening canary: release-active for `qi == 0` and guarded on a +/// non-empty host part, it is the only *in-prove* check that the device m=1 gather is +/// correct. It does not cover DEEP or FRI, which consume separate downstream buffers; +/// those are covered by proof verification (`prover/tests/cuda_d1_path.rs`). Returning +/// empty parts (`vec![Vec::new()]`, as the d=2 device-only arm does) would save the +/// D2H and keep num_parts==1 — but silently disable that canary. /// `None` → the caller downloads `H` and uses it directly as the single host part. -pub(crate) fn try_deinterleave_comp_h_dev( +pub(crate) fn try_comp_h_to_slabs_dev( h: &math_cuda::constraint_interp::GpuCompH, ) -> Option<(Vec>>, math_cuda::lde::GpuLdeExt3)> where F: IsField + 'static, E: IsField + 'static, { - if TypeId::of::() != TypeId::of::() { - return None; - } - if TypeId::of::() != TypeId::of::() { - return None; - } - let lde_size = h.num_rows; - if lde_size < gpu_lde_threshold() || !lde_size.is_power_of_two() { - return None; - } + dev_comp_parts_gate::(h.num_rows)?; // The interleaved `H` download IS the single composition part on the LDE - // coset — same values the slab handle holds, just interleaved. Download - // before building the handle so a decline leaves nothing to unwind. + // coset — same values the slab handle holds, just interleaved. Downloading + // first keeps the blocking D2H off the tail of the de-interleave launch; a + // later handle failure just re-drains in the caller's fallback (both values + // drop by RAII on any early return, in either order). let host = vec![download_comp_h_to_field::(h)?]; let handle = math_cuda::constraint_interp::comp_h_to_slabs(h).ok()?; diff --git a/crypto/stark/src/prover.rs b/crypto/stark/src/prover.rs index 19017aa1c..5078ce290 100644 --- a/crypto/stark/src/prover.rs +++ b/crypto/stark/src/prover.rs @@ -1091,18 +1091,20 @@ pub trait IsStarkProver< // empty constraint set makes `all(end_exemptions == 0)` vacuously // true here but `is_uniform()` false downstream (0 groups). // - Device-only is entered only for the d=2 quotient decomposition, - // checked below once `n` is in hand. The d=1 device R2 path exists - // too but keeps its host trace (never device-only), so it is not a - // concern here. + // checked below once `n` is in hand. A d=1 table also has a device R2 + // path, but the gate below excludes it, so it stays device-additive. if !air.has_aux_trace() || air.constraints_meta().is_empty() { return false; } let n = domain.interpolation_domain_size; - // Device-only is entered only for the d=2 quotient decomposition. The d=1 - // (num_parts==1) device R2 path exists too, but keeps its host trace - // (`want_host` stays true) so the preprocessed main commit still sees real - // data — zeroing a preprocessed table's host trace fails its commitment - // check. So d=1 tables run device-additive, never device-only. + // Only the d=2 quotient decomposition has a device-resident R2 path that + // can serve every downstream consumer from the handle alone. A d=1 table + // does have a device R2 path, but it always drains its single part to host + // (the query-0 composition canary reads it), so it gains nothing from + // dropping the host trace and this gate keeps it device-additive. Any other + // part count has no device R2 path at all and needs the host evaluator, + // which device-only would leave without data until the R2 downgrade + // recovered it. if air.composition_poly_degree_bound(n) / n != 2 { return false; } @@ -1544,11 +1546,13 @@ pub trait IsStarkProver< /// Decompose the resident composition `H` into device-resident parts per the /// AIR's part count: the trivial d=1 de-interleave (`H` is the single part on /// the LDE coset) or the d=2 quotient split H₀/H₁. Both keep the parts - /// device-resident (commit / R3 OOD / R4 DEEP / openings all read the count - /// from the handle); `None` → the caller falls back to the host path. Shared - /// by the R2 producer and the `xcheck` mirror so the two cannot drift. - /// `want_host` gates the d=2 host drain only — d=1 tables are never - /// device-only, so they always keep their host part. + /// device-resident — the commit tree and the R4 openings read `handle.m`, while + /// R3 and R4 DEEP read the host part Vec's length (see + /// [`crate::gpu_lde::try_comp_h_to_slabs_dev`] for the invariant that ties the + /// two together). `None` → the caller falls back to the host path. Shared by the + /// R2 producer and the `xcheck` mirror so the two cannot drift. `want_host` gates + /// the d=2 host drain only — d=1 tables are never device-only, so they always + /// keep their host part. #[cfg(feature = "cuda")] fn decompose_comp_h_dev( number_of_parts: usize, @@ -1561,9 +1565,9 @@ pub trait IsStarkProver< math_cuda::lde::GpuLdeExt3, )> { if number_of_parts == 1 { - // d=1 is never device-only (the degree gate below / `device_only_for` - // only admit d=2), so the single part is always kept on host — - // `want_host` must hold, and the d=1 helper ignores it by design. + // d=1 is never device-only (`device_only_for`'s degree gate admits only + // d=2), so the single part is always kept on host — `want_host` must + // hold, and the d=1 helper ignores it by design. debug_assert!( want_host, "d=1 composition parts are never device-only; want_host must hold" @@ -1576,7 +1580,7 @@ pub trait IsStarkProver< domain.interpolation_domain_size * domain.blowup_factor, "d=1 H row count must equal the LDE domain size" ); - crate::gpu_lde::try_deinterleave_comp_h_dev::(h_dev) + crate::gpu_lde::try_comp_h_to_slabs_dev::(h_dev) } else { crate::gpu_lde::try_decompose_extend_d2_dev::( h_dev, @@ -1743,14 +1747,13 @@ pub trait IsStarkProver< let want_host = !round_1_result.lde_trace.host_trace_empty(); // num_parts==1 de-interleaves `H` (the single part); num_parts==2 // runs the degree-2 quotient split. Both keep the parts resident. - let decomposed = Self::decompose_comp_h_dev( + match Self::decompose_comp_h_dev( number_of_parts, &h_dev, domain, twiddles, want_host, - ); - match decomposed { + ) { Some((parts, handle)) => { gpu_composition_parts = Some(handle); precomputed_parts = Some(parts); diff --git a/crypto/stark/tests/gpu_constraint_interp.rs b/crypto/stark/tests/gpu_constraint_interp.rs index d7adc01d2..eef21953a 100644 --- a/crypto/stark/tests/gpu_constraint_interp.rs +++ b/crypto/stark/tests/gpu_constraint_interp.rs @@ -568,13 +568,6 @@ fn check_composition(prog: &ConstraintProgram, label: &str, seed: u64) dev_raw, gpu, "[{label}] evaluate_dev (keep=true) H != host-drained H, seed {seed:#x}" ); - - // This closes the num_parts==1 device DEEP/FRI path (`gpu_comp_h_slabs_calls`) - // at the unit level. Its end-to-end counterpart is not asserted in - // `prover/tests/cuda_path_integration.rs`: every asm fixture there is a - // `fib_iterative_*` variant whose DECODE ROM is below the GPU LDE threshold, - // so no d1 table engages. The end-to-end d1 path is exercised by real-program - // proves (ethrex) and the GPU bench instead. } #[test] @@ -587,6 +580,12 @@ fn gpu_composition_matches_cpu_oracle_all_ops() { /// num_parts==1 de-risk: the DECODE-shaped (empty-base, LogUp-only) program must /// evaluate on the GPU composition kernel, match the CPU oracle, and produce a /// device-resident `H` bit-identical to the host-drained one. +/// +/// This closes the num_parts==1 device path at the unit level — the `H` the slab +/// de-interleave consumes. The end-to-end counterpart (de-interleave -> commit -> +/// OOD -> DEEP -> FRI -> openings, then verify) is `prover/tests/cuda_d1_path.rs`, +/// which needs a lowered `LAMBDA_VM_GPU_LDE_THRESHOLD` because no fixture crosses +/// the default for a d=1 table; `make test-cuda-d1` runs it. #[test] fn gpu_composition_matches_cpu_oracle_decode_shaped() { for seed in [0x0123_4567_89AB_CDEF, 0xDEAD_BEEF_CAFE_F00D, 7] { diff --git a/prover/tests/cuda_path_integration.rs b/prover/tests/cuda_path_integration.rs index b8e540a3b..dd841d7b7 100644 --- a/prover/tests/cuda_path_integration.rs +++ b/prover/tests/cuda_path_integration.rs @@ -76,11 +76,14 @@ fn gpu_path_fires_end_to_end() { // path. assert!(gpu_bary_calls() > 0, "R3 GPU barycentric did not fire"); - // R2 GPU composition-poly LDE. Fires via one of two paths depending on the + // R2 GPU composition-poly LDE. Fires via one of three paths depending on the // AIR's `number_of_parts`: the fused two-halves quotient decomposition for // the common degree-2 case (`== 2`, counted by `gpu_extend_halves_calls`), - // or the batched parts LDE for `> 2` (counted by `gpu_parts_lde_calls`). - // fib_iterative_1M only exercises the degree-2 path, so assert on either. + // the batched parts LDE for `> 2` (counted by `gpu_parts_lde_calls`), or the + // d=1 de-interleave (`== 1`, counted by `gpu_comp_h_slabs_calls` — covered + // separately by `cuda_d1_path.rs`, since no d=1 table here crosses the + // default LDE threshold). fib_iterative_1M only exercises the degree-2 path, + // so assert on either of the two counted here. assert!( gpu_extend_halves_calls() + gpu_parts_lde_calls() > 0, "R2 GPU composition LDE did not fire (neither two-halves d2 nor parts>2 path)" diff --git a/scripts/gpu_test.sh b/scripts/gpu_test.sh index 44e603635..7d40f9f67 100755 --- a/scripts/gpu_test.sh +++ b/scripts/gpu_test.sh @@ -42,7 +42,7 @@ nvidia-smi --query-gpu=name,driver_version,compute_cap --format=csv,noheader # --- Build the guest ELFs the tests prove --------------------------------------- # math-cuda parity needs none; cuda_path_integration / cuda_fallback prove an asm ELF; the -# prover suite (Groups 4 & 5) proves asm AND rust guests. Build both up front. +# prover suite (Groups 5 & 6) proves asm AND rust guests. Build both up front. log "compiling guest programs (asm + rust)" make compile-programs-asm make compile-programs-rust