FE execute-path host overhead: idempotent set_stream + skip the discarded per-call context

python/cudnn/__init__.py

@@ -21,7 +21,6 @@ def is_windows():
     "backend_version",
     "backend_version_string",
     "get_last_error_string",
-    "destroy_handle",
     "norm_forward_phase",
     "reduction_mode",
     "behavior_note",
@@ -31,7 +30,6 @@ def is_windows():
     "create_device_properties",
     "get_stream",
     "numerical_note",
-    "set_stream",
     "build_plan_policy",
     "data_type",
     "tensor_reordering",
@@ -60,6 +58,36 @@ def is_windows():
     if hasattr(_pybind_module, _optional_symbol):
         globals()[_optional_symbol] = getattr(_pybind_module, _optional_symbol)
 
+
+# The last stream set on each handle, so set_stream() below can skip a redundant backend call.
+_handle_to_stream: dict = {}
+
+
+def set_stream(handle, stream):
+    """Set the CUDA stream a cuDNN handle runs on (wraps the compiled ``cudnnSetStream``).
+
+    ``cudnnSetStream`` is not free: for a non-null stream it issues several CUDA driver queries
+    on every call (green-context detection, stream priority, priority range) to maintain cuDNN's
+    internal per-priority stream pool, even when the stream is unchanged -- ~2.4us/call on
+    Blackwell. Frameworks that call this before every ``execute`` pay it every iteration, so we
+    cache the last stream per handle and skip the backend call when it has not changed; a
+    steady-state loop pays it once. (Assumes a handle is not driven from two streams
+    concurrently, which is the normal single-stream case; a caller that does needs its own
+    handle per stream regardless.)
+    """
+    if _handle_to_stream.get(handle) == stream:
+        return
+    _pybind_module._raw_set_stream(handle, stream)
+    _handle_to_stream[handle] = stream
+
+
+def destroy_handle(handle):
+    """Destroy a cuDNN handle (wraps the compiled binding); forget its cached stream so a
+    reused handle address is not wrongly skipped by set_stream()."""
+    _handle_to_stream.pop(handle, None)
+    return _pybind_module._raw_destroy_handle(handle)
+
+
 from .datatypes import _library_type, _is_torch_tensor
 
 __version__ = "1.27.0"

python/cudnn/_pygraph.py

@@ -1755,9 +1755,12 @@ def execute(
                     ``set_stream`` semantics, both python engines and backend)
             override_uids/shapes/strides: dynamic-shape overrides (backend path)
         """
-        # A JIT engine must compile for the device/stream it will run on.
-        caller_ctx = self._build_context(handle) if handle is not None else None
         if not self._is_built:
+            # A JIT engine must compile for the device/stream it will run on, so
+            # build the caller context here. Only here: a steady-state execute()
+            # otherwise discarded this (a cudnnGetStream round-trip + an
+            # ExecutionContext alloc, ~2.9us) on every already-built call.
+            caller_ctx = self._build_context(handle) if handle is not None else None
             if not self._planning_done:
                 self.create_execution_plans()
             self.build(ctx=caller_ctx)

python/properties.cpp

@@ -285,9 +285,12 @@ init_properties(py::module_& m) {
               &create_device_properties_helper));
 
     m.def("create_handle", &HandleManagement::create_handle);
-    m.def("destroy_handle", &HandleManagement::destroy_handle);
+    // destroy_handle / set_stream are exposed under a raw name and wrapped in Python
+    // (cudnn/__init__.py) to skip a redundant cudnnSetStream when the stream is unchanged,
+    // matching how the graph execute binding (_execute) is wrapped as the public execute().
+    m.def("_raw_destroy_handle", &HandleManagement::destroy_handle);
     m.def("get_stream", &HandleManagement::get_stream);
-    m.def("set_stream", &HandleManagement::set_stream, py::arg("handle"), py::arg("stream"));
+    m.def("_raw_set_stream", &HandleManagement::set_stream, py::arg("handle"), py::arg("stream"));
 
     py::enum_<cudnn_frontend::NormFwdPhase_t>(m, "norm_forward_phase")
         .value("INFERENCE", cudnn_frontend::NormFwdPhase_t::INFERENCE)

test/python/test_set_stream_cache.py

@@ -0,0 +1,38 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""cudnn.set_stream caches the last stream per handle and skips the backend call
+(cudnnSetStream, which re-issues several CUDA driver queries every call) when the
+stream has not changed. These tests mock the raw backend call, so no GPU is needed."""
+
+import pytest
+
+import cudnn
+
+pytestmark = pytest.mark.L0
+
+
+def test_set_stream_skips_backend_call_when_unchanged(monkeypatch):
+    calls = []
+    monkeypatch.setattr(cudnn._pybind_module, "_raw_set_stream", lambda h, s: calls.append((h, s)))
+    cudnn._handle_to_stream.clear()
+
+    cudnn.set_stream(handle=1, stream=100)
+    cudnn.set_stream(handle=1, stream=100)  # unchanged -> skipped
+    cudnn.set_stream(handle=1, stream=200)  # changed -> forwarded
+    cudnn.set_stream(handle=2, stream=100)  # different handle -> forwarded
+
+    assert calls == [(1, 100), (1, 200), (2, 100)]
+
+
+def test_destroy_handle_forgets_cached_stream(monkeypatch):
+    calls = []
+    monkeypatch.setattr(cudnn._pybind_module, "_raw_set_stream", lambda h, s: calls.append((h, s)))
+    monkeypatch.setattr(cudnn._pybind_module, "_raw_destroy_handle", lambda h: None)
+    cudnn._handle_to_stream.clear()
+
+    cudnn.set_stream(handle=7, stream=100)
+    cudnn.destroy_handle(7)
+    cudnn.set_stream(handle=7, stream=100)  # a reused handle address must re-arm the backend
+
+    assert calls == [(7, 100), (7, 100)]