feat: add per-module automatic flash attention selection (#2132)

src/core/ggml_extend.cpp

@@ -761,7 +761,7 @@ ggml_tensor* ggml_ext_attention_ext(ggml_context* ctx,
     }
 #endif
 
-    if (kqv == nullptr && (flash_attn || sage_attn)) {
+    if (kqv == nullptr && flash_attn) {
         // LOG_VERBOSE("attention_ext L_q:%d L_k:%d n_head:%d C:%d d_head:%d N:%d", L_q, L_k, n_head, C, d_head, N);
         bool can_use_flash_attn = true;
         if (mask != nullptr) {

src/core/ggml_extend_backend.cpp

@@ -6,12 +6,17 @@
 #include <mutex>
 #include <sstream>
 #include <stdexcept>
+#include <utility>
 #include <vector>
 
 #ifdef SD_USE_CUDA
 #include <cuda.h>
 #endif
 
+#ifdef SD_USE_HIP
+#include <hip/hip_runtime_api.h>
+#endif
+
 #include "core/util.h"
 #include "ggml-backend-impl.h"
 #include "ggml-impl.h"
@@ -516,6 +521,49 @@ bool sd_backend_supports_cuda_mma(ggml_backend_t backend) {
 #endif
 }
 
+bool sd_backend_supports_hip_mma(ggml_backend_t backend) {
+#ifdef SD_USE_HIP
+    if (!sd_backend_is(backend, "ROCm")) {
+        return false;
+    }
+    auto dev = ggml_backend_get_device(backend);
+    if (dev == nullptr) {
+        return false;
+    }
+    ggml_backend_dev_props props{};
+    ggml_backend_dev_get_props(dev, &props);
+    int device = 0;
+    hipDeviceProp_t hip_props{};
+    if (props.device_id == nullptr ||
+        hipDeviceGetByPCIBusId(&device, props.device_id) != hipSuccess ||
+        hipGetDeviceProperties(&hip_props, device) != hipSuccess) {
+        return false;
+    }
+    const char* arch = hip_props.gcnArchName;
+    if (std::strncmp(arch, "gfx", 3) != 0) {
+        return false;
+    }
+    char* end        = nullptr;
+    const long value = std::strtol(arch + 3, &end, 16);
+    if (end == arch + 3 || (*end != '\0' && *end != ':')) {
+        return false;
+    }
+    // Match the HIP FA kernels; keep unknown architecture families on the probe path.
+    if (value >= 0x1100 && value < 0x1300) {
+        return true;
+    }
+#ifdef SD_HIP_MFMA
+    return value == 0x908 || value == 0x90a ||
+           value == 0x940 || value == 0x941 || value == 0x942 || value == 0x950;
+#else
+    return false;
+#endif
+#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);
@@ -734,6 +782,59 @@ bool sd_parse_backend_assignment(const std::string& spec, SDBackendAssignment* a
     return true;
 }
 
+bool sd_parse_flash_attention_assignment(const std::string& spec, SDBackendAssignment* assignment, std::string* error) {
+    if (assignment == nullptr) {
+        return false;
+    }
+    SDBackendAssignment result;
+    result.set_default("auto");
+    std::istringstream stream(spec);
+    std::string part;
+    while (std::getline(stream, part, ',')) {
+        SDBackendAssignment entry;
+        if (!sd_parse_backend_assignment(part, &entry, error)) {
+            return false;
+        }
+        auto valid_mode = [&](const std::string& mode) {
+            if (mode == "on" || mode == "off" || mode == "auto") {
+                return true;
+            }
+            if (error != nullptr) {
+                *error = "invalid flash attention mode '" + mode + "': expected on, off, or auto";
+            }
+            return false;
+        };
+        if (!entry.default_name.empty()) {
+            if (!valid_mode(entry.default_name)) {
+                return false;
+            }
+            result.set_default(entry.default_name);
+        }
+        for (const auto& setting : entry.module_names) {
+            switch (setting.first) {
+                case SDBackendModule::DIFFUSION:
+                case SDBackendModule::TE:
+                case SDBackendModule::CLIP_VISION:
+                case SDBackendModule::VAE:
+                case SDBackendModule::CONTROL_NET:
+                case SDBackendModule::AUDIO_ENCODER:
+                    break;
+                default:
+                    if (error != nullptr) {
+                        *error = std::string("flash attention is not configurable for module '") + sd_backend_module_name(setting.first) + "'";
+                    }
+                    return false;
+            }
+            if (!valid_mode(setting.second)) {
+                return false;
+            }
+            result.set_module(setting.first, setting.second);
+        }
+    }
+    *assignment = std::move(result);
+    return true;
+}
+
 bool SDBackendAssignment::empty() const {
     return default_name.empty() && module_names.empty();
 }

src/core/ggml_extend_backend.h

@@ -88,6 +88,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);
+bool sd_backend_supports_hip_mma(ggml_backend_t backend);
 ggml_backend_buffer_t sd_backend_dev_buffer_from_host_ptr(ggml_backend_dev_t device,
                                                           void* ptr,
                                                           size_t size,
@@ -100,6 +101,7 @@ ggml_status sd_backend_graph_compute_with_eval_callback(ggml_backend_t backend,
                                                         void* callback_eval_user_data);
 std::string sd_backend_resolve_name(const std::string& name);
 bool sd_parse_backend_assignment(const std::string& spec, SDBackendAssignment* assignment, std::string* error);
+bool sd_parse_flash_attention_assignment(const std::string& spec, SDBackendAssignment* assignment, std::string* error);
 const char* sd_backend_module_name(SDBackendModule module);
 void ggml_ext_im_set_f32_1d(const struct ggml_tensor* tensor, int i, float value);
 bool add_rpc_devices(const std::string& servers);

src/core/probe/backend_probe.cpp

@@ -0,0 +1,193 @@
+#include "core/probe/backend_probe.h"
+
+#include <algorithm>
+#include <chrono>
+#include <cmath>
+#include <exception>
+
+#include "core/ggml_extend_backend.h"
+#include "core/util.h"
+#include "ggml-alloc.h"
+
+BackendProbeGraph::BackendProbeGraph(ggml_backend_t backend, size_t graph_size)
+    : backend_(backend), graph_size_(graph_size) {
+    ctx = ggml_init({graph_size * ggml_tensor_overhead() + ggml_graph_overhead_custom(graph_size, false), nullptr, true});
+}
+
+BackendProbeGraph::~BackendProbeGraph() {
+    ggml_backend_buffer_free(buffer_);
+    ggml_free(ctx);
+}
+
+bool BackendProbeGraph::supported() {
+    if (backend_ == nullptr || ctx == nullptr || output == nullptr) {
+        return false;
+    }
+    if (graph_ == nullptr) {
+        graph_ = ggml_new_graph_custom(ctx, graph_size_, false);
+        ggml_build_forward_expand(graph_, output);
+    }
+    for (int i = 0; i < ggml_graph_n_nodes(graph_); ++i) {
+        if (!ggml_backend_supports_op(backend_, ggml_graph_node(graph_, i))) {
+            return false;
+        }
+    }
+    return true;
+}
+
+bool BackendProbeGraph::allocate() {
+    if (ctx == nullptr || backend_ == nullptr) {
+        return false;
+    }
+    if (buffer_ == nullptr) {
+        buffer_ = ggml_backend_alloc_ctx_tensors(ctx, backend_);
+    }
+    return buffer_ != nullptr;
+}
+
+bool BackendProbeGraph::initialize_f32_inputs() {
+    uint32_t state = 42;
+    for (auto* tensor : inputs) {
+        if (tensor == nullptr || tensor->type != GGML_TYPE_F32 || !ggml_is_contiguous(tensor)) {
+            return false;
+        }
+        std::vector<float> values(static_cast<size_t>(ggml_nelements(tensor)));
+        for (float& value : values) {
+            state = state * 1664525u + 1013904223u;
+            value = static_cast<float>(state >> 8) / 16777216.f - 0.5f;
+        }
+        ggml_backend_tensor_set(tensor, values.data(), 0, ggml_nbytes(tensor));
+    }
+    return true;
+}
+
+double BackendProbeGraph::run() {
+    if (graph_ == nullptr || buffer_ == nullptr) {
+        return -1.;
+    }
+    ggml_backend_synchronize(backend_);
+    const auto start  = std::chrono::steady_clock::now();
+    const auto status = ggml_backend_graph_compute(backend_, graph_);
+    ggml_backend_synchronize(backend_);
+    if (status != GGML_STATUS_SUCCESS) {
+        return -1.;
+    }
+    return std::chrono::duration<double, std::milli>(std::chrono::steady_clock::now() - start).count();
+}
+
+static bool backend_probe_outputs_match(const BackendProbeGraph& reference,
+                                        const BackendProbeGraph& candidate,
+                                        const BackendProbeOptions& options) {
+    if (!ggml_are_same_shape(reference.output, candidate.output) ||
+        reference.output->type != GGML_TYPE_F32 || candidate.output->type != GGML_TYPE_F32 ||
+        !ggml_is_contiguous(reference.output) || !ggml_is_contiguous(candidate.output)) {
+        return false;
+    }
+    const size_t count = static_cast<size_t>(ggml_nelements(reference.output));
+    std::vector<float> expected(count), actual(count);
+    ggml_backend_tensor_get(reference.output, expected.data(), 0, count * sizeof(float));
+    ggml_backend_tensor_get(candidate.output, actual.data(), 0, count * sizeof(float));
+    for (size_t i = 0; i < count; ++i) {
+        if (!std::isfinite(expected[i]) || !std::isfinite(actual[i]) ||
+            std::abs(expected[i] - actual[i]) > options.absolute_tolerance + options.relative_tolerance * std::abs(expected[i])) {
+            return false;
+        }
+    }
+    return true;
+}
+
+BackendProbeResult sd_run_backend_probe(ggml_backend_t backend,
+                                        int n_threads,
+                                        const std::vector<BackendProbeCase>& cases,
+                                        const BackendProbeOptions& options) {
+    BackendProbeResult result;
+    if (backend == nullptr || cases.empty() || options.graph_size == 0 || options.warmup_runs < 1 || options.measurement_runs < 1 ||
+        !std::isfinite(options.time_budget_ms) || options.time_budget_ms < 0. ||
+        !std::isfinite(options.min_improvement) || options.min_improvement < 0. || options.min_improvement >= 1. ||
+        !std::isfinite(options.absolute_tolerance) || options.absolute_tolerance < 0.f ||
+        !std::isfinite(options.relative_tolerance) || options.relative_tolerance < 0.f) {
+        result.error = "invalid probe configuration";
+        return result;
+    }
+    try {
+        if (sd_backend_is_cpu(backend)) {
+            sd_backend_cpu_set_n_threads(backend, n_threads > 0 ? n_threads : sd_get_num_physical_cores());
+        }
+        auto dev = ggml_backend_get_device(backend);
+        if (dev != nullptr && !sd_backend_is_cpu(backend)) {
+            size_t free_bytes = 0, total_bytes = 0;
+            ggml_backend_dev_memory(dev, &free_bytes, &total_bytes);
+            if (free_bytes > total_bytes || free_bytes < options.min_free_device_bytes) {
+                result.error = "insufficient probe memory";
+                return result;
+            }
+        }
+        double timed_ms = 0.;
+        for (const auto& test : cases) {
+            BackendProbeGraph reference(backend, options.graph_size), candidate(backend, options.graph_size);
+            if (!test.build || reference.ctx == nullptr || candidate.ctx == nullptr) {
+                result.error = "probe graph unavailable";
+                return result;
+            }
+            test.build(reference, false);
+            test.build(candidate, true);
+            if (!reference.supported() || !candidate.supported()) {
+                result.error = "unsupported probe graph";
+                return result;
+            }
+            if (!reference.allocate() || !candidate.allocate()) {
+                result.error = "probe allocation failed";
+                return result;
+            }
+            auto initialize = [&](BackendProbeGraph& graph) {
+                return test.initialize ? test.initialize(graph) : graph.initialize_f32_inputs();
+            };
+            if (!initialize(reference) || !initialize(candidate)) {
+                result.error = "probe input initialization failed";
+                return result;
+            }
+            for (int i = 0; i < options.warmup_runs; ++i) {
+                if (reference.run() < 0. || candidate.run() < 0.) {
+                    result.error = "probe warmup failed";
+                    return result;
+                }
+            }
+            if (!(test.compare ? test.compare(reference, candidate) : backend_probe_outputs_match(reference, candidate, options))) {
+                result.error = "probe output check failed";
+                return result;
+            }
+            std::vector<double> reference_times(options.measurement_runs), candidate_times(options.measurement_runs);
+            for (int i = 0; i < options.measurement_runs; ++i) {
+                if (timed_ms >= options.time_budget_ms) {
+                    result.error = "probe time budget exhausted";
+                    return result;
+                }
+                // Alternate order so clock ramp-up does not consistently favor one path.
+                if (i % 2 == 0) {
+                    reference_times[i] = reference.run();
+                    candidate_times[i] = candidate.run();
+                } else {
+                    candidate_times[i] = candidate.run();
+                    reference_times[i] = reference.run();
+                }
+                if (reference_times[i] < 0. || candidate_times[i] < 0.) {
+                    result.error = "probe execution failed";
+                    return result;
+                }
+                timed_ms += reference_times[i] + candidate_times[i];
+            }
+            auto median = [](std::vector<double>& values) {
+                std::sort(values.begin(), values.end());
+                const size_t mid = values.size() / 2;
+                return values.size() % 2 ? values[mid] : (values[mid - 1] + values[mid]) * 0.5;
+            };
+            result.reference_ms += median(reference_times);
+            result.candidate_ms += median(candidate_times);
+        }
+        result.completed        = true;
+        result.candidate_faster = result.candidate_ms < result.reference_ms * (1. - options.min_improvement);
+    } catch (const std::exception& error) {
+        result.error = error.what();
+    }
+    return result;
+}

src/core/probe/backend_probe.h

@@ -0,0 +1,64 @@
+#ifndef __SD_CORE_PROBE_BACKEND_PROBE_H__
+#define __SD_CORE_PROBE_BACKEND_PROBE_H__
+
+#include <functional>
+#include <string>
+#include <vector>
+
+#include "ggml-backend.h"
+
+struct BackendProbeOptions {
+    size_t graph_size            = 128;
+    int warmup_runs              = 1;
+    int measurement_runs         = 5;
+    double time_budget_ms        = 1000.;
+    size_t min_free_device_bytes = 128ULL * 1024 * 1024;
+    double min_improvement       = 0.05;
+    float absolute_tolerance     = 0.001f;
+    float relative_tolerance     = 0.02f;
+};
+
+class BackendProbeGraph {
+    ggml_backend_t backend_;
+    size_t graph_size_;
+    ggml_backend_buffer_t buffer_ = nullptr;
+    ggml_cgraph* graph_           = nullptr;
+
+public:
+    ggml_context* ctx = nullptr;
+    std::vector<ggml_tensor*> inputs;
+    ggml_tensor* output = nullptr;
+
+    explicit BackendProbeGraph(ggml_backend_t backend, size_t graph_size = 128);
+    ~BackendProbeGraph();
+    BackendProbeGraph(const BackendProbeGraph&)            = delete;
+    BackendProbeGraph& operator=(const BackendProbeGraph&) = delete;
+
+    bool supported();
+    bool allocate();
+    bool initialize_f32_inputs();
+    double run();
+};
+
+struct BackendProbeCase {
+    // Variants must be repeatable without mutating inputs; leave output null if unavailable.
+    std::function<void(BackendProbeGraph&, bool candidate)> build;
+    // Optional overrides for non-F32 inputs/outputs or operation-specific data.
+    std::function<bool(BackendProbeGraph&)> initialize;
+    std::function<bool(const BackendProbeGraph&, const BackendProbeGraph&)> compare;
+};
+
+struct BackendProbeResult {
+    bool completed        = false;
+    bool candidate_faster = false;
+    double reference_ms   = 0.;
+    double candidate_ms   = 0.;
+    std::string error;
+};
+
+BackendProbeResult sd_run_backend_probe(ggml_backend_t backend,
+                                        int n_threads,
+                                        const std::vector<BackendProbeCase>& cases,
+                                        const BackendProbeOptions& options = {});
+
+#endif  // __SD_CORE_PROBE_BACKEND_PROBE_H__

src/core/probe/flash_attn.cpp

@@ -0,0 +1,101 @@
+#include "core/probe/flash_attn.h"
+
+#include <array>
+#include <cstring>
+#include <exception>
+#include <sstream>
+#include <utility>
+#include <vector>
+
+#include "core/ggml_extend.h"
+#include "core/ggml_extend_backend.h"
+#include "core/probe/backend_probe.h"
+#include "core/util.h"
+
+static bool is_vulkan_nvidia_rtx(ggml_backend_t backend) {
+    auto device = ggml_backend_get_device(backend);
+    if (device == nullptr) {
+        return false;
+    }
+    auto reg = ggml_backend_dev_backend_reg(device);
+    if (reg == nullptr || std::strcmp(ggml_backend_reg_name(reg), "Vulkan") != 0) {
+        return false;
+    }
+    const char* description = ggml_backend_dev_description(device);
+    if (description == nullptr) {
+        return false;
+    }
+    std::istringstream stream(description);
+    std::string token;
+    bool nvidia = false, rtx = false;
+    while (stream >> token) {
+        nvidia = nvidia || token == "NVIDIA";
+        rtx    = rtx || token == "RTX";
+    }
+    return nvidia && rtx;
+}
+
+static void build_flash_attention_probe(BackendProbeGraph& graph, ggml_backend_t backend, int length, int head_size, bool flash) {
+    constexpr int heads = 4;
+    auto q              = ggml_new_tensor_3d(graph.ctx, GGML_TYPE_F32, head_size * heads, length, 1);
+    auto k              = ggml_dup_tensor(graph.ctx, q);
+    auto v              = ggml_dup_tensor(graph.ctx, q);
+    graph.inputs        = {q, k, v};
+    bool used_flash     = false;
+    auto output         = ggml_ext_attention_ext(graph.ctx, backend, q, k, v, heads, nullptr, false, flash, 1.f, false, &used_flash);
+    graph.output        = flash && !used_flash ? nullptr : output;
+}
+
+bool sd_probe_flash_attention(ggml_backend_t backend, int n_threads) {
+    if (backend == nullptr) {
+        return false;
+    }
+    const char* name = ggml_backend_name(backend);
+    try {
+        BackendProbeGraph capability(backend);
+        if (capability.ctx == nullptr) {
+            LOG_INFO("flash attention auto: %s -> off (probe graph unavailable)", name);
+            return false;
+        }
+        build_flash_attention_probe(capability, backend, 512, 64, true);
+        if (!capability.supported()) {
+            LOG_INFO("flash attention auto: %s -> off (unsupported probe graph)", name);
+            return false;
+        }
+        if (sd_backend_is_cpu(backend)) {
+            LOG_INFO("flash attention auto: %s -> on (CPU)", name);
+            return true;
+        }
+        if (sd_backend_supports_cuda_mma(backend)) {
+            LOG_INFO("flash attention auto: %s -> on (CUDA MMA)", name);
+            return true;
+        }
+        if (sd_backend_supports_hip_mma(backend)) {
+            LOG_INFO("flash attention auto: %s -> on (HIP MMA)", name);
+            return true;
+        }
+        if (is_vulkan_nvidia_rtx(backend)) {
+            LOG_INFO("flash attention auto: %s -> on (Vulkan NVIDIA RTX)", name);
+            return true;
+        }
+        std::vector<BackendProbeCase> cases;
+        for (const auto& shape : {std::array<int, 2>{512, 64}, std::array<int, 2>{1024, 128}}) {
+            BackendProbeCase test;
+            test.build = [backend, shape](BackendProbeGraph& graph, bool candidate) {
+                build_flash_attention_probe(graph, backend, shape[0], shape[1], candidate);
+            };
+            cases.push_back(std::move(test));
+        }
+        const auto result = sd_run_backend_probe(backend, n_threads, cases);
+        if (!result.completed) {
+            LOG_WARN("flash attention auto: %s -> off (%s)", name, result.error.c_str());
+            return false;
+        }
+        LOG_INFO("flash attention auto: %s -> %s (probe: normal %.3f ms, flash %.3f ms)",
+                 name, result.candidate_faster ? "on" : "off", result.reference_ms, result.candidate_ms);
+        return result.candidate_faster;
+    } catch (const std::exception& error) {
+        LOG_WARN("flash attention auto: %s -> off (probe failed: %s)", name, error.what());
+        return false;
+    }
+}

src/core/probe/flash_attn.h

@@ -0,0 +1,8 @@
+#ifndef __SD_CORE_PROBE_FLASH_ATTN_H__
+#define __SD_CORE_PROBE_FLASH_ATTN_H__
+
+#include "ggml-backend.h"
+
+bool sd_probe_flash_attention(ggml_backend_t backend, int n_threads);
+
+#endif  // __SD_CORE_PROBE_FLASH_ATTN_H__

src/pipeline/diffusion_engine.cpp

@@ -18,6 +18,7 @@
 #include "core/ggml_runner.h"
 #include "core/ggml_tensor_utils.h"
 #include "core/layer_split_partition.h"
+#include "core/probe/flash_attn.h"
 #include "model.h"
 
 #include "core/rng.hpp"
@@ -944,6 +945,13 @@ bool StableDiffusionGGML::init(const sd_ctx_params_t* sd_ctx_params) {
         LOG_ERROR("conditioning_cache_size must be non-negative");
         return false;
     }
+    const std::string flash_mode = SAFE_STR(sd_ctx_params->flash_attn_mode);
+    std::string flash_error;
+    if (!sd_parse_flash_attention_assignment(flash_mode, &flash_attn_assignment, &flash_error)) {
+        LOG_ERROR("invalid flash_attn_mode '%s': %s", flash_mode.c_str(), flash_error.c_str());
+        return false;
+    }
+    flash_attn_by_backend.clear();
     conditioning_cache_->set_capacity(static_cast<size_t>(sd_ctx_params->conditioning_cache_size));
     auto configuration = std::make_unique<ModelConfig>(*sd_ctx_params);
     n_threads          = sd_ctx_params->n_threads;
@@ -1182,30 +1190,73 @@ bool StableDiffusionGGML::build_extension_runners() {
     return true;
 }
 
+bool StableDiffusionGGML::flash_attention_for(SDBackendModule module) {
+    const auto& params     = config_->params;
+    const std::string mode = flash_attn_assignment.get(module);
+    if (mode == "on") {
+        return true;
+    }
+    if (mode == "off") {
+        return false;
+    }
+    if (strlen(SAFE_STR(params.flash_attn_mode)) == 0) {
+        if (params.flash_attn) {
+            return true;
+        }
+        if (params.diffusion_flash_attn) {
+            return module == SDBackendModule::DIFFUSION;
+        }
+    }
+    auto backends    = backend_manager.runtime_backends(module);
+    bool enabled     = !backends.empty();
+    bool any_enabled = false;
+    for (auto backend : backends) {
+        auto it = flash_attn_by_backend.find(backend);
+        if (it == flash_attn_by_backend.end()) {
+            it = flash_attn_by_backend.emplace(backend, sd_probe_flash_attention(backend, n_threads)).first;
+        }
+        enabled     = enabled && it->second;
+        any_enabled = any_enabled || it->second;
+    }
+    if (any_enabled && !enabled) {
+        LOG_INFO("flash attention auto: %s -> off (compute devices disagree)", sd_backend_module_name(module));
+    }
+    return enabled;
+}
+
 bool StableDiffusionGGML::validate_and_load_runners() {
     const auto* sd_ctx_params   = &config_->params;
     const bool use_tae          = uses_tae();
     const bool tae_preview_only = this->tae_preview_only();
-    if (sd_ctx_params->flash_attn) {
-        LOG_INFO("Using flash attention");
-        cond_stage_model->set_flash_attention_enabled(true);
-        if (clip_vision) {
-            clip_vision->set_flash_attention_enabled(true);
-        }
+    cond_stage_model->set_flash_attention_enabled(flash_attention_for(SDBackendModule::TE));
+    if (clip_vision) {
+        clip_vision->set_flash_attention_enabled(flash_attention_for(SDBackendModule::CLIP_VISION));
+    }
+    if (first_stage_model || preview_vae || audio_vae_model) {
+        const bool enabled = flash_attention_for(SDBackendModule::VAE);
         if (first_stage_model) {
-            first_stage_model->set_flash_attention_enabled(true);
+            first_stage_model->set_flash_attention_enabled(enabled);
         }
         if (preview_vae) {
-            preview_vae->set_flash_attention_enabled(true);
+            preview_vae->set_flash_attention_enabled(enabled);
+        }
+        if (audio_vae_model) {
+            audio_vae_model->set_flash_attention_enabled(enabled);
         }
     }
-
-    if (sd_ctx_params->flash_attn || sd_ctx_params->diffusion_flash_attn) {
-        LOG_INFO("Using flash attention in the diffusion model");
-        diffusion_model->set_flash_attention_enabled(true);
-        if (high_noise_diffusion_model) {
-            high_noise_diffusion_model->set_flash_attention_enabled(true);
-        }
+    const bool diffusion_fa = flash_attention_for(SDBackendModule::DIFFUSION);
+    diffusion_model->set_flash_attention_enabled(diffusion_fa);
+    if (high_noise_diffusion_model) {
+        high_noise_diffusion_model->set_flash_attention_enabled(diffusion_fa);
+    }
+    if (ip_adapter) {
+        ip_adapter->set_flash_attention_enabled(diffusion_fa);
+    }
+    if (control_net) {
+        control_net->set_flash_attention_enabled(flash_attention_for(SDBackendModule::CONTROL_NET));
+    }
+    if (audio_encoder) {
+        audio_encoder->set_flash_attention_enabled(flash_attention_for(SDBackendModule::AUDIO_ENCODER));
     }
     if (sd_ctx_params->sage_attn && !set_sage_attention_enabled(true)) {
         return false;

src/pipeline/diffusion_engine.h

@@ -11,6 +11,7 @@
 #include <optional>
 #include <set>
 #include <string>
+#include <unordered_map>
 #include <vector>
 
 #include "core/ggml_extend_backend.h"
@@ -48,6 +49,8 @@ static inline bool sd_version_supports_ref_latent_img_cfg(SDVersion version) {
 class StableDiffusionGGML {
 public:
     SDBackendManager backend_manager;
+    std::unordered_map<ggml_backend_t, bool> flash_attn_by_backend;
+    SDBackendAssignment flash_attn_assignment;
 
     SDVersion version;
     bool external_vae_is_invalid = false;
@@ -131,7 +134,7 @@ public:
                                 &sd_ctx_params_t::clip_g_path, &sd_ctx_params_t::clip_vision_path,
                                 &sd_ctx_params_t::t5xxl_path, &sd_ctx_params_t::llm_path,
                                 &sd_ctx_params_t::llm_vision_path, &sd_ctx_params_t::diffusion_model_path,
-                                &sd_ctx_params_t::tokenizer,
+                                &sd_ctx_params_t::tokenizer, &sd_ctx_params_t::flash_attn_mode,
                                 &sd_ctx_params_t::high_noise_diffusion_model_path, &sd_ctx_params_t::uncond_diffusion_model_path,
                                 &sd_ctx_params_t::embeddings_connectors_path, &sd_ctx_params_t::vae_path,
                                 &sd_ctx_params_t::audio_vae_path, &sd_ctx_params_t::taesd_path,
@@ -334,6 +337,7 @@ public:
     bool build_extension_runners();
 
     bool validate_and_load_runners();
+    bool flash_attention_for(SDBackendModule module);
 
     bool build_denoiser();
 

src/stable-diffusion.cpp

@@ -338,6 +338,7 @@ void sd_ctx_params_init(sd_ctx_params_t* sd_ctx_params) {
     sd_ctx_params->batched_cfg               = true;
     sd_ctx_params->eager_load                = false;
     sd_ctx_params->enable_mmap               = false;
+    sd_ctx_params->flash_attn_mode           = nullptr;
     sd_ctx_params->diffusion_flash_attn      = false;
     sd_ctx_params->sage_attn                 = false;
     sd_ctx_params->linear_scale              = 0.f;
@@ -395,6 +396,7 @@ char* sd_ctx_params_to_str(const sd_ctx_params_t* sd_ctx_params) {
              "model_args: %s\n"
              "auto_fit: %s\n"
              "flash_attn: %s\n"
+             "flash_attn_mode: %s\n"
              "diffusion_flash_attn: %s\n"
              "sage_attn: %s\n"
              "linear_scale: %g\n"
@@ -436,6 +438,7 @@ char* sd_ctx_params_to_str(const sd_ctx_params_t* sd_ctx_params) {
              SAFE_STR(sd_ctx_params->model_args),
              BOOL_STR(sd_ctx_params->auto_fit),
              BOOL_STR(sd_ctx_params->flash_attn),
+             sd_ctx_params->flash_attn_mode && sd_ctx_params->flash_attn_mode[0] ? sd_ctx_params->flash_attn_mode : "auto (legacy flags honored)",
              BOOL_STR(sd_ctx_params->diffusion_flash_attn),
              BOOL_STR(sd_ctx_params->sage_attn),
              sd_ctx_params->linear_scale,