Skip to content

Commit a76a9d0

Browse files
fix(cuda): sync memory subtree streams before drop (#2175)
In some edge case where right after we start `build_async` on the memory merkle subtrees, if the program panics, then the order of drop could be that we drop the `initial_memory` buffers on the default stream first, while the `build_async` kernels are still running and using those buffers. This leads to a deadloop. I fixed it by just forcing the drop to drop subtrees first (which should sync their special streams) before dropping `initial_memory`. compare: https://github.com/axiom-crypto/openvm-reth-benchmark/actions/runs/18733111153
1 parent 2a63e2a commit a76a9d0

File tree

1 file changed

+9
-0
lines changed

1 file changed

+9
-0
lines changed

crates/vm/src/system/cuda/memory.rs

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -231,3 +231,12 @@ impl MemoryInventoryGPU {
231231
ret
232232
}
233233
}
234+
235+
impl Drop for PersistentMemoryInventoryGPU {
236+
fn drop(&mut self) {
237+
// Drop merkle subtrees first so their individual streams synchronize before dropping the
238+
// initial memory buffers. This prevents buffers from dropping before build_async completes.
239+
self.merkle_tree.drop_subtrees();
240+
self.initial_memory.clear();
241+
}
242+
}

0 commit comments

Comments
 (0)