Merge pull request #608 from makllama/fix_musa_ext

musa: support bf16

Author: Atream

ktransformers/ktransformers_ext/cpu_backend/vendors/musa.h

@@ -1,7 +1,9 @@
 #pragma once
 
 #include <musa_runtime.h>
+#include <musa_bf16.h>
 
 #define cudaLaunchHostFunc musaLaunchHostFunc
 #define cudaStream_t musaStream_t
-#define cudaHostFn_t musaHostFn_t
+#define cudaHostFn_t musaHostFn_t
+#define nv_bfloat16 mt_bfloat16

setup.py

@@ -350,6 +350,7 @@ def build_extension(self, ext) -> None:
         "at::cuda": "at::musa",
         "#include <ATen/cuda/CUDAContext.h>": "#include \"torch_musa/csrc/aten/musa/MUSAContext.h\"",
         "#include <c10/cuda/CUDAGuard.h>": "#include \"torch_musa/csrc/core/MUSAGuard.h\"",
+        "nv_bfloat16": "mt_bfloat16",
         }).run()
     ops_module = MUSAExtension('KTransformersOps', [
         'ktransformers/ktransformers_ext/cuda_musa/custom_gguf/dequant.mu',