ggml-cuda: generic DiT and VAE fusions for every family, plus a rope PDL race fix

cmake/ggml.cmake

@@ -27,6 +27,7 @@ if(EXISTS "${sd_ggml_header}")
     set_property(DIRECTORY APPEND PROPERTY CMAKE_CONFIGURE_DEPENDS "${sd_ggml_header}")
     file(STRINGS "${sd_ggml_header}" sd_ggml_h3_fused_ops REGEX "ggml_rope_pe_permute\\(")
     file(STRINGS "${sd_ggml_header}" sd_ggml_h3_fused_qk_norm REGEX "ggml_rope_pe_permute_rms\\(")
+    file(STRINGS "${sd_ggml_header}" sd_ggml_rope_pe_interleaved REGEX "ggml_rope_pe_permute_set_interleaved\\(")
 endif()
 if(sd_ggml_h3_fused_ops)
     target_compile_definitions(${SD_LIB} PUBLIC SD_GGML_H3_FUSED_OPS)
@@ -37,6 +38,10 @@ endif()
 if(sd_ggml_h3_fused_qk_norm)
     target_compile_definitions(${SD_LIB} PUBLIC SD_GGML_H3_FUSED_QK_NORM)
 endif()
+# Interleaved pairing for the fused RoPE op (patch 0007): the shared DiT attention (Rope::attention) uses it.
+if(sd_ggml_h3_fused_ops AND sd_ggml_rope_pe_interleaved)
+    target_compile_definitions(${SD_LIB} PUBLIC SD_GGML_FUSED_ROPE_ATTN)
+endif()
 
 if(SD_USE_UPSTREAM_GGML)
     target_compile_definitions(${SD_LIB} PUBLIC SD_USE_UPSTREAM_GGML)

docs/performance.md

@@ -152,3 +152,17 @@ Use `--params-backend diffusion=disk` instead when reducing system RAM residency
 ## Use quantization to reduce memory usage.
 
 [quantization](./quantization_and_gguf.md)
+
+## Fused DiT and VAE kernels (CUDA)
+
+With the fork's ggml patches (`scripts/unsloth/ggml-patches/0007-ggml-cuda-generic-dit-vae.patch`) these run by
+default on CUDA. Each gives the same output bit for bit as the ops it replaces, and each has a switch (`=0` turns it off):
+
+| what | families | switch |
+|---|---|---|
+| Shared DiT attention (`Rope::attention`, flash attention on): RoPE, head-major relayout, kv scale and F16 cast of q / k / v in one op per tensor | FLUX.1, FLUX.2, Chroma, Z-Image, Qwen-Image, Wan, and the other DiTs using `Rope::attention` | `SD_FUSED_ROPE_ATTN` |
+| GroupNorm spread over the whole GPU, fused with the per-channel affine and the SiLU after it | SD 1.x / SDXL U-Net, every KL VAE decode (SD, SDXL, FLUX, Z-Image, Chroma) | `GGML_CUDA_GROUP_NORM_PAR` |
+| Channel RMS norm of the Wan-style VAE computed in place, without the two relayout copies around it (large layers only, so mostly untiled decodes; `=2` forces it at every size) | Qwen-Image, Qwen-Image-2.1, Wan VAE | `GGML_CUDA_NORM_CHANNEL` |
+| `x*s + y (+ z)` elementwise chains (adaLN modulate, gated residual) in one pass | most DiTs | `GGML_CUDA_FUSE_MUL_ADD` |
+
+Convolutions on a batch of one (and the Wan VAE's single-frame 2D convolutions) also skip relayout copies that move no data.

src/core/ggml_extend.cpp

@@ -451,7 +451,14 @@ ggml_tensor* ggml_ext_conv_2d(ggml_context* ctx,
     if (direct) {
         x = ggml_conv_2d_direct(ctx, w, x, s0, s1, p0, p1, d0, d1);
     } else {
-        x = ggml_conv_2d(ctx, w, x, s0, s1, p0, p1, d0, d1);
+        // ggml_conv_2d, but skips the final [OC, N] -> [N, OC] relayout when N == 1 (no data moves).
+        ggml_tensor* im2col = ggml_im2col(ctx, w, x, s0, s1, p0, p1, d0, d1, true,
+                                          w->type == GGML_TYPE_BF16 ? GGML_TYPE_F32 : GGML_TYPE_F16);  // [N, OH, OW, IC * KH * KW]
+        ggml_tensor* out    = ggml_mul_mat(ctx,
+                                           ggml_reshape_2d(ctx, im2col, im2col->ne[0], im2col->ne[3] * im2col->ne[2] * im2col->ne[1]),
+                                           ggml_reshape_2d(ctx, w, (w->ne[0] * w->ne[1] * w->ne[2]), w->ne[3]));
+        out                 = ggml_reshape_4d(ctx, out, im2col->ne[1], im2col->ne[2], im2col->ne[3], w->ne[3]);  // [OC, N, OH, OW]
+        x                   = ggml_ext_cont(ctx, ggml_permute(ctx, out, 0, 1, 3, 2));                       // [N, OC, OH, OW]
     }
     if (scale != 1.f) {
         x = ggml_ext_scale(ctx, x, 1.f / scale);
@@ -831,22 +838,30 @@ ggml_tensor* ggml_ext_attention_prepared(ggml_context* ctx,
                                          ggml_tensor* v,
                                          int64_t n_head,
                                          int64_t N,
-                                         float kv_scale) {
+                                         float kv_scale,
+                                         ggml_tensor* mask) {
     GGML_ASSERT(q->type == GGML_TYPE_F32 && k->type == GGML_TYPE_F16 && v->type == GGML_TYPE_F16);
     const int64_t L_q    = q->ne[1];
     const int64_t d_head = v->ne[0];
-    if (backend == nullptr || d_head < 64) {
+    if (backend == nullptr || d_head < 64 || (mask != nullptr && mask->ne[3] != 1)) {
         return nullptr;
     }
 
     // Same expressions as ggml_ext_attention_ext so the softmax scale rounds identically.
     float scale = (1.0f / sqrt((float)d_head));
 
-    auto kqv = ggml_flash_attn_ext(ctx, q, k, v, nullptr, scale / kv_scale, 0, 0);
+    if (mask != nullptr) {
+        if (mask->ne[1] != L_q) {
+            mask = ggml_repeat(ctx, mask, ggml_new_tensor_4d(ctx, mask->type, mask->ne[0], L_q, mask->ne[2], mask->ne[3]));
+        }
+        mask = ggml_cast(ctx, mask, GGML_TYPE_F16);
+    }
+
+    auto kqv = ggml_flash_attn_ext(ctx, q, k, v, mask, scale / kv_scale, 0, 0);
     if (!ggml_backend_supports_op(backend, kqv)) {
         return nullptr;
     }
-    ggml_flash_attn_ext_set_prec(kqv, GGML_PREC_F32);
+    ggml_prec_set_acc(kqv, GGML_PREC_F32);
     if (kv_scale != 1.0f) {
         kqv = ggml_ext_scale(ctx, kqv, 1.0f / kv_scale);
     }

src/core/ggml_extend.h

@@ -236,6 +236,7 @@ ggml_tensor* ggml_ext_attention_ext(ggml_context* ctx,
 // already applied. Builds the same flash_attn_ext + output scaling as ggml_ext_attention_ext
 // (skip_reshape, unmasked), so the result is identical; returns nullptr when the backend lacks
 // flash attention for these inputs (the caller then uses ggml_ext_attention_ext).
+// mask (optional): as for ggml_ext_attention_ext's flash path (mask->ne[3] == 1).
 // return: [N, L_q, n_head * d_head]
 ggml_tensor* ggml_ext_attention_prepared(ggml_context* ctx,
                                          ggml_backend_t backend,
@@ -244,7 +245,8 @@ ggml_tensor* ggml_ext_attention_prepared(ggml_context* ctx,
                                          ggml_tensor* v,
                                          int64_t n_head,
                                          int64_t N,
-                                         float kv_scale);
+                                         float kv_scale,
+                                         ggml_tensor* mask = nullptr);
 
 ggml_tensor* ggml_ext_layer_norm(ggml_context* ctx,
                                  ggml_tensor* x,

src/model/common/rope.hpp

@@ -8,6 +8,7 @@
 #include <utility>
 #include <vector>
 #include "core/ggml_extend.h"
+#include "core/ggml_extend_backend.h"
 #include "core/ggml_runner.h"
 #include "core/util.h"
 
@@ -1125,6 +1126,74 @@ namespace Rope {
         return x_out;
     }
 
+    __STATIC_INLINE__ bool fused_rope_attention_enabled() {
+        static const bool enabled = [] {
+            const char* e = getenv("SD_FUSED_ROPE_ATTN");
+            return e == nullptr || atoi(e) != 0;
+        }();
+        return enabled;
+    }
+
+    // Fuses RoPE, relayout, kv scale and F16 cast into ggml_rope_pe_permute; must match the unfused
+    // flash path bit for bit (same rounding order as ggml_rope_apply -> ggml_scale -> ggml_cpy).
+    // CUDA / HIP only; SD_FUSED_ROPE_ATTN=0 disables; returns nullptr when it does not apply.
+    __STATIC_INLINE__ ggml_tensor* attention_fused_rope(GGMLRunnerContext* ctx,
+                                                        ggml_tensor* q,
+                                                        ggml_tensor* k,
+                                                        ggml_tensor* v,
+                                                        ggml_tensor* pe,
+                                                        ggml_tensor* mask,
+                                                        float kv_scale,
+                                                        bool rope_interleaved) {
+#ifndef SD_GGML_FUSED_ROPE_ATTN
+        return nullptr;
+#else
+        auto backend = ctx->backend;
+        // same attn_scale override as ggml_ext_attention_ext: keeps large k/v finite in F16
+        if (ctx->attn_scale > 0.f) {
+            kv_scale = ctx->attn_scale;
+        }
+        // interleaved only: non-interleaved ggml_rope_apply reads 2 of 4 table entries, so it would differ
+        if (!rope_interleaved || !ctx->flash_attn_enabled || ctx->sage_attn_enabled || backend == nullptr ||
+            sd_backend_is_cpu(backend) || !fused_rope_attention_enabled()) {
+            return nullptr;
+        }
+        const int64_t d_head = q->ne[0];
+        const int64_t L      = q->ne[2];
+        const int64_t N      = q->ne[3];
+        // ggml_ext_attention_ext pads heads under 64 channels for the CUDA kernel; keep that path for them
+        if (d_head < 64 || d_head % 2 != 0 || k->ne[0] != d_head || v->ne[0] != d_head || k->ne[2] != L || v->ne[2] != L ||
+            k->ne[3] != N || v->ne[3] != N || v->ne[1] != k->ne[1] || q->ne[1] % k->ne[1] != 0) {
+            return nullptr;
+        }
+        if (q->type != GGML_TYPE_F32 || k->type != GGML_TYPE_F32 || v->type != GGML_TYPE_F32 ||
+            q->nb[0] != sizeof(float) || k->nb[0] != sizeof(float) || v->nb[0] != sizeof(float)) {
+            return nullptr;
+        }
+        if (pe == nullptr || pe->type != GGML_TYPE_F32 || !ggml_is_contiguous(pe) || pe->ne[0] != 2 || pe->ne[1] != 2 ||
+            pe->ne[2] != d_head / 2 || pe->ne[3] != L) {
+            return nullptr;
+        }
+        if (mask != nullptr && mask->ne[3] != 1) {
+            return nullptr;
+        }
+        const int n_rot = static_cast<int>(d_head);
+        auto gctx       = ctx->ggml_ctx;
+        auto q_hm       = ggml_rope_pe_permute_set_interleaved(ggml_rope_pe_permute(gctx, q, pe, n_rot, 1.f, GGML_TYPE_F32), true);
+        auto k_hm       = ggml_rope_pe_permute_set_interleaved(ggml_rope_pe_permute(gctx, k, pe, n_rot, kv_scale, GGML_TYPE_F16), true);
+        auto v_hm       = ggml_rope_pe_permute(gctx, v, nullptr, 0, kv_scale, GGML_TYPE_F16);
+        if (!ggml_backend_supports_op(backend, q_hm) || !ggml_backend_supports_op(backend, k_hm) ||
+            !ggml_backend_supports_op(backend, v_hm)) {
+            return nullptr;
+        }
+        // [d_head, L, heads, N] -> [d_head, L, heads * N], as the unfused path
+        q_hm = ggml_reshape_3d(gctx, q_hm, d_head, L, q->ne[1] * N);
+        k_hm = ggml_reshape_3d(gctx, k_hm, d_head, L, k->ne[1] * N);
+        v_hm = ggml_reshape_3d(gctx, v_hm, d_head, L, v->ne[1] * N);
+        return ggml_ext_attention_prepared(gctx, backend, q_hm, k_hm, v_hm, q->ne[1], N, kv_scale, mask);
+#endif
+    }
+
     __STATIC_INLINE__ ggml_tensor* attention(GGMLRunnerContext* ctx,
                                              ggml_tensor* q,
                                              ggml_tensor* k,
@@ -1136,6 +1205,9 @@ namespace Rope {
         // q,k,v: [N, L, n_head, d_head]
         // pe: [L, d_head/2, 2, 2]
         // return: [N, L, n_head*d_head]
+        if (auto fused = attention_fused_rope(ctx, q, k, v, pe, mask, kv_scale, rope_interleaved)) {
+            return fused;
+        }
         int64_t n_head = q->ne[1];
 
         q = apply_rope(ctx, q, pe, rope_interleaved);  // [N*n_head, L, d_head]

src/model/vae/wan_vae.hpp

@@ -147,14 +147,15 @@ namespace WAN {
         using Conv2d::Conv2d;
 
         ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
+            // With T == 1 or C == 1 this swap moves no data, so ggml_ext_cont keeps the view (no copy).
             ggml_tensor* x_swapped = ggml_permute(ctx->ggml_ctx, x, 0, 1, 3, 2);
-            x_swapped              = ggml_cont(ctx->ggml_ctx, x_swapped);
+            x_swapped              = ggml_ext_cont(ctx->ggml_ctx, x_swapped);
 
             ggml_tensor* out = Conv2d::forward(ctx, x_swapped);
 
             ggml_tensor* out_swapped = ggml_permute(ctx->ggml_ctx, out, 0, 1, 3, 2);
 
-            out_swapped = ggml_cont(ctx->ggml_ctx, out_swapped);
+            out_swapped = ggml_ext_cont(ctx->ggml_ctx, out_swapped);
 
             return out_swapped;
         }