From b077aa1b1283989df528c4a8282366396397fdbe Mon Sep 17 00:00:00 2001 From: MauroFab Date: Fri, 21 Aug 2026 16:52:23 -0300 Subject: [PATCH] fix(gpu): review follow-ups on the GPU grinding PR MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Route the GPU dispatch and its tests through one inner-hash-to-lanes conversion. The tests built their own copy, so the line the prover actually runs was executed by nothing: swapping it to from_be_bytes would have kept every test green while is_valid_nonce rejected every device nonce at runtime and the search sat on the CPU fallback forever. stark::grinding:: inner_hash_lanes is now the single entry point, which also lets get_inner_hash go back to private. Report that fallback on stderr instead of log::warn. The CLI initialises env_logger with no default filter, so a warn-level line never prints unless RUST_LOG is set — and it is the only signal that the kernel has started returning garbage. The other device-decline paths already use eprintln with a [gpu] prefix. Wrap test-math-cuda in GPU_TEST_TIMEOUT. It was the only one of the five GPU targets without it, and it is Group 1 of gpu_test.sh, so a hang there costs Groups 2-5 as well and a job timeout yields `cancelled`, which skips the run-summary step and leaves no readable output. Document LAMBDA_VM_NO_GPU_GRIND in the profiling README's knob list. Drop the "Parity" framing from the test module: there is nothing to be at parity with, since any valid nonce is acceptable and the CPU's find_any does not agree with itself between runs. What is pinned is validity, plus the search completeness that minimality stands in for — noted as a probe rather than a contract, so a future kernel that deliberately returns any valid nonce relaxes the assertion instead of being treated as broken. Same for the doc on generate_nonce_maybe_gpu, which claimed "smallest" for both arms. --- Makefile | 5 ++-- crypto/math-cuda/src/grinding.rs | 3 +- crypto/math-cuda/tests/grinding.rs | 34 +++++++++++++++------- crypto/stark/src/grinding.rs | 46 +++++++++++++++++++----------- scripts/profiling/README.md | 4 +++ 5 files changed, 62 insertions(+), 30 deletions(-) diff --git a/Makefile b/Makefile index c19ea0da0..f11ed8581 100644 --- a/Makefile +++ b/Makefile @@ -573,9 +573,10 @@ test-disk-spill: # timeout's 124 exit fails the target so gpu_test.sh reports the group as failed. GPU_TEST_TIMEOUT := timeout -k 30 2700 -# math-cuda parity tests (requires NVIDIA GPU + nvcc) +# 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. test-math-cuda: - cargo test -p math-cuda --release + $(GPU_TEST_TIMEOUT) cargo test -p math-cuda --release # End-to-end cuda dispatch coverage (requires NVIDIA GPU + nvcc). # Asserts the R1-R4 GPU dispatch counters fired on a real prove. diff --git a/crypto/math-cuda/src/grinding.rs b/crypto/math-cuda/src/grinding.rs index 6c5f539c2..fe7803eb9 100644 --- a/crypto/math-cuda/src/grinding.rs +++ b/crypto/math-cuda/src/grinding.rs @@ -20,7 +20,8 @@ const GRIND_MIN_FACTOR: u8 = 12; /// is unavailable/errors (the caller then runs the CPU search). /// /// `inner_lanes` are the four little-endian-read u64 lanes of the 32-byte -/// `inner_hash` (`get_inner_hash` on the host). `grinding_factor` (1..=64) +/// inner hash — build them with `stark::grinding::inner_hash_lanes`, which is +/// what the prover and the tests here both call. `grinding_factor` (1..=64) /// fixes `limit = 1 << (64 - grinding_factor)` and sizes the search: the /// expected first valid nonce is ~`2^grinding_factor`, so each launch scans a /// contiguous block several times that, from 0 upward, and the first block that diff --git a/crypto/math-cuda/tests/grinding.rs b/crypto/math-cuda/tests/grinding.rs index 76db77cf0..84bc5e624 100644 --- a/crypto/math-cuda/tests/grinding.rs +++ b/crypto/math-cuda/tests/grinding.rs @@ -1,5 +1,10 @@ -//! Parity: the GPU proof-of-work nonce search must agree with the host -//! predicate. Runs on the merge-queue GPU box via `make test-math-cuda` +//! The GPU nonce search must produce nonces the host predicate accepts. There +//! is nothing to compare against the CPU search itself — any nonce satisfying +//! `is_valid_nonce` is as good as any other, and the CPU's `find_any` does not +//! even agree with itself between runs — so what is pinned here is validity, +//! plus the search completeness that minimality stands in for. +//! +//! Runs on the merge-queue GPU box via `make test-math-cuda` //! (`cargo test -p math-cuda --release`) — `device::backend()` inside //! `generate_nonce_gpu` requires a real GPU, like the other tests here. //! @@ -7,21 +12,28 @@ //! suite only exercises `grinding_factor: 1`, where `limit = 1 << 63` lets a //! broken kernel return an accepted nonce ~half the time; these factors make a //! wrong kernel fail deterministically. +//! +//! The lanes come from `stark::grinding::inner_hash_lanes`, the same call the +//! prover makes — building them here instead would leave the production +//! conversion untested. -use stark::grinding::{get_inner_hash, is_valid_nonce}; - -fn lanes_for(seed: &[u8; 32], factor: u8) -> [u64; 4] { - let inner = get_inner_hash(seed, factor); - core::array::from_fn(|i| u64::from_le_bytes(inner[i * 8..i * 8 + 8].try_into().unwrap())) -} +use stark::grinding::{inner_hash_lanes, is_valid_nonce}; /// At a moderate factor the kernel returns a valid nonce, and it is the /// smallest one (the exhaustive CPU scan below it is cheap at factor 14). +/// +/// Minimality is not a contract — any valid nonce would do — but it is a cheap +/// probe of search completeness: a stride or bounds bug that skipped part of +/// the range would still return a *valid* nonce, just not the first one, and +/// plain validity checking would miss that. Deterministic despite the grid +/// being parallel, because `atomicMin` is an order-independent reduction. If a +/// future kernel drops minimality deliberately, relax this to validity rather +/// than treating the red as a defect. #[test] fn gpu_grind_returns_smallest_valid_nonce() { let seed = [14u8; 32]; let factor = 14u8; - let nonce = math_cuda::grinding::generate_nonce_gpu(&lanes_for(&seed, factor), factor) + let nonce = math_cuda::grinding::generate_nonce_gpu(&inner_hash_lanes(&seed, factor), factor) .expect("GPU grind (needs a GPU)"); assert!( is_valid_nonce(&seed, nonce, factor), @@ -39,7 +51,7 @@ fn gpu_grind_returns_smallest_valid_nonce() { fn gpu_grind_valid_at_production_factor() { let seed = [20u8; 32]; let factor = 20u8; - let nonce = math_cuda::grinding::generate_nonce_gpu(&lanes_for(&seed, factor), factor) + let nonce = math_cuda::grinding::generate_nonce_gpu(&inner_hash_lanes(&seed, factor), factor) .expect("GPU grind (needs a GPU)"); assert!( is_valid_nonce(&seed, nonce, factor), @@ -53,7 +65,7 @@ fn gpu_grind_valid_at_production_factor() { fn gpu_grind_declines_below_min_factor() { let seed = [1u8; 32]; assert!( - math_cuda::grinding::generate_nonce_gpu(&lanes_for(&seed, 1), 1).is_none(), + math_cuda::grinding::generate_nonce_gpu(&inner_hash_lanes(&seed, 1), 1).is_none(), "GPU grind should decline factor 1" ); } diff --git a/crypto/stark/src/grinding.rs b/crypto/stark/src/grinding.rs index b04fda912..adb7601b6 100644 --- a/crypto/stark/src/grinding.rs +++ b/crypto/stark/src/grinding.rs @@ -78,10 +78,7 @@ fn is_valid_nonce_for_inner_hash(inner_hash: &[u8; 32], candidate_nonce: u64, li /// Returns the bit-string constructed as /// Hash(prefix || seed || grinding_factor) /// `prefix` is the bit-string `0x123456789abcded` -/// -/// Public so the GPU parity test can build the same inner-hash lanes the -/// device kernel searches over. -pub fn get_inner_hash(seed: &[u8; 32], grinding_factor: u8) -> [u8; 32] { +fn get_inner_hash(seed: &[u8; 32], grinding_factor: u8) -> [u8; 32] { let mut inner_data = [0u8; 41]; inner_data[0..8].copy_from_slice(&PREFIX); inner_data[8..40].copy_from_slice(seed); @@ -91,13 +88,27 @@ pub fn get_inner_hash(seed: &[u8; 32], grinding_factor: u8) -> [u8; 32] { digest[..32].try_into().unwrap() } +/// The inner hash as the four little-endian u64 lanes Keccak absorbs it into — +/// the form the device nonce search takes as input. +/// +/// The GPU dispatch and its test both go through here rather than each doing +/// their own byte-to-lane conversion: a second copy would let this one drift +/// (`from_le_bytes` → `from_be_bytes` reads identically at a glance) with every +/// test still green, while at runtime `is_valid_nonce` rejected every device +/// nonce and the search silently sat on the CPU fallback forever. +pub fn inner_hash_lanes(seed: &[u8; 32], grinding_factor: u8) -> [u64; 4] { + let inner_hash = get_inner_hash(seed, grinding_factor); + core::array::from_fn(|i| u64::from_le_bytes(inner_hash[i * 8..i * 8 + 8].try_into().unwrap())) +} + /// Grind on the GPU when a CUDA backend is up, falling back to the CPU search -/// otherwise (or on any device error). The nonce is the smallest valid one in -/// the searched range, which — like the CPU's — the verifier accepts by -/// checking `is_valid_nonce`; nothing downstream depends on which valid nonce -/// is chosen. The heavy per-table-per-epoch ~2^grinding_factor hashing is the -/// prover's dominant CPU cost, so this moves it off the 16 cores onto the idle -/// GPU. +/// otherwise (or on any device error). Which valid nonce comes back depends on +/// the arm: the device search returns the smallest in the range it scanned, +/// while the CPU's `find_any` returns an arbitrary one. Neither is a contract — +/// the verifier accepts any nonce passing `is_valid_nonce`, and nothing +/// downstream depends on the choice. The heavy per-table-per-epoch +/// ~2^grinding_factor hashing is the prover's dominant CPU cost, so this moves +/// it off the 16 cores onto the idle GPU. #[cfg(feature = "cuda")] pub fn generate_nonce_maybe_gpu(seed: &[u8; 32], grinding_factor: u8) -> Option { debug_assert!( @@ -111,11 +122,7 @@ pub fn generate_nonce_maybe_gpu(seed: &[u8; 32], grinding_factor: u8) -> Option< if *GPU_DISABLED.get_or_init(|| std::env::var_os("LAMBDA_VM_NO_GPU_GRIND").is_some()) { return generate_nonce(seed, grinding_factor); } - let inner_hash = get_inner_hash(seed, grinding_factor); - // Keccak reads the 32-byte inner hash as four little-endian lanes. - let inner_lanes: [u64; 4] = core::array::from_fn(|i| { - u64::from_le_bytes(inner_hash[i * 8..i * 8 + 8].try_into().unwrap()) - }); + let inner_lanes = inner_hash_lanes(seed, grinding_factor); if let Some(nonce) = math_cuda::grinding::generate_nonce_gpu(&inner_lanes, grinding_factor) { // Validate unconditionally (one host hash against the ~2^grinding_factor // device search): a kernel/driver defect must degrade to the CPU search, @@ -125,7 +132,14 @@ pub fn generate_nonce_maybe_gpu(seed: &[u8; 32], grinding_factor: u8) -> Option< crate::gpu_lde::GPU_GRIND_CALLS.fetch_add(1, std::sync::atomic::Ordering::Relaxed); return Some(nonce); } - log::warn!("GPU grind returned an invalid nonce ({nonce}); falling back to CPU search"); + // eprintln, not log::warn: the CLI initialises env_logger with no + // default filter, so a warn-level line is invisible unless RUST_LOG is + // set — and this is the only signal that the kernel has started + // returning garbage and the feature has silently reverted to the CPU + // search. Matches the `[gpu]` prefix the other device-decline paths use. + eprintln!( + "[gpu] grind returned an invalid nonce ({nonce}); falling back to the CPU search" + ); } generate_nonce(seed, grinding_factor) } diff --git a/scripts/profiling/README.md b/scripts/profiling/README.md index f4ad4d57b..bad7962ef 100644 --- a/scripts/profiling/README.md +++ b/scripts/profiling/README.md @@ -122,6 +122,10 @@ Useful prover knobs for A/B experiments (pre-existing, see plan §11): `LAMBDA_VM_GPU_BARY_THRESHOLD`, `LAMBDA_VM_VRAM_BUDGET_MB`, `TABLE_PARALLELISM`. +| var | effect | +|---|---| +| `LAMBDA_VM_NO_GPU_GRIND=1` | force the round-4 proof-of-work nonce search onto the CPU (presence-based, like `LAMBDA_VM_NO_GPU_LOGUP`). The production escape hatch if the device search ever misbehaves; also the way to A/B the grind on its own. Below grinding factor 12 the GPU path declines regardless, so wrap and recursion proves (factor 1) never use it | + ## Continuations: per-epoch data for parallelization `prove_continuation` is instrumented independently of the monolithic path