Fix tensor bridge DLL import failure on Windows

aoti_torch_get_current_cuda_stream lives in torch_cuda.dll, not
torch_cpu.dll. The stub import library pointed at the wrong DLL,
causing "The specified procedure could not be found" on Windows.

- Move aoti_torch_get_current_cuda_stream from aoti_shim.def
  (torch_cpu.dll) to new aoti_shim_cuda.def (torch_cuda.dll)
- Update build_hooks.py to generate stub libs for both DLLs
  via a loop
- Add torch_cuda.dll to delvewheel exclude list

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: leofang

cuda_core/build_hooks.py

@@ -183,19 +183,22 @@ def get_sources(mod_name):
         # related to free-threading builds.
         extra_compile_args += ["-DCYTHON_TRACE_NOGIL=1", "-DCYTHON_USE_SYS_MONITORING=0"]
 
-    # On Windows, _tensor_bridge.pyx needs a stub import library so the MSVC
-    # linker can resolve the AOTI symbols (they live in torch_cpu.dll at
-    # runtime).  We generate the .lib from a .def file at build time.
+    # On Windows, _tensor_bridge.pyx needs stub import libraries so the MSVC
+    # linker can resolve the AOTI symbols at link time.  At runtime the symbols
+    # resolve from the actual DLLs loaded by 'import torch'.
+    #   - aoti_shim.def       -> torch_cpu.dll   (dtype, device, tensor metadata)
+    #   - aoti_shim_cuda.def  -> torch_cuda.dll   (CUDA stream access)
     _aoti_extra_link_args = []
     if sys.platform == "win32":
-        _def_file = os.path.join("cuda", "core", "_include", "aoti_shim.def")
-        _lib_file = os.path.join("build", "aoti_shim.lib")
         os.makedirs("build", exist_ok=True)
-        subprocess.check_call(  # noqa: S603
-            ["lib", f"/DEF:{_def_file}", f"/OUT:{_lib_file}", "/MACHINE:X64"],  # noqa: S607
-            stdout=subprocess.DEVNULL,
-        )
-        _aoti_extra_link_args = [_lib_file]
+        for def_name in ("aoti_shim", "aoti_shim_cuda"):
+            def_file = os.path.join("cuda", "core", "_include", f"{def_name}.def")
+            lib_file = os.path.join("build", f"{def_name}.lib")
+            subprocess.check_call(  # noqa: S603
+                ["lib", f"/DEF:{def_file}", f"/OUT:{lib_file}", "/MACHINE:X64"],  # noqa: S607
+                stdout=subprocess.DEVNULL,
+            )
+            _aoti_extra_link_args.append(lib_file)
 
     def get_extra_link_args(mod_name):
         if mod_name == "_tensor_bridge" and _aoti_extra_link_args:

cuda_core/cuda/core/_include/aoti_shim.def

@@ -34,4 +34,3 @@ EXPORTS
     aoti_torch_get_device_index
     aoti_torch_device_type_cpu
     aoti_torch_device_type_cuda
-    aoti_torch_get_current_cuda_stream

cuda_core/cuda/core/_include/aoti_shim_cuda.def

@@ -0,0 +1,5 @@
+; Stub import library for CUDA-specific AOTI symbols (torch_cuda.dll).
+; See aoti_shim.def for the torch_cpu.dll counterpart.
+LIBRARY torch_cuda.dll
+EXPORTS
+    aoti_torch_get_current_cuda_stream

cuda_core/pyproject.toml

@@ -117,4 +117,4 @@ archs = "native"
 [tool.cibuildwheel.windows]
 archs = "AMD64"
 before-build = "pip install delvewheel"
-repair-wheel-command = "delvewheel repair --namespace-pkg cuda --exclude \"torch_cpu.dll;torch_python.dll\" -w {dest_dir} {wheel}"
+repair-wheel-command = "delvewheel repair --namespace-pkg cuda --exclude \"torch_cpu.dll;torch_python.dll;torch_cuda.dll\" -w {dest_dir} {wheel}"