fix(gpu): review follow-ups on the GPU grinding PR

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.

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

crypto/math-cuda/tests/grinding.rs

@@ -1,27 +1,39 @@
-//! 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.
 //!
 //! 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::{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"
     );
 }

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

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