torch.accelerator: A Unified, Device-Agnostic Runtime API for Stream-Based Accelerators - Yu Guangye

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