From 3374af49277d2f90be0f5862413013fabde9ca07 Mon Sep 17 00:00:00 2001 From: Omar Atie Date: Tue, 25 Aug 2026 17:19:30 -0700 Subject: [PATCH 1/6] feat(cuda.core): add cluster scheduling policy to LaunchConfig Expose CU_LAUNCH_ATTRIBUTE_CLUSTER_SCHEDULING_POLICY_PREFERENCE on LaunchConfig via ClusterSchedulingPolicyType, with validation, CC >= 9.0 gating, and tests mapping to the native launch attribute. Closes #2629 Co-authored-by: Cursor Signed-off-by: Omar Atie --- cuda_core/cuda/core/_launch_config.pxd | 1 + cuda_core/cuda/core/_launch_config.pyi | 13 ++- cuda_core/cuda/core/_launch_config.pyx | 46 ++++++++- cuda_core/cuda/core/typing.py | 18 ++++ cuda_core/tests/test_launcher.py | 118 ++++++++++++++++++++++- cuda_core/tests/test_object_protocols.py | 3 +- 6 files changed, 192 insertions(+), 7 deletions(-) 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 a731f2999ff..85ed687d48a 100644 --- a/cuda_core/cuda/core/_launch_config.pyi +++ b/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,6 +41,9 @@ 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, ...] @@ -46,8 +51,9 @@ class LaunchConfig: 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,6 +70,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, ...]: ... def __repr__(self) -> str: @@ -71,6 +79,7 @@ class 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. 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 2ab766cc2f0..eb9243672a7 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,120 @@ 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_launch_cluster_scheduling_policy_smoke(get_saxpy_kernel_cubin): + """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") + + kernel, _ = get_saxpy_kernel_cubin + stream = dev.default_stream + n = np.int32(4) + a = np.float32(2.0) + x = np.from_dlpack(dev.allocate(16, stream=stream)).view(np.float32) + y = np.from_dlpack(dev.allocate(16, stream=stream)).view(np.float32) + x[:] = 1.0 + y[:] = 0.0 + + launch_config = LaunchConfig( + grid=1, + block=32, + cluster=(2, 1, 1), + cluster_scheduling_policy_preference=ClusterSchedulingPolicyType.LOAD_BALANCING, + ) + launch(stream, launch_config, kernel, n, a, x.ctypes.data, y.ctypes.data) + stream.sync() + np.testing.assert_allclose(y, 2.0) + + 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 a5b30e9e5ba..17ef2a879b8 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)\)", + r"programmatic_stream_serialization=(?:True|False), " + r"cluster_scheduling_policy_preference=.+\)", ), ("sample_kernel", r""), # ObjectCode variations (by code_type) From b6b0fbbc831d119c4986b4b30a32576c4772c6b0 Mon Sep 17 00:00:00 2001 From: Omar Atie Date: Tue, 25 Aug 2026 17:19:31 -0700 Subject: [PATCH 2/6] test(cuda.core): cover cluster dimension + scheduling policy attrs Assert LaunchConfig emits both CU_LAUNCH_ATTRIBUTE_CLUSTER_DIMENSION and CU_LAUNCH_ATTRIBUTE_CLUSTER_SCHEDULING_POLICY_PREFERENCE when set. Co-authored-by: Cursor Signed-off-by: Omar Atie --- cuda_core/tests/test_launcher.py | 25 +++++++++++++++++++++++++ 1 file changed, 25 insertions(+) diff --git a/cuda_core/tests/test_launcher.py b/cuda_core/tests/test_launcher.py index eb9243672a7..ba1cfb41b1d 100644 --- a/cuda_core/tests/test_launcher.py +++ b/cuda_core/tests/test_launcher.py @@ -451,6 +451,31 @@ class _FakeDev: 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(get_saxpy_kernel_cubin): """Smoke-test launching with cluster scheduling policy on Hopper+.""" From 18c95f2210bdbff698092a68ee8e0f9d4b32eeba Mon Sep 17 00:00:00 2001 From: Omar Atie Date: Tue, 25 Aug 2026 17:19:31 -0700 Subject: [PATCH 3/6] test(cuda.core): fix cluster scheduling policy smoke test fixture Use init_cuda with an inline noop kernel so test_launcher.py does not depend on the get_saxpy_kernel_cubin fixture from test_module.py. Co-authored-by: Cursor Signed-off-by: Omar Atie --- cuda_core/tests/test_launcher.py | 15 +++++---------- 1 file changed, 5 insertions(+), 10 deletions(-) diff --git a/cuda_core/tests/test_launcher.py b/cuda_core/tests/test_launcher.py index ba1cfb41b1d..b3616ef3ee8 100644 --- a/cuda_core/tests/test_launcher.py +++ b/cuda_core/tests/test_launcher.py @@ -477,20 +477,16 @@ class _FakeDev: @pytest.mark.agent_authored(model="composer-2.5-fast") -def test_launch_cluster_scheduling_policy_smoke(get_saxpy_kernel_cubin): +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") - kernel, _ = get_saxpy_kernel_cubin + prog = Program('extern "C" __global__ void noop() {}', SourceCodeType.CXX) + mod = prog.compile(ObjectCodeFormatType.CUBIN) + kernel = mod.get_kernel("noop") stream = dev.default_stream - n = np.int32(4) - a = np.float32(2.0) - x = np.from_dlpack(dev.allocate(16, stream=stream)).view(np.float32) - y = np.from_dlpack(dev.allocate(16, stream=stream)).view(np.float32) - x[:] = 1.0 - y[:] = 0.0 launch_config = LaunchConfig( grid=1, @@ -498,9 +494,8 @@ def test_launch_cluster_scheduling_policy_smoke(get_saxpy_kernel_cubin): cluster=(2, 1, 1), cluster_scheduling_policy_preference=ClusterSchedulingPolicyType.LOAD_BALANCING, ) - launch(stream, launch_config, kernel, n, a, x.ctypes.data, y.ctypes.data) + launch(stream, launch_config, kernel) stream.sync() - np.testing.assert_allclose(y, 2.0) def test_launch_invalid_values(init_cuda): From c6c283bb5d3fce902539bfeb05960c11d61e9e99 Mon Sep 17 00:00:00 2001 From: Omar Atie Date: Tue, 25 Aug 2026 17:19:31 -0700 Subject: [PATCH 4/6] test(cuda.core): cover cluster policy getter/setter and launch run-through Align #2629 tests with reviewer guidance: round-trip each policy on LaunchConfig and exercise launch() for DEFAULT/SPREAD/LOAD_BALANCING. Co-authored-by: Cursor Signed-off-by: Omar Atie --- cuda_core/tests/test_launcher.py | 45 +++++++++++++++++++++++++++++--- 1 file changed, 41 insertions(+), 4 deletions(-) diff --git a/cuda_core/tests/test_launcher.py b/cuda_core/tests/test_launcher.py index b3616ef3ee8..96d1a55d79f 100644 --- a/cuda_core/tests/test_launcher.py +++ b/cuda_core/tests/test_launcher.py @@ -476,9 +476,46 @@ class _FakeDev: 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+.""" +@pytest.mark.parametrize( + "policy", + [ + ClusterSchedulingPolicyType.DEFAULT, + ClusterSchedulingPolicyType.SPREAD, + ClusterSchedulingPolicyType.LOAD_BALANCING, + ], +) +@pytest.mark.agent_authored(model="cursor-grok-4.6") +def test_launch_config_cluster_scheduling_policy_getter_setter(monkeypatch, policy): + """Getter/setter round-trip for each cluster scheduling policy.""" + from cuda.core import _launch_config as _lc_mod + + class _FakeDev: + compute_capability = (9, 0) + + monkeypatch.setattr(_lc_mod, "Device", lambda: _FakeDev()) + + 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 + + cfg2 = LaunchConfig(grid=1, block=1, cluster_scheduling_policy_preference=policy) + assert cfg2.cluster_scheduling_policy_preference is policy + cfg2.cluster_scheduling_policy_preference = None + assert cfg2.cluster_scheduling_policy_preference is None + + +@pytest.mark.parametrize( + "policy", + [ + ClusterSchedulingPolicyType.DEFAULT, + ClusterSchedulingPolicyType.SPREAD, + ClusterSchedulingPolicyType.LOAD_BALANCING, + ], +) +@pytest.mark.agent_authored(model="cursor-grok-4.6") +def test_launch_cluster_scheduling_policy_smoke(init_cuda, policy): + """Application code runs through launch() for each policy on Hopper+.""" dev = Device() if dev.compute_capability < (9, 0): pytest.skip("Cluster scheduling policy requires compute capability >= 9.0") @@ -492,7 +529,7 @@ def test_launch_cluster_scheduling_policy_smoke(init_cuda): grid=1, block=32, cluster=(2, 1, 1), - cluster_scheduling_policy_preference=ClusterSchedulingPolicyType.LOAD_BALANCING, + cluster_scheduling_policy_preference=policy, ) launch(stream, launch_config, kernel) stream.sync() From 39c9fc0843a82b03b86af59729822b4a5064e87f Mon Sep 17 00:00:00 2001 From: Omar Atie Date: Thu, 27 Aug 2026 14:51:09 -0700 Subject: [PATCH 5/6] docs(cuda.core): document ClusterSchedulingPolicyType in API RST CI test_subpackage_exports_are_documented[typing] failed because the new typing export was missing from docs/source/*.rst. Signed-off-by: Omar Atie --- cuda_core/docs/source/api_private.rst | 1 + 1 file changed, 1 insertion(+) diff --git a/cuda_core/docs/source/api_private.rst b/cuda_core/docs/source/api_private.rst index 80675799c07..001e718e820 100644 --- a/cuda_core/docs/source/api_private.rst +++ b/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 From 1a78ad5cd1cf0294d91994b0214ec71bdf9bdd2f Mon Sep 17 00:00:00 2001 From: Omar Atie Date: Thu, 27 Aug 2026 17:27:38 -0700 Subject: [PATCH 6/6] test(cuda.core): collapse cluster policy tests to three cases Merge seven overlapping LaunchConfig cluster-scheduling tests into mapping, rejection, and Hopper launch smoke to cut maintenance. Signed-off-by: Omar Atie --- cuda_core/tests/test_launcher.py | 184 ++++++++++--------------------- 1 file changed, 57 insertions(+), 127 deletions(-) diff --git a/cuda_core/tests/test_launcher.py b/cuda_core/tests/test_launcher.py index 96d1a55d79f..748e55cf76e 100644 --- a/cuda_core/tests/test_launcher.py +++ b/cuda_core/tests/test_launcher.py @@ -364,46 +364,16 @@ 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, - ], +_CLUSTER_SCHED_POLICIES = ( + 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.""" +@pytest.mark.agent_authored(model="cursor-grok-4.6") +def test_launch_config_cluster_scheduling_policy(monkeypatch): + """Ctor, getter/setter, driver enum, and native attrs for all policies.""" from cuda.bindings import driver from cuda.core import _launch_config as _lc_mod from cuda.core._launch_config import _to_native_launch_config @@ -412,36 +382,56 @@ class _FakeDev: compute_capability = (9, 0) monkeypatch.setattr(_lc_mod, "Device", lambda: _FakeDev()) - - config = LaunchConfig( + pref = driver.CUlaunchAttributeID.CU_LAUNCH_ATTRIBUTE_CLUSTER_SCHEDULING_POLICY_PREFERENCE + + for policy in _CLUSTER_SCHED_POLICIES: + 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(policy) + cfg.cluster_scheduling_policy_preference = None + assert cfg.cluster_scheduling_policy_preference is None + + cfg = 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) - + assert cfg.cluster_scheduling_policy_preference == 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) + cfg = LaunchConfig( + grid=(2, 1, 1), + block=32, + cluster=(2, 1, 1), + cluster_scheduling_policy_preference=ClusterSchedulingPolicyType.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="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.""" +@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 ClusterSchedulingPolicyType"): + LaunchConfig(grid=1, block=1, cluster_scheduling_policy_preference=999) + 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, @@ -451,87 +441,27 @@ class _FakeDev: 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.parametrize( - "policy", - [ - ClusterSchedulingPolicyType.DEFAULT, - ClusterSchedulingPolicyType.SPREAD, - ClusterSchedulingPolicyType.LOAD_BALANCING, - ], -) @pytest.mark.agent_authored(model="cursor-grok-4.6") -def test_launch_config_cluster_scheduling_policy_getter_setter(monkeypatch, policy): - """Getter/setter round-trip for each cluster scheduling policy.""" - from cuda.core import _launch_config as _lc_mod - - class _FakeDev: - compute_capability = (9, 0) - - monkeypatch.setattr(_lc_mod, "Device", lambda: _FakeDev()) - - 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 - - cfg2 = LaunchConfig(grid=1, block=1, cluster_scheduling_policy_preference=policy) - assert cfg2.cluster_scheduling_policy_preference is policy - cfg2.cluster_scheduling_policy_preference = None - assert cfg2.cluster_scheduling_policy_preference is None - - -@pytest.mark.parametrize( - "policy", - [ - ClusterSchedulingPolicyType.DEFAULT, - ClusterSchedulingPolicyType.SPREAD, - ClusterSchedulingPolicyType.LOAD_BALANCING, - ], -) -@pytest.mark.agent_authored(model="cursor-grok-4.6") -def test_launch_cluster_scheduling_policy_smoke(init_cuda, policy): - """Application code runs through launch() for each policy on Hopper+.""" +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) - mod = prog.compile(ObjectCodeFormatType.CUBIN) - kernel = mod.get_kernel("noop") + kernel = prog.compile(ObjectCodeFormatType.CUBIN).get_kernel("noop") stream = dev.default_stream - - launch_config = LaunchConfig( - grid=1, - block=32, - cluster=(2, 1, 1), - cluster_scheduling_policy_preference=policy, - ) - launch(stream, launch_config, kernel) + for policy in _CLUSTER_SCHED_POLICIES: + launch( + stream, + LaunchConfig( + grid=1, + block=32, + cluster=(2, 1, 1), + cluster_scheduling_policy_preference=policy, + ), + kernel, + ) stream.sync()