torch.accelerator: A Unified, Device-Agnostic Runtime API for Stream-Based Accelerators - Yu Guangye, Intel
Motivation
PyTorch supports a wide range of acceleration hardware beyond CPUs, including CUDA, XPU, MPS, NPU, HPU, and more. Its architecture allows new backend integration through two key components: ATen operators and device runtime.
While ATen operators are device-agnostic, the runtime remains device-specific, relying on APIs like torch.cuda and torch.xpu. This fragmentation complicates writing portable, hardware-agnostic code across PyTorch and its ecosystem.
To address this challenge, we propose torch.accelerator: a unified, device-agnostic runtime API for stream-based accelerators.
Design
An Accelerator refers to a device that collaborates with the CPU to accelerate computation, typically via asynchronous execution using Stream and Event for synchronization. Our design assumes a single active accelerator per host.
The torch.accelerator API provides a consistent interface for device and stream management, with backend support integrated via the existing DeviceGuardImplInterface registration mechanism.
Further Work
We are actively working on a unified device memory API. These will streamline the library, model, and UTs.
Reference - https://github.com/pytorch/pytorch/pull/132204