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
 
         vector[cydriver.CUlaunchAttribute] _attrs
         object __weakref__

cuda_core/cuda/core/_launch_config.pyi

@@ -2,7 +2,10 @@
 
 from typing import Any
 
-_LAUNCH_CONFIG_ATTRS = ('grid', 'cluster', 'block', 'shmem_size', 'is_cooperative', 'programmatic_stream_serialization')
+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')
+_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}
 __all__ = ['LaunchConfig']
 
 class LaunchConfig:
@@ -39,15 +42,20 @@ 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 omitted, the driver uses
+        the kernel function's default policy.
     """
     grid: tuple[Any, ...]
     cluster: tuple[Any, ...]
     block: tuple[Any, ...]
     shmem_size: int
     is_cooperative: bool
     programmatic_stream_serialization: bool
+    cluster_scheduling_policy_preference: str
 
-    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) -> 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) -> None:
         """Initialize LaunchConfig with validation.
 
         Parameters
@@ -64,12 +72,17 @@ 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"`` (default: None)
         """
     def _identity(self) -> tuple[Any, ...]: ...
     def __repr__(self) -> str:
         """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

@@ -12,6 +12,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',
@@ -20,8 +21,15 @@ _LAUNCH_CONFIG_ATTRS = (
     'shmem_size',
     'is_cooperative',
     'programmatic_stream_serialization',
+    'cluster_scheduling_policy_preference',
 )
 
+_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,
+}
+
 __all__ = ['LaunchConfig']
 
 
@@ -59,6 +67,10 @@ 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 omitted, the driver uses
+        the kernel function's default policy.
     """
 
     # TODO: expand LaunchConfig to include other attributes
@@ -72,6 +84,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,
     ) -> None:
         """Initialize LaunchConfig with validation.
 
@@ -89,21 +102,31 @@ 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"`` (default: None)
         """
         # Convert and validate grid and block dimensions
         self.grid = cast_to_3_tuple("LaunchConfig.grid", grid)
         self.block = cast_to_3_tuple("LaunchConfig.block", block)
 
+        validated_policy = 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.
         # thread block clusters are supported starting H100
-        if cluster is not None:
+        cc = None
+        if cluster is not None or validated_policy is not None:
             cc = Device().compute_capability
             if cc < (9, 0):
                 raise CUDAError(
-                    f"thread block clusters are not supported on devices with compute capability < 9.0 (got {cc})"
+                    "cluster launch attributes are not supported on devices with "
+                    f"compute capability < 9.0 (got {cc})"
                 )
+        if cluster is not None:
             self.cluster = cast_to_3_tuple("LaunchConfig.cluster", cluster)
         else:
             self.cluster = None
@@ -116,6 +139,7 @@ cdef class LaunchConfig:
 
         self.is_cooperative = is_cooperative
         self.programmatic_stream_serialization = programmatic_stream_serialization
+        self.cluster_scheduling_policy_preference = validated_policy
 
         if self.is_cooperative and not Device().properties.cooperative_launch:
             raise CUDAError("cooperative kernels are not supported on this device")
@@ -136,6 +160,25 @@ 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 _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(_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 _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
@@ -169,6 +212,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)
+
         drv_cfg.numAttrs = self._attrs.size()
         drv_cfg.attrs = self._attrs.data()
 
@@ -230,6 +280,13 @@ 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
+        # 13.0.2 setter reads .value; pass FastEnum, not a raw int.
+        attr.value.clusterSchedulingPolicyPreference = config._cluster_sched_policy_driver_value()
+        attrs.append(attr)
+
     drv_cfg.numAttrs = len(attrs)
     drv_cfg.attrs = attrs
 

cuda_core/tests/test_launcher.py

@@ -324,7 +324,7 @@ class _FakeDev:
     looked_up = []
     monkeypatch.setattr(_lc_mod, "Device", lambda: looked_up.append(1) or _FakeDev())
 
-    with pytest.raises(CUDAError, match="thread block clusters are not supported"):
+    with pytest.raises(CUDAError, match="cluster launch attributes are not supported"):
         LaunchConfig(grid=2, cluster=2, block=32)
     assert looked_up, "Device was not looked up via the module global; mock did not take effect"
 
@@ -364,6 +364,97 @@ 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 _CLUSTER_SCHED_POLICY_TO_DRIVER, _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 _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)."""
+    from cuda.core._launch_config import _CLUSTER_SCHED_POLICY_TO_DRIVER
+
+    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 _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)\)",
+        r"programmatic_stream_serialization=(?:True|False), "
+        r"cluster_scheduling_policy_preference=.+\)",
     ),
     ("sample_kernel", r"<Kernel handle=0x[0-9a-f]+>"),
     # ObjectCode variations (by code_type)