Implement a thread-local means to access kernel launch config. #288

tpn · 2025-06-10T22:46:29Z

This allows downstream passes, such as rewriting, to access information about the kernel launch for which they have been enlisted to participate.

Posting this as a PR now to get feedback on the overall approach. Assuming this solution is acceptable, I'll follow up with tests and docs.

copy-pr-bot · 2025-06-10T22:46:32Z

Auto-sync is disabled for ready for review pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

tpn · 2025-06-10T22:55:38Z

Will close #280.

gmarkall · 2025-06-11T22:00:28Z

/ok to test

gmarkall · 2025-06-11T22:07:45Z

Using the simple benchmark from numba/numba#3003 (comment):

from numba import cuda
import numpy as np


@cuda.jit('void()')
def some_kernel_1():
    return

@cuda.jit('void(float32[:])')
def some_kernel_2(arr1):
    return

@cuda.jit('void(float32[:],float32[:])')
def some_kernel_3(arr1,arr2):
    return

@cuda.jit('void(float32[:],float32[:],float32[:])')
def some_kernel_4(arr1,arr2,arr3):
    return

@cuda.jit('void(float32[:],float32[:],float32[:],float32[:])')
def some_kernel_5(arr1,arr2,arr3,arr4):
    return

arr = cuda.device_array(10000, dtype=np.float32)

%timeit some_kernel_1[1, 1]()
%timeit some_kernel_2[1, 1](arr)
%timeit some_kernel_3[1, 1](arr,arr)
%timeit some_kernel_4[1, 1](arr,arr,arr)
%timeit some_kernel_5[1, 1](arr,arr,arr,arr)

On main:

6.6 μs ± 11.3 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)
10.6 μs ± 79.9 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)
14.1 μs ± 65.8 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)
17.4 μs ± 236 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)
20.3 μs ± 38.5 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)

On this branch:

8.6 μs ± 70.7 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)
13 μs ± 41.3 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)
17.2 μs ± 152 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)
20.3 μs ± 35.9 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)
24.4 μs ± 44.4 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)

From this crude benchmark, it appears to add 2-4μs per launch, or 17-23% overhead. I don't yet know how we should consider this (or whether the benchmark is appropriate) but I think it's a consideration to keep in mind when thinking about this approach.

gmarkall · 2025-06-11T22:08:54Z

numba_cuda/numba/cuda/dispatcher.py

+            stream=stream,
+            sharedmem=sharedmem,
+        ):
+            if self.specialized:


Specialized kernels cannot be recompiled, so a new launch configuration would not be able to affect the compilation of a new version - so this check could be kept outside the context manager.

gmarkall · 2025-06-11T22:11:12Z

numba_cuda/numba/cuda/launchconfig.py

+
+
+@dataclass(frozen=True, slots=True)
+class LaunchConfig:


There seems to be quite some overlap with dispatcher._LaunchConfiguration in this class (an observation at this point - I don't know whether it makes sense to combine them)

kkraus14 · 2025-06-12T01:22:13Z

Using the simple benchmark from numba/numba#3003 (comment):

from numba import cuda
import numpy as np


@cuda.jit('void()')
def some_kernel_1():
    return

@cuda.jit('void(float32[:])')
def some_kernel_2(arr1):
    return

@cuda.jit('void(float32[:],float32[:])')
def some_kernel_3(arr1,arr2):
    return

@cuda.jit('void(float32[:],float32[:],float32[:])')
def some_kernel_4(arr1,arr2,arr3):
    return

@cuda.jit('void(float32[:],float32[:],float32[:],float32[:])')
def some_kernel_5(arr1,arr2,arr3,arr4):
    return

arr = cuda.device_array(10000, dtype=np.float32)

%timeit some_kernel_1[1, 1]()
%timeit some_kernel_2[1, 1](arr)
%timeit some_kernel_3[1, 1](arr,arr)
%timeit some_kernel_4[1, 1](arr,arr,arr)
%timeit some_kernel_5[1, 1](arr,arr,arr,arr)

On main:

6.6 μs ± 11.3 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)
10.6 μs ± 79.9 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)
14.1 μs ± 65.8 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)
17.4 μs ± 236 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)
20.3 μs ± 38.5 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)

On this branch:

8.6 μs ± 70.7 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)
13 μs ± 41.3 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)
17.2 μs ± 152 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)
20.3 μs ± 35.9 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)
24.4 μs ± 44.4 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)

From this crude benchmark, it appears to add 2-4μs per launch, or 17-23% overhead. I don't yet know how we should consider this (or whether the benchmark is appropriate) but I think it's a consideration to keep in mind when thinking about this approach.

Kernel launch latency is something we need to care about moving forward and is becoming more and more of a bottleneck in important workloads. Adding 2-4μs per launch for this is probably unacceptable, but that being said our launch latency right now is much higher than we'd like in general where we probably need to rework the entire launch path at some point in the not too distant future.

gmarkall · 2025-06-12T10:27:59Z

we probably need to rework the entire launch path at some point in the not too distant future.

I think that's another thing that concerns me - if we implement something like this, then it constrains how we can rework the launch path in future if we have to go on supporting it.

This allows downstream passes, such as rewriting, to access information about the kernel launch for which they have been enlisted to participate.

This routine raises an error if no launch config is set, which is inevitably going to be the preferred way of obtaining the current launch config.

This is required by cuda.coop in order to pass two-phase primitive instances as kernel parameters without having to call the @cuda.jit decorator with extensions=[...] up-front.

greptile-apps · 2026-01-12T17:01:08Z

Greptile Overview

Greptile Summary

Overview

This PR implements a thread-local mechanism to access kernel launch configuration information using Python's contextvars.ContextVar. The implementation adds:

New file launchconfig.py: Defines LaunchConfig dataclass and context manager launch_config_ctx for storing launch parameters
Modified dispatcher.py: Wraps kernel calls in launch_config_ctx and invokes pre-launch callbacks

Architecture

The design uses ContextVar for thread-safe and asyncio-task-safe storage of launch configuration, which is appropriate for this use case. The context manager properly handles nested kernel launches by using tokens to restore previous values.

The pre_launch_callbacks mechanism allows downstream passes (like rewriting) to register callbacks that execute just before kernel launch, enabling dynamic kernel modifications.

Key Concerns

1. Frozen Dataclass with Mutable Field

The LaunchConfig dataclass is marked frozen=True but contains a mutable List[Callable] field. This creates an inconsistent immutability guarantee and could surprise users who expect frozen dataclasses to be truly immutable.

2. Missing Exception Handling

Pre-launch callbacks are invoked without try-except blocks. If a callback fails, the error message won't indicate which callback caused the problem, making debugging difficult.

3. Performance Considerations

The __str__ method converts all args and callbacks to strings without truncation, which could be expensive for large arrays or cause memory issues.

4. Missing Tests and Documentation

As noted by the author, tests and documentation are not yet included. Given this is a core dispatch mechanism change, comprehensive testing is critical before merge.

Recommendations

Add exception handling around callback invocation with clear error messages
Clarify the design intent for the frozen dataclass with mutable list (or change to non-frozen)
Optimize __str__ to handle large objects gracefully
Add comprehensive tests covering:
- Nested kernel launches
- Callback exceptions
- Thread safety
- Callback registration and execution order
Document the callback mechanism including signature requirements and execution guarantees

Confidence Score: 3/5

This PR introduces a solid architectural pattern but has design inconsistencies and missing error handling that should be addressed before merge.
Score of 3 reflects that while the core implementation using ContextVar is sound and thread-safe, there are several style and robustness issues: (1) frozen dataclass with mutable list creates misleading immutability guarantees, (2) missing exception handling for callbacks could cause unclear failures, (3) str method could have performance issues, and (4) no tests or documentation yet. The functionality appears correct but needs polish before production use.
Both files need attention: launchconfig.py for design clarification on immutability, and dispatcher.py for adding error handling around callback invocation.

Important Files Changed

File Analysis

Filename	Score	Overview
numba_cuda/numba/cuda/launchconfig.py	3/5	New file introducing thread-local launch config context. Design issues with frozen dataclass containing mutable list, and str method could be optimized.
numba_cuda/numba/cuda/dispatcher.py	3/5	Integrates launch config context into kernel dispatch. Missing exception handling for pre-launch callbacks could cause unclear errors.

greptile-apps

_{3 files reviewed, 3 comments}

_{Edit Code Review Agent Settings | Greptile}

greptile-apps · 2026-01-12T17:01:12Z

numba_cuda/numba/cuda/launchconfig.py

+@dataclass(frozen=True, slots=True)
+class LaunchConfig:
+    """
+    Helper class used to encapsulate kernel launch configuration for storing
+    and retrieving from a thread-local ContextVar.
+    """
+
+    dispatcher: "CUDADispatcher"
+    args: Tuple[Any, ...]
+    griddim: Tuple[int, int, int]
+    blockdim: Tuple[int, int, int]
+    stream: Any
+    sharedmem: int
+    pre_launch_callbacks: List[Callable[["_Kernel", "LaunchConfig"], None]]


The LaunchConfig dataclass is marked as frozen=True, but contains a mutable List field pre_launch_callbacks. This creates an inconsistent state where the dataclass appears immutable but the list can still be modified via append(), clear(), etc.

While this may be intentional to allow dynamic callback registration, it violates the principle of least surprise. Consider one of these alternatives:

Remove frozen=True if mutability is intended

Use tuple instead of List for pre_launch_callbacks to maintain true immutability

Add a comment explaining why frozen + mutable list is intentional

The current design could lead to subtle bugs if code assumes the frozen dataclass is truly immutable.

greptile-apps · 2026-01-12T17:01:13Z

numba_cuda/numba/cuda/dispatcher.py

+            for callback in launch_config.pre_launch_callbacks:
+                callback(kernel, launch_config)


Pre-launch callbacks are invoked without any exception handling. If a callback raises an exception:

The kernel launch will be aborted (which may be intended)

The error won't indicate which callback failed or provide context

If earlier callbacks modified state, there's no rollback mechanism

Consider wrapping the callback invocation in a try-except block to provide better error messages:

Suggested change

for callback in launch_config.pre_launch_callbacks:

callback(kernel, launch_config)

for callback in launch_config.pre_launch_callbacks:

try:

callback(kernel, launch_config)

except Exception as e:

raise RuntimeError(

f"Pre-launch callback {callback} failed: {e}"

) from e

Alternatively, if callback failures should not prevent kernel launch, you could log the error and continue.

_{Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!}

greptile-apps · 2026-01-12T17:01:14Z

numba_cuda/numba/cuda/launchconfig.py

+    def __str__(self) -> str:
+        a = ", ".join(map(str, self.args))
+        g = "×".join(map(str, self.griddim))
+        b = "×".join(map(str, self.blockdim))
+        cb = ", ".join(map(str, self.pre_launch_callbacks))


The __str__ method has potential performance and usability issues:

Line 42: Converting all args to strings could be expensive if args contains large arrays or GPU memory. Consider truncating or limiting the representation.

Line 45: Function objects will display as <function name at 0x...> which isn't very informative for debugging.

Consider improving the string representation:

Suggested change

def __str__(self) -> str:

a = ", ".join(map(str, self.args))

g = "×".join(map(str, self.griddim))

b = "×".join(map(str, self.blockdim))

cb = ", ".join(map(str, self.pre_launch_callbacks))

def __str__(self) -> str:

# Truncate args representation if too long

args_str = ", ".join(str(arg)[:50] + "..." if len(str(arg)) > 50 else str(arg) for arg in self.args)

g = "×".join(map(str, self.griddim))

b = "×".join(map(str, self.blockdim))

# Show callback names instead of memory addresses

cb = ", ".join(getattr(c, '__name__', repr(c)) for c in self.pre_launch_callbacks)

return (

This provides more useful debug output without performance issues.

tpn mentioned this pull request Jun 10, 2025

[FEA] Plumb kernel launch configuration such that it's accessible from rewriting and lowering phases #280

Open

gmarkall reviewed Jun 11, 2025

View reviewed changes

gmarkall added the 2 - In Progress Currently a work in progress label Jun 12, 2025

tpn force-pushed the 280-launch-config-contextvar branch from 019d60d to bba597d Compare July 14, 2025 02:36

leofang mentioned this pull request Sep 19, 2025

Reduce kernel launch overhead #477

Open

tpn force-pushed the 280-launch-config-contextvar branch from bba597d to fea9f79 Compare September 29, 2025 22:02

NVIDIA deleted a comment from CLAassistant Oct 1, 2025

leofang mentioned this pull request Oct 7, 2025

perf: speed up kernel launch #510

Merged

gmarkall mentioned this pull request Nov 17, 2025

[BUG] Warp vote operations must use a constant int for the mode #592

Closed

tpn added 4 commits January 12, 2026 08:55

Implement a thread-local means to access kernel launch config.

edcc5ed

This allows downstream passes, such as rewriting, to access information about the kernel launch for which they have been enlisted to participate.

Add args and dispatcher to LaunchConfig.

992fb99

Implement ensure_current_launch_config().

c6377db

This routine raises an error if no launch config is set, which is inevitably going to be the preferred way of obtaining the current launch config.

Add support for pre-kernel-launch callbacks to launch config.

23405f3

This is required by cuda.coop in order to pass two-phase primitive instances as kernel parameters without having to call the @cuda.jit decorator with extensions=[...] up-front.

tpn force-pushed the 280-launch-config-contextvar branch from fea9f79 to 23405f3 Compare January 12, 2026 16:56

greptile-apps bot reviewed Jan 12, 2026

View reviewed changes

		for callback in launch_config.pre_launch_callbacks:
		callback(kernel, launch_config)

Implement a thread-local means to access kernel launch config. #288

Are you sure you want to change the base?

Implement a thread-local means to access kernel launch config. #288

Uh oh!

Conversation

tpn commented Jun 10, 2025

Uh oh!

copy-pr-bot bot commented Jun 10, 2025

Uh oh!

tpn commented Jun 10, 2025

Uh oh!

gmarkall commented Jun 11, 2025

Uh oh!

gmarkall commented Jun 11, 2025

Uh oh!

gmarkall Jun 11, 2025

Choose a reason for hiding this comment

Uh oh!

gmarkall Jun 11, 2025

Choose a reason for hiding this comment

Uh oh!

kkraus14 commented Jun 12, 2025

Uh oh!

gmarkall commented Jun 12, 2025

Uh oh!

greptile-apps bot commented Jan 12, 2026

Greptile Overview

Greptile Summary

Overview

Architecture

Key Concerns

1. Frozen Dataclass with Mutable Field

2. Missing Exception Handling

3. Performance Considerations

4. Missing Tests and Documentation

Recommendations

Confidence Score: 3/5

Important Files Changed

Uh oh!

greptile-apps bot left a comment

Choose a reason for hiding this comment

Uh oh!

greptile-apps bot Jan 12, 2026

Choose a reason for hiding this comment

Uh oh!

greptile-apps bot Jan 12, 2026

Choose a reason for hiding this comment

Uh oh!

greptile-apps bot Jan 12, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants