fix(graph): correct MEM_FREE node dependencies for substream allocations

Captured cudaGraphAddMemFreeNode was inheriting the begin-stream's
frontier instead of the substream's, causing use-after-free on replay.

Also adds reproducer.py with CPU fallback for GPU-less environments.

Signed-off-by: theonlychant <sacehenry@gmail.com>

Author: theonlychant

reproducer.py

@@ -0,0 +1,67 @@
+import warp as wp
+
+# CPU-mode reproducer adapted from your CUDA snippet
+
+
+def main():
+    # Try CUDA first, fall back to CPU
+    try:
+        device = wp.get_device("cuda:0")
+    except Exception:
+        device = wp.get_device("cpu")
+
+    try:
+        wp.load_module(device=device)
+    except Exception:
+        if not device.is_cpu:
+            raise
+        # CPU load_module failures are expected, ignore
+
+    @wp.kernel
+    def touch(x: wp.array(dtype=wp.float32)):
+        i = wp.tid()
+        if i < x.shape[0]:
+            x[i] = x[i] + 1.0
+
+    # CPU devices do not have CUDA streams; run a CPU-friendly capture path.
+    if device.is_cpu:
+        wp.capture_begin(device=device, force_module_load=False)
+        try:
+            for _ in range(4):
+                t = wp.empty(4096, dtype=wp.float32, device=device)
+                wp.launch(touch, dim=4096, inputs=[t])
+                del t
+        finally:
+            g = wp.capture_end(device=device)
+
+        for _ in range(8):
+            wp.capture_launch(g)
+    else:
+        main_stream = wp.get_stream(device)
+        sub_stream = wp.Stream(device)
+
+        wp.capture_begin(device=device, stream=main_stream, force_module_load=False)
+        try:
+            sub_stream.wait_stream(main_stream)
+            with wp.ScopedStream(sub_stream, sync_enter=False):
+                for _ in range(4):
+                    t = wp.empty(4096, dtype=wp.float32, device=device)
+                    wp.launch(touch, dim=4096, inputs=[t], stream=sub_stream)
+                    del t
+            main_stream.wait_stream(sub_stream)
+        finally:
+            g = wp.capture_end(device=device, stream=main_stream)
+
+        replay = wp.Stream(device)
+        for _ in range(8):
+            wp.capture_launch(g, stream=replay)
+        wp.synchronize_stream(replay)
+
+
+if __name__ == "__main__":
+    try:
+        main()
+        print("Reproducer finished successfully")
+    except Exception as e:
+        print("Reproducer failed:", e)
+        raise

warp/native/warp.cu

@@ -827,18 +827,51 @@ void wp_free_device_async(void* context, void* ptr)
         // check if the capture is still active
         auto capture_iter = g_captures.find(capture_id);
         if (capture_iter != g_captures.end()) {
-            // Add a mem free node.  Use all current leaf nodes as dependencies to ensure that all prior
-            // work completes before deallocating.  This works with both Warp-initiated and external captures
-            // and avoids the need to explicitly track all streams used during the capture.
+            // Add a mem free node. Prefer per-caller-stream capture dependencies so frees
+            // are ordered with respect to work recorded on the stream where the free
+            // occurs (handles forked substreams correctly).  Fall back to using the
+            // graph's leaf nodes if per-stream info isn't available.
             CaptureInfo* capture = capture_iter->second;
-            cudaGraph_t graph = get_capture_graph(capture->stream);
-            std::vector<cudaGraphNode_t> leaf_nodes;
-            if (graph && get_graph_leaf_nodes(graph, leaf_nodes)) {
+            cudaGraph_t begin_graph = get_capture_graph(capture->stream);
+
+            // get the caller stream (the stream on which wp_free_device_async was invoked)
+            CUstream caller_cuda_stream = get_current_stream();
+
+            cudaGraph_t caller_graph = nullptr;
+            const cudaGraphNode_t* capture_deps = nullptr;
+            size_t dep_count = 0;
+            CUstreamCaptureStatus capture_status = CU_STREAM_CAPTURE_STATUS_NONE;
+
+            bool added = false;
+
+            // Query per-stream capture info into separate variables so we don't
+            // clobber the begin_graph value used for the fallback.
+            if (begin_graph
+                && check_cu(cuStreamGetCaptureInfo_f(
+                    caller_cuda_stream, &capture_status, nullptr, &caller_graph, &capture_deps, &dep_count
+                ))
+                && caller_graph && capture_status == CU_STREAM_CAPTURE_STATUS_ACTIVE) {
                 cudaGraphNode_t free_node;
-                if (check_cuda(cudaGraphAddMemFreeNode(&free_node, graph, leaf_nodes.data(), leaf_nodes.size(), ptr))) {
+                if (check_cuda(cudaGraphAddMemFreeNode(&free_node, caller_graph, capture_deps, dep_count, ptr))) {
                     check_cu(cuStreamUpdateCaptureDependencies_f(
-                        capture->stream, &free_node, 1, cudaStreamSetCaptureDependencies
+                        caller_cuda_stream, &free_node, 1, cudaStreamSetCaptureDependencies
                     ));
+                    added = true;
+                }
+            }
+
+            // Fallback: use the graph leaf nodes from the original capture if needed
+            if (!added && begin_graph) {
+                std::vector<cudaGraphNode_t> leaf_nodes;
+                if (get_graph_leaf_nodes(begin_graph, leaf_nodes)) {
+                    cudaGraphNode_t free_node;
+                    if (check_cuda(
+                            cudaGraphAddMemFreeNode(&free_node, begin_graph, leaf_nodes.data(), leaf_nodes.size(), ptr)
+                        )) {
+                        check_cu(cuStreamUpdateCaptureDependencies_f(
+                            capture->stream, &free_node, 1, cudaStreamSetCaptureDependencies
+                        ));
+                    }
                 }
             }