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

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();