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/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..fe7803eb9 --- /dev/null +++ b/crypto/math-cuda/src/grinding.rs @@ -0,0 +1,81 @@ +//! 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 — 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 +/// 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..84bc5e624 --- /dev/null +++ b/crypto/math-cuda/tests/grinding.rs @@ -0,0 +1,71 @@ +//! 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. +//! +//! 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. +//! +//! 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::{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(&inner_hash_lanes(&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(&inner_hash_lanes(&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(&inner_hash_lanes(&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..adb7601b6 100644 --- a/crypto/stark/src/grinding.rs +++ b/crypto/stark/src/grinding.rs @@ -87,3 +87,64 @@ fn get_inner_hash(seed: &[u8; 32], grinding_factor: u8) -> [u8; 32] { let digest = Keccak256::digest(inner_data); 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). 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!( + (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_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, + // 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); + } + // 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) +} + +#[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"); 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