From fd319a6b6a674be813012de6e3dbc36578f7cf97 Mon Sep 17 00:00:00 2001
From: Jonathan Wang <31040440+jonathanpwang@users.noreply.github.com>
Date: Thu, 23 Oct 2025 16:26:44 -0700
Subject: [PATCH 01/10] feat(cli): add `SegmentationArgs` to `prove` command

---
 benchmarks/prove/src/util.rs                  |  2 +-
 crates/cli/src/commands/prove.rs              | 65 +++++++++++++++++--
 .../execution_mode/metered/segment_ctx.rs     |  4 +-
 3 files changed, 63 insertions(+), 8 deletions(-)

diff --git a/benchmarks/prove/src/util.rs b/benchmarks/prove/src/util.rs
index 24a43c890f..7e73d664e6 100644
--- a/benchmarks/prove/src/util.rs
+++ b/benchmarks/prove/src/util.rs
@@ -66,7 +66,7 @@ pub struct BenchmarkCli {
     #[arg(long, alias = "max_segment_length")]
     pub max_segment_length: Option<u32>,
 
-    /// Max cells per chip in segment for continuations
+    /// Total cells used in all chips in segment for continuations
     #[arg(long)]
     pub segment_max_cells: Option<usize>,
 
diff --git a/crates/cli/src/commands/prove.rs b/crates/cli/src/commands/prove.rs
index 8ae0e61a70..ea8b73c210 100644
--- a/crates/cli/src/commands/prove.rs
+++ b/crates/cli/src/commands/prove.rs
@@ -2,9 +2,14 @@ use std::path::PathBuf;
 
 use clap::Parser;
 use eyre::Result;
-use openvm_circuit::arch::instructions::exe::VmExe;
+use openvm_circuit::arch::{
+    execution_mode::metered::segment_ctx::{
+        SegmentationLimits, DEFAULT_MAX_CELLS, DEFAULT_MAX_TRACE_HEIGHT,
+    },
+    instructions::exe::VmExe,
+};
 use openvm_sdk::{
-    config::{AggregationTreeConfig, SdkVmConfig},
+    config::{AggregationTreeConfig, AppConfig, SdkVmConfig},
     fs::{encode_to_file, read_object_from_file, write_to_file_json},
     keygen::AppProvingKey,
     types::VersionedVmStarkProof,
@@ -52,6 +57,9 @@ enum ProveSubCommand {
 
         #[command(flatten)]
         cargo_args: RunCargoArgs,
+
+        #[command(flatten)]
+        segmentation_args: SegmentationArgs,
     },
     Stark {
         #[arg(
@@ -76,6 +84,9 @@ enum ProveSubCommand {
         #[command(flatten)]
         cargo_args: RunCargoArgs,
 
+        #[command(flatten)]
+        segmentation_args: SegmentationArgs,
+
         #[command(flatten)]
         agg_tree_config: AggregationTreeConfig,
     },
@@ -103,11 +114,26 @@ enum ProveSubCommand {
         #[command(flatten)]
         cargo_args: RunCargoArgs,
 
+        #[command(flatten)]
+        segmentation_args: SegmentationArgs,
+
         #[command(flatten)]
         agg_tree_config: AggregationTreeConfig,
     },
 }
 
+#[derive(Clone, Copy, Parser)]
+pub struct SegmentationArgs {
+    /// Trace height threshold per chip for triggering segmentation for continuations in the app
+    /// proof. Note that these thresholds are not absolute limits.
+    #[arg(long, default_value_t = DEFAULT_MAX_TRACE_HEIGHT)]
+    pub segment_max_height: u32,
+    /// Total cells used across all chips for triggering segmentation for continuations in the app
+    /// proof. Note that these thresholds are not absolute limits.
+    #[arg(long, default_value_t = DEFAULT_MAX_CELLS)]
+    pub segment_max_cells: usize,
+}
+
 impl ProveCmd {
     pub fn run(&self) -> Result<()> {
         match &self.command {
@@ -116,9 +142,11 @@ impl ProveCmd {
                 proof,
                 run_args,
                 cargo_args,
+                segmentation_args,
             } => {
                 let app_pk = load_app_pk(app_pk, cargo_args)?;
-                let sdk = Sdk::new(app_pk.app_config())?.with_app_pk(app_pk);
+                let app_config = get_app_config(&app_pk, segmentation_args);
+                let sdk = Sdk::new(app_config)?.with_app_pk(app_pk);
                 let (exe, target_name) = load_or_build_exe(run_args, cargo_args)?;
 
                 let app_proof = sdk
@@ -141,6 +169,7 @@ impl ProveCmd {
                 proof,
                 run_args,
                 cargo_args,
+                segmentation_args,
                 agg_tree_config,
             } => {
                 let app_pk = load_app_pk(app_pk, cargo_args)?;
@@ -149,7 +178,8 @@ impl ProveCmd {
                 let agg_pk = read_object_from_file(default_agg_stark_pk_path()).map_err(|e| {
                     eyre::eyre!("Failed to read aggregation proving key: {}\nPlease run 'cargo openvm setup' first", e)
                 })?;
-                let sdk = Sdk::new(app_pk.app_config())?
+                let app_config = get_app_config(&app_pk, segmentation_args);
+                let sdk = Sdk::new(app_config)?
                     .with_agg_tree_config(*agg_tree_config)
                     .with_app_pk(app_pk)
                     .with_agg_pk(agg_pk);
@@ -178,6 +208,7 @@ impl ProveCmd {
                 proof,
                 run_args,
                 cargo_args,
+                segmentation_args,
                 agg_tree_config,
             } => {
                 let app_pk = load_app_pk(app_pk, cargo_args)?;
@@ -187,7 +218,8 @@ impl ProveCmd {
                 let (agg_pk, halo2_pk) = read_default_agg_and_halo2_pk().map_err(|e| {
                     eyre::eyre!("Failed to read aggregation proving key: {}\nPlease run 'cargo openvm setup' first", e)
                 })?;
-                let sdk = Sdk::new(app_pk.app_config())?
+                let app_config = get_app_config(&app_pk, segmentation_args);
+                let sdk = Sdk::new(app_config)?
                     .with_agg_tree_config(*agg_tree_config)
                     .with_app_pk(app_pk)
                     .with_agg_pk(agg_pk)
@@ -253,3 +285,26 @@ pub(crate) fn load_or_build_exe(
         exe_path.file_stem().unwrap().to_string_lossy().into_owned(),
     ))
 }
+
+fn get_app_config(
+    app_pk: &AppProvingKey<SdkVmConfig>,
+    segmentation_args: &SegmentationArgs,
+) -> AppConfig<SdkVmConfig> {
+    let mut app_config = app_pk.app_config();
+    app_config
+        .app_vm_config
+        .system
+        .config
+        .set_segmentation_limits((*segmentation_args).into());
+    app_config
+}
+
+impl From<SegmentationArgs> for SegmentationLimits {
+    fn from(args: SegmentationArgs) -> Self {
+        SegmentationLimits {
+            max_trace_height: args.segment_max_height,
+            max_cells: args.segment_max_cells,
+            ..Default::default()
+        }
+    }
+}
diff --git a/crates/vm/src/arch/execution_mode/metered/segment_ctx.rs b/crates/vm/src/arch/execution_mode/metered/segment_ctx.rs
index 12dad38ece..1b3a15c000 100644
--- a/crates/vm/src/arch/execution_mode/metered/segment_ctx.rs
+++ b/crates/vm/src/arch/execution_mode/metered/segment_ctx.rs
@@ -5,8 +5,8 @@ use serde::{Deserialize, Serialize};
 
 pub const DEFAULT_SEGMENT_CHECK_INSNS: u64 = 1000;
 
-pub const DEFAULT_MAX_TRACE_HEIGHT: u32 = 1 << 23;
-pub const DEFAULT_MAX_CELLS: usize = 2_000_000_000; // 2B
+pub const DEFAULT_MAX_TRACE_HEIGHT: u32 = 1 << 22;
+pub const DEFAULT_MAX_CELLS: usize = 1_200_000_000; // 1.2B
 const DEFAULT_MAX_INTERACTIONS: usize = BabyBear::ORDER_U32 as usize;
 
 #[derive(derive_new::new, Clone, Debug, Serialize, Deserialize)]

From bbcdb8938a81ba636bb2c07b74eb0a4f84a6e2bb Mon Sep 17 00:00:00 2001
From: Jonathan Wang <31040440+jonathanpwang@users.noreply.github.com>
Date: Thu, 23 Oct 2025 16:30:17 -0700
Subject: [PATCH 02/10] chore: update changelog

---
 CHANGELOG.md | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 48f06435f8..3a2144736d 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,10 +7,14 @@ and this project follows a versioning principles documented in [VERSIONING.md](.
 
 ## v1.4.1 (Unreleased)
 
+### Added
+- (CLI) Add configurable segmentation arguments to `cargo openvm prove` commands.
+
 ### Changed
 - (Toolchain) Update `cargo openvm build` to use Rust nightly version `nightly-2025-08-02`.
 - (Primitives Library) Fix in `openvm-algebra-complex-macros` to ensure `const` byte arrays have proper memory alignment.
 - (Executor) Modified `VirtualMachine::build_metered_ctx` to take the program (`&VmExe<Val<E::SC>>`) as an argument.
+- (Prover/Executor) Changed default segmentation limits.
 
 ## v1.4.0 (2025-09-01)
 

From f115cfadecd385ee082815cd7ca5c62395e176a7 Mon Sep 17 00:00:00 2001
From: Jonathan Wang <31040440+jonathanpwang@users.noreply.github.com>
Date: Thu, 23 Oct 2025 16:30:55 -0700
Subject: [PATCH 03/10] Update crates/cli/src/commands/prove.rs

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
---
 crates/cli/src/commands/prove.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/crates/cli/src/commands/prove.rs b/crates/cli/src/commands/prove.rs
index ea8b73c210..1686829867 100644
--- a/crates/cli/src/commands/prove.rs
+++ b/crates/cli/src/commands/prove.rs
@@ -124,7 +124,7 @@ enum ProveSubCommand {
 
 #[derive(Clone, Copy, Parser)]
 pub struct SegmentationArgs {
-    /// Trace height threshold per chip for triggering segmentation for continuations in the app
+    /// Trace height threshold across all chips for triggering segmentation for continuations in the app
     /// proof. Note that these thresholds are not absolute limits.
     #[arg(long, default_value_t = DEFAULT_MAX_TRACE_HEIGHT)]
     pub segment_max_height: u32,

From 44d647cb6fd4312332d4cac35cde5813fae17da7 Mon Sep 17 00:00:00 2001
From: Jonathan Wang <31040440+jonathanpwang@users.noreply.github.com>
Date: Thu, 23 Oct 2025 16:35:10 -0700
Subject: [PATCH 04/10] wip

---
 crates/cli/src/commands/prove.rs                         | 5 +++--
 crates/vm/src/arch/execution_mode/metered/segment_ctx.rs | 3 ++-
 2 files changed, 5 insertions(+), 3 deletions(-)

diff --git a/crates/cli/src/commands/prove.rs b/crates/cli/src/commands/prove.rs
index 1686829867..23f2630ad4 100644
--- a/crates/cli/src/commands/prove.rs
+++ b/crates/cli/src/commands/prove.rs
@@ -5,6 +5,7 @@ use eyre::Result;
 use openvm_circuit::arch::{
     execution_mode::metered::segment_ctx::{
         SegmentationLimits, DEFAULT_MAX_CELLS, DEFAULT_MAX_TRACE_HEIGHT,
+        DEFAULT_MAX_TRACE_HEIGHT_BITS,
     },
     instructions::exe::VmExe,
 };
@@ -126,8 +127,8 @@ enum ProveSubCommand {
 pub struct SegmentationArgs {
     /// Trace height threshold across all chips for triggering segmentation for continuations in the app
     /// proof. Note that these thresholds are not absolute limits.
-    #[arg(long, default_value_t = DEFAULT_MAX_TRACE_HEIGHT)]
-    pub segment_max_height: u32,
+    #[arg(long, default_value_t = DEFAULT_MAX_TRACE_HEIGHT_BITS)]
+    pub segment_max_height_bits: u8,
     /// Total cells used across all chips for triggering segmentation for continuations in the app
     /// proof. Note that these thresholds are not absolute limits.
     #[arg(long, default_value_t = DEFAULT_MAX_CELLS)]
diff --git a/crates/vm/src/arch/execution_mode/metered/segment_ctx.rs b/crates/vm/src/arch/execution_mode/metered/segment_ctx.rs
index 1b3a15c000..795466cb29 100644
--- a/crates/vm/src/arch/execution_mode/metered/segment_ctx.rs
+++ b/crates/vm/src/arch/execution_mode/metered/segment_ctx.rs
@@ -5,7 +5,8 @@ use serde::{Deserialize, Serialize};
 
 pub const DEFAULT_SEGMENT_CHECK_INSNS: u64 = 1000;
 
-pub const DEFAULT_MAX_TRACE_HEIGHT: u32 = 1 << 22;
+pub const DEFAULT_MAX_TRACE_HEIGHT_BITS: u8 = 22;
+pub const DEFAULT_MAX_TRACE_HEIGHT: u32 = 1 << DEFAULT_MAX_TRACE_HEIGHT_BITS;
 pub const DEFAULT_MAX_CELLS: usize = 1_200_000_000; // 1.2B
 const DEFAULT_MAX_INTERACTIONS: usize = BabyBear::ORDER_U32 as usize;
 

From 902f4c54e10d2355c9c00fadbffe6e3f7a2dd019 Mon Sep 17 00:00:00 2001
From: Jonathan Wang <31040440+jonathanpwang@users.noreply.github.com>
Date: Thu, 23 Oct 2025 16:37:37 -0700
Subject: [PATCH 05/10] chore: switch to using height_bits

---
 crates/cli/src/commands/prove.rs | 11 ++++++-----
 1 file changed, 6 insertions(+), 5 deletions(-)

diff --git a/crates/cli/src/commands/prove.rs b/crates/cli/src/commands/prove.rs
index 23f2630ad4..1b3a87cb23 100644
--- a/crates/cli/src/commands/prove.rs
+++ b/crates/cli/src/commands/prove.rs
@@ -4,8 +4,7 @@ use clap::Parser;
 use eyre::Result;
 use openvm_circuit::arch::{
     execution_mode::metered::segment_ctx::{
-        SegmentationLimits, DEFAULT_MAX_CELLS, DEFAULT_MAX_TRACE_HEIGHT,
-        DEFAULT_MAX_TRACE_HEIGHT_BITS,
+        SegmentationLimits, DEFAULT_MAX_CELLS, DEFAULT_MAX_TRACE_HEIGHT_BITS,
     },
     instructions::exe::VmExe,
 };
@@ -125,8 +124,8 @@ enum ProveSubCommand {
 
 #[derive(Clone, Copy, Parser)]
 pub struct SegmentationArgs {
-    /// Trace height threshold across all chips for triggering segmentation for continuations in the app
-    /// proof. Note that these thresholds are not absolute limits.
+    /// Trace height threshold, in bits, across all chips for triggering segmentation for
+    /// continuations in the app proof. Note that these thresholds are not absolute limits.
     #[arg(long, default_value_t = DEFAULT_MAX_TRACE_HEIGHT_BITS)]
     pub segment_max_height_bits: u8,
     /// Total cells used across all chips for triggering segmentation for continuations in the app
@@ -303,7 +302,9 @@ fn get_app_config(
 impl From<SegmentationArgs> for SegmentationLimits {
     fn from(args: SegmentationArgs) -> Self {
         SegmentationLimits {
-            max_trace_height: args.segment_max_height,
+            max_trace_height: 1u32
+                .checked_shl(args.segment_max_height_bits as u32)
+                .expect("segment_max_height_bits too large"),
             max_cells: args.segment_max_cells,
             ..Default::default()
         }

From 7382274722e036858e38605d44e43fb96232b61c Mon Sep 17 00:00:00 2001
From: Jonathan Wang <31040440+jonathanpwang@users.noreply.github.com>
Date: Thu, 23 Oct 2025 16:48:36 -0700
Subject: [PATCH 06/10] fix: help_heading

---
 crates/cli/src/commands/prove.rs | 12 ++++++++++--
 1 file changed, 10 insertions(+), 2 deletions(-)

diff --git a/crates/cli/src/commands/prove.rs b/crates/cli/src/commands/prove.rs
index 1b3a87cb23..d29beb4128 100644
--- a/crates/cli/src/commands/prove.rs
+++ b/crates/cli/src/commands/prove.rs
@@ -126,11 +126,19 @@ enum ProveSubCommand {
 pub struct SegmentationArgs {
     /// Trace height threshold, in bits, across all chips for triggering segmentation for
     /// continuations in the app proof. Note that these thresholds are not absolute limits.
-    #[arg(long, default_value_t = DEFAULT_MAX_TRACE_HEIGHT_BITS)]
+    #[arg(
+        long,
+        default_value_t = DEFAULT_MAX_TRACE_HEIGHT_BITS,
+        help_heading = "OpenVM Options"
+    )]
     pub segment_max_height_bits: u8,
     /// Total cells used across all chips for triggering segmentation for continuations in the app
     /// proof. Note that these thresholds are not absolute limits.
-    #[arg(long, default_value_t = DEFAULT_MAX_CELLS)]
+    #[arg(
+        long,
+        default_value_t = DEFAULT_MAX_CELLS,
+        help_heading = "OpenVM Options"
+    )]
     pub segment_max_cells: usize,
 }
 

From dc351aedacdaa0fc97d205de79a574f0509608db Mon Sep 17 00:00:00 2001
From: Jonathan Wang <31040440+jonathanpwang@users.noreply.github.com>
Date: Thu, 23 Oct 2025 21:07:46 -0700
Subject: [PATCH 07/10] fix: need to update within app_pk

---
 crates/cli/src/commands/prove.rs | 26 ++++++++++++++------------
 1 file changed, 14 insertions(+), 12 deletions(-)

diff --git a/crates/cli/src/commands/prove.rs b/crates/cli/src/commands/prove.rs
index d29beb4128..a2238a2c17 100644
--- a/crates/cli/src/commands/prove.rs
+++ b/crates/cli/src/commands/prove.rs
@@ -1,4 +1,4 @@
-use std::path::PathBuf;
+use std::{path::PathBuf, sync::Arc};
 
 use clap::Parser;
 use eyre::Result;
@@ -152,8 +152,8 @@ impl ProveCmd {
                 cargo_args,
                 segmentation_args,
             } => {
-                let app_pk = load_app_pk(app_pk, cargo_args)?;
-                let app_config = get_app_config(&app_pk, segmentation_args);
+                let mut app_pk = load_app_pk(app_pk, cargo_args)?;
+                let app_config = get_app_config(&mut app_pk, segmentation_args);
                 let sdk = Sdk::new(app_config)?.with_app_pk(app_pk);
                 let (exe, target_name) = load_or_build_exe(run_args, cargo_args)?;
 
@@ -180,13 +180,13 @@ impl ProveCmd {
                 segmentation_args,
                 agg_tree_config,
             } => {
-                let app_pk = load_app_pk(app_pk, cargo_args)?;
+                let mut app_pk = load_app_pk(app_pk, cargo_args)?;
                 let (exe, target_name) = load_or_build_exe(run_args, cargo_args)?;
 
                 let agg_pk = read_object_from_file(default_agg_stark_pk_path()).map_err(|e| {
                     eyre::eyre!("Failed to read aggregation proving key: {}\nPlease run 'cargo openvm setup' first", e)
                 })?;
-                let app_config = get_app_config(&app_pk, segmentation_args);
+                let app_config = get_app_config(&mut app_pk, segmentation_args);
                 let sdk = Sdk::new(app_config)?
                     .with_agg_tree_config(*agg_tree_config)
                     .with_app_pk(app_pk)
@@ -219,14 +219,14 @@ impl ProveCmd {
                 segmentation_args,
                 agg_tree_config,
             } => {
-                let app_pk = load_app_pk(app_pk, cargo_args)?;
+                let mut app_pk = load_app_pk(app_pk, cargo_args)?;
                 let (exe, target_name) = load_or_build_exe(run_args, cargo_args)?;
 
                 println!("Generating EVM proof, this may take a lot of compute and memory...");
                 let (agg_pk, halo2_pk) = read_default_agg_and_halo2_pk().map_err(|e| {
                     eyre::eyre!("Failed to read aggregation proving key: {}\nPlease run 'cargo openvm setup' first", e)
                 })?;
-                let app_config = get_app_config(&app_pk, segmentation_args);
+                let app_config = get_app_config(&mut app_pk, segmentation_args);
                 let sdk = Sdk::new(app_config)?
                     .with_agg_tree_config(*agg_tree_config)
                     .with_app_pk(app_pk)
@@ -294,17 +294,19 @@ pub(crate) fn load_or_build_exe(
     ))
 }
 
+/// Should only be called when `app_pk` has only a single reference internally.
+/// Mutates the `SystemConfig` within `app_pk` and then returns the updated `AppConfig`.
 fn get_app_config(
-    app_pk: &AppProvingKey<SdkVmConfig>,
+    app_pk: &mut AppProvingKey<SdkVmConfig>,
     segmentation_args: &SegmentationArgs,
 ) -> AppConfig<SdkVmConfig> {
-    let mut app_config = app_pk.app_config();
-    app_config
-        .app_vm_config
+    Arc::get_mut(&mut app_pk.app_vm_pk)
+        .unwrap()
+        .vm_config
         .system
         .config
         .set_segmentation_limits((*segmentation_args).into());
-    app_config
+    app_pk.app_config()
 }
 
 impl From<SegmentationArgs> for SegmentationLimits {

From 27691371341518c7a2c783b5875b5c6dc47cd4d4 Mon Sep 17 00:00:00 2001
From: Jonathan Wang <31040440+jonathanpwang@users.noreply.github.com>
Date: Fri, 24 Oct 2025 17:14:29 -0700
Subject: [PATCH 08/10] chore: debug and warn in segmentation_ctx

---
 crates/cli/src/commands/prove.rs                         | 5 +++--
 crates/vm/src/arch/execution_mode/metered/segment_ctx.rs | 8 ++++++++
 2 files changed, 11 insertions(+), 2 deletions(-)

diff --git a/crates/cli/src/commands/prove.rs b/crates/cli/src/commands/prove.rs
index a2238a2c17..9d0d5dd1bb 100644
--- a/crates/cli/src/commands/prove.rs
+++ b/crates/cli/src/commands/prove.rs
@@ -125,7 +125,8 @@ enum ProveSubCommand {
 #[derive(Clone, Copy, Parser)]
 pub struct SegmentationArgs {
     /// Trace height threshold, in bits, across all chips for triggering segmentation for
-    /// continuations in the app proof. Note that these thresholds are not absolute limits.
+    /// continuations in the app proof. These thresholds are not exceeded except when they are too
+    /// small.
     #[arg(
         long,
         default_value_t = DEFAULT_MAX_TRACE_HEIGHT_BITS,
@@ -133,7 +134,7 @@ pub struct SegmentationArgs {
     )]
     pub segment_max_height_bits: u8,
     /// Total cells used across all chips for triggering segmentation for continuations in the app
-    /// proof. Note that these thresholds are not absolute limits.
+    /// proof. These thresholds are not exceeded except when they are too small.
     #[arg(
         long,
         default_value_t = DEFAULT_MAX_CELLS,
diff --git a/crates/vm/src/arch/execution_mode/metered/segment_ctx.rs b/crates/vm/src/arch/execution_mode/metered/segment_ctx.rs
index 795466cb29..a8004676f0 100644
--- a/crates/vm/src/arch/execution_mode/metered/segment_ctx.rs
+++ b/crates/vm/src/arch/execution_mode/metered/segment_ctx.rs
@@ -100,6 +100,10 @@ impl SegmentationCtx {
     }
 
     pub fn set_max_trace_height(&mut self, max_trace_height: u32) {
+        debug_assert!(
+            max_trace_height.is_power_of_two(),
+            "max_trace_height should be a power of two"
+        );
         self.segmentation_limits.max_trace_height = max_trace_height;
     }
 
@@ -258,6 +262,10 @@ impl SegmentationCtx {
                 self.checkpoint_trace_heights.clone(),
             )
         } else {
+            tracing::warn!(
+                "No valid checkpoint, creating segment using instret={instret}, trace_heights={:?}",
+                trace_heights
+            );
             // No valid checkpoint, use current values
             (instret, trace_heights.to_vec())
         };

From f0b8f2be08e5bb851d41c69eeac34a82b6cef72a Mon Sep 17 00:00:00 2001
From: Jonathan Wang <31040440+jonathanpwang@users.noreply.github.com>
Date: Fri, 24 Oct 2025 17:15:28 -0700
Subject: [PATCH 09/10] chore: bump rc.5

---
 Cargo.toml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/Cargo.toml b/Cargo.toml
index a994989de7..7574c9bcb5 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -1,5 +1,5 @@
 [workspace.package]
-version = "1.4.1-rc.3"
+version = "1.4.1-rc.5"
 edition = "2021"
 rust-version = "1.86.0"
 authors = ["OpenVM Authors"]

From 0aee951f02d3b2ea48feceb83ad03fdc00255418 Mon Sep 17 00:00:00 2001
From: Jonathan Wang <31040440+jonathanpwang@users.noreply.github.com>
Date: Sat, 25 Oct 2025 01:02:44 +0000
Subject: [PATCH 10/10] fix: more explicit stream sync in merkle tree drop

---
 Cargo.lock                          | 122 ++++++++++++++--------------
 crates/vm/src/system/cuda/memory.rs |   6 +-
 2 files changed, 66 insertions(+), 62 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index bfb1f78004..48280eb5a3 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1779,7 +1779,7 @@ dependencies = [
 
 [[package]]
 name = "cargo-openvm"
-version = "1.4.1-rc.3"
+version = "1.4.1-rc.5"
 dependencies = [
  "aws-config",
  "aws-sdk-s3",
@@ -5104,7 +5104,7 @@ checksum = "d05e27ee213611ffe7d6348b942e8f942b37114c00cc03cec254295a4a17852e"
 
 [[package]]
 name = "openvm"
-version = "1.4.1-rc.3"
+version = "1.4.1-rc.5"
 dependencies = [
  "bytemuck",
  "chrono",
@@ -5119,7 +5119,7 @@ dependencies = [
 
 [[package]]
 name = "openvm-algebra-circuit"
-version = "1.4.1-rc.3"
+version = "1.4.1-rc.5"
 dependencies = [
  "cfg-if",
  "derive-new 0.6.0",
@@ -5152,7 +5152,7 @@ dependencies = [
 
 [[package]]
 name = "openvm-algebra-complex-macros"
-version = "1.4.1-rc.3"
+version = "1.4.1-rc.5"
 dependencies = [
  "openvm-macros-common",
  "quote",
@@ -5161,7 +5161,7 @@ dependencies = [
 
 [[package]]
 name = "openvm-algebra-guest"
-version = "1.4.1-rc.3"
+version = "1.4.1-rc.5"
 dependencies = [
  "halo2curves-axiom",
  "num-bigint 0.4.6",
@@ -5176,7 +5176,7 @@ dependencies = [
 
 [[package]]
 name = "openvm-algebra-moduli-macros"
-version = "1.4.1-rc.3"
+version = "1.4.1-rc.5"
 dependencies = [
  "num-bigint 0.4.6",
  "num-prime",
@@ -5187,7 +5187,7 @@ dependencies = [
 
 [[package]]
 name = "openvm-algebra-tests"
-version = "1.4.1-rc.3"
+version = "1.4.1-rc.5"
 dependencies = [
  "eyre",
  "num-bigint 0.4.6",
@@ -5204,7 +5204,7 @@ dependencies = [
 
 [[package]]
 name = "openvm-algebra-transpiler"
-version = "1.4.1-rc.3"
+version = "1.4.1-rc.5"
 dependencies = [
  "openvm-algebra-guest",
  "openvm-instructions",
@@ -5217,7 +5217,7 @@ dependencies = [
 
 [[package]]
 name = "openvm-benchmarks-execute"
-version = "1.4.1-rc.3"
+version = "1.4.1-rc.5"
 dependencies = [
  "bitcode",
  "clap",
@@ -5255,7 +5255,7 @@ dependencies = [
 
 [[package]]
 name = "openvm-benchmarks-prove"
-version = "1.4.1-rc.3"
+version = "1.4.1-rc.5"
 dependencies = [
  "clap",
  "derive_more 1.0.0",
@@ -5281,7 +5281,7 @@ dependencies = [
 
 [[package]]
 name = "openvm-benchmarks-utils"
-version = "1.4.1-rc.3"
+version = "1.4.1-rc.5"
 dependencies = [
  "bitcode",
  "cargo_metadata",
@@ -5303,7 +5303,7 @@ dependencies = [
 
 [[package]]
 name = "openvm-bigint-circuit"
-version = "1.4.1-rc.3"
+version = "1.4.1-rc.5"
 dependencies = [
  "alloy-primitives 1.3.1",
  "cfg-if",
@@ -5330,7 +5330,7 @@ dependencies = [
 
 [[package]]
 name = "openvm-bigint-guest"
-version = "1.4.1-rc.3"
+version = "1.4.1-rc.5"
 dependencies = [
  "openvm-platform",
  "strum_macros",
@@ -5338,7 +5338,7 @@ dependencies = [
 
 [[package]]
 name = "openvm-bigint-transpiler"
-version = "1.4.1-rc.3"
+version = "1.4.1-rc.5"
 dependencies = [
  "openvm-bigint-guest",
  "openvm-instructions",
@@ -5352,7 +5352,7 @@ dependencies = [
 
 [[package]]
 name = "openvm-build"
-version = "1.4.1-rc.3"
+version = "1.4.1-rc.5"
 dependencies = [
  "cargo_metadata",
  "eyre",
@@ -5363,7 +5363,7 @@ dependencies = [
 
 [[package]]
 name = "openvm-circuit"
-version = "1.4.1-rc.3"
+version = "1.4.1-rc.5"
 dependencies = [
  "backtrace",
  "cfg-if",
@@ -5406,7 +5406,7 @@ dependencies = [
 
 [[package]]
 name = "openvm-circuit-derive"
-version = "1.4.1-rc.3"
+version = "1.4.1-rc.5"
 dependencies = [
  "itertools 0.14.0",
  "proc-macro2",
@@ -5416,7 +5416,7 @@ dependencies = [
 
 [[package]]
 name = "openvm-circuit-primitives"
-version = "1.4.1-rc.3"
+version = "1.4.1-rc.5"
 dependencies = [
  "derive-new 0.6.0",
  "itertools 0.14.0",
@@ -5435,7 +5435,7 @@ dependencies = [
 
 [[package]]
 name = "openvm-circuit-primitives-derive"
-version = "1.4.1-rc.3"
+version = "1.4.1-rc.5"
 dependencies = [
  "itertools 0.14.0",
  "quote",
@@ -5444,7 +5444,7 @@ dependencies = [
 
 [[package]]
 name = "openvm-continuations"
-version = "1.4.1-rc.3"
+version = "1.4.1-rc.5"
 dependencies = [
  "derivative",
  "openvm-circuit",
@@ -5522,7 +5522,7 @@ dependencies = [
 
 [[package]]
 name = "openvm-ecc-circuit"
-version = "1.4.1-rc.3"
+version = "1.4.1-rc.5"
 dependencies = [
  "cfg-if",
  "derive-new 0.6.0",
@@ -5554,7 +5554,7 @@ dependencies = [
 
 [[package]]
 name = "openvm-ecc-guest"
-version = "1.4.1-rc.3"
+version = "1.4.1-rc.5"
 dependencies = [
  "ecdsa 0.16.9",
  "elliptic-curve 0.13.8",
@@ -5572,7 +5572,7 @@ dependencies = [
 
 [[package]]
 name = "openvm-ecc-integration-tests"
-version = "1.4.1-rc.3"
+version = "1.4.1-rc.5"
 dependencies = [
  "eyre",
  "halo2curves-axiom",
@@ -5594,7 +5594,7 @@ dependencies = [
 
 [[package]]
 name = "openvm-ecc-sw-macros"
-version = "1.4.1-rc.3"
+version = "1.4.1-rc.5"
 dependencies = [
  "openvm-macros-common",
  "quote",
@@ -5603,7 +5603,7 @@ dependencies = [
 
 [[package]]
 name = "openvm-ecc-transpiler"
-version = "1.4.1-rc.3"
+version = "1.4.1-rc.5"
 dependencies = [
  "openvm-ecc-guest",
  "openvm-instructions",
@@ -5616,7 +5616,7 @@ dependencies = [
 
 [[package]]
 name = "openvm-ff-derive"
-version = "1.4.1-rc.3"
+version = "1.4.1-rc.5"
 dependencies = [
  "addchain",
  "eyre",
@@ -5639,7 +5639,7 @@ dependencies = [
 
 [[package]]
 name = "openvm-instructions"
-version = "1.4.1-rc.3"
+version = "1.4.1-rc.5"
 dependencies = [
  "backtrace",
  "bitcode",
@@ -5659,7 +5659,7 @@ dependencies = [
 
 [[package]]
 name = "openvm-instructions-derive"
-version = "1.4.1-rc.3"
+version = "1.4.1-rc.5"
 dependencies = [
  "openvm-instructions",
  "quote",
@@ -5670,7 +5670,7 @@ dependencies = [
 
 [[package]]
 name = "openvm-keccak256"
-version = "1.4.1-rc.3"
+version = "1.4.1-rc.5"
 dependencies = [
  "eyre",
  "openvm-circuit",
@@ -5687,7 +5687,7 @@ dependencies = [
 
 [[package]]
 name = "openvm-keccak256-circuit"
-version = "1.4.1-rc.3"
+version = "1.4.1-rc.5"
 dependencies = [
  "cfg-if",
  "derive-new 0.6.0",
@@ -5716,14 +5716,14 @@ dependencies = [
 
 [[package]]
 name = "openvm-keccak256-guest"
-version = "1.4.1-rc.3"
+version = "1.4.1-rc.5"
 dependencies = [
  "openvm-platform",
 ]
 
 [[package]]
 name = "openvm-keccak256-transpiler"
-version = "1.4.1-rc.3"
+version = "1.4.1-rc.5"
 dependencies = [
  "openvm-instructions",
  "openvm-instructions-derive",
@@ -5736,14 +5736,14 @@ dependencies = [
 
 [[package]]
 name = "openvm-macros-common"
-version = "1.4.1-rc.3"
+version = "1.4.1-rc.5"
 dependencies = [
  "syn 2.0.106",
 ]
 
 [[package]]
 name = "openvm-mod-circuit-builder"
-version = "1.4.1-rc.3"
+version = "1.4.1-rc.5"
 dependencies = [
  "cuda-runtime-sys",
  "halo2curves-axiom",
@@ -5765,7 +5765,7 @@ dependencies = [
 
 [[package]]
 name = "openvm-native-circuit"
-version = "1.4.1-rc.3"
+version = "1.4.1-rc.5"
 dependencies = [
  "cfg-if",
  "derive-new 0.6.0",
@@ -5800,7 +5800,7 @@ dependencies = [
 
 [[package]]
 name = "openvm-native-compiler"
-version = "1.4.1-rc.3"
+version = "1.4.1-rc.5"
 dependencies = [
  "backtrace",
  "itertools 0.14.0",
@@ -5823,7 +5823,7 @@ dependencies = [
 
 [[package]]
 name = "openvm-native-compiler-derive"
-version = "1.4.1-rc.3"
+version = "1.4.1-rc.5"
 dependencies = [
  "quote",
  "syn 2.0.106",
@@ -5831,7 +5831,7 @@ dependencies = [
 
 [[package]]
 name = "openvm-native-recursion"
-version = "1.4.1-rc.3"
+version = "1.4.1-rc.5"
 dependencies = [
  "bitcode",
  "cfg-if",
@@ -5861,7 +5861,7 @@ dependencies = [
 
 [[package]]
 name = "openvm-native-transpiler"
-version = "1.4.1-rc.3"
+version = "1.4.1-rc.5"
 dependencies = [
  "openvm-instructions",
  "openvm-transpiler",
@@ -5870,7 +5870,7 @@ dependencies = [
 
 [[package]]
 name = "openvm-pairing"
-version = "1.4.1-rc.3"
+version = "1.4.1-rc.5"
 dependencies = [
  "eyre",
  "group 0.13.0",
@@ -5908,7 +5908,7 @@ dependencies = [
 
 [[package]]
 name = "openvm-pairing-circuit"
-version = "1.4.1-rc.3"
+version = "1.4.1-rc.5"
 dependencies = [
  "cfg-if",
  "derive-new 0.6.0",
@@ -5938,7 +5938,7 @@ dependencies = [
 
 [[package]]
 name = "openvm-pairing-guest"
-version = "1.4.1-rc.3"
+version = "1.4.1-rc.5"
 dependencies = [
  "halo2curves-axiom",
  "hex-literal 0.4.1",
@@ -5959,7 +5959,7 @@ dependencies = [
 
 [[package]]
 name = "openvm-pairing-transpiler"
-version = "1.4.1-rc.3"
+version = "1.4.1-rc.5"
 dependencies = [
  "openvm-instructions",
  "openvm-pairing-guest",
@@ -5971,7 +5971,7 @@ dependencies = [
 
 [[package]]
 name = "openvm-platform"
-version = "1.4.1-rc.3"
+version = "1.4.1-rc.5"
 dependencies = [
  "critical-section",
  "embedded-alloc",
@@ -5982,7 +5982,7 @@ dependencies = [
 
 [[package]]
 name = "openvm-poseidon2-air"
-version = "1.4.1-rc.3"
+version = "1.4.1-rc.5"
 dependencies = [
  "derivative",
  "lazy_static",
@@ -6001,7 +6001,7 @@ dependencies = [
 
 [[package]]
 name = "openvm-prof"
-version = "1.4.1-rc.3"
+version = "1.4.1-rc.5"
 dependencies = [
  "clap",
  "eyre",
@@ -6014,7 +6014,7 @@ dependencies = [
 
 [[package]]
 name = "openvm-rv32-adapters"
-version = "1.4.1-rc.3"
+version = "1.4.1-rc.5"
 dependencies = [
  "derive-new 0.6.0",
  "itertools 0.14.0",
@@ -6030,7 +6030,7 @@ dependencies = [
 
 [[package]]
 name = "openvm-rv32im-circuit"
-version = "1.4.1-rc.3"
+version = "1.4.1-rc.5"
 dependencies = [
  "cfg-if",
  "derive-new 0.6.0",
@@ -6057,7 +6057,7 @@ dependencies = [
 
 [[package]]
 name = "openvm-rv32im-guest"
-version = "1.4.1-rc.3"
+version = "1.4.1-rc.5"
 dependencies = [
  "openvm-custom-insn",
  "p3-field",
@@ -6066,7 +6066,7 @@ dependencies = [
 
 [[package]]
 name = "openvm-rv32im-integration-tests"
-version = "1.4.1-rc.3"
+version = "1.4.1-rc.5"
 dependencies = [
  "eyre",
  "openvm",
@@ -6085,7 +6085,7 @@ dependencies = [
 
 [[package]]
 name = "openvm-rv32im-transpiler"
-version = "1.4.1-rc.3"
+version = "1.4.1-rc.5"
 dependencies = [
  "openvm-instructions",
  "openvm-instructions-derive",
@@ -6100,7 +6100,7 @@ dependencies = [
 
 [[package]]
 name = "openvm-scripts"
-version = "1.4.1-rc.3"
+version = "1.4.1-rc.5"
 dependencies = [
  "clap",
  "eyre",
@@ -6112,7 +6112,7 @@ dependencies = [
 
 [[package]]
 name = "openvm-sdk"
-version = "1.4.1-rc.3"
+version = "1.4.1-rc.5"
 dependencies = [
  "alloy-sol-types",
  "bitcode",
@@ -6171,7 +6171,7 @@ dependencies = [
 
 [[package]]
 name = "openvm-sha2"
-version = "1.4.1-rc.3"
+version = "1.4.1-rc.5"
 dependencies = [
  "eyre",
  "openvm-circuit",
@@ -6188,7 +6188,7 @@ dependencies = [
 
 [[package]]
 name = "openvm-sha256-air"
-version = "1.4.1-rc.3"
+version = "1.4.1-rc.5"
 dependencies = [
  "openvm-circuit",
  "openvm-circuit-primitives",
@@ -6200,7 +6200,7 @@ dependencies = [
 
 [[package]]
 name = "openvm-sha256-circuit"
-version = "1.4.1-rc.3"
+version = "1.4.1-rc.5"
 dependencies = [
  "cfg-if",
  "derive-new 0.6.0",
@@ -6226,14 +6226,14 @@ dependencies = [
 
 [[package]]
 name = "openvm-sha256-guest"
-version = "1.4.1-rc.3"
+version = "1.4.1-rc.5"
 dependencies = [
  "openvm-platform",
 ]
 
 [[package]]
 name = "openvm-sha256-transpiler"
-version = "1.4.1-rc.3"
+version = "1.4.1-rc.5"
 dependencies = [
  "openvm-instructions",
  "openvm-instructions-derive",
@@ -6313,7 +6313,7 @@ dependencies = [
 
 [[package]]
 name = "openvm-toolchain-tests"
-version = "1.4.1-rc.3"
+version = "1.4.1-rc.5"
 dependencies = [
  "derive_more 1.0.0",
  "eyre",
@@ -6339,7 +6339,7 @@ dependencies = [
 
 [[package]]
 name = "openvm-transpiler"
-version = "1.4.1-rc.3"
+version = "1.4.1-rc.5"
 dependencies = [
  "elf",
  "eyre",
@@ -6353,7 +6353,7 @@ dependencies = [
 
 [[package]]
 name = "openvm-verify-stark"
-version = "1.4.1-rc.3"
+version = "1.4.1-rc.5"
 dependencies = [
  "eyre",
  "openvm-circuit",
diff --git a/crates/vm/src/system/cuda/memory.rs b/crates/vm/src/system/cuda/memory.rs
index 297d7c8624..8c1d64558c 100644
--- a/crates/vm/src/system/cuda/memory.rs
+++ b/crates/vm/src/system/cuda/memory.rs
@@ -234,8 +234,12 @@ impl MemoryInventoryGPU {
 
 impl Drop for PersistentMemoryInventoryGPU {
     fn drop(&mut self) {
-        // Drop merkle subtrees first so their individual streams synchronize before dropping the
+        // Force synchronize all streams in merkle tree before dropping the
         // initial memory buffers. This prevents buffers from dropping before build_async completes.
+        for s in &self.merkle_tree.subtrees {
+            s.stream.synchronize().unwrap();
+        }
+        self.merkle_tree.stream.synchronize().unwrap();
         self.merkle_tree.drop_subtrees();
         self.initial_memory.clear();
     }