feat: add native CUDA Sol-Attn support

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(&params);
+// 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(&params);
+```
+
+`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.

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.

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;

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;
 

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 {

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) {

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,

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;

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;