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

cuda_core/cuda/core/_launch_config.pyi

@@ -4,6 +4,8 @@ from __future__ import annotations
 
 from typing import Any
 
+from cuda.core.typing import ClusterSchedulingPolicyType
+
 
 class LaunchConfig:
     """Customizable launch options.
@@ -39,9 +41,12 @@ 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 : ClusterSchedulingPolicyType, optional
+        Cluster scheduling policy for the launch. When omitted, the driver uses
+        the kernel function's default policy.
     """
 
-    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: ClusterSchedulingPolicyType | None=None) -> None:
         """Initialize LaunchConfig with validation.
 
         Parameters
@@ -58,6 +63,8 @@ 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 : ClusterSchedulingPolicyType, optional
+            Cluster scheduling policy for the launch (default: None)
         """
 
     def _identity(self) -> tuple[Any, ...]:
@@ -71,7 +78,7 @@ class LaunchConfig:
 
     def __hash__(self) -> int:
         ...
-_LAUNCH_CONFIG_ATTRS = ('grid', 'cluster', 'block', 'shmem_size', 'is_cooperative', 'programmatic_stream_serialization')
+_LAUNCH_CONFIG_ATTRS = ('grid', 'cluster', 'block', 'shmem_size', 'is_cooperative', 'programmatic_stream_serialization', 'cluster_scheduling_policy_preference')
 __all__ = ['LaunchConfig']
 
 def _to_native_launch_config(config: LaunchConfig) -> object:
@@ -86,4 +93,7 @@ def _to_native_launch_config(config: LaunchConfig) -> object:
     -------
     driver.CUlaunchConfig
         Native CUDA driver launch configuration
-    """
+    """
+
+def _validate_cluster_scheduling_policy_preference(value):
+    ...

cuda_core/cuda/core/_launch_config.pyx

@@ -12,6 +12,8 @@ from cuda.core._utils.cuda_utils import (
     cast_to_3_tuple,
     driver,
 )
+from cuda.core._utils.validators import format_or_list
+from cuda.core.typing import ClusterSchedulingPolicyType
 
 _LAUNCH_CONFIG_ATTRS = (
     'grid',
@@ -20,8 +22,23 @@ _LAUNCH_CONFIG_ATTRS = (
     'shmem_size',
     'is_cooperative',
     'programmatic_stream_serialization',
+    'cluster_scheduling_policy_preference',
 )
 
+
+def _validate_cluster_scheduling_policy_preference(value):
+    if value is None:
+        return None
+    if isinstance(value, ClusterSchedulingPolicyType):
+        return value
+    try:
+        return ClusterSchedulingPolicyType(int(value))
+    except (TypeError, ValueError):
+        valid = format_or_list(ClusterSchedulingPolicyType)
+        raise ValueError(
+            f"{value!r} is not a valid ClusterSchedulingPolicyType. Must be {valid}"
+        ) from None
+
 __all__ = ['LaunchConfig']
 
 
@@ -59,6 +76,9 @@ 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 : ClusterSchedulingPolicyType, optional
+        Cluster scheduling policy for the launch. When omitted, the driver uses
+        the kernel function's default policy.
     """
 
     # TODO: expand LaunchConfig to include other attributes
@@ -72,6 +92,7 @@ cdef class LaunchConfig:
         shmem_size: int | None = None,
         is_cooperative: bool = False,
         programmatic_stream_serialization: bool = False,
+        cluster_scheduling_policy_preference: ClusterSchedulingPolicyType | None = None,
     ) -> None:
         """Initialize LaunchConfig with validation.
 
@@ -89,21 +110,30 @@ 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 : ClusterSchedulingPolicyType, optional
+            Cluster scheduling policy for the launch (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 = _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 +146,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")
@@ -169,6 +200,11 @@ 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_scheduling_policy_preference)
+            self._attrs.push_back(attr)
+
         drv_cfg.numAttrs = self._attrs.size()
         drv_cfg.attrs = self._attrs.data()
 
@@ -230,6 +266,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 = int(config.cluster_scheduling_policy_preference)
+        attrs.append(attr)
+
     drv_cfg.numAttrs = len(attrs)
     drv_cfg.attrs = attrs
 

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
@@ -34,6 +35,7 @@ class StrEnum(str, Enum):
 __all__ = [
     "AddressModeType",
     "ArrayFormatType",
+    "ClusterSchedulingPolicyType",
     "CompilerBackendType",
     "DevicePointerType",
     "DeviceResourcesType",
@@ -64,6 +66,22 @@ class StrEnum(str, Enum):
 ProcessStateType = _Literal["running", "locked", "checkpointed", "failed"]
 
 
+class ClusterSchedulingPolicyType(IntEnum):
+    """Cluster scheduling policy for :class:`~cuda.core.LaunchConfig`.
+
+    Corresponds to ``CUclusterSchedulingPolicy`` from the CUDA driver API.
+    Valid for graph nodes and kernel launches on Hopper+ (compute capability >= 9.0).
+
+    * ``DEFAULT`` — driver default scheduling within a cluster.
+    * ``SPREAD`` — spread blocks within a cluster across SMs.
+    * ``LOAD_BALANCING`` — allow hardware load-balancing of cluster blocks.
+    """
+
+    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
+
+
 class SourceCodeType(StrEnum):
     """Source language passed to :class:`~cuda.core.Program`.
 

cuda_core/tests/test_launcher.py

@@ -26,7 +26,7 @@
 )
 from cuda.core._memory._legacy import _SynchronousMemoryResource
 from cuda.core._utils.cuda_utils import CUDAError
-from cuda.core.typing import ObjectCodeFormatType, SourceCodeType
+from cuda.core.typing import ClusterSchedulingPolicyType, ObjectCodeFormatType, SourceCodeType
 
 
 def test_launch_config_init(init_cuda):
@@ -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,140 @@ 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.parametrize(
+    "policy",
+    [
+        ClusterSchedulingPolicyType.DEFAULT,
+        ClusterSchedulingPolicyType.SPREAD,
+        ClusterSchedulingPolicyType.LOAD_BALANCING,
+    ],
+)
+@pytest.mark.agent_authored(model="composer-2.5-fast")
+def test_to_native_launch_config_cluster_scheduling_policy(monkeypatch, policy):
+    """LaunchConfig(cluster_scheduling_policy_preference=...) maps to the native attribute."""
+    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())
+
+    config = LaunchConfig(
+        grid=2,
+        block=4,
+        cluster_scheduling_policy_preference=policy,
+    )
+    native = _to_native_launch_config(config)
+    assert native.numAttrs == 1
+    attr = native.attrs[0]
+    assert attr.id == driver.CUlaunchAttributeID.CU_LAUNCH_ATTRIBUTE_CLUSTER_SCHEDULING_POLICY_PREFERENCE, (
+        f"Expected CU_LAUNCH_ATTRIBUTE_CLUSTER_SCHEDULING_POLICY_PREFERENCE, got {attr.id}"
+    )
+    assert int(attr.value.clusterSchedulingPolicyPreference) == int(policy), (
+        f"Expected clusterSchedulingPolicyPreference={int(policy)!r}, "
+        f"got {int(attr.value.clusterSchedulingPolicyPreference)!r}"
+    )
+
+
+@pytest.mark.agent_authored(model="composer-2.5-fast")
+def test_to_native_launch_config_cluster_scheduling_policy_accepts_driver_enum(monkeypatch):
+    """LaunchConfig accepts cuda.bindings.driver.CUclusterSchedulingPolicy values."""
+    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())
+
+    config = LaunchConfig(
+        grid=1,
+        block=1,
+        cluster_scheduling_policy_preference=driver.CUclusterSchedulingPolicy.CU_CLUSTER_SCHEDULING_POLICY_SPREAD,
+    )
+    assert config.cluster_scheduling_policy_preference == ClusterSchedulingPolicyType.SPREAD
+    native = _to_native_launch_config(config)
+    assert native.numAttrs == 1
+    assert int(native.attrs[0].value.clusterSchedulingPolicyPreference) == int(ClusterSchedulingPolicyType.SPREAD)
+
+
+@pytest.mark.agent_authored(model="composer-2.5-fast")
+def test_launch_config_cluster_scheduling_policy_invalid():
+    """LaunchConfig rejects invalid cluster scheduling policy values."""
+    with pytest.raises(ValueError, match="not a valid ClusterSchedulingPolicyType"):
+        LaunchConfig(grid=1, block=1, cluster_scheduling_policy_preference=999)
+
+
+@pytest.mark.agent_authored(model="composer-2.5-fast")
+def test_launch_config_cluster_scheduling_policy_rejects_pre_hopper_cc(monkeypatch):
+    """LaunchConfig(cluster_scheduling_policy_preference=...) raises on CC < 9.0."""
+    from cuda.core import _launch_config as _lc_mod
+
+    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=ClusterSchedulingPolicyType.SPREAD,
+        )
+    assert looked_up, "Device was not looked up via the module global; mock did not take effect"
+
+
+@pytest.mark.agent_authored(model="composer-2.5-fast")
+def test_to_native_launch_config_cluster_and_policy(monkeypatch):
+    """Cluster dimension and scheduling policy produce two launch attributes."""
+    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())
+
+    config = LaunchConfig(
+        grid=(2, 1, 1),
+        block=32,
+        cluster=(2, 1, 1),
+        cluster_scheduling_policy_preference=ClusterSchedulingPolicyType.SPREAD,
+    )
+    native = _to_native_launch_config(config)
+    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 driver.CUlaunchAttributeID.CU_LAUNCH_ATTRIBUTE_CLUSTER_SCHEDULING_POLICY_PREFERENCE in attr_ids
+
+
+@pytest.mark.agent_authored(model="composer-2.5-fast")
+def test_launch_cluster_scheduling_policy_smoke(init_cuda):
+    """Smoke-test launching with cluster scheduling policy on Hopper+."""
+    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)
+    mod = prog.compile(ObjectCodeFormatType.CUBIN)
+    kernel = mod.get_kernel("noop")
+    stream = dev.default_stream
+
+    launch_config = LaunchConfig(
+        grid=1,
+        block=32,
+        cluster=(2, 1, 1),
+        cluster_scheduling_policy_preference=ClusterSchedulingPolicyType.LOAD_BALANCING,
+    )
+    launch(stream, launch_config, 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

@@ -685,7 +685,8 @@ def sample_switch_node_alt(sample_graphdef):
         "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)