diff --git a/cuda_core/cuda/core/_launch_config.pxd b/cuda_core/cuda/core/_launch_config.pxd index 892a73f8efc..b26610374b0 100644 --- a/cuda_core/cuda/core/_launch_config.pxd +++ b/cuda_core/cuda/core/_launch_config.pxd @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # SPDX-License-Identifier: Apache-2.0 @@ -16,6 +16,7 @@ cdef class LaunchConfig: public int shmem_size public bint is_cooperative public bint programmatic_stream_serialization + public int priority 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..3a778202768 100644 --- a/cuda_core/cuda/core/_launch_config.pyi +++ b/cuda_core/cuda/core/_launch_config.pyi @@ -2,7 +2,7 @@ from typing import Any -_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', 'priority') __all__ = ['LaunchConfig'] class LaunchConfig: @@ -39,6 +39,15 @@ 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. + 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. """ grid: tuple[Any, ...] cluster: tuple[Any, ...] @@ -46,8 +55,9 @@ class LaunchConfig: shmem_size: int is_cooperative: bool programmatic_stream_serialization: bool + 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) -> 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, priority: int | None=None) -> None: """Initialize LaunchConfig with validation. Parameters @@ -64,6 +74,15 @@ class LaunchConfig: Whether to launch as cooperative kernel (default: False) programmatic_stream_serialization : bool, optional Whether to allow programmatic stream serialization / PDL (default: False) + 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. """ def _identity(self) -> tuple[Any, ...]: ... def __repr__(self) -> str: diff --git a/cuda_core/cuda/core/_launch_config.pyx b/cuda_core/cuda/core/_launch_config.pyx index adbf9a16c5d..2f665369296 100644 --- a/cuda_core/cuda/core/_launch_config.pyx +++ b/cuda_core/cuda/core/_launch_config.pyx @@ -7,6 +7,7 @@ from libc.string cimport memset from typing import Any from cuda.core._device import Device +from cuda.core._utils.cuda_utils cimport HANDLE_RETURN from cuda.core._utils.cuda_utils import ( CUDAError, cast_to_3_tuple, @@ -20,6 +21,7 @@ _LAUNCH_CONFIG_ATTRS = ( 'shmem_size', 'is_cooperative', 'programmatic_stream_serialization', + 'priority', ) __all__ = ['LaunchConfig'] @@ -59,6 +61,15 @@ 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. + 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. """ # TODO: expand LaunchConfig to include other attributes @@ -72,6 +83,7 @@ cdef class LaunchConfig: shmem_size: int | None = None, is_cooperative: bool = False, programmatic_stream_serialization: bool = False, + priority: int | None = None, ) -> None: """Initialize LaunchConfig with validation. @@ -89,6 +101,15 @@ 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) + 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. """ # Convert and validate grid and block dimensions self.grid = cast_to_3_tuple("LaunchConfig.grid", grid) @@ -117,6 +138,26 @@ cdef class LaunchConfig: self.is_cooperative = is_cooperative self.programmatic_stream_serialization = programmatic_stream_serialization + # priority=0 is treated the same as an unset priority (see + # _to_native_launch_config), so only nonzero values need validating + # against the device's stream priority range. + cdef int high, low + cdef cydriver.CUresult res_code + cdef int prio + if priority: + with nogil: + res_code = cydriver.cuCtxGetStreamPriorityRange(&high, &low) + if res_code != cydriver.CUresult.CUDA_SUCCESS: + if res_code == cydriver.CUresult.CUDA_ERROR_INVALID_CONTEXT: + raise RuntimeError( + "No current CUDA context. Call dev.set_current() before creating a LaunchConfig with a priority." + ) + HANDLE_RETURN(res_code) + prio = priority + if not (low <= prio <= high): + raise ValueError(f"{priority=} is out of range {[low, high]}") + self.priority = prio + if self.is_cooperative and not Device().properties.cooperative_launch: raise CUDAError("cooperative kernels are not supported on this device") @@ -169,6 +210,11 @@ cdef class LaunchConfig: attr.value.programmaticStreamSerializationAllowed = 1 self._attrs.push_back(attr) + if self.priority: + attr.id = cydriver.CUlaunchAttributeID.CU_LAUNCH_ATTRIBUTE_PRIORITY + attr.value.priority = self.priority + self._attrs.push_back(attr) + drv_cfg.numAttrs = self._attrs.size() drv_cfg.attrs = self._attrs.data() @@ -230,6 +276,12 @@ cpdef object _to_native_launch_config(LaunchConfig config): attr.value.programmaticStreamSerializationAllowed = 1 attrs.append(attr) + if config.priority: + attr = driver.CUlaunchAttribute() + attr.id = driver.CUlaunchAttributeID.CU_LAUNCH_ATTRIBUTE_PRIORITY + attr.value.priority = config.priority + attrs.append(attr) + drv_cfg.numAttrs = len(attrs) drv_cfg.attrs = attrs diff --git a/cuda_core/tests/test_launcher.py b/cuda_core/tests/test_launcher.py index 2ab766cc2f0..7d3e94ee04d 100644 --- a/cuda_core/tests/test_launcher.py +++ b/cuda_core/tests/test_launcher.py @@ -202,6 +202,47 @@ def test_to_native_launch_config_pdl(): ) +@pytest.mark.parametrize( + ("initial_priority", "updated_priority"), + ((-1, 0), (0, -1), (0, 5)), +) +def test_launch_config_priority_getter_setter(init_cuda, initial_priority, updated_priority): + """Direct attribute assignment (unlike __init__) is not range-checked.""" + config = LaunchConfig(grid=1, block=1, priority=initial_priority) + + assert config.priority == initial_priority + config.priority = updated_priority + assert config.priority == updated_priority + + +@pytest.mark.parametrize( + ("priority", "expected_num_attrs"), + ((None, 0), (0, 0), (-1, 1)), +) +def test_to_native_launch_config_priority(init_cuda, priority, expected_num_attrs): + """LaunchConfig priority maps to the native attribute for nonzero values. + + priority=0 (and the None default, which is stored as 0) is treated the + same as unset (numAttrs=0), matching the truthy check used both here and + in the bound LaunchConfig._to_native_launch_config method. + """ + from cuda.bindings import driver + from cuda.core._launch_config import _to_native_launch_config + + config = LaunchConfig(grid=2, block=4, priority=priority) + native = _to_native_launch_config(config) + + assert config.priority == (priority or 0) + assert native.numAttrs == expected_num_attrs + if expected_num_attrs == 0: + assert list(native.attrs) == [] + return + + attr = native.attrs[0] + assert attr.id == driver.CUlaunchAttributeID.CU_LAUNCH_ATTRIBUTE_PRIORITY + assert attr.value.priority == priority + + @skipif_need_cuda_headers def test_pdl_primary_secondary_overlap_same_stream(): """Primary + secondary PDL launch on one stream can overlap on Hopper+. diff --git a/cuda_core/tests/test_object_protocols.py b/cuda_core/tests/test_object_protocols.py index a5b30e9e5ba..0e6af02e128 100644 --- a/cuda_core/tests/test_object_protocols.py +++ b/cuda_core/tests/test_object_protocols.py @@ -704,7 +704,7 @@ 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), priority=-?\d+\)", ), ("sample_kernel", r""), # ObjectCode variations (by code_type)