feat(cuda.core): add Stream.synchronization_policy via cuStreamSetAttribute

cuda_core/cuda/core/_stream.pxd

@@ -13,6 +13,7 @@ cdef class Stream:
         int _device_id
         int _nonblocking
         int _priority
+        int _synchronization_policy
         object __weakref__
 
     @staticmethod

cuda_core/cuda/core/_stream.pyi

@@ -25,10 +25,15 @@ class StreamOptions:
     priority : int, optional
         Stream priority where lower number represents a
         higher priority. (Default to lowest priority)
+    synchronization_policy : SynchronizationPolicyType | None, optional
+        CPU wait policy applied when synchronizing this stream from the host.
+        Maps to ``CU_LAUNCH_ATTRIBUTE_SYNCHRONIZATION_POLICY`` via
+        ``cuStreamSetAttribute``. (Default to None, leaving the driver default)
 
     """
     nonblocking: cython.bint = True
     priority: int | None = None
+    synchronization_policy: object = None
 
 class IsStreamType(Protocol):
 
@@ -110,6 +115,18 @@ class Stream:
     def priority(self) -> int:
         """Return the stream priority."""
 
+    @property
+    def synchronization_policy(self):
+        """Return the stream's CPU wait policy for host-side synchronization.
+
+        Maps to ``CU_LAUNCH_ATTRIBUTE_SYNCHRONIZATION_POLICY``, queried via
+        ``cuStreamGetAttribute``.
+        """
+
+    @synchronization_policy.setter
+    def synchronization_policy(self, policy: object):
+        """Set the stream's CPU wait policy for subsequent work on this stream."""
+
     def sync(self) -> None:
         """Synchronize the stream."""
 

cuda_core/cuda/core/_stream.pyx

@@ -49,6 +49,60 @@ if TYPE_CHECKING:
 
 __all__ = ['LEGACY_DEFAULT_STREAM', 'PER_THREAD_DEFAULT_STREAM', 'Stream', 'StreamOptions']
 
+cdef int _SYNC_POLICY_ATTR_ID = (
+    cydriver.CUlaunchAttributeID_enum.CU_LAUNCH_ATTRIBUTE_SYNCHRONIZATION_POLICY
+)
+
+
+cdef object _validate_synchronization_policy(object policy):
+    from cuda.core.typing import SynchronizationPolicyType
+
+    if policy is None:
+        raise TypeError("synchronization_policy must not be None")
+    if isinstance(policy, SynchronizationPolicyType):
+        return policy
+    try:
+        value = int(policy)
+    except (TypeError, ValueError) as exc:
+        raise TypeError(
+            "Stream.synchronization_policy must be a SynchronizationPolicyType, "
+            f"cuda.bindings.driver.CUsynchronizationPolicy, or int; got {type(policy).__name__}"
+        ) from exc
+    try:
+        return SynchronizationPolicyType(value)
+    except ValueError as exc:
+        raise ValueError(
+            f"Stream.synchronization_policy must be one of "
+            f"{[member.name for member in SynchronizationPolicyType]}; got {policy!r}"
+        ) from exc
+
+
+cdef int Stream_query_synchronization_policy(Stream self) except?-1 nogil:
+    cdef cydriver.CUstreamAttrValue value
+    HANDLE_RETURN(cydriver.cuStreamGetAttribute(
+        as_cu(self._h_stream),
+        <cydriver.CUstreamAttrID>_SYNC_POLICY_ATTR_ID,
+        &value,
+    ))
+    return <int>value.syncPolicy
+
+
+cdef int Stream_normalize_queried_sync_policy(int policy) nogil:
+    # cuStreamGetAttribute may return -1 when the attribute was never set.
+    if policy < 0:
+        return <int>cydriver.CUsynchronizationPolicy_enum.CU_SYNC_POLICY_AUTO
+    return policy
+
+
+cdef void Stream_apply_synchronization_policy(Stream self, int policy) except * nogil:
+    cdef cydriver.CUstreamAttrValue value
+    value.syncPolicy = <cydriver.CUsynchronizationPolicy>policy
+    HANDLE_RETURN(cydriver.cuStreamSetAttribute(
+        as_cu(self._h_stream),
+        <cydriver.CUstreamAttrID>_SYNC_POLICY_ATTR_ID,
+        &value,
+    ))
+
 
 @dataclass
 cdef class StreamOptions:
@@ -61,11 +115,16 @@ cdef class StreamOptions:
     priority : int, optional
         Stream priority where lower number represents a
         higher priority. (Default to lowest priority)
+    synchronization_policy : SynchronizationPolicyType | None, optional
+        CPU wait policy applied when synchronizing this stream from the host.
+        Maps to ``CU_LAUNCH_ATTRIBUTE_SYNCHRONIZATION_POLICY`` via
+        ``cuStreamSetAttribute``. (Default to None, leaving the driver default)
 
     """
 
     nonblocking : cython.bint = True
     priority: int | None = None
+    synchronization_policy: object = None
 
 
 class IsStreamType(Protocol):
@@ -110,6 +169,7 @@ cdef class Stream:
         s._device_id = -1  # lazy init'd (invalid sentinel)
         s._nonblocking = -1  # lazy init'd
         s._priority = INT32_MIN  # lazy init'd
+        s._synchronization_policy = INT32_MIN  # lazy init'd
         return s
 
     @classmethod
@@ -146,12 +206,14 @@ cdef class Stream:
         cdef StreamOptions opts = check_or_create_options(StreamOptions, options, "Stream options")
         nonblocking = opts.nonblocking
         priority = opts.priority
+        sync_policy_opt = opts.synchronization_policy
 
         cdef unsigned int flags = (cydriver.CUstream_flags.CU_STREAM_NON_BLOCKING if nonblocking
                                    else cydriver.CUstream_flags.CU_STREAM_DEFAULT)
         # TODO: we might want to consider memoizing high/low per CUDA context and avoid this call
         cdef int high, low
         cdef cydriver.CUresult res_code
+        cdef int sync_policy
         with nogil:
             res_code = cydriver.cuCtxGetStreamPriorityRange(&high, &low)
         if res_code != cydriver.CUresult.CUDA_SUCCESS:
@@ -190,6 +252,12 @@ cdef class Stream:
         cdef Stream self = Stream._from_handle(cls, h_stream)
         self._nonblocking = int(nonblocking)
         self._priority = prio
+        if sync_policy_opt is not None:
+            validated = _validate_synchronization_policy(sync_policy_opt)
+            sync_policy = int(validated)
+            with nogil:
+                Stream_apply_synchronization_policy(self, sync_policy)
+            self._synchronization_policy = sync_policy
         if device_id is not None:
             self._device_id = device_id
         return self
@@ -251,6 +319,36 @@ cdef class Stream:
             self._priority = prio
         return self._priority
 
+    @property
+    def synchronization_policy(self):
+        """Return the stream's CPU wait policy for host-side synchronization.
+
+        Maps to ``CU_LAUNCH_ATTRIBUTE_SYNCHRONIZATION_POLICY``, queried via
+        ``cuStreamGetAttribute``.
+        """
+        from cuda.core.typing import SynchronizationPolicyType
+
+        cdef int policy
+        if self._synchronization_policy == INT32_MIN or Stream_is_default_token(self):
+            with nogil:
+                policy = Stream_query_synchronization_policy(self)
+                policy = Stream_normalize_queried_sync_policy(policy)
+            if not Stream_is_default_token(self):
+                self._synchronization_policy = policy
+        else:
+            policy = self._synchronization_policy
+        return SynchronizationPolicyType(policy)
+
+    @synchronization_policy.setter
+    def synchronization_policy(self, object policy):
+        """Set the stream's CPU wait policy for subsequent work on this stream."""
+        cdef object validated = _validate_synchronization_policy(policy)
+        cdef int value = int(validated)
+        with nogil:
+            Stream_apply_synchronization_policy(self, value)
+        if not Stream_is_default_token(self):
+            self._synchronization_policy = value
+
     def sync(self) -> None:
         """Synchronize the stream."""
         with nogil:

cuda_core/cuda/core/typing.py

@@ -5,6 +5,7 @@
 """Public type aliases, protocols, and enumerations used in cuda.core API signatures."""
 
 import sys
+from enum import IntEnum
 from typing import TYPE_CHECKING
 from typing import Literal as _Literal
 from typing import TypeAlias as _TypeAlias
@@ -47,6 +48,7 @@ class StrEnum(str, Enum):
     "ProcessStateType",
     "ReadModeType",
     "SourceCodeType",
+    "SynchronizationPolicyType",
     "VirtualMemoryAccessType",
     "VirtualMemoryAllocationType",
     "VirtualMemoryGranularityType",
@@ -124,6 +126,24 @@ class PCHStatusType(StrEnum):
     FAILED = "failed"
 
 
+class SynchronizationPolicyType(IntEnum):
+    """CPU wait policy when synchronizing a stream from the host.
+
+    Maps to ``cuda.bindings.driver.CUsynchronizationPolicy`` and is applied via
+    ``cuStreamSetAttribute`` with ``CU_LAUNCH_ATTRIBUTE_SYNCHRONIZATION_POLICY``.
+
+    * ``AUTO`` — driver selects the policy (default for new streams).
+    * ``SPIN`` — busy-wait on the CPU (lowest latency).
+    * ``YIELD`` — yield the CPU while waiting.
+    * ``BLOCKING_SYNC`` — block in the OS scheduler while waiting.
+    """
+
+    AUTO = driver.CUsynchronizationPolicy.CU_SYNC_POLICY_AUTO
+    SPIN = driver.CUsynchronizationPolicy.CU_SYNC_POLICY_SPIN
+    YIELD = driver.CUsynchronizationPolicy.CU_SYNC_POLICY_YIELD
+    BLOCKING_SYNC = driver.CUsynchronizationPolicy.CU_SYNC_POLICY_BLOCKING_SYNC
+
+
 class GraphConditionalType(StrEnum):
     """Conditional node flavor for :class:`~cuda.core.graph.GraphBuilder`.
 

cuda_core/docs/source/api_private.rst

@@ -35,6 +35,7 @@ CUDA runtime
    typing.ProcessStateType
    typing.ReadModeType
    typing.SourceCodeType
+   typing.SynchronizationPolicyType
    typing.VirtualMemoryAccessType
    typing.VirtualMemoryAllocationType
    typing.VirtualMemoryGranularityType

cuda_core/tests/test_stream.py

@@ -469,3 +469,88 @@ def test_created_stream_keeps_its_own_context():
     finally:
         dev0.set_current()
         stream.close()
+
+
+# ============================================================================
+# Stream synchronization policy tests
+# ============================================================================
+
+
+@pytest.mark.parametrize(
+    "policy",
+    [
+        "AUTO",
+        "SPIN",
+        "YIELD",
+        "BLOCKING_SYNC",
+    ],
+)
+def test_stream_synchronization_policy_set_get(init_cuda, policy):
+    from cuda.core.typing import SynchronizationPolicyType
+
+    stream = Device().create_stream()
+    enum_policy = SynchronizationPolicyType[policy]
+    stream.synchronization_policy = enum_policy
+    assert stream.synchronization_policy == enum_policy
+    stream.close()
+
+
+@pytest.mark.parametrize("invalid_policy", ["spin", -1, 99])
+def test_stream_synchronization_policy_invalid(init_cuda, invalid_policy):
+    stream = Device().create_stream()
+    with pytest.raises((TypeError, ValueError)):
+        stream.synchronization_policy = invalid_policy
+    stream.close()
+
+
+def test_stream_synchronization_policy_default(init_cuda):
+    from cuda.core.typing import SynchronizationPolicyType
+
+    stream = Device().create_stream()
+    assert stream.synchronization_policy == SynchronizationPolicyType.AUTO
+    stream.close()
+
+
+def test_stream_options_synchronization_policy(init_cuda):
+    from cuda.core.typing import SynchronizationPolicyType
+
+    stream = Device().create_stream(options=StreamOptions(synchronization_policy=SynchronizationPolicyType.YIELD))
+    assert stream.synchronization_policy == SynchronizationPolicyType.YIELD
+    stream.close()
+
+
+@pytest.mark.parametrize(
+    "policy",
+    [
+        "AUTO",
+        "SPIN",
+        "YIELD",
+        "BLOCKING_SYNC",
+    ],
+)
+@pytest.mark.agent_authored(model="composer-2.5")
+def test_stream_synchronization_policy_launch_and_sync(init_cuda, policy):
+    """Stream-level sync policy is applied via cuStreamSetAttribute and survives launch+sync."""
+    import cuda.bindings
+    from cuda.core import Device, LaunchConfig, Program, ProgramOptions, launch
+    from cuda.core._utils.version import driver_version
+    from cuda.core.typing import ObjectCodeFormatType, SourceCodeType, SynchronizationPolicyType
+
+    if int(cuda.bindings.__version__.split(".")[0]) >= 13 and driver_version()[0] < 13:
+        pytest.skip("CUDA 13 bindings produce modules incompatible with CUDA 12 drivers")
+
+    dev = Device()
+    dev.set_current()
+    stream = dev.create_stream()
+    stream.synchronization_policy = SynchronizationPolicyType[policy]
+
+    code = 'extern "C" __global__ void noop() {}'
+    arch = "".join(f"{i}" for i in dev.compute_capability)
+    program = Program(code, SourceCodeType.CXX, options=ProgramOptions(arch=f"sm_{arch}"))
+    mod = program.compile(ObjectCodeFormatType.CUBIN)
+    ker = mod.get_kernel("noop")
+
+    launch(stream, LaunchConfig(grid=1, block=1), ker)
+    stream.sync()
+    assert stream.synchronization_policy == SynchronizationPolicyType[policy]
+    stream.close()