From 2754433123a63c526b4eac8ce8fdec3e05a57b34 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 11 Oct 2026 08:28:22 -0700 Subject: [PATCH 01/10] ggml-cuda: generic DiT and VAE fusions (interleaved rope_pe_permute, parallel group norm, channel rms norm, mul+add) Adds ggml patch 0007 on top of 0001-0006: - rope_pe_permute gains an interleaved pairing mode (Z-Image, Flux, Chroma rope) - rope_apply_f32 waits on the PDL predecessor (fixes nondeterministic output on sm90+) - parallel GROUP_NORM for groups >= 1024 elements, fused with the affine MUL/ADD and SILU - channel RMS norm fusion for CONT(PERMUTE) -> RMS_NORM -> MUL -> PERMUTE -> CONT (Wan/Qwen-Image VAE) - MUL -> ADD [-> ADD] broadcast fusion for modulation All fusions reproduce the stock rounding and reduction order, so outputs are bit-identical. Each has a kill switch (GGML_CUDA_GROUP_NORM_PAR, GGML_CUDA_NORM_CHANNEL, GGML_CUDA_FUSE_MUL_ADD). cmake detects the interleaved setter and defines SD_GGML_FUSED_ROPE_ATTN. --- cmake/ggml.cmake | 5 + .../0007-ggml-cuda-generic-dit-vae.patch | 1252 +++++++++++++++++ 2 files changed, 1257 insertions(+) create mode 100644 scripts/unsloth/ggml-patches/0007-ggml-cuda-generic-dit-vae.patch diff --git a/cmake/ggml.cmake b/cmake/ggml.cmake index 82d4814933..d3e952dca4 100644 --- a/cmake/ggml.cmake +++ b/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) diff --git a/scripts/unsloth/ggml-patches/0007-ggml-cuda-generic-dit-vae.patch b/scripts/unsloth/ggml-patches/0007-ggml-cuda-generic-dit-vae.patch new file mode 100644 index 0000000000..9741f98b52 --- /dev/null +++ b/scripts/unsloth/ggml-patches/0007-ggml-cuda-generic-dit-vae.patch @@ -0,0 +1,1252 @@ +diff --git a/include/ggml.h b/include/ggml.h +index 41ac28bb..fa83c5e6 100644 +--- a/include/ggml.h ++++ b/include/ggml.h +@@ -2576,6 +2576,14 @@ extern "C" { + // ggml_rope_pe_permute of ggml_rms_norm(a, eps) (no weight): every head row of `a` is + // RMS-normalised first, with the same arithmetic as GGML_OP_RMS_NORM on the same backend, so + // the result matches the two-op chain bit for bit where the backend supports it. ++ // ggml_rope_pe_permute (or _rms) with interleaved pairing: rotated element i < n_rot pairs ++ // a[2*j] with a[2*j + 1] (j = i / 2, r = i % 2): y = a[2*j]*pe[0, r, j, t] + a[2*j + 1]*pe[1, r, j, t], ++ // the pairing ggml_rope_apply(interleaved = true) uses, with the same rounding (bit-identical ++ // to ggml_rope_apply -> ggml_scale -> ggml_cpy). Returns `t` (a ggml_rope_pe_permute result). ++ GGML_API struct ggml_tensor * ggml_rope_pe_permute_set_interleaved( ++ struct ggml_tensor * t, ++ bool interleaved); ++ + GGML_API struct ggml_tensor * ggml_rope_pe_permute_rms( + struct ggml_context * ctx, + struct ggml_tensor * a, +diff --git a/src/ggml-cpu/ops.cpp b/src/ggml-cpu/ops.cpp +index bdd36aea..cadb25ac 100644 +--- a/src/ggml-cpu/ops.cpp ++++ b/src/ggml-cpu/ops.cpp +@@ -12397,6 +12397,7 @@ void ggml_compute_forward_rope_pe_permute( + const float scale = ggml_get_op_params_f32(dst, 1); + const bool norm = ggml_get_op_params_i32(dst, 2) != 0; + const float eps = ggml_get_op_params_f32(dst, 3); ++ const bool inter = ggml_get_op_params_i32(dst, 4) != 0; + const int half = n_rot / 2; + + const int64_t D = src0->ne[0]; +@@ -12436,14 +12437,14 @@ void ggml_compute_forward_rope_pe_permute( + if (n_rot > 0) { + const float * t = (const float *) pe->data + l * (int64_t) half * 4; + for (int i = 0; i < n_rot; ++i) { +- const int r = i / half; +- const int j = i - r * half; +- p0[i] = x[j] * t[4 * j + 2 * r + 0]; ++ const int r = inter ? i % 2 : i / half; ++ const int j = inter ? i / 2 : i - r * half; ++ p0[i] = x[inter ? 2 * j : j] * t[4 * j + 2 * r + 0]; + } + for (int i = 0; i < n_rot; ++i) { +- const int r = i / half; +- const int j = i - r * half; +- p1[i] = x[j + half] * t[4 * j + 2 * r + 1]; ++ const int r = inter ? i % 2 : i / half; ++ const int j = inter ? i / 2 : i - r * half; ++ p1[i] = x[inter ? 2 * j + 1 : j + half] * t[4 * j + 2 * r + 1]; + } + ggml_vec_add_f32(n_rot, y.data(), p0.data(), p1.data()); + } +diff --git a/src/ggml-cuda/ggml-cuda.cu b/src/ggml-cuda/ggml-cuda.cu +index 831d0de8..2df04786 100644 +--- a/src/ggml-cuda/ggml-cuda.cu ++++ b/src/ggml-cuda/ggml-cuda.cu +@@ -37,6 +37,9 @@ + #include "ggml-cuda/mmvq.cuh" + #include "ggml-cuda/moe-weighted-reduction.cuh" + #include "ggml-cuda/norm.cuh" ++#include "ggml-cuda/group-norm-par.cuh" ++#include "ggml-cuda/norm-channel.cuh" ++#include "ggml-cuda/mul-add-fused.cuh" + #include "ggml-cuda/opt-step-adamw.cuh" + #include "ggml-cuda/opt-step-sgd.cuh" + #include "ggml-cuda/out-prod.cuh" +@@ -4313,6 +4316,71 @@ static int ggml_cuda_try_fuse_cublas_epilogue(ggml_backend_cuda_context & ctx, g + return consumed; + } + ++// GROUP_NORM -> MUL(w) -> ADD(b) [-> SILU] with w, b per channel ([1, 1, C, 1], F32, contiguous): the ++// shape ggml_ext_group_norm-style VAE / U-Net code builds. Returns the number of extra nodes consumed. ++static int ggml_cuda_try_fuse_group_norm(ggml_backend_cuda_context & ctx, ggml_cgraph * cgraph, int i) { ++ ggml_tensor * gn = cgraph->nodes[i]; ++ if (gn->op != GGML_OP_GROUP_NORM || !ggml_cuda_group_norm_par_supported(gn)) { ++ return 0; ++ } ++ // the per-channel operands are usually reshaped 1-D weights; their RESHAPE nodes sit between the ++ // ops of the chain, so step over view nodes that do not read the chain itself ++ auto next_op = [&](int from, const ggml_tensor * prev) { ++ for (int j = from + 1; j < cgraph->n_nodes; ++j) { ++ const ggml_tensor * n = cgraph->nodes[j]; ++ const bool view = n->op == GGML_OP_RESHAPE || n->op == GGML_OP_VIEW || n->op == GGML_OP_PERMUTE || ++ n->op == GGML_OP_TRANSPOSE; ++ if (view && n->src[0] != prev && n->view_src != prev) { ++ continue; ++ } ++ return j; ++ } ++ return -1; ++ }; ++ const int i_mul = next_op(i, gn); ++ const int i_add = i_mul < 0 ? -1 : next_op(i_mul, cgraph->nodes[i_mul]); ++ if (i_add < 0) { ++ return 0; ++ } ++ int idx[4] = { i, i_mul, i_add, -1 }; ++ const enum ggml_op ops[4] = { GGML_OP_GROUP_NORM, GGML_OP_MUL, GGML_OP_ADD, GGML_OP_UNARY }; ++ if (!ggml_can_fuse_ext(cgraph, idx, ops, 3)) { ++ return 0; ++ } ++ ggml_tensor * mul = cgraph->nodes[i_mul]; ++ ggml_tensor * add = cgraph->nodes[i_add]; ++ const ggml_tensor * w = mul->src[0] == gn ? mul->src[1] : mul->src[0]; ++ const ggml_tensor * b = add->src[0] == mul ? add->src[1] : add->src[0]; ++ const int64_t C = gn->ne[2]; ++ auto per_channel = [&](const ggml_tensor * t) { ++ return t != nullptr && t->type == GGML_TYPE_F32 && ggml_is_contiguous(t) && t->ne[0] == 1 && t->ne[1] == 1 && ++ t->ne[2] == C && t->ne[3] == 1; ++ }; ++ if (!per_channel(w) || !per_channel(b) || mul->type != GGML_TYPE_F32 || add->type != GGML_TYPE_F32 || ++ !ggml_is_contiguous(mul) || !ggml_is_contiguous(add)) { ++ return 0; ++ } ++ int n_fused = 3; ++ const int i_u = next_op(i_add, add); ++ if (i_u > 0) { ++ const ggml_tensor * u = cgraph->nodes[i_u]; ++ idx[3] = i_u; ++ if (u->op == GGML_OP_UNARY && ggml_get_unary_op(u) == GGML_UNARY_OP_SILU && u->type == GGML_TYPE_F32 && ++ ggml_is_contiguous(u) && ggml_can_fuse_ext(cgraph, idx, ops, 4)) { ++ n_fused = 4; ++ } ++ } ++ ggml_tensor * out = cgraph->nodes[idx[n_fused - 1]]; ++ // every node in [i, last] is either fused here or a view the loop would skip anyway ++ const int last = idx[n_fused - 1]; ++ int out_nodes[] = { last }; ++ if (!ggml_cuda_check_fusion_memory_ranges(cgraph, i, last - i + 1, out_nodes, 1)) { ++ return 0; ++ } ++ ggml_cuda_op_group_norm_par(ctx, gn, w, b, n_fused == 4, out); ++ return last - i; ++} ++ + static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph * cgraph, int i) { + + static bool disable_fusion = getenv("GGML_CUDA_DISABLE_FUSION") != nullptr && std::atoi(getenv("GGML_CUDA_DISABLE_FUSION")); +@@ -4322,6 +4390,20 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph + + ggml_tensor * node = cgraph->nodes[i]; + ++ if (node->op == GGML_OP_GROUP_NORM) { ++ const int consumed = ggml_cuda_try_fuse_group_norm(*cuda_ctx, cgraph, i); ++ if (consumed > 0) { ++ return consumed; ++ } ++ } ++ ++ if (node->op == GGML_OP_CONT) { ++ const int consumed = ggml_cuda_try_fuse_rms_norm_channel(*cuda_ctx, cgraph, i); ++ if (consumed > 0) { ++ return consumed; ++ } ++ } ++ + if (node->op == GGML_OP_MUL) { + ggml_cuda_moe_weighted_reduction_match match; + if (ggml_cuda_match_moe_weighted_reduction(cgraph, i, match)) { +@@ -4334,6 +4416,13 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph + } + } + ++ if (node->op == GGML_OP_MUL) { ++ const int consumed = ggml_cuda_try_fuse_mul_add(*cuda_ctx, cgraph, i); ++ if (consumed > 0) { ++ return consumed; ++ } ++ } ++ + // gated_delta_net -> cpy: scatter recurrent-state snapshots into the cache + if (node->op == GGML_OP_GATED_DELTA_NET) { + ggml_cuda_gated_delta_net_fused_cache fused_state_cpy; +diff --git a/src/ggml-cuda/group-norm-par.cu b/src/ggml-cuda/group-norm-par.cu +new file mode 100644 +index 00000000..3e640e62 +--- /dev/null ++++ b/src/ggml-cuda/group-norm-par.cu +@@ -0,0 +1,173 @@ ++#include "group-norm-par.cuh" ++#include "unary.cuh" ++ ++// The stock group_norm_f32<1024> runs one 1024-thread block per group, so a VAE decode with 32 ++// groups keeps 32 SMs busy and walks the group three times (sum, centred sum of squares, scale), ++// then the affine and the SiLU are three more full passes. Here the same reductions run spread over ++// the GPU without changing a single rounding: ++// - thread t (0..1023) of the stock block sums x[g*G + t + 1024*k] in increasing k; each such ++// chain is computed whole by one thread here too, 32 consecutive t per warp, so the warp sums ++// (xor butterfly) are the same values the stock kernel writes to shared memory; ++// - the 32 warp sums are combined by the same warp butterfly (stock: warp lanes read shared[lane]); ++// - mean = sum / G and the centred pass use the same expressions, then scale = rsqrtf(var + eps); ++// - y = (x - mean) * scale, then * w, + b, then silu, each rounded separately as in the ++// unfused MUL / ADD / UNARY kernels. ++// Kernels: partial sums, partial centred squares, per-group stats, one elementwise apply pass. ++ ++static constexpr int GN_PAR_THREADS = 1024; // the stock block size whose arithmetic is reproduced ++static constexpr int GN_PAR_WARPS = GN_PAR_THREADS / WARP_SIZE; ++static constexpr int GN_PAR_WPB = 4; // warps per block in the reduction kernels ++ ++bool ggml_cuda_group_norm_par_enabled() { ++#if defined(GGML_USE_HIP) || defined(GGML_USE_MUSA) ++ return false; // warp size / reduction differ from what is reproduced here ++#else ++ static const bool enabled = [] { ++ const char * e = getenv("GGML_CUDA_GROUP_NORM_PAR"); ++ return e == nullptr || atoi(e) != 0; ++ }(); ++ return enabled && WARP_SIZE == 32; ++#endif ++} ++ ++static int64_t gn_par_group_size(const ggml_tensor * gn) { ++ const ggml_tensor * src0 = gn->src[0]; ++ const int num_groups = gn->op_params[0]; ++ return src0->ne[0] * src0->ne[1] * ((src0->ne[2] + num_groups - 1) / num_groups); ++} ++ ++bool ggml_cuda_group_norm_par_supported(const ggml_tensor * gn) { ++ if (!ggml_cuda_group_norm_par_enabled() || gn->op != GGML_OP_GROUP_NORM) { ++ return false; ++ } ++ const ggml_tensor * src0 = gn->src[0]; ++ if (src0->type != GGML_TYPE_F32 || gn->type != GGML_TYPE_F32 || !ggml_is_contiguous(src0) || !ggml_is_contiguous(gn)) { ++ return false; ++ } ++ const int64_t group_size = gn_par_group_size(gn); ++ // the stock kernel uses its 1024-thread variant only from 1024 elements per group; it indexes with int ++ return group_size >= GN_PAR_THREADS && ggml_nelements(src0) < INT_MAX; ++} ++ ++// one warp per (group, chain block of 32 threads): sum of x over the warp's 32 chains ++static __global__ void gn_par_sum(const float * __restrict__ x, float * __restrict__ part, const int group_size, ++ const int ne_elements) { ++ const int g = blockIdx.x; ++ const int w = blockIdx.y * GN_PAR_WPB + threadIdx.x / WARP_SIZE; ++ const int t = w * WARP_SIZE + threadIdx.x % WARP_SIZE; ++ const int start = g * group_size + t; ++ const int end = min(g * group_size + group_size, ne_elements); ++ float tmp = 0.0f; ++ for (int j = start; j < end; j += GN_PAR_THREADS) { ++ tmp += x[j]; ++ } ++ tmp = warp_reduce_sum(tmp); ++ if (threadIdx.x % WARP_SIZE == 0) { ++ part[g * GN_PAR_WARPS + w] = tmp; ++ } ++} ++ ++static __device__ __forceinline__ float gn_par_combine(const float * part, const int g) { ++ // stock: lanes read shared_vals[lane] (one value per warp) and reduce ++ return warp_reduce_sum(part[g * GN_PAR_WARPS + threadIdx.x % WARP_SIZE]); ++} ++ ++static __global__ void gn_par_sqdev(const float * __restrict__ x, const float * __restrict__ part, float * __restrict__ part2, ++ const int group_size, const int ne_elements) { ++ const int g = blockIdx.x; ++ const int w = blockIdx.y * GN_PAR_WPB + threadIdx.x / WARP_SIZE; ++ const int t = w * WARP_SIZE + threadIdx.x % WARP_SIZE; ++ const int start = g * group_size + t; ++ const int end = min(g * group_size + group_size, ne_elements); ++ ++ float tmp = gn_par_combine(part, g); ++ const float mean = tmp / group_size; ++ tmp = 0.0f; ++ for (int j = start; j < end; j += GN_PAR_THREADS) { ++ const float xi = x[j] - mean; ++ tmp += xi * xi; ++ } ++ tmp = warp_reduce_sum(tmp); ++ if (threadIdx.x % WARP_SIZE == 0) { ++ part2[g * GN_PAR_WARPS + w] = tmp; ++ } ++} ++ ++static __global__ void gn_par_stats(const float * __restrict__ part, const float * __restrict__ part2, float2 * __restrict__ stats, ++ const int group_size, const float eps) { ++ const int g = blockIdx.x; ++ float tmp = gn_par_combine(part, g); ++ const float mean = tmp / group_size; ++ tmp = gn_par_combine(part2, g); ++ const float variance = tmp / group_size; ++ const float scale = rsqrtf(variance + eps); ++ if (threadIdx.x == 0) { ++ stats[g] = make_float2(mean, scale); ++ } ++} ++ ++template ++static __global__ void gn_par_apply(const float * x, const float2 * __restrict__ stats, const float * __restrict__ w, ++ const float * __restrict__ b, float * dst, const int64_t ne_elements, ++ const int group_size, const int64_t plane, const int64_t channels) { ++ const int64_t stride = (int64_t) gridDim.x * blockDim.x; ++ for (int64_t j = (int64_t) blockIdx.x * blockDim.x + threadIdx.x; j < ne_elements; j += stride) { ++ const float2 st = stats[j / group_size]; ++ float y = __fmul_rn(__fsub_rn(x[j], st.x), st.y); ++ if constexpr (affine) { ++ const int64_t c = (j / plane) % channels; ++ y = __fmul_rn(y, w[c]); ++ y = __fadd_rn(y, b[c]); ++ } ++ if constexpr (silu) { ++ y = ggml_cuda_op_silu_single(y); ++ } ++ dst[j] = y; ++ } ++} ++ ++void ggml_cuda_op_group_norm_par(ggml_backend_cuda_context & ctx, const ggml_tensor * gn, const ggml_tensor * w, ++ const ggml_tensor * b, bool silu, ggml_tensor * out) { ++ const ggml_tensor * src0 = gn->src[0]; ++ const float * x = (const float *) src0->data; ++ float * dst = (float *) out->data; ++ cudaStream_t stream = ctx.stream(); ++ ++ const int num_groups = gn->op_params[0]; ++ float eps; ++ memcpy(&eps, gn->op_params + 1, sizeof(float)); ++ ++ const int group_size = (int) gn_par_group_size(gn); ++ const int ne_elements = (int) ggml_nelements(src0); ++ const int n_groups = num_groups * (int) src0->ne[3]; ++ ++ ggml_cuda_pool_alloc part(ctx.pool(), (size_t) n_groups * GN_PAR_WARPS); ++ ggml_cuda_pool_alloc part2(ctx.pool(), (size_t) n_groups * GN_PAR_WARPS); ++ ggml_cuda_pool_alloc stats(ctx.pool(), (size_t) n_groups); ++ ++ const dim3 rgrid(n_groups, GN_PAR_WARPS / GN_PAR_WPB, 1); ++ const dim3 rblock(GN_PAR_WPB * WARP_SIZE, 1, 1); ++ gn_par_sum<<>>(x, part.get(), group_size, ne_elements); ++ gn_par_sqdev<<>>(x, part.get(), part2.get(), group_size, ne_elements); ++ gn_par_stats<<>>(part.get(), part2.get(), stats.get(), group_size, eps); ++ ++ const int threads = 256; ++ const int64_t want = ((int64_t) ne_elements + threads - 1) / threads; ++ const int blocks = (int) std::min(want, (int64_t) ggml_cuda_info().devices[ctx.device].nsm * 32); ++ const int64_t plane = src0->ne[0] * src0->ne[1]; ++ const float * wd = w ? (const float *) w->data : nullptr; ++ const float * bd = b ? (const float *) b->data : nullptr; ++ if (w) { ++ if (silu) { ++ gn_par_apply<<>>(x, stats.get(), wd, bd, dst, ne_elements, group_size, plane, src0->ne[2]); ++ } else { ++ gn_par_apply<<>>(x, stats.get(), wd, bd, dst, ne_elements, group_size, plane, src0->ne[2]); ++ } ++ } else { ++ if (silu) { ++ gn_par_apply<<>>(x, stats.get(), wd, bd, dst, ne_elements, group_size, plane, src0->ne[2]); ++ } else { ++ gn_par_apply<<>>(x, stats.get(), wd, bd, dst, ne_elements, group_size, plane, src0->ne[2]); ++ } ++ } ++} +diff --git a/src/ggml-cuda/group-norm-par.cuh b/src/ggml-cuda/group-norm-par.cuh +new file mode 100644 +index 00000000..eb6298ea +--- /dev/null ++++ b/src/ggml-cuda/group-norm-par.cuh +@@ -0,0 +1,15 @@ ++#pragma once ++#include "common.cuh" ++ ++// Parallel GROUP_NORM with the stock kernel's exact arithmetic, optionally fused with the ++// per-channel affine (MUL by w, ADD of b, both [1, 1, C, 1]) and a SiLU that follow it in the graph. ++// Returns false (nothing launched) when the inputs are outside what it handles; the caller then ++// runs the stock ops. GGML_CUDA_GROUP_NORM_PAR=0 turns it off. ++bool ggml_cuda_group_norm_par_enabled(); ++ ++bool ggml_cuda_group_norm_par_supported(const ggml_tensor * gn); ++ ++// gn: the GROUP_NORM node; w / b: affine operands or nullptr; silu: apply SiLU; out: the tensor ++// that receives the result (the last fused node). ++void ggml_cuda_op_group_norm_par(ggml_backend_cuda_context & ctx, const ggml_tensor * gn, const ggml_tensor * w, ++ const ggml_tensor * b, bool silu, ggml_tensor * out); +diff --git a/src/ggml-cuda/mul-add-fused.cu b/src/ggml-cuda/mul-add-fused.cu +new file mode 100644 +index 00000000..ff4e1dd0 +--- /dev/null ++++ b/src/ggml-cuda/mul-add-fused.cu +@@ -0,0 +1,171 @@ ++#include "mul-add-fused.cuh" ++#include "ggml-impl.h" ++ ++static bool mul_add_fusion_enabled() { ++ static const bool enabled = [] { ++ const char * e = getenv("GGML_CUDA_FUSE_MUL_ADD"); ++ return e == nullptr || atoi(e) != 0; ++ }(); ++ return enabled; ++} ++ ++struct mad_operand { ++ const float * p; ++ int64_t ne[4]; ++ int64_t nb[4]; // in elements ++}; ++ ++static __device__ __forceinline__ const float * mad_row(const mad_operand & o, const int64_t i1, const int64_t i2, const int64_t i3) { ++ return o.p + (o.ne[1] == 1 ? 0 : i1 * o.nb[1]) + (o.ne[2] == 1 ? 0 : i2 * o.nb[2]) + (o.ne[3] == 1 ? 0 : i3 * o.nb[3]); ++} ++ ++static __device__ __forceinline__ float mad_at(const float * row, const mad_operand & o, const int64_t i0) { ++ return row[o.ne[0] == 1 ? 0 : i0 * o.nb[0]]; ++} ++ ++// one block per output row; y = (a*b) + c [+ d], b, c, d broadcast like binbcast (dims of size 1) ++template ++static __global__ void mul_add_fused_f32(const mad_operand a, const mad_operand b, const mad_operand c, const mad_operand d, ++ float * __restrict__ dst, const int64_t ne0, const int64_t ne1, const int64_t ne2) { ++ const int64_t r = blockIdx.x; ++ const int64_t i1 = r % ne1; ++ const int64_t i2 = (r / ne1) % ne2; ++ const int64_t i3 = r / (ne1 * ne2); ++ const float * ra = mad_row(a, i1, i2, i3); ++ const float * rb = mad_row(b, i1, i2, i3); ++ const float * rc = mad_row(c, i1, i2, i3); ++ const float * rd = has_d ? mad_row(d, i1, i2, i3) : nullptr; ++ float * out = dst + r * ne0; ++ for (int64_t i0 = threadIdx.x; i0 < ne0; i0 += blockDim.x) { ++ float y = __fmul_rn(mad_at(ra, a, i0), mad_at(rb, b, i0)); ++ y = __fadd_rn(y, mad_at(rc, c, i0)); ++ if constexpr (has_d) { ++ y = __fadd_rn(y, mad_at(rd, d, i0)); ++ } ++ out[i0] = y; ++ } ++} ++ ++static bool mad_ok_operand(const ggml_tensor * t, const ggml_tensor * out) { ++ if (t == nullptr || t->type != GGML_TYPE_F32 || t->nb[0] != sizeof(float)) { ++ return false; ++ } ++ for (int k = 0; k < 4; ++k) { ++ if (t->ne[k] != out->ne[k] && t->ne[k] != 1) { ++ return false; ++ } ++ if (t->nb[k] % sizeof(float) != 0) { ++ return false; ++ } ++ } ++ return true; ++} ++ ++static mad_operand mad_make(const ggml_tensor * t) { ++ mad_operand o; ++ o.p = (const float *) t->data; ++ for (int k = 0; k < 4; ++k) { ++ o.ne[k] = t->ne[k]; ++ o.nb[k] = t->nb[k] / (int64_t) sizeof(float); ++ } ++ return o; ++} ++ ++static bool mad_is_skippable_view(const ggml_tensor * n, const ggml_tensor * prev) { ++ const bool view = n->op == GGML_OP_RESHAPE || n->op == GGML_OP_VIEW || n->op == GGML_OP_PERMUTE || n->op == GGML_OP_TRANSPOSE; ++ return view && n->src[0] != prev && n->view_src != prev; ++} ++ ++int ggml_cuda_try_fuse_mul_add(ggml_backend_cuda_context & ctx, ggml_cgraph * cgraph, int i) { ++ if (!mul_add_fusion_enabled()) { ++ return 0; ++ } ++ ggml_tensor * mul = cgraph->nodes[i]; ++ if (mul->op != GGML_OP_MUL || mul->type != GGML_TYPE_F32 || !ggml_is_contiguous(mul) || !ggml_node_has_n_uses(cgraph, i, 1)) { ++ return 0; ++ } ++ auto next_op = [&](int from) { ++ for (int j = from + 1; j < cgraph->n_nodes; ++j) { ++ if (!mad_is_skippable_view(cgraph->nodes[j], cgraph->nodes[from])) { ++ return j; ++ } ++ } ++ return -1; ++ }; ++ // the full-shape factor goes first so the output shape is the MUL's shape ++ const ggml_tensor * a = mul->src[0]; ++ const ggml_tensor * b = mul->src[1]; ++ if (!ggml_are_same_shape(a, mul)) { ++ std::swap(a, b); ++ } ++ if (!ggml_are_same_shape(a, mul) || !mad_ok_operand(a, mul) || !mad_ok_operand(b, mul)) { ++ return 0; ++ } ++ const int i_add = next_op(i); ++ if (i_add < 0) { ++ return 0; ++ } ++ ggml_tensor * add = cgraph->nodes[i_add]; ++ if (add->op != GGML_OP_ADD || add->type != GGML_TYPE_F32 || !ggml_is_contiguous(add) || !ggml_are_same_shape(add, mul) || ++ (add->src[0] != mul && add->src[1] != mul)) { ++ return 0; ++ } ++ const ggml_tensor * c = add->src[0] == mul ? add->src[1] : add->src[0]; ++ if (c == mul || !mad_ok_operand(c, add)) { ++ return 0; ++ } ++ int last = i_add; ++ const ggml_tensor * d = nullptr; ++ if (ggml_node_has_n_uses(cgraph, i_add, 1)) { ++ const int i_add2 = next_op(i_add); ++ if (i_add2 > 0) { ++ ggml_tensor * add2 = cgraph->nodes[i_add2]; ++ if (add2->op == GGML_OP_ADD && add2->type == GGML_TYPE_F32 && ggml_is_contiguous(add2) && ggml_are_same_shape(add2, add) && ++ (add2->src[0] == add || add2->src[1] == add)) { ++ const ggml_tensor * dd = add2->src[0] == add ? add2->src[1] : add2->src[0]; ++ if (dd != add && dd != mul && mad_ok_operand(dd, add2)) { ++ d = dd; ++ last = i_add2; ++ } ++ } ++ } ++ } ++ for (int j = i + 1; j < last; ++j) { ++ const ggml_tensor * n = cgraph->nodes[j]; ++ if (j != i_add && !(n->op == GGML_OP_RESHAPE || n->op == GGML_OP_VIEW || n->op == GGML_OP_PERMUTE || n->op == GGML_OP_TRANSPOSE)) { ++ return 0; ++ } ++ } ++ ggml_tensor * out = cgraph->nodes[last]; ++ // the output may only alias the elementwise inputs at the same element (in-place adds) ++ auto overlaps = [&](const ggml_tensor * t) { ++ if (t == nullptr || t->data == nullptr) { ++ return false; ++ } ++ const char * o0 = (const char *) out->data; ++ const char * o1 = o0 + ggml_nbytes(out); ++ const char * t0 = (const char *) t->data; ++ const char * t1 = t0 + ggml_nbytes(t); ++ return t0 < o1 && o0 < t1; ++ }; ++ auto same_layout = [&](const ggml_tensor * t) { ++ return t->data == out->data && ggml_are_same_shape(t, out) && ggml_is_contiguous(t); ++ }; ++ for (const ggml_tensor * t : { a, b, c, d }) { ++ if (overlaps(t) && !same_layout(t)) { ++ return 0; ++ } ++ } ++ const int64_t nrows = ggml_nrows(out); ++ if (nrows > INT_MAX) { ++ return 0; ++ } ++ const int threads = (int) std::min(256, std::max(32, (out->ne[0] + 31) / 32 * 32)); ++ const mad_operand oa = mad_make(a), ob = mad_make(b), oc = mad_make(c), od = d ? mad_make(d) : oc; ++ if (d) { ++ mul_add_fused_f32<<<(unsigned int) nrows, threads, 0, ctx.stream()>>>(oa, ob, oc, od, (float *) out->data, out->ne[0], out->ne[1], out->ne[2]); ++ } else { ++ mul_add_fused_f32<<<(unsigned int) nrows, threads, 0, ctx.stream()>>>(oa, ob, oc, od, (float *) out->data, out->ne[0], out->ne[1], out->ne[2]); ++ } ++ return last - i; ++} +diff --git a/src/ggml-cuda/mul-add-fused.cuh b/src/ggml-cuda/mul-add-fused.cuh +new file mode 100644 +index 00000000..8d9752ca +--- /dev/null ++++ b/src/ggml-cuda/mul-add-fused.cuh +@@ -0,0 +1,8 @@ ++#pragma once ++#include "common.cuh" ++ ++// MUL -> ADD [-> ADD] chains of F32 elementwise ops with broadcast operands (adaLN modulate ++// x + x*scale (+ shift), gated residual r + y*gate): one pass instead of two or three, each ++// product and sum still rounded separately (bit-identical to the unfused MUL / ADD kernels). ++// Returns the number of extra nodes consumed. GGML_CUDA_FUSE_MUL_ADD=0 turns it off. ++int ggml_cuda_try_fuse_mul_add(ggml_backend_cuda_context & ctx, ggml_cgraph * cgraph, int i); +diff --git a/src/ggml-cuda/norm-channel.cu b/src/ggml-cuda/norm-channel.cu +new file mode 100644 +index 00000000..84a33c9c +--- /dev/null ++++ b/src/ggml-cuda/norm-channel.cu +@@ -0,0 +1,228 @@ ++#include "norm-channel.cuh" ++#include "ggml-impl.h" ++ ++// The unfused RMS_NORM over rows of C < 1024 channels runs rms_norm_f32<256> (1024 threads from 1024 ++// channels; that tree is reproduced too) (or the narrow-block ++// variant, documented bit-identical to it): thread t sums squares of columns t, t+256, ... (in that ++// order), each warp reduces its 32 partials with the xor butterfly, and the 8 warp sums plus 24 zero ++// partials are reduced by the same butterfly. Here one thread owns one spatial position (so loads ++// across a warp are contiguous) and evaluates exactly that reduction tree in registers. ++ ++static bool norm_channel_enabled() { ++#if defined(GGML_USE_HIP) || defined(GGML_USE_MUSA) ++ return false; ++#else ++ static const bool enabled = [] { ++ const char * e = getenv("GGML_CUDA_NORM_CHANNEL"); ++ return e == nullptr || atoi(e) != 0; ++ }(); ++ return enabled && WARP_SIZE == 32; ++#endif ++} ++ ++static __device__ __forceinline__ float nc_butterfly32(float * v) { ++ // value every lane holds after warp_reduce_sum: lane l adds lane l^off at each step ++#pragma unroll ++ for (int off = 16; off > 0; off >>= 1) { ++#pragma unroll ++ for (int l = 0; l < off; ++l) { ++ v[l] = v[l] + v[l + off]; ++ } ++ } ++ return v[0]; ++} ++ ++// BLOCK: the thread count of the rms_norm_f32 launch this reproduces (256 below 1024 columns, else 1024) ++template ++static __global__ void rms_norm_channel_f32(const float * x, const float * __restrict__ w, float * dst, ++ const int C, const int64_t P, const int64_t inner, const float eps) { ++ const int64_t q = (int64_t) blockIdx.x * blockDim.x + threadIdx.x; ++ if (q >= P) { ++ return; ++ } ++ // position q = (outer index, inner index); its channel c sits at base + c*inner ++ const int64_t p = (q / inner) * inner * C + q % inner; ++ const float * xp = x + p; ++ float wsum[32]; ++#pragma unroll ++ for (int i = 0; i < 32; ++i) { ++ wsum[i] = 0.0f; ++ } ++#pragma unroll 1 ++ for (int wi = 0; wi < BLOCK / 32; ++wi) { ++ float lane[32]; ++#pragma unroll ++ for (int l = 0; l < 32; ++l) { ++ float tmp = 0.0f; ++ for (int col = wi * 32 + l; col < C; col += BLOCK) { ++ const float xi = xp[(int64_t) col * inner]; ++ tmp += xi * xi; ++ } ++ lane[l] = tmp; ++ } ++ wsum[wi] = nc_butterfly32(lane); ++ } ++ const float tmp = nc_butterfly32(wsum); ++ const float mean = tmp / C; ++ const float scale = rsqrtf(mean + eps); ++ for (int col = 0; col < C; ++col) { ++ dst[p + (int64_t) col * inner] = scale * xp[(int64_t) col * inner] * w[col]; ++ } ++} ++ ++static bool nc_debug() { ++ static const bool on = getenv("GGML_CUDA_NORM_CHANNEL_DEBUG") != nullptr; ++ return on; ++} ++#define NC_REJECT(why) do { if (nc_debug()) { GGML_LOG_INFO("norm-channel: node %d (%s) rejected: %s\n", i, cont1->name, why); } return 0; } while (0) ++ ++static bool is_view_op(const ggml_tensor * t) { ++ return t->op == GGML_OP_PERMUTE || t->op == GGML_OP_VIEW || t->op == GGML_OP_RESHAPE || t->op == GGML_OP_TRANSPOSE; ++} ++ ++int ggml_cuda_try_fuse_rms_norm_channel(ggml_backend_cuda_context & ctx, ggml_cgraph * cgraph, int i) { ++ if (!norm_channel_enabled()) { ++ return 0; ++ } ++ // node i: CONT whose source is a permuted view of a contiguous F32 x: cont[c, a, b, d] = x[a, b, d, c] ++ ggml_tensor * cont1 = cgraph->nodes[i]; ++ if (cont1->op != GGML_OP_CONT || cont1->type != GGML_TYPE_F32) { ++ return 0; ++ } ++ const ggml_tensor * pv = cont1->src[0]; ++ if (pv == nullptr || pv->op != GGML_OP_PERMUTE || pv->type != GGML_TYPE_F32) { ++ return 0; ++ } ++ const ggml_tensor * x = pv->src[0]; ++ if (x == nullptr || x->type != GGML_TYPE_F32 || !ggml_is_contiguous(x) || pv->view_offs != 0 || pv->data != x->data) { ++ NC_REJECT("check 1"); ++ } ++ // pv dim 0 (the normalised axis) must be one axis k of x; the other axes may come in any order ++ const int64_t C = pv->ne[0]; ++ int k = -1; ++ for (int m = 0; m < GGML_MAX_DIMS; ++m) { ++ if (x->ne[m] == C && x->nb[m] == pv->nb[0] && (k < 0 || x->ne[m] > 1)) { ++ k = m; ++ } ++ } ++ // map every pv axis j > 0 of size > 1 to the x axis it reads ++ int pv_from_x[GGML_MAX_DIMS] = { -1, -1, -1, -1 }; ++ pv_from_x[0] = k; ++ bool layout_ok = k >= 0 && C >= 1 && ggml_nelements(pv) == ggml_nelements(x); ++ for (int j = 1; j < GGML_MAX_DIMS && layout_ok; ++j) { ++ if (pv->ne[j] == 1) { ++ continue; ++ } ++ for (int m = 0; m < GGML_MAX_DIMS; ++m) { ++ if (m != k && x->ne[m] == pv->ne[j] && x->nb[m] == pv->nb[j]) { ++ pv_from_x[j] = m; ++ } ++ } ++ layout_ok = pv_from_x[j] >= 0; ++ } ++ if (!layout_ok) { ++ NC_REJECT("check 2"); ++ } ++ const int64_t inner = x->nb[k] / (int64_t) sizeof(float); // elements below the channel axis ++ const int64_t P = ggml_nelements(x) / C; // positions normalised independently ++ // then RMS_NORM(cont1) and MUL(rms, w[C]) in that order (views in between are skipped) ++ auto next_op = [&](int from) { ++ for (int j = from + 1; j < cgraph->n_nodes; ++j) { ++ if (!is_view_op(cgraph->nodes[j]) || cgraph->nodes[j]->src[0] == cgraph->nodes[from]) { ++ return j; ++ } ++ } ++ return -1; ++ }; ++ const int i_rms = next_op(i); ++ if (i_rms < 0) { ++ NC_REJECT("check 3"); ++ } ++ ggml_tensor * rms = cgraph->nodes[i_rms]; ++ const int i_mul = next_op(i_rms); ++ if (i_mul < 0) { ++ NC_REJECT("check 4"); ++ } ++ ggml_tensor * mul = cgraph->nodes[i_mul]; ++ if (rms->op != GGML_OP_RMS_NORM) { ++ if (nc_debug()) { GGML_LOG_INFO("norm-channel: C=%lld next op %s (%s)\n", (long long) C, ggml_op_desc(rms), rms->name); } ++ NC_REJECT("check 5a"); ++ } ++ if (rms->src[0] != cont1) { ++ NC_REJECT("check 5b"); ++ } ++ if (mul->op != GGML_OP_MUL || mul->src[0] != rms) { ++ if (nc_debug()) { GGML_LOG_INFO("norm-channel: C=%lld after rms %s (%s)\n", (long long) C, ggml_op_desc(mul), mul->name); } ++ NC_REJECT("check 5c"); ++ } ++ if (!ggml_node_has_n_uses(cgraph, i, 1) || !ggml_node_has_n_uses(cgraph, i_rms, 1) || !ggml_node_has_n_uses(cgraph, i_mul, 1)) { ++ NC_REJECT("check 5d"); ++ } ++ const ggml_tensor * w = mul->src[1]; ++ if (w == nullptr || w->type != GGML_TYPE_F32 || !ggml_is_contiguous(w) || w->ne[0] != C || ggml_nelements(w) != C || ++ mul->type != GGML_TYPE_F32 || !ggml_are_same_shape(mul, rms)) { ++ NC_REJECT("check 6"); ++ } ++ // then PERMUTE(mul) back to x's layout and its CONT ++ const int i_pv2 = i_mul + 1; ++ if (i_pv2 + 1 >= cgraph->n_nodes) { ++ NC_REJECT("check 7"); ++ } ++ ggml_tensor * pv2 = cgraph->nodes[i_pv2]; ++ ggml_tensor * cont2 = cgraph->nodes[i_pv2 + 1]; ++ bool back_ok = pv2->op == GGML_OP_PERMUTE && pv2->src[0] == mul && cont2->op == GGML_OP_CONT && cont2->src[0] == pv2 && ++ cont2->type == GGML_TYPE_F32 && ggml_are_same_shape(cont2, x) && ++ ggml_is_contiguous(cont2) && ggml_is_contiguous(mul); ++ // the graph does not count uses of views, so check by hand that cont2 is the only reader of pv2 ++ for (int j = 0; j < cgraph->n_nodes && back_ok; ++j) { ++ const ggml_tensor * n = cgraph->nodes[j]; ++ for (int s = 0; s < GGML_MAX_SRC && back_ok; ++s) { ++ back_ok = n == cont2 || n->src[s] != pv2; ++ } ++ } ++ // pv2 must be the inverse permutation: x axis m reads the mul axis that was built from x axis m ++ for (int j = 0; j < GGML_MAX_DIMS && back_ok; ++j) { ++ if (pv->ne[j] == 1 && j != 0) { ++ continue; ++ } ++ const int m = pv_from_x[j]; ++ back_ok = x->ne[m] == 1 || pv2->nb[m] == mul->nb[j]; ++ } ++ if (!back_ok) { ++ NC_REJECT("check 8"); ++ } ++ // every node from i to the second CONT is part of the pattern or a view the loop skips ++ for (int j = i + 1; j < i_pv2 + 1; ++j) { ++ const ggml_tensor * n = cgraph->nodes[j]; ++ if (j != i_rms && j != i_mul && j != i_pv2 && !is_view_op(n)) { ++ NC_REJECT("check 9"); ++ } ++ } ++ // The result may reuse x's memory (x is often dead after this norm): with the same layout every ++ // thread reads all channels of its position before writing them, so only partial overlap is unsafe. ++ const char * x0 = (const char *) x->data; ++ const char * d0 = (const char *) cont2->data; ++ const bool partial_overlap = d0 != x0 && d0 < x0 + ggml_nbytes(x) && x0 < d0 + ggml_nbytes(cont2); ++ if (P > INT64_C(1) << 40 || partial_overlap) { ++ NC_REJECT("check 10"); ++ } ++ float eps; ++ memcpy(&eps, rms->op_params, sizeof(float)); ++ ++ const int threads = 128; ++ const int64_t blocks = (P + threads - 1) / threads; ++ if (blocks > INT_MAX) { ++ NC_REJECT("check 11"); ++ } ++ if (C < 1024) { ++ rms_norm_channel_f32<256><<<(unsigned int) blocks, threads, 0, ctx.stream()>>>( ++ (const float *) x->data, (const float *) w->data, (float *) cont2->data, (int) C, P, inner, eps); ++ } else { ++ rms_norm_channel_f32<1024><<<(unsigned int) blocks, threads, 0, ctx.stream()>>>( ++ (const float *) x->data, (const float *) w->data, (float *) cont2->data, (int) C, P, inner, eps); ++ } ++ if (nc_debug()) { ++ GGML_LOG_INFO("norm-channel: fused C=%lld P=%lld\n", (long long) C, (long long) P); ++ } ++ return i_pv2 + 1 - i; ++} +diff --git a/src/ggml-cuda/norm-channel.cuh b/src/ggml-cuda/norm-channel.cuh +new file mode 100644 +index 00000000..da57e590 +--- /dev/null ++++ b/src/ggml-cuda/norm-channel.cuh +@@ -0,0 +1,9 @@ ++#pragma once ++#include "common.cuh" ++ ++// CONT(PERMUTE(x)) -> RMS_NORM -> MUL(w) -> CONT(PERMUTE back): an RMS norm over the outermost ++// (channel) axis of a contiguous x, the way channel-first video VAEs (Wan) build it. Computes the ++// result straight from x into the final tensor, without the two relayout copies, bit-identical ++// to the unfused ops. Returns the number of extra nodes consumed (0: not this pattern / unsupported). ++// GGML_CUDA_NORM_CHANNEL=0 turns it off. ++int ggml_cuda_try_fuse_rms_norm_channel(ggml_backend_cuda_context & ctx, ggml_cgraph * cgraph, int i); +diff --git a/src/ggml-cuda/norm.cu b/src/ggml-cuda/norm.cu +index 9c9385f8..ae2def32 100644 +--- a/src/ggml-cuda/norm.cu ++++ b/src/ggml-cuda/norm.cu +@@ -1,4 +1,5 @@ + #include "norm.cuh" ++#include "group-norm-par.cuh" + #include + + template +@@ -534,6 +535,11 @@ void ggml_cuda_op_group_norm(ggml_backend_cuda_context & ctx, ggml_tensor * dst) + memcpy(&eps, dst->op_params + 1, sizeof(float)); + GGML_ASSERT(eps >= 0.0f); + ++ if (ggml_cuda_group_norm_par_supported(dst)) { ++ ggml_cuda_op_group_norm_par(ctx, dst, nullptr, nullptr, false, dst); ++ return; ++ } ++ + int group_size = src0->ne[0] * src0->ne[1] * ((src0->ne[2] + num_groups - 1) / num_groups); + group_norm_f32_cuda(src0_d, dst_d, num_groups * src0->ne[3], eps, group_size, ggml_nelements(src0), stream); + } +diff --git a/src/ggml-cuda/rope-pe.cu b/src/ggml-cuda/rope-pe.cu +index b6dd7a69..2fe9a599 100644 +--- a/src/ggml-cuda/rope-pe.cu ++++ b/src/ggml-cuda/rope-pe.cu +@@ -6,7 +6,7 @@ + // Rounding matches the unfused graph exactly: __fmul_rn/__fadd_rn keep the two products and the + // sum as separately rounded operations (no FMA contraction), the scale is the same `scale*y + bias` + // expression the SCALE kernel evaluates, and the conversion uses the same ggml_cuda_cast as CPY. +-template ++template + static __global__ void rope_pe_permute_kernel( + const char * __restrict__ x, const float * __restrict__ pe, dst_t * __restrict__ dst, + const int D, const int H, const int64_t L, const int n_rot, +@@ -21,9 +21,13 @@ static __global__ void rope_pe_permute_kernel( + int j = i; + float c0 = 0.0f; + float c1 = 0.0f; ++ int ja = i; ++ int jb = i; + if (rot) { +- const int r = i >= half; +- j = i - r * half; ++ const int r = interleaved ? (i & 1) : (i >= half); ++ j = interleaved ? (i >> 1) : i - r * half; ++ ja = interleaved ? 2 * j : j; ++ jb = interleaved ? 2 * j + 1 : j + half; + const float * t = pe + (l * half + j) * 4 + 2 * r; + c0 = t[0]; + c1 = t[1]; +@@ -32,7 +36,7 @@ static __global__ void rope_pe_permute_kernel( + const float * xr = (const float *) (x + h * nb01 + l * nb02 + n * nb03); + float y; + if (rot) { +- y = __fadd_rn(__fmul_rn(xr[j], c0), __fmul_rn(xr[j + half], c1)); ++ y = __fadd_rn(__fmul_rn(xr[ja], c0), __fmul_rn(xr[jb], c1)); + } else { + y = xr[i]; + } +@@ -48,7 +52,7 @@ static __global__ void rope_pe_permute_kernel( + // per (token, batch), the size and reduction rms_norm_f32_cuda's small-row path uses for this row + // length, so each normalised value is bit-identical to GGML_OP_RMS_NORM's; the normalised row is + // staged in shared memory for the RoPE partner reads. +-template ++template + static __global__ void rope_pe_permute_rms_kernel( + const char * __restrict__ x, const float * __restrict__ pe, dst_t * __restrict__ dst, + const int D, const int H, const int64_t L, const int n_rot, +@@ -64,11 +68,15 @@ static __global__ void rope_pe_permute_rms_kernel( + + const bool rot = i < n_rot; + int j = i; ++ int ja = i; ++ int jb = i; + float c0 = 0.0f; + float c1 = 0.0f; + if (rot) { +- const int r = i >= half; +- j = i - r * half; ++ const int r = interleaved ? (i & 1) : (i >= half); ++ j = interleaved ? (i >> 1) : i - r * half; ++ ja = interleaved ? 2 * j : j; ++ jb = interleaved ? 2 * j + 1 : j + half; + const float * t = pe + (l * half + j) * 4 + 2 * r; + c0 = t[0]; + c1 = t[1]; +@@ -91,7 +99,7 @@ static __global__ void rope_pe_permute_rms_kernel( + if (i < D) { + float y; + if (rot) { +- y = __fadd_rn(__fmul_rn(xs[j], c0), __fmul_rn(xs[j + half], c1)); ++ y = __fadd_rn(__fmul_rn(xs[ja], c0), __fmul_rn(xs[jb], c1)); + } else { + y = xs[i]; + } +@@ -117,7 +125,7 @@ static bool rope_pe_rms_small_rows_enabled() { + #endif + } + +-template ++template + static void rope_pe_permute_rms_launch(const ggml_tensor * src0, const float * pe, dst_t * dst, const int n_rot, + const float scale, const float eps, cudaStream_t stream) { + const int D = (int) src0->ne[0]; +@@ -128,29 +136,29 @@ static void rope_pe_permute_rms_launch(const ggml_tensor * src0, const float * p + const dim3 block(block_size, 1, 1); + const char * x = (const char *) src0->data; + if (scale != 1.0f) { +- rope_pe_permute_rms_kernel<<>>( ++ rope_pe_permute_rms_kernel<<>>( + x, pe, dst, D, H, L, n_rot, src0->nb[1], src0->nb[2], src0->nb[3], scale, 0.0f, eps); + } else { +- rope_pe_permute_rms_kernel<<>>( ++ rope_pe_permute_rms_kernel<<>>( + x, pe, dst, D, H, L, n_rot, src0->nb[1], src0->nb[2], src0->nb[3], scale, 0.0f, eps); + } + } + +-template ++template + static void rope_pe_permute_rms_cuda(const ggml_tensor * src0, const float * pe, dst_t * dst, const int n_rot, + const float scale, const float eps, cudaStream_t stream) { + const int D = (int) src0->ne[0]; + GGML_ASSERT(D <= 128); + if (D <= 32) { +- rope_pe_permute_rms_launch(src0, pe, dst, n_rot, scale, eps, stream); ++ rope_pe_permute_rms_launch(src0, pe, dst, n_rot, scale, eps, stream); + } else if (D <= 64) { +- rope_pe_permute_rms_launch(src0, pe, dst, n_rot, scale, eps, stream); ++ rope_pe_permute_rms_launch(src0, pe, dst, n_rot, scale, eps, stream); + } else { +- rope_pe_permute_rms_launch(src0, pe, dst, n_rot, scale, eps, stream); ++ rope_pe_permute_rms_launch(src0, pe, dst, n_rot, scale, eps, stream); + } + } + +-template ++template + static void rope_pe_permute_cuda(const ggml_tensor * src0, const float * pe, dst_t * dst, const int n_rot, + const float scale, cudaStream_t stream) { + const int D = (int) src0->ne[0]; +@@ -162,10 +170,10 @@ static void rope_pe_permute_cuda(const ggml_tensor * src0, const float * pe, dst + const dim3 block(D >= 128 ? 128 : ((D + 31) / 32) * 32, 1, 1); + const char * x = (const char *) src0->data; + if (scale != 1.0f) { +- rope_pe_permute_kernel<<>>( ++ rope_pe_permute_kernel<<>>( + x, pe, dst, D, H, L, n_rot, src0->nb[1], src0->nb[2], src0->nb[3], scale, 0.0f); + } else { +- rope_pe_permute_kernel<<>>( ++ rope_pe_permute_kernel<<>>( + x, pe, dst, D, H, L, n_rot, src0->nb[1], src0->nb[2], src0->nb[3], scale, 0.0f); + } + } +@@ -184,22 +192,41 @@ void ggml_cuda_op_rope_pe_permute(ggml_backend_cuda_context & ctx, ggml_tensor * + cudaStream_t stream = ctx.stream(); + const float * pe_d = pe ? (const float *) pe->data : nullptr; + ++ // interleaved pairing (2j, 2j+1) only matters when something rotates ++ const bool interleaved = n_rot > 0 && ggml_get_op_params_i32(dst, 4) != 0; ++ + if (ggml_get_op_params_i32(dst, 2) != 0) { + const float eps = ggml_get_op_params_f32(dst, 3); + if (dst->type == GGML_TYPE_F16) { +- rope_pe_permute_rms_cuda(src0, pe_d, (half *) dst->data, n_rot, scale, eps, stream); ++ if (interleaved) { ++ rope_pe_permute_rms_cuda(src0, pe_d, (half *) dst->data, n_rot, scale, eps, stream); ++ } else { ++ rope_pe_permute_rms_cuda(src0, pe_d, (half *) dst->data, n_rot, scale, eps, stream); ++ } + } else { + GGML_ASSERT(dst->type == GGML_TYPE_F32); +- rope_pe_permute_rms_cuda(src0, pe_d, (float *) dst->data, n_rot, scale, eps, stream); ++ if (interleaved) { ++ rope_pe_permute_rms_cuda(src0, pe_d, (float *) dst->data, n_rot, scale, eps, stream); ++ } else { ++ rope_pe_permute_rms_cuda(src0, pe_d, (float *) dst->data, n_rot, scale, eps, stream); ++ } + } + return; + } + + if (dst->type == GGML_TYPE_F16) { +- rope_pe_permute_cuda(src0, pe_d, (half *) dst->data, n_rot, scale, stream); ++ if (interleaved) { ++ rope_pe_permute_cuda(src0, pe_d, (half *) dst->data, n_rot, scale, stream); ++ } else { ++ rope_pe_permute_cuda(src0, pe_d, (half *) dst->data, n_rot, scale, stream); ++ } + } else { + GGML_ASSERT(dst->type == GGML_TYPE_F32); +- rope_pe_permute_cuda(src0, pe_d, (float *) dst->data, n_rot, scale, stream); ++ if (interleaved) { ++ rope_pe_permute_cuda(src0, pe_d, (float *) dst->data, n_rot, scale, stream); ++ } else { ++ rope_pe_permute_cuda(src0, pe_d, (float *) dst->data, n_rot, scale, stream); ++ } + } + } + +diff --git a/src/ggml-cuda/rope.cu b/src/ggml-cuda/rope.cu +index d332618b..fb5f62bb 100644 +--- a/src/ggml-cuda/rope.cu ++++ b/src/ggml-cuda/rope.cu +@@ -13,6 +13,9 @@ static __global__ void rope_apply_f32( + const float * x, const float * pe, float * dst, + int64_t half_dim, int64_t heads, int64_t tokens, int64_t pairs, + RopeApplyStrides strides) { ++ // Launched with programmatic dependent launch (ggml_cuda_kernel_launch): wait for the producer of x ++ // before reading it. Without this the kernel could read q/k while the preceding norm still wrote them. ++ ggml_cuda_pdl_sync(); + const int64_t index = (int64_t) blockIdx.x*blockDim.x + threadIdx.x; + if (index >= pairs) { + return; +diff --git a/src/ggml.c b/src/ggml.c +index 1fe59a0e..016995c0 100644 +--- a/src/ggml.c ++++ b/src/ggml.c +@@ -5751,6 +5751,14 @@ struct ggml_tensor * ggml_rope_pe_permute_rms( + return result; + } + ++struct ggml_tensor * ggml_rope_pe_permute_set_interleaved( ++ struct ggml_tensor * t, ++ bool interleaved) { ++ GGML_ASSERT(t->op == GGML_OP_ROPE_PE_PERMUTE); ++ ggml_set_op_params_i32(t, 4, interleaved ? 1 : 0); ++ return t; ++} ++ + static struct ggml_tensor * ggml_modulate_rows_impl( + struct ggml_context * ctx, + struct ggml_tensor * acc, +diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp +index 778bed20..34417288 100644 +--- a/tests/test-backend-ops.cpp ++++ b/tests/test-backend-ops.cpp +@@ -2196,15 +2196,16 @@ struct test_rope_pe_permute : public test_case { + const float scale; + const bool view; // a is a strided view (one third of a fused qkv projection) + const float eps; // >= 0: ggml_rope_pe_permute_rms (rows RMS-normalised first) ++ const bool inter; // interleaved pairing (2j, 2j+1) + + std::string vars() override { +- return VARS_TO_STR6(type, ne_a, n_rot, scale, view, eps); ++ return VARS_TO_STR7(type, ne_a, n_rot, scale, view, eps, inter); + } + + test_rope_pe_permute(ggml_type type = GGML_TYPE_F32, + std::array ne_a = {128, 4, 9, 1}, +- int n_rot = 96, float scale = 1.0f, bool view = false, float eps = -1.0f) +- : type(type), ne_a(ne_a), n_rot(n_rot), scale(scale), view(view), eps(eps) {} ++ int n_rot = 96, float scale = 1.0f, bool view = false, float eps = -1.0f, bool inter = false) ++ : type(type), ne_a(ne_a), n_rot(n_rot), scale(scale), view(view), eps(eps), inter(inter) {} + + ggml_tensor * build_graph(ggml_context * ctx) override { + ggml_tensor * a; +@@ -2225,6 +2226,9 @@ struct test_rope_pe_permute : public test_case { + } + ggml_tensor * out = eps >= 0.0f ? ggml_rope_pe_permute_rms(ctx, a, pe, n_rot, scale, type, eps) + : ggml_rope_pe_permute(ctx, a, pe, n_rot, scale, type); ++ if (inter) { ++ ggml_rope_pe_permute_set_interleaved(out, true); ++ } + ggml_set_name(out, "out"); + return out; + } +@@ -7530,6 +7534,129 @@ struct test_group_norm_mul_add : public test_case { + } + }; + ++// GGML_OP_GROUP_NORM + per-channel MUL + per-channel ADD (+ SILU): the VAE / U-Net resnet norm ++struct test_group_norm_affine : public test_case { ++ const std::array ne; ++ const int num_groups; ++ const bool silu; ++ const bool reshape; // 1-D weights reshaped to [1, 1, C, 1] after the norm is built (sd.cpp's order) ++ ++ std::string op_desc(ggml_tensor * t) override { ++ GGML_UNUSED(t); ++ return "GROUP_NORM_AFFINE"; ++ } ++ ++ bool run_whole_graph() override { return true; } ++ ++ std::string vars() override { ++ return VARS_TO_STR4(ne, num_groups, silu, reshape); ++ } ++ ++ test_group_norm_affine(std::array ne = {64, 64, 320, 1}, int num_groups = 32, bool silu = true, bool reshape = false) ++ : ne(ne), num_groups(num_groups), silu(silu), reshape(reshape) {} ++ ++ ggml_tensor * build_graph(ggml_context * ctx) override { ++ ggml_tensor * a = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne.data()); ++ ggml_tensor * w = reshape ? ggml_new_tensor_1d(ctx, GGML_TYPE_F32, ne[2]) : ggml_new_tensor_4d(ctx, GGML_TYPE_F32, 1, 1, ne[2], 1); ++ ggml_tensor * b = reshape ? ggml_new_tensor_1d(ctx, GGML_TYPE_F32, ne[2]) : ggml_new_tensor_4d(ctx, GGML_TYPE_F32, 1, 1, ne[2], 1); ++ ggml_set_name(a, "a"); ggml_set_name(w, "w"); ggml_set_name(b, "b"); ++ if (reshape) { ++ w = ggml_reshape_4d(ctx, w, 1, 1, ne[2], 1); ++ b = ggml_reshape_4d(ctx, b, 1, 1, ne[2], 1); ++ } ++ ggml_tensor * out = ggml_group_norm(ctx, a, num_groups, 1e-6f); ++ out = ggml_mul_inplace(ctx, out, w); ++ out = ggml_add_inplace(ctx, out, b); ++ if (silu) { ++ out = ggml_silu_inplace(ctx, out); ++ } ++ ggml_set_name(out, "out"); ++ return out; ++ } ++}; ++ ++// RMS norm over the outermost (channel) axis via permute + cont, as the Wan VAE builds it ++struct test_rms_norm_channel : public test_case { ++ const std::array ne; // [W, H, T, C] ++ const bool reshape_w; ++ ++ std::string op_desc(ggml_tensor * t) override { ++ GGML_UNUSED(t); ++ return "RMS_NORM_CHANNEL"; ++ } ++ ++ bool run_whole_graph() override { return true; } ++ ++ std::string vars() override { ++ return VARS_TO_STR2(ne, reshape_w); ++ } ++ ++ test_rms_norm_channel(std::array ne = {32, 24, 1, 96}, bool reshape_w = true) ++ : ne(ne), reshape_w(reshape_w) {} ++ ++ ggml_tensor * build_graph(ggml_context * ctx) override { ++ ggml_tensor * x = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne.data()); ++ ggml_tensor * w = reshape_w ? ggml_new_tensor_3d(ctx, GGML_TYPE_F32, ne[3], 1, 1) : ggml_new_tensor_1d(ctx, GGML_TYPE_F32, ne[3]); ++ ggml_set_name(x, "x"); ggml_set_name(w, "w"); ++ if (reshape_w) { ++ w = ggml_reshape_1d(ctx, w, ne[3]); ++ } ++ ggml_tensor * h = ggml_cont(ctx, ggml_permute(ctx, x, 1, 2, 3, 0)); // [C, W, H, T] ++ h = ggml_rms_norm(ctx, h, 1e-12f); ++ h = ggml_mul(ctx, h, w); ++ h = ggml_cont(ctx, ggml_permute(ctx, h, 3, 0, 1, 2)); // [W, H, T, C] ++ ggml_set_name(h, "out"); ++ return h; ++ } ++}; ++ ++// MUL -> ADD (-> ADD) with broadcast operands: adaLN modulate x + x*scale + shift, gated residual r + y*g ++struct test_mul_add_chain : public test_case { ++ const std::array ne; ++ const int mode; // 0: x + x*s, 1: x + x*s + shift, 2: r + y*g, 3: (y*g) + r with views of a fused qkv-like tensor ++ ++ std::string op_desc(ggml_tensor * t) override { ++ GGML_UNUSED(t); ++ return "MUL_ADD_CHAIN"; ++ } ++ ++ bool run_whole_graph() override { return true; } ++ ++ std::string vars() override { ++ return VARS_TO_STR2(ne, mode); ++ } ++ ++ test_mul_add_chain(std::array ne = {64, 33, 2, 1}, int mode = 1) : ne(ne), mode(mode) {} ++ ++ ggml_tensor * build_graph(ggml_context * ctx) override { ++ ggml_tensor * x = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne.data()); ++ ggml_tensor * s = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, ne[0], 1, ne[2], ne[3]); ++ ggml_set_name(x, "x"); ggml_set_name(s, "s"); ++ ggml_tensor * out; ++ if (mode <= 1) { ++ out = ggml_add(ctx, x, ggml_mul(ctx, x, s)); ++ if (mode == 1) { ++ ggml_tensor * sh = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, ne[0], 1, ne[2], ne[3]); ++ ggml_set_name(sh, "shift"); ++ out = ggml_add(ctx, out, sh); ++ } ++ } else { ++ ggml_tensor * r = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne.data()); ++ ggml_set_name(r, "r"); ++ ggml_tensor * g = s; ++ if (mode == 3) { ++ ggml_tensor * big = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, 3 * ne[0], 1, ne[2], ne[3]); ++ ggml_set_name(big, "big"); ++ g = ggml_view_4d(ctx, big, ne[0], 1, ne[2], ne[3], big->nb[1], big->nb[2], big->nb[3], ne[0] * sizeof(float)); ++ } ++ ggml_tensor * m = ggml_mul(ctx, x, g); ++ out = mode == 2 ? ggml_add(ctx, r, m) : ggml_add(ctx, m, r); ++ } ++ ggml_set_name(out, "out"); ++ return out; ++ } ++}; ++ + // GGML_OP_L2_NORM x N: independent same-shape norms in one graph (strided qkv views or + // contiguous), consuming adds nested so the norms stay adjacent in the graph. + struct test_l2_norm_batch : public test_case { +@@ -9233,6 +9360,12 @@ static std::vector> make_test_cases_eval() { + test_cases.emplace_back(new test_rope_pe_permute(type, {128, 4, 9, 1}, 96, scale, view, 1e-6f)); + test_cases.emplace_back(new test_rope_pe_permute(type, {32, 2, 5, 1}, 16, scale, view, 1e-5f)); + test_cases.emplace_back(new test_rope_pe_permute(type, {48, 2, 5, 1}, 32, scale, view, 1e-5f)); ++ // interleaved pairing (Flux / Qwen-Image / Z-Image style DiT attention) ++ test_cases.emplace_back(new test_rope_pe_permute(type, {128, 4, 9, 1}, 128, scale, view, -1.0f, true)); ++ test_cases.emplace_back(new test_rope_pe_permute(type, {128, 3, 5, 2}, 96, scale, view, -1.0f, true)); ++ test_cases.emplace_back(new test_rope_pe_permute(type, {64, 3, 7, 2}, 64, scale, view, -1.0f, true)); ++ test_cases.emplace_back(new test_rope_pe_permute(type, {64, 32, 37, 1}, 48, scale, view, 1e-5f, true)); ++ test_cases.emplace_back(new test_rope_pe_permute(type, {128, 4, 9, 1}, 128, scale, view, 1e-6f, true)); + } + } + } +@@ -10815,6 +10948,27 @@ static std::vector> make_test_cases_eval() { + test_cases.emplace_back(new test_group_norm(GGML_TYPE_F32, {9, 9, 1280, 1})); + test_cases.emplace_back(new test_group_norm_mul_add(GGML_TYPE_F32, {64, 64, 320, 1})); + test_cases.emplace_back(new test_group_norm_mul_add(GGML_TYPE_F32, {9, 9, 1280, 1})); ++ test_cases.emplace_back(new test_group_norm(GGML_TYPE_F32, {32, 48, 128, 2})); ++ test_cases.emplace_back(new test_group_norm(GGML_TYPE_F32, {33, 17, 100, 1})); ++ for (bool silu : {false, true}) { ++ test_cases.emplace_back(new test_group_norm_affine({64, 64, 320, 1}, 32, silu)); ++ test_cases.emplace_back(new test_group_norm_affine({32, 48, 128, 2}, 32, silu)); ++ test_cases.emplace_back(new test_group_norm_affine({9, 9, 1280, 1}, 32, silu)); ++ test_cases.emplace_back(new test_group_norm_affine({33, 17, 100, 1}, 32, silu)); ++ test_cases.emplace_back(new test_group_norm_affine({16, 16, 64, 1}, 32, silu)); ++ test_cases.emplace_back(new test_group_norm_affine({128, 96, 128, 1}, 32, silu)); ++ test_cases.emplace_back(new test_group_norm_affine({64, 64, 320, 1}, 32, silu, true)); ++ test_cases.emplace_back(new test_group_norm_affine({33, 17, 100, 2}, 32, silu, true)); ++ } ++ for (int mode : {0, 1, 2, 3}) { ++ test_cases.emplace_back(new test_mul_add_chain({64, 33, 2, 1}, mode)); ++ test_cases.emplace_back(new test_mul_add_chain({3072, 17, 1, 1}, mode)); ++ test_cases.emplace_back(new test_mul_add_chain({129, 5, 3, 2}, mode)); ++ } ++ for (int64_t C : {32, 96, 128, 144, 192, 288, 384, 576, 1000, 1024, 1152, 2100}) { ++ test_cases.emplace_back(new test_rms_norm_channel({32, 24, 1, C}, true)); ++ test_cases.emplace_back(new test_rms_norm_channel({17, 9, 3, C}, false)); ++ } + test_cases.emplace_back(new test_acc(GGML_TYPE_F32, {256, 17, 1, 1}, {256, 16, 1, 1}, -1)); + test_cases.emplace_back(new test_acc(GGML_TYPE_F32, {256, 17, 2, 3}, {256, 16, 2, 3}, -1)); + test_cases.emplace_back(new test_acc(GGML_TYPE_F32, {256, 17, 2, 3}, {128, 16, 2, 3}, -1)); From f81f82b18983844dbdbb27a05750d01e56fdb709 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 11 Oct 2026 08:28:22 -0700 Subject: [PATCH 02/10] Use fused rope+permute for interleaved DiT attention; skip no-op VAE conts Rope::attention builds q/k/v through rope_pe_permute (interleaved mode) when flash attention is on and the shapes match (Z-Image, Flux, Flux2, Chroma, Qwen-Image), falling back to the existing path otherwise. ggml_ext_attention_prepared accepts an optional mask. SD_FUSED_ROPE_ATTN=0 disables it. Conv2d im2col output and Wan Conv2dBut3d skip the final cont when the permute is a no-op (N == 1 / T == 1). --- src/core/ggml_extend.cpp | 25 +++++++++++--- src/core/ggml_extend.h | 5 ++- src/model/common/rope.hpp | 72 +++++++++++++++++++++++++++++++++++++++ src/model/vae/wan_vae.hpp | 6 ++-- 4 files changed, 101 insertions(+), 7 deletions(-) diff --git a/src/core/ggml_extend.cpp b/src/core/ggml_extend.cpp index 585e542c4e..c136be009a 100644 --- a/src/core/ggml_extend.cpp +++ b/src/core/ggml_extend.cpp @@ -451,7 +451,15 @@ 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, except that the final [OC, N] -> [N, OC] relayout is skipped when it moves no data + // (a batch of one, every VAE decode and most U-Net steps): same values, one full copy fewer per conv. + 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,18 +839,27 @@ 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) { + // same mask preparation as ggml_ext_attention_ext's flash path + 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; } diff --git a/src/core/ggml_extend.h b/src/core/ggml_extend.h index e94dd1451e..512bafe9b5 100644 --- a/src/core/ggml_extend.h +++ b/src/core/ggml_extend.h @@ -236,6 +236,8 @@ 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); it gets the +// same query-dimension repeat and F16 cast. // return: [N, L_q, n_head * d_head] ggml_tensor* ggml_ext_attention_prepared(ggml_context* ctx, ggml_backend_t backend, @@ -244,7 +246,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, diff --git a/src/model/common/rope.hpp b/src/model/common/rope.hpp index 93e457d4b0..ce7ed2778a 100644 --- a/src/model/common/rope.hpp +++ b/src/model/common/rope.hpp @@ -8,6 +8,7 @@ #include #include #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; } + // SD_FUSED_ROPE_ATTN=0 turns the fused attention prologue below off. + __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; + } + + // Flash attention with the RoPE, head-major relayout, kv scale and F16 cast of q/k/v done by one + // ggml_rope_pe_permute op per tensor, instead of RoPE + relayout, then scale, then cast (k), and + // relayout, scale, cast (v). Same values as the unfused flash path: the op rounds each product, + // the sum, the scale and the conversion separately, exactly like ggml_rope_apply -> ggml_scale -> + // ggml_cpy. GPU backends only (CUDA / HIP); 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; + // Interleaved pairing only: the non-interleaved form of ggml_rope_apply reads two of the four table + // entries, so it matches the fused op only for exact rotation tables. + 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(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], the shapes the unfused path hands flash attention + 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] diff --git a/src/model/vae/wan_vae.hpp b/src/model/vae/wan_vae.hpp index 9dd30fdff6..3dea88a82e 100644 --- a/src/model/vae/wan_vae.hpp +++ b/src/model/vae/wan_vae.hpp @@ -147,14 +147,16 @@ namespace WAN { using Conv2d::Conv2d; ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { + // Swapping the time and channel axes moves no data when one of them is 1 (a single frame, the + // image VAE case): ggml_ext_cont then keeps the view instead of copying the whole tensor. 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; } From 24b8c3717d844b1c547678bc2597e18a5bd22be8 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 11 Oct 2026 08:28:22 -0700 Subject: [PATCH 03/10] docs: fused DiT and VAE kernel switches --- docs/performance.md | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/docs/performance.md b/docs/performance.md index 028a6409a5..525eea92ba 100644 --- a/docs/performance.md +++ b/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 | 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. From f170601cf81eb1988fb8c0eaa372d65d75b87d0a Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 11 Oct 2026 08:57:51 -0700 Subject: [PATCH 04/10] ggml-cuda: fuse the channel RMS norm only where it beats the unfused chain The fused kernel walks every channel in one thread per position, so it is slower than CONT + RMS_NORM + MUL + CONT on small layers (VAE tiles). Fuse from 16384 positions when C < 1024 and from 65536 when C >= 1024; GGML_CUDA_NORM_CHANNEL=2 forces it at every size. Adds VAE-shaped perf cases for RMS_NORM_CHANNEL. --- docs/performance.md | 2 +- .../0007-ggml-cuda-generic-dit-vae.patch | 36 +++++++++++++++++-- 2 files changed, 34 insertions(+), 4 deletions(-) diff --git a/docs/performance.md b/docs/performance.md index 525eea92ba..b95bd45975 100644 --- a/docs/performance.md +++ b/docs/performance.md @@ -162,7 +162,7 @@ default on CUDA. Each gives the same output bit for bit as the ops it replaces, |---|---|---| | 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 | Qwen-Image, Qwen-Image-2.1, Wan VAE | `GGML_CUDA_NORM_CHANNEL` | +| 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. diff --git a/scripts/unsloth/ggml-patches/0007-ggml-cuda-generic-dit-vae.patch b/scripts/unsloth/ggml-patches/0007-ggml-cuda-generic-dit-vae.patch index 9741f98b52..02d8f82681 100644 --- a/scripts/unsloth/ggml-patches/0007-ggml-cuda-generic-dit-vae.patch +++ b/scripts/unsloth/ggml-patches/0007-ggml-cuda-generic-dit-vae.patch @@ -564,10 +564,10 @@ index 00000000..8d9752ca +int ggml_cuda_try_fuse_mul_add(ggml_backend_cuda_context & ctx, ggml_cgraph * cgraph, int i); diff --git a/src/ggml-cuda/norm-channel.cu b/src/ggml-cuda/norm-channel.cu new file mode 100644 -index 00000000..84a33c9c +index 00000000..1b921122 --- /dev/null +++ b/src/ggml-cuda/norm-channel.cu -@@ -0,0 +1,228 @@ +@@ -0,0 +1,244 @@ +#include "norm-channel.cuh" +#include "ggml-impl.h" + @@ -590,6 +590,15 @@ index 00000000..84a33c9c +#endif +} + ++// GGML_CUDA_NORM_CHANNEL=2 fuses at every size (tests); by default only where it beats the unfused chain. ++static bool norm_channel_forced() { ++ static const bool forced = [] { ++ const char * e = getenv("GGML_CUDA_NORM_CHANNEL"); ++ return e != nullptr && atoi(e) == 2; ++ }(); ++ return forced; ++} ++ +static __device__ __forceinline__ float nc_butterfly32(float * v) { + // value every lane holds after warp_reduce_sum: lane l adds lane l^off at each step +#pragma unroll @@ -776,6 +785,13 @@ index 00000000..84a33c9c + if (P > INT64_C(1) << 40 || partial_overlap) { + NC_REJECT("check 10"); + } ++ // One thread per position walks every channel, so the kernel needs many positions to fill the GPU. ++ // Measured on B200 against CONT + RMS_NORM + MUL + CONT (Qwen-Image VAE shapes, whole image and tiles): ++ // with C < 1024 it wins from 16384 positions (173 vs 245 us), with C >= 1024 from 65536 (1.2 vs 3.3 ms), ++ // and loses below (P = 4096, C = 1152: 656 vs 209 us). ++ if (!norm_channel_forced() && P < (C < 1024 ? 16384 : 65536)) { ++ NC_REJECT("check 12"); ++ } + float eps; + memcpy(&eps, rms->op_params, sizeof(float)); + @@ -1046,7 +1062,7 @@ index 1fe59a0e..016995c0 100644 struct ggml_context * ctx, struct ggml_tensor * acc, diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp -index 778bed20..34417288 100644 +index 778bed20..e4438d10 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -2196,15 +2196,16 @@ struct test_rope_pe_permute : public test_case { @@ -1250,3 +1266,17 @@ index 778bed20..34417288 100644 test_cases.emplace_back(new test_acc(GGML_TYPE_F32, {256, 17, 1, 1}, {256, 16, 1, 1}, -1)); test_cases.emplace_back(new test_acc(GGML_TYPE_F32, {256, 17, 2, 3}, {256, 16, 2, 3}, -1)); test_cases.emplace_back(new test_acc(GGML_TYPE_F32, {256, 17, 2, 3}, {128, 16, 2, 3}, -1)); +@@ -11303,6 +11457,13 @@ static std::vector> make_test_cases_eval() { + static std::vector> make_test_cases_perf() { + std::vector> test_cases; + ++ // Wan / Qwen-Image VAE channel RMS norm at whole-image and tile sizes ++ for (int64_t C : {144, 288, 576, 1152}) { ++ for (int64_t S : {16, 32, 64, 128, 256}) { ++ test_cases.emplace_back(new test_rms_norm_channel({S, S, 1, C}, true)); ++ } ++ } ++ + // SWIGLU at a 27B-class FFN width, fused [gate|up] vs split operands + // note: same bytes either way, so a backend that indexes them differently shows it here + for (ggml_type type : {GGML_TYPE_F16, GGML_TYPE_F32}) { From aeb58f026f0f6718b1111d04edc3201783bc9d35 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 11 Oct 2026 09:26:22 -0700 Subject: [PATCH 05/10] ggml-cuda: skip the channel RMS norm fusion when its output overlaps the weight --- .../ggml-patches/0007-ggml-cuda-generic-dit-vae.patch | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/scripts/unsloth/ggml-patches/0007-ggml-cuda-generic-dit-vae.patch b/scripts/unsloth/ggml-patches/0007-ggml-cuda-generic-dit-vae.patch index 02d8f82681..9279aeeb7c 100644 --- a/scripts/unsloth/ggml-patches/0007-ggml-cuda-generic-dit-vae.patch +++ b/scripts/unsloth/ggml-patches/0007-ggml-cuda-generic-dit-vae.patch @@ -564,10 +564,10 @@ index 00000000..8d9752ca +int ggml_cuda_try_fuse_mul_add(ggml_backend_cuda_context & ctx, ggml_cgraph * cgraph, int i); diff --git a/src/ggml-cuda/norm-channel.cu b/src/ggml-cuda/norm-channel.cu new file mode 100644 -index 00000000..1b921122 +index 00000000..0ff8e91e --- /dev/null +++ b/src/ggml-cuda/norm-channel.cu -@@ -0,0 +1,244 @@ +@@ -0,0 +1,248 @@ +#include "norm-channel.cuh" +#include "ggml-impl.h" + @@ -782,7 +782,11 @@ index 00000000..1b921122 + const char * x0 = (const char *) x->data; + const char * d0 = (const char *) cont2->data; + const bool partial_overlap = d0 != x0 && d0 < x0 + ggml_nbytes(x) && x0 < d0 + ggml_nbytes(cont2); -+ if (P > INT64_C(1) << 40 || partial_overlap) { ++ // Unfused, w is read by the MUL before the second CONT writes, so the allocator may hand w's buffer to ++ // cont2; the fused kernel reads w while writing cont2, so any overlap with w is unsafe. ++ const char * w0 = (const char *) w->data; ++ const bool w_overlap = d0 < w0 + ggml_nbytes(w) && w0 < d0 + ggml_nbytes(cont2); ++ if (P > INT64_C(1) << 40 || partial_overlap || w_overlap) { + NC_REJECT("check 10"); + } + // One thread per position walks every channel, so the kernel needs many positions to fill the GPU. From 0f53e720725b2245d07dd35de7029fa24d3de19c Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 11 Oct 2026 16:21:15 -0700 Subject: [PATCH 06/10] fused attention: honour the attention-scale override; fusions: skip nodes outside a selected graph branch - Rope::attention's fused RoPE + flash path now applies ctx->attn_scale like ggml_ext_attention_ext, so --attn-scale keeps large K/V finite in F16 on both paths and they stay bit-identical - ggml_ext_attention_prepared: ggml_prec_set_acc instead of the deprecated ggml_flash_attn_ext_set_prec - 0007: the MUL+ADD(+ADD) and channel RMS norm matchers require GGML_TENSOR_FLAG_COMPUTE on every node they consume (graphs from ggml_build_forward_select), as ggml_can_fuse_ext does for GROUP_NORM - 0007: CPU interleaved ROPE_PE_PERMUTE evaluates the same expression as ROPE_APPLY so FMA contraction cannot make them differ - 0007: tests/test-fusion-select.cpp covers the selected-branch graphs and the CPU RoPE parity --- .../0007-ggml-cuda-generic-dit-vae.patch | 336 +++++++++++++++++- src/core/ggml_extend.cpp | 2 +- src/model/common/rope.hpp | 5 + 3 files changed, 323 insertions(+), 20 deletions(-) diff --git a/scripts/unsloth/ggml-patches/0007-ggml-cuda-generic-dit-vae.patch b/scripts/unsloth/ggml-patches/0007-ggml-cuda-generic-dit-vae.patch index 9279aeeb7c..6347ab6da6 100644 --- a/scripts/unsloth/ggml-patches/0007-ggml-cuda-generic-dit-vae.patch +++ b/scripts/unsloth/ggml-patches/0007-ggml-cuda-generic-dit-vae.patch @@ -18,7 +18,7 @@ index 41ac28bb..fa83c5e6 100644 struct ggml_context * ctx, struct ggml_tensor * a, diff --git a/src/ggml-cpu/ops.cpp b/src/ggml-cpu/ops.cpp -index bdd36aea..cadb25ac 100644 +index bdd36aea..2883f5df 100644 --- a/src/ggml-cpu/ops.cpp +++ b/src/ggml-cpu/ops.cpp @@ -12397,6 +12397,7 @@ void ggml_compute_forward_rope_pe_permute( @@ -29,27 +29,45 @@ index bdd36aea..cadb25ac 100644 const int half = n_rot / 2; const int64_t D = src0->ne[0]; -@@ -12436,14 +12437,14 @@ void ggml_compute_forward_rope_pe_permute( +@@ -12435,17 +12436,28 @@ void ggml_compute_forward_rope_pe_permute( + if (n_rot > 0) { const float * t = (const float *) pe->data + l * (int64_t) half * 4; - for (int i = 0; i < n_rot; ++i) { +- for (int i = 0; i < n_rot; ++i) { - const int r = i / half; - const int j = i - r * half; - p0[i] = x[j] * t[4 * j + 2 * r + 0]; -+ const int r = inter ? i % 2 : i / half; -+ const int j = inter ? i / 2 : i - r * half; -+ p0[i] = x[inter ? 2 * j : j] * t[4 * j + 2 * r + 0]; - } - for (int i = 0; i < n_rot; ++i) { +- } +- for (int i = 0; i < n_rot; ++i) { - const int r = i / half; - const int j = i - r * half; - p1[i] = x[j + half] * t[4 * j + 2 * r + 1]; -+ const int r = inter ? i % 2 : i / half; -+ const int j = inter ? i / 2 : i - r * half; -+ p1[i] = x[inter ? 2 * j + 1 : j + half] * t[4 * j + 2 * r + 1]; ++ if (inter) { ++ // the expression ggml_compute_forward_rope_apply (interleaved) evaluates, so the compiler ++ // contracts it (or not) the same way and the two ops stay bit-identical on this backend ++ for (int j = 0; j < half; ++j) { ++ const float a = x[2 * j]; ++ const float b = x[2 * j + 1]; ++ y[2 * j] = a*t[4 * j + 0] + b*t[4 * j + 1]; ++ y[2 * j + 1] = a*t[4 * j + 2] + b*t[4 * j + 3]; ++ } ++ } else { ++ for (int i = 0; i < n_rot; ++i) { ++ const int r = i / half; ++ const int j = i - r * half; ++ p0[i] = x[j] * t[4 * j + 2 * r + 0]; ++ } ++ for (int i = 0; i < n_rot; ++i) { ++ const int r = i / half; ++ const int j = i - r * half; ++ p1[i] = x[j + half] * t[4 * j + 2 * r + 1]; ++ } ++ ggml_vec_add_f32(n_rot, y.data(), p0.data(), p1.data()); } - ggml_vec_add_f32(n_rot, y.data(), p0.data(), p1.data()); +- ggml_vec_add_f32(n_rot, y.data(), p0.data(), p1.data()); } + if (D > n_rot) { + memcpy(y.data() + n_rot, x + n_rot, (D - n_rot) * sizeof(float)); diff --git a/src/ggml-cuda/ggml-cuda.cu b/src/ggml-cuda/ggml-cuda.cu index 831d0de8..2df04786 100644 --- a/src/ggml-cuda/ggml-cuda.cu @@ -373,10 +391,10 @@ index 00000000..eb6298ea + const ggml_tensor * b, bool silu, ggml_tensor * out); diff --git a/src/ggml-cuda/mul-add-fused.cu b/src/ggml-cuda/mul-add-fused.cu new file mode 100644 -index 00000000..ff4e1dd0 +index 00000000..29484ce4 --- /dev/null +++ b/src/ggml-cuda/mul-add-fused.cu -@@ -0,0 +1,171 @@ +@@ -0,0 +1,172 @@ +#include "mul-add-fused.cuh" +#include "ggml-impl.h" + @@ -485,8 +503,9 @@ index 00000000..ff4e1dd0 + return 0; + } + ggml_tensor * add = cgraph->nodes[i_add]; -+ if (add->op != GGML_OP_ADD || add->type != GGML_TYPE_F32 || !ggml_is_contiguous(add) || !ggml_are_same_shape(add, mul) || -+ (add->src[0] != mul && add->src[1] != mul)) { ++ // under ggml_build_forward_select only GGML_TENSOR_FLAG_COMPUTE nodes run: never fuse an inactive node ++ if (add->op != GGML_OP_ADD || (add->flags & GGML_TENSOR_FLAG_COMPUTE) == 0 || add->type != GGML_TYPE_F32 || ++ !ggml_is_contiguous(add) || !ggml_are_same_shape(add, mul) || (add->src[0] != mul && add->src[1] != mul)) { + return 0; + } + const ggml_tensor * c = add->src[0] == mul ? add->src[1] : add->src[0]; @@ -499,7 +518,7 @@ index 00000000..ff4e1dd0 + const int i_add2 = next_op(i_add); + if (i_add2 > 0) { + ggml_tensor * add2 = cgraph->nodes[i_add2]; -+ if (add2->op == GGML_OP_ADD && add2->type == GGML_TYPE_F32 && ggml_is_contiguous(add2) && ggml_are_same_shape(add2, add) && ++ if (add2->op == GGML_OP_ADD && (add2->flags & GGML_TENSOR_FLAG_COMPUTE) && add2->type == GGML_TYPE_F32 && ggml_is_contiguous(add2) && ggml_are_same_shape(add2, add) && + (add2->src[0] == add || add2->src[1] == add)) { + const ggml_tensor * dd = add2->src[0] == add ? add2->src[1] : add2->src[0]; + if (dd != add && dd != mul && mad_ok_operand(dd, add2)) { @@ -564,10 +583,10 @@ index 00000000..8d9752ca +int ggml_cuda_try_fuse_mul_add(ggml_backend_cuda_context & ctx, ggml_cgraph * cgraph, int i); diff --git a/src/ggml-cuda/norm-channel.cu b/src/ggml-cuda/norm-channel.cu new file mode 100644 -index 00000000..0ff8e91e +index 00000000..9c0b61fd --- /dev/null +++ b/src/ggml-cuda/norm-channel.cu -@@ -0,0 +1,248 @@ +@@ -0,0 +1,254 @@ +#include "norm-channel.cuh" +#include "ggml-impl.h" + @@ -770,6 +789,12 @@ index 00000000..0ff8e91e + if (!back_ok) { + NC_REJECT("check 8"); + } ++ // under ggml_build_forward_select only GGML_TENSOR_FLAG_COMPUTE nodes run: never fuse an inactive node ++ for (const ggml_tensor * n : { (const ggml_tensor *) cont1, (const ggml_tensor *) rms, (const ggml_tensor *) mul, (const ggml_tensor *) cont2 }) { ++ if ((n->flags & GGML_TENSOR_FLAG_COMPUTE) == 0) { ++ NC_REJECT("check 8b"); ++ } ++ } + // every node from i to the second CONT is part of the pattern or a view the loop skips + for (int j = i + 1; j < i_pv2 + 1; ++j) { + const ggml_tensor * n = cgraph->nodes[j]; @@ -1065,6 +1090,25 @@ index 1fe59a0e..016995c0 100644 static struct ggml_tensor * ggml_modulate_rows_impl( struct ggml_context * ctx, struct ggml_tensor * acc, +diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt +index a4fd2f88..1bb0ec64 100644 +--- a/tests/CMakeLists.txt ++++ b/tests/CMakeLists.txt +@@ -280,6 +280,14 @@ if (NOT GGML_BACKEND_DL) + target_link_libraries(${TEST_TARGET} PRIVATE ggml) + add_test(NAME ${TEST_TARGET} COMMAND $) + ++ # ++ # test-fusion-select ++ ++ set(TEST_TARGET test-fusion-select) ++ add_executable(${TEST_TARGET} ${TEST_TARGET}.cpp) ++ target_link_libraries(${TEST_TARGET} PRIVATE ggml) ++ add_test(NAME ${TEST_TARGET} COMMAND $) ++ + # + # test-conv-transpose + diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 778bed20..e4438d10 100644 --- a/tests/test-backend-ops.cpp @@ -1284,3 +1328,257 @@ index 778bed20..e4438d10 100644 // SWIGLU at a 27B-class FFN width, fused [gate|up] vs split operands // note: same bytes either way, so a backend that indexes them differently shows it here for (ggml_type type : {GGML_TYPE_F16, GGML_TYPE_F32}) { +diff --git a/tests/test-fusion-select.cpp b/tests/test-fusion-select.cpp +new file mode 100644 +index 00000000..16746c56 +--- /dev/null ++++ b/tests/test-fusion-select.cpp +@@ -0,0 +1,248 @@ ++// Checks for the fused CUDA paths of the generic DiT / VAE patch: ++// 1. graphs built with ggml_build_forward_select: a backend may only fuse nodes that are flagged ++// GGML_TENSOR_FLAG_COMPUTE, so the selected tensors must equal a graph that holds only them and the ++// inactive tensors must stay untouched (MUL -> ADD -> ADD chain and the channel RMS norm chain). ++// 2. CPU: ggml_rope_pe_permute with interleaved pairing is bit-identical to ggml_rope_apply(interleaved) ++// -> ggml_scale -> ggml_cpy, the chain it replaces. ++// Returns 0 when every check passes (backends without the ops skip their checks). ++ ++#include "ggml.h" ++#include "ggml-alloc.h" ++#include "ggml-backend.h" ++#include "ggml-cpu.h" ++ ++#include ++#include ++#include ++#include ++#include ++#include ++#include ++#include ++ ++static int n_fail = 0; ++static int n_pass = 0; ++ ++static void check(bool ok, const std::string & what) { ++ printf(" %-72s %s\n", what.c_str(), ok ? "OK" : "FAIL"); ++ (ok ? n_pass : n_fail)++; ++} ++ ++static std::vector rand_vec(size_t n, uint32_t seed, float lo = -2.0f, float hi = 2.0f) { ++ std::mt19937 rng(seed); ++ std::uniform_real_distribution d(lo, hi); ++ std::vector v(n); ++ for (auto & x : v) { ++ x = d(rng); ++ } ++ return v; ++} ++ ++static std::vector get_f32(const ggml_tensor * t) { ++ std::vector v(ggml_nelements(t)); ++ ggml_backend_tensor_get(t, v.data(), 0, ggml_nbytes(t)); ++ return v; ++} ++ ++static void fill(ggml_tensor * t, float value) { ++ std::vector v(ggml_nelements(t), value); ++ ggml_backend_tensor_set(t, v.data(), 0, ggml_nbytes(t)); ++} ++ ++// A graph whose tensors each own their memory (no allocator reuse), so an unwritten output keeps its sentinel. ++struct graph_case { ++ ggml_context * ctx = nullptr; ++ ggml_backend_buffer_t buf = nullptr; ++ std::vector inputs; ++ std::vector outs; // candidates for ggml_build_forward_select ++ ++ ~graph_case() { ++ if (buf) { ++ ggml_backend_buffer_free(buf); ++ } ++ if (ctx) { ++ ggml_free(ctx); ++ } ++ } ++}; ++ ++using builder_t = std::function &, std::vector &)>; ++ ++// run `build` on `backend`: select output `idx` (n_sel = number of candidates passed to the select) or, with ++// idx < 0, expand only candidate `ref`. Returns the candidate values after compute (all candidates). ++static std::vector> run_case(ggml_backend_t backend, const builder_t & build, int idx, int ref, ++ const std::vector> & input_data) { ++ graph_case g; ++ ggml_init_params ip = { 64 * ggml_tensor_overhead() + ggml_graph_overhead(), nullptr, true }; ++ g.ctx = ggml_init(ip); ++ build(g.ctx, g.inputs, g.outs); ++ g.buf = ggml_backend_alloc_ctx_tensors(g.ctx, backend); ++ for (size_t k = 0; k < g.inputs.size(); ++k) { ++ ggml_backend_tensor_set(g.inputs[k], input_data[k].data(), 0, ggml_nbytes(g.inputs[k])); ++ } ++ for (ggml_tensor * t = ggml_get_first_tensor(g.ctx); t; t = ggml_get_next_tensor(g.ctx, t)) { ++ if (t->op != GGML_OP_NONE && t->view_src == nullptr) { ++ fill(t, -77.0f); ++ } ++ } ++ ggml_cgraph * gf = ggml_new_graph(g.ctx); ++ if (idx >= 0) { ++ ggml_build_forward_select(gf, g.outs.data(), (int) g.outs.size(), idx); ++ } else { ++ ggml_build_forward_expand(gf, g.outs[ref]); ++ } ++ GGML_ASSERT(ggml_backend_graph_compute(backend, gf) == GGML_STATUS_SUCCESS); ++ std::vector> res; ++ for (ggml_tensor * t : g.outs) { ++ res.push_back(get_f32(t)); ++ } ++ return res; ++} ++ ++static bool bit_equal(const std::vector & a, const std::vector & b) { ++ return a.size() == b.size() && memcmp(a.data(), b.data(), a.size() * sizeof(float)) == 0; ++} ++ ++static bool all_sentinel(const std::vector & a) { ++ for (float x : a) { ++ if (x != -77.0f) { ++ return false; ++ } ++ } ++ return true; ++} ++ ++// every select index: the selected candidate equals the expand-only reference and every candidate not ++// needed by it keeps the sentinel. `needed[idx]` lists the candidates the selected one depends on. ++static void check_select(ggml_backend_t backend, const char * name, const builder_t & build, ++ const std::vector> & input_data, const std::vector> & needed) { ++ for (int idx = 0; idx < (int) needed.size(); ++idx) { ++ const auto ref = run_case(backend, build, -1, idx, input_data); ++ const auto sel = run_case(backend, build, idx, idx, input_data); ++ bool inactive_ok = true; ++ for (int k = 0; k < (int) sel.size(); ++k) { ++ bool need = false; ++ for (int n : needed[idx]) { ++ need |= n == k; ++ } ++ if (!need) { ++ inactive_ok &= all_sentinel(sel[k]); ++ } ++ } ++ check(bit_equal(ref[idx], sel[idx]), std::string(name) + " select " + std::to_string(idx) + ": selected output computed"); ++ check(inactive_ok, std::string(name) + " select " + std::to_string(idx) + ": inactive outputs untouched"); ++ } ++} ++ ++static void test_select(ggml_backend_t backend) { ++ // MUL -> ADD -> ADD: x*s + x + shift (adaLN modulate), candidates {mul, add, add2} ++ { ++ const int64_t ne0 = 64, ne1 = 33; ++ builder_t build = [&](ggml_context * ctx, std::vector & in, std::vector & out) { ++ ggml_tensor * x = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, ne0, ne1); ++ ggml_tensor * s = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, ne0, 1); ++ ggml_tensor * sh = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, ne0, 1); ++ ggml_tensor * m = ggml_mul(ctx, x, s); ++ ggml_tensor * a = ggml_add(ctx, m, x); ++ ggml_tensor * a2 = ggml_add(ctx, a, sh); ++ in = { x, s, sh }; ++ out = { m, a, a2 }; ++ }; ++ const std::vector> data = { rand_vec(ne0 * ne1, 1), rand_vec(ne0, 2), rand_vec(ne0, 3) }; ++ check_select(backend, "MUL_ADD_CHAIN", build, data, { { 0 }, { 0, 1 }, { 0, 1, 2 } }); ++ } ++ // channel RMS norm: CONT(PERMUTE(x)) -> RMS_NORM -> MUL(w) -> CONT(PERMUTE), P = 16384 positions so the ++ // fusion engages at its default size threshold; candidates {cont1, rms, mul, cont2} ++ { ++ const int64_t W = 128, H = 128, C = 8; ++ builder_t build = [&](ggml_context * ctx, std::vector & in, std::vector & out) { ++ ggml_tensor * x = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, W, H, 1, C); ++ ggml_tensor * w = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, C); ++ ggml_tensor * c1 = ggml_cont(ctx, ggml_permute(ctx, x, 1, 2, 3, 0)); ++ ggml_tensor * r = ggml_rms_norm(ctx, c1, 1e-6f); ++ ggml_tensor * m = ggml_mul(ctx, r, w); ++ ggml_tensor * c2 = ggml_cont(ctx, ggml_permute(ctx, m, 3, 0, 1, 2)); ++ in = { x, w }; ++ out = { c1, r, m, c2 }; ++ }; ++ const std::vector> data = { rand_vec(W * H * C, 4), rand_vec(C, 5, 0.5f, 1.5f) }; ++ check_select(backend, "RMS_NORM_CHANNEL", build, data, { { 0 }, { 0, 1 }, { 0, 1, 2 }, { 0, 1, 2, 3 } }); ++ } ++} ++ ++// CPU: interleaved ggml_rope_pe_permute vs ggml_rope_apply -> ggml_scale -> ggml_cpy, bitwise ++static void test_rope_interleaved_cpu(ggml_backend_t cpu) { ++ const int64_t shapes[][4] = { { 64, 3, 7, 1 }, { 128, 2, 5, 2 }, { 80, 4, 3, 1 } }; ++ const float scales[] = { 1.0f, 1.0f / 128.0f }; ++ const ggml_type types[] = { GGML_TYPE_F32, GGML_TYPE_F16 }; ++ uint32_t seed = 100; ++ for (const auto & ne : shapes) { ++ for (float scale : scales) { ++ for (ggml_type type : types) { ++ const int64_t D = ne[0], Hh = ne[1], L = ne[2], N = ne[3]; ++ const auto xd = rand_vec(D * Hh * L * N, seed++); ++ const auto ped = rand_vec(4 * (D / 2) * L, seed++, -1.0f, 1.0f); // not an exact rotation ++ std::vector outs[2]; ++ for (int arm = 0; arm < 2; ++arm) { ++ ggml_init_params ip = { 16 * ggml_tensor_overhead() + ggml_graph_overhead(), nullptr, true }; ++ ggml_context * ctx = ggml_init(ip); ++ ggml_tensor * x = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, D, Hh, L, N); ++ ggml_tensor * pe = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, 2, 2, D / 2, L); ++ ggml_tensor * y; ++ if (arm == 0) { ++ y = ggml_rope_apply(ctx, x, pe, true); ++ if (scale != 1.0f) { ++ y = ggml_scale(ctx, y, scale); ++ } ++ if (type != GGML_TYPE_F32) { ++ y = ggml_cpy(ctx, y, ggml_new_tensor(ctx, type, GGML_MAX_DIMS, y->ne)); ++ } ++ } else { ++ y = ggml_rope_pe_permute_set_interleaved(ggml_rope_pe_permute(ctx, x, pe, (int) D, scale, type), true); ++ } ++ ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors(ctx, cpu); ++ ggml_backend_tensor_set(x, xd.data(), 0, ggml_nbytes(x)); ++ ggml_backend_tensor_set(pe, ped.data(), 0, ggml_nbytes(pe)); ++ ggml_cgraph * gf = ggml_new_graph(ctx); ++ ggml_build_forward_expand(gf, y); ++ GGML_ASSERT(ggml_backend_graph_compute(cpu, gf) == GGML_STATUS_SUCCESS); ++ outs[arm].resize(ggml_nbytes(y)); ++ ggml_backend_tensor_get(y, outs[arm].data(), 0, ggml_nbytes(y)); ++ ggml_backend_buffer_free(buf); ++ ggml_free(ctx); ++ } ++ char name[160]; ++ snprintf(name, sizeof(name), "CPU ROPE_PE_PERMUTE interleaved [%lld,%lld,%lld,%lld] scale=%g %s", ++ (long long) D, (long long) Hh, (long long) L, (long long) N, scale, ggml_type_name(type)); ++ check(outs[0].size() == outs[1].size() && memcmp(outs[0].data(), outs[1].data(), outs[0].size()) == 0, name); ++ } ++ } ++ } ++} ++ ++int main() { ++ ggml_backend_load_all(); ++ ++ ggml_backend_t cpu = ggml_backend_init_by_type(GGML_BACKEND_DEVICE_TYPE_CPU, nullptr); ++ GGML_ASSERT(cpu != nullptr); ++ printf("backend CPU\n"); ++ test_rope_interleaved_cpu(cpu); ++ test_select(cpu); ++ ggml_backend_free(cpu); ++ ++ for (size_t i = 0; i < ggml_backend_dev_count(); ++i) { ++ ggml_backend_dev_t dev = ggml_backend_dev_get(i); ++ if (ggml_backend_dev_type(dev) == GGML_BACKEND_DEVICE_TYPE_CPU) { ++ continue; ++ } ++ ggml_backend_t be = ggml_backend_dev_init(dev, nullptr); ++ if (be == nullptr) { ++ continue; ++ } ++ printf("backend %s\n", ggml_backend_name(be)); ++ test_select(be); ++ ggml_backend_free(be); ++ } ++ printf("%d passed, %d failed\n", n_pass, n_fail); ++ return n_fail == 0 ? 0 : 1; ++} diff --git a/src/core/ggml_extend.cpp b/src/core/ggml_extend.cpp index c136be009a..fd98a73f71 100644 --- a/src/core/ggml_extend.cpp +++ b/src/core/ggml_extend.cpp @@ -863,7 +863,7 @@ ggml_tensor* ggml_ext_attention_prepared(ggml_context* ctx, 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); } diff --git a/src/model/common/rope.hpp b/src/model/common/rope.hpp index ce7ed2778a..f1e9d7b61a 100644 --- a/src/model/common/rope.hpp +++ b/src/model/common/rope.hpp @@ -1152,6 +1152,11 @@ namespace Rope { return nullptr; #else auto backend = ctx->backend; + // the same attention-scale override ggml_ext_attention_ext(GGMLRunnerContext*, ...) applies + // (it keeps large k/v finite in F16), so both paths compute identical values + if (ctx->attn_scale > 0.f) { + kv_scale = ctx->attn_scale; + } // Interleaved pairing only: the non-interleaved form of ggml_rope_apply reads two of the four table // entries, so it matches the fused op only for exact rotation tables. if (!rope_interleaved || !ctx->flash_attn_enabled || ctx->sage_attn_enabled || backend == nullptr || From 3fb027e9d7497cd5cb85ea456f8aa3d533cd7d6e Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 11 Oct 2026 18:04:04 -0700 Subject: [PATCH 07/10] chore: keep the fusion tests out of the 0007 ggml patch --- .../0007-ggml-cuda-generic-dit-vae.patch | 492 ------------------ 1 file changed, 492 deletions(-) diff --git a/scripts/unsloth/ggml-patches/0007-ggml-cuda-generic-dit-vae.patch b/scripts/unsloth/ggml-patches/0007-ggml-cuda-generic-dit-vae.patch index 6347ab6da6..3c86471ca3 100644 --- a/scripts/unsloth/ggml-patches/0007-ggml-cuda-generic-dit-vae.patch +++ b/scripts/unsloth/ggml-patches/0007-ggml-cuda-generic-dit-vae.patch @@ -1090,495 +1090,3 @@ index 1fe59a0e..016995c0 100644 static struct ggml_tensor * ggml_modulate_rows_impl( struct ggml_context * ctx, struct ggml_tensor * acc, -diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt -index a4fd2f88..1bb0ec64 100644 ---- a/tests/CMakeLists.txt -+++ b/tests/CMakeLists.txt -@@ -280,6 +280,14 @@ if (NOT GGML_BACKEND_DL) - target_link_libraries(${TEST_TARGET} PRIVATE ggml) - add_test(NAME ${TEST_TARGET} COMMAND $) - -+ # -+ # test-fusion-select -+ -+ set(TEST_TARGET test-fusion-select) -+ add_executable(${TEST_TARGET} ${TEST_TARGET}.cpp) -+ target_link_libraries(${TEST_TARGET} PRIVATE ggml) -+ add_test(NAME ${TEST_TARGET} COMMAND $) -+ - # - # test-conv-transpose - -diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp -index 778bed20..e4438d10 100644 ---- a/tests/test-backend-ops.cpp -+++ b/tests/test-backend-ops.cpp -@@ -2196,15 +2196,16 @@ struct test_rope_pe_permute : public test_case { - const float scale; - const bool view; // a is a strided view (one third of a fused qkv projection) - const float eps; // >= 0: ggml_rope_pe_permute_rms (rows RMS-normalised first) -+ const bool inter; // interleaved pairing (2j, 2j+1) - - std::string vars() override { -- return VARS_TO_STR6(type, ne_a, n_rot, scale, view, eps); -+ return VARS_TO_STR7(type, ne_a, n_rot, scale, view, eps, inter); - } - - test_rope_pe_permute(ggml_type type = GGML_TYPE_F32, - std::array ne_a = {128, 4, 9, 1}, -- int n_rot = 96, float scale = 1.0f, bool view = false, float eps = -1.0f) -- : type(type), ne_a(ne_a), n_rot(n_rot), scale(scale), view(view), eps(eps) {} -+ int n_rot = 96, float scale = 1.0f, bool view = false, float eps = -1.0f, bool inter = false) -+ : type(type), ne_a(ne_a), n_rot(n_rot), scale(scale), view(view), eps(eps), inter(inter) {} - - ggml_tensor * build_graph(ggml_context * ctx) override { - ggml_tensor * a; -@@ -2225,6 +2226,9 @@ struct test_rope_pe_permute : public test_case { - } - ggml_tensor * out = eps >= 0.0f ? ggml_rope_pe_permute_rms(ctx, a, pe, n_rot, scale, type, eps) - : ggml_rope_pe_permute(ctx, a, pe, n_rot, scale, type); -+ if (inter) { -+ ggml_rope_pe_permute_set_interleaved(out, true); -+ } - ggml_set_name(out, "out"); - return out; - } -@@ -7530,6 +7534,129 @@ struct test_group_norm_mul_add : public test_case { - } - }; - -+// GGML_OP_GROUP_NORM + per-channel MUL + per-channel ADD (+ SILU): the VAE / U-Net resnet norm -+struct test_group_norm_affine : public test_case { -+ const std::array ne; -+ const int num_groups; -+ const bool silu; -+ const bool reshape; // 1-D weights reshaped to [1, 1, C, 1] after the norm is built (sd.cpp's order) -+ -+ std::string op_desc(ggml_tensor * t) override { -+ GGML_UNUSED(t); -+ return "GROUP_NORM_AFFINE"; -+ } -+ -+ bool run_whole_graph() override { return true; } -+ -+ std::string vars() override { -+ return VARS_TO_STR4(ne, num_groups, silu, reshape); -+ } -+ -+ test_group_norm_affine(std::array ne = {64, 64, 320, 1}, int num_groups = 32, bool silu = true, bool reshape = false) -+ : ne(ne), num_groups(num_groups), silu(silu), reshape(reshape) {} -+ -+ ggml_tensor * build_graph(ggml_context * ctx) override { -+ ggml_tensor * a = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne.data()); -+ ggml_tensor * w = reshape ? ggml_new_tensor_1d(ctx, GGML_TYPE_F32, ne[2]) : ggml_new_tensor_4d(ctx, GGML_TYPE_F32, 1, 1, ne[2], 1); -+ ggml_tensor * b = reshape ? ggml_new_tensor_1d(ctx, GGML_TYPE_F32, ne[2]) : ggml_new_tensor_4d(ctx, GGML_TYPE_F32, 1, 1, ne[2], 1); -+ ggml_set_name(a, "a"); ggml_set_name(w, "w"); ggml_set_name(b, "b"); -+ if (reshape) { -+ w = ggml_reshape_4d(ctx, w, 1, 1, ne[2], 1); -+ b = ggml_reshape_4d(ctx, b, 1, 1, ne[2], 1); -+ } -+ ggml_tensor * out = ggml_group_norm(ctx, a, num_groups, 1e-6f); -+ out = ggml_mul_inplace(ctx, out, w); -+ out = ggml_add_inplace(ctx, out, b); -+ if (silu) { -+ out = ggml_silu_inplace(ctx, out); -+ } -+ ggml_set_name(out, "out"); -+ return out; -+ } -+}; -+ -+// RMS norm over the outermost (channel) axis via permute + cont, as the Wan VAE builds it -+struct test_rms_norm_channel : public test_case { -+ const std::array ne; // [W, H, T, C] -+ const bool reshape_w; -+ -+ std::string op_desc(ggml_tensor * t) override { -+ GGML_UNUSED(t); -+ return "RMS_NORM_CHANNEL"; -+ } -+ -+ bool run_whole_graph() override { return true; } -+ -+ std::string vars() override { -+ return VARS_TO_STR2(ne, reshape_w); -+ } -+ -+ test_rms_norm_channel(std::array ne = {32, 24, 1, 96}, bool reshape_w = true) -+ : ne(ne), reshape_w(reshape_w) {} -+ -+ ggml_tensor * build_graph(ggml_context * ctx) override { -+ ggml_tensor * x = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne.data()); -+ ggml_tensor * w = reshape_w ? ggml_new_tensor_3d(ctx, GGML_TYPE_F32, ne[3], 1, 1) : ggml_new_tensor_1d(ctx, GGML_TYPE_F32, ne[3]); -+ ggml_set_name(x, "x"); ggml_set_name(w, "w"); -+ if (reshape_w) { -+ w = ggml_reshape_1d(ctx, w, ne[3]); -+ } -+ ggml_tensor * h = ggml_cont(ctx, ggml_permute(ctx, x, 1, 2, 3, 0)); // [C, W, H, T] -+ h = ggml_rms_norm(ctx, h, 1e-12f); -+ h = ggml_mul(ctx, h, w); -+ h = ggml_cont(ctx, ggml_permute(ctx, h, 3, 0, 1, 2)); // [W, H, T, C] -+ ggml_set_name(h, "out"); -+ return h; -+ } -+}; -+ -+// MUL -> ADD (-> ADD) with broadcast operands: adaLN modulate x + x*scale + shift, gated residual r + y*g -+struct test_mul_add_chain : public test_case { -+ const std::array ne; -+ const int mode; // 0: x + x*s, 1: x + x*s + shift, 2: r + y*g, 3: (y*g) + r with views of a fused qkv-like tensor -+ -+ std::string op_desc(ggml_tensor * t) override { -+ GGML_UNUSED(t); -+ return "MUL_ADD_CHAIN"; -+ } -+ -+ bool run_whole_graph() override { return true; } -+ -+ std::string vars() override { -+ return VARS_TO_STR2(ne, mode); -+ } -+ -+ test_mul_add_chain(std::array ne = {64, 33, 2, 1}, int mode = 1) : ne(ne), mode(mode) {} -+ -+ ggml_tensor * build_graph(ggml_context * ctx) override { -+ ggml_tensor * x = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne.data()); -+ ggml_tensor * s = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, ne[0], 1, ne[2], ne[3]); -+ ggml_set_name(x, "x"); ggml_set_name(s, "s"); -+ ggml_tensor * out; -+ if (mode <= 1) { -+ out = ggml_add(ctx, x, ggml_mul(ctx, x, s)); -+ if (mode == 1) { -+ ggml_tensor * sh = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, ne[0], 1, ne[2], ne[3]); -+ ggml_set_name(sh, "shift"); -+ out = ggml_add(ctx, out, sh); -+ } -+ } else { -+ ggml_tensor * r = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne.data()); -+ ggml_set_name(r, "r"); -+ ggml_tensor * g = s; -+ if (mode == 3) { -+ ggml_tensor * big = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, 3 * ne[0], 1, ne[2], ne[3]); -+ ggml_set_name(big, "big"); -+ g = ggml_view_4d(ctx, big, ne[0], 1, ne[2], ne[3], big->nb[1], big->nb[2], big->nb[3], ne[0] * sizeof(float)); -+ } -+ ggml_tensor * m = ggml_mul(ctx, x, g); -+ out = mode == 2 ? ggml_add(ctx, r, m) : ggml_add(ctx, m, r); -+ } -+ ggml_set_name(out, "out"); -+ return out; -+ } -+}; -+ - // GGML_OP_L2_NORM x N: independent same-shape norms in one graph (strided qkv views or - // contiguous), consuming adds nested so the norms stay adjacent in the graph. - struct test_l2_norm_batch : public test_case { -@@ -9233,6 +9360,12 @@ static std::vector> make_test_cases_eval() { - test_cases.emplace_back(new test_rope_pe_permute(type, {128, 4, 9, 1}, 96, scale, view, 1e-6f)); - test_cases.emplace_back(new test_rope_pe_permute(type, {32, 2, 5, 1}, 16, scale, view, 1e-5f)); - test_cases.emplace_back(new test_rope_pe_permute(type, {48, 2, 5, 1}, 32, scale, view, 1e-5f)); -+ // interleaved pairing (Flux / Qwen-Image / Z-Image style DiT attention) -+ test_cases.emplace_back(new test_rope_pe_permute(type, {128, 4, 9, 1}, 128, scale, view, -1.0f, true)); -+ test_cases.emplace_back(new test_rope_pe_permute(type, {128, 3, 5, 2}, 96, scale, view, -1.0f, true)); -+ test_cases.emplace_back(new test_rope_pe_permute(type, {64, 3, 7, 2}, 64, scale, view, -1.0f, true)); -+ test_cases.emplace_back(new test_rope_pe_permute(type, {64, 32, 37, 1}, 48, scale, view, 1e-5f, true)); -+ test_cases.emplace_back(new test_rope_pe_permute(type, {128, 4, 9, 1}, 128, scale, view, 1e-6f, true)); - } - } - } -@@ -10815,6 +10948,27 @@ static std::vector> make_test_cases_eval() { - test_cases.emplace_back(new test_group_norm(GGML_TYPE_F32, {9, 9, 1280, 1})); - test_cases.emplace_back(new test_group_norm_mul_add(GGML_TYPE_F32, {64, 64, 320, 1})); - test_cases.emplace_back(new test_group_norm_mul_add(GGML_TYPE_F32, {9, 9, 1280, 1})); -+ test_cases.emplace_back(new test_group_norm(GGML_TYPE_F32, {32, 48, 128, 2})); -+ test_cases.emplace_back(new test_group_norm(GGML_TYPE_F32, {33, 17, 100, 1})); -+ for (bool silu : {false, true}) { -+ test_cases.emplace_back(new test_group_norm_affine({64, 64, 320, 1}, 32, silu)); -+ test_cases.emplace_back(new test_group_norm_affine({32, 48, 128, 2}, 32, silu)); -+ test_cases.emplace_back(new test_group_norm_affine({9, 9, 1280, 1}, 32, silu)); -+ test_cases.emplace_back(new test_group_norm_affine({33, 17, 100, 1}, 32, silu)); -+ test_cases.emplace_back(new test_group_norm_affine({16, 16, 64, 1}, 32, silu)); -+ test_cases.emplace_back(new test_group_norm_affine({128, 96, 128, 1}, 32, silu)); -+ test_cases.emplace_back(new test_group_norm_affine({64, 64, 320, 1}, 32, silu, true)); -+ test_cases.emplace_back(new test_group_norm_affine({33, 17, 100, 2}, 32, silu, true)); -+ } -+ for (int mode : {0, 1, 2, 3}) { -+ test_cases.emplace_back(new test_mul_add_chain({64, 33, 2, 1}, mode)); -+ test_cases.emplace_back(new test_mul_add_chain({3072, 17, 1, 1}, mode)); -+ test_cases.emplace_back(new test_mul_add_chain({129, 5, 3, 2}, mode)); -+ } -+ for (int64_t C : {32, 96, 128, 144, 192, 288, 384, 576, 1000, 1024, 1152, 2100}) { -+ test_cases.emplace_back(new test_rms_norm_channel({32, 24, 1, C}, true)); -+ test_cases.emplace_back(new test_rms_norm_channel({17, 9, 3, C}, false)); -+ } - test_cases.emplace_back(new test_acc(GGML_TYPE_F32, {256, 17, 1, 1}, {256, 16, 1, 1}, -1)); - test_cases.emplace_back(new test_acc(GGML_TYPE_F32, {256, 17, 2, 3}, {256, 16, 2, 3}, -1)); - test_cases.emplace_back(new test_acc(GGML_TYPE_F32, {256, 17, 2, 3}, {128, 16, 2, 3}, -1)); -@@ -11303,6 +11457,13 @@ static std::vector> make_test_cases_eval() { - static std::vector> make_test_cases_perf() { - std::vector> test_cases; - -+ // Wan / Qwen-Image VAE channel RMS norm at whole-image and tile sizes -+ for (int64_t C : {144, 288, 576, 1152}) { -+ for (int64_t S : {16, 32, 64, 128, 256}) { -+ test_cases.emplace_back(new test_rms_norm_channel({S, S, 1, C}, true)); -+ } -+ } -+ - // SWIGLU at a 27B-class FFN width, fused [gate|up] vs split operands - // note: same bytes either way, so a backend that indexes them differently shows it here - for (ggml_type type : {GGML_TYPE_F16, GGML_TYPE_F32}) { -diff --git a/tests/test-fusion-select.cpp b/tests/test-fusion-select.cpp -new file mode 100644 -index 00000000..16746c56 ---- /dev/null -+++ b/tests/test-fusion-select.cpp -@@ -0,0 +1,248 @@ -+// Checks for the fused CUDA paths of the generic DiT / VAE patch: -+// 1. graphs built with ggml_build_forward_select: a backend may only fuse nodes that are flagged -+// GGML_TENSOR_FLAG_COMPUTE, so the selected tensors must equal a graph that holds only them and the -+// inactive tensors must stay untouched (MUL -> ADD -> ADD chain and the channel RMS norm chain). -+// 2. CPU: ggml_rope_pe_permute with interleaved pairing is bit-identical to ggml_rope_apply(interleaved) -+// -> ggml_scale -> ggml_cpy, the chain it replaces. -+// Returns 0 when every check passes (backends without the ops skip their checks). -+ -+#include "ggml.h" -+#include "ggml-alloc.h" -+#include "ggml-backend.h" -+#include "ggml-cpu.h" -+ -+#include -+#include -+#include -+#include -+#include -+#include -+#include -+#include -+ -+static int n_fail = 0; -+static int n_pass = 0; -+ -+static void check(bool ok, const std::string & what) { -+ printf(" %-72s %s\n", what.c_str(), ok ? "OK" : "FAIL"); -+ (ok ? n_pass : n_fail)++; -+} -+ -+static std::vector rand_vec(size_t n, uint32_t seed, float lo = -2.0f, float hi = 2.0f) { -+ std::mt19937 rng(seed); -+ std::uniform_real_distribution d(lo, hi); -+ std::vector v(n); -+ for (auto & x : v) { -+ x = d(rng); -+ } -+ return v; -+} -+ -+static std::vector get_f32(const ggml_tensor * t) { -+ std::vector v(ggml_nelements(t)); -+ ggml_backend_tensor_get(t, v.data(), 0, ggml_nbytes(t)); -+ return v; -+} -+ -+static void fill(ggml_tensor * t, float value) { -+ std::vector v(ggml_nelements(t), value); -+ ggml_backend_tensor_set(t, v.data(), 0, ggml_nbytes(t)); -+} -+ -+// A graph whose tensors each own their memory (no allocator reuse), so an unwritten output keeps its sentinel. -+struct graph_case { -+ ggml_context * ctx = nullptr; -+ ggml_backend_buffer_t buf = nullptr; -+ std::vector inputs; -+ std::vector outs; // candidates for ggml_build_forward_select -+ -+ ~graph_case() { -+ if (buf) { -+ ggml_backend_buffer_free(buf); -+ } -+ if (ctx) { -+ ggml_free(ctx); -+ } -+ } -+}; -+ -+using builder_t = std::function &, std::vector &)>; -+ -+// run `build` on `backend`: select output `idx` (n_sel = number of candidates passed to the select) or, with -+// idx < 0, expand only candidate `ref`. Returns the candidate values after compute (all candidates). -+static std::vector> run_case(ggml_backend_t backend, const builder_t & build, int idx, int ref, -+ const std::vector> & input_data) { -+ graph_case g; -+ ggml_init_params ip = { 64 * ggml_tensor_overhead() + ggml_graph_overhead(), nullptr, true }; -+ g.ctx = ggml_init(ip); -+ build(g.ctx, g.inputs, g.outs); -+ g.buf = ggml_backend_alloc_ctx_tensors(g.ctx, backend); -+ for (size_t k = 0; k < g.inputs.size(); ++k) { -+ ggml_backend_tensor_set(g.inputs[k], input_data[k].data(), 0, ggml_nbytes(g.inputs[k])); -+ } -+ for (ggml_tensor * t = ggml_get_first_tensor(g.ctx); t; t = ggml_get_next_tensor(g.ctx, t)) { -+ if (t->op != GGML_OP_NONE && t->view_src == nullptr) { -+ fill(t, -77.0f); -+ } -+ } -+ ggml_cgraph * gf = ggml_new_graph(g.ctx); -+ if (idx >= 0) { -+ ggml_build_forward_select(gf, g.outs.data(), (int) g.outs.size(), idx); -+ } else { -+ ggml_build_forward_expand(gf, g.outs[ref]); -+ } -+ GGML_ASSERT(ggml_backend_graph_compute(backend, gf) == GGML_STATUS_SUCCESS); -+ std::vector> res; -+ for (ggml_tensor * t : g.outs) { -+ res.push_back(get_f32(t)); -+ } -+ return res; -+} -+ -+static bool bit_equal(const std::vector & a, const std::vector & b) { -+ return a.size() == b.size() && memcmp(a.data(), b.data(), a.size() * sizeof(float)) == 0; -+} -+ -+static bool all_sentinel(const std::vector & a) { -+ for (float x : a) { -+ if (x != -77.0f) { -+ return false; -+ } -+ } -+ return true; -+} -+ -+// every select index: the selected candidate equals the expand-only reference and every candidate not -+// needed by it keeps the sentinel. `needed[idx]` lists the candidates the selected one depends on. -+static void check_select(ggml_backend_t backend, const char * name, const builder_t & build, -+ const std::vector> & input_data, const std::vector> & needed) { -+ for (int idx = 0; idx < (int) needed.size(); ++idx) { -+ const auto ref = run_case(backend, build, -1, idx, input_data); -+ const auto sel = run_case(backend, build, idx, idx, input_data); -+ bool inactive_ok = true; -+ for (int k = 0; k < (int) sel.size(); ++k) { -+ bool need = false; -+ for (int n : needed[idx]) { -+ need |= n == k; -+ } -+ if (!need) { -+ inactive_ok &= all_sentinel(sel[k]); -+ } -+ } -+ check(bit_equal(ref[idx], sel[idx]), std::string(name) + " select " + std::to_string(idx) + ": selected output computed"); -+ check(inactive_ok, std::string(name) + " select " + std::to_string(idx) + ": inactive outputs untouched"); -+ } -+} -+ -+static void test_select(ggml_backend_t backend) { -+ // MUL -> ADD -> ADD: x*s + x + shift (adaLN modulate), candidates {mul, add, add2} -+ { -+ const int64_t ne0 = 64, ne1 = 33; -+ builder_t build = [&](ggml_context * ctx, std::vector & in, std::vector & out) { -+ ggml_tensor * x = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, ne0, ne1); -+ ggml_tensor * s = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, ne0, 1); -+ ggml_tensor * sh = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, ne0, 1); -+ ggml_tensor * m = ggml_mul(ctx, x, s); -+ ggml_tensor * a = ggml_add(ctx, m, x); -+ ggml_tensor * a2 = ggml_add(ctx, a, sh); -+ in = { x, s, sh }; -+ out = { m, a, a2 }; -+ }; -+ const std::vector> data = { rand_vec(ne0 * ne1, 1), rand_vec(ne0, 2), rand_vec(ne0, 3) }; -+ check_select(backend, "MUL_ADD_CHAIN", build, data, { { 0 }, { 0, 1 }, { 0, 1, 2 } }); -+ } -+ // channel RMS norm: CONT(PERMUTE(x)) -> RMS_NORM -> MUL(w) -> CONT(PERMUTE), P = 16384 positions so the -+ // fusion engages at its default size threshold; candidates {cont1, rms, mul, cont2} -+ { -+ const int64_t W = 128, H = 128, C = 8; -+ builder_t build = [&](ggml_context * ctx, std::vector & in, std::vector & out) { -+ ggml_tensor * x = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, W, H, 1, C); -+ ggml_tensor * w = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, C); -+ ggml_tensor * c1 = ggml_cont(ctx, ggml_permute(ctx, x, 1, 2, 3, 0)); -+ ggml_tensor * r = ggml_rms_norm(ctx, c1, 1e-6f); -+ ggml_tensor * m = ggml_mul(ctx, r, w); -+ ggml_tensor * c2 = ggml_cont(ctx, ggml_permute(ctx, m, 3, 0, 1, 2)); -+ in = { x, w }; -+ out = { c1, r, m, c2 }; -+ }; -+ const std::vector> data = { rand_vec(W * H * C, 4), rand_vec(C, 5, 0.5f, 1.5f) }; -+ check_select(backend, "RMS_NORM_CHANNEL", build, data, { { 0 }, { 0, 1 }, { 0, 1, 2 }, { 0, 1, 2, 3 } }); -+ } -+} -+ -+// CPU: interleaved ggml_rope_pe_permute vs ggml_rope_apply -> ggml_scale -> ggml_cpy, bitwise -+static void test_rope_interleaved_cpu(ggml_backend_t cpu) { -+ const int64_t shapes[][4] = { { 64, 3, 7, 1 }, { 128, 2, 5, 2 }, { 80, 4, 3, 1 } }; -+ const float scales[] = { 1.0f, 1.0f / 128.0f }; -+ const ggml_type types[] = { GGML_TYPE_F32, GGML_TYPE_F16 }; -+ uint32_t seed = 100; -+ for (const auto & ne : shapes) { -+ for (float scale : scales) { -+ for (ggml_type type : types) { -+ const int64_t D = ne[0], Hh = ne[1], L = ne[2], N = ne[3]; -+ const auto xd = rand_vec(D * Hh * L * N, seed++); -+ const auto ped = rand_vec(4 * (D / 2) * L, seed++, -1.0f, 1.0f); // not an exact rotation -+ std::vector outs[2]; -+ for (int arm = 0; arm < 2; ++arm) { -+ ggml_init_params ip = { 16 * ggml_tensor_overhead() + ggml_graph_overhead(), nullptr, true }; -+ ggml_context * ctx = ggml_init(ip); -+ ggml_tensor * x = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, D, Hh, L, N); -+ ggml_tensor * pe = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, 2, 2, D / 2, L); -+ ggml_tensor * y; -+ if (arm == 0) { -+ y = ggml_rope_apply(ctx, x, pe, true); -+ if (scale != 1.0f) { -+ y = ggml_scale(ctx, y, scale); -+ } -+ if (type != GGML_TYPE_F32) { -+ y = ggml_cpy(ctx, y, ggml_new_tensor(ctx, type, GGML_MAX_DIMS, y->ne)); -+ } -+ } else { -+ y = ggml_rope_pe_permute_set_interleaved(ggml_rope_pe_permute(ctx, x, pe, (int) D, scale, type), true); -+ } -+ ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors(ctx, cpu); -+ ggml_backend_tensor_set(x, xd.data(), 0, ggml_nbytes(x)); -+ ggml_backend_tensor_set(pe, ped.data(), 0, ggml_nbytes(pe)); -+ ggml_cgraph * gf = ggml_new_graph(ctx); -+ ggml_build_forward_expand(gf, y); -+ GGML_ASSERT(ggml_backend_graph_compute(cpu, gf) == GGML_STATUS_SUCCESS); -+ outs[arm].resize(ggml_nbytes(y)); -+ ggml_backend_tensor_get(y, outs[arm].data(), 0, ggml_nbytes(y)); -+ ggml_backend_buffer_free(buf); -+ ggml_free(ctx); -+ } -+ char name[160]; -+ snprintf(name, sizeof(name), "CPU ROPE_PE_PERMUTE interleaved [%lld,%lld,%lld,%lld] scale=%g %s", -+ (long long) D, (long long) Hh, (long long) L, (long long) N, scale, ggml_type_name(type)); -+ check(outs[0].size() == outs[1].size() && memcmp(outs[0].data(), outs[1].data(), outs[0].size()) == 0, name); -+ } -+ } -+ } -+} -+ -+int main() { -+ ggml_backend_load_all(); -+ -+ ggml_backend_t cpu = ggml_backend_init_by_type(GGML_BACKEND_DEVICE_TYPE_CPU, nullptr); -+ GGML_ASSERT(cpu != nullptr); -+ printf("backend CPU\n"); -+ test_rope_interleaved_cpu(cpu); -+ test_select(cpu); -+ ggml_backend_free(cpu); -+ -+ for (size_t i = 0; i < ggml_backend_dev_count(); ++i) { -+ ggml_backend_dev_t dev = ggml_backend_dev_get(i); -+ if (ggml_backend_dev_type(dev) == GGML_BACKEND_DEVICE_TYPE_CPU) { -+ continue; -+ } -+ ggml_backend_t be = ggml_backend_dev_init(dev, nullptr); -+ if (be == nullptr) { -+ continue; -+ } -+ printf("backend %s\n", ggml_backend_name(be)); -+ test_select(be); -+ ggml_backend_free(be); -+ } -+ printf("%d passed, %d failed\n", n_pass, n_fail); -+ return n_fail == 0 ? 0 : 1; -+} From 122b7151efaf1cec608d5a160fa78ca4e7174460 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Sun, 11 Oct 2026 18:13:19 -0700 Subject: [PATCH 08/10] Trim comments in fused RoPE attention and conv paths --- src/core/ggml_extend.cpp | 4 +--- src/core/ggml_extend.h | 3 +-- src/model/common/rope.hpp | 17 ++++++----------- src/model/vae/wan_vae.hpp | 3 +-- 4 files changed, 9 insertions(+), 18 deletions(-) diff --git a/src/core/ggml_extend.cpp b/src/core/ggml_extend.cpp index fd98a73f71..947eee854d 100644 --- a/src/core/ggml_extend.cpp +++ b/src/core/ggml_extend.cpp @@ -451,8 +451,7 @@ 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 { - // ggml_conv_2d, except that the final [OC, N] -> [N, OC] relayout is skipped when it moves no data - // (a batch of one, every VAE decode and most U-Net steps): same values, one full copy fewer per conv. + // 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, @@ -852,7 +851,6 @@ ggml_tensor* ggml_ext_attention_prepared(ggml_context* ctx, float scale = (1.0f / sqrt((float)d_head)); if (mask != nullptr) { - // same mask preparation as ggml_ext_attention_ext's flash path 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])); } diff --git a/src/core/ggml_extend.h b/src/core/ggml_extend.h index 512bafe9b5..165eb673c9 100644 --- a/src/core/ggml_extend.h +++ b/src/core/ggml_extend.h @@ -236,8 +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); it gets the -// same query-dimension repeat and F16 cast. +// 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, diff --git a/src/model/common/rope.hpp b/src/model/common/rope.hpp index f1e9d7b61a..a0b342240d 100644 --- a/src/model/common/rope.hpp +++ b/src/model/common/rope.hpp @@ -1126,7 +1126,6 @@ namespace Rope { return x_out; } - // SD_FUSED_ROPE_ATTN=0 turns the fused attention prologue below off. __STATIC_INLINE__ bool fused_rope_attention_enabled() { static const bool enabled = [] { const char* e = getenv("SD_FUSED_ROPE_ATTN"); @@ -1135,11 +1134,9 @@ namespace Rope { return enabled; } - // Flash attention with the RoPE, head-major relayout, kv scale and F16 cast of q/k/v done by one - // ggml_rope_pe_permute op per tensor, instead of RoPE + relayout, then scale, then cast (k), and - // relayout, scale, cast (v). Same values as the unfused flash path: the op rounds each product, - // the sum, the scale and the conversion separately, exactly like ggml_rope_apply -> ggml_scale -> - // ggml_cpy. GPU backends only (CUDA / HIP); returns nullptr when it does not apply. + // 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, @@ -1152,13 +1149,11 @@ namespace Rope { return nullptr; #else auto backend = ctx->backend; - // the same attention-scale override ggml_ext_attention_ext(GGMLRunnerContext*, ...) applies - // (it keeps large k/v finite in F16), so both paths compute identical values + // 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 pairing only: the non-interleaved form of ggml_rope_apply reads two of the four table - // entries, so it matches the fused op only for exact rotation tables. + // 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; @@ -1191,7 +1186,7 @@ namespace Rope { !ggml_backend_supports_op(backend, v_hm)) { return nullptr; } - // [d_head, L, heads, N] -> [d_head, L, heads * N], the shapes the unfused path hands flash attention + // [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); diff --git a/src/model/vae/wan_vae.hpp b/src/model/vae/wan_vae.hpp index 3dea88a82e..9a05839cfd 100644 --- a/src/model/vae/wan_vae.hpp +++ b/src/model/vae/wan_vae.hpp @@ -147,8 +147,7 @@ namespace WAN { using Conv2d::Conv2d; ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { - // Swapping the time and channel axes moves no data when one of them is 1 (a single frame, the - // image VAE case): ggml_ext_cont then keeps the view instead of copying the whole tensor. + // 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_ext_cont(ctx->ggml_ctx, x_swapped); From f017bf3ddae72a895ee76b5809cd8d887b7e6cde Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 11 Oct 2026 19:04:43 -0700 Subject: [PATCH 09/10] ggml-cuda: tighten the GroupNorm and channel RMS norm fusions - fused GroupNorm and channel RMS norm assert eps >= 0 like the unfused ops - channel RMS norm: both CONTs must keep their PERMUTE's shape (a CONT_4D may reshape) - GroupNorm affine: w and b may not overlap the fused output, x only element for element (the generic range check skips leaf sources) - GroupNorm affine (+ SILU): accept the in-place chain sd.cpp builds (MUL / ADD / SILU are whole views of the GROUP_NORM result), keeping the COMPUTE, OUTPUT and use-count checks --- .../0007-ggml-cuda-generic-dit-vae.patch | 66 +++++++++++++++---- 1 file changed, 55 insertions(+), 11 deletions(-) diff --git a/scripts/unsloth/ggml-patches/0007-ggml-cuda-generic-dit-vae.patch b/scripts/unsloth/ggml-patches/0007-ggml-cuda-generic-dit-vae.patch index 3c86471ca3..7f3e7af43c 100644 --- a/scripts/unsloth/ggml-patches/0007-ggml-cuda-generic-dit-vae.patch +++ b/scripts/unsloth/ggml-patches/0007-ggml-cuda-generic-dit-vae.patch @@ -69,7 +69,7 @@ index bdd36aea..2883f5df 100644 if (D > n_rot) { memcpy(y.data() + n_rot, x + n_rot, (D - n_rot) * sizeof(float)); diff --git a/src/ggml-cuda/ggml-cuda.cu b/src/ggml-cuda/ggml-cuda.cu -index 831d0de8..2df04786 100644 +index 831d0de8..558767ff 100644 --- a/src/ggml-cuda/ggml-cuda.cu +++ b/src/ggml-cuda/ggml-cuda.cu @@ -37,6 +37,9 @@ @@ -82,7 +82,7 @@ index 831d0de8..2df04786 100644 #include "ggml-cuda/opt-step-adamw.cuh" #include "ggml-cuda/opt-step-sgd.cuh" #include "ggml-cuda/out-prod.cuh" -@@ -4313,6 +4316,71 @@ static int ggml_cuda_try_fuse_cublas_epilogue(ggml_backend_cuda_context & ctx, g +@@ -4313,6 +4316,109 @@ static int ggml_cuda_try_fuse_cublas_epilogue(ggml_backend_cuda_context & ctx, g return consumed; } @@ -114,7 +114,33 @@ index 831d0de8..2df04786 100644 + } + int idx[4] = { i, i_mul, i_add, -1 }; + const enum ggml_op ops[4] = { GGML_OP_GROUP_NORM, GGML_OP_MUL, GGML_OP_ADD, GGML_OP_UNARY }; -+ if (!ggml_can_fuse_ext(cgraph, idx, ops, 3)) { ++ // ggml_can_fuse_ext without its view rule: sd.cpp builds the affine and SiLU in place ++ // (ggml_mul_inplace(gn, w) ...), so MUL / ADD / SILU are whole views of the GROUP_NORM result. ++ // Such a view writes gn's memory, every reader of that memory is a counted use of a chain node, ++ // and the fused kernel leaves the same final values there. ++ auto chain_ok = [&](int n) { ++ for (int k = 0; k < n; ++k) { ++ const ggml_tensor * node = cgraph->nodes[idx[k]]; ++ if (node->op != ops[k] || (node->flags & GGML_TENSOR_FLAG_COMPUTE) == 0) { ++ return false; ++ } ++ if (k == 0 ? node->view_src != nullptr ++ : node->view_src != nullptr && (node->view_src != gn || node->view_offs != 0 || node->data != gn->data)) { ++ return false; ++ } ++ if (k < n - 1 && (ggml_node_get_use_count(cgraph, idx[k]) != 1 || (node->flags & GGML_TENSOR_FLAG_OUTPUT))) { ++ return false; ++ } ++ if (k > 0) { ++ const ggml_tensor * prev = cgraph->nodes[idx[k - 1]]; ++ if ((node->src[0] != prev && node->src[1] != prev) || !ggml_are_same_shape(node, prev)) { ++ return false; ++ } ++ } ++ } ++ return true; ++ }; ++ if (!chain_ok(3)) { + return 0; + } + ggml_tensor * mul = cgraph->nodes[i_mul]; @@ -136,7 +162,7 @@ index 831d0de8..2df04786 100644 + const ggml_tensor * u = cgraph->nodes[i_u]; + idx[3] = i_u; + if (u->op == GGML_OP_UNARY && ggml_get_unary_op(u) == GGML_UNARY_OP_SILU && u->type == GGML_TYPE_F32 && -+ ggml_is_contiguous(u) && ggml_can_fuse_ext(cgraph, idx, ops, 4)) { ++ ggml_is_contiguous(u) && chain_ok(4)) { + n_fused = 4; + } + } @@ -147,6 +173,18 @@ index 831d0de8..2df04786 100644 + if (!ggml_cuda_check_fusion_memory_ranges(cgraph, i, last - i + 1, out_nodes, 1)) { + return 0; + } ++ // the check above skips leaf sources: w and b are read while out is written, so they may not ++ // overlap it at all; x only at the same element (the last pass reads x[j] before writing out[j]) ++ auto overlaps = [&](const ggml_tensor * t) { ++ const char * o0 = (const char *) out->data; ++ const char * t0 = (const char *) t->data; ++ return t0 < o0 + ggml_nbytes(out) && o0 < t0 + ggml_nbytes(t); ++ }; ++ const ggml_tensor * x = gn->src[0]; ++ if (overlaps(w) || overlaps(b) || ++ (overlaps(x) && !(x->data == out->data && ggml_are_same_shape(x, out) && ggml_is_contiguous(x)))) { ++ return 0; ++ } + ggml_cuda_op_group_norm_par(ctx, gn, w, b, n_fused == 4, out); + return last - i; +} @@ -154,7 +192,7 @@ index 831d0de8..2df04786 100644 static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph * cgraph, int i) { static bool disable_fusion = getenv("GGML_CUDA_DISABLE_FUSION") != nullptr && std::atoi(getenv("GGML_CUDA_DISABLE_FUSION")); -@@ -4322,6 +4390,20 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph +@@ -4322,6 +4428,20 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph ggml_tensor * node = cgraph->nodes[i]; @@ -175,7 +213,7 @@ index 831d0de8..2df04786 100644 if (node->op == GGML_OP_MUL) { ggml_cuda_moe_weighted_reduction_match match; if (ggml_cuda_match_moe_weighted_reduction(cgraph, i, match)) { -@@ -4334,6 +4416,13 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph +@@ -4334,6 +4454,13 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph } } @@ -191,10 +229,10 @@ index 831d0de8..2df04786 100644 ggml_cuda_gated_delta_net_fused_cache fused_state_cpy; diff --git a/src/ggml-cuda/group-norm-par.cu b/src/ggml-cuda/group-norm-par.cu new file mode 100644 -index 00000000..3e640e62 +index 00000000..dbc652df --- /dev/null +++ b/src/ggml-cuda/group-norm-par.cu -@@ -0,0 +1,173 @@ +@@ -0,0 +1,174 @@ +#include "group-norm-par.cuh" +#include "unary.cuh" + @@ -333,6 +371,7 @@ index 00000000..3e640e62 + const int num_groups = gn->op_params[0]; + float eps; + memcpy(&eps, gn->op_params + 1, sizeof(float)); ++ GGML_ASSERT(eps >= 0.0f); + + const int group_size = (int) gn_par_group_size(gn); + const int ne_elements = (int) ggml_nelements(src0); @@ -583,10 +622,10 @@ index 00000000..8d9752ca +int ggml_cuda_try_fuse_mul_add(ggml_backend_cuda_context & ctx, ggml_cgraph * cgraph, int i); diff --git a/src/ggml-cuda/norm-channel.cu b/src/ggml-cuda/norm-channel.cu new file mode 100644 -index 00000000..9c0b61fd +index 00000000..c1a8819d --- /dev/null +++ b/src/ggml-cuda/norm-channel.cu -@@ -0,0 +1,254 @@ +@@ -0,0 +1,259 @@ +#include "norm-channel.cuh" +#include "ggml-impl.h" + @@ -746,6 +785,10 @@ index 00000000..9c0b61fd + if (nc_debug()) { GGML_LOG_INFO("norm-channel: C=%lld next op %s (%s)\n", (long long) C, ggml_op_desc(rms), rms->name); } + NC_REJECT("check 5a"); + } ++ // the CONT must only make the permuted view contiguous (a CONT_4D may reshape it) ++ if (!ggml_are_same_shape(cont1, pv)) { ++ NC_REJECT("check 5e"); ++ } + if (rms->src[0] != cont1) { + NC_REJECT("check 5b"); + } @@ -769,7 +812,7 @@ index 00000000..9c0b61fd + ggml_tensor * pv2 = cgraph->nodes[i_pv2]; + ggml_tensor * cont2 = cgraph->nodes[i_pv2 + 1]; + bool back_ok = pv2->op == GGML_OP_PERMUTE && pv2->src[0] == mul && cont2->op == GGML_OP_CONT && cont2->src[0] == pv2 && -+ cont2->type == GGML_TYPE_F32 && ggml_are_same_shape(cont2, x) && ++ cont2->type == GGML_TYPE_F32 && ggml_are_same_shape(cont2, x) && ggml_are_same_shape(cont2, pv2) && + ggml_is_contiguous(cont2) && ggml_is_contiguous(mul); + // the graph does not count uses of views, so check by hand that cont2 is the only reader of pv2 + for (int j = 0; j < cgraph->n_nodes && back_ok; ++j) { @@ -823,6 +866,7 @@ index 00000000..9c0b61fd + } + float eps; + memcpy(&eps, rms->op_params, sizeof(float)); ++ GGML_ASSERT(eps >= 0.0f); + + const int threads = 128; + const int64_t blocks = (P + threads - 1) / threads; From 5da68dc309ee2858bbe522cd2f55fba28dc12172 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 11 Oct 2026 22:39:50 -0700 Subject: [PATCH 10/10] Trim comments --- .../0007-ggml-cuda-generic-dit-vae.patch | 17 +++++++---------- 1 file changed, 7 insertions(+), 10 deletions(-) diff --git a/scripts/unsloth/ggml-patches/0007-ggml-cuda-generic-dit-vae.patch b/scripts/unsloth/ggml-patches/0007-ggml-cuda-generic-dit-vae.patch index 7f3e7af43c..b7fd0dfcb4 100644 --- a/scripts/unsloth/ggml-patches/0007-ggml-cuda-generic-dit-vae.patch +++ b/scripts/unsloth/ggml-patches/0007-ggml-cuda-generic-dit-vae.patch @@ -69,7 +69,7 @@ index bdd36aea..2883f5df 100644 if (D > n_rot) { memcpy(y.data() + n_rot, x + n_rot, (D - n_rot) * sizeof(float)); diff --git a/src/ggml-cuda/ggml-cuda.cu b/src/ggml-cuda/ggml-cuda.cu -index 831d0de8..558767ff 100644 +index 831d0de8..8f4e6f3d 100644 --- a/src/ggml-cuda/ggml-cuda.cu +++ b/src/ggml-cuda/ggml-cuda.cu @@ -37,6 +37,9 @@ @@ -82,7 +82,7 @@ index 831d0de8..558767ff 100644 #include "ggml-cuda/opt-step-adamw.cuh" #include "ggml-cuda/opt-step-sgd.cuh" #include "ggml-cuda/out-prod.cuh" -@@ -4313,6 +4316,109 @@ static int ggml_cuda_try_fuse_cublas_epilogue(ggml_backend_cuda_context & ctx, g +@@ -4313,6 +4316,106 @@ static int ggml_cuda_try_fuse_cublas_epilogue(ggml_backend_cuda_context & ctx, g return consumed; } @@ -114,10 +114,8 @@ index 831d0de8..558767ff 100644 + } + int idx[4] = { i, i_mul, i_add, -1 }; + const enum ggml_op ops[4] = { GGML_OP_GROUP_NORM, GGML_OP_MUL, GGML_OP_ADD, GGML_OP_UNARY }; -+ // ggml_can_fuse_ext without its view rule: sd.cpp builds the affine and SiLU in place -+ // (ggml_mul_inplace(gn, w) ...), so MUL / ADD / SILU are whole views of the GROUP_NORM result. -+ // Such a view writes gn's memory, every reader of that memory is a counted use of a chain node, -+ // and the fused kernel leaves the same final values there. ++ // ggml_can_fuse_ext minus its view rule: sd.cpp's in-place MUL/ADD/SILU are views of the GROUP_NORM ++ // result, every reader of that memory is a counted chain use, and the fused kernel writes the same values. + auto chain_ok = [&](int n) { + for (int k = 0; k < n; ++k) { + const ggml_tensor * node = cgraph->nodes[idx[k]]; @@ -173,8 +171,7 @@ index 831d0de8..558767ff 100644 + if (!ggml_cuda_check_fusion_memory_ranges(cgraph, i, last - i + 1, out_nodes, 1)) { + return 0; + } -+ // the check above skips leaf sources: w and b are read while out is written, so they may not -+ // overlap it at all; x only at the same element (the last pass reads x[j] before writing out[j]) ++ // leaf sources skipped above: w/b must not overlap out; x only elementwise (x[j] read before out[j] written) + auto overlaps = [&](const ggml_tensor * t) { + const char * o0 = (const char *) out->data; + const char * t0 = (const char *) t->data; @@ -192,7 +189,7 @@ index 831d0de8..558767ff 100644 static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph * cgraph, int i) { static bool disable_fusion = getenv("GGML_CUDA_DISABLE_FUSION") != nullptr && std::atoi(getenv("GGML_CUDA_DISABLE_FUSION")); -@@ -4322,6 +4428,20 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph +@@ -4322,6 +4425,20 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph ggml_tensor * node = cgraph->nodes[i]; @@ -213,7 +210,7 @@ index 831d0de8..558767ff 100644 if (node->op == GGML_OP_MUL) { ggml_cuda_moe_weighted_reduction_match match; if (ggml_cuda_match_moe_weighted_reduction(cgraph, i, match)) { -@@ -4334,6 +4454,13 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph +@@ -4334,6 +4451,13 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph } }