Rate this Page

★ ★ ★ ★ ★

torch.cuda#

Created On: Dec 23, 2016 | Last Updated On: Aug 05, 2026

This package adds support for CUDA tensor types.

It implements the same function as CPU tensors, but they utilize GPUs for computation.

It is lazily initialized, so you can always import it, and use is_available() to determine if your system supports CUDA.

CUDA semantics has more details about working with CUDA.

`StreamContext`	Context-manager that selects a given stream.
`can_device_access_peer`	Check if peer access between two devices is possible.
`check_error`	Raise an error if the result of a CUDA runtime API call is not success.
`current_blas_handle`	Return cublasHandle_t pointer to current cuBLAS handle
`current_solver_handle`	Return cusolverDnHandle_t pointer to current cuSOLVER handle
`current_device`	Return the index of a currently selected device.
`current_stream`	Return the currently selected `Stream` for a given device.
`cudart`	Retrieves the CUDA runtime API module.
`default_stream`	Return the default `Stream` for a given device.
`device`	Context-manager that changes the selected device.
`device_count`	Return the number of GPUs available.
`device_memory_used`	Return used global (device) memory in bytes as given by nvidia-smi or amd-smi.
`device_of`	Context-manager that changes the current device to that of given object.
`get_arch_list`	Return list CUDA architectures this library was compiled for.
`get_device_capability`	Get the cuda capability of a device.
`get_device_name`	Get the name of a device.
`get_device_properties`	Get the properties of a device.
`get_gencode_flags`	Return NVCC gencode flags this library was compiled with.
`get_stream_from_external`	Return a `Stream` from an externally allocated CUDA stream.
`get_sync_debug_mode`	Return current value of debug mode for cuda synchronizing operations.
`init`	Initialize PyTorch's CUDA state.
`ipc_collect`	Force collects GPU memory after it has been released by CUDA IPC.
`is_available`	Return a bool indicating if CUDA is currently available.
`is_bf16_supported`	Return a bool indicating if the current CUDA/ROCm device supports dtype bfloat16.
`is_initialized`	Return whether PyTorch's CUDA state has been initialized.
`is_tf32_supported`	Return a bool indicating if the current CUDA/ROCm device supports dtype tf32.
`memory_usage`	Return the percent of time over the past sample period during which global (device) memory was being read or written as given by nvidia-smi.
`set_device`	Set the current device.
`set_stream`	Set the current stream. This is a wrapper API to set the stream.
`set_sync_debug_mode`	Set the debug mode for cuda synchronizing operations.
`stream`	Wrap around the Context-manager StreamContext that selects a given stream.
`synchronize`	Wait for all kernels in all streams on a CUDA device to complete.
`utilization`	Return the percent of time over the past sample period during which one or more kernels was executing on the GPU as given by nvidia-smi.
`temperature`	Return the average temperature of the GPU sensor in Degrees C (Centigrades).
`power_draw`	Return the average power draw of the GPU sensor in mW (MilliWatts)
`clock_rate`	Return the clock speed of the GPU SM in MHz (megahertz) over the past sample period as given by nvidia-smi.
`AcceleratorError`	Exception raised while executing on device
`OutOfMemoryError`	Exception raised when device is out of memory

Random Number Generator#

`get_rng_state`	Return the random number generator state of the specified GPU as a ByteTensor.
`get_rng_state_all`	Return a list of ByteTensor representing the random number states of all devices.
`set_rng_state`	Set the random number generator state of the specified GPU.
`set_rng_state_all`	Set the random number generator state of all devices.
`manual_seed`	Set the seed for generating random numbers for the current GPU.
`manual_seed_all`	Set the seed for generating random numbers on all GPUs.
`seed`	Set the seed for generating random numbers to a random number for the current GPU.
`seed_all`	Set the seed for generating random numbers to a random number on all GPUs.
`initial_seed`	Return the current random seed of the current GPU.

Communication collectives#

`comm.broadcast`	Broadcasts a tensor to specified GPU devices.
`comm.broadcast_coalesced`	Broadcast a sequence of tensors to the specified GPUs.
`comm.reduce_add`	Sum tensors from multiple GPUs.
`comm.reduce_add_coalesced`	Sum tensors from multiple GPUs.
`comm.scatter`	Scatters tensor across multiple GPUs.
`comm.gather`	Gathers tensors from multiple GPU devices.

Streams and events#

Stream

Wrapper around a CUDA stream.

ExternalStream

Wrapper around an externally allocated CUDA stream.

Event

Wrapper around a CUDA event.

Graphs (beta)#

`is_current_stream_capturing`	Return True if CUDA graph capture is underway on the current CUDA stream, False otherwise.
`graph_pool_handle`	Return an opaque token representing the id of a graph memory pool.
`CUDAGraph`	Wrapper around a CUDA graph.
`graph`	Context-manager that captures CUDA work into a `torch.cuda.CUDAGraph` object for later replay.
`make_graphed_callables`	Accept callables (functions or `nn.Module`s) and returns graphed versions.
`export_dot`	Return a capture-end hook that dumps the captured graph to `path` in Graphviz DOT format.
`export_graph_data`	Return a post-instantiate hook that pickles `CUDAGraph.get_graph_data()` to `path`.

CUDA graph lifecycle hooks#

Register callbacks that fire at each point in any CUDA graph’s lifecycle – capture start, capture end, instantiate, each replay, and destroy – for example, a profiler observing graph lifecycle without the graph code carrying any consumer knowledge, and including graphs the consumer did not build. Registering a hook is the opt-in and the whole API – the graph fires them; with none registered they are no-ops. Each has a per-graph counterpart on torch.cuda.CUDAGraph. Live in torch.cuda.graphs.

`register_graph_capture_start_hook`	Register a hook run with each CUDA graph as its capture begins.
`register_graph_capture_end_hook`	Register a hook run with each CUDA graph when its capture ends, while the captured `cudaGraph_t` is still live (see `CUDAGraph.register_capture_end_hook()`).
`register_graph_instantiate_hook`	Register a hook run with each CUDA graph right after it is instantiated.
`register_graph_replay_start_hook`	Register a hook run with each CUDA graph at the start of every replay, just before it is launched.
`register_graph_replay_end_hook`	Register a hook run with each CUDA graph at the end of every replay, once the replay is enqueued (the launch is asynchronous, so the GPU work has not completed).
`register_graph_destroy_hook`	Register `fn(exec_ids)` to run when a CUDA graph is destroyed.

Graph Kernel Annotations (prototype)#

torch.cuda.graph_annotations annotates the kernels captured in a CUDA graph with user metadata, keyed so the annotations can be joined against the graph node id field on kernel events in profiler traces. Enable recording per capture with the enable_annotations argument of torch.cuda.graph, then wrap regions of the captured workload in mark_kernels() scopes.

These APIs require the cuda-bindings package and a CUDA driver that supports cudaGraphNodeGetToolsId (CUDA 13.1 or newer, or an equivalent cuda-compat package); recording silently degrades to a no-op otherwise, and is not supported on ROCm. Use is_available() to check support.

The end-to-end workflow: annotate during capture, profile the replay, then merge the annotations into the exported trace and view it in Perfetto. During capture:

import torch
from torch.cuda.graph_annotations import mark_kernels, get_kernel_annotations

x = torch.randn(1024, 1024, device="cuda")

# Warmup: run the workload once outside capture so lazy initialization
# (e.g. cuBLAS handles) does not end up in -- or invalidate -- the capture.
y = x @ x.t()
z = torch.relu(y) @ x
torch.cuda.synchronize()

g = torch.cuda.CUDAGraph()
with torch.cuda.graph(g, enable_annotations=True):
    with mark_kernels("attention"):
        y = x @ x.t()
    with mark_kernels({"name": "mlp", "layer": 3}):
        z = torch.relu(y) @ x

with torch.profiler.profile() as prof:
    g.replay()
    torch.cuda.synchronize()
prof.export_chrome_trace("trace.json")

Each kernel event in the exported trace carries a graph node id in its args; the recorded annotations are keyed by the same ids, so merging them into the trace is a dictionary lookup:

import json

annotations = get_kernel_annotations()
with open("trace.json") as f:
    trace = json.load(f)
for event in trace["traceEvents"]:
    node_id = event.get("args", {}).get("graph node id")
    for ann in annotations.get(node_id, []):
        event["args"].update(ann)
with open("trace_annotated.json", "w") as f:
    json.dump(trace, f)

Opening trace_annotated.json in Perfetto (or chrome://tracing) and clicking a kernel from the graph replay now shows the annotation fields – name: attention or name: mlp, layer: 3 – alongside the kernel’s grid and block sizes, identifying which region of the captured workload each kernel came from.

Because annotations live in a process-global registry keyed by ids that match the profiler’s, the pickle of dict(get_kernel_annotations()) can equally be saved next to a trace and joined offline.

`is_available`	Return whether CUDA graph annotation recording is supported.
`mark_kernels`	Context manager that annotates GPU work captured within its scope.
`get_kernel_annotations`	Return the live registry of recorded kernel annotations.
`clear_kernel_annotations`	Clear all recorded kernel annotations.

This package adds support for device memory management implemented in CUDA.

Memory management#

`empty_cache`	Release all unoccupied cached memory currently held by the caching allocator so that those can be used in other GPU application and visible in nvidia-smi.
`get_per_process_memory_fraction`	Get memory fraction for a process.
`list_gpu_processes`	Return a human-readable printout of the running processes and their GPU memory use for a given device.
`mem_get_info`	Return the global free and total GPU memory for a given device using cudaMemGetInfo.
`memory_stats`	Return a dictionary of CUDA memory allocator statistics for a given device.
`memory_stats_as_nested_dict`	Return the result of `memory_stats()` as a nested dictionary.
`reset_accumulated_memory_stats`	Reset the "accumulated" (historical) stats tracked by the CUDA memory allocator.
`host_memory_stats`	Return a dictionary of pinned (host) allocator statistics.
`host_memory_stats_as_nested_dict`	Return the result of `host_memory_stats()` as a nested dictionary.
`reset_accumulated_host_memory_stats`	Reset the "accumulated" (historical) stats tracked by the host memory allocator.
`memory_summary`	Return a human-readable printout of the current memory allocator statistics for a given device.
`memory_snapshot`	Return a snapshot of the CUDA memory allocator state across all devices.
`memory_allocated`	Return the current GPU memory occupied by tensors in bytes for a given device.
`max_memory_allocated`	Return the maximum GPU memory occupied by tensors in bytes for a given device.
`reset_max_memory_allocated`	Reset the starting point in tracking maximum GPU memory occupied by tensors for a given device.
`memory_reserved`	Return the current GPU memory managed by the caching allocator in bytes for a given device.
`max_memory_reserved`	Return the maximum GPU memory managed by the caching allocator in bytes for a given device.
`set_per_process_memory_fraction`	Set memory fraction for a process.
`memory_cached`	Deprecated; see `memory_reserved()`.
`max_memory_cached`	Deprecated; see `max_memory_reserved()`.
`reset_max_memory_cached`	Reset the starting point in tracking maximum GPU memory managed by the caching allocator for a given device.
`reset_peak_memory_stats`	Reset the "peak" stats tracked by the CUDA memory allocator.
`reset_peak_host_memory_stats`	Reset the "peak" stats tracked by the host memory allocator.
`caching_allocator_alloc`	Perform a memory allocation using the CUDA memory allocator.
`caching_allocator_delete`	Delete memory allocated using the CUDA memory allocator.
`get_allocator_backend`	Return a string describing the active allocator backend as set by `PYTORCH_ALLOC_CONF`.
`CUDAPluggableAllocator`	CUDA memory allocator loaded from a so file.
`change_current_allocator`	Change the currently used memory allocator to be the one provided.
`MemPool`	MemPool represents a pool of memory in a caching allocator.

caching_allocator_disabled

Context manager that temporarily disables the CUDA caching allocator.

caching_allocator_enable

Enable or disable the CUDA memory allocator.

class torch.cuda.use_mem_pool(pool, device=None)[source]#

A context manager that routes allocations to a given pool.

Parameters:

pool (torch.cuda.MemPool) – a MemPool object to be made active so that allocations route to this pool.
device (torch.device or int, optional) – selected device. Uses MemPool on the current device, given by current_device(), if device is None (default).

Note

This context manager makes only current thread’s allocations route to the given pool. If a new thread is spawned inside the context manager (e.g. by calling backward) the allocations in that thread will not route to the given pool.

Note

When used during CUDAGraph capture, the graph retains the pool until the graph is reset or destroyed.

torch.cuda.nccl.version()[source]#

Returns the version of the NCCL.

This function returns a tuple containing the major, minor, and patch version numbers of the NCCL. The suffix is also included in the tuple if a version suffix exists. :returns: The version information of the NCCL. :rtype: tuple

profile

Enable profiling.

start

Starts cuda profiler data collection.

stop

Stops cuda profiler data collection.

NVIDIA Tools Extension (NVTX)#

`nvtx.mark`	Describe an instantaneous event that occurred at some point.
`nvtx.range_push`	Push a range onto a stack of nested range span.
`nvtx.range_pop`	Pop a range off of a stack of nested range spans.
`nvtx.range`	Context manager / decorator that pushes an NVTX range at the beginning of its scope, and pops it at the end.
`nvtx.range_end`	Mark the end of a range for a given range_id.
`nvtx.range_start`	Mark the start of a range with string message.

Jiterator (beta)#

`jiterator._create_jit_fn`	Create a jiterator-generated cuda kernel for an elementwise op.
`jiterator._create_multi_output_jit_fn`	Create a jiterator-generated cuda kernel for an elementwise op that supports returning one or more outputs.

TunableOp#

Some operations could be implemented using more than one library or more than one technique. For example, a GEMM could be implemented for CUDA or ROCm using either the cublas/cublasLt libraries or hipblas/hipblasLt libraries, respectively. How does one know which implementation is the fastest and should be chosen? That’s what TunableOp provides. Certain operators have been implemented using multiple strategies as Tunable Operators. At runtime, all strategies are profiled and the fastest is selected for all subsequent operations.

See the documentation for information on how to use it.

Stream Sanitizer (prototype)#

CUDA Sanitizer is a prototype tool for detecting synchronization errors between streams in PyTorch. See the documentation for information on how to use it.

GPUDirect Storage (prototype)#

The APIs in torch.cuda.gds provide thin wrappers around certain cuFile APIs that allow direct memory access transfers between GPU memory and storage, avoiding a bounce buffer in the CPU. See the cufile api documentation for more details.

These APIs can be used in versions greater than or equal to CUDA 12.6. In order to use these APIs, one must ensure that their system is appropriately configured to use GPUDirect Storage per the GPUDirect Storage documentation.

See the docs for GdsFile for an example of how to use these.

gds_register_buffer

Registers a storage on a CUDA device as a cufile buffer.

gds_deregister_buffer

Deregisters a previously registered storage on a CUDA device as a cufile buffer.

GdsFile

Wrapper around cuFile.

Green Contexts (experimental)#

torch.cuda.green_contexts provides thin wrappers around the CUDA Green Context APIs to enable more general carveout of SM resources for CUDA kernels.

These APIs require the cuda.bindings package and can be used in PyTorch with CUDA versions greater than or equal to 12.8. Workqueue configuration requires CUDA 13.1 or newer.

Install instructions for cuda.bindings can be found here: https://nvidia.github.io/cuda-python/

Create streams from the green context and use them like other custom CUDA streams:

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

Synchronization between green-context streams and other streams is the user’s responsibility. Use CUDA events to order work, just as you would for any other custom stream.

The GreenContext.set_context() and GreenContext.pop_context() methods are deprecated compatibility APIs.

GreenContext

Wrapper around a CUDA green context.

torch.cuda.nccl.is_available(tensors)[source]#

This package adds support for NVIDIA Tools Extension (NVTX) used in profiling.