diff --git a/cuda_core/cuda/core/_launch_config.pxd b/cuda_core/cuda/core/_launch_config.pxd index b26610374b0..dad1f1a9e9f 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 str cluster_scheduling_policy_preference public int priority vector[cydriver.CUlaunchAttribute] _attrs diff --git a/cuda_core/cuda/core/_launch_config.pyi b/cuda_core/cuda/core/_launch_config.pyi index 5fd38195318..e970adcbc53 100644 --- a/cuda_core/cuda/core/_launch_config.pyi +++ b/cuda_core/cuda/core/_launch_config.pyi @@ -2,62 +2,24 @@ 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: - """Customizable launch options. - - Note - ---- - When cluster is specified, the grid parameter represents the number of - clusters (not blocks). The hierarchy is: grid (clusters) -> cluster (blocks) -> - block (threads). Each dimension in grid specifies clusters in the grid, each dimension in - cluster specifies blocks per cluster, and each dimension in block specifies - threads per block. - - Attributes - ---------- - grid : tuple | int - Collection of threads that will execute a kernel function. When cluster - is not specified, this represents the number of blocks, otherwise - this represents the number of clusters. - cluster : tuple | int - Group of blocks (Thread Block Cluster) that will execute on the same - GPU Processing Cluster (GPC). Blocks within a cluster have access to - distributed shared memory and can be explicitly synchronized. - block : tuple | int - Group of threads (Thread Block) that will execute on the same - streaming multiprocessor (SM). Threads within a thread blocks have - access to shared memory and can be explicitly synchronized. - shmem_size : int, optional - Dynamic shared-memory size per thread block in bytes. - (Default to size 0) - is_cooperative : bool, optional - Whether this config can be used to launch a cooperative kernel. - programmatic_stream_serialization : bool, optional - 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. - priority : int, optional - Execution priority of the kernel. Lower numbers represent higher - priorities. The meaningful range of values is device-specific, - given by ``[greatestPriority, leastPriority]`` as returned by - ``cuCtxGetStreamPriorityRange`` (the same range used by - :attr:`~cuda.core.StreamOptions.priority`); both bounds are 0 on - a device that does not support multiple stream priorities. A - nonzero value outside this range raises :class:`ValueError`. - When omitted (or 0), the launch uses the stream's priority. - """ + """Customizable launch options.""" + _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 +36,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 +56,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. diff --git a/cuda_core/cuda/core/_launch_config.pyx b/cuda_core/cuda/core/_launch_config.pyx index 2f665369296..e6c29aa9064 100644 --- a/cuda_core/cuda/core/_launch_config.pyx +++ b/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', ) @@ -28,6 +30,11 @@ __all__ = ['LaunchConfig'] cdef class LaunchConfig: + _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, + } """Customizable launch options. Note @@ -61,6 +68,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, @@ -83,6 +97,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 +116,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 +135,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 +203,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 +257,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 +330,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 diff --git a/cuda_core/tests/test_launcher.py b/cuda_core/tests/test_launcher.py index 8f83435764c..66d9c7b99b7 100644 --- a/cuda_core/tests/test_launcher.py +++ b/cuda_core/tests/test_launcher.py @@ -365,7 +365,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" @@ -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) diff --git a/cuda_core/tests/test_object_protocols.py b/cuda_core/tests/test_object_protocols.py index 12492b8ce33..2bfe660acba 100644 --- a/cuda_core/tests/test_object_protocols.py +++ b/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""), # ObjectCode variations (by code_type)