CENO_GPU_MEM_TRACKING

Author: Velaciela

ceno_zkvm/src/scheme/gpu/memory.rs

@@ -14,30 +14,44 @@ use gkr_iop::gpu::{
     gpu_prover::{BB31Ext, CudaHalBB31, MemTracker},
 };
 use mpcs::PolynomialCommitmentScheme;
+use std::sync::OnceLock;
 
 use crate::scheme::scheduler::{ChipProvingMode, get_chip_proving_mode};
 
-const ESTIMATION_TOLERANCE_BYTES: usize = 1024 * 1024; // max estimation error: 1 MB
-const ESTIMATION_SAFETY_MARGIN_BYTES: usize = 5 * 1024 * 1024; // reserved headroom: 5 MB, 1MB for each sub-stage
+#[cfg(feature = "gpu")]
+static MEM_TRACKING_MODE: OnceLock<bool> = OnceLock::new();
+
+#[cfg(feature = "gpu")]
+pub fn get_mem_tracking_mode() -> bool {
+    *MEM_TRACKING_MODE.get_or_init(|| match std::env::var("CENO_GPU_MEM_TRACKING").as_deref() {
+        Ok("1") => true,
+        _ => false,
+    })
+}
 
 #[cfg(feature = "gpu")]
 pub fn start_gpu_mem_tracking<'a>(
     cuda_hal: &'a CudaHalBB31,
     label: &'static str,
 ) -> Option<MemTracker<'a>> {
-    if get_chip_proving_mode() == ChipProvingMode::Sequential {
+    let is_sequential = get_chip_proving_mode() == ChipProvingMode::Sequential;
+    let is_mem_tracking = get_mem_tracking_mode();
+    if is_sequential && is_mem_tracking {
         Some(cuda_hal.inner.mem_tracker(label))
     } else {
         None
     }
 }
 
+const ESTIMATION_TOLERANCE_BYTES: usize = 1024 * 1024; // max estimation error: 1 MB
+const ESTIMATION_SAFETY_MARGIN_BYTES: usize = 5 * 1024 * 1024; // reserved headroom: 5 MB, 1MB for each sub-stage
+
 /// Validate that the estimated GPU memory matches actual usage within tolerance.
 /// - Under-estimate (actual > estimated): diff must be <= `ESTIMATION_TOLERANCE_BYTES`
 /// - Over-estimate (estimated > actual): diff must be <= `ESTIMATION_SAFETY_MARGIN_BYTES`
 #[cfg(feature = "gpu")]
 pub fn check_gpu_mem_estimation(mem_tracker: Option<MemTracker>, estimated_bytes: usize) {
-    // `mem_tracker will` be Some only in sequential mode, so if it's None, do nothing
+    // `mem_tracker will` be Some only in sequential mode with mem tracking enabled, so if it's None, do nothing
     if let Some(mem_tracker) = mem_tracker {
         const ONE_MB: usize = 1024 * 1024;
         let label = mem_tracker.name();

ceno_zkvm/src/scheme/gpu/mod.rs

@@ -720,9 +720,6 @@ where
     E: ExtensionField,
     PCS: PolynomialCommitmentScheme<E>,
 {
-    let (resident, temporary) = estimate_trace_extraction_bytes(expected_num, num_vars);
-    let estimated_bytes = resident + temporary;
-
     let pcs_data_basefold: &BasefoldCommitmentWithWitnessGpu<
         BB31Base,
         BufferImpl<BB31Base>,
@@ -740,7 +737,8 @@ where
         .get_trace(&cuda_hal, pcs_data_basefold, trace_idx, stream.as_ref())
         .unwrap_or_else(|err| panic!("Failed to extract trace {trace_idx}: {err}"));
 
-    check_gpu_mem_estimation(gpu_mem_tracker, estimated_bytes);
+    let (resident, temporary) = estimate_trace_extraction_bytes(expected_num, num_vars);
+    check_gpu_mem_estimation(gpu_mem_tracker, resident + temporary);
 
     let mles: Vec<Arc<MultilinearExtensionGpu<'a, E>>> = poly_group
         .into_iter()