From 7eb089cf6c1c305d4fefa3764f40c469ddc6b9a0 Mon Sep 17 00:00:00 2001 From: MauroFab Date: Mon, 10 Aug 2026 15:04:32 -0300 Subject: [PATCH] docs(gpu): align device-only and downgrade docs with the recovery semantics MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This branch turned two of the device-only hard-aborts into downloads that recover and continue host-backed, but the surrounding docs still describe the old contract: "every host read hard-aborts", "the prove aborts loudly", "a mis-gate panics one of the guards". Rewrite those to say what the code now does — R2 and the R1 resident-aux commit recover and bump GPU_DEVICE_ONLY_DOWNGRADES, R3/R4 still abort, and the R3 guards check the individual buffer so mixed states are legal. Also correct the R2 lock comment (it serializes submission, not execution, for device-only tables), note that the numeric gate is not the complete predicate on its own, broaden the downgrade counter's doc to cover resident-aux declines on tables that were never device-only, and drop the false "only" from materialize_lde_trace_host's failure list. Comments, doc comments, two assertion message strings and one doc-comment run command (--test-threads=1, matching the Makefile target). No behavior changes. --- crypto/stark/src/gpu_lde.rs | 75 +++++++++++++++++++-------- crypto/stark/src/prover.rs | 49 +++++++++++------ crypto/stark/src/trace.rs | 32 ++++++++---- prover/tests/cuda_path_integration.rs | 15 ++++-- 4 files changed, 118 insertions(+), 53 deletions(-) diff --git a/crypto/stark/src/gpu_lde.rs b/crypto/stark/src/gpu_lde.rs index 2a9547153..e1d7865ef 100644 --- a/crypto/stark/src/gpu_lde.rs +++ b/crypto/stark/src/gpu_lde.rs @@ -54,13 +54,24 @@ fn gpu_lde_threshold() -> usize { }) } -/// Serialize the device R2 window (constraint eval + decompose) across -/// tables. Concurrent R2 windows under VRAM pressure can transiently corrupt -/// a whole H buffer (root mechanism unidentified; reruns on the same resident -/// inputs come out correct), yielding a proof that fails verification. -/// Serializing only this window eliminates it at negligible cost — the -/// windows rarely overlap. `LAMBDA_VM_GPU_SERIALIZE_R2=0` disables the lock -/// (e.g. to bisect or once the underlying race is fixed). +/// Serialize the SUBMISSION of the device R2 window (constraint eval + +/// decompose) across tables. Concurrent R2 windows under VRAM pressure can +/// transiently corrupt a whole H buffer (root mechanism unidentified; reruns +/// on the same resident inputs come out correct), yielding a proof that fails +/// verification. Holding this lock empirically suppresses that at negligible +/// cost — the windows rarely overlap. +/// +/// How much it enforces depends on the table. One that keeps its host trace +/// ends the window in a blocking D2H (the `want_host` arm of +/// [`try_decompose_extend_d2_dev`]), so the guard is held until that table's +/// kernels have completed — a real execution barrier. A device-only table's +/// window is enqueue-only, so two tables' R2 kernels can still overlap on +/// device; what the lock orders there is submission and allocation, which is +/// enough to suppress the corruption in practice but is not a guarantee that +/// R2 kernels never run concurrently. +/// +/// `LAMBDA_VM_GPU_SERIALIZE_R2=0` disables the lock (e.g. to bisect or once +/// the underlying race is fixed). pub(crate) fn r2_serialize_guard() -> Option> { static ENABLED: OnceLock = OnceLock::new(); static LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(()); @@ -193,12 +204,23 @@ pub(crate) fn device_only_disabled() -> bool { /// 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. +/// guaranteed to take its GPU path. A strict AND of the numeric and shape +/// preconditions that imply the R2 composition, R3 barycentric, R4 DEEP, and +/// R4 opening GPU paths all fire and read the device LDE — but not the whole +/// predicate on its own: the caller `IsStarkProver::device_only_for` +/// (prover.rs) adds the AIR-level preconditions this signature does not +/// carry, notably the d=2 quotient part count the device-resident R2 path +/// requires. +/// +/// If a precondition is nonetheless violated at runtime (mis-gate or +/// transient GPU error), what happens depends on the round. R2 and the R1 +/// resident-aux commit recover: they download what the host arms need (the +/// resident LDEs at R2, the resident aux trace plus the main LDE at R1), bump +/// [`GPU_DEVICE_ONLY_DOWNGRADES`] and continue host-backed — slower, never +/// wrong — aborting only when the resident handles cannot serve the data. R3 +/// and R4 have no such recovery: the R3 barycentric arms assert on the buffer +/// they are about to read and the R4 guards on `host_trace_empty`, both +/// failing loudly rather than reading an empty host trace. /// /// `zerofier_uniform` must be the R1-derived conservative form (all constraints /// share `end_exemptions == 0`), which implies `ZerofierEvaluations::is_uniform` @@ -206,8 +228,11 @@ pub(crate) fn device_only_disabled() -> bool { /// /// 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. +/// arms). A fallback condition added to a dispatch without a mirror here +/// costs every gate-true table either a hard-abort at R3/R4 — loud, but an +/// avoidable crash — or, at R2 and the R1 resident-aux commit, a silent +/// downgrade to the host path, which is what [`GPU_DEVICE_ONLY_DOWNGRADES`] +/// exists to surface. pub(crate) fn device_only_gate( lde_size: usize, n: usize, @@ -1435,11 +1460,16 @@ pub fn gpu_fri_calls() -> u64 { /// are counted here, so a single failed dispatch does not necessarily lower /// the total; R3's fallbacks are CPU-only, so a failure there does. pub(crate) static GPU_BATCH_INVERT_CALLS: AtomicU64 = AtomicU64::new(0); -/// Times a device-only table had to be downgraded back to a host trace -/// because a downstream device path missed at runtime (see -/// [`materialize_lde_trace_host`]). Nonzero values mean the device-only gate -/// admitted a table some dispatch later declined — correct but slower, and -/// worth mirroring the missing condition into the gate. +/// Times a table had to fall back to a host trace whose data first had to be +/// downloaded off the device, because a device path declined at runtime (see +/// [`materialize_lde_trace_host`] and [`materialize_aux_trace_host`]). +/// Nonzero means a device dispatch declined and the table continued +/// host-backed — correct but slower. Not every one is a gate miss: the R1 +/// resident-aux site is entered whenever `aux_resident()` is set, whatever +/// the device-only gate said, so it also counts declines on tables that were +/// never device-only. Mirroring the missing condition into the gate is the +/// fix for the device-only case; a resident-aux decline is usually transient +/// VRAM pressure instead. pub(crate) static GPU_DEVICE_ONLY_DOWNGRADES: AtomicU64 = AtomicU64::new(0); pub fn gpu_device_only_downgrades() -> u64 { GPU_DEVICE_ONLY_DOWNGRADES.load(Ordering::Relaxed) @@ -1454,8 +1484,9 @@ pub fn gpu_device_only_downgrades() -> u64 { /// reason a dynamic dispatch might decline (kernel eligibility, transient /// errors, shapes a new workload brings), so any miss lands here and degrades /// to a slower-but-correct CPU round instead of a hard abort. Returns false -/// (→ the caller's abort) only when a missing side has no handle or a -/// download fails. +/// (→ the caller's abort) when the resident handles cannot serve the data: a +/// missing handle or bound stream, a handle whose shape disagrees with the +/// trace, a failed download or sync, or a field tower with no CUDA lowering. pub(crate) fn materialize_lde_trace_host( lde_trace: &mut crate::trace::LDETraceTable, ) -> bool diff --git a/crypto/stark/src/prover.rs b/crypto/stark/src/prover.rs index d5cca05e1..20083bb05 100644 --- a/crypto/stark/src/prover.rs +++ b/crypto/stark/src/prover.rs @@ -310,8 +310,15 @@ where // 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. 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. + // empty, the other full) still sets the flag, and is legal rather than + // an error: the aux commit may be more conservative than the main one + // (never less), so an aux side that kept its host copy can sit next to + // a device-only main. The R3 barycentric arms therefore guard on the + // individual buffer — the side that still holds host data stays + // readable — while the flag keeps the R4 and host-evaluator guards + // armed. Reading the real state also picks up an R1 resident-aux + // downgrade: it repopulates the host buffers before this point, so the + // flag simply comes out false. #[cfg(feature = "cuda")] let main_empty = num_main_cols > 0 && main_data.is_empty(); #[cfg(feature = "cuda")] @@ -1021,20 +1028,25 @@ pub trait IsStarkProver< 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: + // below does not capture. A table missing any of them would pass the + // gate and skip its host D2H, leaving round 2 to recover through + // `materialize_lde_trace_host` — correct, but a downgrade, and an + // abort if the resident handles cannot serve the data: // - 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). + // - The device-resident R2 path exists only for the d=2 quotient + // decomposition, checked below once `n` is in hand. if !air.has_aux_trace() || air.constraints_meta().is_empty() { return false; } let n = domain.interpolation_domain_size; // The device-resident R2 path only exists for the d=2 quotient // decomposition; any other part count skips it entirely and needs the - // host evaluator, which device-only leaves without data. + // 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; } @@ -1602,10 +1614,13 @@ pub trait IsStarkProver< #[cfg(feature = "cuda")] if number_of_parts == 2 { // Serializing this window across tables (device constraint eval + - // decompose, where H is born) eliminates a transient whole-buffer - // H corruption seen under concurrent R2 windows on VRAM pressure. - // The commit, the host decompose of a downloaded `H` and every - // host arm run outside the lock. + // decompose, where H is born) empirically eliminates a transient + // whole-buffer H corruption seen under concurrent R2 windows on + // VRAM pressure. What the guard orders is submission: a + // device-only table's window is enqueue-only, so its kernels may + // still overlap another table's on device. The commit, the host + // decompose of a downloaded `H` and every host arm run outside + // the lock. let _r2_serial_guard = crate::gpu_lde::r2_serialize_guard(); if let Some(h_dev) = evaluator.evaluate_dev( air, @@ -1647,16 +1662,16 @@ pub trait IsStarkProver< // Every arm below runs the HOST evaluator, which reads `get_main` / // `get_aux`. Under device-only those buffers are intentionally empty, // so landing here means the device decompose AND the `H` download both - // failed. Abort with the device-only contract's message rather than a - // bare index-out-of-bounds from somewhere inside the evaluator. + // failed. The gate is a static predicate and cannot mirror every + // dynamic decline, so recover rather than abort: download the resident + // LDEs into the host buffers (which also clears the device-only flag) + // and let the host arms run — slower for this table, never wrong. The + // assert is left for the case where the handles themselves cannot + // serve the data, so that failure carries the device-only contract's + // message rather than a bare index-out-of-bounds from somewhere inside + // the evaluator. #[cfg(feature = "cuda")] if precomputed_parts.is_none() && round_1_result.lde_trace.host_trace_empty() { - // The device R2 path missed on a device-only table. The gate is a - // static predicate and cannot mirror every dynamic decline, so - // recover instead of aborting: download the resident LDEs from - // the device handles and continue on the host path — slower for - // this table, never wrong. The abort remains only for the case - // where the handles themselves cannot serve the data. let recovered = crate::gpu_lde::materialize_lde_trace_host(&mut round_1_result.lde_trace); assert!( diff --git a/crypto/stark/src/trace.rs b/crypto/stark/src/trace.rs index c19d2be16..ccf35cca5 100644 --- a/crypto/stark/src/trace.rs +++ b/crypto/stark/src/trace.rs @@ -328,12 +328,18 @@ where 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. Set by `build_round1` when the device-only gate kept this table's - /// round-1 LDE on the GPU. + /// skipped and at least one of `main_data`/`aux_data` is empty — those + /// columns are read off the device instead. Set by `build_round1` when the + /// device-only gate kept this table's round-1 LDE on the GPU, and cleared + /// again by `set_host_data` once a downgrade has downloaded the resident + /// LDEs back into the host buffers. + /// + /// The R4 and host-evaluator guards 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. The R3 barycentric + /// arms instead check the individual buffer they are about to read: mixed + /// states (one side host-backed, the other device-only) are valid, and the + /// populated side stays readable. #[cfg(feature = "cuda")] pub(crate) host_trace_empty: bool, /// Per table GPU residency session: owns this table's device LDE buffers @@ -525,8 +531,11 @@ where } /// 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. + /// device-only path): the round-1 D2H was skipped, so the R4 and + /// host-evaluator reads hard-abort on the flag instead of indexing the + /// empty buffers, while the R3 arms consult the individual buffer. Cleared + /// by [`Self::set_host_data`] once a downgrade has downloaded the resident + /// LDEs back to the host. #[cfg(feature = "cuda")] pub fn set_host_trace_empty(&mut self, empty: bool) { self.host_trace_empty = empty; @@ -563,8 +572,11 @@ where } /// 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`. + /// [`Self::set_host_trace_empty`]). The R4 and host-evaluator fallbacks + /// check this before touching `main_data`/`aux_data`; the R3 barycentric + /// arms check the individual buffer instead, since a mixed state leaves + /// one side readable. False again once a downgrade has repopulated the + /// buffers through [`Self::set_host_data`]. #[cfg(feature = "cuda")] pub fn host_trace_empty(&self) -> bool { self.host_trace_empty diff --git a/prover/tests/cuda_path_integration.rs b/prover/tests/cuda_path_integration.rs index 859f361ec..5b5e52fbc 100644 --- a/prover/tests/cuda_path_integration.rs +++ b/prover/tests/cuda_path_integration.rs @@ -5,7 +5,9 @@ //! regressions (GPU path fired but produced output that fails verification). //! //! `#[ignore]`'d so the no-GPU CI path skips it. Run via `make test-cuda-integration` -//! or `cargo test -p lambda-vm-prover --release --features cuda --test cuda_path_integration -- --ignored --nocapture`. +//! or `cargo test -p lambda-vm-prover --release --features cuda --test cuda_path_integration -- --ignored --nocapture --test-threads=1`. +//! The single test thread is not optional: the counters these tests assert on +//! are process-global, so parallel proves in one process cross-contaminate them. #![cfg(feature = "cuda")] use lambda_vm_prover::test_utils::asm_elf_bytes; @@ -183,7 +185,10 @@ fn gpu_opening_gather_fires_and_verifies() { /// 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. +/// verify). A mis-gate that forces a host fallback shows up one of two ways: +/// at R3/R4 it panics one of the guards, while at R2 and the R1 resident-aux +/// commit it recovers silently and is caught by the downgrade-counter +/// assertion below. #[test] #[ignore = "requires GPU; run with --ignored --nocapture"] fn gpu_device_only_residency_fires_and_verifies() { @@ -197,8 +202,10 @@ fn gpu_device_only_residency_fires_and_verifies() { assert_eq!( stark::gpu_lde::gpu_device_only_downgrades(), 0, - "a device-only table was downgraded back to host on the happy path \ - (a device dispatch declined that the gate should mirror)" + "a table was downgraded back to a host trace on the happy path \ + (a device dispatch declined at runtime: on a device-only table the \ + gate should mirror the missing condition; a resident-aux decline is \ + usually VRAM pressure)" ); assert!( verify(&proof, &elf).expect("verify"),