vrr/kvcache-ai-ktransformers
2
0
Fork 0
mirror of https://github.com/kvcache-ai/ktransformers.git synced 2025-09-06 20:49:55 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 2

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
Show stats Download patch file Download diff file Expand all files Collapse all files

2
ktransformers/ktransformers_ext/cpu_backend/vendors/musa.h vendored
View file

@ -1,7 +1,9 @@
#pragma once #pragma once
#include <musa_runtime.h> #include <musa_runtime.h>
#include <musa_bf16.h>
#define cudaLaunchHostFunc musaLaunchHostFunc #define cudaLaunchHostFunc musaLaunchHostFunc
#define cudaStream_t musaStream_t #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", "at::cuda": "at::musa",
"#include <ATen/cuda/CUDAContext.h>": "#include \"torch_musa/csrc/aten/musa/MUSAContext.h\"", "#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\"", "#include <c10/cuda/CUDAGuard.h>": "#include \"torch_musa/csrc/core/MUSAGuard.h\"",
"nv_bfloat16": "mt_bfloat16",
}).run() }).run()
ops_module = MUSAExtension('KTransformersOps', [ ops_module = MUSAExtension('KTransformersOps', [
'ktransformers/ktransformers_ext/cuda_musa/custom_gguf/dequant.mu', 'ktransformers/ktransformers_ext/cuda_musa/custom_gguf/dequant.mu',