feat: 增加摩尔线程大模型算子和部分传统模型算子

operatorspy/tests/add.py

@@ -115,6 +115,16 @@ def test_bang(lib, test_cases):
         test(lib, handle, "mlu", c_shape, a_shape, b_shape, tensor_dtype=torch.float32, inplace=inplace)
     destroy_handle(lib, handle)
 
+def test_musa(lib, test_cases):
+    import torch_musa
+
+    device = DeviceEnum.DEVICE_MUSA
+    handle = create_handle(lib, device)
+    for c_shape, a_shape, b_shape, inplace in test_cases:
+        test(lib, handle, "musa", c_shape, a_shape, b_shape, tensor_dtype=torch.float16, inplace=inplace)
+        test(lib, handle, "musa", c_shape, a_shape, b_shape, tensor_dtype=torch.float32, inplace=inplace)
+    destroy_handle(lib, handle)
+
 
 if __name__ == "__main__":
     test_cases = [
@@ -163,6 +173,8 @@ def test_bang(lib, test_cases):
         test_cuda(lib, test_cases)
     if args.bang:
         test_bang(lib, test_cases)
-    if not (args.cpu or args.cuda or args.bang):
+    if args.musa:
+        test_musa(lib, test_cases)
+    if not (args.cpu or args.cuda or args.bang or args.musa):
         test_cpu(lib, test_cases)
     print("\033[92mTest passed!\033[0m")

operatorspy/tests/causal_softmax.py

@@ -119,6 +119,16 @@ def test_maca(lib, test_cases):
 
     destroy_handle(lib, handle)
 
+def test_musa(lib, test_cases):
+    import torch_musa
+    device = DeviceEnum.DEVICE_MUSA
+
+    handle = create_handle(lib, device)
+    for x_shape, x_stride in test_cases:
+        test(lib, handle, "musa", x_shape, x_stride)
+
+    destroy_handle(lib, handle)
+
 if __name__ == "__main__":
     test_cases = [
         # x_shape, x_stride
@@ -161,6 +171,8 @@ def test_maca(lib, test_cases):
         test_ascend(lib, test_cases)
     if args.maca:
         test_maca(lib, test_cases)
-    if not (args.cpu or args.cuda or args.bang or args.ascend or args.maca):
+    if args.musa:
+        test_musa(lib, test_cases)
+    if not (args.cpu or args.cuda or args.bang or args.ascend or args.maca or args.musa):
         test_cpu(lib, test_cases)
     print("\033[92mTest passed!\033[0m")

operatorspy/tests/expand.py

@@ -133,6 +133,16 @@ def test_bang(lib, test_cases):
         test(lib, handle, "mlu", y_shape, x_shape, y_stride, x_stride, tensor_dtype=torch.float32)
     destroy_handle(lib, handle)
 
+def test_musa(lib, test_cases):
+    import torch_musa
+
+    device = DeviceEnum.DEVICE_MUSA
+    handle = create_handle(lib, device)
+    for y_shape, x_shape, y_stride, x_stride in test_cases:
+        test(lib, handle, "musa", y_shape, x_shape, y_stride, x_stride, tensor_dtype=torch.float16)
+        test(lib, handle, "musa", y_shape, x_shape, y_stride, x_stride, tensor_dtype=torch.float32)
+    destroy_handle(lib, handle)
+
 
 if __name__ == "__main__":
     test_cases = [
@@ -174,6 +184,8 @@ def test_bang(lib, test_cases):
         test_cuda(lib, test_cases)
     if args.bang:
         test_bang(lib, test_cases)
-    if not (args.cpu or args.cuda or args.bang):
+    if args.musa:
+        test_musa(lib, test_cases)
+    if not (args.cpu or args.cuda or args.bang or args.musa):
         test_cpu(lib, test_cases)
     print("\033[92mTest passed!\033[0m")

operatorspy/tests/matmul.py

@@ -325,6 +325,37 @@ def test_maca(lib, test_cases):
 
     destroy_handle(lib, handle)
 
+def test_musa(lib, test_cases):
+    import torch_musa
+
+    device = DeviceEnum.DEVICE_MUSA
+    handle = create_handle(lib, device)
+    for (
+        alpha,
+        beta,
+        a_shape,
+        b_shape,
+        c_shape,
+        a_stride,
+        b_stride,
+        c_stride,
+        dtype,
+    ) in test_cases:
+        test(
+            lib,
+            handle,
+            "musa",
+            alpha,
+            beta,
+            a_shape,
+            b_shape,
+            c_shape,
+            a_stride,
+            b_stride,
+            c_stride,
+            dtype,
+        )
+
 if __name__ == "__main__":
     test_cases = [
         # alpha, beta, a_shape, b_shape, c_shape, a_stride, b_stride, c_stride, dtype
@@ -387,6 +418,8 @@ def test_maca(lib, test_cases):
         test_ascend(lib, test_cases)
     if args.maca:
         test_maca(lib, test_cases)
-    if not (args.cpu or args.cuda or args.bang or args.ascend or args.maca):
+    if args.musa:
+        test_musa(lib, test_cases)
+    if not (args.cpu or args.cuda or args.bang or args.ascend or args.maca or args.musa):
         test_cpu(lib, test_cases)
     print("\033[92mTest passed!\033[0m")

operatorspy/tests/random_sample.py

@@ -94,7 +94,7 @@ def test(lib, handle, torch_device, voc, random_val, topp, topk, temperature, x_
     if(torch_device == 'maca'):
         indices = torch.zeros([1], dtype = torch.int64).to('cuda')
     else:
-        indices = torch.zeros([1], dtype = torch.uint64).to(torch_device)
+        indices = torch.zeros([1], dtype = torch.int64).to(torch_device)
     x_tensor = to_tensor(data, lib)
     indices_tensor = to_tensor(indices, lib)
     indices_tensor.descriptor.contents.dt = U64  # treat int64 as uint64
@@ -170,7 +170,7 @@ def test_ascend(lib, test_cases):
     for (voc, random_val, topp, topk, temperature) in test_cases:
         test(lib, handle, "npu", voc, random_val, topp, topk, temperature)
     destroy_handle(lib, handle) 
-
+    
 def test_maca(lib, test_cases):
     device = DeviceEnum.DEVICE_MACA
     handle = create_handle(lib, device)
@@ -179,6 +179,13 @@ def test_maca(lib, test_cases):
     destroy_handle(lib, handle) 
 
 
+def test_musa(lib, test_cases):
+    import torch_musa
+    device = DeviceEnum.DEVICE_MUSA
+    handle = create_handle(lib, device)
+    for (voc, random_val, topp, topk, temperature) in test_cases:
+        test(lib, handle, "musa", voc, random_val, topp, topk, temperature)
+    destroy_handle(lib, handle) 
 
 if __name__ == "__main__":
     test_cases = [
@@ -236,6 +243,8 @@ def test_maca(lib, test_cases):
         test_ascend(lib, test_cases)
     if args.maca:
         test_maca(lib, test_cases)
-    if not (args.cpu or args.cuda or args.bang or args.ascend or args.maca):
+    if args.musa:
+        test_musa(lib, test_cases)
+    if not (args.cpu or args.cuda or args.bang or args.ascend or args.maca or args.musa):
         test_cpu(lib, test_cases)
     print("\033[92mTest passed!\033[0m")

operatorspy/tests/rearrange.py

@@ -117,6 +117,26 @@ def test_maca(lib, test_cases):
         test(lib, handle, "cuda", x_shape, x_stride, y_shape, y_stride)
     destroy_handle(lib, handle) 
 
+def test_musa(lib, test_cases):
+    import torch_musa
+    device = DeviceEnum.DEVICE_MUSA
+    handle = create_handle(lib, device)
+    for test_case in test_cases:
+        x_shape, x_stride = test_case[0]
+        y_shape, y_stride = test_case[1]
+        test(lib, handle, "musa", x_shape, x_stride, y_shape, y_stride)
+    destroy_handle(lib, handle)
+
+def test_musa(lib, test_cases):
+    import torch_musa
+    device = DeviceEnum.DEVICE_MUSA
+    handle = create_handle(lib, device)
+    for test_case in test_cases:
+        x_shape, x_stride = test_case[0]
+        y_shape, y_stride = test_case[1]
+        test(lib, handle, "musa", x_shape, x_stride, y_shape, y_stride)
+    destroy_handle(lib, handle)
+
 if __name__ == "__main__":
     args = get_args()
     test_cases = [
@@ -156,4 +176,6 @@ def test_maca(lib, test_cases):
         test_ascend(lib, test_cases)
     if args.maca:
         test_maca(lib, test_cases)
+    if args.musa:
+        test_musa(lib, test_cases)
     print("\033[92mTest passed!\033[0m")

operatorspy/tests/relu.py

@@ -132,6 +132,16 @@ def test_bang(lib, test_cases):
         test(lib, handle, "mlu", tensor_shape, tensor_dtype=torch.float32, inplace=inplace)
     destroy_handle(lib, handle)
 
+def test_musa(lib, test_cases):
+    import torch_musa
+
+    device = DeviceEnum.DEVICE_MUSA
+    handle = create_handle(lib, device)
+    for tensor_shape, inplace in test_cases:
+        test(lib, handle, "musa", tensor_shape, tensor_dtype=torch.float16, inplace=inplace)
+        test(lib, handle, "musa", tensor_shape, tensor_dtype=torch.float32, inplace=inplace)
+    destroy_handle(lib, handle)
+
 
 if __name__ == "__main__":
     test_cases = [
@@ -172,6 +182,8 @@ def test_bang(lib, test_cases):
         test_cuda(lib, test_cases)
     if args.bang:
         test_bang(lib, test_cases)
-    if not (args.cpu or args.cuda or args.bang):
+    if args.musa:
+        test_musa(lib, test_cases)
+    if not (args.cpu or args.cuda or args.bang or args.musa):
         test_cpu(lib, test_cases)
     print("\033[92mTest passed!\033[0m")

operatorspy/tests/rms_norm.py

@@ -125,6 +125,14 @@ def test_maca(lib, test_cases):
 
     destroy_handle(lib, handle)
 
+def test_musa(lib, test_cases):
+    import torch_musa
+    device = DeviceEnum.DEVICE_MUSA
+    handle = create_handle(lib, device)
+    for (y_shape, x_shape, w_shape, dtype, w_dtype) in test_cases:
+        test(lib, handle, "musa", y_shape, x_shape, w_shape, dtype, w_dtype)
+    destroy_handle(lib, handle)
+
 if __name__ == "__main__":
     test_cases = [
         # y_shape, x_shape, w_shape, dtype, w_dtype
@@ -174,6 +182,8 @@ def test_maca(lib, test_cases):
         test_ascend(lib, test_cases)
     if args.maca:
         test_maca(lib, test_cases)
-    if not (args.cpu or args.cuda or args.bang or args.ascend or args.maca):
+    if args.musa:
+        test_musa(lib, test_cases)
+    if not (args.cpu or args.cuda or args.bang or args.ascend or args.maca or args.musa):
         test_cpu(lib, test_cases)
     print("\033[92mTest passed!\033[0m")

operatorspy/tests/rotary_embedding.py

@@ -77,7 +77,7 @@ def test(lib, handle, torch_device, shape, strides=None, dtype=torch.float16):
         pos[2 * i] = posTmp[i]
         pos[2 * i + 1] = 0
     theta = 1e4
-    if torch_device == 'mlu' or torch_device == 'npu':
+    if torch_device == 'mlu' or torch_device == 'npu' or torch_device == 'musa':
         ans = rotary_embedding(t, posTmp, theta, "cpu").to(torch_device)
         pos = pos.to(torch_device)
         t = t.to(torch_device)
@@ -181,6 +181,14 @@ def test_maca(lib, test_cases) :
         test(lib, handle, "maca", shape, strides, dtype)
     destroy_handle(lib, handle)
 
+def test_musa(lib, test_cases) :
+    import torch_musa
+    device = DeviceEnum.DEVICE_MUSA
+    handle = create_handle(lib, device)
+    for shape, strides, dtype in test_cases:
+        test(lib, handle, "musa", shape, strides, dtype)
+    destroy_handle(lib, handle)
+
 if __name__ == "__main__":
     test_cases = [
         ((1, 32, 128), None, torch.float16),
@@ -233,6 +241,8 @@ def test_maca(lib, test_cases) :
         test_ascend(lib, test_cases)
     if args.maca:
         test_maca(lib, test_cases)
-    if not (args.cpu or args.cuda or args.bang or args.ascend or args.maca):
+    if args.musa:
+        test_musa(lib, test_cases)
+    if not (args.cpu or args.cuda or args.bang or args.ascend or args.maca or args.musa):
         test_cpu(lib, test_cases)
     print("\033[92mTest passed!\033[0m")

operatorspy/tests/swiglu.py

@@ -262,6 +262,20 @@ def test_maca(lib, test_cases):
 
     destroy_handle(lib, handle) 
 
+def test_musa(lib, test_cases):
+    import torch_musa
+    device = DeviceEnum.DEVICE_MUSA
+    handle = create_handle(lib, device)
+
+    for shape, a_stride, b_stride, c_stride, dtype in test_cases:
+        test_out_of_place(
+            lib, handle, "musa", shape, a_stride, b_stride, c_stride, dtype
+        )
+        test_in_place1(lib, handle, "musa", shape, a_stride, b_stride, dtype)
+        test_in_place2(lib, handle, "musa", shape, a_stride, b_stride, dtype)
+
+    destroy_handle(lib, handle) 
+
 
 if __name__ == "__main__":
     test_cases = [
@@ -307,4 +321,6 @@ def test_maca(lib, test_cases):
         test_ascend(lib, test_cases)
     if args.maca:
         test_maca(lib, test_cases)
+    if args.musa:
+        test_musa(lib, test_cases)
     print("\033[92mTest passed!\033[0m")

operatorspy/tests/test_utils.py

@@ -32,6 +32,11 @@ def get_args():
         action="store_true",
         help="Run ASCEND NPU test",
     )
+    parser.add_argument(
+        "--musa",
+        action="store_true",
+        help="Run MUSA test",
+    )
 
     return parser.parse_args()
 

src/devices/handle.cc

@@ -14,6 +14,9 @@
 #ifdef ENABLE_METAX_GPU
 #include "./maca/maca_handle.h"
 #endif
+#ifdef ENABLE_MTHREADS_GPU
+#include "./musa/musa_handle.h"
+#endif
 
 
 __C infiniopStatus_t infiniopCreateHandle(infiniopHandle_t *handle_ptr, Device device, int device_id) {
@@ -48,6 +51,11 @@ __C infiniopStatus_t infiniopCreateHandle(infiniopHandle_t *handle_ptr, Device d
         case DevMetaxGpu: {
             return createMacaHandle((MacaHandle_t *) handle_ptr, device_id);
         }
+#endif
+#ifdef ENABLE_MTHREADS_GPU
+        case DevMthreadsGpu: {
+            return createMusaHandle((MusaHandle_t *) handle_ptr, device_id);
+        }
 #endif
     }
     return STATUS_BAD_DEVICE;
@@ -81,6 +89,12 @@ __C infiniopStatus_t infiniopDestroyHandle(infiniopHandle_t handle) {
         case DevMetaxGpu: {
             return deleteMacaHandle((MacaHandle_t) handle);
         }
+#endif
+#ifdef ENABLE_MTHREADS_GPU
+        case DevMthreadsGpu: {
+            deleteMusaHandle((MusaHandle_t) handle);
+            return STATUS_SUCCESS;
+        }
 #endif
     }
     return STATUS_BAD_DEVICE;