Add CPU callbacks for stream capture & Cythonize GraphBuilder (#1814)

* Cythonize _graph/_graph_builder (move from pure Python to .pyx)

Move the GraphBuilder/Graph/GraphCompleteOptions/GraphDebugPrintOptions
implementation out of _graph/__init__.py into _graph/_graph_builder.pyx
so it is compiled by Cython. A thin __init__.py re-exports the public
names so all existing import sites continue to work unchanged.

Cython compatibility adjustments:
- Remove `from __future__ import annotations` (unsupported by Cython)
- Remove TYPE_CHECKING guard; quote annotations that reference Stream
  (circular import), forward-reference GraphBuilder/Graph, or use
  X | None union syntax
- Update _graphdef.pyx lazy imports to point directly at _graph_builder

No build_hooks.py changes needed — the build system auto-discovers .pyx
files via glob.

Ref: #1076
Made-with: Cursor

* Remove _lazy_init from _graph_builder; add cached get_driver_version

Replace the per-module _lazy_init / _inited / _driver_ver / _py_major_minor
pattern in _graph_builder.pyx with direct calls to centralized cached
functions in cuda_utils:

- Add get_driver_version() with @functools.cache alongside get_binding_version
- Switch get_binding_version from @functools.lru_cache to @functools.cache
  (cleaner for nullary functions)
- Fix split() to return tuple(result) — Cython enforces return type
  annotations unlike pure Python
- Fix _cond_with_params annotation from -> GraphBuilder to -> tuple
  to match actual return value

Made-with: Cursor

* Add CPU callbacks for stream capture (GraphBuilder.callback)

Implements #1328: host callbacks during stream capture via
cuLaunchHostFunc, mirroring the existing GraphDef.callback API.

Extracts shared callback infrastructure (_attach_user_object,
_attach_host_callback_to_graph, trampoline/destructor) into a new
_graph/_utils.pyx module to avoid circular imports between
_graph_builder and _graphdef.

Made-with: Cursor

Author: Andy-Jost

cuda_core/cuda/core/_graph/__init__.py

@@ -30,8 +30,6 @@ GraphNode hierarchy:
 
 from __future__ import annotations
 
-from cpython.ref cimport Py_INCREF
-
 from libc.stddef cimport size_t
 from libc.stdint cimport uintptr_t
 from libc.stdlib cimport malloc, free
@@ -102,16 +100,11 @@ cdef bint _check_node_get_params():
     return _has_cuGraphNodeGetParams
 
 
-cdef extern from "Python.h":
-    void _py_decref "Py_DECREF" (void*)
-
-
-cdef void _py_host_trampoline(void* data) noexcept with gil:
-    (<object>data)()
-
-
-cdef void _py_host_destructor(void* data) noexcept with gil:
-    _py_decref(data)
+from cuda.core._graph._utils cimport (
+    _attach_host_callback_to_graph,
+    _attach_user_object,
+    _is_py_host_trampoline,
+)
 
 
 cdef void _destroy_event_handle_copy(void* ptr) noexcept nogil:
@@ -124,30 +117,6 @@ cdef void _destroy_kernel_handle_copy(void* ptr) noexcept nogil:
     del p
 
 
-cdef void _attach_user_object(
-        cydriver.CUgraph graph, void* ptr,
-        cydriver.CUhostFn destroy) except *:
-    """Create a CUDA user object and transfer ownership to the graph.
-
-    On success the graph owns the resource (via MOVE semantics).
-    On failure the destroy callback is invoked to clean up ptr,
-    then a CUDAError is raised — callers need no try/except.
-    """
-    cdef cydriver.CUuserObject user_obj = NULL
-    cdef cydriver.CUresult ret
-    with nogil:
-        ret = cydriver.cuUserObjectCreate(
-            &user_obj, ptr, destroy, 1,
-            cydriver.CU_USER_OBJECT_NO_DESTRUCTOR_SYNC)
-        if ret == cydriver.CUDA_SUCCESS:
-            ret = cydriver.cuGraphRetainUserObject(
-                graph, user_obj, 1, cydriver.CU_GRAPH_USER_OBJECT_MOVE)
-            if ret != cydriver.CUDA_SUCCESS:
-                cydriver.cuUserObjectRelease(user_obj, 1)
-    if ret != cydriver.CUDA_SUCCESS:
-        if user_obj == NULL:
-            destroy(ptr)
-        HANDLE_RETURN(ret)
 
 
 cdef class Condition:
@@ -470,7 +439,7 @@ cdef class GraphDef:
         Graph
             An executable graph that can be launched on a stream.
         """
-        from cuda.core._graph import _instantiate_graph
+        from cuda.core._graph._graph_builder import _instantiate_graph
 
         return _instantiate_graph(
             driver.CUgraph(as_intptr(self._h_graph)), options)
@@ -485,7 +454,7 @@ cdef class GraphDef:
         options : GraphDebugPrintOptions, optional
             Customizable options for the debug print.
         """
-        from cuda.core._graph import GraphDebugPrintOptions
+        from cuda.core._graph._graph_builder import GraphDebugPrintOptions
 
         cdef unsigned int flags = 0
         if options is not None:
@@ -1270,56 +1239,20 @@ cdef class GraphNode:
         cdef cydriver.CUgraphNode pred_node = as_cu(self._h_node)
         cdef cydriver.CUgraphNode* deps = NULL
         cdef size_t num_deps = 0
-        cdef void* c_user_data = NULL
-        cdef object callable_obj = None
-        cdef void* fn_pyobj = NULL
 
         if pred_node != NULL:
             deps = &pred_node
             num_deps = 1
 
-        if isinstance(fn, ct._CFuncPtr):
-            Py_INCREF(fn)
-            fn_pyobj = <void*>fn
-            _attach_user_object(
-                as_cu(h_graph), fn_pyobj,
-                <cydriver.CUhostFn>_py_host_destructor)
-            node_params.fn = <cydriver.CUhostFn><uintptr_t>ct.cast(
-                fn, ct.c_void_p).value
-
-            if user_data is not None:
-                if isinstance(user_data, int):
-                    c_user_data = <void*><uintptr_t>user_data
-                else:
-                    buf = bytes(user_data)
-                    c_user_data = malloc(len(buf))
-                    if c_user_data == NULL:
-                        raise MemoryError(
-                            "failed to allocate user_data buffer")
-                    c_memcpy(c_user_data, <const char*>buf, len(buf))
-                    _attach_user_object(
-                        as_cu(h_graph), c_user_data,
-                        <cydriver.CUhostFn>free)
-
-            node_params.userData = c_user_data
-        else:
-            if user_data is not None:
-                raise ValueError(
-                    "user_data is only supported with ctypes "
-                    "function pointers")
-            callable_obj = fn
-            Py_INCREF(fn)
-            fn_pyobj = <void*>fn
-            node_params.fn = <cydriver.CUhostFn>_py_host_trampoline
-            node_params.userData = fn_pyobj
-            _attach_user_object(
-                as_cu(h_graph), fn_pyobj,
-                <cydriver.CUhostFn>_py_host_destructor)
+        _attach_host_callback_to_graph(
+            as_cu(h_graph), fn, user_data,
+            &node_params.fn, &node_params.userData)
 
         with nogil:
             HANDLE_RETURN(cydriver.cuGraphAddHostNode(
                 &new_node, as_cu(h_graph), deps, num_deps, &node_params))
 
+        cdef object callable_obj = fn if not isinstance(fn, ct._CFuncPtr) else None
         self._succ_cache = None
         return HostCallbackNode._create_with_params(
             create_graph_node_handle(new_node, h_graph), callable_obj,
@@ -1947,7 +1880,7 @@ cdef class HostCallbackNode(GraphNode):
             HANDLE_RETURN(cydriver.cuGraphHostNodeGetParams(node, &params))
 
         cdef object callable_obj = None
-        if params.fn == <cydriver.CUhostFn>_py_host_trampoline:
+        if _is_py_host_trampoline(params.fn):
             callable_obj = <object>params.userData
 
         return HostCallbackNode._create_with_params(

cuda_core/cuda/core/_graph/_graph_builder.pyx

@@ -0,0 +1,16 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from cuda.bindings cimport cydriver
+
+
+cdef bint _is_py_host_trampoline(cydriver.CUhostFn fn) noexcept nogil
+
+cdef void _attach_user_object(
+    cydriver.CUgraph graph, void* ptr,
+    cydriver.CUhostFn destroy) except *
+
+cdef void _attach_host_callback_to_graph(
+    cydriver.CUgraph graph, object fn, object user_data,
+    cydriver.CUhostFn* out_fn, void** out_user_data) except *

cuda_core/cuda/core/_graph/_graphdef.pyx

@@ -0,0 +1,106 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from cpython.ref cimport Py_INCREF
+
+from libc.stdint cimport uintptr_t
+from libc.stdlib cimport malloc, free
+from libc.string cimport memcpy as c_memcpy
+
+from cuda.bindings cimport cydriver
+
+from cuda.core._utils.cuda_utils cimport HANDLE_RETURN
+
+
+cdef extern from "Python.h":
+    void _py_decref "Py_DECREF" (void*)
+
+
+cdef void _py_host_trampoline(void* data) noexcept with gil:
+    (<object>data)()
+
+
+cdef void _py_host_destructor(void* data) noexcept with gil:
+    _py_decref(data)
+
+
+cdef bint _is_py_host_trampoline(cydriver.CUhostFn fn) noexcept nogil:
+    return fn == <cydriver.CUhostFn>_py_host_trampoline
+
+
+cdef void _attach_user_object(
+        cydriver.CUgraph graph, void* ptr,
+        cydriver.CUhostFn destroy) except *:
+    """Create a CUDA user object and transfer ownership to the graph.
+
+    On success the graph owns the resource (via MOVE semantics).
+    On failure the destroy callback is invoked to clean up ptr,
+    then a CUDAError is raised — callers need no try/except.
+    """
+    cdef cydriver.CUuserObject user_obj = NULL
+    cdef cydriver.CUresult ret
+    with nogil:
+        ret = cydriver.cuUserObjectCreate(
+            &user_obj, ptr, destroy, 1,
+            cydriver.CU_USER_OBJECT_NO_DESTRUCTOR_SYNC)
+        if ret == cydriver.CUDA_SUCCESS:
+            ret = cydriver.cuGraphRetainUserObject(
+                graph, user_obj, 1, cydriver.CU_GRAPH_USER_OBJECT_MOVE)
+            if ret != cydriver.CUDA_SUCCESS:
+                cydriver.cuUserObjectRelease(user_obj, 1)
+    if ret != cydriver.CUDA_SUCCESS:
+        if user_obj == NULL:
+            destroy(ptr)
+        HANDLE_RETURN(ret)
+
+
+cdef void _attach_host_callback_to_graph(
+        cydriver.CUgraph graph, object fn, object user_data,
+        cydriver.CUhostFn* out_fn, void** out_user_data) except *:
+    """Resolve a Python callable or ctypes CFuncPtr into a C callback pair.
+
+    Handles Py_INCREF, user-object attachment for lifetime management,
+    and user_data copying.  On return, *out_fn and *out_user_data are
+    ready to pass to cuGraphAddHostNode or cuLaunchHostFunc.
+    """
+    import ctypes as ct
+
+    cdef void* fn_pyobj = NULL
+
+    if isinstance(fn, ct._CFuncPtr):
+        Py_INCREF(fn)
+        fn_pyobj = <void*>fn
+        _attach_user_object(
+            graph, fn_pyobj,
+            <cydriver.CUhostFn>_py_host_destructor)
+        out_fn[0] = <cydriver.CUhostFn><uintptr_t>ct.cast(
+            fn, ct.c_void_p).value
+
+        if user_data is not None:
+            if isinstance(user_data, int):
+                out_user_data[0] = <void*><uintptr_t>user_data
+            else:
+                buf = bytes(user_data)
+                out_user_data[0] = malloc(len(buf))
+                if out_user_data[0] == NULL:
+                    raise MemoryError(
+                        "failed to allocate user_data buffer")
+                c_memcpy(out_user_data[0], <const char*>buf, len(buf))
+                _attach_user_object(
+                    graph, out_user_data[0],
+                    <cydriver.CUhostFn>free)
+        else:
+            out_user_data[0] = NULL
+    else:
+        if user_data is not None:
+            raise ValueError(
+                "user_data is only supported with ctypes "
+                "function pointers")
+        Py_INCREF(fn)
+        fn_pyobj = <void*>fn
+        out_fn[0] = <cydriver.CUhostFn>_py_host_trampoline
+        out_user_data[0] = fn_pyobj
+        _attach_user_object(
+            graph, fn_pyobj,
+            <cydriver.CUhostFn>_py_host_destructor)

cuda_core/cuda/core/_graph/_utils.pxd

@@ -298,14 +298,18 @@ def is_nested_sequence(obj):
     return is_sequence(obj) and any(is_sequence(elem) for elem in obj)
 
 
-@functools.lru_cache
+@functools.cache
 def get_binding_version():
     try:
         major_minor = importlib.metadata.version("cuda-bindings").split(".")[:2]
     except importlib.metadata.PackageNotFoundError:
         major_minor = importlib.metadata.version("cuda-python").split(".")[:2]
     return tuple(int(v) for v in major_minor)
 
+@functools.cache
+def get_driver_version():
+    return handle_return(driver.cuDriverGetVersion())
+
 
 class Transaction:
     """

cuda_core/cuda/core/_graph/_utils.pyx

@@ -163,3 +163,45 @@ def test_graph_capture_errors(init_cuda):
     with pytest.raises(RuntimeError, match="^Graph has not finished building."):
         gb.complete()
     gb.end_building().complete()
+
+
+def test_graph_capture_callback_python(init_cuda):
+    results = []
+
+    def my_callback():
+        results.append(42)
+
+    launch_stream = Device().create_stream()
+    gb = launch_stream.create_graph_builder().begin_building()
+
+    with pytest.raises(ValueError, match="user_data is only supported"):
+        gb.callback(my_callback, user_data=b"hello")
+
+    gb.callback(my_callback)
+    graph = gb.end_building().complete()
+
+    graph.launch(launch_stream)
+    launch_stream.sync()
+
+    assert results == [42]
+
+
+def test_graph_capture_callback_ctypes(init_cuda):
+    import ctypes
+
+    CALLBACK = ctypes.CFUNCTYPE(None, ctypes.c_void_p)
+    result = [0]
+
+    @CALLBACK
+    def read_byte(data):
+        result[0] = ctypes.cast(data, ctypes.POINTER(ctypes.c_uint8))[0]
+
+    launch_stream = Device().create_stream()
+    gb = launch_stream.create_graph_builder().begin_building()
+    gb.callback(read_byte, user_data=bytes([0xAB]))
+    graph = gb.end_building().complete()
+
+    graph.launch(launch_stream)
+    launch_stream.sync()
+
+    assert result[0] == 0xAB