import torchada # ← Add this one line at the top import torch # Your existing CUDA code works unchanged: x = torch.randn(10, 10).cuda() print(torch.cuda.device_count()) torch.cuda.synchronize()

import torchada import torch model = MyModel().cuda() scaler = torch.cuda.amp.GradScaler() with torch.cuda.amp.autocast(): output = model(data.cuda()) loss = criterion(output, target.cuda()) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()

import torchada import torch.distributed as dist # 'nccl' is automatically mapped to 'mccl' on MUSA dist.init_process_group(backend='nccl')

import torchada import torch g = torch.cuda.CUDAGraph() with torch.cuda.graph(cuda_graph=g): # cuda_graph= keyword works on MUSA y = model(x)

import torchada import torch compiled_model = torch.compile(model.cuda(), backend='inductor')

import torchada # Must import before torch.utils.cpp_extension from torch.utils.cpp_extension import CUDAExtension, BuildExtension # Standard CUDAExtension works — torchada handles CUDA→MUSA translation ext = CUDAExtension("my_ext", sources=["kernel.cu"])

import torchada import torch my_lib = torch.library.Library("my_lib", "DEF") my_lib.define("my_op(Tensor x) -> Tensor") my_lib.impl("my_op", my_func, "CUDA") # Works on MUSA!

import torchada import torch # ProfilerActivity.CUDA works on MUSA with torch.profiler.profile( activities=[torch.profiler.ProfilerActivity.CPU, torch.profiler.ProfilerActivity.CUDA] ) as prof: model(x)

import torchada import ctypes # Load MUSA runtime library with CUDA function names lib = ctypes.CDLL("libmusart.so") func = lib.cudaMalloc # Automatically translates to musaMalloc # Works with MCCL too nccl_lib = ctypes.CDLL("libmccl.so") func = nccl_lib.ncclAllReduce # Automatically translates to mcclAllReduce

import torchada from torchada import detect_platform, Platform platform = detect_platform() if platform == Platform.MUSA: print("Running on Moore Threads GPU") elif platform == Platform.CUDA: print("Running on NVIDIA GPU") # Or use torch.version-based detection def is_musa(): import torch return hasattr(torch.version, 'musa') and torch.version.musa is not None

device = torch.device("cuda:0") # On MUSA, this becomes musa:0 device.type == "cuda" # Returns False!

import torchada if torchada.is_gpu_device(device): # Works on both CUDA and MUSA ... # Or: device.type in ("cuda", "musa")

# At your application entry point def is_musa(): import torch return hasattr(torch.version, "musa") and torch.version.musa is not None if is_musa(): import torchada # noqa: F401 # Rest of your code uses torch.cuda.* as normal

# Include MUSA in GPU capability checks if is_nvidia() or is_musa(): ENABLE_FLASH_ATTENTION = True

# Instead of: device.type == "cuda" # Use: device.type in ("cuda", "musa") # Or: torchada.is_gpu_device(device)