feat(gpu): bytes-aware GPU admission, loud abort inside the device envelope, BALU row-chunking policy

crypto/math-cuda/src/blake3.rs

@@ -603,6 +603,10 @@ pub fn build_comp_poly_tree_from_slabs_dev(
     m: usize,
     lde_size: usize,
 ) -> Result<crate::lde::GpuMerkleTree> {
+    // Same sticky hook as the keccak twin: the comp-tree cliff test arms one
+    // counter and must reach it under whichever hash the build pins.
+    #[cfg(feature = "test-faults")]
+    crate::faults::check_sticky(&crate::faults::FAULT_COMP_TREE_STICKY)?;
     assert!(m > 0);
     assert!(lde_size.is_power_of_two() && lde_size >= 2);
     assert_eq!(buf.len(), 3 * m * lde_size, "slab buffer shape");
@@ -647,6 +651,8 @@ pub fn build_comp_poly_tree_from_slabs_dev(
 pub fn build_comp_poly_tree_from_evals_ext3_keep(
     parts_interleaved: &[&[u64]],
 ) -> Result<crate::lde::GpuMerkleTree> {
+    #[cfg(feature = "test-faults")]
+    crate::faults::check_sticky(&crate::faults::FAULT_COMP_TREE_STICKY)?;
     assert!(!parts_interleaved.is_empty());
     let m = parts_interleaved.len();
     let ext3_elems = parts_interleaved[0].len() / 3;

crypto/math-cuda/src/device.rs

@@ -285,42 +285,102 @@ pub struct Backend {
     inv_twiddles: Mutex<Vec<Option<Arc<CudaSlice<u64>>>>>,
 }
 
-/// Raise the device default memory pool's release threshold so freed
-/// stream-ordered allocations are kept for reuse instead of returned to the OS
-/// at each sync. Best-effort: any failure (e.g. a device/driver without
-/// stream-ordered allocator support) leaves the default behaviour untouched.
-fn retain_default_mempool(ctx: &CudaContext) {
+/// The environment knob for the device default memory pool's release
+/// threshold, in MiB: the bytes of freed stream-ordered memory the pool keeps
+/// before handing memory back to the OS at the next sync. Unset means
+/// [`DEFAULT_MEMPOOL_RELEASE_THRESHOLD_BYTES`]. The VRAM sampler runs set it
+/// to `0`, so `total - free` reads the live working set and not the retained
+/// pool.
+pub const MEMPOOL_RELEASE_ENV: &str = "LAMBDA_VM_MEMPOOL_RELEASE_MB";
+
+/// Retain every freed block (`u64::MAX`): a same-shape allocation skips the
+/// driver and reuses the block, which is what the per-table pipeline's
+/// repeated LDE/FRI buffers want.
+///
+/// Measured, not guessed (RTX 5090, 2026-09-07, `one_lde_buffer::vram_arm`
+/// at 2^21 × 316 @ blowup 2, five commits, in-process 1 kHz peak): retain-all
+/// 1176.9 / 1057.9 / 1055.3 / 1056.2 / 1055.6 ms against release-0
+/// 1178.5 / 1057.2 / 1077.5 / 1081.2 / 1079.3 ms — retention ≈2% faster once
+/// the first commit has populated the pool — and a peak of 15.67 GiB under
+/// both: a same-shape allocation reuses the retained block, so retention adds
+/// nothing to the peak. The unequal-shape case is covered at block scale by
+/// the multi-table q=41 wrap rung under this default (VRAM peak 28,976 MiB, no
+/// device decline): the stream-ordered allocator serves a new request from the
+/// physical chunks it retains, and the release threshold governs only what a
+/// sync hands back to the OS. The explicit release for a moment reuse cannot
+/// serve is [`Backend::trim_mempool_to`]; the sampler runs set the knob to `0`
+/// so `total - free` reads the live set rather than the pool.
+pub const DEFAULT_MEMPOOL_RELEASE_THRESHOLD_BYTES: u64 = u64::MAX;
+
+/// The effective release threshold in bytes: the knob when set and parseable,
+/// the default otherwise. Read once per process; the prover's diagnostics
+/// print it so every box log states the posture its run had.
+pub fn mempool_release_threshold_bytes() -> u64 {
+    static CACHED: OnceLock<u64> = OnceLock::new();
+    *CACHED.get_or_init(|| {
+        std::env::var(MEMPOOL_RELEASE_ENV)
+            .ok()
+            .and_then(|s| s.parse::<u64>().ok())
+            .map(|mb| mb.saturating_mul(1024 * 1024))
+            .unwrap_or(DEFAULT_MEMPOOL_RELEASE_THRESHOLD_BYTES)
+    })
+}
+
+/// The device default memory pool, or `None` on a device/driver without
+/// stream-ordered allocator support.
+///
+/// # Safety
+///
+/// `ctx` must be a live context; its device is queried directly.
+unsafe fn default_mempool(ctx: &CudaContext) -> Option<cudarc::driver::sys::CUmemoryPool> {
     use cudarc::driver::sys;
-    // SAFETY: raw CUDA driver calls. `ctx.cu_device()` is a valid device for
-    // the just-created context; the out-pointers are valid stack slots; the
-    // threshold is read as a u64 by the driver. Errors are swallowed.
+    let mut pool: sys::CUmemoryPool = std::ptr::null_mut();
+    // SAFETY: the out-pointer is a valid stack slot; the device is the
+    // context's own.
     unsafe {
-        let dev = ctx.cu_device();
-        let mut pool: sys::CUmemoryPool = std::ptr::null_mut();
-        if sys::cuDeviceGetDefaultMemPool(&mut pool as *mut _, dev)
+        sys::cuDeviceGetDefaultMemPool(&mut pool as *mut _, ctx.cu_device())
             .result()
-            .is_err()
-        {
-            return;
-        }
-        // Default: retain freed stream-ordered blocks indefinitely (u64::MAX)
-        // for reuse. `LAMBDA_VM_MEMPOOL_RELEASE_MB` overrides the cap (bytes the
-        // pool keeps before returning memory to the OS) when retained-pool
-        // growth needs bounding.
-        let threshold: u64 = std::env::var("LAMBDA_VM_MEMPOOL_RELEASE_MB")
             .ok()
-            .and_then(|s| s.parse::<u64>().ok())
-            .map(|mb| mb.saturating_mul(1024 * 1024))
-            .unwrap_or(u64::MAX);
-        let _ = sys::cuMemPoolSetAttribute(
-            pool,
-            sys::CUmemPool_attribute_enum::CU_MEMPOOL_ATTR_RELEASE_THRESHOLD,
-            &threshold as *const u64 as *mut core::ffi::c_void,
-        )
-        .result();
+            .map(|()| pool)
     }
 }
 
+/// Set the device default memory pool's release threshold
+/// ([`mempool_release_threshold_bytes`]) so freed stream-ordered allocations
+/// are kept for reuse instead of returned to the OS at each sync. Best-effort:
+/// any failure leaves the driver default (release everything) untouched, and
+/// the one-line report says so.
+fn retain_default_mempool(ctx: &CudaContext) {
+    use cudarc::driver::sys;
+    let threshold = mempool_release_threshold_bytes();
+    // SAFETY: raw CUDA driver calls on the just-created context's device; the
+    // threshold is read as a u64 by the driver. Errors are swallowed.
+    let set = unsafe {
+        default_mempool(ctx).is_some_and(|pool| {
+            sys::cuMemPoolSetAttribute(
+                pool,
+                sys::CUmemPool_attribute_enum::CU_MEMPOOL_ATTR_RELEASE_THRESHOLD,
+                &threshold as *const u64 as *mut core::ffi::c_void,
+            )
+            .result()
+            .is_ok()
+        })
+    };
+    // One line per process, so the box log states the posture the run had.
+    eprintln!(
+        "[gpu] mempool release threshold: {}{}",
+        match threshold {
+            u64::MAX => "retain all freed blocks".to_string(),
+            t => format!("{} MiB", t >> 20),
+        },
+        if set {
+            ""
+        } else {
+            " (driver refused; the release-on-sync default stays)"
+        }
+    );
+}
+
 /// Device VRAM budget in bytes for table session admission control.
 ///
 /// LAMBDA_VM_VRAM_BUDGET_MB overrides it (used to force the throttle in tests).
@@ -557,6 +617,36 @@ impl Backend {
         self.vram_budget_bytes
     }
 
+    /// Live `(free, total)` device memory in bytes, for diagnostics — the
+    /// admission gates never read it (they must answer the same at R1 and at
+    /// R4). `None` when the query fails.
+    pub fn device_mem_info(&self) -> Option<(u64, u64)> {
+        self.ctx
+            .mem_get_info()
+            .ok()
+            .map(|(free, total)| (free as u64, total as u64))
+    }
+
+    /// Hand the default memory pool's unused reserved memory back to the OS,
+    /// keeping at most `keep_bytes` (`cuMemPoolTrimTo`). Under the retained
+    /// posture ([`mempool_release_threshold_bytes`]) a sync never releases;
+    /// this is the explicit release for the moments reuse cannot serve — a
+    /// differently-shaped table after a large one, or a sampler that must read
+    /// the live working set. Best effort: `false` when the pool cannot be
+    /// queried or the trim fails.
+    pub fn trim_mempool_to(&self, keep_bytes: u64) -> bool {
+        use cudarc::driver::sys;
+        // SAFETY: raw driver calls on this backend's live context; the trim
+        // takes a plain byte count.
+        unsafe {
+            default_mempool(&self.ctx).is_some_and(|pool| {
+                sys::cuMemPoolTrimTo(pool, keep_bytes as usize)
+                    .result()
+                    .is_ok()
+            })
+        }
+    }
+
     /// Round-robin over the stream pool. Concurrent callers get different
     /// streams so their kernel launches overlap on the GPU.
     pub fn next_stream(&self) -> Arc<CudaStream> {