diff --git a/docs/sol_attention.md b/docs/sol_attention.md new file mode 100644 index 000000000..8604a3e65 --- /dev/null +++ b/docs/sol_attention.md @@ -0,0 +1,106 @@ +# Sol-Attn + +`--sol-attn` enables native CUDA Sol-Attn in the diffusion model, including the +high-noise diffusion model when present. It uses the shared attention dispatcher +without classifying tokens as text, images, or video. Python, PyTorch, Triton, +and CuTe DSL are not needed to build or run it. + +This implementation follows the diagonal-threshold algorithm in +[NVlabs/Sana's Sol-Attn](https://github.com/NVlabs/Sana/tree/sol-engine/techniques/sparse_backends/sol_attn). +It summarizes 64-token KV blocks, selects exact blocks using proxy scores and +an online threshold, and approximates the remaining blocks using their K means +and V sums. Adjacent blocks remain exact. Both contributions share an online +softmax normalizer. Q/K/V and probability tiles use BF16 Tensor Cores with FP32 +accumulation; the BF16 result is returned through the existing FP32 interface. + +## Build + +Use patched GGML, CUDA Toolkit 12.0 or newer, and an NVIDIA GPU with compute +capability 8.0 or newer. Compile kernels for the target GPU: + +```sh +cmake -S . -B build -DSD_CUDA=ON -DSD_USE_UPSTREAM_GGML=OFF +cmake --build build --config Release +``` + +The feature is compiled with the CUDA backend; no separate build option is +required. Upstream GGML and non-CUDA backends do not support it. A system GGML +must provide the matching patched API and CUDA implementation. Tensor-parallel +row splitting is not supported; layer splitting requires supported devices. + +## Use + +Add `--sol-attn` to an existing generation command: + +```sh +sd-cli ... --sol-attn +sd-cli ... --sol-attn --sol-attn-tau 1.0 +``` + +The default threshold coefficient is `1.0`. Larger coefficients select fewer +blocks for exact attention. The coefficient must be finite; zero does not mean +dense attention. Omit `--sol-attn` to disable the feature. + +The native kernel supports unmasked, noncausal attention with head dimension +128, equal Q/K/V sequence lengths and head counts, and multiple batches. Other +attention operations fall back to FlashAttention when available, then ordinary +attention. Existing attention scaling overrides remain effective. `--fa` and +`--diffusion-fa` may be used together with Sol-Attn; `--sage-attn` is mutually +exclusive. Text encoders and VAEs retain their existing attention selection. + +Initialization reports an error if the requested diffusion backend cannot run +Sol-Attn. Graph logs report the number of Sol-Attn and FlashAttention nodes and +warn when no Sol-Attn nodes are selected. CUDA execution errors are not silently +converted into dense attention. + +This is approximate attention. Validate quality and end-to-end speed with the +same prompt, seed, dimensions, frame count, and sampling settings. Include +packing, preprocessing, offload, and decode time in comparisons. Short sequences +may not benefit. Upstream combined pipeline speedups are not measurements of +this native kernel. Exact-covariance thresholds, text sinks, Morton ordering, +and step/layer schedules are not implemented. + +## Validation + +On an RTX 4090 with CUDA 12.4, Wan 2.1 T2V 1.3B was tested at 832x480, +33 frames, 20 Euler steps, seed 42, CFG 6, and flow shift 3, using the prompt +`a lovely cat` and the same negative prompt for every run: + +| Attention | Sampling time | Total process time | +| --- | ---: | ---: | +| FlashAttention | 45.73 s | 74.63 s | +| Sol-Attn, tau 1 | 37.17 s | 66.20 s | +| Sol-Attn, tau 0 | 40.66 s | 68.50 s | + +These are single-run measurements. The graph selected 30 Sol-Attn nodes and +30 FlashAttention nodes. At tau 1, sampled video frames showed washed-out +colors and reduced detail. Tau 0 improved clarity in this example, but still +changed the composition. Neither setting guarantees the baseline's quality. +For this Wan command, `--sol-attn --sol-attn-tau 0` is a more conservative +starting point. In the one-frame case, tau 1 increased warm sampling time from +0.140 to 0.148 seconds per step. + +Validation also covered 15 numerical reference cases, 11 layout/scaling/fallback +cases, CUDA memory checking, and 36 existing SageAttention regression cases. +CLI and server CUDA builds and the upstream GGML CPU library build passed. +Other GPU architectures, multi-GPU execution, and other models have not been +tested. + +## Library API + +Configure Sol-Attn in `sd_ctx_params_t` before creating the context: + +```cpp +sd_ctx_params_t params; +sd_ctx_params_init(¶ms); +// Set model paths and other context options here. +params.sol_attn = true; +params.sol_attn_tau = 1.0f; +sd_ctx_t* ctx = new_sd_ctx(¶ms); +``` + +`sd_ctx_params_init` defaults `sol_attn` to false and `sol_attn_tau` to 1.0. +`new_sd_ctx` returns null for a nonfinite threshold, unavailable requested +backends, or a conflict with SageAttention. The context owns a copy of these +settings; changing the input structure after creation does not reconfigure it. +Applications must be rebuilt against the updated `sd_ctx_params_t` definition. diff --git a/examples/cli/README.md b/examples/cli/README.md index fe47122b6..85f43040a 100644 --- a/examples/cli/README.md +++ b/examples/cli/README.md @@ -25,3 +25,7 @@ Metadata mode inspects PNG/JPEG container metadata without loading any model: For completely black or white images or videos, NaNs, and the `--linear-scale` / `--attn-scale` workaround, see [Troubleshooting](../../docs/troubleshooting.md). + +For native CUDA sparse attention in the diffusion model, use `--sol-attn`. +See [Sol-Attn](../../docs/sol_attention.md) for requirements, supported shapes, +and the `--sol-attn-tau` threshold coefficient. diff --git a/examples/common/common.cpp b/examples/common/common.cpp index afea3afa0..3c151188e 100644 --- a/examples/common/common.cpp +++ b/examples/common/common.cpp @@ -622,6 +622,10 @@ ArgOptions SDContextParams::get_options() { "--sage-attn", "use native CUDA SageAttention in the diffusion model, with flash/default attention fallback", true, &sage_attn}, + {"", + "--sol-attn", + "use native CUDA Sol-Attn in the diffusion model, with flash/default attention fallback", + true, &sol_attn}, {"", "--diffusion-conv-direct", "use ggml_conv2d_direct in the diffusion model", @@ -719,6 +723,8 @@ ArgOptions SDContextParams::get_options() { return 1; }; + options.float_options.push_back({"", "--sol-attn-tau", "Sol-Attn routing threshold coefficient (default: 1; higher selects fewer exact blocks)", &sol_attn_tau}); + options.manual_options = { {"", "--linear-scale", @@ -822,6 +828,14 @@ bool SDContextParams::resolve(SDMode mode) { } bool SDContextParams::validate(SDMode mode) { + if (sol_attn && sage_attn) { + LOG_ERROR("--sol-attn and --sage-attn cannot be enabled together"); + return false; + } + if (!std::isfinite(sol_attn_tau)) { + LOG_ERROR("--sol-attn-tau must be finite"); + return false; + } if (mode == CONVERT) { const bool has_convert_input = model_path.length() != 0 || clip_l_path.length() != 0 || @@ -943,6 +957,8 @@ std::string SDContextParams::to_string() const { << " flash_attn: " << (flash_attn ? "true" : "false") << ",\n" << " diffusion_flash_attn: " << (diffusion_flash_attn ? "true" : "false") << ",\n" << " sage_attn: " << (sage_attn ? "true" : "false") << ",\n" + << " sol_attn: " << (sol_attn ? "true" : "false") << ",\n" + << " sol_attn_tau: " << sol_attn_tau << ",\n" << " linear_scale: " << linear_scale << ",\n" << " attn_scale: " << attn_scale << ",\n" << " diffusion_conv_direct: " << (diffusion_conv_direct ? "true" : "false") << ",\n" @@ -1001,6 +1017,8 @@ sd_ctx_params_t SDContextParams::to_sd_ctx_params_t(bool taesd_preview) { 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.sol_attn = sol_attn; + sd_ctx_params.sol_attn_tau = sol_attn_tau; 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 0e5c50e3c..cf22ac65e 100644 --- a/examples/common/common.h +++ b/examples/common/common.h @@ -171,6 +171,8 @@ struct SDContextParams { bool flash_attn = false; bool diffusion_flash_attn = false; bool sage_attn = false; + bool sol_attn = false; + float sol_attn_tau = 1.f; bool diffusion_conv_direct = false; bool vae_conv_direct = false; diff --git a/ggml b/ggml index f583f393c..223feb34a 160000 --- a/ggml +++ b/ggml @@ -1 +1 @@ -Subproject commit f583f393cd5dfdc129360bbf75cb3d49ccc837a8 +Subproject commit 223feb34abee6289e275746201ba0a0b53e8c494 diff --git a/include/stable-diffusion.h b/include/stable-diffusion.h index 9173240be..3211f19fa 100644 --- a/include/stable-diffusion.h +++ b/include/stable-diffusion.h @@ -247,6 +247,8 @@ typedef struct { 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; + bool sol_attn; + float sol_attn_tau; } sd_ctx_params_t; typedef struct { diff --git a/src/core/ggml_extend.cpp b/src/core/ggml_extend.cpp index 74324169e..5945907d4 100644 --- a/src/core/ggml_extend.cpp +++ b/src/core/ggml_extend.cpp @@ -623,7 +623,9 @@ ggml_tensor* ggml_ext_attention_ext(ggml_context* ctx, bool skip_reshape, bool flash_attn, float kv_scale, - bool sage_attn) { // avoid overflow + bool sage_attn, + bool sol_attn, + float sol_attn_tau) { // avoid overflow int64_t L_q; int64_t L_k; int64_t C; @@ -715,7 +717,23 @@ ggml_tensor* ggml_ext_attention_ext(ggml_context* ctx, }; #ifndef SD_USE_UPSTREAM_GGML - if (sage_attn && mask == nullptr && d_head > 0 && d_head <= 128) { + if (sol_attn && mask == nullptr && d_head == 128 && L_q == L_k && n_head == n_kv_head) { + 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)); + if (v_in->type != GGML_TYPE_F32) { + v_in = ggml_cast(ctx, v_in, GGML_TYPE_F32); + } + 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); + } + auto out = ggml_sol_attn(ctx, q_in, k_in, v_in, scale / kv_scale, sol_attn_tau); + if (ggml_backend_supports_op(backend, out)) { + kqv = kv_scale != 1.0f ? ggml_ext_scale(ctx, out, 1.0f / kv_scale) : out; + } + } + if (kqv == nullptr && 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)); @@ -744,7 +762,7 @@ ggml_tensor* ggml_ext_attention_ext(ggml_context* ctx, } #endif - if (kqv == nullptr && (flash_attn || sage_attn)) { + if (kqv == nullptr && (flash_attn || sage_attn || sol_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 33d8c3bf2..b274ff223 100644 --- a/src/core/ggml_extend.h +++ b/src/core/ggml_extend.h @@ -217,11 +217,13 @@ ggml_tensor* ggml_ext_attention_ext(ggml_context* ctx, ggml_tensor* k, ggml_tensor* v, int64_t n_head, - ggml_tensor* mask = nullptr, - bool skip_reshape = false, - bool flash_attn = false, - float kv_scale = 1.0f, - bool sage_attn = false); + ggml_tensor* mask = nullptr, + bool skip_reshape = false, + bool flash_attn = false, + float kv_scale = 1.0f, + bool sage_attn = false, + bool sol_attn = false, + float sol_attn_tau = 1.0f); 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 abfa8639b..5ca718f0a 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, ctx->sage_attn_enabled); + 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, ctx->sol_attn_enabled, ctx->sol_attn_tau); } void GGMLRunner::alloc_params_ctx() { @@ -164,6 +164,22 @@ ggml_cgraph* GGMLRunner::get_compute_graph(get_graph_cb_t get_graph) { } } prepare_build_in_tensor_after(gf); +#ifndef SD_USE_UPSTREAM_GGML + if (sol_attn_enabled && !sol_attn_graph_logged) { + int sol_nodes = 0; + int flash_nodes = 0; + for (int i = 0; i < ggml_graph_n_nodes(gf); ++i) { + const auto op = ggml_graph_node(gf, i)->op; + sol_nodes += op == GGML_OP_SOL_ATTN; + flash_nodes += op == GGML_OP_FLASH_ATTN_EXT; + } + LOG_INFO("Sol-Attn graph: %d Sol-Attn nodes, %d FlashAttention nodes", sol_nodes, flash_nodes); + if (sol_nodes == 0) { + LOG_WARN("This graph has no attention operations supported by Sol-Attn"); + } + sol_attn_graph_logged = true; + } +#endif return gf; } @@ -521,6 +537,8 @@ GGMLRunnerContext GGMLRunner::get_context() { runner_ctx.backend = runtime_backend; runner_ctx.flash_attn_enabled = flash_attn_enabled; runner_ctx.sage_attn_enabled = sage_attn_enabled; + runner_ctx.sol_attn_enabled = sol_attn_enabled; + runner_ctx.sol_attn_tau = sol_attn_tau; 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 29a6808bd..bef7402fd 100644 --- a/src/core/ggml_runner.h +++ b/src/core/ggml_runner.h @@ -69,6 +69,8 @@ struct GGMLRunnerContext { ggml_context* ggml_ctx = nullptr; bool flash_attn_enabled = false; bool sage_attn_enabled = false; + bool sol_attn_enabled = false; + float sol_attn_tau = 1.f; float linear_scale = 0.f; float attn_scale = 0.f; bool conv2d_direct_enabled = false; @@ -178,6 +180,9 @@ struct GGMLRunner { bool flash_attn_enabled = false; bool sage_attn_enabled = false; + bool sol_attn_enabled = false; + float sol_attn_tau = 1.f; + bool sol_attn_graph_logged = false; float linear_scale = 0.f; float attn_scale = 0.f; bool conv2d_direct_enabled = false; @@ -347,6 +352,16 @@ struct GGMLRunner { } } + void set_sol_attention_enabled(bool enabled, float tau) { + if (sol_attn_enabled != enabled || sol_attn_tau != tau) { + free_cache_ctx_and_buffer(); + graph_cut_plan_cache_.graph_cut_plans.clear(); + sol_attn_enabled = enabled; + sol_attn_tau = tau; + sol_attn_graph_logged = false; + } + } + 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 e41c1151a..ff7e08576 100644 --- a/src/pipeline/diffusion_engine.cpp +++ b/src/pipeline/diffusion_engine.cpp @@ -899,7 +899,60 @@ bool StableDiffusionGGML::set_sage_attention_enabled(bool enabled) { return true; } +bool StableDiffusionGGML::set_sol_attention_enabled(bool enabled, float tau) { + if (!diffusion_model || !std::isfinite(tau)) { + LOG_ERROR("Sol-Attn requires a diffusion model and finite tau"); + return false; + } + if (enabled) { + if (config_->params.sage_attn) { + LOG_ERROR("Sol-Attn and SageAttention cannot be enabled together"); + return false; + } +#ifndef SD_USE_UPSTREAM_GGML + auto* ctx = ggml_init({4 * ggml_tensor_overhead(), nullptr, true}); + if (!ctx) { + 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_F32, 128, 128, 1, 1); + auto* op = ggml_sol_attn(ctx, q, k, v, 1.f / sqrtf(128.f), tau); + bool supported = true; + for (auto backend : backend_manager.runtime_backends(SDBackendModule::DIFFUSION)) { + if (!ggml_backend_supports_op(backend, op)) { + LOG_ERROR("Sol-Attn is unavailable on %s; it requires patched GGML, CUDA 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("Sol-Attn requires -DSD_USE_UPSTREAM_GGML=OFF and a CUDA backend"); + return false; +#endif + } + diffusion_model->set_sol_attention_enabled(enabled, tau); + if (high_noise_diffusion_model) { + high_noise_diffusion_model->set_sol_attention_enabled(enabled, tau); + } + if (enabled) { + LOG_INFO("Using Sol-Attn (tau=%g, diagonal threshold) in diffusion; unsupported attention uses flash/default attention", tau); + } + return true; +} + bool StableDiffusionGGML::init(const sd_ctx_params_t* sd_ctx_params) { + if (sd_ctx_params->sol_attn && sd_ctx_params->sage_attn) { + LOG_ERROR("Sol-Attn and SageAttention cannot be enabled together"); + return false; + } + if (!std::isfinite(sd_ctx_params->sol_attn_tau)) { + LOG_ERROR("Sol-Attn tau must be finite"); + return false; + } #ifdef SD_USE_UPSTREAM_GGML LOG_WARN( "Using upstream GGML: INT8 tensorwise/convrot is disabled and FP8 weights are " @@ -1180,6 +1233,9 @@ bool StableDiffusionGGML::validate_and_load_runners() { if (sd_ctx_params->sage_attn && !set_sage_attention_enabled(true)) { return false; } + if (sd_ctx_params->sol_attn && !set_sol_attention_enabled(true, sd_ctx_params->sol_attn_tau)) { + 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 f2f8a8ced..6a5b39f9a 100644 --- a/src/pipeline/diffusion_engine.h +++ b/src/pipeline/diffusion_engine.h @@ -313,6 +313,7 @@ class StableDiffusionGGML { bool init(const sd_ctx_params_t* sd_ctx_params); bool set_sage_attention_enabled(bool enabled); + bool set_sol_attention_enabled(bool enabled, float tau); bool uses_tae() const; diff --git a/src/stable-diffusion.cpp b/src/stable-diffusion.cpp index a8f83336d..68c930029 100644 --- a/src/stable-diffusion.cpp +++ b/src/stable-diffusion.cpp @@ -338,6 +338,8 @@ void sd_ctx_params_init(sd_ctx_params_t* sd_ctx_params) { sd_ctx_params->enable_mmap = false; sd_ctx_params->diffusion_flash_attn = false; sd_ctx_params->sage_attn = false; + sd_ctx_params->sol_attn = false; + sd_ctx_params->sol_attn_tau = 1.f; sd_ctx_params->linear_scale = 0.f; sd_ctx_params->attn_scale = 0.f; sd_ctx_params->vae_format = SD_VAE_FORMAT_AUTO; @@ -394,6 +396,8 @@ char* sd_ctx_params_to_str(const sd_ctx_params_t* sd_ctx_params) { "flash_attn: %s\n" "diffusion_flash_attn: %s\n" "sage_attn: %s\n" + "sol_attn: %s\n" + "sol_attn_tau: %g\n" "linear_scale: %g\n" "attn_scale: %g\n" "vae_format: %s\n", @@ -434,6 +438,8 @@ char* sd_ctx_params_to_str(const sd_ctx_params_t* sd_ctx_params) { BOOL_STR(sd_ctx_params->flash_attn), BOOL_STR(sd_ctx_params->diffusion_flash_attn), BOOL_STR(sd_ctx_params->sage_attn), + BOOL_STR(sd_ctx_params->sol_attn), + sd_ctx_params->sol_attn_tau, sd_ctx_params->linear_scale, sd_ctx_params->attn_scale, sd_vae_format_name(sd_ctx_params->vae_format));