fix: bind a CUDA context on the calling thread — the right one, and on both sides of the boundary

AGENTS.md

@@ -45,6 +45,7 @@ Python (editable; compiles the pybind11 extension via CMake):
 ```bash
 pip install -e .              # core graph API only
 pip install -e ".[cutedsl]"   # + OSS CuTeDSL kernels (nvidia-cutlass-dsl, cuda-python, tvm-ffi; framework-neutral)
+pip install -e ".[cutile]"    # + the cuTile linear-attention engines (cuda-tile; needs a system tileiras)
 pip install --group torch      # + torch for the CuTeDSL APIs (torch, torch-c-dlpack-ext)
 pip install --group jax        # + jax for the CuTeDSL APIs (jax >= 0.5; XLA entry points via cutlass.jax)
 ```

include/cudnn_frontend/graph_interface.h

@@ -1454,6 +1454,17 @@ class Graph : public ICudnn, public INode {
                           std::vector<int64_t> const &override_uids                 = {},
                           std::vector<std::vector<int64_t>> const &override_shapes  = {},
                           std::vector<std::vector<int64_t>> const &override_strides = {}) const {
+        // The driver-API engines read the calling THREAD's context stack, and a
+        // thread that has done no CUDA work yet (a PyTorch autograd worker) has
+        // none. All roads reach here, so every execute is covered; the steady
+        // state is one cuCtxGetCurrent, and the stream query only runs when a
+        // context actually has to be established.
+        if (!detail::has_current_context()) {
+            cudaStream_t stream = nullptr;
+            detail::get_stream(handle, &stream);
+            detail::ensure_current_context(stream);
+        }
+
         // Lazy init: prepare template if not done (e.g. deserialized graphs, build_plan_at_index)
         if (!varpack_prep_state->prepared.load(std::memory_order_acquire)) {
             CHECK_CUDNN_FRONTEND_ERROR(const_cast<Graph *>(this)->prepare_variant_pack_template());

include/cudnn_frontend_shim.h

@@ -417,6 +417,140 @@ cuda_get_device(int *device) {
     NV_FE_CALL_TO_CUDA(cuda_get_device, cudaGetDevice, device);
 }
 
+// Driver entry points resolved through the runtime, so the front end never links
+// libcuda.so -- the same approach as cu_tensor_map_encode_tiled(). Returns
+// nullptr when the entry point is unavailable.
+inline void *
+get_driver_entry_point(const char *name) {
+#if defined(CUDART_VERSION) && CUDART_VERSION >= 12000
+    void *pfn = nullptr;
+    cudaDriverEntryPointQueryResult query_result;
+    cudaError_t err;
+#if defined NV_CUDNN_FRONTEND_USE_DYNAMIC_LOADING
+#if CUDART_VERSION >= 12080
+    using GetEntryPointFn =
+        cudaError_t (*)(const char *, void **, unsigned int, int, cudaDriverEntryPointQueryResult *);
+    const char *resolver = "cudaGetDriverEntryPointByVersion";
+#else
+    using GetEntryPointFn =
+        cudaError_t (*)(const char *, void **, unsigned long long, cudaDriverEntryPointQueryResult *);
+    const char *resolver = "cudaGetDriverEntryPoint";
+#endif
+    GetEntryPointFn get_entry_point = nullptr;
+    // get_cuda_symbol throws when the library or symbol is missing; this lookup is
+    // best-effort and runs inside a static initializer, so it must not escape.
+#ifndef NV_CUDNN_DISABLE_EXCEPTION
+    try {
+        get_entry_point = reinterpret_cast<GetEntryPointFn>(get_cuda_symbol(CudaLibrary::CUDART, resolver));
+    } catch (...) {
+        return nullptr;
+    }
+#else
+    get_entry_point = reinterpret_cast<GetEntryPointFn>(get_cuda_symbol(CudaLibrary::CUDART, resolver));
+#endif
+    if (get_entry_point == nullptr) {
+        return nullptr;
+    }
+#if CUDART_VERSION >= 12080
+    err = get_entry_point(name, &pfn, 12000, cudaEnableDefault, &query_result);
+#else
+    err = get_entry_point(name, &pfn, cudaEnableDefault, &query_result);
+#endif
+#elif CUDART_VERSION >= 12080
+    err = cudaGetDriverEntryPointByVersion(name, &pfn, 12000, cudaEnableDefault, &query_result);
+#else
+    err = cudaGetDriverEntryPoint(name, &pfn, cudaEnableDefault, &query_result);
+#endif
+    if (err != cudaSuccess || query_result != cudaDriverEntryPointSuccess) {
+        return nullptr;
+    }
+    return pfn;
+#else
+    (void)name;
+    return nullptr;
+#endif
+}
+
+// True when the calling thread already has a context. One driver call, so the
+// execute path can skip the rest -- including the stream query -- in the steady
+// state.
+inline bool
+has_current_context() {
+#if defined(CUDART_VERSION) && CUDART_VERSION >= 12000
+    using PfnCtxGetCurrent            = CUresult(CUDAAPI *)(CUcontext *);
+    static const auto ctx_get_current = reinterpret_cast<PfnCtxGetCurrent>(get_driver_entry_point("cuCtxGetCurrent"));
+    if (ctx_get_current == nullptr) {
+        return true;  // cannot tell; leave the thread alone
+    }
+    CUcontext current = nullptr;
+    return ctx_get_current(&current) == CUDA_SUCCESS && current != nullptr;
+#else
+    return true;
+#endif
+}
+
+// Bind a driver context to the calling thread when it has none: a driver-API
+// launch reads that stack, and a thread that has done no CUDA work has nothing
+// on it (a PyTorch autograd worker). A bound context is left alone -- it is
+// process-wide and the caller chose it. A real stream names the right context;
+// the default-stream handles name none, so the runtime's device decides there.
+// Best-effort: what this cannot establish fails at the launch.
+inline void
+ensure_current_context(cudaStream_t stream) {
+#if defined(CUDART_VERSION) && CUDART_VERSION >= 12000
+    using PfnCtxGetCurrent    = CUresult(CUDAAPI *)(CUcontext *);
+    using PfnCtxSetCurrent    = CUresult(CUDAAPI *)(CUcontext);
+    using PfnStreamGetCtx     = CUresult(CUDAAPI *)(CUstream, CUcontext *);
+    using PfnDeviceGet        = CUresult(CUDAAPI *)(CUdevice *, int);
+    using PfnPrimaryCtxRetain = CUresult(CUDAAPI *)(CUcontext *, CUdevice);
+
+    // Thread-safe static initialization (C++11): the lookups happen once.
+    static const auto ctx_get_current = reinterpret_cast<PfnCtxGetCurrent>(get_driver_entry_point("cuCtxGetCurrent"));
+    static const auto ctx_set_current = reinterpret_cast<PfnCtxSetCurrent>(get_driver_entry_point("cuCtxSetCurrent"));
+    static const auto stream_get_ctx  = reinterpret_cast<PfnStreamGetCtx>(get_driver_entry_point("cuStreamGetCtx"));
+    static const auto device_get      = reinterpret_cast<PfnDeviceGet>(get_driver_entry_point("cuDeviceGet"));
+    static const auto primary_retain =
+        reinterpret_cast<PfnPrimaryCtxRetain>(get_driver_entry_point("cuDevicePrimaryCtxRetain"));
+    if (ctx_get_current == nullptr || ctx_set_current == nullptr) {
+        return;
+    }
+
+    CUcontext current = nullptr;
+    if (ctx_get_current(&current) == CUDA_SUCCESS && current != nullptr) {
+        return;
+    }
+
+    const bool names_a_context = stream != nullptr && stream != reinterpret_cast<cudaStream_t>(CU_STREAM_LEGACY) &&
+                                 stream != reinterpret_cast<cudaStream_t>(CU_STREAM_PER_THREAD);
+    if (names_a_context && stream_get_ctx != nullptr) {
+        CUcontext stream_ctx = nullptr;
+        if (stream_get_ctx(reinterpret_cast<CUstream>(stream), &stream_ctx) == CUDA_SUCCESS && stream_ctx != nullptr) {
+            ctx_set_current(stream_ctx);
+            return;
+        }
+    }
+
+    if (device_get == nullptr || primary_retain == nullptr) {
+        return;
+    }
+    int ordinal = 0;
+    if (cuda_get_device(&ordinal) != cudaSuccess) {
+        return;
+    }
+    CUdevice device = 0;
+    if (device_get(&device, ordinal) != CUDA_SUCCESS) {
+        return;
+    }
+    // Retained for the process lifetime, like the primary context itself.
+    CUcontext primary = nullptr;
+    if (primary_retain(&primary, device) == CUDA_SUCCESS) {
+        ctx_set_current(primary);
+    }
+#else
+    (void)stream;
+#endif
+}
+
 inline cudaError_t
 cuda_pointer_get_attributes(cudaPointerAttributes *attributes, const void *ptr) {
     NV_FE_CALL_TO_CUDA(cuda_pointer_get_attributes, cudaPointerGetAttributes, attributes, ptr);

pyproject.toml

@@ -68,6 +68,16 @@ cutedsl = [
     "cuda-python",
     "apache-tvm-ffi>=0.1.11",
 ]
+cutile = [
+    # The cuTile linear-attention engines. Base cuda-tile only -- its [tileiras]
+    # extra pins cuda-toolkit>=13.2,<13.4, and that upper bound would cap the
+    # whole environment's toolkit and shut out CUDA 12 entirely. Without it,
+    # cuda.tile falls back to a system `tileiras`, the same way this package
+    # already leaves GPU wheels to the user. Engines that cannot import the
+    # runtime decline in check_support, so a missing compiler costs those
+    # engines and nothing else.
+    "cuda-tile>=1.4; python_version >= '3.10'",
+]
 
 [dependency-groups]
 dev = [

python/cudnn/_device.py

@@ -59,38 +59,76 @@ def _device_handle(device: int):
     return _ck(*drv.cuDeviceGet(device))
 
 
-def ensure_current_context(stream=None) -> None:
-    """Bind a driver context to the calling thread when none is bound.
-
-    JIT engines talk to the driver directly, which reads the calling THREAD's
-    context stack — and an autograd backward runs on a worker thread where
-    ``cudaSetDevice`` has only moved the runtime's thread-local slot. Prefer
-    the stream's context, else retain the runtime-current device's primary one
-    (the retain ref is held for the process lifetime, like the primary context
-    itself). Best-effort: a context this cannot establish fails at the launch,
-    with the launch's own diagnostics."""
+@functools.lru_cache(maxsize=1)
+def _default_streams() -> frozenset:
+    """Handles that name no context: ``cuStreamGetCtx`` answers for the calling
+    thread's current context on all of them."""
+    drv = _driver()
+    if drv is None:
+        return frozenset({0})
+    return frozenset({0, int(drv.CU_STREAM_LEGACY), int(drv.CU_STREAM_PER_THREAD)})
+
+
+@functools.lru_cache(maxsize=None)
+def _primary_context(device: int):
+    """The retained primary context for ``device``, or ``None``. Once per
+    ordinal: a retain per execute would grow the usage count without bound."""
+    drv = _driver()
+    if drv is None:
+        return None
+    err, handle = drv.cuDeviceGet(device)
+    if int(err) != 0:
+        return None
+    err, primary = drv.cuDevicePrimaryCtxRetain(handle)
+    return primary if int(err) == 0 else None
+
+
+def _runtime_device():
+    """Ordinal the CUDA *runtime* holds current on this thread, or ``None``.
+    The driver cannot see that slot."""
     try:
-        drv = _driver()
-        if drv is None:
-            return
-        err, cur = drv.cuCtxGetCurrent()
-        if int(err) == 0 and int(cur) != 0:
-            return
-        if stream:
-            err, sctx = drv.cuStreamGetCtx(int(stream))
-            if int(err) == 0:
-                drv.cuCtxSetCurrent(sctx)
-                return
         import cuda.bindings.runtime as rt
+    except ImportError:
+        return None
+
+    err, device = rt.cudaGetDevice()
+    return int(device) if int(err) == 0 else None
+
 
-        err, device = rt.cudaGetDevice()
-        if int(err) != 0:
+def ensure_current_context(stream=None, device=None) -> None:
+    """Bind the context this work runs in to the calling thread.
+
+    A driver-API launch reads the calling thread's context stack; an autograd
+    backward runs on a worker whose stack is empty. A real stream names the
+    right context; the default-stream handles name none, so ``device`` decides
+    there. With no ``device``, a bound context is authoritative and only a cold
+    thread is given one (``frost.device.ambient_device``'s rung order).
+    Best-effort: what this cannot establish fails at the launch."""
+    drv = _driver()
+    if drv is None:
+        return
+    err, cur = drv.cuCtxGetCurrent()
+    cur = int(cur) if int(err) == 0 else 0
+    stream = 0 if stream is None else int(stream)
+    if stream not in _default_streams():
+        err, stream_ctx = drv.cuStreamGetCtx(stream)
+        if int(err) == 0 and int(stream_ctx) != 0:
+            if int(stream_ctx) != cur:
+                drv.cuCtxSetCurrent(stream_ctx)
+            return
+    if device is None:
+        if cur:
+            return
+        device = _runtime_device()
+        if device is None:
+            return
+    elif cur:
+        err, cur_device = drv.cuCtxGetDevice()
+        if int(err) == 0 and int(cur_device) == int(device):
             return
-        err, pctx = drv.cuDevicePrimaryCtxRetain(_device_handle(int(device)))
-        if int(err) == 0:
-            drv.cuCtxSetCurrent(pctx)
-    except Exception:  # noqa: BLE001
-        pass
+    primary = _primary_context(int(device))
+    if primary is not None:
+        drv.cuCtxSetCurrent(primary)
 
 
 class DeviceInfo:

python/cudnn/_pygraph.py

@@ -1777,10 +1777,10 @@ def execute(
         if eng is not None:  # python engine (plan id in the reserved region)
             h = handle if handle is not None else self._handle
             ctx = ExecutionContext(handle=h, stream=self._resolve_stream(h), workspace=workspace)
-            # A JIT engine talks to the driver directly, which reads the calling
-            # THREAD's context stack -- and an autograd backward runs on a worker
-            # thread that has none. Once here, so every python engine is covered.
-            ensure_current_context(ctx.stream)
+            # A JIT engine launches through the driver, which reads the calling
+            # thread's context stack; an autograd worker has none. The handle's
+            # device decides when the stream names no context.
+            ensure_current_context(ctx.stream, h.device.ordinal if h is not None else None)
             if self._plan_index not in self._compiled_plans:
                 # compile with the CALLER's context (execute-supplied handle
                 # and its stream reach the JIT build)