From af12f01c3f2be8ca9d02bcb88929203a84ef8d1e Mon Sep 17 00:00:00 2001 From: leejet Date: Sun, 20 Sep 2026 21:00:57 +0800 Subject: [PATCH] feat: add native CUDA SageAttention support --- docs/build.md | 4 ++ docs/sage_attention.md | 68 +++++++++++++++++++++++++++++++ examples/common/common.cpp | 6 +++ examples/common/common.h | 1 + ggml | 2 +- include/stable-diffusion.h | 1 + src/core/ggml_extend.cpp | 35 +++++++++++++++- src/core/ggml_extend.h | 3 +- src/core/ggml_runner.cpp | 3 +- src/core/ggml_runner.h | 10 +++++ src/pipeline/diffusion_engine.cpp | 44 ++++++++++++++++++++ src/pipeline/diffusion_engine.h | 1 + src/stable-diffusion.cpp | 3 ++ 13 files changed, 176 insertions(+), 5 deletions(-) create mode 100644 docs/sage_attention.md diff --git a/docs/build.md b/docs/build.md index a50ec9bfb..53729d00a 100644 --- a/docs/build.md +++ b/docs/build.md @@ -88,6 +88,10 @@ cmake --build . --config Release ## Build with CUDA +Native SageAttention is included when using CUDA with patched GGML +(`SD_USE_UPSTREAM_GGML=OFF`). +See [SageAttention](sage_attention.md) for GPU requirements and `--sage-attn` usage. + This provides GPU acceleration using NVIDIA GPU. Make sure to have the CUDA toolkit installed. You can download it from your Linux distro's package manager (e.g. `apt install nvidia-cuda-toolkit`) or from here: [CUDA Toolkit](https://developer.nvidia.com/cuda-downloads). Recommended to have at least 4 GB of VRAM. ```shell diff --git a/docs/sage_attention.md b/docs/sage_attention.md new file mode 100644 index 000000000..c051f3b70 --- /dev/null +++ b/docs/sage_attention.md @@ -0,0 +1,68 @@ +# SageAttention + +`--sage-attn` enables native CUDA SageAttention in the diffusion model, including +the high-noise diffusion model when present. Python, PyTorch, and Triton are not +required at build time or runtime. + +The CUDA backend automatically selects a kernel supported by both the GPU and +the compiled CUDA toolkit: + +| GPU / toolkit | Implementation | +| --- | --- | +| SM89 or newer, CUDA 12.8 or newer (except SM90) | SageAttention2++: per-thread INT8 Q/K, FP8 PV, FP16 instruction accumulation with an FP32 buffer | +| SM89 or newer, CUDA 12.4 or newer; SM90 also uses this path with newer toolkits | SageAttention2: per-thread INT8 Q/K, FP8 PV, two-level FP32 accumulation | +| SM80 or newer, CUDA 12.0 or newer | INT8 Q/K, FP16 PV compatibility path | + +The FP8 paths smooth K, quantize V per channel, and pad and permute V for FP8 +Tensor Cores. The 2++ path uses the upstream V scale limit of 2.25 to avoid +overflow in its FP16 instruction accumulator. The public output remains FP32. +These are the upstream **INT8** SageAttention2/2++ variants; the paper's INT4 +variant and Hopper-specific WGMMA kernel are not implemented here. + +## Build + +Use the bundled patched GGML, CUDA Toolkit 12.0 or newer, and an NVIDIA GPU with +compute capability 8.0 or newer. Compile kernels for the GPU being used. + +```sh +cmake -S . -B build -DSD_CUDA=ON -DSD_USE_UPSTREAM_GGML=OFF +cmake --build build --config Release +``` + +No separate SageAttention build option is needed. Upstream GGML builds do not +support it. A system GGML must include the matching patched API and CUDA +backend. Enabling `--sage-attn` with an unavailable build or diffusion device +reports an error. Building with CUDA 12.4 selects SageAttention2 on an RTX 4090; +rebuild with CUDA 12.8 or newer to use SageAttention2++. + +## Use + +Replace `--diffusion-fa` with `--sage-attn` in an existing command. For example, +from the build directory: + +```powershell +.\bin\Release\sd-cli.exe -M vid_gen --diffusion-model ..\models\diffusion_models\Wan2.2-T2V-A14B-LowNoise-Q8_0.gguf --high-noise-diffusion-model ..\models\diffusion_models\Wan2.2-T2V-A14B-HighNoise-Q8_0.gguf --vae ..\models\vae\wan_2.1_vae.safetensors --t5xxl ..\models\text_encoders\umt5-xxl-encoder-Q8_0.gguf -p "a lovely cat" --cfg-scale 3.5 --sampling-method euler --steps 10 --high-noise-cfg-scale 3.5 --high-noise-sampling-method euler --high-noise-steps 8 -v -n "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的， +形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走" -W 832 -H 480 --diffusion-fa --offload-to-cpu --video-frames 33 --sage-attn +``` + +SageAttention currently handles unmasked attention with head dimensions from +1 through 128, including grouped-query attention, different query/key lengths, +and multiple batches. Dimensions below 64 are zero-padded to 64; dimensions +between 65 and 127 are zero-padded to 128. The original softmax scale is preserved, +and the output is cropped back to the original dimension. Other attention +operations fall back to FlashAttention when supported, then ordinary attention. +SageAttention takes precedence in diffusion +when combined with `--fa` or `--diffusion-fa`; `--fa` continues to control other +modules. Existing attention scaling overrides remain effective. + +Attention quantization changes numerical results. Compare image quality and +end-to-end generation time using the same seed, dimensions, and sampling +settings. Compare sampling steps after the first step for warmed-up inference +speed, and report model loading and first-step initialization separately. +Quantization, smoothing, and format conversion costs are included in generation +time, so short sequences may not benefit. + +Library callers set `sd_ctx_params_t.sage_attn = true` before `new_sd_ctx()`, +like `diffusion_flash_attn`. Context creation fails if the requested feature is +unavailable. Initialize the parameter structure with `sd_ctx_params_init()`. +Rebuild library callers against the updated public header. diff --git a/examples/common/common.cpp b/examples/common/common.cpp index 933c8622e..fad9003d1 100644 --- a/examples/common/common.cpp +++ b/examples/common/common.cpp @@ -618,6 +618,10 @@ ArgOptions SDContextParams::get_options() { "--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", + true, &sage_attn}, {"", "--diffusion-conv-direct", "use ggml_conv2d_direct in the diffusion model", @@ -938,6 +942,7 @@ std::string SDContextParams::to_string() const { << " 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" + << " sage_attn: " << (sage_attn ? "true" : "false") << ",\n" << " linear_scale: " << linear_scale << ",\n" << " attn_scale: " << attn_scale << ",\n" << " diffusion_conv_direct: " << (diffusion_conv_direct ? "true" : "false") << ",\n" @@ -995,6 +1000,7 @@ sd_ctx_params_t SDContextParams::to_sd_ctx_params_t(bool taesd_preview) { 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.sage_attn = sage_attn; sd_ctx_params.linear_scale = linear_scale; sd_ctx_params.attn_scale = attn_scale; sd_ctx_params.tae_preview_only = taesd_preview; diff --git a/examples/common/common.h b/examples/common/common.h index 3b66ae4ab..0e5c50e3c 100644 --- a/examples/common/common.h +++ b/examples/common/common.h @@ -170,6 +170,7 @@ struct SDContextParams { bool vae_on_cpu = false; bool flash_attn = false; bool diffusion_flash_attn = false; + bool sage_attn = false; bool diffusion_conv_direct = false; bool vae_conv_direct = false; diff --git a/ggml b/ggml index c6632cd90..f583f393c 160000 --- a/ggml +++ b/ggml @@ -1 +1 @@ -Subproject commit c6632cd905401abc58b6f5cdd52d228aa7ca1b88 +Subproject commit f583f393cd5dfdc129360bbf75cb3d49ccc837a8 diff --git a/include/stable-diffusion.h b/include/stable-diffusion.h index 9bbf8c757..041a9a36f 100644 --- a/include/stable-diffusion.h +++ b/include/stable-diffusion.h @@ -245,6 +245,7 @@ typedef struct { float linear_scale; // Override linear input scaling; 0 keeps the model default float attn_scale; // Override flash-attention K/V scaling; 0 keeps the model default const char* tokenizer; // tokenizer.json path or main=FILE,clip-l=FILE,clip-g=FILE assignments; required for PiD and Lens + bool sage_attn; } sd_ctx_params_t; typedef struct { diff --git a/src/core/ggml_extend.cpp b/src/core/ggml_extend.cpp index fd65f41c2..74324169e 100644 --- a/src/core/ggml_extend.cpp +++ b/src/core/ggml_extend.cpp @@ -622,7 +622,8 @@ ggml_tensor* ggml_ext_attention_ext(ggml_context* ctx, ggml_tensor* mask, bool skip_reshape, bool flash_attn, - float kv_scale) { // avoid overflow + float kv_scale, + bool sage_attn) { // avoid overflow int64_t L_q; int64_t L_k; int64_t C; @@ -713,7 +714,37 @@ ggml_tensor* ggml_ext_attention_ext(ggml_context* ctx, return out; }; - if (flash_attn) { +#ifndef SD_USE_UPSTREAM_GGML + if (sage_attn && mask == nullptr && d_head > 0 && d_head <= 128) { + auto q_in = ggml_reshape_4d(ctx, ggml_ext_cont(ctx, q->type == GGML_TYPE_F32 ? q : ggml_cast(ctx, q, GGML_TYPE_F32)), d_head, L_q, n_head, N); + auto k_in = ggml_reshape_4d(ctx, ggml_ext_cont(ctx, k->type == GGML_TYPE_F32 ? k : ggml_cast(ctx, k, GGML_TYPE_F32)), d_head, L_k, n_kv_head, N); + auto v_in = ggml_ext_cont(ctx, ggml_permute(ctx, v, 0, 2, 1, 3)); + const int64_t padded_head = d_head <= 64 ? 64 : 128; + if ((padded_head != d_head || kv_scale != 1.0f) && v_in->type != GGML_TYPE_F32) { + v_in = ggml_cast(ctx, v_in, GGML_TYPE_F32); + } + if (padded_head != d_head) { + // Keep the original head's softmax scale when padding for the CUDA kernel. + q_in = ggml_pad(ctx, q_in, padded_head - d_head, 0, 0, 0); + k_in = ggml_pad(ctx, k_in, padded_head - d_head, 0, 0, 0); + v_in = ggml_pad(ctx, v_in, padded_head - d_head, 0, 0, 0); + } + if (kv_scale != 1.0f) { + k_in = ggml_ext_scale(ctx, k_in, kv_scale); + v_in = ggml_ext_scale(ctx, v_in, kv_scale); + } + v_in = ggml_cast(ctx, v_in, GGML_TYPE_F16); + auto out = ggml_sage_attn(ctx, q_in, k_in, v_in, scale / kv_scale, GGML_SAGE_ATTN_AUTO); + if (ggml_backend_supports_op(backend, out)) { + kqv = kv_scale != 1.0f ? ggml_ext_scale(ctx, out, 1.0f / kv_scale) : out; + if (padded_head != d_head) { + kqv = ggml_ext_slice(ctx, kqv, 0, 0, d_head); + } + } + } +#endif + + if (kqv == nullptr && (flash_attn || sage_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) { diff --git a/src/core/ggml_extend.h b/src/core/ggml_extend.h index e843cbee0..33d8c3bf2 100644 --- a/src/core/ggml_extend.h +++ b/src/core/ggml_extend.h @@ -220,7 +220,8 @@ ggml_tensor* ggml_ext_attention_ext(ggml_context* ctx, ggml_tensor* mask = nullptr, bool skip_reshape = false, bool flash_attn = false, - float kv_scale = 1.0f); + float kv_scale = 1.0f, + bool sage_attn = false); ggml_tensor* ggml_ext_layer_norm(ggml_context* ctx, ggml_tensor* x, diff --git a/src/core/ggml_runner.cpp b/src/core/ggml_runner.cpp index 50173aab7..abfa8639b 100644 --- a/src/core/ggml_runner.cpp +++ b/src/core/ggml_runner.cpp @@ -25,7 +25,7 @@ ggml_tensor* ggml_ext_attention_ext(GGMLRunnerContext* ctx, if (ctx->attn_scale > 0.f) { kv_scale = ctx->attn_scale; } - return ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, n_head, mask, skip_reshape, flash_attn, kv_scale); + return ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, n_head, mask, skip_reshape, flash_attn, kv_scale, ctx->sage_attn_enabled); } void GGMLRunner::alloc_params_ctx() { @@ -520,6 +520,7 @@ GGMLRunnerContext GGMLRunner::get_context() { runner_ctx.ggml_ctx = compute_ctx; runner_ctx.backend = runtime_backend; runner_ctx.flash_attn_enabled = flash_attn_enabled; + runner_ctx.sage_attn_enabled = sage_attn_enabled; runner_ctx.linear_scale = linear_scale; runner_ctx.attn_scale = attn_scale; runner_ctx.conv2d_direct_enabled = conv2d_direct_enabled; diff --git a/src/core/ggml_runner.h b/src/core/ggml_runner.h index 2152ec8a2..29a6808bd 100644 --- a/src/core/ggml_runner.h +++ b/src/core/ggml_runner.h @@ -68,6 +68,7 @@ struct GGMLRunnerContext { ggml_backend_t backend = nullptr; ggml_context* ggml_ctx = nullptr; bool flash_attn_enabled = false; + bool sage_attn_enabled = false; float linear_scale = 0.f; float attn_scale = 0.f; bool conv2d_direct_enabled = false; @@ -176,6 +177,7 @@ struct GGMLRunner { const std::string final_result_name = "ggml_runner_final_result_tensor"; bool flash_attn_enabled = false; + bool sage_attn_enabled = false; float linear_scale = 0.f; float attn_scale = 0.f; bool conv2d_direct_enabled = false; @@ -337,6 +339,14 @@ struct GGMLRunner { flash_attn_enabled = enabled; } + void set_sage_attention_enabled(bool enabled) { + if (sage_attn_enabled != enabled) { + free_cache_ctx_and_buffer(); + graph_cut_plan_cache_.graph_cut_plans.clear(); + sage_attn_enabled = enabled; + } + } + void set_scale_overrides(float linear_scale, float attn_scale) { this->linear_scale = linear_scale; this->attn_scale = attn_scale; diff --git a/src/pipeline/diffusion_engine.cpp b/src/pipeline/diffusion_engine.cpp index e0663c655..0326f08cd 100644 --- a/src/pipeline/diffusion_engine.cpp +++ b/src/pipeline/diffusion_engine.cpp @@ -856,6 +856,47 @@ bool StableDiffusionGGML::init_model_loader(ModelLoader& model_loader, ModelConf return true; } +bool StableDiffusionGGML::set_sage_attention_enabled(bool enabled) { + if (!diffusion_model) { + return false; + } + if (enabled) { +#ifndef SD_USE_UPSTREAM_GGML + auto* ctx = ggml_init({4 * ggml_tensor_overhead(), nullptr, true}); + if (ctx == nullptr) { + return false; + } + auto* q = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, 128, 128, 1, 1); + auto* k = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, 128, 128, 1, 1); + auto* v = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, 128, 128, 1, 1); + auto* op = ggml_sage_attn(ctx, q, k, v, 1.f / sqrtf(128.f), GGML_SAGE_ATTN_AUTO); + bool supported = true; + for (auto backend : backend_manager.runtime_backends(SDBackendModule::DIFFUSION)) { + if (!ggml_backend_supports_op(backend, op)) { + LOG_ERROR("SageAttention is unavailable on %s; it requires patched GGML, CUDA Toolkit 12.0 or newer, and SM80 or newer kernels", + ggml_backend_name(backend)); + supported = false; + } + } + ggml_free(ctx); + if (!supported) { + return false; + } +#else + LOG_ERROR("SageAttention requires -DSD_USE_UPSTREAM_GGML=OFF and a CUDA backend"); + return false; +#endif + } + diffusion_model->set_sage_attention_enabled(enabled); + if (high_noise_diffusion_model) { + high_noise_diffusion_model->set_sage_attention_enabled(enabled); + } + if (enabled) { + LOG_INFO("Using SageAttention in the diffusion model; CUDA selects the supported kernel, unsupported layers use flash/default attention"); + } + return true; +} + bool StableDiffusionGGML::init(const sd_ctx_params_t* sd_ctx_params) { #ifdef SD_USE_UPSTREAM_GGML LOG_WARN( @@ -1133,6 +1174,9 @@ bool StableDiffusionGGML::validate_and_load_runners() { high_noise_diffusion_model->set_flash_attention_enabled(true); } } + if (sd_ctx_params->sage_attn && !set_sage_attention_enabled(true)) { + return false; + } LOG_VERBOSE("validating model metadata"); std::set ignore_tensors; diff --git a/src/pipeline/diffusion_engine.h b/src/pipeline/diffusion_engine.h index 682825fc3..26fe98a18 100644 --- a/src/pipeline/diffusion_engine.h +++ b/src/pipeline/diffusion_engine.h @@ -312,6 +312,7 @@ class StableDiffusionGGML { bool init_model_loader(ModelLoader& model_loader, ModelConfig& configuration); bool init(const sd_ctx_params_t* sd_ctx_params); + bool set_sage_attention_enabled(bool enabled); bool uses_tae() const; diff --git a/src/stable-diffusion.cpp b/src/stable-diffusion.cpp index 8f6032099..4597cdb89 100644 --- a/src/stable-diffusion.cpp +++ b/src/stable-diffusion.cpp @@ -336,6 +336,7 @@ void sd_ctx_params_init(sd_ctx_params_t* sd_ctx_params) { sd_ctx_params->eager_load = false; sd_ctx_params->enable_mmap = false; sd_ctx_params->diffusion_flash_attn = false; + sd_ctx_params->sage_attn = false; sd_ctx_params->linear_scale = 0.f; sd_ctx_params->attn_scale = 0.f; sd_ctx_params->vae_format = SD_VAE_FORMAT_AUTO; @@ -391,6 +392,7 @@ char* sd_ctx_params_to_str(const sd_ctx_params_t* sd_ctx_params) { "auto_fit: %s\n" "flash_attn: %s\n" "diffusion_flash_attn: %s\n" + "sage_attn: %s\n" "linear_scale: %g\n" "attn_scale: %g\n" "vae_format: %s\n", @@ -430,6 +432,7 @@ char* sd_ctx_params_to_str(const sd_ctx_params_t* sd_ctx_params) { BOOL_STR(sd_ctx_params->auto_fit), BOOL_STR(sd_ctx_params->flash_attn), BOOL_STR(sd_ctx_params->diffusion_flash_attn), + BOOL_STR(sd_ctx_params->sage_attn), sd_ctx_params->linear_scale, sd_ctx_params->attn_scale, sd_vae_format_name(sd_ctx_params->vae_format));