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

@@ -2,7 +2,9 @@
 
 from typing import Any
 
-_LAUNCH_CONFIG_ATTRS = ('grid', 'cluster', 'block', 'shmem_size', 'is_cooperative', 'programmatic_stream_serialization')
+from cuda.core.typing import ClusterSchedulingPolicyType
+
+_LAUNCH_CONFIG_ATTRS = ('grid', 'cluster', 'block', 'shmem_size', 'is_cooperative', 'programmatic_stream_serialization', 'cluster_scheduling_policy_preference')
 __all__ = ['LaunchConfig']
 
 class LaunchConfig:
@@ -39,15 +41,19 @@ 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.
     """
     grid: tuple[Any, ...]
     cluster: tuple[Any, ...]
     block: tuple[Any, ...]
     shmem_size: int
     is_cooperative: bool
     programmatic_stream_serialization: bool
+    cluster_scheduling_policy_preference: object
 
-    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
@@ -64,13 +70,16 @@ 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, ...]: ...
     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(value): ...
 def _to_native_launch_config(config: LaunchConfig) -> object:
     """Convert LaunchConfig to native driver CUlaunchConfig.
 

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/docs/source/api_private.rst

@@ -23,6 +23,7 @@ CUDA runtime
    _module.ParamInfo
    typing.AddressModeType
    typing.ArrayFormatType
+   typing.ClusterSchedulingPolicyType
    typing.CompilerBackendType
    typing.DevicePointerType
    typing.DeviceResourcesType