feat(cuda.core): add cluster scheduling policy to LaunchConfig

cuda_core/cuda/core/_launch_config.pxd

@@ -16,6 +16,7 @@ cdef class LaunchConfig:
         public int shmem_size
         public bint is_cooperative
         public bint programmatic_stream_serialization
+        public str cluster_scheduling_policy_preference
         public int priority
 
         vector[cydriver.CUlaunchAttribute] _attrs

cuda_core/cuda/core/_launch_config.pyi

@@ -2,7 +2,9 @@
 
 from typing import Any
 
-_LAUNCH_CONFIG_ATTRS = ('grid', 'cluster', 'block', 'shmem_size', 'is_cooperative', 'programmatic_stream_serialization', 'priority')
+from cuda.core._utils.cuda_utils import driver
+
+_LAUNCH_CONFIG_ATTRS = ('grid', 'cluster', 'block', 'shmem_size', 'is_cooperative', 'programmatic_stream_serialization', 'cluster_scheduling_policy_preference', 'priority')
 __all__ = ['LaunchConfig']
 
 class LaunchConfig:
@@ -39,6 +41,13 @@ class LaunchConfig:
         Whether to allow programmatic stream serialization (PDL). When True,
         the kernel may overlap with a previous kernel in the same stream that
         signals completion via programmatic means.
+    cluster_scheduling_policy_preference : str, optional
+        Cluster scheduling policy for the launch. One of ``"DEFAULT"``,
+        ``"SPREAD"``, or ``"LOAD_BALANCING"``.
+        When ``None`` (default), the launch attribute is omitted and the
+        driver applies the kernel function's default policy.
+        Passing ``"DEFAULT"`` explicitly sets the driver default via the
+        launch attribute.
     priority : int, optional
         Execution priority of the kernel. Lower numbers represent higher
         priorities. The meaningful range of values is device-specific,
@@ -49,15 +58,17 @@ class LaunchConfig:
         nonzero value outside this range raises :class:`ValueError`.
         When omitted (or 0), the launch uses the stream's priority.
     """
+    _CLUSTER_SCHED_POLICY_TO_DRIVER = {'DEFAULT': driver.CUclusterSchedulingPolicy.CU_CLUSTER_SCHEDULING_POLICY_DEFAULT, 'SPREAD': driver.CUclusterSchedulingPolicy.CU_CLUSTER_SCHEDULING_POLICY_SPREAD, 'LOAD_BALANCING': driver.CUclusterSchedulingPolicy.CU_CLUSTER_SCHEDULING_POLICY_LOAD_BALANCING}
     grid: tuple[Any, ...]
     cluster: tuple[Any, ...]
     block: tuple[Any, ...]
     shmem_size: int
     is_cooperative: bool
     programmatic_stream_serialization: bool
+    cluster_scheduling_policy_preference: str
     priority: int
 
-    def __init__(self, grid: int | tuple[int, ...] | None=None, cluster: int | tuple[int, ...] | None=None, block: int | tuple[int, ...] | None=None, shmem_size: int | None=None, is_cooperative: bool=False, programmatic_stream_serialization: bool=False, priority: int | None=None) -> None:
+    def __init__(self, grid: int | tuple[int, ...] | None=None, cluster: int | tuple[int, ...] | None=None, block: int | tuple[int, ...] | None=None, shmem_size: int | None=None, is_cooperative: bool=False, programmatic_stream_serialization: bool=False, cluster_scheduling_policy_preference: str | None=None, priority: int | None=None) -> None:
         """Initialize LaunchConfig with validation.
 
         Parameters
@@ -74,6 +85,11 @@ class LaunchConfig:
             Whether to launch as cooperative kernel (default: False)
         programmatic_stream_serialization : bool, optional
             Whether to allow programmatic stream serialization / PDL (default: False)
+        cluster_scheduling_policy_preference : str, optional
+            Cluster scheduling policy for the launch: ``"DEFAULT"``,
+            ``"SPREAD"``, or ``"LOAD_BALANCING"``.
+            ``None`` (default) omits the launch attribute; ``"DEFAULT"``
+            sets the driver default explicitly.
         priority : int, optional
             Execution priority of the kernel. Lower numbers represent higher
             priorities. The meaningful range of values is device-specific,
@@ -89,6 +105,8 @@ class LaunchConfig:
         """Return string representation of LaunchConfig."""
     def __eq__(self, other: object) -> bool: ...
     def __hash__(self) -> int: ...
+    def _validate_cluster_scheduling_policy_preference(self, value): ...
+    def _cluster_sched_policy_driver_value(self): ...
 
 def _to_native_launch_config(config: LaunchConfig) -> object:
     """Convert LaunchConfig to native driver CUlaunchConfig.

cuda_core/cuda/core/_launch_config.pyx

@@ -13,6 +13,7 @@ from cuda.core._utils.cuda_utils import (
     cast_to_3_tuple,
     driver,
 )
+from cuda.core._utils.validators import format_or_list
 
 _LAUNCH_CONFIG_ATTRS = (
     'grid',
@@ -21,6 +22,7 @@ _LAUNCH_CONFIG_ATTRS = (
     'shmem_size',
     'is_cooperative',
     'programmatic_stream_serialization',
+    'cluster_scheduling_policy_preference',
     'priority',
 )
 
@@ -61,6 +63,13 @@ cdef class LaunchConfig:
         Whether to allow programmatic stream serialization (PDL). When True,
         the kernel may overlap with a previous kernel in the same stream that
         signals completion via programmatic means.
+    cluster_scheduling_policy_preference : str, optional
+        Cluster scheduling policy for the launch. One of ``"DEFAULT"``,
+        ``"SPREAD"``, or ``"LOAD_BALANCING"``.
+        When ``None`` (default), the launch attribute is omitted and the
+        driver applies the kernel function's default policy.
+        Passing ``"DEFAULT"`` explicitly sets the driver default via the
+        launch attribute.
     priority : int, optional
         Execution priority of the kernel. Lower numbers represent higher
         priorities. The meaningful range of values is device-specific,
@@ -72,6 +81,12 @@ cdef class LaunchConfig:
         When omitted (or 0), the launch uses the stream's priority.
     """
 
+    _CLUSTER_SCHED_POLICY_TO_DRIVER = {
+        "DEFAULT": driver.CUclusterSchedulingPolicy.CU_CLUSTER_SCHEDULING_POLICY_DEFAULT,
+        "SPREAD": driver.CUclusterSchedulingPolicy.CU_CLUSTER_SCHEDULING_POLICY_SPREAD,
+        "LOAD_BALANCING": driver.CUclusterSchedulingPolicy.CU_CLUSTER_SCHEDULING_POLICY_LOAD_BALANCING,
+    }
+
     # TODO: expand LaunchConfig to include other attributes
     # Note: attributes are declared in _launch_config.pxd
 
@@ -83,6 +98,7 @@ cdef class LaunchConfig:
         shmem_size: int | None = None,
         is_cooperative: bool = False,
         programmatic_stream_serialization: bool = False,
+        cluster_scheduling_policy_preference: str | None = None,
         priority: int | None = None,
     ) -> None:
         """Initialize LaunchConfig with validation.
@@ -101,6 +117,11 @@ cdef class LaunchConfig:
             Whether to launch as cooperative kernel (default: False)
         programmatic_stream_serialization : bool, optional
             Whether to allow programmatic stream serialization / PDL (default: False)
+        cluster_scheduling_policy_preference : str, optional
+            Cluster scheduling policy for the launch: ``"DEFAULT"``,
+            ``"SPREAD"``, or ``"LOAD_BALANCING"``.
+            ``None`` (default) omits the launch attribute; ``"DEFAULT"``
+            sets the driver default explicitly.
         priority : int, optional
             Execution priority of the kernel. Lower numbers represent higher
             priorities. The meaningful range of values is device-specific,
@@ -115,6 +136,12 @@ cdef class LaunchConfig:
         self.grid = cast_to_3_tuple("LaunchConfig.grid", grid)
         self.block = cast_to_3_tuple("LaunchConfig.block", block)
 
+        self.cluster_scheduling_policy_preference = (
+            self._validate_cluster_scheduling_policy_preference(
+                cluster_scheduling_policy_preference
+            )
+        )
+
         # FIXME: Calling Device() strictly speaking is not quite right; we should instead
         # look up the device from stream. We probably need to defer the checks related to
         # device compute capability or attributes.
@@ -177,6 +204,27 @@ cdef class LaunchConfig:
     def __hash__(self) -> int:
         return hash(self._identity())
 
+    def _validate_cluster_scheduling_policy_preference(self, value):
+        if value is None:
+            return None
+        if isinstance(value, str) and value in LaunchConfig._CLUSTER_SCHED_POLICY_TO_DRIVER:
+            cc = Device().compute_capability
+            if cc < (9, 0):
+                raise CUDAError(
+                    "cluster launch attributes are not supported on devices with "
+                    f"compute capability < 9.0 (got {cc})"
+                )
+            return value
+        valid = format_or_list(LaunchConfig._CLUSTER_SCHED_POLICY_TO_DRIVER.keys())
+        raise ValueError(
+            f"{value!r} is not a valid cluster_scheduling_policy_preference. Must be {valid}"
+        )
+
+    def _cluster_sched_policy_driver_value(self):
+        return LaunchConfig._CLUSTER_SCHED_POLICY_TO_DRIVER[
+            self.cluster_scheduling_policy_preference
+        ]
+
     cdef cydriver.CUlaunchConfig _to_native_launch_config(self):
         cdef cydriver.CUlaunchConfig drv_cfg
         cdef cydriver.CUlaunchAttribute attr
@@ -210,6 +258,13 @@ cdef class LaunchConfig:
             attr.value.programmaticStreamSerializationAllowed = 1
             self._attrs.push_back(attr)
 
+        if self.cluster_scheduling_policy_preference is not None:
+            attr.id = cydriver.CUlaunchAttributeID.CU_LAUNCH_ATTRIBUTE_CLUSTER_SCHEDULING_POLICY_PREFERENCE
+            attr.value.clusterSchedulingPolicyPreference = int(
+                self._cluster_sched_policy_driver_value()
+            )
+            self._attrs.push_back(attr)
+
         if self.priority:
             attr.id = cydriver.CUlaunchAttributeID.CU_LAUNCH_ATTRIBUTE_PRIORITY
             attr.value.priority = self.priority
@@ -276,6 +331,12 @@ cpdef object _to_native_launch_config(LaunchConfig config):
         attr.value.programmaticStreamSerializationAllowed = 1
         attrs.append(attr)
 
+    if config.cluster_scheduling_policy_preference is not None:
+        attr = driver.CUlaunchAttribute()
+        attr.id = driver.CUlaunchAttributeID.CU_LAUNCH_ATTRIBUTE_CLUSTER_SCHEDULING_POLICY_PREFERENCE
+        attr.value.clusterSchedulingPolicyPreference = config._cluster_sched_policy_driver_value()
+        attrs.append(attr)
+
     if config.priority:
         attr = driver.CUlaunchAttribute()
         attr.id = driver.CUlaunchAttributeID.CU_LAUNCH_ATTRIBUTE_PRIORITY

cuda_core/tests/test_launcher.py

@@ -405,6 +405,95 @@ def test_to_native_launch_config_cluster_branch():
     assert (attr.value.clusterDim.x, attr.value.clusterDim.y, attr.value.clusterDim.z) == (2, 2, 2)
 
 
+@pytest.mark.agent_authored(model="cursor-grok-4.6")
+def test_launch_config_cluster_scheduling_policy(monkeypatch):
+    """Ctor, getter/setter, and native attrs for all policy strings."""
+    from cuda.bindings import driver
+    from cuda.core import _launch_config as _lc_mod
+    from cuda.core._launch_config import _to_native_launch_config
+
+    class _FakeDev:
+        compute_capability = (9, 0)
+
+    monkeypatch.setattr(_lc_mod, "Device", lambda: _FakeDev())
+    pref = driver.CUlaunchAttributeID.CU_LAUNCH_ATTRIBUTE_CLUSTER_SCHEDULING_POLICY_PREFERENCE
+
+    for policy in LaunchConfig._CLUSTER_SCHED_POLICY_TO_DRIVER:
+        cfg = LaunchConfig(grid=1, block=1)
+        assert cfg.cluster_scheduling_policy_preference is None
+        cfg.cluster_scheduling_policy_preference = policy
+        assert cfg.cluster_scheduling_policy_preference is policy
+
+        cfg = LaunchConfig(grid=2, block=4, cluster_scheduling_policy_preference=policy)
+        assert cfg.cluster_scheduling_policy_preference is policy
+        native = _to_native_launch_config(cfg)
+        assert native.numAttrs == 1
+        assert native.attrs[0].id == pref
+        assert int(native.attrs[0].value.clusterSchedulingPolicyPreference) == int(
+            getattr(driver.CUclusterSchedulingPolicy, f"CU_CLUSTER_SCHEDULING_POLICY_{policy}")
+        )
+        cfg.cluster_scheduling_policy_preference = None
+        assert cfg.cluster_scheduling_policy_preference is None
+
+    cfg = LaunchConfig(
+        grid=(2, 1, 1),
+        block=32,
+        cluster=(2, 1, 1),
+        cluster_scheduling_policy_preference="SPREAD",
+    )
+    native = _to_native_launch_config(cfg)
+    assert native.numAttrs == 2
+    attr_ids = {attr.id for attr in native.attrs}
+    assert driver.CUlaunchAttributeID.CU_LAUNCH_ATTRIBUTE_CLUSTER_DIMENSION in attr_ids
+    assert pref in attr_ids
+
+
+@pytest.mark.agent_authored(model="cursor-grok-4.6")
+def test_launch_config_cluster_scheduling_policy_rejected(monkeypatch):
+    """Invalid values and pre-Hopper devices are rejected."""
+    from cuda.core import _launch_config as _lc_mod
+
+    with pytest.raises(ValueError, match="not a valid cluster_scheduling_policy_preference"):
+        LaunchConfig(grid=1, block=1, cluster_scheduling_policy_preference="NOT_A_POLICY")
+
+    class _FakeDev:
+        compute_capability = (8, 6)
+
+    looked_up = []
+    monkeypatch.setattr(_lc_mod, "Device", lambda: looked_up.append(1) or _FakeDev())
+    with pytest.raises(CUDAError, match="cluster launch attributes are not supported"):
+        LaunchConfig(
+            grid=2,
+            block=32,
+            cluster_scheduling_policy_preference="SPREAD",
+        )
+    assert looked_up, "Device was not looked up via the module global; mock did not take effect"
+
+
+@pytest.mark.agent_authored(model="cursor-grok-4.6")
+def test_launch_cluster_scheduling_policy_smoke(init_cuda):
+    """launch() accepts each policy on Hopper+ (skip on CC < 9.0)."""
+    dev = Device()
+    if dev.compute_capability < (9, 0):
+        pytest.skip("Cluster scheduling policy requires compute capability >= 9.0")
+
+    prog = Program('extern "C" __global__ void noop() {}', SourceCodeType.CXX)
+    kernel = prog.compile(ObjectCodeFormatType.CUBIN).get_kernel("noop")
+    stream = dev.default_stream
+    for policy in LaunchConfig._CLUSTER_SCHED_POLICY_TO_DRIVER:
+        launch(
+            stream,
+            LaunchConfig(
+                grid=1,
+                block=32,
+                cluster=(2, 1, 1),
+                cluster_scheduling_policy_preference=policy,
+            ),
+            kernel,
+        )
+    stream.sync()
+
+
 def test_launch_invalid_values(init_cuda):
     code = 'extern "C" __global__ void my_kernel() {}'
     program = Program(code, SourceCodeType.CXX)

cuda_core/tests/test_object_protocols.py

@@ -704,7 +704,8 @@ def sample_object_b(request):
         "sample_launch_config",
         r"LaunchConfig\(grid=\(\d+, \d+, \d+\), cluster=.+, block=\(\d+, \d+, \d+\), "
         r"shmem_size=\d+, is_cooperative=(?:True|False), "
-        r"programmatic_stream_serialization=(?:True|False), priority=-?\d+\)",
+        r"programmatic_stream_serialization=(?:True|False), "
+        r"cluster_scheduling_policy_preference=.+, priority=-?\d+\)",
     ),
     ("sample_kernel", r"<Kernel handle=0x[0-9a-f]+>"),
     # ObjectCode variations (by code_type)