perf(gpu): grind the proof-of-work nonce on the GPU

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.
 //

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")?,

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<u64> {
+    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)?;
+    }
+}

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;

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

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

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<u64> {
+    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<bool> = 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<u64> {
+    generate_nonce(seed, grinding_factor)
+}

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

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