$ show MooreThreads/torch_musa
This repository implements a full PyTorch backend for Moore Threads MUSA GPUs, providing a drop-in replacement for `torch.cuda`. It enables PyTorch models to run on MUSA hardware with minimal code changes by registering a `PrivateUse1` device backend. The architecture consists of C++ ATen operator implementations (over 200 ops), wrappers around the proprietary muDNN deep neural network library, a core runtime layer (streams, events, caching allocator, IPC), distributed training via MCCL (MUSA Collective Communications Library), an Inductor compiler backend with Triton support, quantization (W4A16, int8), and profiling tools. Key entry points include `torch_musa/__init__.py` for autoload, `setup.py` for building, `torch_musa/csrc/core/MUSAHooksInterface.cpp` for device registration, and `torch_musa/csrc/aten/ops/` for operator implementations. Data flows follow PyTorch's dispatch mechanism: operations are routed to MUSA-specific kernels via the `PrivateUse1` dispatch key, with muDNN handling most deep learning primitives. External dependencies include PyTorch (patched for version 2.11.0), muDNN, MCCL, musolver, muFFT, and Triton. The repository also includes automated CUDA-to-MUSA porting tools (`tools/cuda_porting/`) and a comprehensive test suite covering operators, distributed training, quantization, and sparse tensors.
The LLM identified 4 files worth highlighting for reuse:
Automated script that transforms CUDA source files (`.cu`, `.cuh`) into MUSA equivalents (`.mu`, `.muh`) using regex-based replacement rules and an automaton for context-sensitive rewrites. It handles API renaming, header substitution, and macro adaptation. This tool is self-contained and could be reused for porting CUDA code to other GPU backends with a similar mapping approach.
Provides `OpTest`, a generic test harness for comparing CPU and MUSA tensor operation results. It supports configurable comparators (`DefaultComparator` with absolute/relative tolerance), gradient checks, in-place variants, and multiple dtypes. This framework is reusable for testing any backend that implements PyTorch operators, requiring only a device name and comparator customization.
Implements all binary operations (add, sub, mul, div, comparison, logical, pow, etc.) for MUSA using muDNN. Uses `DEFINE_BINARY_OP` and `DEFINE_BINARY_GRAD_OP` macros to reduce boilerplate for tensor/tensor_, out variants. Includes an optimization that redirects scalar-second operands to a unary call for performance. The macro pattern and scalar optimization logic are reusable for other backends.
Defines the interface and data structures for MUSA's caching memory allocator, including block pools, stream associations, and snapshot recording. Supports configurable policies (expandable segments, garbage collection) and a pluggable allocator interface. While MUSA-specific, the design pattern of a caching allocator with stream-aware recycling and memory statistics is reusable for other GPU-like devices.