Inside the AI Runtime: PyTorch, C++, CUDA and Beyond

Home » Python » Inside the AI Runtime: PyTorch, C++, CUDA and Beyond

Modern AI runtimes sit at the crossroads of Python, C++, JavaScript and CUDA-powered GPUs, and PyTorch has become the de facto example of how to glue all these worlds together efficiently. Behind the apparently simple Python APIs lies a sophisticated C++ core, automatic differentiation engine, JIT compiler, CUDA integration and even experimental runtimes like NVIDIA’s VibeTensor that explore how far AI-generated code can go.

If you are building high‑performance AI systems, understanding how PyTorch’s C++ runtime, CUDA backend, TorchScript, C++ extensions and memory management work is the difference between “it runs” and “it scales and flies”. In this guide we walk through those layers, explain how they relate to Python and JavaScript environments, and connect them with CUDA tooling and new approaches such as AI‑assisted runtimes.

PyTorch as a runtime: from Python API to C++ and CUDA

At first glance PyTorch looks like a pure Python library, but in reality it is a layered runtime whose core is written in C++ with tight integration to CUDA and optimized libraries such as cuDNN, NCCL, oneDNN and Intel MKL. Python mostly orchestrates operations, while the heavy lifting is done by native code on CPU and GPU.

From a high level, you can think of PyTorch as composed of several key components: the tensor library, the autograd engine, neural‑network utilities, the JIT compiler (TorchScript), multiprocessing helpers and a toolbox of utilities for data loading and serialization. Together they provide both a NumPy‑like tensor interface and a flexible deep‑learning platform.

The public C++ API roughly mirrors this architecture and is split into ATen, Autograd, the C++ Frontend, TorchScript bindings and C++ extensions, each addressing a specific layer of the computation stack. Understanding these pieces is essential if you plan to go beyond Python, integrate with custom C++ or CUDA kernels, or expose your models to environments like Node.js or browser runtimes via bindings.

Underneath everything, PyTorch leverages the NVIDIA CUDA Toolkit, which delivers a compiler, GPU‑accelerated math libraries, debugging and profiling tools, and the CUDA runtime used by torch.cuda and by custom extensions. This toolkit is what ultimately lets PyTorch saturate modern GPUs in data centers, workstations and embedded systems.

ATen: the tensor and operator backbone

ATen is the low‑level tensor library that backs almost every tensor operation in PyTorch, whether you call it from Python or from C++. It defines the core at::Tensor type plus hundreds of mathematical, indexing, linear‑algebra and reduction operations implemented for both CPU and GPU.

Each ATen tensor carries device and dtype metadata, and dispatching to CPU or CUDA implementations happens dynamically based on that metadata, so the same C++ operator symbol can seamlessly run on different backends. This is what allows you to move data via .to("cuda") or .cuda() in Python without changing the actual operation you call.

Although ATen can be used directly through the C++ API (everything lives in the at:: namespace), in typical workflows you only touch ATen indirectly through higher‑level modules such as torch, torch.nn or the C++ Frontend. Still, for performance‑critical code or custom kernels, working at the ATen level is often where you get the most control.

ATen’s design emphasizes fast vectorized kernels, strong GPU acceleration and good CPU performance, making it suitable both for research prototyping and production‑grade workloads. Internally, it leans on vendor libraries (MKL, cuBLAS, cuDNN, etc.) wherever that makes sense, instead of reinventing well‑optimized math.

Autograd: automatic differentiation on top of tensors

Autograd is the subsystem that augments ATen tensors with gradient tracking, turning pure tensor math into differentiable computation graphs. When autograd is enabled, operations on tensors generate an internal graph that can be traversed backward to compute gradients for training.

The autograd engine is built around a concept often called a “tape” or reverse‑mode differentiation graph, in which intermediate operations are recorded during the forward pass and replayed backwards on demand. Calling .backward() on a leaf tensor with a gradient seeds this graph traversal and accumulates derivatives into the parameters.

An important subtlety is that the raw at::Tensor type from ATen is not differentiable by default, so for gradient‑aware tensors in C++ you use the torch:: namespace factory functions (like torch::ones) instead of the at:: factories. The former creates tensors wired into the autograd system, whereas the latter stays purely numerical.

This separation is intentional: it lets you choose whether autograd overhead should exist, which is crucial for inference‑only or low‑latency paths where gradients are unnecessary. The same distinction is preserved in the C++ Frontend and in TorchScript graphs.

C++ Frontend: building and training models without Python

The PyTorch C++ Frontend is a high‑level API that mirrors much of torch.nn, torch.optim, torch.utils.data and related Python modules, but is implemented as idiomatic modern C++. It is designed for scenarios where you need native performance, tighter integration with existing C++ systems, or a Python‑free deployment environment.

Using the C++ Frontend, you define models as hierarchies of modules in the same spirit as torch.nn.Module, with parameters registered inside classes and composed into larger networks. There is a “standard library” of layers: convolutions, RNNs, batch normalization, linear layers and many others that behave almost identically to their Python counterparts.

The frontend also provides an optimization API with classic algorithms like SGD, Adam and RMSprop, dataset abstractions and data loaders capable of streaming data from many CPU threads, and utilities for serialization of checkpoints. Essentially, if you know how to write a training loop in Python (and avoid issues like overfitting vs underfitting), porting it to the C++ Frontend is mostly mechanical.

On multi‑GPU systems, the C++ Frontend exposes helpers for automatic model parallelization over several CUDA devices in a style similar to torch.nn.parallel.DataParallel. It additionally includes support code for binding C++ models back into Python with pybind11 when you want a hybrid workflow.

TorchScript: JIT compilation and Python‑free execution

TorchScript is a statically analyzable subset of Python plus the PyTorch API that can be compiled, optimized and serialized for deployment. Conceptually, it is a small programming language designed around tensor operations and control flow constructs commonly used in models.

From the C++ side, TorchScript exposes three primary capabilities: loading and running serialized models created in Python, defining custom operators that extend the TorchScript standard library, and compiling TorchScript source directly from C++. All of this is accessible via the torch::jit namespace.

A common production pattern is to author and train models entirely in Python, export them as TorchScript modules, and then embed those modules into a C++ service that performs low‑latency inference without depending on a Python runtime. This is particularly attractive in containerized microservices and performance‑sensitive backends.

Advanced users can register custom operators, including CUDA‑accelerated kernels, that TorchScript can invoke as if they were built‑in ops; these custom ops are serializable and work in both Python and C++ contexts when properly registered. Finally, functions like torch::jit::compile let you create TorchScript modules on the fly from C++ strings or AST‑like constructs.

C++ and CUDA extensions: plugging in your own kernels

PyTorch’s C++ extension mechanism offers a straightforward way to inject custom C++ and CUDA code into regular Python workflows while still taking advantage of ATen, autograd and the dispatcher. You typically use these extensions to implement custom operators or accelerate specialized workloads that are not well‑served by the built‑in operators.

The extension API itself doesn’t introduce new tensor semantics; instead, it wires your code into Python’s packaging system (via setuptools or alternative backends) and into PyTorch’s JIT compilation tools so that kernels compile with the correct ABI and link flags. Under the hood, helpers like CUDAExtension add boilerplate for nvcc, include paths and link options.

On the binding side, pybind11 is commonly used to expose C++ functions or classes as Python symbols, and those bindings interact with torch::Tensor objects that share storage with Python tensors. This keeps data copies to a minimum and allows autograd to see the new operators if you implement backward passes correctly.

Once compiled, these extensions behave like ordinary Python modules, can be imported from training scripts, and even play nicely with TorchScript if you register the ops for JIT usage. For CUDA‑heavy research, this route is often the quickest way to experiment with new kernels while preserving the productivity of Python.

Bridging Torch and CUDA libraries: a GEMM example with CUTLASS

When you need to exploit the full performance of NVIDIA GPUs, it is common to connect PyTorch’s tensor runtime with highly tuned CUDA C++ libraries such as CUTLASS or custom kernels that aren’t available as built‑in ops. A classic example is writing an optimized GEMM (matrix multiplication) operator.

The general pattern is to write a C++ wrapper function that accepts torch::Tensor inputs, extracts their shapes, dtypes and data pointers, and forwards them into CUTLASS or another CUDA library function. The tensor dimensions are obtained via .sizes(), data types via .dtype(), and raw pointers via .data_ptr() (optionally combined with reinterpret_cast if you need custom types like cutlass::half_t).

Because template parameters in C++ are resolved at compile time, yet tensor data types are only known at runtime, the wrapper usually contains conditional dispatch logic that selects the correct template instantiation (for float16, float32, etc.) based on the tensors’ dtypes. For complex template hierarchies, people often generate that boilerplate using Python scripts to avoid hand‑writing many branches.

Input validation is crucial: tensors need to have compatible shapes for GEMM, reside on a CUDA device, and be contiguous in memory because CUTLASS expects adjacent elements to be laid out sequentially. You can check contiguity with .is_contiguous() and fix it using .contiguous(), copying back results into an original tensor if necessary when in‑place semantics are desired.

To align with PyTorch’s style, you typically make output tensors optional using c10::optional<torch::Tensor>, and if none is provided, you allocate a fresh one with the right device and dtype using ATen factory functions. Returning that tensor keeps the API symmetric with built‑in operators like torch.mm.

Binding and building CUDA extensions with PyBind11 and setuptools

Once the C++ wrapper is ready, you must bind it to Python and compile it, which is where pybind11 and PyTorch’s build utilities come into play. A pybind11 module typically declares a function like m.def("cutlass_gemm", &cutlass_gemm, "GEMM with CUTLASS", py::arg("A"), py::arg("B"), py::arg("out") = py::none()); to expose your C++ function to Python code.

On the build side, plain setuptools does not natively understand nvcc, so PyTorch ships a helper class called CUDAExtension that configures include paths, CUDA compilation flags and linkage against libtorch automatically. You pass your .cpp and .cu sources to CUDAExtension much like a standard setuptools Extension.

When the extension is installed, your Python code can import it as a regular module, and from that point a call to your custom op looks almost identical to a call to built‑in PyTorch operators. You retain the ability to interoperate with autograd, TorchScript and CUDA streams as long as you respect PyTorch’s conventions.

If you prefer not to depend on PyTorch for the build itself (for example, if your CUDA library targets multiple frameworks), you can use alternative backends like scikit‑build‑core with CMake, or hand‑configure nvcc integration in setuptools, but CUDAExtension is by far the simplest for PyTorch users.

The CUDA Toolkit: compiler, libraries and runtime layer

The NVIDIA CUDA Toolkit underpins PyTorch’s GPU runtime by providing a C/C++ compiler, GPU‑optimized numerical libraries, debugging and profiling tools, and the low‑level runtime used by torch.cuda and custom extensions. It targets a broad range of platforms, from embedded boards to cloud clusters.

By delegating heavy linear algebra, convolution and FFT workloads to CUDA libraries, PyTorch leverages highly tuned implementations that know about each GPU architecture’s particularities and can unlock performance well beyond naïve GPU code. For many AI workloads this is the main reason PyTorch scales as well as it does on new GPU generations.

Developers writing CUDA extensions still typically write C++ kernels compiled with nvcc, and binding them back to PyTorch means they can be orchestrated from Python, C++ or even JavaScript runtimes that call into shared libraries. This is the sweet spot for teams that have legacy C++ or CUDA codebases but want to expose them through PyTorch‑style APIs.

torch.cuda: device selection, precision modes and execution model

The torch.cuda module is the user‑facing gateway to CUDA from PyTorch, handling device selection, stream management, memory allocation and runtime configuration. All CUDA tensors default to the currently selected device, which you can change using torch.cuda.device as a context manager or by targeting specific devices explicitly.

Cross‑GPU operations are intentionally restricted: aside from copy‑like operations (such as copy_(), to(), cuda()), most ops cannot span devices unless peer‑to‑peer memory access is enabled, to prevent subtle performance and correctness issues. Tensors keep track of their device, and outputs stay on the same device as their inputs.

PyTorch exposes detailed controls over computation precision on CUDA backends, including TensorFloat‑32 (TF32) on Ampere and later GPUs, reduced precision reductions for FP16 and BF16 GEMMs, and options for full FP16 accumulation when supported by the hardware. These flags can be tuned per backend (CUDA, cuDNN) and even per operator, letting you trade numerical accuracy for speed.

GPU operations are asynchronous by default: when you call a CUDA op, it is enqueued on a device stream but might execute later, allowing overlap of CPU computation, transfers and kernels. Timing such code correctly requires either explicit torch.cuda.synchronize() calls or the use of CUDA events; environment variables like CUDA_LAUNCH_BLOCKING=1 are handy for debugging but not for performance runs.

For more advanced use cases, PyTorch exposes CUDA streams and events, allowing you to orchestrate multiple parallel streams, synchronize selectively, and ensure tensors are not deallocated before all pending work on them finishes. Functions like record_stream() and wait_stream() are central in these patterns.

CUDA memory management, allocators and tuning knobs

PyTorch uses a caching allocator on CUDA devices to make frequent allocations and deallocations fast without constant synchronization with the GPU (see how memory works in C++ for related concepts). Instead of calling cudaMalloc and cudaFree for every tensor, it keeps pools of memory blocks that can be reused across allocations.

As a result, tools like nvidia-smi often show more memory “in use” than is actually occupied by tensors, because some of that memory is reserved by the allocator but currently unassigned. Functions such as memory_allocated(), max_memory_allocated(), memory_reserved() and max_memory_reserved() help distinguish between live tensor usage and cached capacity.

You can release unused cached memory with torch.cuda.empty_cache(), which hands blocks back to the CUDA driver but does not free memory still owned by live tensors. For deeper inspection, memory_stats() and memory_snapshot() provide low‑level allocation information that can be crucial when chasing fragmentation or OOMs.

The allocator behavior can be tuned via the PYTORCH_ALLOC_CONF (or its alias PYTORCH_CUDA_ALLOC_CONF) environment variable, which lets you pick the backend implementation, tweak split sizes, round‑up strategies, garbage‑collection thresholds and more. There is also an option to use CUDA’s cudaMallocAsync‑based allocator as an alternative backend on supported toolkits.

On top of the built‑in allocator, PyTorch supports pluggable CUDA allocators via shared libraries written in C or C++. These can integrate with external systems such as NCCL’s NVLink Switch Reductions or custom CPU‑GPU memory placement strategies, and are exposed in Python through torch.cuda.memory.CUDAPluggableAllocator and torch.cuda.MemPool.

CUDA Graphs and performance‑oriented best practices

CUDA Graphs are a powerful feature for reducing CPU overhead by capturing a batch of GPU operations as a graph and replaying them with a single launch call. PyTorch integrates this feature via torch.cuda.CUDAGraph, the torch.cuda.graph context manager and the torch.cuda.make_graphed_callables() helper.

Graph capture works best when your workload has static shapes, deterministic control flow and no CPU‑GPU synchronization in the hot path. During capture, GPU work is recorded rather than executed, and on replay the exact same sequence of kernels runs, reading from and writing to the same virtual addresses.

Because the captured graph assumes fixed tensor layouts, PyTorch allocates graph‑private memory pools whose lifetimes are tied to the CUDAGraph object and to the tensors created inside the capture. These pools can be shared across related graphs to conserve memory, provided you guarantee a consistent execution order and no concurrency.

In distributed setups, especially with DistributedDataParallel and NCCL, CUDA Graph integration requires some care, such as ensuring compatible NCCL versions, disabling certain async error handlers, and handling warmup steps to allow collectives to be captured safely. Done correctly, it can dramatically increase throughput for stable, production‑grade training loops.

Beyond graphs, PyTorch documentation recommends using pinned CPU memory for faster host‑to‑GPU transfers, preferring DistributedDataParallel over DataParallel or naive multiprocessing for multi‑GPU training, and writing device‑agnostic code that switches cleanly between CPU and GPU via torch.device. These practices collectively help you hit the sweet spot between ergonomics and performance.

Installation, build options and hardware backends

PyTorch can be installed via Conda, pip wheels or from source, with prebuilt binaries available for typical Linux, macOS and Windows setups and specialized builds for NVIDIA Jetson platforms. Each distribution targets particular CUDA versions according to a published support matrix.

When building from source, you need Python 3.10 or newer, a modern C++17‑capable compiler (for example GCC 9.4+ on Linux) and the appropriate toolchain on Windows (Visual Studio or the standalone Build Tools). You can compile with or without CUDA, ROCm (for AMD GPUs) or Intel GPU support by installing the corresponding SDKs and setting flags like USE_CUDA, USE_ROCM and USE_XPU.

For CPU‑only builds, special care may be required to link the desired OpenMP implementation (often Intel OpenMP) and BLAS libraries, especially on Windows where CMake may otherwise fall back to the default MSVC runtime. For CUDA builds, additional libraries like Magma or oneDNN can be brought in to accelerate specific operations.

PyTorch also offers Docker images with preconfigured CUDA and cuDNN environments (see an introduction to containerization); these images rely on shared memory segments for multiprocessing data loaders, so you often need to increase shared memory via --ipc=host or --shm-size on docker run. Ensuring your Docker and NVIDIA driver versions are compatible with the CUDA toolkit version is a must.

The documentation itself is built with Sphinx and a custom theme, and can be generated in HTML or PDF form if you install the necessary Python dependencies, TeX tooling and PyTorch package in a local environment. This is useful when you add new modules or docstrings and want to preview them before contributing upstream.

VibeTensor: AI‑generated runtime inspired by PyTorch

NVIDIA’s VibeTensor project illustrates a different angle on runtimes: it is an experimental execution environment, conceptually similar to PyTorch, whose code base was largely generated by AI agents under human supervision. The idea is to see how far “vibe coding” — relying heavily on AI assistants to write code — can be pushed for a complex system.

VibeTensor’s architecture combines a Python and JavaScript‑friendly API with a C++ runtime core, custom tensor storage allocators, an autograd engine, a dispatcher, an advanced indexing subsystem and a CUDA‑driven memory cache, all targeting Linux x86_64 and NVIDIA GPUs. There is even an experimental Fabric subsystem that uses CUDA P2P for multi‑GPU execution.

The project also supports external GPU plugins, such as a backend for the upcoming Blackwell architecture (SM100/SM103), and demonstrates that these plugins themselves can be bootstrapped with AI‑generated code as long as humans provide constraints and validation. While performance and features do not currently rival PyTorch, the project serves as a proof of concept for AI‑assisted systems programming.

During the roughly two months of development, human engineers mostly focused on defining tasks, constraints and review loops, while AI agents iteratively produced code, ran comparisons, compiled, tested and refined implementations. This hybrid workflow underscores both the power and the limitations of AI: the bulk of boilerplate and repetitive patterns can be automated, but correctness and architecture still demand human judgment.

For practitioners, VibeTensor is less a drop‑in replacement for PyTorch and more an exploration of what future runtimes might look like when AI is deeply involved in their design and evolution, especially in environments that mix Python, JavaScript, C++ and CUDA. It hints at a world where specialized runtimes for niche workloads can be spun up far faster than before.

Putting all these pieces together — PyTorch’s tensor and autograd core, the C++ Frontend, TorchScript, C++ and CUDA extensions, finely tuned CUDA integration and experimental efforts like VibeTensor — you get a picture of an AI runtime ecosystem where Python, C++, JavaScript and CUDA are tightly interwoven, and where developers can fluidly choose the layer that best balances productivity, control and raw performance for their particular workload.