kvcache-ai/ktransformers
1
0
Fork 0
mirror of https://github.com/kvcache-ai/ktransformers.git synced 2026-03-15 02:47:22 +00:00
Code Issues Packages Projects Releases Wiki Activity

musa: support bf16

Browse Source 
Signed-off-by: Xiaodong Ye <xiaodong.ye@mthreads.com>
This commit is contained in:
Xiaodong Ye
2025-02-23 10:19:19 +08:00
parent 94ab2de3b9
commit 18b1d18367
2 changed files with 4 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
ktransformers/ktransformers_ext/cpu_backend/vendors/musa.h vendored
View File

@@ -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

1
setup.py
View File

@@ -350,6 +350,7 @@ elif MUSA_HOME is not 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',