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

This commit is contained in:
leejet authored and GitHub committed 2026-10-13 22:58:45 +08:00
1 parent 60423ab5f6
commit d763dae3b8
53 files changed
+805 -125

No files matched your search

+35 -13
View File
@@ -632,14 +632,6 @@ ArgOptions SDContextParams::get_options() {
"--vae-on-cpu",
"deprecated; use --backend vae=cpu",
true, &vae_on_cpu},
{"",
"--fa",
"use flash attention",
true, &flash_attn},
{"",
"--diffusion-fa",
"use flash attention in the diffusion model only",
true, &diffusion_flash_attn},
{"",
"--sage-attn",
"use native CUDA SageAttention in the diffusion model, with flash/default attention fallback",
@@ -733,10 +725,42 @@ ArgOptions SDContextParams::get_options() {
}},
{"",
"--attn-scale",
"flash-attention K/V scale override (float, default: 0 = model default, 1 = no scaling); requires --fa or --diffusion-fa",
"flash-attention K/V scale override (float, default: 0 = model default, 1 = no scaling); applies when flash attention is selected",
[this](int argc, const char** argv, int index) {
return parse_scale_override(argc, argv, index, attn_scale);
}},
{"",
"--flash-attn",
"on|off|auto or module assignments (default: auto), e.g. auto,te=off,vae=on. Last repeated option wins",
[this](int argc, const char** argv, int index) {
if (++index >= argc) {
LOG_ERROR("--flash-attn requires on, off, or auto");
return -1;
}
const std::string value = argv[index];
if (value.empty()) {
LOG_ERROR("--flash-attn requires a nonempty value");
return -1;
}
flash_attn_mode = value;
return 1;
}},
{"",
"--fa",
"deprecated; use --flash-attn on",
[this](int, const char**, int) {
LOG_WARN("--fa is deprecated; use --flash-attn on instead.");
flash_attn_mode = "on";
return 0;
}},
{"",
"--diffusion-fa",
"deprecated; use --flash-attn off,diffusion=on",
[this](int, const char**, int) {
LOG_WARN("--diffusion-fa is deprecated; use --flash-attn off,diffusion=on instead.");
flash_attn_mode = "off,diffusion=on";
return 0;
}},
{"",
"--auto-fit",
"on|off (default: on). Preserve --backend (otherwise select one GPU) and place weights on the compute GPU, "
@@ -959,8 +983,7 @@ std::string SDContextParams::to_string() const {
<< " control_net_cpu: " << (control_net_cpu ? "true" : "false") << ",\n"
<< " clip_on_cpu: " << (clip_on_cpu ? "true" : "false") << ",\n"
<< " vae_on_cpu: " << (vae_on_cpu ? "true" : "false") << ",\n"
<< " flash_attn: " << (flash_attn ? "true" : "false") << ",\n"
<< " diffusion_flash_attn: " << (diffusion_flash_attn ? "true" : "false") << ",\n"
<< " flash_attn_mode: " << flash_attn_mode << ",\n"
<< " sage_attn: " << (sage_attn ? "true" : "false") << ",\n"
<< " linear_scale: " << linear_scale << ",\n"
<< " attn_scale: " << attn_scale << ",\n"
@@ -1018,8 +1041,7 @@ sd_ctx_params_t SDContextParams::to_sd_ctx_params_t(bool taesd_preview) {
sd_ctx_params.prediction = prediction;
sd_ctx_params.lora_apply_mode = lora_apply_mode;
sd_ctx_params.enable_mmap = enable_mmap;
sd_ctx_params.flash_attn = flash_attn;
sd_ctx_params.diffusion_flash_attn = diffusion_flash_attn;
sd_ctx_params.flash_attn_mode = flash_attn_mode.c_str();
sd_ctx_params.sage_attn = sage_attn;
sd_ctx_params.linear_scale = linear_scale;
sd_ctx_params.attn_scale = attn_scale;