GreenContext

class torch.cuda.green_contexts.GreenContext(*, num_sms=None, sm_partition=None, workqueue_scope=None, workqueue_concurrency_limit=None, device_id=None)

Wrapper around a CUDA green context.

Warning

This API is in beta and may change in future releases.

CUDA work should be placed on streams created from the green context:

ctx = GreenContext(...)
stream = ctx.Stream()
with torch.cuda.stream(stream):
    # torch operations here are using resources from `ctx`
    pass

Green-context streams are custom CUDA streams. Synchronization with other streams is the user’s responsibility and should be handled with CUDA events, as with any other custom stream.

Stream()

Return a CUDA stream associated with this green context.

Use the returned stream with torch.cuda.stream() to run work on the green context. Synchronization with other streams is not automatic; use CUDA events as with any other custom stream.

Return type:

Stream

static create(*, num_sms=None, sm_partition=None, workqueue_scope=None, workqueue_concurrency_limit=None, device_id=None)

Create a CUDA green context.

Kept for compatibility, see GreenContext constructor.

Return type:

GreenContext

property device_id: int

The device index of this green context.

property locality_domain_id: int | None

The locality domain reported by CUDA for this context, if specified.

Requires CUDA driver and bindings 13.4+.

static max_workqueue_concurrency(device_id=None)

Return the maximum workqueue concurrency limit for the device.

This queries the device for the default number of concurrent stream-ordered workloads supported by workqueue configuration resources.

Parameters:

device_id (int, optional) – The device index to query. When None, the current device is used.

Return type:

int

pop_context()

Assuming the green context is the current context, pop it from the context stack and restore the previous context.

Deprecated. Create streams with Stream() and use torch.cuda.stream() instead.

set_context()

Make the green context the current context.

Deprecated. Create streams with Stream() and use torch.cuda.stream() instead.

property sm_count: int

The actual number of SMs available to this context.

property sm_partition: SMPartition

The context’s actual SM resource, which can be subdivided.

The returned resource keeps this context alive while it is in use.

static split(*, num_sms=0, coscheduled_sm_count=0, preferred_coscheduled_sm_count=0, backfill=False, locality_domain_ids=None, workqueue_scope=None, workqueue_concurrency_limit=None, device_id=None)

Create contexts backed by disjoint SM partitions of a device.

Partition options are those of SMPartition.split(). Workqueue options are applied to each context. Unassigned SMs are unused; use SMPartition.split() to retain the remainder for later use.

Example:

>>> a, b = GreenContext.split(
...     num_sms=(24, 40), coscheduled_sm_count=8, device_id=0
... )

Return type:

tuple[GreenContext, …]