Commit 984eb06 for stable-diffusion.cpp
commit 984eb0686cbf5b7fc4f0063d9e0b1fb38d5c4f82
Author: leejet <leejet714@gmail.com>
Date: Wed Nov 4 02:01:41 2026 +0800
perf(cuda): fuse MOSS modulation and direct RoPE F16 output (#2269)
Co-authored-by: mirek190 <mirek190@gmail.com>
diff --git a/src/core/cuda/rms_affine.cu b/src/core/cuda/rms_affine.cu
index d9ac351..9f4192b 100644
--- a/src/core/cuda/rms_affine.cu
+++ b/src/core/cuda/rms_affine.cu
@@ -2,6 +2,7 @@
#include <cstdlib>
#include <cstring>
#include "bf16.h"
+#include "core/ggml_extensions.h"
#include "rms_affine.h"
namespace SD_CUDA_NAMESPACE {
@@ -18,6 +19,28 @@ namespace SD_CUDA_NAMESPACE {
return t && t->type == GGML_TYPE_F32 && ggml_is_contiguous(t);
}
+ static int modulate_match(ggml_backend_dev_t, const ggml_cgraph* graph, int first, void*) {
+ auto* g = const_cast<ggml_cgraph*>(graph);
+ if (disabled() || first + 2 >= ggml_graph_n_nodes(g))
+ return 0;
+ const auto* norm = ggml_graph_node(g, first);
+ const auto* affine = ggml_graph_node(g, first + 1);
+ const auto* modulate = ggml_graph_node(g, first + 2);
+ if (norm->op != GGML_OP_RMS_NORM || affine->op != GGML_OP_MUL || affine->src[0] != norm ||
+ !sd_ggml_modulate_describe(modulate) || modulate->src[0] != affine)
+ return 0;
+ const auto* x = norm->src[0];
+ if (!f32(x) || !f32(norm) || !f32(affine) || x->ne[0] < 32 || x->ne[0] > 8192 ||
+ ggml_nrows(x) > INT_MAX || !ggml_are_same_shape(x, norm) || !ggml_are_same_shape(x, affine))
+ return 0;
+ for (const auto* parameter : {affine->src[1], modulate->src[1], modulate->src[2]}) {
+ if (parameter && (!f32(parameter) || parameter == norm || parameter == affine ||
+ parameter->ne[0] != x->ne[0] || ggml_nelements(parameter) != x->ne[0]))
+ return 0;
+ }
+ return affine->src[1] ? 3 : 0;
+ }
+
static int match(ggml_backend_dev_t, const ggml_cgraph* graph, int first, void*) {
auto* g = const_cast<ggml_cgraph*>(graph);
if (disabled() || first + 3 >= ggml_graph_n_nodes(g))
@@ -44,13 +67,17 @@ namespace SD_CUDA_NAMESPACE {
static void optimize(ggml_backend_dev_t, ggml_cgraph* graph, void*) {
if (disabled())
return;
- for (int i = 2; i + 1 < ggml_graph_n_nodes(graph); ++i) {
+ for (int i = 2; i < ggml_graph_n_nodes(graph); ++i) {
const auto* mul = ggml_graph_node(graph, i);
- const auto* add = ggml_graph_node(graph, i + 1);
- if (mul->op != GGML_OP_MUL || add->op != GGML_OP_ADD || (add->src[0] != mul && add->src[1] != mul))
- continue;
+ if (!sd_ggml_modulate_describe(mul)) {
+ if (i + 1 >= ggml_graph_n_nodes(graph))
+ continue;
+ const auto* add = ggml_graph_node(graph, i + 1);
+ if (mul->op != GGML_OP_MUL || add->op != GGML_OP_ADD || (add->src[0] != mul && add->src[1] != mul))
+ continue;
+ }
const auto* affine = mul->src[0];
- if (affine->op != GGML_OP_MUL || affine->src[0]->op != GGML_OP_RMS_NORM)
+ if (!affine || affine->op != GGML_OP_MUL || !affine->src[0] || affine->src[0]->op != GGML_OP_RMS_NORM)
continue;
for (int j = i - 2; j >= 0 && j >= i - 32; --j) {
if (ggml_graph_node(graph, j) == affine->src[0] && ggml_graph_node(graph, j + 1) == affine) {
@@ -68,64 +95,136 @@ namespace SD_CUDA_NAMESPACE {
return value;
}
- template <int Block, bool Modulate, typename T = float>
- static __global__ void kernel(const float* x, const float* weight, const float* gate, const float* residual, T* dst, int width, float eps) {
+ enum class Kind { Residual,
+ SelfResidual,
+ Modulate };
+
+ struct Operation {
+ Kind kind = Kind::Residual;
+ const ggml_tensor* input;
+ const ggml_tensor* weight;
+ const ggml_tensor* gate;
+ const ggml_tensor* residual = nullptr;
+ const ggml_tensor* shift = nullptr;
+ ggml_tensor* output;
+ float eps;
+ };
+
+ static Operation describe(ggml_tensor* const* nodes, int count) {
+ Operation op;
+ op.input = nodes[0]->src[0];
+ op.weight = nodes[1]->src[1];
+ op.gate = nodes[2]->src[1];
+ op.output = nodes[count - 1];
+ std::memcpy(&op.eps, nodes[0]->op_params, sizeof(op.eps));
+ if (count == 3) {
+ op.kind = Kind::Modulate;
+ op.shift = op.output->src[2];
+ } else {
+ op.residual = op.output->src[0] == nodes[2] ? op.output->src[1] : op.output->src[0];
+ if (op.residual == nodes[1]) {
+ op.kind = Kind::SelfResidual;
+ op.residual = nullptr;
+ }
+ }
+ return op;
+ }
+
+ template <int Block, int Groups, Kind Op, typename T>
+ static __global__ void kernel(const float* x, const float* weight, const float* gate, const float* residual, const float* shift, T* dst, int width, float eps) {
const int col0 = threadIdx.x;
const int64_t base = int64_t(blockIdx.x) * width;
- float values[8], sum = 0.f;
- for (int j = 0, col = col0; col < width; ++j, col += Block) {
- values[j] = x[base + col];
- sum = __fmaf_rn(values[j], values[j], sum);
- }
- sum = warp_sum(sum);
+ float values[Groups][8];
__shared__ float sums[32];
const int lane = col0 % 32;
- if (lane == 0)
- sums[col0 / 32] = sum;
+#pragma unroll
+ for (int group = 0; group < Groups; ++group) {
+ float sum = 0.f;
+ for (int j = 0, col = col0 + group * Block; col < width; ++j, col += Block * Groups) {
+ values[group][j] = x[base + col];
+ sum = __fmaf_rn(values[group][j], values[group][j], sum);
+ }
+ sum = warp_sum(sum);
+ if (lane == 0)
+ sums[(col0 + group * Block) / 32] = sum;
+ }
+ // Groups preserve the native reduction tree when using fewer physical threads.
__syncthreads();
- sum = warp_sum(lane < Block / 32 ? sums[lane] : 0.f);
+ const float sum = warp_sum(lane < Block * Groups / 32 ? sums[lane] : 0.f);
const float scale = rsqrtf(sum / width + eps);
- for (int j = 0, col = col0; col < width; ++j, col += Block) {
- const float value = __fmul_rn(__fmul_rn(scale, values[j]), weight[col]);
- const float r = Modulate ? value : residual[base + col];
- dst[base + col] = T(__fadd_rn(__fmul_rn(value, gate[col]), r));
+#pragma unroll
+ for (int group = 0; group < Groups; ++group) {
+ for (int j = 0, col = col0 + group * Block; col < width; ++j, col += Block * Groups) {
+ float value = __fmul_rn(__fmul_rn(scale, values[group][j]), weight[col]);
+ if constexpr (Op == Kind::Modulate) {
+ value = __fadd_rn(value, __fmul_rn(value, gate[col]));
+ if (shift)
+ value = __fadd_rn(value, shift[col]);
+ } else {
+ const float r = Op == Kind::SelfResidual ? value : residual[base + col];
+ value = __fadd_rn(__fmul_rn(value, gate[col]), r);
+ }
+ dst[base + col] = T(value);
+ }
}
}
static bool overlaps(const ggml_tensor* a, const ggml_tensor* b) {
+ if (!b)
+ return false;
const auto x = reinterpret_cast<uintptr_t>(a->data), y = reinterpret_cast<uintptr_t>(b->data);
return x <= y ? y - x < ggml_nbytes(a) : x - y < ggml_nbytes(b);
}
- static ggml_status prepare(ggml_backend_t, ggml_tensor* const* nodes, int, void**, size_t* workspace, void*) {
- *workspace = ggml_nbytes(nodes[3]);
+ static ggml_status prepare(ggml_backend_t, ggml_tensor* const* nodes, int count, void**, size_t* workspace, void*) {
+ *workspace = ggml_nbytes(nodes[count - 1]);
return GGML_STATUS_SUCCESS;
}
static void destroy(void*, void*) {}
- static ggml_status execute(const ggml_ext_execution_context* execution, ggml_tensor* const* nodes, int, void*, void*) {
+ template <typename T, Kind Op>
+ static CUDAError launch(CUDAStream stream, const Operation& op, T* output) {
+ auto compute = kernel<256, 1, Op, T>;
+ int threads = 256;
+ if (op.input->ne[0] >= 1024) {
+ if constexpr (Op == Kind::Modulate) {
+ compute = kernel<512, 2, Op, T>;
+ threads = 512;
+ } else {
+ compute = kernel<1024, 1, Op, T>;
+ threads = 1024;
+ }
+ }
+ compute<<<unsigned(ggml_nrows(op.input)), threads, 0, stream>>>(
+ static_cast<const float*>(op.input->data), static_cast<const float*>(op.weight->data),
+ static_cast<const float*>(op.gate->data), op.residual ? static_cast<const float*>(op.residual->data) : nullptr,
+ op.shift ? static_cast<const float*>(op.shift->data) : nullptr, output, int(op.input->ne[0]), op.eps);
+ return cuda_last_error();
+ }
+
+ static ggml_status execute(const ggml_ext_execution_context* execution, ggml_tensor* const* nodes, int count, void*, void*) {
ggml_ext_gpu_context context;
if (!cuda_context(execution->backend, &context))
return GGML_STATUS_FAILED;
- const auto* x = nodes[0]->src[0];
- const auto* w = nodes[1]->src[1];
- const auto* gate = nodes[2]->src[1];
- auto* dst = nodes[3];
- const auto* residual = dst->src[0] == nodes[2] ? dst->src[1] : dst->src[0];
- const bool modulation = residual == nodes[1];
- const bool alias = (overlaps(dst, x) && dst->data != x->data) || overlaps(dst, w) || overlaps(dst, gate) ||
- (!modulation && overlaps(dst, residual) && dst->data != residual->data);
- auto* output = static_cast<float*>(alias ? ggml_backend_buffer_get_base(execution->workspace) : dst->data);
- float eps;
- std::memcpy(&eps, nodes[0]->op_params, sizeof(eps));
- auto launch = x->ne[0] < 1024 ? (modulation ? kernel<256, true> : kernel<256, false>)
- : (modulation ? kernel<1024, true> : kernel<1024, false>);
+ const auto op = describe(nodes, count);
+ auto* dst = op.output;
+ const bool alias = (overlaps(dst, op.input) && dst->data != op.input->data) ||
+ (overlaps(dst, op.residual) && dst->data != op.residual->data) ||
+ overlaps(dst, op.weight) || overlaps(dst, op.gate) || overlaps(dst, op.shift);
+ auto* output = static_cast<float*>(alias ? ggml_backend_buffer_get_base(execution->workspace) : dst->data);
const auto stream = static_cast<CUDAStream>(context.stream);
- const int threads = x->ne[0] < 1024 ? 256 : 1024;
- launch<<<unsigned(ggml_nrows(x)), threads, 0, stream>>>(
- static_cast<const float*>(x->data), static_cast<const float*>(w->data), static_cast<const float*>(gate->data),
- modulation ? nullptr : static_cast<const float*>(residual->data), output, int(x->ne[0]), eps);
- auto error = cuda_last_error();
+ CUDAError error = cuda_invalid_value;
+ switch (op.kind) {
+ case Kind::Residual:
+ error = launch<float, Kind::Residual>(stream, op, output);
+ break;
+ case Kind::SelfResidual:
+ error = launch<float, Kind::SelfResidual>(stream, op, output);
+ break;
+ case Kind::Modulate:
+ error = launch<float, Kind::Modulate>(stream, op, output);
+ break;
+ }
if (error == cuda_success && alias)
error = cuda_copy_async(dst->data, output, ggml_nbytes(dst), stream);
return error == cuda_success ? GGML_STATUS_SUCCESS : GGML_STATUS_FAILED;
@@ -137,19 +236,12 @@ namespace SD_CUDA_NAMESPACE {
}
ggml_status rms_affine_pack(CUDAStream stream, ggml_tensor* const* nodes, void* output) {
- const auto* x = nodes[0]->src[0];
- const auto* residual = nodes[3]->src[0] == nodes[2] ? nodes[3]->src[1] : nodes[3]->src[0];
- const bool modulation = residual == nodes[1];
- auto launch = x->ne[0] < 1024 ? (modulation ? sd_rms_affine::kernel<256, true, CUDABFloat16> : sd_rms_affine::kernel<256, false, CUDABFloat16>)
- : (modulation ? sd_rms_affine::kernel<1024, true, CUDABFloat16> : sd_rms_affine::kernel<1024, false, CUDABFloat16>);
- float eps;
- std::memcpy(&eps, nodes[0]->op_params, sizeof(eps));
- const int threads = x->ne[0] < 1024 ? 256 : 1024;
- launch<<<unsigned(ggml_nrows(x)), threads, 0, stream>>>(
- static_cast<const float*>(x->data), static_cast<const float*>(nodes[1]->src[1]->data),
- static_cast<const float*>(nodes[2]->src[1]->data), modulation ? nullptr : static_cast<const float*>(residual->data),
- static_cast<CUDABFloat16*>(output), int(x->ne[0]), eps);
- return cuda_last_error() == cuda_success ? GGML_STATUS_SUCCESS : GGML_STATUS_FAILED;
+ const auto op = sd_rms_affine::describe(nodes, 4);
+ auto* dst = static_cast<CUDABFloat16*>(output);
+ const auto error = op.kind == sd_rms_affine::Kind::SelfResidual
+ ? sd_rms_affine::launch<CUDABFloat16, sd_rms_affine::Kind::SelfResidual>(stream, op, dst)
+ : sd_rms_affine::launch<CUDABFloat16, sd_rms_affine::Kind::Residual>(stream, op, dst);
+ return error == cuda_success ? GGML_STATUS_SUCCESS : GGML_STATUS_FAILED;
}
void register_rms_affine(ggml_backend_dev_t device) {
@@ -164,5 +256,11 @@ namespace SD_CUDA_NAMESPACE {
desc.destroy = sd_rms_affine::destroy;
desc.execute = sd_rms_affine::execute;
GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
+ desc.name = "sd.rms_affine_modulate." SD_CUDA_NAME ".v1";
+ // Prefer the existing fusion that also produces ConvRot input when available.
+ desc.priority = 200;
+ desc.match = sd_rms_affine::modulate_match;
+ desc.optimize = nullptr;
+ GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
}
}
diff --git a/src/core/cuda/rope.cu b/src/core/cuda/rope.cu
index 1b724b0..99cc824 100644
--- a/src/core/cuda/rope.cu
+++ b/src/core/cuda/rope.cu
@@ -169,7 +169,7 @@ namespace SD_CUDA_NAMESPACE {
return false;
const auto* x = norm->src[0];
const auto* weight = mul->src[1];
- return x && weight && x->ne[0] >= 32 && x->ne[0] <= 256 && x->ne[0] % 32 == 0 &&
+ return x && weight && weight != norm && x->ne[0] >= 32 && x->ne[0] <= 256 && x->ne[0] % 32 == 0 &&
x->type == GGML_TYPE_F32 && weight->type == GGML_TYPE_F32 && norm->type == GGML_TYPE_F32 &&
ggml_are_same_shape(norm, mul) && ggml_are_same_shape(x, norm) &&
weight->ne[0] == x->ne[0] && ggml_nelements(weight) == x->ne[0] &&
@@ -232,6 +232,13 @@ namespace SD_CUDA_NAMESPACE {
ggml_are_same_shape(rope, scale) && ggml_are_same_shape(rope, copy) && ggml_is_contiguous(copy))
return 5;
}
+ if (first + 3 < ggml_graph_n_nodes(g)) {
+ const auto* rope = ggml_graph_node(g, first + 2);
+ const auto* copy = ggml_graph_node(g, first + 3);
+ if (copy->op == GGML_OP_CPY && copy->src[0] == rope && copy->type == GGML_TYPE_F16 &&
+ ggml_are_same_shape(rope, copy) && ggml_is_contiguous(copy))
+ return 4;
+ }
return 3;
}
@@ -293,7 +300,7 @@ namespace SD_CUDA_NAMESPACE {
void* output = alias ? ggml_backend_buffer_get_base(execution->workspace) : dst->data;
CUDAError error;
if (count > 1) {
- error = rms_rope_launch(stream, nodes, output, count == 5, count == 5 ? param_f32(nodes[3], 0) : 1.f);
+ error = rms_rope_launch(stream, nodes, output, count > 3, count == 5 ? param_f32(nodes[3], 0) : 1.f);
} else {
RopeApplyStrides strides;
const auto* pe = rope->src[1];
diff --git a/src/model/diffusion/moss_image.hpp b/src/model/diffusion/moss_image.hpp
index 1f5dc16..a0fae74 100644
--- a/src/model/diffusion/moss_image.hpp
+++ b/src/model/diffusion/moss_image.hpp
@@ -116,7 +116,7 @@ namespace MOSSImage {
auto norm = [&](const char* name, ggml_tensor* h) { return std::dynamic_pointer_cast<RMSNorm>(blocks[name])->forward(ctx, h); };
auto part = [&](int i) { return mod[i]; };
auto y = norm("norm1", x);
- y = ggml_add(c, ggml_add(c, y, ggml_mul(c, y, part(1))), part(0));
+ y = ggml_ext_modulate(ctx, y, part(0), part(1));
auto qkv = linear("proj.qkv", y);
auto split = [&](int i) {
return ggml_view_4d(c, qkv, config.head_dim, config.dim / config.head_dim, x->ne[1], x->ne[2],
@@ -130,7 +130,7 @@ namespace MOSSImage {
y = ggml_mul(c, y, part(2));
x = ggml_add(c, x, y);
y = norm("norm2", x);
- y = ggml_add(c, ggml_add(c, y, ggml_mul(c, y, part(4))), part(3));
+ y = ggml_ext_modulate(ctx, y, part(3), part(4));
auto up_gate = linear("proj.ffn_up_gate", y);
y = ggml_ext_situ(ctx, up_gate);
y = linear("proj.ffn_down", y);
@@ -191,7 +191,7 @@ namespace MOSSImage {
x = std::dynamic_pointer_cast<RMSNorm>(blocks["final_norm_video"])->forward(ctx, x);
auto shift = ggml_ext_slice(c, head, 0, 0, config.dim);
auto scale = ggml_ext_slice(c, head, 0, config.dim, 2 * config.dim);
- x = ggml_add(c, ggml_add(c, x, ggml_mul(c, x, scale)), shift);
+ x = ggml_ext_modulate(ctx, x, shift, scale);
x = linear("final_linear_video", x);
return DiT::unpatchify(c, x, image->ne[1], image->ne[0], 1, 1);
}