DarthAffe/stable-diffusion.cpp
1
0 Fork 0
mirror of https://github.com/leejet/stable-diffusion.cpp.git synced 2026-11-11 03:24:30 +00:00
Code Issues Packages Projects Releases Wiki Activity

perf(cuda): fuse MOSS modulation and direct RoPE F16 output (#2269)

Browse Source 
Co-authored-by: mirek190 <mirek190@gmail.com>
This commit is contained in:
leejetandmirek190 authored and GitHub committed 2026-11-04 02:01:41 +08:00
1 parent d30854463d
commit 984eb0686c
3 files changed
+164 -59
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

src/core/cuda/rms_affine.cu
+152 -54
View File
@@ -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));
}
}
src/core/cuda/rope.cu
+9 -2
View File
@@ -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];