From adcac69650769411e702c440f03257b22c02973e Mon Sep 17 00:00:00 2001 From: leejet Date: Fri, 18 Sep 2026 23:59:58 +0800 Subject: [PATCH] perf: pad small attention heads to 64 for MMA Flash Attention (#1992) --- CMakeLists.txt | 5 ++ cmake/stable-diffusion-config.cmake.in | 7 +++ src/core/ggml_extend.cpp | 14 ++++++ src/core/ggml_extend_backend.cpp | 68 ++++++++++++++++++++++++++ src/core/ggml_extend_backend.h | 1 + 5 files changed, 95 insertions(+) diff --git a/CMakeLists.txt b/CMakeLists.txt index 39e639f3..2c4b8a69 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -343,6 +343,11 @@ add_subdirectory(thirdparty) target_sources(${SD_LIB} PRIVATE $) target_link_libraries(${SD_LIB} PUBLIC ggml) target_link_libraries(${SD_LIB} PRIVATE onig sd-utf8proc) +if (SD_CUDA) + find_package(CUDAToolkit REQUIRED) + target_link_libraries(${SD_LIB} PRIVATE CUDA::cuda_driver) + set_property(SOURCE src/core/ggml_extend_backend.cpp APPEND PROPERTY COMPILE_DEFINITIONS SD_USE_CUDA) +endif() target_include_directories(${SD_LIB} PUBLIC . src include) target_include_directories(${SD_LIB} PRIVATE src/core) target_include_directories(${SD_LIB} PUBLIC . thirdparty) diff --git a/cmake/stable-diffusion-config.cmake.in b/cmake/stable-diffusion-config.cmake.in index c3d7433c..b31f4416 100644 --- a/cmake/stable-diffusion-config.cmake.in +++ b/cmake/stable-diffusion-config.cmake.in @@ -10,6 +10,9 @@ set(SD_BIN_DIR "@PACKAGE_SD_BIN_INSTALL_DIR@") include(CMakeFindDependencyMacro) find_dependency(ggml REQUIRED HINTS "${SD_LIB_DIR}/cmake") +if(@SD_CUDA@ AND NOT SD_SHARED_LIB) + find_dependency(CUDAToolkit REQUIRED) +endif() if(NOT TARGET stable-diffusion) find_library(stable-diffusion_LIBRARY stable-diffusion @@ -28,6 +31,10 @@ if(NOT TARGET stable-diffusion) INTERFACE_COMPILE_FEATURES "c_std_11;cxx_std_17" POSITION_INDEPENDENT_CODE ON) + if(@SD_CUDA@ AND NOT SD_SHARED_LIB) + set_property(TARGET stable-diffusion APPEND PROPERTY INTERFACE_LINK_LIBRARIES CUDA::cuda_driver) + endif() + if(SD_SHARED_LIB) target_compile_definitions(stable-diffusion INTERFACE SD_BUILD_SHARED_LIB) diff --git a/src/core/ggml_extend.cpp b/src/core/ggml_extend.cpp index 79e3ea9b..36f70d79 100644 --- a/src/core/ggml_extend.cpp +++ b/src/core/ggml_extend.cpp @@ -643,6 +643,14 @@ ggml_tensor* ggml_ext_attention_ext(ggml_context* ctx, ggml_tensor* kqv = nullptr; auto build_kqv = [&](ggml_tensor* q_in, ggml_tensor* k_in, ggml_tensor* v_in, ggml_tensor* mask_in) -> ggml_tensor* { + const bool pad_head = d_head > 0 && d_head < 64 && q_in->ne[0] == d_head && k_in->ne[0] == d_head && + q_in->type == GGML_TYPE_F32 && k_in->type == GGML_TYPE_F32 && + v_in->type == GGML_TYPE_F32 && sd_backend_supports_cuda_mma(backend); + if (pad_head) { + // CUDA FA MMA starts at 64 channels; keep the original head's attention scale. + q_in = ggml_pad(ctx, q_in, 64 - d_head, 0, 0, 0); + k_in = ggml_pad(ctx, k_in, 64 - d_head, 0, 0, 0); + } if (kv_scale != 1.0f) { k_in = ggml_ext_scale(ctx, k_in, kv_scale); } @@ -650,6 +658,9 @@ ggml_tensor* ggml_ext_attention_ext(ggml_context* ctx, v_in = ggml_ext_cont(ctx, ggml_permute(ctx, v_in, 0, 2, 1, 3)); v_in = ggml_reshape_3d(ctx, v_in, d_head, L_k, n_kv_head * N); + if (pad_head) { + v_in = ggml_pad(ctx, v_in, 64 - d_head, 0, 0, 0); + } if (kv_scale != 1.0f) { v_in = ggml_ext_scale(ctx, v_in, kv_scale); } @@ -679,6 +690,9 @@ ggml_tensor* ggml_ext_attention_ext(ggml_context* ctx, if (kv_scale != 1.0f) { out = ggml_ext_scale(ctx, out, 1.0f / kv_scale); } + if (pad_head) { + out = ggml_ext_slice(ctx, out, 0, 0, d_head); + } return out; }; diff --git a/src/core/ggml_extend_backend.cpp b/src/core/ggml_extend_backend.cpp index b85fc5cb..e6fd496a 100644 --- a/src/core/ggml_extend_backend.cpp +++ b/src/core/ggml_extend_backend.cpp @@ -8,6 +8,10 @@ #include #include +#ifdef SD_USE_CUDA +#include +#endif + #include "core/util.h" #include "ggml/src/ggml-impl.h" #include "stable-diffusion.h" @@ -429,6 +433,70 @@ bool sd_backend_is_cpu(ggml_backend_t backend) { return dev != nullptr && ggml_backend_dev_type(dev) == GGML_BACKEND_DEVICE_TYPE_CPU; } +bool sd_backend_supports_cuda_mma(ggml_backend_t backend) { +#ifdef SD_USE_CUDA + if (!sd_backend_is(backend, "CUDA")) { + return false; + } + auto dev = ggml_backend_get_device(backend); + if (dev == nullptr) { + return false; + } + static std::mutex mutex; + static std::unordered_map cache; + std::lock_guard lock(mutex); + auto it = cache.find(dev); + if (it != cache.end()) { + return it->second; + } + const bool supported = [&]() { + ggml_backend_dev_props props{}; + ggml_backend_dev_get_props(dev, &props); + CUdevice device; + int major = 0, minor = 0; + if (props.device_id == nullptr || cuInit(0) != CUDA_SUCCESS || + cuDeviceGetByPCIBusId(&device, props.device_id) != CUDA_SUCCESS || + cuDeviceGetAttribute(&major, CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR, device) != CUDA_SUCCESS || + cuDeviceGetAttribute(&minor, CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MINOR, device) != CUDA_SUCCESS) { + return false; + } + auto reg = ggml_backend_dev_backend_reg(dev); + auto get_features = reinterpret_cast( + ggml_backend_reg_get_proc_address(reg, "ggml_backend_get_features")); + if (get_features == nullptr) { + return false; + } + // Match ggml's highest compiled architecture for this device, including PTX fallback. + const int cc = 100 * major + 10 * minor; + int compiled_arch = 0; + for (auto feature = get_features(reg); feature != nullptr && feature->name != nullptr; ++feature) { + if (std::strcmp(feature->name, "ARCHS") != 0 || feature->value == nullptr) { + continue; + } + const char* arch = feature->value; + while (*arch != '\0') { + char* end = nullptr; + const long value = std::strtol(arch, &end, 10); + if (end == arch) { + ++arch; + continue; + } + if (value <= cc && value > compiled_arch) { + compiled_arch = static_cast(value); + } + arch = end; + } + } + return compiled_arch == 700 || compiled_arch >= 750; + }(); + cache.emplace(dev, supported); + return supported; +#else + (void)backend; + return false; +#endif +} + ggml_backend_t sd_backend_cpu_init() { ggml_backend_load_all_once(); return ggml_backend_init_by_type(GGML_BACKEND_DEVICE_TYPE_CPU, nullptr); diff --git a/src/core/ggml_extend_backend.h b/src/core/ggml_extend_backend.h index a66bf51c..b2ebd7eb 100644 --- a/src/core/ggml_extend_backend.h +++ b/src/core/ggml_extend_backend.h @@ -87,6 +87,7 @@ private: bool sd_backend_is(ggml_backend_t backend, const std::string& name); bool sd_backend_is_cpu(ggml_backend_t backend); +bool sd_backend_supports_cuda_mma(ggml_backend_t backend); ggml_backend_t sd_backend_cpu_init(); bool sd_backend_cpu_set_n_threads(ggml_backend_t backend_cpu, int n_threads); ggml_status sd_backend_graph_compute_with_eval_callback(ggml_backend_t backend,