Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 2 additions & 1 deletion cuda_core/cuda/core/_launch_config.pxd
Original file line number Diff line number Diff line change
@@ -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

Expand All @@ -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__
Expand Down
11 changes: 9 additions & 2 deletions cuda_core/cuda/core/_launch_config.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -39,15 +39,19 @@ 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

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Do we have a range for the value we can set? It cannot be any number.
CUDA doc should provide guidance on what value to set. Let's include the guidance in docstring.

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, ...]
block: tuple[Any, ...]
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
Expand All @@ -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:
Expand Down
20 changes: 20 additions & 0 deletions cuda_core/cuda/core/_launch_config.pyx
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@ _LAUNCH_CONFIG_ATTRS = (
'shmem_size',
'is_cooperative',
'programmatic_stream_serialization',
'priority',
)

__all__ = ['LaunchConfig']
Expand Down Expand Up @@ -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
Expand All @@ -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.

Expand All @@ -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)
Expand Down Expand Up @@ -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")
Expand Down Expand Up @@ -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()

Expand Down Expand Up @@ -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

Expand Down
35 changes: 35 additions & 0 deletions cuda_core/tests/test_launcher.py
Original file line number Diff line number Diff line change
Expand Up @@ -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+.
Expand Down
2 changes: 1 addition & 1 deletion cuda_core/tests/test_object_protocols.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"<Kernel handle=0x[0-9a-f]+>"),
# ObjectCode variations (by code_type)
Expand Down
Loading