diff --git a/cuda_core/cuda/core/_launch_config.pxd b/cuda_core/cuda/core/_launch_config.pxd index 892a73f8efc..5f3681ce5ac 100644 --- a/cuda_core/cuda/core/_launch_config.pxd +++ b/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__ diff --git a/cuda_core/cuda/core/_launch_config.pyi b/cuda_core/cuda/core/_launch_config.pyi index 47187fb03d6..36686dc172b 100644 --- a/cuda_core/cuda/core/_launch_config.pyi +++ b/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 - """ \ No newline at end of file + """ + +def _validate_cluster_scheduling_policy_preference(value): + ... \ No newline at end of file diff --git a/cuda_core/cuda/core/_launch_config.pyx b/cuda_core/cuda/core/_launch_config.pyx index adbf9a16c5d..312f8bc4d3a 100644 --- a/cuda_core/cuda/core/_launch_config.pyx +++ b/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 diff --git a/cuda_core/cuda/core/typing.py b/cuda_core/cuda/core/typing.py index 1bf9bb7c0d2..2980341ce22 100644 --- a/cuda_core/cuda/core/typing.py +++ b/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`. diff --git a/cuda_core/tests/test_launcher.py b/cuda_core/tests/test_launcher.py index e5cf05b435d..0f2ff6767ff 100644 --- a/cuda_core/tests/test_launcher.py +++ b/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) diff --git a/cuda_core/tests/test_object_protocols.py b/cuda_core/tests/test_object_protocols.py index baf790abea8..ad17cebbe04 100644 --- a/cuda_core/tests/test_object_protocols.py +++ b/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""), # ObjectCode variations (by code_type)