CUDA backend

Active contributors: eqy, syed-ahmed, Aidyn-A

Purpose

PyTorch's CUDA support spans three layers of the codebase: low-level runtime helpers in c10/cuda/, kernel libraries in aten/src/ATen/cuda/ and aten/src/ATen/native/cuda/, and the Python-level torch.cuda module. Together they implement the GPU execution path: streams, events, the caching allocator, OOM diagnostics, kernel launches, and CUDA Graphs.

This page is the cross-cutting tour of the CUDA backend; for the dispatcher mechanics see Dispatcher, and for torch.compile GPU codegen see Inductor.

Directory layout

Path	Contents
c10/cuda/	Low-level runtime helpers shared with all CUDA code
c10/cuda/CUDACachingAllocator.cpp	The infamous caching allocator
c10/cuda/CUDAStream.cpp	Stream wrapper
c10/cuda/CUDAGuard.h	RAII device guard
c10/cuda/CUDAFunctions.cpp	Thin wrappers over CUDA Runtime calls
aten/src/ATen/cuda/	ATen-level CUDA helpers (cuBLAS, cuSPARSE, cuSolver wrappers, RNG)
aten/src/ATen/cuda/CUDAGraph.cpp	CUDA Graphs
aten/src/ATen/cuda/Atomic.cuh	Atomic ops shared by kernels
aten/src/ATen/native/cuda/	The bulk of CUDA op kernels
aten/src/ATen/native/cudnn/	cuDNN-backed convolutions, RNNs, attention
aten/src/ATen/native/sparse/cuda/	Sparse CUDA
aten/src/ATen/native/transformers/cuda/	Flash attention, scaled-dot-product attention
torch/cuda/	Python torch.cuda package
torch/csrc/cuda/	Python <-> C++ glue for torch.cuda
torch/_inductor/codegen/triton.py	Triton GPU codegen for torch.compile
third_party/nccl/, third_party/cutlass/, third_party/cudnn-frontend/	Vendor submodules

Key abstractions

Type	File	Purpose
c10::cuda::CUDACachingAllocator	c10/cuda/CUDACachingAllocator.cpp	Caches freed blocks instead of cudaFree
c10::cuda::CUDAStream	c10/cuda/CUDAStream.h	A stream + guard helper
c10::cuda::CUDAGuard	c10/cuda/CUDAGuard.h	RAII device-context switcher
at::cuda::CUDAGraph	aten/src/ATen/cuda/CUDAGraph.h	CUDA Graphs capture/replay
at::cuda::CUDABlasHandle	aten/src/ATen/cuda/CUDABlas.h	Cached cuBLAS handle per stream
at::cuda::detail::PhiloxCudaState	aten/src/ATen/cuda/PhiloxUtils.cuh	Counter-based RNG state

How it works

Caching allocator

The allocator (CUDACachingAllocator.cpp, ~5K lines) is the most important piece. It:

• Allocates large blocks via cudaMalloc and splits them into smaller blocks on demand, kept in size-bucketed free lists.
• Returns freed blocks to the pool (so subsequent allocations are O(1)) instead of cudaFree (which would synchronize the device).
• Tracks usage per stream so concurrent kernels don't trample each other.
• Optionally uses expandable segments (PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True) — virtual address ranges that grow with cuMemMap instead of repeatedly fragmenting the heap.
• Exposes hooks (set_allocator_settings, allocator traces, memory snapshots) used by the memory profiler and the OOM trace dumper.

PYTORCH_CUDA_ALLOC_CONF is the configuration knob. Its full schema is documented in c10/cuda/CUDAAllocatorConfig.cpp.

Streams and events

c10::cuda::CUDAStream wraps a cudaStream_t with the device id, an optional priority, and a "stream type" (default, current, custom). PyTorch maintains a per-thread current stream per device that ATen kernels submit to. c10::cuda::CUDAGuard and c10::cuda::CUDAStreamGuard are the RAII helpers for switching device/stream temporarily.

Events (c10::cuda::CUDAEvent) wrap cudaEvent_t and are used for cross-stream and cross-device synchronization, both internally (the allocator uses them) and externally (torch.cuda.Event, torch.cuda.synchronize).

Kernel launch path

When dispatch lands on a CUDA kernel (typically in aten/src/ATen/native/cuda/):

1. TensorIterator computes broadcasting + dtype promotion (CPU side).
2. The kernel allocates outputs via the caching allocator.
3. It launches a grid + block via <<<...>>> or cuLaunchKernel. Grid sizes use helpers from aten/src/ATen/cuda/detail/KernelUtils.h.
4. The kernel runs on the current stream of the current device.
5. Async operations are recorded against any saved tensors so the allocator knows when memory becomes free.

Vectorized loops are encoded via aten/src/ATen/native/cuda/Loops.cuh and similar; reductions through aten/src/ATen/native/cuda/Reduce.cuh.

cuBLAS / cuDNN / cuSPARSE / cuSolver

• cuBLAS: aten/src/ATen/cuda/CUDABlas.cpp wraps cuBLAS handles per stream. Used by mm, bmm, gemm, addmm. The at::globalContext().blasPreferredBackend() lets users pick between cuBLAS and cuBLASLt.
• cuDNN: aten/src/ATen/native/cudnn/ holds convolution, RNN, batch-norm, and (newer) attention bindings. cuDNN-frontend is in third_party/cudnn-frontend/.
• cuSPARSE: aten/src/ATen/cuda/CUDASparseDescriptors.h etc. Used by sparse matmul.
• cuSolver: aten/src/ATen/cuda/CUDASolver.cpp. Used by linalg.solve, eigendecomp, etc.
• NCCL: torch/csrc/distributed/c10d/ProcessGroupNCCL.cpp. See Distributed.

Mixed precision and autocast

aten/src/ATen/autocast_mode.cpp registers Autocast* dispatch keys for CUDA, CPU, MPS, XPU. Inside an autocast context the keys are added to TLS, which causes the autocast kernels to wrap each op call: cast-down inputs to float16/bfloat16 for ops that benefit, leave others alone, and run the underlying kernel.

CUDA Graphs

at::cuda::CUDAGraph (aten/src/ATen/cuda/CUDAGraph.cpp) wraps cudaGraph capture/instantiate/replay. torch.cuda.graph(...) is the user-facing API; torch._inductor.cudagraph_trees is the auto-applied integration for torch.compile. CUDA Graphs replay a fixed sequence of kernel launches with no per-launch CPU overhead.

Flash / SDPA

aten/src/ATen/native/transformers/cuda/ ships a curated implementation of scaled-dot-product attention with three backends: math (fallback), mem_efficient_attention (xFormers/Cutlass), flash_attention (FlashAttention v2). The chooser is in sdp_utils_cpp.h.

ROCm

ROCm support uses hipify — a script-driven CUDA→HIP source translation — to produce HIP versions of every CUDA kernel. The translator lives at torch/utils/hipify/ and tools/amd_build/. ROCm-specific divergences live under aten/src/ATen/hip/ and c10/hip/.

Integration points

• torch.cuda Python API. Backed by torch/csrc/cuda/Module.cpp and friends.
• Distributed / NCCL. See Distributed.
• Compile. Inductor emits Triton kernels that launch on the same streams. CUDA Graphs are integrated automatically. See Inductor.
• MPS / XPU / MTIA. Other accelerators mirror this layout but with their own runtime helpers under c10/<backend>/ and kernels under aten/src/ATen/native/<backend>/. See MPS backend.

Entry points for modification

• New CUDA op → add a kernel under aten/src/ATen/native/cuda/, dispatch entry in native_functions.yaml, derivative in derivatives.yaml if needed.
• Caching allocator behaviour → c10/cuda/CUDACachingAllocator.cpp (this is sensitive code; expect a long review).
• Streams/events runtime → c10/cuda/CUDAStream.cpp, c10/cuda/CUDAEvent.cpp.
• New cuDNN binding → aten/src/ATen/native/cudnn/.
• For debugging memory: torch.cuda.memory._record_memory_history() then visualize at https://docs.pytorch.org/memory_viz.

Key source files

File	Purpose
c10/cuda/CUDACachingAllocator.cpp	Caching allocator
c10/cuda/CUDAStream.cpp	Stream wrapper
c10/cuda/CUDAFunctions.cpp	Runtime call helpers
aten/src/ATen/cuda/CUDABlas.cpp	cuBLAS bindings
aten/src/ATen/cuda/CUDAGraph.cpp	CUDA Graphs
aten/src/ATen/native/cuda/	Op kernels
aten/src/ATen/native/cudnn/	cuDNN-backed kernels
aten/src/ATen/native/transformers/cuda/	SDPA / flash attention
torch/cuda/__init__.py	Python torch.cuda
torch/csrc/cuda/Module.cpp	C++ binding for torch.cuda