From 5a895bc6e022ea2961e2b9dc5f3ca8d4da2a1457 Mon Sep 17 00:00:00 2001 From: Joaquin Carletti Date: Tue, 18 Aug 2026 15:38:58 -0300 Subject: [PATCH 1/2] perf(gpu): grind the proof-of-work nonce on the GPU MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Grinding (generate_nonce) runs a ~2^grinding_factor parallel Keccak search per table per epoch and is the prover's dominant CPU cost — 64.7% of on-CPU time in a 100tx flamegraph, on the 16 cores while the GPU sits ~66% idle. Add a keccak nonce-search kernel (each thread strides a nonce block, atomicMin keeps the smallest valid nonce), a math-cuda wrapper that searches in expanding blocks from 0, and a stark dispatch that computes the inner hash on the host, validates the device result unconditionally, and falls back to the CPU search on any device miss or invalid nonce. Result-valid: the verifier only checks is_valid_nonce, so any valid nonce works. A device launch is skipped below a minimum grinding factor (tiny factors are faster on the CPU), and LAMBDA_VM_NO_GPU_GRIND forces the CPU path. GPU_GRIND_CALLS counts the dispatches so a silent fallback is caught by the integration test. 100tx e20 (ABBA, same binary): 18.89s -> 13.10s = -30.6%. --- crypto/math-cuda/kernels/keccak.cu | 58 +++++++++++++++++++ crypto/math-cuda/src/device.rs | 2 + crypto/math-cuda/src/grinding.rs | 80 +++++++++++++++++++++++++++ crypto/math-cuda/src/lib.rs | 1 + crypto/math-cuda/tests/grinding.rs | 59 ++++++++++++++++++++ crypto/stark/src/gpu_lde.rs | 10 ++++ crypto/stark/src/grinding.rs | 49 +++++++++++++++- crypto/stark/src/prover.rs | 5 +- prover/tests/cuda_path_integration.rs | 14 ++++- 9 files changed, 273 insertions(+), 5 deletions(-) create mode 100644 crypto/math-cuda/src/grinding.rs create mode 100644 crypto/math-cuda/tests/grinding.rs diff --git a/crypto/math-cuda/kernels/keccak.cu b/crypto/math-cuda/kernels/keccak.cu index b026ff2b6..2762d7469 100644 --- a/crypto/math-cuda/kernels/keccak.cu +++ b/crypto/math-cuda/kernels/keccak.cu @@ -137,6 +137,64 @@ __device__ __forceinline__ void finalize_keccak256(uint64_t st[25], } } +// --------------------------------------------------------------------------- +// Proof-of-work grinding search. +// +// Mirrors the host `grinding::is_valid_nonce_for_inner_hash`: a nonce is valid +// when the big-endian u64 of the first 8 bytes of +// Keccak256(inner_hash[32] || nonce.to_be_bytes()[8]) +// is `< limit`. The 40-byte message is exactly five Keccak lanes, so there is +// no intermediate block permute — st[0..3] hold the inner hash (passed as four +// LE-read lanes), st[4] holds the nonce lane (`bswap64(nonce)`, since the nonce +// is serialised big-endian and Keccak reads lanes little-endian), padding lands +// in st[5] and st[16], and the head we compare is `bswap64(st[0])` after one +// permutation (the host takes `from_be_bytes(digest[..8])`, i.e. the byte-swap +// of the first squeezed lane). +// +// Each thread strides over `[base, base+count)` and `atomicMin`s the smallest +// valid nonce it finds into `*result` (initialised to U64_MAX by the caller), +// so the launch returns the globally smallest valid nonce in the searched +// block — deterministic, and any valid nonce satisfies the verifier. +extern "C" __global__ void grind_search(const uint64_t *inner_lanes, + uint64_t limit, + uint64_t base, + uint64_t count, + volatile unsigned long long *result) { + uint64_t tid = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; + uint64_t stride = (uint64_t)gridDim.x * blockDim.x; + uint64_t h0 = inner_lanes[0], h1 = inner_lanes[1], h2 = inner_lanes[2], + h3 = inner_lanes[3]; + for (uint64_t i = tid; i < count; i += stride) { + uint64_t nonce = base + i; + // Guard the u64 wrap on the final block (the host bounds the search to + // ~2^36 launches, so this is unreachable in practice): a wrapped nonce + // is < base, so stop rather than re-scan from 0. + if (nonce < base) break; + // A thread's nonces only increase, so once a smaller valid one is known + // this thread can never beat it — stop scanning. `result` is volatile + // so this load re-reads L2 (where the atomicMin writes land) instead of + // being hoisted into a register or served stale from L1; the early exit + // depends on that, though correctness does not. + if (nonce >= (uint64_t)*result) break; + uint64_t st[25]; + #pragma unroll + for (int k = 0; k < 25; ++k) st[k] = 0; + st[0] = h0; + st[1] = h1; + st[2] = h2; + st[3] = h3; + st[4] = bswap64(nonce); + // Keccak (0x01) padding for a 40-byte message: 0x01 at byte 40 (lane 5) + // and 0x80 at byte 135 (top of lane 16). + st[5] ^= (uint64_t)0x01; + st[16] ^= ((uint64_t)0x80) << 56; + keccak_f1600(st); + if (bswap64(st[0]) < limit) { + atomicMin((unsigned long long *)result, (unsigned long long)nonce); + } + } +} + // --------------------------------------------------------------------------- // Goldilocks BASE-FIELD leaf hashing. // diff --git a/crypto/math-cuda/src/device.rs b/crypto/math-cuda/src/device.rs index a7c129cc8..e45ad05dc 100644 --- a/crypto/math-cuda/src/device.rs +++ b/crypto/math-cuda/src/device.rs @@ -196,6 +196,7 @@ pub struct Backend { pub keccak256_leaves_base_batched: CudaFunction, pub keccak256_leaves_base_row_pair_batched: CudaFunction, pub keccak256_leaves_ext3_batched: CudaFunction, + pub grind_search: CudaFunction, pub keccak_comp_poly_leaves_ext3: CudaFunction, pub keccak_fri_leaves_ext3: CudaFunction, pub keccak_merkle_level: CudaFunction, @@ -427,6 +428,7 @@ impl Backend { keccak256_leaves_base_row_pair_batched: keccak .load_function("keccak256_leaves_base_row_pair_batched")?, keccak256_leaves_ext3_batched: keccak.load_function("keccak256_leaves_ext3_batched")?, + grind_search: keccak.load_function("grind_search")?, keccak_comp_poly_leaves_ext3: keccak.load_function("keccak_comp_poly_leaves_ext3")?, keccak_fri_leaves_ext3: keccak.load_function("keccak_fri_leaves_ext3")?, keccak_merkle_level: keccak.load_function("keccak_merkle_level")?, diff --git a/crypto/math-cuda/src/grinding.rs b/crypto/math-cuda/src/grinding.rs new file mode 100644 index 000000000..6c5f539c2 --- /dev/null +++ b/crypto/math-cuda/src/grinding.rs @@ -0,0 +1,80 @@ +//! GPU proof-of-work grinding: a parallel Keccak nonce search that mirrors the +//! host `stark::grinding::generate_nonce`, offloading the ~2^grinding_factor +//! hashes it does per table per epoch from the CPU (where they dominate the +//! prove) to the otherwise-idle GPU. + +use cudarc::driver::{LaunchConfig, PushKernelArg}; + +use crate::device::backend; + +const BLOCK_DIM: u32 = 256; +const GRID_DIM: u32 = 1024; + +/// Below this grinding factor the CPU search finds a valid nonce in well under +/// a microsecond, so a device launch + shared-stream `synchronize` (which also +/// stalls whatever a rayon peer queued on that stream) is pure loss. Bounce +/// those to the CPU. The production factor is 20; only tests use tiny factors. +const GRIND_MIN_FACTOR: u8 = 12; + +/// Smallest nonce whose grind head is `< limit`, or `None` when the CUDA path +/// 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) +/// 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 +/// hits yields the globally smallest valid nonce (the kernel `atomicMin`s it). +pub fn generate_nonce_gpu(inner_lanes: &[u64; 4], grinding_factor: u8) -> Option { + if !(GRIND_MIN_FACTOR..=64).contains(&grinding_factor) { + return None; + } + let limit: u64 = 1u64 << (64 - grinding_factor); + + let be = backend().ok()?; + let stream = be.next_stream(); + let inner_dev = stream.clone_htod(inner_lanes.as_slice()).ok()?; + + // Per-launch block size: ~8× the expected hit distance, clamped so tiny + // factors still launch a full grid and huge factors don't ask for an + // absurd single block. `2^grinding_factor` can overflow u64 (factor 64), so + // saturate. + let expected = 1u64.checked_shl(grinding_factor as u32).unwrap_or(u64::MAX); + let count = expected.saturating_mul(8).clamp(1 << 18, 1 << 28); + + let cfg = LaunchConfig { + grid_dim: (GRID_DIM, 1, 1), + block_dim: (BLOCK_DIM, 1, 1), + shared_mem_bytes: 0, + }; + + // One reusable device slot for the running minimum, reset to the sentinel + // (U64_MAX) before each block rather than reallocated every iteration. + // `sentinel` is a named binding so it outlives every async H2D below. + let sentinel = [u64::MAX]; + let mut result_dev = stream.clone_htod(&sentinel).ok()?; + + let mut base: u64 = 0; + loop { + stream.memcpy_htod(&sentinel, &mut result_dev).ok()?; + unsafe { + stream + .launch_builder(&be.grind_search) + .arg(&inner_dev) + .arg(&limit) + .arg(&base) + .arg(&count) + .arg(&mut result_dev) + .launch(cfg) + .ok()?; + } + let host = stream.clone_dtoh(&result_dev).ok()?; + stream.synchronize().ok()?; + if host[0] != u64::MAX { + return Some(host[0]); + } + // Nothing in `[base, base+count)` — advance. Bail (→ CPU fallback) if + // the block would run past u64, matching the host search's finite range. + base = base.checked_add(count)?; + } +} diff --git a/crypto/math-cuda/src/lib.rs b/crypto/math-cuda/src/lib.rs index d6f19b7c7..838bf9044 100644 --- a/crypto/math-cuda/src/lib.rs +++ b/crypto/math-cuda/src/lib.rs @@ -12,6 +12,7 @@ pub mod device; #[cfg(feature = "test-faults")] pub mod faults; pub mod fri; +pub mod grinding; pub mod inverse; pub mod lde; pub mod logup; diff --git a/crypto/math-cuda/tests/grinding.rs b/crypto/math-cuda/tests/grinding.rs new file mode 100644 index 000000000..76db77cf0 --- /dev/null +++ b/crypto/math-cuda/tests/grinding.rs @@ -0,0 +1,59 @@ +//! 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` +//! (`cargo test -p math-cuda --release`) — `device::backend()` inside +//! `generate_nonce_gpu` requires a real GPU, like the other tests here. +//! +//! Uses real grinding factors (>= the min-factor gate). The end-to-end prover +//! 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. + +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())) +} + +/// 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). +#[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) + .expect("GPU grind (needs a GPU)"); + assert!( + is_valid_nonce(&seed, nonce, factor), + "GPU nonce {nonce} fails is_valid_nonce (factor {factor})" + ); + assert!( + (0..nonce).all(|n| !is_valid_nonce(&seed, n, factor)), + "GPU nonce {nonce} is not the smallest valid nonce (factor {factor})" + ); +} + +/// At the production factor the kernel returns a valid nonce (validity only — +/// scanning 0..nonce would be ~2^20 hashes). +#[test] +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) + .expect("GPU grind (needs a GPU)"); + assert!( + is_valid_nonce(&seed, nonce, factor), + "GPU nonce {nonce} fails is_valid_nonce (factor {factor})" + ); +} + +/// Below the min-factor gate the GPU path declines (→ CPU search), so the tiny +/// factors every non-GPU-benchmark test uses never pay a launch. +#[test] +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(), + "GPU grind should decline factor 1" + ); +} diff --git a/crypto/stark/src/gpu_lde.rs b/crypto/stark/src/gpu_lde.rs index a1ec18fa7..52faa8d3e 100644 --- a/crypto/stark/src/gpu_lde.rs +++ b/crypto/stark/src/gpu_lde.rs @@ -118,6 +118,16 @@ pub fn reset_all_gpu_call_counters() { GPU_RESIDENT_AUX_RETRIES.store(0, Ordering::Relaxed); GPU_RESIDENT_AUX_DOWNGRADES.store(0, Ordering::Relaxed); GPU_COMPOSITION_PARTS_DOWNLOADS.store(0, Ordering::Relaxed); + GPU_GRIND_CALLS.store(0, Ordering::Relaxed); +} + +/// Successful GPU proof-of-work grind dispatches — one per table whose round-4 +/// nonce search ran on device and produced a nonce that passed the host +/// validity check (a device miss or an invalid kernel result falls back to the +/// CPU search and is not counted). +pub(crate) static GPU_GRIND_CALLS: AtomicU64 = AtomicU64::new(0); +pub fn gpu_grind_calls() -> u64 { + GPU_GRIND_CALLS.load(Ordering::Relaxed) } pub(crate) static GPU_EXTEND_HALVES_CALLS: AtomicU64 = AtomicU64::new(0); diff --git a/crypto/stark/src/grinding.rs b/crypto/stark/src/grinding.rs index 4666b7946..b04fda912 100644 --- a/crypto/stark/src/grinding.rs +++ b/crypto/stark/src/grinding.rs @@ -78,7 +78,10 @@ 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` -fn get_inner_hash(seed: &[u8; 32], grinding_factor: u8) -> [u8; 32] { +/// +/// 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] { let mut inner_data = [0u8; 41]; inner_data[0..8].copy_from_slice(&PREFIX); inner_data[8..40].copy_from_slice(seed); @@ -87,3 +90,47 @@ fn get_inner_hash(seed: &[u8; 32], grinding_factor: u8) -> [u8; 32] { let digest = Keccak256::digest(inner_data); digest[..32].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. +#[cfg(feature = "cuda")] +pub fn generate_nonce_maybe_gpu(seed: &[u8; 32], grinding_factor: u8) -> Option { + debug_assert!( + (1..=64).contains(&grinding_factor), + "grinding_factor must be in 1..=64, got {grinding_factor}" + ); + // Kill switch (presence-based, matching `LAMBDA_VM_NO_GPU_LOGUP`): + // `LAMBDA_VM_NO_GPU_GRIND` forces the CPU search — a production escape hatch + // and fallback-path coverage. Cached; read once. + static GPU_DISABLED: std::sync::OnceLock = std::sync::OnceLock::new(); + 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()) + }); + 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, + // never append an unverifiable nonce to the transcript. This runs in + // release too — the cost is negligible next to the grind it replaces. + if is_valid_nonce(seed, nonce, grinding_factor) { + 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"); + } + generate_nonce(seed, grinding_factor) +} + +#[cfg(not(feature = "cuda"))] +pub fn generate_nonce_maybe_gpu(seed: &[u8; 32], grinding_factor: u8) -> Option { + generate_nonce(seed, grinding_factor) +} diff --git a/crypto/stark/src/prover.rs b/crypto/stark/src/prover.rs index f67fea4e6..f31e6c1c1 100644 --- a/crypto/stark/src/prover.rs +++ b/crypto/stark/src/prover.rs @@ -2203,8 +2203,9 @@ pub trait IsStarkProver< let security_bits = air.context().proof_options.grinding_factor; let mut nonce = None; if security_bits > 0 { - let nonce_value = grinding::generate_nonce(&transcript.state(), security_bits) - .expect("nonce not found"); + let nonce_value = + grinding::generate_nonce_maybe_gpu(&transcript.state(), security_bits) + .expect("nonce not found"); transcript.append_bytes(&nonce_value.to_be_bytes()); nonce = Some(nonce_value); } diff --git a/prover/tests/cuda_path_integration.rs b/prover/tests/cuda_path_integration.rs index 29f0070d8..b8e540a3b 100644 --- a/prover/tests/cuda_path_integration.rs +++ b/prover/tests/cuda_path_integration.rs @@ -14,8 +14,9 @@ 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_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, + gpu_deep_calls, gpu_device_only_calls, gpu_extend_halves_calls, gpu_fri_calls, gpu_grind_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 @@ -108,6 +109,15 @@ fn gpu_path_fires_end_to_end() { "GPU batch-invert dispatch did not fire on R3 + R4" ); + // R4 proof-of-work grind: with_blowup(2) grinds at factor 20 (above the + // GPU min-factor gate), so the device search fires for every table and a + // valid nonce is served. A silent CPU fallback (or an invalid kernel result + // rejected by the host check) would drop this to zero. + assert!( + gpu_grind_calls() > 0, + "R4 GPU proof-of-work grind did not fire" + ); + // Counters only prove the dispatches ran; this checks the GPU proof // actually satisfies the verifier. let ok = verify(&proof, &elf).expect("verify"); From f324dfdaac940861718708549499d9a21c2031c5 Mon Sep 17 00:00:00 2001 From: Mauro Toscano <12560266+MauroToscano@users.noreply.github.com> Date: Fri, 21 Aug 2026 17:05:00 -0300 Subject: [PATCH 2/2] fix(gpu): review follow-ups on the GPU grinding PR (#945) 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