diff --git a/cuda_core/cuda/core/_launch_config.pxd b/cuda_core/cuda/core/_launch_config.pxd index 892a73f8efc..92a73a0fb6e 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 object 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..271d5206751 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,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. + priority : int, optional + Execution priority of the kernel. Lower numbers represent higher + priorities. When omitted, the launch uses the stream's priority. """ grid: tuple[Any, ...] cluster: tuple[Any, ...] @@ -46,8 +49,9 @@ class LaunchConfig: shmem_size: int is_cooperative: bool programmatic_stream_serialization: bool + priority: 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, priority: int | None=None) -> None: """Initialize LaunchConfig with validation. Parameters @@ -64,6 +68,9 @@ 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. When omitted, 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..f6966a8bc17 100644 --- a/cuda_core/cuda/core/_launch_config.pyx +++ b/cuda_core/cuda/core/_launch_config.pyx @@ -20,6 +20,7 @@ _LAUNCH_CONFIG_ATTRS = ( 'shmem_size', 'is_cooperative', 'programmatic_stream_serialization', + 'priority', ) __all__ = ['LaunchConfig'] @@ -59,6 +60,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. + priority : int, optional + Execution priority of the kernel. Lower numbers represent higher + priorities. When omitted, the launch uses the stream's priority. """ # TODO: expand LaunchConfig to include other attributes @@ -72,6 +76,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 +94,9 @@ 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. When omitted, the launch uses the stream's priority. """ # Convert and validate grid and block dimensions self.grid = cast_to_3_tuple("LaunchConfig.grid", grid) @@ -116,6 +124,7 @@ cdef class LaunchConfig: self.is_cooperative = is_cooperative self.programmatic_stream_serialization = programmatic_stream_serialization + self.priority = priority if self.is_cooperative and not Device().properties.cooperative_launch: raise CUDAError("cooperative kernels are not supported on this device") @@ -169,6 +178,11 @@ cdef class LaunchConfig: attr.value.programmaticStreamSerializationAllowed = 1 self._attrs.push_back(attr) + if self.priority is not None: + 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 +244,12 @@ cpdef object _to_native_launch_config(LaunchConfig config): attr.value.programmaticStreamSerializationAllowed = 1 attrs.append(attr) + if config.priority is not None: + 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..ce9e544fd65 100644 --- a/cuda_core/tests/test_launcher.py +++ b/cuda_core/tests/test_launcher.py @@ -202,6 +202,41 @@ def test_to_native_launch_config_pdl(): ) +@pytest.mark.parametrize( + ("initial_priority", "updated_priority"), + ((-1, 0), (0, 1), (1, -1)), +) +def test_launch_config_priority_getter_setter(initial_priority, updated_priority): + 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, 1), (-1, 1), (1, 1)), +) +def test_to_native_launch_config_priority(priority, expected_num_attrs): + """LaunchConfig priority maps to the native attribute, including zero.""" + 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 + assert native.numAttrs == expected_num_attrs + if priority is None: + 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..95ad09ac42e 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=(?:None|-?\d+)\)", ), ("sample_kernel", r""), # ObjectCode variations (by code_type)