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

feat: add custom modulate operator and ConvRot fusion (#2258)

Browse Source
This commit is contained in:
leejet authored and GitHub committed 2026-11-02 01:19:28 +08:00
1 parent 133087deb8
commit 290148d2e3
36 files changed
+655 -161
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

src/core/cuda/CMakeLists.txt
+1 -1
View File
@@ -2,7 +2,7 @@ cmake_policy(SET CMP0079 NEW)
function(sd_cuda_sources VENDOR LANGUAGE)
set(SOURCES)
foreach(OP rope int8 convrot affine registration)
foreach(OP rope int8 convrot affine modulate registration)
set(SOURCE "${CMAKE_CURRENT_BINARY_DIR}/${OP}_${VENDOR}.cu")
string(TOUPPER "${VENDOR}" UPPER_VENDOR)
file(GENERATE OUTPUT "${SOURCE}" CONTENT "#define SD_CUDA_VENDOR_${UPPER_VENDOR} 1\n#include \"core/cuda/${OP}.cu\"\n")
src/core/cuda/convrot.cu
+32 -22
View File
@@ -27,6 +27,7 @@ namespace SD_CUDA_NAMESPACE {
}
enum Kind { SCALE,
NORM,
MOD,
RMS_MOD,
SWIGLU,
SILU,
@@ -65,16 +66,20 @@ namespace SD_CUDA_NAMESPACE {
op.eps = param_f32(first, 0);
if (op.k < 1024 || !vector(op.modulation, op.k))
return false;
} else if (count == 6 && first->op == GGML_OP_RMS_NORM && nodes[1]->op == GGML_OP_MUL && nodes[1]->src[0] == first &&
nodes[2]->op == GGML_OP_MUL && nodes[2]->src[0] == nodes[1] &&
nodes[3]->op == GGML_OP_ADD && nodes[3]->src[0] == nodes[1] && nodes[3]->src[1] == nodes[2] &&
nodes[4]->op == GGML_OP_ADD && nodes[4]->src[0] == nodes[3] && value == nodes[4]) {
} else if (count == 2 && sd_ggml_modulate_describe(first) && value == first) {
op.kind = MOD;
op.modulation = first->src[1];
op.shift = first->src[2];
if (!vector(op.modulation, op.k) || (op.shift && !vector(op.shift, op.k)))
return false;
} else if (count == 4 && first->op == GGML_OP_RMS_NORM && nodes[1]->op == GGML_OP_MUL && nodes[1]->src[0] == first &&
sd_ggml_modulate_describe(nodes[2]) && nodes[2]->src[0] == nodes[1] && value == nodes[2]) {
op.kind = RMS_MOD;
op.weight = nodes[1]->src[1];
op.modulation = nodes[2]->src[1];
op.shift = nodes[4]->src[1];
op.shift = nodes[2]->src[2];
op.eps = param_f32(first, 0);
if (op.k < 1024 || !vector(op.weight, op.k) || !vector(op.modulation, op.k) || !vector(op.shift, op.k))
if (op.k < 1024 || !vector(op.weight, op.k) || !vector(op.modulation, op.k) || (op.shift && !vector(op.shift, op.k)))
return false;
} else if (count == 2 && first->op == GGML_OP_GLU && ggml_get_glu_op(first) == GGML_GLU_OP_SWIGLU && value == first) {
op.kind = SWIGLU;
@@ -119,7 +124,7 @@ namespace SD_CUDA_NAMESPACE {
auto* g = const_cast<ggml_cgraph*>(graph);
const int available = ggml_graph_n_nodes(g) - first;
const auto* node = ggml_graph_node(g, first);
const int count = node->op == GGML_OP_RMS_NORM ? 6 : node->op == GGML_OP_NORM || node->op == GGML_OP_UNARY ? 3
const int count = node->op == GGML_OP_RMS_NORM ? 4 : node->op == GGML_OP_NORM || node->op == GGML_OP_UNARY ? 3
: 2;
if (available < count)
return 0;
@@ -132,7 +137,7 @@ namespace SD_CUDA_NAMESPACE {
static void optimize(ggml_backend_dev_t device, ggml_cgraph* graph, void*) {
if (fusion_disabled())
return;
for (int i = 2; i < ggml_graph_n_nodes(graph); ++i) {
for (int i = 1; i < ggml_graph_n_nodes(graph); ++i) {
auto* quant = ggml_graph_node(graph, i);
int group;
if (!sd_ggml_convrot_describe(quant, &group) || group != 256)
@@ -140,17 +145,18 @@ namespace SD_CUDA_NAMESPACE {
ggml_tensor* chain[6] = {};
int count = 0;
auto* value = quant->src[0];
if (value->op == GGML_OP_ADD && value->src[0]->op == GGML_OP_ADD) {
auto* add = value->src[0];
auto* weight = add->src[0];
if (weight->op == GGML_OP_MUL) {
if (sd_ggml_modulate_describe(value)) {
auto* weight = value->src[0];
if (weight->op == GGML_OP_MUL && weight->src[0]->op == GGML_OP_RMS_NORM) {
chain[0] = weight->src[0];
chain[1] = weight;
chain[2] = add->src[1];
chain[3] = add;
chain[4] = value;
chain[5] = quant;
count = 6;
chain[2] = value;
chain[3] = quant;
count = 4;
} else {
chain[0] = value;
chain[1] = quant;
count = 2;
}
} else if (value->op == GGML_OP_MUL) {
for (int src = 0; src < 2; ++src)
@@ -210,13 +216,13 @@ namespace SD_CUDA_NAMESPACE {
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);
}
template <bool Swiglu = false, bool Normalize = false, bool Scaled = false, bool Sigmoid = false, bool RMS = false>
template <bool Swiglu = false, bool Normalize = false, bool Scaled = false, bool Sigmoid = false, bool RMS = false, bool Modulate = false>
static CUDAError launch(Plan& p, const Operation& op, int8_t* output, CUDAStream stream) {
constexpr bool warp = Scaled || RMS;
constexpr bool warp = Scaled || RMS || Modulate;
constexpr int extra = Normalize ? 384 : RMS ? 256
: 128;
if (!p.configured) {
auto status = cuda_set_shared_memory(quantize<Swiglu, false, Normalize, Scaled, warp, float, Sigmoid, RMS>,
auto status = cuda_set_shared_memory(quantize<Swiglu, false, Normalize, Scaled, warp, float, Sigmoid, RMS, Modulate>,
int(p.device.shared_memory - extra));
if (status != cuda_success)
return status;
@@ -229,7 +235,7 @@ namespace SD_CUDA_NAMESPACE {
: 256;
const int64_t padded = (op.rows + 3) / 4 * 4;
auto data = [](const ggml_tensor* t) { return t ? static_cast<const float*>(t->data) : nullptr; };
quantize<Swiglu, false, Normalize, Scaled, warp, float, Sigmoid, RMS><<<unsigned(op.rows), threads, size_t(op.k) * sizeof(float), stream>>>(
quantize<Swiglu, false, Normalize, Scaled, warp, float, Sigmoid, RMS, Modulate><<<unsigned(op.rows), threads, size_t(op.k) * sizeof(float), stream>>>(
data(op.input) + op.gate_offset, output, reinterpret_cast<float*>(output + op.k * padded), op.k, op.rows, op.k,
op.up ? data(op.up) + op.up_offset : nullptr, op.stride, op.up_stride, nullptr, nullptr, data(op.modulation), op.eps, op.scale, op.bias, data(op.weight), data(op.shift));
return cuda_last_error();
@@ -260,6 +266,9 @@ namespace SD_CUDA_NAMESPACE {
case RMS_MOD:
error = launch<false, false, false, false, true>(plan, op, output, stream);
break;
case MOD:
error = launch<false, false, false, false, false, true>(plan, op, output, stream);
break;
case SWIGLU:
case SILU:
error = launch<true>(plan, op, output, stream);
@@ -277,12 +286,13 @@ namespace SD_CUDA_NAMESPACE {
}
void register_convrot(ggml_backend_dev_t device) {
for (auto op : {GGML_OP_SCALE, GGML_OP_NORM, GGML_OP_RMS_NORM, GGML_OP_GLU, GGML_OP_UNARY}) {
for (auto op : {GGML_OP_SCALE, GGML_OP_NORM, GGML_OP_RMS_NORM, GGML_OP_GLU, GGML_OP_UNARY, GGML_OP_CUSTOM}) {
ggml_ext_kernel_desc desc = {};
desc.abi_version = GGML_EXT_ABI_VERSION;
const std::string name = std::string("sd.convrot_fusion.") + ggml_op_name(op) + "." SD_CUDA_NAME ".v1";
desc.name = name.c_str();
desc.op = op;
desc.custom_op = op == GGML_OP_CUSTOM ? sd_ggml_modulate_op() : nullptr;
desc.priority = 225;
desc.optimize = op == GGML_OP_SCALE ? sd_convrot::optimize : nullptr;
desc.match = sd_convrot::match;
src/core/cuda/convrot_kernel.h
+7 -2
View File
@@ -40,7 +40,7 @@ namespace SD_CUDA_NAMESPACE {
return __fadd_rn(__fadd_rn(__fadd_rn(a, b), c), d);
}
template <bool swiglu, bool dequantize = false, bool normalize = false, bool scaled = false, bool warp_rotation = false, typename T = float, bool sigmoid_gate = false, bool rms_modulation = false>
template <bool swiglu, bool dequantize = false, bool normalize = false, bool scaled = false, bool warp_rotation = false, typename T = float, bool sigmoid_gate = false, bool rms_modulation = false, bool modulate = false>
__global__ void quantize(
const T* src,
int8_t* dst,
@@ -103,7 +103,12 @@ namespace SD_CUDA_NAMESPACE {
float value = src_row[i];
if constexpr (rms_modulation) {
value = __fmul_rn(__fmul_rn(value, rms_scale), rms_weight[i]);
value = __fadd_rn(__fadd_rn(value, __fmul_rn(value, modulation[i])), rms_shift[i]);
}
if constexpr (rms_modulation || modulate) {
value = __fadd_rn(value, __fmul_rn(value, modulation[i]));
if (rms_shift) {
value = __fadd_rn(value, rms_shift[i]);
}
}
if constexpr (scaled) {
value = __fmaf_rn(value, input_scale, input_bias);
src/core/cuda/modulate.cu
+95
View File
@@ -0,0 +1,95 @@
#include <algorithm>
#include "core/ggml_extensions.h"
#include "runtime.h"
namespace SD_CUDA_NAMESPACE {
namespace sd_modulate {
struct Parameters {
int64_t ne[4], scale_ne[4], shift_ne[4];
size_t x_stride[4], scale_stride[4], shift_stride[4];
int64_t count;
};
static __global__ void modulate(const float* x, const float* scale, const float* shift, float* dst, Parameters p) {
for (int64_t i = int64_t(blockIdx.x) * blockDim.x + threadIdx.x; i < p.count; i += int64_t(blockDim.x) * gridDim.x) {
int64_t index = i;
size_t xi = 0, si = 0, hi = 0;
#pragma unroll
for (int d = 0; d < 4; ++d) {
const int64_t coord = index % p.ne[d];
index /= p.ne[d];
xi += coord * p.x_stride[d];
si += (coord % p.scale_ne[d]) * p.scale_stride[d];
if (shift) {
hi += (coord % p.shift_ne[d]) * p.shift_stride[d];
}
}
float value = __fadd_rn(x[xi], __fmul_rn(x[xi], scale[si]));
dst[i] = shift ? __fadd_rn(value, shift[hi]) : value;
}
}
static bool supports(ggml_backend_dev_t, const ggml_tensor* node, void*) {
return sd_ggml_modulate_describe(node);
}
static ggml_status prepare(ggml_backend_t, ggml_tensor* const* nodes, int, void**, size_t* bytes, void*) {
*bytes = sd_ggml_modulate_workspace_size(nodes[0]);
return GGML_STATUS_SUCCESS;
}
static void destroy(void*, void*) {}
static ggml_status execute(const ggml_ext_execution_context* e, ggml_tensor* const* nodes, int, void*, void*) {
ggml_ext_gpu_context context;
if (!cuda_context(e->backend, &context)) {
return GGML_STATUS_FAILED;
}
auto* dst = nodes[0];
const auto* x = dst->src[0];
const auto* scale = dst->src[1];
const auto* shift = dst->src[2];
Parameters p = {};
p.count = ggml_nelements(dst);
bool alias = false;
const auto address = reinterpret_cast<uintptr_t>(dst->data);
for (const auto* input : {x, scale, shift}) {
if (input) {
const auto src = reinterpret_cast<uintptr_t>(input->data);
alias |= address < src + ggml_nbytes(input) && src < address + ggml_nbytes(dst);
}
}
for (int d = 0; d < 4; ++d) {
p.ne[d] = x->ne[d];
p.scale_ne[d] = scale->ne[d];
p.shift_ne[d] = shift ? shift->ne[d] : 1;
p.x_stride[d] = x->nb[d] / sizeof(float);
p.scale_stride[d] = scale->nb[d] / sizeof(float);
p.shift_stride[d] = shift ? shift->nb[d] / sizeof(float) : 0;
}
auto stream = static_cast<CUDAStream>(context.stream);
auto* output = static_cast<float*>(alias ? ggml_backend_buffer_get_base(e->workspace) : dst->data);
const unsigned blocks = unsigned(std::min<int64_t>((p.count + 255) / 256, 65535));
modulate<<<blocks, 256, 0, stream>>>(static_cast<const float*>(x->data), static_cast<const float*>(scale->data),
shift ? static_cast<const float*>(shift->data) : nullptr, output, p);
if (cuda_last_error() != cuda_success || (alias && cuda_copy_async(dst->data, output, ggml_nbytes(dst), stream) != cuda_success)) {
return GGML_STATUS_FAILED;
}
return GGML_STATUS_SUCCESS;
}
}
void register_modulate(ggml_backend_dev_t device) {
ggml_ext_kernel_desc desc = {};
desc.abi_version = GGML_EXT_ABI_VERSION;
desc.name = "sd.modulate." SD_CUDA_NAME ".v1";
desc.op = GGML_OP_CUSTOM;
desc.custom_op = sd_ggml_modulate_op();
desc.priority = 100;
desc.supports = sd_modulate::supports;
desc.prepare = sd_modulate::prepare;
desc.destroy = sd_modulate::destroy;
desc.execute = sd_modulate::execute;
GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
}
}
src/core/cuda/registration.cu
+2
View File
@@ -11,6 +11,7 @@ namespace SD_CUDA_NAMESPACE {
void register_i8(ggml_backend_dev_t device);
void register_convrot(ggml_backend_dev_t device);
void register_affine(ggml_backend_dev_t device);
void register_modulate(ggml_backend_dev_t device);
#if !defined(SD_CUDA_VENDOR_HIP) && !defined(SD_CUDA_VENDOR_MUSA)
void register_sage_attention(ggml_backend_dev_t device);
void register_sol_attention(ggml_backend_dev_t device);
@@ -20,6 +21,7 @@ namespace SD_CUDA_NAMESPACE {
if (cuda_device_info(device).device < 0)
return;
register_affine(device);
register_modulate(device);
register_i8(device);
register_convrot(device);
register_rope(device);