perf(prover): device-resident rounds 2-4 and fused NTT for GPU continuations

crypto/crypto/src/merkle_tree/merkle.rs

@@ -168,6 +168,19 @@ where
         })
     }
 
+    /// True when this tree carries only its root (the nodes live elsewhere,
+    /// e.g. device-resident): openings must not walk this tree.
+    pub fn is_root_only(&self) -> bool {
+        #[cfg(feature = "disk-spill")]
+        {
+            self.nodes.is_empty() && self.mmap_backing.is_none()
+        }
+        #[cfg(not(feature = "disk-spill"))]
+        {
+            self.nodes.is_empty()
+        }
+    }
+
     /// Create a root only Merkle tree placeholder: stores the commitment root
     /// but no nodes. Used when paths are gathered from a device resident copy
     /// (GPU) instead of this host tree, so the host nodes are never built.
@@ -253,7 +266,14 @@ where
     /// Returns a Merkle proof for the element/s at position pos
     /// For example, give me an inclusion proof for the 3rd element in the
     /// Merkle tree
+    ///
+    /// Returns `None` on a root-only tree ([`from_root`](Self::from_root)):
+    /// its nodes live elsewhere (e.g. device-resident), so a host path would
+    /// be a silently-empty bogus proof rather than an inclusion witness.
     pub fn get_proof_by_pos(&self, pos: usize) -> Option<Proof<B::Node>> {
+        if self.is_root_only() {
+            return None;
+        }
         let pos = pos + self.node_count() / 2;
         let Ok(merkle_path) = self.build_merkle_path(pos) else {
             return None;

crypto/math-cuda/kernels/fri.cu

@@ -59,3 +59,20 @@ extern "C" __global__ void fri_update_twiddles(
     uint64_t old = tw_in[2 * j];
     tw_out[j] = goldilocks::mul(old, old);
 }
+
+// Gather interleaved ext3 elements at arbitrary positions: one thread per
+// query copies evals[positions[i]] (3 u64) into out[i]. Serves the FRI query
+// phase's symmetric-eval reads off the resident layer buffers.
+extern "C" __global__ void gather_ext3_at(
+    const uint64_t *evals,
+    const uint32_t *positions,
+    uint64_t q,
+    uint64_t *out
+) {
+    uint64_t i = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x;
+    if (i >= q) return;
+    uint64_t p = positions[i];
+    out[i * 3]     = evals[p * 3];
+    out[i * 3 + 1] = evals[p * 3 + 1];
+    out[i * 3 + 2] = evals[p * 3 + 2];
+}

crypto/math-cuda/kernels/inverse.cu

@@ -309,3 +309,38 @@ extern "C" __global__ void batch_inverse_combine_ext3(
     out_base[1] = res.b;
     out_base[2] = res.c;
 }
+
+// ---------------------------------------------------------------------------
+// 7. invert_total_ext3
+//
+// One-thread Fermat inversion of the scan total: out = src[n-1]^(p^3 - 2).
+// Replaces the host round-trip (D2H + host Fermat + H2D + stream sync) so the
+// whole batch inverse stays stream-ordered. The 192-bit exponent arrives as
+// three little-endian u64 limbs.
+// ---------------------------------------------------------------------------
+extern "C" __global__ void invert_total_ext3(
+    const uint64_t *src,   // 3 * n u64 (reads element n-1)
+    uint64_t n,
+    uint64_t e0,           // exponent limbs, little-endian
+    uint64_t e1,
+    uint64_t e2,
+    uint64_t *out          // 3 u64
+) {
+    if (blockIdx.x != 0 || threadIdx.x != 0) return;
+    const uint64_t *base = src + (n - 1) * 3;
+    ext3::Fe3 a = {base[0], base[1], base[2]};
+    ext3::Fe3 r = ext3::one();
+    uint64_t limbs[3] = {e0, e1, e2};
+    for (int li = 2; li >= 0; --li) {
+        uint64_t bits = limbs[li];
+        for (int b = 63; b >= 0; --b) {
+            r = ext3::mul(r, r);
+            if ((bits >> b) & 1) {
+                r = ext3::mul(r, a);
+            }
+        }
+    }
+    out[0] = r.a;
+    out[1] = r.b;
+    out[2] = r.c;
+}

crypto/math-cuda/kernels/keccak.cu

@@ -366,13 +366,11 @@ extern "C" __global__ void keccak_fri_leaves_ext3(
 // concatenation of two 32-byte siblings, identical to
 // `FieldElementVectorBackend::hash_new_parent` on host.
 // ---------------------------------------------------------------------------
-extern "C" __global__ void keccak_merkle_level(
+__device__ __forceinline__ void hash_merkle_parent(
     uint8_t *nodes,
     uint64_t parent_begin,     // node index (counted in 32-byte nodes)
-    uint64_t n_pairs) {
-    uint64_t tid = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x;
-    if (tid >= n_pairs) return;
-
+    uint64_t n_pairs,
+    uint64_t tid) {
     uint64_t st[25];
     #pragma unroll
     for (int i = 0; i < 25; ++i) st[i] = 0;
@@ -393,6 +391,35 @@ extern "C" __global__ void keccak_merkle_level(
     finalize_keccak256(st, rate_pos, nodes + (parent_begin + tid) * 32);
 }
 
+extern "C" __global__ void keccak_merkle_level(
+    uint8_t *nodes,
+    uint64_t parent_begin,     // node index (counted in 32-byte nodes)
+    uint64_t n_pairs) {
+    uint64_t tid = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x;
+    if (tid >= n_pairs) return;
+    hash_merkle_parent(nodes, parent_begin, n_pairs, tid);
+}
+
+// Build every remaining level (from `level_begin` up to the root) in ONE
+// single-block launch: each level's pairs are grid-strided over the block,
+// with a __syncthreads() barrier between levels. Replaces log2 launches of
+// `keccak_merkle_level` for the small top levels of the tree, whose per-level
+// work is dwarfed by launch overhead.
+extern "C" __global__ void keccak_merkle_tail(
+    uint8_t *nodes,
+    uint64_t level_begin) {
+    uint64_t lb = level_begin;
+    while (lb != 0) {
+        uint64_t nb = lb / 2;
+        uint64_t n_pairs = lb - nb;
+        for (uint64_t tid = threadIdx.x; tid < n_pairs; tid += blockDim.x) {
+            hash_merkle_parent(nodes, nb, n_pairs, tid);
+        }
+        __syncthreads();
+        lb = nb;
+    }
+}
+
 // Gather Merkle authentication paths for a batch of leaf positions, reading the
 // resident tree `nodes` (32-byte nodes; layout: inner nodes [0..leaves_len-1],
 // root at 0, leaves at [leaves_len-1..]). One thread per query walks leaf->root,

crypto/math-cuda/kernels/ntt.cu

@@ -411,3 +411,67 @@ extern "C" __global__ void matrix_transpose_strided(
         __syncthreads();
     }
 }
+
+// First-8-levels fused DIT on row-major data: one block stages 256 consecutive
+// rows x blockDim.x columns in shmem and runs levels 0..min(8,log_n) with
+// __syncthreads between levels (row-major analog of ntt_dit_8_levels_batched
+// with base_step == 0, whose twiddle math this reuses verbatim). Grid:
+// x = column tiles, y = n/256 row blocks. Requires n >= 256. Shmem tile is
+// padded (pitch = T+1) to break bank conflicts on the butterfly accesses.
+extern "C" __global__ void ntt_dit_8_levels_row_major(uint64_t *data,
+                                                      const uint64_t *tw,
+                                                      uint64_t n,
+                                                      uint64_t log_n,
+                                                      uint64_t m)
+{
+    extern __shared__ uint64_t tile[];
+    uint32_t T = blockDim.x;
+    uint32_t pitch = T + 1;
+    uint64_t col = (uint64_t)blockIdx.x * T + threadIdx.x;
+    bool live = col < m;
+
+    uint32_t n_loc_steps = (uint32_t)min((uint64_t)8, log_n);
+    uint32_t remaining_high_bits = (uint32_t)(log_n - 1);
+    uint32_t high_mask = (1u << remaining_high_bits) - 1u;
+
+    // Grid-stride over 256-row blocks: gridDim.y caps at 65535, so lde sizes
+    // >= 2^24 need more than one row block per y-slot. The trip count is
+    // uniform across the block, keeping every __syncthreads converged.
+    for (uint64_t rb = blockIdx.y; rb < (n >> 8); rb += gridDim.y) {
+        uint64_t row_base = rb * 256;
+
+        for (uint32_t r = threadIdx.y; r < 256; r += blockDim.y) {
+            if (live) tile[r * pitch + threadIdx.x] = data[(row_base + r) * m + col];
+        }
+        __syncthreads();
+
+        for (uint32_t loc_step = 0; loc_step < n_loc_steps; ++loc_step) {
+            for (uint32_t i = threadIdx.y; i < 128; i += blockDim.y) {
+                uint32_t half    = 1u << loc_step;
+                uint32_t grp     = i >> loc_step;
+                uint32_t grp_pos = i & (half - 1);
+                uint32_t idx1 = (grp << (loc_step + 1)) + grp_pos;
+                uint32_t idx2 = idx1 + half;
+
+                uint32_t gs  = loc_step;
+                uint32_t ggp = ((uint32_t)rb << 7) + i;
+                ggp = (ggp & high_mask) + (ggp >> remaining_high_bits);
+                ggp = ggp & ((1u << gs) - 1u);
+                uint64_t factor = tw[(uint64_t)ggp * (n >> (gs + 1))];
+
+                if (live) {
+                    uint64_t u = tile[idx1 * pitch + threadIdx.x];
+                    uint64_t v = mul(tile[idx2 * pitch + threadIdx.x], factor);
+                    tile[idx1 * pitch + threadIdx.x] = add(u, v);
+                    tile[idx2 * pitch + threadIdx.x] = sub(u, v);
+                }
+            }
+            __syncthreads();
+        }
+
+        for (uint32_t r = threadIdx.y; r < 256; r += blockDim.y) {
+            if (live) data[(row_base + r) * m + col] = tile[r * pitch + threadIdx.x];
+        }
+        __syncthreads();
+    }
+}

crypto/math-cuda/src/device.rs

@@ -186,6 +186,7 @@ pub struct Backend {
     // row-major NTT kernels
     pub bit_reverse_row_major: CudaFunction,
     pub ntt_dit_level_row_major: CudaFunction,
+    pub ntt_dit_8_levels_row_major: CudaFunction,
     pub pointwise_mul_row_major: CudaFunction,
     pub matrix_transpose_strided: CudaFunction,
 
@@ -198,6 +199,7 @@ pub struct Backend {
     pub keccak_comp_poly_leaves_ext3: CudaFunction,
     pub keccak_fri_leaves_ext3: CudaFunction,
     pub keccak_merkle_level: CudaFunction,
+    pub keccak_merkle_tail: CudaFunction,
     pub merkle_gather_paths: CudaFunction,
 
     // barycentric.cubin
@@ -214,6 +216,7 @@ pub struct Backend {
 
     // fri.cubin
     pub fri_fold_ext3: CudaFunction,
+    pub gather_ext3_at: CudaFunction,
     pub fri_update_twiddles: CudaFunction,
 
     // inverse.cubin
@@ -223,6 +226,7 @@ pub struct Backend {
     pub block_inclusive_scan_rev_ext3: CudaFunction,
     pub apply_block_offsets_rev_ext3: CudaFunction,
     pub batch_inverse_combine_ext3: CudaFunction,
+    pub invert_total_ext3: CudaFunction,
     pub logup_fingerprint_ext3: CudaFunction,
     pub logup_term_ext3: CudaFunction,
     pub logup_row_sum_ext3: CudaFunction,
@@ -412,6 +416,7 @@ impl Backend {
             scalar_mul_batched: ntt.load_function("scalar_mul_batched")?,
             bit_reverse_row_major: ntt.load_function("bit_reverse_row_major")?,
             ntt_dit_level_row_major: ntt.load_function("ntt_dit_level_row_major")?,
+            ntt_dit_8_levels_row_major: ntt.load_function("ntt_dit_8_levels_row_major")?,
             pointwise_mul_row_major: ntt.load_function("pointwise_mul_row_major")?,
             matrix_transpose_strided: ntt.load_function("matrix_transpose_strided")?,
             keccak256_leaves_base_row_major_row_pair: keccak
@@ -425,6 +430,7 @@ impl Backend {
             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")?,
+            keccak_merkle_tail: keccak.load_function("keccak_merkle_tail")?,
             merkle_gather_paths: keccak.load_function("merkle_gather_paths")?,
             barycentric_base_batched: bary.load_function("barycentric_base_batched")?,
             barycentric_ext3_batched: bary.load_function("barycentric_ext3_batched")?,
@@ -437,6 +443,7 @@ impl Backend {
             deep_composition_ext3_row: deep.load_function("deep_composition_ext3_row")?,
             bit_reverse_ext3_kernel: deep.load_function("bit_reverse_ext3_interleaved")?,
             fri_fold_ext3: fri.load_function("fri_fold_ext3")?,
+            gather_ext3_at: fri.load_function("gather_ext3_at")?,
             fri_update_twiddles: fri.load_function("fri_update_twiddles")?,
             compute_denoms_ext3: inverse.load_function("compute_denoms_ext3")?,
             block_inclusive_scan_fwd_ext3: inverse
@@ -446,6 +453,7 @@ impl Backend {
                 .load_function("block_inclusive_scan_rev_ext3")?,
             apply_block_offsets_rev_ext3: inverse.load_function("apply_block_offsets_rev_ext3")?,
             batch_inverse_combine_ext3: inverse.load_function("batch_inverse_combine_ext3")?,
+            invert_total_ext3: inverse.load_function("invert_total_ext3")?,
             logup_fingerprint_ext3: logup.load_function("logup_fingerprint_ext3")?,
             logup_term_ext3: logup.load_function("logup_term_ext3")?,
             logup_row_sum_ext3: logup.load_function("logup_row_sum_ext3")?,
@@ -603,7 +611,9 @@ pub fn backend() -> Result<&'static Backend> {
 ///
 /// Holding this value keeps the staging slot's mutex locked, which is what
 /// makes the whole scheme safe: no other caller (and no capacity growth) can
-/// touch the slab while the DMA is in flight.
+/// touch the slab while the DMA is in flight. Corollary: never call
+/// `htod_via`/`async_dtoh_via` on the same slot from the thread holding a
+/// live `PendingD2H` — the non-reentrant slot mutex self-deadlocks.
 pub struct PendingD2H<'a> {
     staging: std::sync::MutexGuard<'a, PinnedStaging>,
     n_bytes: usize,
@@ -620,6 +630,91 @@ impl Drop for PendingD2H<'_> {
     }
 }
 
+/// Chunk size for [`htod_via`]'s staged upload — the upper bound a single H2D
+/// puts on a staging slot's page-locked footprint. 64 MB is large enough to
+/// amortize the per-chunk DMA launch + event sync, small enough to keep the
+/// pinned slab independent of trace size.
+const HTOD_CHUNK_BYTES: usize = 64 << 20; // 64 MB
+
+/// Host→device copy staged through the pinned slot, in fixed-size chunks: each
+/// chunk is one host memcpy into pinned memory + one async DMA, instead of the
+/// driver's internal pageable staging (small chunks; 2-3x slower for
+/// multi-hundred-MB traces and it convoys under multi-thread load). Blocks
+/// until the last DMA lands, so the slot and `src_host` are both reusable on
+/// return.
+///
+/// Chunking caps the slot's page-locked footprint at [`HTOD_CHUNK_BYTES`]
+/// regardless of trace size. This matters on the device-only path
+/// (`retain_host_lde = false`): there is no [`async_dtoh_via`] drain to size
+/// the slot, so `htod_via` is its only writer — an uncapped copy would grow
+/// the per-worker slab to a whole trace and, being grow-only, never shrink it.
+/// The host-retaining path is unaffected: its later `async_dtoh_via` grows the
+/// same slot to the full LDE anyway, and we simply reuse the first chunk of it.
+pub fn htod_via<T: cudarc::driver::DeviceRepr>(
+    stream: &Arc<CudaStream>,
+    slot: &Mutex<PinnedStaging>,
+    ctx: &CudaContext,
+    src_host: &[T],
+    dst: &mut cudarc::driver::CudaViewMut<'_, T>,
+) -> Result<()> {
+    use cudarc::driver::DevicePtrMut;
+    assert!(
+        dst.len() >= src_host.len(),
+        "htod_via: destination shorter than source"
+    );
+    let n_bytes = std::mem::size_of_val(src_host);
+    if n_bytes == 0 {
+        return Ok(());
+    }
+    let elem_size = std::mem::size_of::<T>();
+    // Chunk in whole elements so a `T` never straddles a chunk boundary.
+    let chunk_elems = (HTOD_CHUNK_BYTES / elem_size.max(1)).max(1);
+
+    let mut staging = slot.lock().unwrap();
+    // Only ask for a chunk's worth of pinned memory (or the whole copy when
+    // smaller). If another path (`async_dtoh_via` on the host-retaining flow)
+    // already grew this slot larger, it stays larger — grow-only — and we just
+    // use the first chunk of it.
+    let want_u64 = (chunk_elems * elem_size)
+        .div_ceil(8)
+        .min(n_bytes.div_ceil(8));
+    staging.ensure_capacity(want_u64, ctx)?;
+    ctx.bind_to_thread()?;
+
+    // SAFETY: `device_ptr_mut` yields the destination base pointer and orders
+    // the device writes on `stream`; `dst.len() >= src_host.len()` (asserted),
+    // so every chunk's byte range stays within `dst`.
+    let (dst_base, _record) = dst.device_ptr_mut(stream);
+    let src = src_host.as_ptr() as *const u8;
+    let n_elems = src_host.len();
+    let mut elem_off = 0usize;
+    while elem_off < n_elems {
+        let this_elems = (n_elems - elem_off).min(chunk_elems);
+        let this_bytes = this_elems * elem_size;
+        let byte_off = elem_off * elem_size;
+        // SAFETY: the pinned slab holds at least `chunk_elems * elem_size`
+        // bytes (or the whole copy when smaller). The previous chunk's DMA is
+        // synced below before this memcpy overwrites the slab, so the slab is
+        // never read (by an in-flight DMA) and written at the same time.
+        unsafe {
+            std::ptr::copy_nonoverlapping(src.add(byte_off), staging.ptr as *mut u8, this_bytes);
+            cudarc::driver::sys::cuMemcpyHtoDAsync_v2(
+                dst_base + byte_off as u64,
+                staging.ptr as *const core::ffi::c_void,
+                this_bytes,
+                stream.cu_stream(),
+            )
+            .result()?;
+        }
+        // Single-buffered: wait for this chunk's DMA before the next memcpy
+        // reuses the slab.
+        staging.record_event(stream)?;
+        staging.sync_event()?;
+        elem_off += this_elems;
+    }
+    Ok(())
+}
+
 /// Enqueue an async D2H of `n_elems` of `src` into the pinned slab of `slot`,
 /// without synchronizing the stream. Unlike `stream.memcpy_dtoh` into a plain
 /// (pageable) slice — which the driver services synchronously — this returns