perf: accelerate INT8 diffusion inference with CUDA fusions (#2235)

.dockerignore

@@ -1,7 +1,8 @@
 build*/
 docs/
 test/
+tests/
 
 .cache/
 *.swp
-models/
+models/

.github/workflows/build.yml

@@ -439,6 +439,9 @@ jobs:
             ls .\build\bin
             throw "Can't find stable-diffusion.dll"
           }
+          if ("${{ matrix.build }}" -eq "cuda12") {
+            Copy-Item ggml/src/ggml-cuda/vendors/cutlass/LICENSE.txt (Join-Path (Split-Path -Parent $filePath) 'cutlass.txt')
+          }
           7z a sd-${{ env.BRANCH_NAME }}-${{ steps.commit.outputs.short }}-bin-win-${{ matrix.build }}-x64.zip $filePath .\start-server.sh .\start-server.bat
 
       - name: Copy and pack Cuda runtime

.gitignore

@@ -2,6 +2,7 @@ build*/
 cmake-build-*/
 test/
 tests/
+logs/
 .vscode/
 .idea/
 .cache/

docker/Dockerfile.cuda

@@ -1,7 +1,7 @@
 ARG CUDA_VERSION=12.6.3
 ARG UBUNTU_VERSION=24.04
 
-FROM nvidia/cuda:${CUDA_VERSION}-cudnn-devel-ubuntu${UBUNTU_VERSION} AS build
+FROM nvidia/cuda:${CUDA_VERSION}-devel-ubuntu${UBUNTU_VERSION} AS build
 
 # sd-server embeds the web UI at build time, so the build image needs Node/pnpm.
 RUN apt-get update && apt-get install -y --no-install-recommends build-essential git ccache cmake ca-certificates curl gnupg && \
@@ -46,13 +46,14 @@ RUN set -- \
     cmake . -B ./build "$@"
 RUN cmake --build ./build --config Release -j$(nproc)
 
-FROM nvidia/cuda:${CUDA_VERSION}-cudnn-runtime-ubuntu${UBUNTU_VERSION} AS runtime
+FROM nvidia/cuda:${CUDA_VERSION}-runtime-ubuntu${UBUNTU_VERSION} AS runtime
 
 RUN apt-get update && \
     apt-get install --yes --no-install-recommends libgomp1 && \
     apt-get clean
 
 COPY --from=build /sd.cpp/build/bin /sd.cpp/bin
+COPY --from=build /sd.cpp/ggml/src/ggml-cuda/vendors/cutlass/LICENSE.txt /sd.cpp/bin/cutlass.txt
 RUN printf '#!/bin/sh\nexec /sd.cpp/bin/sd-cli "$@"\n' > /sd-cli && \
     printf '#!/bin/sh\nexec /sd.cpp/bin/sd-server "$@"\n' > /sd-server && \
     chmod +x /sd-cli /sd-server

docs/build.md

@@ -2,18 +2,23 @@
 
 ## Get the Code
 
-```
+```shell
 git clone --recursive https://github.com/leejet/stable-diffusion.cpp
 cd stable-diffusion.cpp
 ```
 
-- If you have already cloned the repository, you can use the following command to update the repository to the latest code.
+If you already cloned the repository without submodules, initialize them from the repository root:
 
+```shell
+git submodule update --init --recursive
 ```
+
+To update an existing checkout and its submodules:
+
+```shell
 cd stable-diffusion.cpp
 git pull origin master
-git submodule init
+git submodule update --init --recursive
-git submodule update
 ```
 
 ## Selecting a GGML source tree
@@ -102,6 +107,25 @@ cmake .. -DSD_CUDA=ON
 cmake --build . --config Release
 ```
 
+CUDA builds with the bundled patched GGML use the CUTLASS submodule by default
+for eligible INT8 GEMM and CuTe attention operations. If its headers are missing,
+CMake stops with an initialization hint; it does not download them automatically.
+Run `git submodule update --init --recursive` from the repository root before
+configuring.
+
+To build with the existing CUDA implementations without CUTLASS, configure from
+the repository root with:
+
+```shell
+cmake -S . -B build -DSD_CUDA=ON -DGGML_CUDA_CUTLASS_PATH:PATH=
+cmake --build build --config Release
+```
+
+Alternatively, set `GGML_CUDA_CUTLASS_PATH` to another local CUTLASS source
+directory (tested with 3.9.2). Existing CMake caches retain their configured path;
+pass `-UGGML_CUDA_CUTLASS_PATH` when reconfiguring to restore the bundled default.
+Non-CUDA builds do not require CUTLASS.
+
 ## Build with HipBLAS
 
 This provides GPU acceleration using AMD GPU. Make sure to have the ROCm toolkit installed.

docs/docker.md

@@ -18,12 +18,26 @@ docker run --rm --init -v /path/to/models:/models -v /path/to/output/:/output -p
 
 ## Building using Docker
 
+Initialize the submodules from the repository root before building. Docker copies
+the local checkout, so its build context must include the submodule contents:
+
 ```shell
+git submodule update --init --recursive
 docker build -f docker/Dockerfile -t sd .
 ```
 
 ## Building variants using Docker
 
+### CUDA (NVIDIA GPU):
+
+```shell
+docker build -f docker/Dockerfile.cuda -t sd-cuda .
+```
+
+To select GPU architectures, pass `--build-arg "CUDA_ARCHITECTURES=89"` (for
+example, RTX 4090). Run the image on a host configured with NVIDIA Container
+Toolkit, using `docker run --rm --gpus all ... sd-cuda [args...]`.
+
 ### Vulkan:
 
 ```shell
@@ -67,4 +81,4 @@ docker run --rm --device=/dev/dri --device=/dev/kfd --group-add video \
 #   -v ./models:/models -v ./build:/output \
 #   sd -m /models/sd-v1-4.ckpt -v -o /output/output.png \
 #   -p "a lovely cat"
-```
+```

src/core/compute_workspace.cpp

@@ -1,16 +1,13 @@
 #include "core/compute_workspace.h"
 
 #include <algorithm>
-#include <cstring>
 #include <exception>
 #include <map>
 #include <unordered_map>
 #include <unordered_set>
 
 #include "core/ggml_extend_backend.h"
-#include "core/ggml_graph_cut.h"
 #include "core/util.h"
-#include "ggml-cpu.h"
 #include "ggml-impl.h"
 
 namespace sd {
@@ -28,18 +25,6 @@ namespace sd {
         }
     }
 
-    bool ComputeWorkspace::needs_scheduler(ggml_cgraph* graph) const {
-        if (!extra_backends_.empty()) {
-            return true;
-        }
-        for (int i = 0; i < ggml_graph_n_nodes(graph); ++i) {
-            if (!ggml_backend_supports_op(backend_, ggml_graph_node(graph, i))) {
-                return true;
-            }
-        }
-        return false;
-    }
-
     ggml_backend_sched_t ComputeWorkspace::make_scheduler(size_t graph_size) {
         std::vector<ggml_backend_t> backends{backend_};
         backends.insert(backends.end(), extra_backends_.begin(), extra_backends_.end());
@@ -63,10 +48,6 @@ namespace sd {
                                       graph_size, false, false);
     }
 
-    bool ComputeWorkspace::measurement_matches(ggml_cgraph* graph, const Measurement& measurement) const {
-        return measurement.scheduler == needs_scheduler(graph);
-    }
-
     bool ComputeWorkspace::prepare(const Measurement& measurement) {
         GGML_ASSERT(!active_);
         if (measurement.buffers.empty()) {
@@ -74,7 +55,7 @@ namespace sd {
         }
         const bool grows = std::any_of(measurement.buffers.begin(), measurement.buffers.end(),
                                        [&](const BackendBufferSize& size) { return size.bytes > bytes(size.backend); });
-        if (measurement.scheduler != (scheduler_ != nullptr) || grows) {
+        if (grows) {
             release();
         }
         return true;
@@ -88,47 +69,21 @@ namespace sd {
 
     bool ComputeWorkspace::allocate(ggml_cgraph* graph, const AssignNodes& assign_nodes) {
         GGML_ASSERT(!active_);
-        const bool use_scheduler = needs_scheduler(graph);
+        const size_t capacity = static_cast<size_t>(graph->n_nodes + graph->n_leafs) + 8;
-        if (use_scheduler) {
+        if (scheduler_ == nullptr || capacity > scheduler_capacity_) {
-            if (allocator_ != nullptr) {
+            release();
-                release();
+            scheduler_          = make_scheduler(capacity);
-            }
+            scheduler_capacity_ = capacity;
-            const size_t capacity = static_cast<size_t>(graph->n_nodes + graph->n_leafs) + 8;
+        }
-            if (scheduler_ == nullptr || capacity > scheduler_capacity_) {
+        if (scheduler_ == nullptr) {
-                release();
+            return false;
-                scheduler_          = make_scheduler(capacity);
+        }
-                scheduler_capacity_ = capacity;
+        ggml_backend_sched_reset(scheduler_);
-            }
+        assign_nodes(scheduler_, graph);
-            if (scheduler_ == nullptr) {
+        // Scheduler allocation rewrites sources. Split the execution graph only once.
-                return false;
+        if (!ggml_backend_sched_alloc_graph(scheduler_, graph)) {
-            }
+            release();
-            ggml_backend_sched_reset(scheduler_);
+            return false;
-            assign_nodes(scheduler_, graph);
-            // Scheduler allocation rewrites sources. Split the execution graph only once.
-            if (!ggml_backend_sched_alloc_graph(scheduler_, graph)) {
-                release();
-                return false;
-            }
-        } else {
-            if (scheduler_ != nullptr) {
-                release();
-            }
-            if (allocator_ == nullptr) {
-                allocator_ = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend_));
-            }
-            auto signature = ggml_graph_cut::graph_layout(graph, true);
-            if (signature != reservation_) {
-                if (!ggml_gallocr_reserve(allocator_, graph)) {
-                    release();
-                    return false;
-                }
-                reservation_ = std::move(signature);
-                ++reservations_;
-            }
-            if (!ggml_gallocr_alloc_graph(allocator_, graph)) {
-                release();
-                return false;
-            }
         }
         active_ = true;
         return true;
@@ -136,12 +91,8 @@ namespace sd {
 
     ComputeWorkspace::Measurement ComputeWorkspace::measure(
         ggml_cgraph* graph,
-        size_t direct_bytes,
         const std::function<ggml_backend_t(const ggml_tensor*)>& external_backend,
         const AssignNodes& assign_nodes) {
-        if (!needs_scheduler(graph)) {
-            return {{{backend_, direct_bytes}}, false};
-        }
         std::vector<const ggml_tensor*> tensors;
         std::unordered_set<const ggml_tensor*> seen;
         auto visit = [&](const ggml_tensor* tensor) {
@@ -202,9 +153,17 @@ namespace sd {
         for (int i = 0; i < graph->n_leafs; ++i) {
             copy_graph->leafs[i] = copies.at(graph->leafs[i]);
         }
+        copy_graph->order = graph->order;
+        // Preserve fusion metadata, including uses outside a graph view, without rebuilding its node order.
+        for (const auto& entry : copies) {
+            const size_t source_index = ggml_hash_find(&graph->visited_hash_set, entry.first);
+            if (source_index != GGML_HASHSET_FULL && ggml_bitset_get(graph->visited_hash_set.used, source_index)) {
+                const size_t copy_index = ggml_hash_insert(&copy_graph->visited_hash_set, entry.second);
+                copy_graph->use_counts[copy_index] = graph->use_counts[source_index];
+            }
+        }
         Measurement result;
-        result.scheduler = true;
+        auto scheduler = make_scheduler(graph_size);
-        auto scheduler   = make_scheduler(graph_size);
         if (scheduler != nullptr) {
             assign_nodes(scheduler, copy_graph);
             std::vector<size_t> sizes(extra_backends_.size() + 2);
@@ -253,12 +212,9 @@ namespace sd {
         if (active_) {
             return false;
         }
-        ggml_gallocr_free(allocator_);
-        allocator_ = nullptr;
         ggml_backend_sched_free(scheduler_);
         scheduler_          = nullptr;
         scheduler_capacity_ = 0;
-        reservation_.clear();
         return true;
     }
 
@@ -266,6 +222,6 @@ namespace sd {
         if (scheduler_ != nullptr) {
             return ggml_backend_sched_get_buffer_size(scheduler_, backend);
         }
-        return allocator_ != nullptr && backend == backend_ ? ggml_gallocr_get_buffer_size(allocator_, 0) : 0;
+        return 0;
     }
 }

src/core/compute_workspace.h

@@ -4,7 +4,6 @@
 #include <functional>
 #include <vector>
 
-#include "ggml-alloc.h"
 #include "ggml-backend.h"
 
 namespace sd {
@@ -17,20 +16,15 @@ namespace sd {
         ggml_backend_t backend_;
         std::vector<ggml_backend_t> extra_backends_;
         ggml_backend_t cpu_backend_     = nullptr;
-        ggml_gallocr_t allocator_       = nullptr;
         ggml_backend_sched_t scheduler_ = nullptr;
         size_t scheduler_capacity_      = 0;
-        std::vector<uint64_t> reservation_;
+        bool active_                    = false;
-        bool active_         = false;
-        size_t reservations_ = 0;
 
         ggml_backend_sched_t make_scheduler(size_t graph_size);
-        bool needs_scheduler(ggml_cgraph* graph) const;
 
     public:
         struct Measurement {
             std::vector<BackendBufferSize> buffers;
-            bool scheduler = false;
         };
         using AssignNodes = std::function<void(ggml_backend_sched_t, ggml_cgraph*)>;
 
@@ -41,13 +35,11 @@ namespace sd {
         ComputeWorkspace& operator=(const ComputeWorkspace&) = delete;
 
         void set_extra_backends(const std::vector<ggml_backend_t>& backends);
-        bool measurement_matches(ggml_cgraph* graph, const Measurement& measurement) const;
         bool prepare(const Measurement& measurement);
         bool release_excess(const Measurement& measurement);
         bool allocate(ggml_cgraph* graph, const AssignNodes& assign_nodes);
         Measurement measure(
             ggml_cgraph* graph,
-            size_t direct_bytes,
             const std::function<ggml_backend_t(const ggml_tensor*)>& external_backend,
             const AssignNodes& assign_nodes);
         void synchronize() const;
@@ -57,7 +49,6 @@ namespace sd {
         ggml_backend_sched_t scheduler() const { return scheduler_; }
         ggml_backend_t cpu_backend() const { return cpu_backend_; }
         size_t bytes(ggml_backend_t backend) const;
-        size_t reservation_count() const { return reservations_; }
     };
 }
 

src/core/ggml_runner.cpp

@@ -52,8 +52,13 @@ void GGMLRunner::free_params_ctx() {
 
 void GGMLRunner::alloc_compute_ctx() {
     ggml_init_params params;
-    params.mem_size   = static_cast<size_t>(ggml_tensor_overhead() * MAX_GRAPH_SIZE + ggml_graph_overhead());
+    params.mem_size = static_cast<size_t>(ggml_tensor_overhead() * MAX_GRAPH_SIZE + ggml_graph_overhead());
-    params.mem_buffer = nullptr;
+    if (compute_ctx_buffer_ == nullptr) {
+        const size_t count = (params.mem_size + sizeof(std::max_align_t) - 1) / sizeof(std::max_align_t);
+        compute_ctx_buffer_.reset(new std::max_align_t[count]);
+    }
+    // Keep tensor addresses stable between steps so CUDA graphs can finish warming up.
+    params.mem_buffer = compute_ctx_buffer_.get();
     params.no_alloc   = true;
 
     compute_ctx = ggml_init(params);
@@ -482,6 +487,8 @@ void GGMLRunner::runner_end() {
         return;
     }
     workspace_.release();
+    free_compute_ctx();
+    compute_ctx_buffer_.reset();
     cache_.clear();
     logged_compute_bytes_.clear();
     logged_segment_count_ = 0;
@@ -705,7 +712,7 @@ static size_t add_bytes(size_t a, size_t b) {
     return b > SIZE_MAX - a ? SIZE_MAX : a + b;
 }
 
-ComputeWorkspace::Measurement GGMLRunner::measure(ggml_cgraph* graph, size_t direct_bytes) {
+ComputeWorkspace::Measurement GGMLRunner::measure(ggml_cgraph* graph) {
     auto external_backend = [&](const ggml_tensor* tensor) -> ggml_backend_t {
         if (!params_tensor_set_.count(tensor)) {
             return nullptr;
@@ -716,7 +723,7 @@ ComputeWorkspace::Measurement GGMLRunner::measure(ggml_cgraph* graph, size_t dir
     auto assign_nodes = [&](ggml_backend_sched_t scheduler, ggml_cgraph* copy) {
         pin_multi_device_nodes(scheduler, copy, graph);
     };
-    return workspace_.measure(graph, direct_bytes, external_backend, assign_nodes);
+    return workspace_.measure(graph, external_backend, assign_nodes);
 }
 
 std::vector<DeviceMemoryRequest> GGMLRunner::memory_requests(
@@ -767,18 +774,9 @@ bool GGMLRunner::execute_segment(ggml_cgraph* graph, int n_threads) {
         sd_backend_cpu_set_n_threads(workspace_.cpu_backend(), n_threads);
     }
     auto scheduler = workspace_.scheduler();
-    ggml_status status;
+    GGML_ASSERT(scheduler != nullptr);
-    if (scheduler != nullptr) {
+    ggml_backend_sched_set_eval_callback(scheduler, sd_get_backend_eval_callback(), sd_get_backend_eval_callback_data());
-        if (sd_get_backend_eval_callback() != nullptr && !multi_device_eval_callback_warned) {
+    const auto status = ggml_backend_sched_graph_compute(scheduler, graph);
-            LOG_WARN("%s: eval callback is not supported with the backend scheduler; ignoring", get_desc().c_str());
-            multi_device_eval_callback_warned = true;
-        }
-        status = ggml_backend_sched_graph_compute(scheduler, graph);
-    } else {
-        status = sd_backend_graph_compute_with_eval_callback(runtime_backend, graph,
-                                                             sd_get_backend_eval_callback(),
-                                                             sd_get_backend_eval_callback_data());
-    }
     workspace_.synchronize();
     if (status != GGML_STATUS_SUCCESS) {
         last_compute_status_ = status;
@@ -832,7 +830,7 @@ std::optional<Tensor<float>> GGMLRunner::execute_graph(ggml_cgraph* graph, int n
     }
     const auto params           = collect_used_param_tensors(graph);
     const auto& cached_plan     = resolve_graph_cut_plan(graph);
-    const auto full_measurement = measure(graph, cached_plan.compute_buffer_size);
+    const auto full_measurement = measure(graph);
     if (full_measurement.buffers.empty()) {
         last_compute_status_ = GGML_STATUS_ALLOC_FAILED;
         return std::nullopt;
@@ -924,7 +922,7 @@ std::optional<Tensor<float>> GGMLRunner::execute_graph(ggml_cgraph* graph, int n
             }
         } segment_cleanup{*this, weights, bindings, segment_context};
 
-        auto measurement = segmented ? measure(segment_graph, segment.compute_buffer_size) : full_measurement;
+        auto measurement = segmented ? measure(segment_graph) : full_measurement;
         if (!workspace_.prepare(measurement)) {
             last_compute_status_ = GGML_STATUS_ALLOC_FAILED;
             return fail_segment("workspace preparation");
@@ -951,9 +949,6 @@ std::optional<Tensor<float>> GGMLRunner::execute_graph(ggml_cgraph* graph, int n
             return fail_segment("weight preparation");
         }
         // Preparing weights can execute LoRA graphs and reclaim an idle workspace.
-        if (!workspace_.measurement_matches(segment_graph, measurement)) {
-            measurement = measure(segment_graph, segment.compute_buffer_size);
-        }
         if (!workspace_.prepare(measurement)) {
             last_compute_status_ = GGML_STATUS_ALLOC_FAILED;
             return fail_segment("workspace preparation");

src/core/ggml_runner.h

@@ -141,7 +141,7 @@ private:
     size_t logged_segment_count_     = 0;
     ggml_status last_compute_status_ = GGML_STATUS_SUCCESS;
 
-    sd::ComputeWorkspace::Measurement measure(ggml_cgraph* graph, size_t direct_bytes);
+    sd::ComputeWorkspace::Measurement measure(ggml_cgraph* graph);
     std::vector<DeviceMemoryRequest> memory_requests(const std::vector<sd::BackendBufferSize>& sizes,
                                                      size_t pending_cache_bytes) const;
     bool fits(const std::vector<DeviceMemoryRequest>& requests,
@@ -161,6 +161,7 @@ protected:
     sd::GraphCutTensorCache cut_cache_;
     sd::ComputeWorkspace workspace_;
     ggml_context* compute_ctx = nullptr;
+    std::unique_ptr<std::max_align_t[]> compute_ctx_buffer_;
     bool runner_started_      = false;
     bool graph_active_        = false;
 
@@ -169,7 +170,6 @@ protected:
     std::vector<size_t> graph_cut_layer_split_backend_vram_limits_;
 
     std::vector<ggml_backend_t> extra_runtime_backends;  // borrowed (SDBackendManager-owned)
-    bool multi_device_eval_callback_warned = false;
 
     std::shared_ptr<WeightAdapter> weight_adapter = nullptr;
     std::weak_ptr<DeviceResidencyManager> residency_manager;

src/model/diffusion/qwen_image_2_1.hpp

@@ -266,16 +266,18 @@ namespace Qwen {
             blocks["img_mlp.out"] = std::make_shared<Linear>(config.intermediate_size, config.hidden_size, false);
         }
 
-        static ggml_tensor* modulate(ggml_context* ctx, ggml_tensor* x, ggml_tensor* params, int64_t prefix_length, bool gate = false) {
+        static ggml_tensor* modulate(GGMLRunnerContext* runner_ctx, ggml_tensor* x, ggml_tensor* params, int64_t prefix_length, bool gate = false) {
+            auto ctx   = runner_ctx->ggml_ctx;
             auto rows  = ggml_ext_chunk(ctx, params, 2, 1);
             auto apply = [&](ggml_tensor* part, ggml_tensor* row) {
                 row = gate ? ggml_tanh(ctx, row) : ggml_scale_bias(ctx, row, 1.f, 1.f);
+                runner_ctx->expand_graph(row);
                 return ggml_mul(ctx, part, row);
             };
-            auto target = apply(ggml_ext_slice(ctx, x, 1, prefix_length, x->ne[1]), rows[0]);
             if (prefix_length == 0) {
-                return target;
+                return apply(x, rows[0]);
             }
+            auto target = apply(ggml_ext_slice(ctx, x, 1, prefix_length, x->ne[1]), rows[0]);
             auto prefix = apply(ggml_ext_slice(ctx, x, 1, 0, prefix_length), rows[1]);
             return ggml_concat(ctx, prefix, target, 1);
         }
@@ -283,25 +285,22 @@ namespace Qwen {
         ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, const std::vector<ggml_tensor*>& modulation, ggml_tensor* pe, const QwenImage21Layout& layout, const std::vector<ggml_tensor*>& masks, const QwenImage21PrefixCache& cache) {
             const int64_t prefix_length = cache.mode == QwenImage21PrefixCache::Mode::REUSE ? 0 : layout.prefix_length;
             auto h                      = std::dynamic_pointer_cast<LayerNorm>(blocks["img_norm1"])->forward(ctx, x);
-            h                           = modulate(ctx->ggml_ctx, h, modulation[0], prefix_length);
+            h                           = modulate(ctx, h, modulation[0], prefix_length);
             h                           = std::dynamic_pointer_cast<QwenImage21Attention>(blocks["attn"])->forward(ctx, h, pe, layout.segments, masks, cache);
-            x                           = ggml_add(ctx->ggml_ctx, x, modulate(ctx->ggml_ctx, h, modulation[1], prefix_length, true));
+            x                           = ggml_add(ctx->ggml_ctx, x, modulate(ctx, h, modulation[1], prefix_length, true));
             h                           = std::dynamic_pointer_cast<LayerNorm>(blocks["img_norm2"])->forward(ctx, x);
-            h                           = modulate(ctx->ggml_ctx, h, modulation[2], prefix_length);
+            h                           = modulate(ctx, h, modulation[2], prefix_length);
-            ggml_tensor* gate;
+            auto fused                  = blocks.find("img_mlp.gate_up");
-            auto fused = blocks.find("img_mlp.gate_up");
             if (fused != blocks.end()) {
                 auto gate_up = std::dynamic_pointer_cast<Linear>(fused->second)->forward(ctx, h);
-                auto parts   = ggml_ext_chunk(ctx->ggml_ctx, gate_up, 2, 0);
+                h            = ggml_swiglu(ctx->ggml_ctx, gate_up);
-                gate         = parts[0];
-                h            = parts[1];
             } else {
-                gate = std::dynamic_pointer_cast<Linear>(blocks["img_mlp.gate_layer"])->forward(ctx, h);
+                auto gate = std::dynamic_pointer_cast<Linear>(blocks["img_mlp.gate_layer"])->forward(ctx, h);
-                h    = std::dynamic_pointer_cast<Linear>(blocks["img_mlp.proj"])->forward(ctx, h);
+                h         = std::dynamic_pointer_cast<Linear>(blocks["img_mlp.proj"])->forward(ctx, h);
+                h         = ggml_swiglu_split(ctx->ggml_ctx, gate, h);
             }
-            h = ggml_mul(ctx->ggml_ctx, h, ggml_silu(ctx->ggml_ctx, gate));
             h = std::dynamic_pointer_cast<Linear>(blocks["img_mlp.out"])->forward(ctx, h);
-            return ggml_add(ctx->ggml_ctx, x, modulate(ctx->ggml_ctx, h, modulation[3], prefix_length, true));
+            return ggml_add(ctx->ggml_ctx, x, modulate(ctx, h, modulation[3], prefix_length, true));
         }
     };