refactor: register CuTe attention and fusion in sd.cpp

This commit is contained in:
leejet committed 2026-10-31 19:00:24 +08:00
1 parent c6a4666289
commit 4b274821e9
5 files changed
+946 -10

No files matched your search

+15 -3
View File
@@ -8,6 +8,10 @@
#include <mutex>
#include <unordered_set>
#ifdef SD_CUDA_EXTENSIONS
void sd_ggml_register_cuda_attention(ggml_backend_dev_t device);
#endif
static bool rope_infer(ggml_tensor* const* inputs, int count, const void*, ggml_ext_tensor_desc* out) {
if (count != 2) {
return false;
@@ -155,9 +159,11 @@ void sd_ggml_register_extensions(ggml_backend_t backend) {
if (!backend) {
return;
}
auto device = ggml_backend_get_device(backend);
auto reg = ggml_backend_dev_backend_reg(device);
if (!ggml_backend_reg_get_proc_address(reg, "ggml_ext_cpu_parallel_v1")) {
auto device = ggml_backend_get_device(backend);
auto reg = ggml_backend_dev_backend_reg(device);
const bool cpu = ggml_backend_reg_get_proc_address(reg, "ggml_ext_cpu_parallel_v1") != nullptr;
const bool cuda = ggml_backend_reg_get_proc_address(reg, "ggml_ext_cuda_context_v1") != nullptr;
if (!cpu && !cuda) {
return;
}
static std::mutex mutex;
@@ -166,6 +172,12 @@ void sd_ggml_register_extensions(ggml_backend_t backend) {
if (!registered.insert(device).second) {
return;
}
if (cuda) {
#ifdef SD_CUDA_EXTENSIONS
sd_ggml_register_cuda_attention(device);
#endif
return;
}
ggml_ext_kernel_desc desc = {};
desc.abi_version = GGML_EXT_ABI_VERSION;
desc.name = "sd.rope_apply.cpu.v1";