Add green context support by leofang · Pull Request #1976 · NVIDIA/cuda-python

leofang · 2026-04-25T03:56:04Z

Close #1563. Close #112.

Summary

Add green context support to cuda.core — the explicit-model API for querying device resources, splitting SMs, creating green contexts, and using them without touching the thread-local context stack.

Design

See the companion design doc for full rationale. Key decisions:

Unified Context type — no user-visible GreenContext subclass. A single Context wraps either a primary CUcontext or a CUgreenCtx + derived CUcontext. ctx.is_green distinguishes them. Inspired by the CUDA runtime's execution-context (EC) abstraction.
dev.resources namespace — DeviceResources groups hardware resource queries (dev.resources.sm, dev.resources.workqueue). Follows the existing "plural = namespace" pattern (dev.properties, kernel.attributes).
ctx.resources / stream.resources — same DeviceResources type, but queries the context's provisioned resources (cuCtxGetDevResource / cuGreenCtxGetDevResource) instead of the full device.
SMResourceOptions with SoA broadcasting — single dataclass for SMResource.split(). Scalar fields broadcast; count drives the group count. count=None means discovery mode (translated to smCount=0 internally).
Merged workqueue types — WorkqueueResource merges CU_DEV_RESOURCE_TYPE_WORKQUEUE_CONFIG and CU_DEV_RESOURCE_TYPE_WORKQUEUE under one user-facing class. Strings for option values (e.g. sharing_scope="green_ctx_balanced").
ContextOptions(resources=[...]) → dev.create_context() — resource descriptor generation and cuGreenCtxCreate are internal. The user passes pre-split resource objects.
Explicit model — ctx.create_stream() creates streams bound to a green context without calling dev.set_current(). The C++ handle layer auto-dispatches between cuGreenCtxStreamCreate and cuStreamCreateWithPriority based on the context type. Green context streams must be non-blocking.
ctx.close() does not manage the context stack — closing a current context raises RuntimeError. dev.set_current(green_ctx) still works for backward compatibility but is not the recommended path.

New public API

Device.resources → DeviceResources (namespace: .sm, .workqueue)
Context.resources → DeviceResources (context-level query of provisioned resources)
Stream.resources → DeviceResources (delegates to the stream's parent context)
Context.create_stream(options) → Stream (green contexts only; raises on primary)
Context.is_green → bool
SMResource — properties: sm_count, min_partition_size, coscheduled_alignment, flags, handle; method: split(options, *, dry_run=False)
SMResourceOptions — count, coscheduled_sm_count, preferred_coscheduled_sm_count
WorkqueueResource — method: configure(options)
WorkqueueResourceOptions — sharing_scope
ContextOptions.resources — accepts Sequence[SMResource | WorkqueueResource]

Implementation details

C++ handle layer (resource_handles.hpp/cpp):

GreenCtxHandle (shared_ptr<const CUgreenCtx>) — owning handle; destructor calls cuGreenCtxDestroy.
ContextBox gains a GreenCtxHandle field so the derived CUcontext keeps the green ctx alive. get_context_green_ctx() provides reverse lookup.
create_green_ctx_handle() combines cuDevResourceGenerateDesc + cuGreenCtxCreate in one call — the descriptor is transient (no DevResourceDescHandle needed since CUDA has no explicit destroy for it).
create_stream_handle() auto-dispatches: checks get_context_green_ctx() on the provided ContextHandle and calls cuGreenCtxStreamCreate for green contexts, cuStreamCreateWithPriority for primary. Returns CUDA_ERROR_NOT_SUPPORTED if the context is green but cuGreenCtxStreamCreate is unavailable (CUDA < 12.5).
context_registry / stream_registry (HandleRegistry) deduplicate handles by raw CUDA pointer, enabling identity-preserving set_current swaps.

Bug fix — stream context tracking:

StreamBox now carries a ContextHandle dependency, populated at creation time.
get_stream_context() returns it without a driver call.
Stream._from_handle and Stream_ensure_ctx prefer the registry-backed handle before falling back to cuStreamGetCtx. This fixes a latent issue where streams created in a green context would lose their context association after a set_current swap.

Error handling:

dev.create_context() without resources raises ValueError with a clear message.
Green context stream creation with nonblocking=False is caught by the driver (CUDA_ERROR_INVALID_VALUE) and re-raised as ValueError with a helpful message.
cuCtxGetStreamPriorityRange failure (CUDA_ERROR_INVALID_CONTEXT) raises "Call dev.set_current() before creating streams."

Version guards:

Compile-time: IF CUDA_CORE_BUILD_MAJOR >= 13 gates cuDevSmResourceSplit (the general/structured form).
Runtime: trinary cached checks (supported / unsupported / unchecked). cy_driver_version() >= (12, 4, 0) for all green ctx APIs; >= (13, 1, 0) for structured splits. Raises ValueError when unsupported.
CUDA 12.x fallback: cuDevSmResourceSplitByCount for basic (homogeneous) splits. Per-group coscheduled_sm_count and heterogeneous counts require 13.1+ and raise NotImplementedError on 12.x.
Green ctx function pointers loaded via _get_optional_driver_fn — graceful NULL when bindings lack the symbol.

Test coverage

33 tests in test_green_context.py, organized with proper pytest fixtures and classes:

Fixtures: sm_resource, wq_resource, green_ctx (with CUDAError → skip), fill_kernel
_use_green_ctx context manager for safe push/pop in set_current regression tests
TestSMResourceQuery — properties, arch constraints (pre-Hopper vs Hopper+)
TestWorkqueueResource — query, configure valid/invalid
TestSMResourceSplitValidation — scalar/Sequence mismatch, negative count, dry-run blocked
TestSMResourceSplit — single/two-group splits with arch-aligned counts, discovery mode, alignment, dry-run parity
TestGreenContextLifecycle — is_green, create_stream on primary raises, blocking stream raises, explicit stream creation, stream/event context tracking, close-while-current guard, set_current regression
TestContextResources — green ctx SM resources are subset of device, two contexts have disjoint partitions, stream.resources matches ctx.resources (SM + workqueue)
TestGreenContextKernelLaunch — compile + launch + host-verify via ctx.create_stream(), two independent green contexts with different fill values, SM + workqueue combined

Validation

CUDA_HOME=... pip install -e . --no-build-isolation
python -m pytest tests/test_green_context.py -v                          # 32 passed, 1 skipped (arch)
python -m pytest tests/test_device.py tests/test_stream.py tests/test_event.py tests/test_context.py -v  # no regressions (257 total passed)

-- Leo's bot

copy-pr-bot · 2026-04-25T03:56:07Z

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

Restructure tests into fixtures + classes with full resource cleanup: - Fixtures: sm_resource, wq_resource, green_ctx (with CUDAError skip), green_ctx_active (with try/finally restore), fill_kernel - _use_green_ctx context manager for safe push/pop in all tests - TestSMResourceQuery: properties, arch constraints per CC - TestSMResourceSplit: single/two-group splits, discovery, alignment, dry-run vs real parity - TestGreenContextKernelLaunch: compile + launch + verify in green ctx, two independent green contexts, SM + workqueue combined All set_current calls are paired with restore in finally blocks to prevent context stack leaks on test failure. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

leofang · 2026-04-25T04:44:20Z

/ok to test ac5c0fc

github-actions · 2026-04-25T05:02:45Z

Doc Preview CI
🚀 View preview at https://nvidia.github.io/cuda-python/pr-preview/pr-1976/
https://nvidia.github.io/cuda-python/pr-preview/pr-1976/cuda-core/
https://nvidia.github.io/cuda-python/pr-preview/pr-1976/cuda-bindings/
https://nvidia.github.io/cuda-python/pr-preview/pr-1976/cuda-pathfinder/
Preview will be ready when the GitHub Pages deployment is complete.

- Convert ContextOptions and SMResourceOptions/WorkqueueResourceOptions to cdef dataclasses for check_or_create_options compatibility. - Cache SM metadata in typed cdef fields; fall back to arch-based granularity on CUDA 12.x where CUdevSmResource lacks minSmPartitionSize/smCoscheduledAlignment. - Simplify Context to hold only ContextHandle (remove _h_green_ctx and _is_green fields). Green ctx association lives in ContextBox; is_green queries get_context_green_ctx() on demand. - ContextOptions.resources accepts Sequence only (no bare resource). Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Switch from the push model (dev.set_current + dev.create_stream) to the explicit model (ctx.create_stream + ctx.resources) as the primary way to use green contexts. Context.create_stream(options): - Only supported on green contexts (raises on primary contexts). - Delegates to Stream._init, which calls create_stream_handle in C++. - C++ create_stream_handle auto-dispatches: checks get_context_green_ctx and calls cuGreenCtxStreamCreate for green contexts, or cuStreamCreateWithPriority for primary. Single function, no duplication. Context.resources: - Returns a DeviceResources namespace querying this context's resources (cuCtxGetDevResource / cuGreenCtxGetDevResource), not the full device. dev.set_current(green_ctx) still works but is not the recommended path. Tests rewritten to use the explicit model throughout. Push-model set_current kept as regression tests with _use_green_ctx helper. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

- Let the driver validate the nonblocking flag for green context streams: cuGreenCtxStreamCreate rejects CU_STREAM_DEFAULT. On failure, check if the context is green + nonblocking is False and raise a clear ValueError. - cuCtxGetStreamPriorityRange failure (CUDA_ERROR_INVALID_CONTEXT) now raises: "No current CUDA context. Call dev.set_current() before creating streams." - C++ create_stream_handle returns CUDA_ERROR_NOT_SUPPORTED if the context is green but cuGreenCtxStreamCreate is unavailable (CUDA < 12.5), instead of falling through to cuStreamCreateWithPriority. - ctx.resources.workqueue now dispatches to cuGreenCtxGetDevResource for green contexts, matching the SM query path. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

leofang · 2026-04-27T04:38:28Z

/ok to test 2812c5b

Stream.resources delegates to DeviceResources._init_from_ctx via the stream's tracked context handle, returning the same resource view as ctx.resources for the stream's parent context. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

leofang · 2026-04-27T13:12:11Z

/ok to test 5b3c610

- dev.create_context raises ValueError (not NotImplementedError) when options or resources are missing. - Cache version checks (_check_green_ctx_support, _check_workqueue_support) at module level; raise ValueError instead of NotImplementedError. - Simplify _device_resources.pyx: merge _as_uint and _count_to_sm_count into _to_sm_count; inline unsigned int casts for coscheduled params. - Add green context classes to api.rst (Context, ContextOptions, DeviceResources, SMResource, SMResourceOptions, WorkqueueResource, WorkqueueResourceOptions). - Update all docstrings to NumPy style with Attributes/Parameters/Returns sections matching the existing codebase convention. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

leofang · 2026-04-27T14:11:29Z

/ok to test fa254a5

Andy-Jost

I see a few issues with the registries. It appears the context registry needs to be checked in one place.

I have a bigger concern with the stream registry. I don't see what problem it solves, and it appears to be corruptible through the user API.

Andy-Jost · 2026-04-29T17:57:21Z

+            context_registry.unregister_handle(b->resource);
            GILReleaseGuard gil;


nit: put the GIL release first for consistency in this file and to make it easier to spot.

For my own understanding: this seems to be consistent with the rest of the file? In the shared pointer deleter we:

unregister a handle from the registry

get GIL guard

call C API

Are you saying we should swap 1 and 2 for the whole file, including code that is not touched by this PR?

Andy-Jost · 2026-04-29T17:59:17Z

+            stream_registry.unregister_handle(b->resource);
            GILReleaseGuard gil;


nit: same as above

Andy-Jost · 2026-04-29T18:05:21Z

+            stream_registry.unregister_handle(b->resource);
            GILAcquireGuard gil;


nit: as above

Andy-Jost · 2026-04-29T19:08:47Z

 StreamHandle create_stream_handle_with_owner(CUstream stream, PyObject* owner) {
+    if (auto h = stream_registry.lookup(stream)) {
+        // Reuse handles that already carry structural context metadata, e.g.
+        // cuda-core-owned streams. Owner-backed foreign streams still need a
+        // fresh handle so the supplied owner is retained.
+        if (get_box(h)->h_context) {
+            return h;
+        }
+    }


This stream registry only allows one entry per unique CUstream, but there can be multiple PyObject owners, which will compete for the one and only cache slot. It can easily lead to registry corruption.

At a minimum, I don't think the stream + Python owner path should use a registry at all, since it is really just tying a Py_INCREF/Py_DECREF pair to a CUstream lifetime, which can be stacked any number of times.

More broadly, I'm having trouble understanding why a stream registry is needed at all. The PR description says:

Stream._from_handle and Stream_ensure_ctx prefer the registry-backed handle before falling back to cuStreamGetCtx. This fixes a latent issue where streams created in a green context would lose their context association after a set_current swap.

I don't see how this comes into play. A registry allows a raw handle such as CUstream to round-trip through the driver without losing metadata, but set_current is about CUcontext, not CUstream. Where is the path that needs to match up a raw CUstream with an existing StreamHandle?

@Andy-Jost correct me if I am wrong, I would think every type we own should have its own registry, why would we differentiate CUcontext from others? It is not only confusing but also unsafe in general, considering that we do want to handle the general cases for each object via from_handle() (#1989).

For example, our Buffer internally tracks the allocation Stream, so having a buffer registry would help express the intent (though for buffer we may have other considerations such as performance).

Where is the path that needs to match up a raw CUstream with an existing StreamHandle?

On way to trigger this is:

# cuda.core creates a green stream (cuGreenCtxStreamCreate under the hood) ctx = dev.create_context(...) s = ctx.create_stream() # User extracts the raw handle and re-wraps it (round-trip through int) raw = s.handle s2 = Stream.from_handle(raw) # Without registry: s2 is a fresh handle with no context dep # With registry: s2 reuses the existing StreamHandle, preserving # the ContextHandle → GreenCtxHandle dependency chain s2.context.is_green # True with registry, might fail without

so in this case a (stream) registry lookup is indeed useful.

It can easily lead to registry corruption.

I think we do want to address the corruption issue. Would it help if we keep the registry but exclude the owner path (i.e. foreign streams don't register)?

Andy-Jost · 2026-04-29T19:13:06Z

+        """True if this context was created from device resources."""
+        if not self._h_context:
+            return False
+        return get_context_green_ctx(self._h_context).get() != NULL


nit: consider putting this in the handle API as is_green(self._h_context).

…td::vector Review comment 1: Consolidate create_context_handle_from_green_ctx with create_context_handle_ref by adding a private overload that takes an optional GreenCtxHandle. The green ctx path now delegates to it after calling cuCtxFromGreenCtx, ensuring registry lookup and deduplication. Review comments 2-4: Move GILReleaseGuard to the first line in create_green_ctx_handle and create_context_handle_from_green_ctx for consistency with the rest of the file. Review comment 6: Keep is_green check inline in _context.pyx using get_context_green_ctx (cannot add a C++ is_green function across separate .so boundaries without linker issues). Review comment 8: Replace malloc/free with std::vector<CUdevResource> in Device.create_context for automatic cleanup. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

leofang · 2026-05-02T02:23:40Z

/ok to test 41cf1de

leofang added 5 commits April 24, 2026 23:16

Implement green context v1 API

353a382

Refine green context split compatibility

b33c381

Encode green context handle dependencies

faf0d17

Simplify green context view handles

7c35ef3

Simplify green context descriptor handling

da58b7d

github-actions Bot added the cuda.core Everything related to the cuda.core module label Apr 25, 2026

leofang changed the title ~~Add cuda.core green context v1 API~~ Add green context support Apr 25, 2026

leofang added P0 High priority - Must do! feature New feature or request labels Apr 25, 2026

leofang self-assigned this Apr 25, 2026

leofang added this to the cuda.core v1.0.0 milestone Apr 25, 2026

leofang mentioned this pull request Apr 25, 2026

Implement Context.__init__() in cuda.core.experimental._context #189

Closed

leofang force-pushed the leof/green-ctx-v1 branch 4 times, most recently from 17c2be2 to 08d52d1 Compare April 27, 2026 02:56

leofang force-pushed the leof/green-ctx-v1 branch from 08d52d1 to 3013fe8 Compare April 27, 2026 03:06

leofang force-pushed the leof/green-ctx-v1 branch 6 times, most recently from 62e4883 to 3287204 Compare April 27, 2026 04:21

leofang force-pushed the leof/green-ctx-v1 branch from 3287204 to 2812c5b Compare April 27, 2026 04:23

Add Stream.resources property

340506e

Stream.resources delegates to DeviceResources._init_from_ctx via the stream's tracked context handle, returning the same resource view as ctx.resources for the stream's parent context. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

leofang force-pushed the leof/green-ctx-v1 branch 3 times, most recently from d5a7297 to 5b3c610 Compare April 27, 2026 13:07

leofang force-pushed the leof/green-ctx-v1 branch 2 times, most recently from eebd4cf to 5989fd1 Compare April 27, 2026 13:53

leofang force-pushed the leof/green-ctx-v1 branch from 5989fd1 to fa254a5 Compare April 27, 2026 14:08

leofang requested a review from Andy-Jost April 27, 2026 16:48

leofang marked this pull request as ready for review April 27, 2026 18:24

leofang mentioned this pull request Apr 27, 2026

Support interop with active green contexts NVIDIA/numba-cuda#853

Draft

Andy-Jost requested changes Apr 29, 2026

View reviewed changes

leofang and others added 2 commits May 1, 2026 21:57

Merge remote-tracking branch 'origin/main' into leof/green-ctx-v1

37ebe2a

		context_registry.unregister_handle(b->resource);
		GILReleaseGuard gil;

		stream_registry.unregister_handle(b->resource);
		GILReleaseGuard gil;

		stream_registry.unregister_handle(b->resource);
		GILAcquireGuard gil;

Conversation

leofang commented Apr 25, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Design

New public API

Implementation details

Test coverage

Validation

Uh oh!

copy-pr-bot Bot commented Apr 25, 2026

Uh oh!

leofang commented Apr 25, 2026

Uh oh!

github-actions Bot commented Apr 25, 2026

Preview will be ready when the GitHub Pages deployment is complete.

Uh oh!

leofang commented Apr 27, 2026

Uh oh!

leofang commented Apr 27, 2026

Uh oh!

leofang commented Apr 27, 2026

Uh oh!

Andy-Jost left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Andy-Jost Apr 29, 2026

Choose a reason for hiding this comment

Uh oh!

leofang May 2, 2026

Choose a reason for hiding this comment

Uh oh!

Andy-Jost Apr 29, 2026

Choose a reason for hiding this comment

Uh oh!

Andy-Jost Apr 29, 2026

Choose a reason for hiding this comment

Uh oh!

Andy-Jost Apr 29, 2026

Choose a reason for hiding this comment

Uh oh!

leofang May 1, 2026

Choose a reason for hiding this comment

Uh oh!

Andy-Jost Apr 29, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

leofang commented May 2, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

leofang commented Apr 25, 2026 •

edited

Loading