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");