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.

Author: theonlychant

reproducer.py

@@ -0,0 +1,60 @@
+import warp as wp
+
+# CPU-mode reproducer adapted from your CUDA snippet
+
+def main():
+    device = wp.get_device("cpu")
+    try:
+        wp.load_module(device=device)
+    except Exception:
+        # load_module may be unnecessary on CPU; ignore failures
+        pass
+
+    @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,49 @@ 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. Use the caller stream's 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
+            // global graph leaf nodes if capture info isn't available for the caller.
             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)) {
-                cudaGraphNode_t free_node;
-                if (check_cuda(cudaGraphAddMemFreeNode(&free_node, graph, leaf_nodes.data(), leaf_nodes.size(), ptr))) {
-                    check_cu(cuStreamUpdateCaptureDependencies_f(
-                        capture->stream, &free_node, 1, cudaStreamSetCaptureDependencies
-                    ));
+
+            // get the caller stream (the stream on which wp_free_device_async was invoked)
+            CUstream caller_cuda_stream = get_current_stream();
+
+            const cudaGraphNode_t* capture_deps = nullptr;
+            size_t dep_count = 0;
+            CUstreamCaptureStatus capture_status = CU_STREAM_CAPTURE_STATUS_NONE;
+
+            bool added = false;
+
+            // Try to get per-stream capture dependencies for the caller stream and use
+            // them as predecessors for the memfree node. This ensures the memfree will
+            // be ordered after work recorded on the caller stream (including forked
+            // substreams brought into the capture via wait_stream/wait_event).
+            if (graph && check_cu(cuStreamGetCaptureInfo_f(caller_cuda_stream, &capture_status, nullptr, &graph, &capture_deps, &dep_count))) {
+                if (graph && (capture_deps != nullptr || dep_count > 0)) {
+                    cudaGraphNode_t free_node;
+                    if (check_cuda(cudaGraphAddMemFreeNode(&free_node, graph, capture_deps, dep_count, ptr))) {
+                        check_cu(cuStreamUpdateCaptureDependencies_f(
+                            caller_cuda_stream, &free_node, 1, cudaStreamSetCaptureDependencies
+                        ));
+                        added = true;
+                    }
+                }
+            }
+
+            // Fallback: if we couldn't obtain per-stream capture dependencies, use
+            // the graph leaf nodes as before.
+            if (!added && graph) {
+                std::vector<cudaGraphNode_t> leaf_nodes;
+                if (get_graph_leaf_nodes(graph, leaf_nodes)) {
+                    cudaGraphNode_t free_node;
+                    if (check_cuda(cudaGraphAddMemFreeNode(&free_node, graph, leaf_nodes.data(), leaf_nodes.size(), ptr))) {
+                        check_cu(cuStreamUpdateCaptureDependencies_f(
+                            capture->stream, &free_node, 1, cudaStreamSetCaptureDependencies
+                        ));
+                    }
                 }
             }