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

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/cpu/modulate.cpp
+82
View File
@@ -0,0 +1,82 @@
#include "core/cpu/registration.h"
#include "core/ggml_extensions.h"
#ifndef SD_USE_UPSTREAM_GGML
#include <cstring>
namespace sd_cpu {
static bool modulate_supports(ggml_backend_dev_t, const ggml_tensor* node, void*) {
return sd_ggml_modulate_describe(node);
}
static float modulate_value(const ggml_tensor* t, const int64_t* index) {
size_t offset = 0;
for (int d = 0; d < GGML_MAX_DIMS; ++d) {
offset += (index[d] % t->ne[d]) * t->nb[d];
}
return *reinterpret_cast<const float*>(static_cast<const char*>(t->data) + offset);
}
static void modulate_compute(ggml_tensor* dst, int ith, int nth, void*) {
auto* output = static_cast<float*>(dst->data);
const int64_t rows = ggml_nrows(dst);
for (int64_t row = ith; row < rows; row += nth) {
int64_t index[] = {0, row % dst->ne[1], (row / dst->ne[1]) % dst->ne[2], row / (dst->ne[1] * dst->ne[2])};
for (int64_t col = 0; col < dst->ne[0]; ++col) {
index[0] = col;
const float x = modulate_value(dst->src[0], index);
// Match separate ggml MUL/ADD rounding even with contraction enabled.
volatile float product = x * modulate_value(dst->src[1], index);
volatile float scaled = x + product;
output[row * dst->ne[0] + col] = dst->src[2] ? scaled + modulate_value(dst->src[2], index) : scaled;
}
}
}
static ggml_status modulate_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 modulate_destroy(void*, void*) {}
static ggml_status modulate_execute(const ggml_ext_execution_context* context, ggml_tensor* const* nodes, int, void*, void*) {
auto reg = ggml_backend_dev_backend_reg(ggml_backend_get_device(context->backend));
auto run = reinterpret_cast<ggml_ext_cpu_parallel_fn>(ggml_backend_reg_get_proc_address(reg, "ggml_ext_cpu_parallel_v1"));
GGML_ASSERT(run);
auto* dst = nodes[0];
const auto address = reinterpret_cast<uintptr_t>(dst->data);
bool alias = false;
for (const auto* input : dst->src) {
if (input) {
const auto src = reinterpret_cast<uintptr_t>(input->data);
alias |= address < src + ggml_nbytes(input) && src < address + ggml_nbytes(dst);
}
}
if (!alias) {
return run(context->backend, dst, modulate_compute, nullptr);
}
ggml_tensor temporary = *dst;
temporary.data = ggml_backend_buffer_get_base(context->workspace);
const auto status = run(context->backend, &temporary, modulate_compute, nullptr);
if (status == GGML_STATUS_SUCCESS) {
std::memcpy(dst->data, temporary.data, ggml_nbytes(dst));
}
return status;
}
void register_modulate(ggml_backend_dev_t device) {
ggml_ext_kernel_desc desc = {};
desc.abi_version = GGML_EXT_ABI_VERSION;
desc.name = "sd.modulate.cpu.v1";
desc.op = GGML_OP_CUSTOM;
desc.custom_op = sd_ggml_modulate_op();
desc.priority = 100;
desc.supports = modulate_supports;
desc.prepare = modulate_prepare;
desc.destroy = modulate_destroy;
desc.execute = modulate_execute;
GGML_ASSERT(ggml_ext_register_kernel(device, &desc));
}
}
#endif
src/core/cpu/registration.cpp
+1
View File
@@ -6,6 +6,7 @@ namespace sd_cpu {
void register_extensions(ggml_backend_dev_t device) {
register_i8(device);
register_rope(device);
register_modulate(device);
}
}
src/core/cpu/registration.h
+1
View File
@@ -7,6 +7,7 @@
namespace sd_cpu {
void register_extensions(ggml_backend_dev_t device);
void register_rope(ggml_backend_dev_t device);
void register_modulate(ggml_backend_dev_t device);
void register_i8(ggml_backend_dev_t device);
}
#endif
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);
src/core/ggml_extend.cpp
+21
View File
@@ -6,6 +6,7 @@
#include <vector>
#include "core/ggml_extend_backend.h"
#include "core/ggml_runner.h"
ggml_tensor* ggml_ext_mul_n_mode(ggml_context* ctx, ggml_tensor* a, ggml_tensor* b, int mode) {
// reshape A
@@ -155,6 +156,26 @@ static bool ggml_ext_is_padded_1d(const ggml_tensor* x) {
x->nb[3] == x->nb[2] * x->ne[2];
}
ggml_tensor* ggml_ext_modulate(GGMLRunnerContext* ctx,
ggml_tensor* x,
ggml_tensor* shift,
ggml_tensor* scale) {
#ifndef SD_USE_UPSTREAM_GGML
if (x->type == GGML_TYPE_F32 && scale->type == GGML_TYPE_F32 && (!shift || shift->type == GGML_TYPE_F32) &&
(sd_backend_is_cpu(ctx->backend) || sd_backend_is(ctx->backend, "CUDA") ||
sd_backend_is(ctx->backend, "ROCm") || sd_backend_is(ctx->backend, "MUSA") || sd_backend_is(ctx->backend, "Vulkan"))) {
ggml_tensor* inputs[] = {x, scale, shift};
return ggml_ext_op_create(ctx->ggml_ctx, sd_ggml_modulate_op(), inputs, shift ? 3 : 2, nullptr);
}
#endif
// Preserve separate multiply/add nodes for rounding and backend fusion matching.
x = ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, x, scale));
if (shift != nullptr) {
x = ggml_add(ctx->ggml_ctx, x, shift);
}
return x;
}
ggml_tensor* ggml_ext_scale(ggml_context* ctx,
ggml_tensor* x,
float factor,
src/core/ggml_extend.h
+9
View File
@@ -8,6 +8,8 @@
#include "ggml-backend.h"
#include "ggml.h"
struct GGMLRunnerContext;
#define EPS 1e-05f
static_assert(GGML_MAX_NAME >= 160, "GGML_MAX_NAME must be at least 160");
@@ -58,6 +60,13 @@ ggml_tensor* ggml_ext_scale(ggml_context* ctx,
float factor,
bool inplace = false);
// scale and optional shift must already be broadcastable to x in ggml layout.
// Computes (x + x * scale) + shift; F32 uses sd.modulate.v1 when supported by ctx->backend.
ggml_tensor* ggml_ext_modulate(GGMLRunnerContext* ctx,
ggml_tensor* x,
ggml_tensor* shift,
ggml_tensor* scale);
ggml_tensor* ggml_ext_gelu(ggml_context* ctx,
ggml_tensor* x,
bool inplace = false);
src/core/ggml_extensions.cpp
+77
View File
@@ -15,6 +15,83 @@
#include "sd_backend_extensions.h"
#endif
static bool modulate_infer(ggml_tensor* const* inputs, int count, const void*, ggml_ext_tensor_desc* out) {
if (count != 2 && count != 3) {
return false;
}
const auto* x = inputs[0];
for (int i = 0; i < count; ++i) {
const auto* t = inputs[i];
if (!t || t->type != GGML_TYPE_F32) {
return false;
}
for (int d = 0; d < GGML_MAX_DIMS; ++d) {
if (t->ne[d] <= 0 || x->ne[d] % t->ne[d] != 0 || t->nb[d] % sizeof(float)) {
return false;
}
}
}
out->type = GGML_TYPE_F32;
std::memcpy(out->ne, x->ne, sizeof(out->ne));
return true;
}
ggml_ext_op_t sd_ggml_modulate_op() {
static const auto op = [] {
const ggml_ext_op_desc desc = {GGML_EXT_ABI_VERSION, "sd.modulate.v1", 0, modulate_infer};
auto result = ggml_ext_register_op(&desc);
GGML_ASSERT(result);
return result;
}();
return op;
}
bool sd_ggml_modulate_describe(const ggml_tensor* node) {
if (!node || ggml_ext_op_get(node) != sd_ggml_modulate_op() || !ggml_is_contiguous(node) || node->src[3]) {
return false;
}
ggml_ext_tensor_desc out;
return modulate_infer(node->src, node->src[2] ? 3 : 2, nullptr, &out) &&
node->type == out.type && std::memcmp(node->ne, out.ne, sizeof(out.ne)) == 0;
}
size_t sd_ggml_modulate_workspace_size(const ggml_tensor* node) {
// Owning custom outputs cannot reuse live input storage. Segment outputs may
// already be allocated, so reserve scratch only if their ranges can overlap.
if (!node->data && !node->view_src) {
return 0;
}
if (node->view_src || !node->buffer) {
return ggml_nbytes(node);
}
for (const auto* input : node->src) {
if (!input) {
continue;
}
const size_t bytes = ggml_nbytes(input);
size_t offset = 0;
while (!input->data && input->view_src) {
offset += input->view_offs;
input = input->view_src;
}
if (!input->data) {
continue;
}
if (!input->buffer) {
return ggml_nbytes(node);
}
if (node->buffer != input->buffer) {
continue;
}
const auto dst = reinterpret_cast<uintptr_t>(node->data);
const auto src = reinterpret_cast<uintptr_t>(input->data) + offset;
if (dst < src + bytes && src < dst + ggml_nbytes(node)) {
return ggml_nbytes(node);
}
}
return 0;
}
static bool rope_infer(ggml_tensor* const* inputs, int count, const void*, ggml_ext_tensor_desc* out) {
if (count != 2) {
return false;
src/core/ggml_extensions.h
+4
View File
@@ -8,6 +8,10 @@ void sd_ggml_register_extensions(ggml_backend_t backend);
#ifndef SD_USE_UPSTREAM_GGML
#include "ggml-ext.h"
ggml_ext_op_t sd_ggml_modulate_op();
bool sd_ggml_modulate_describe(const ggml_tensor* node);
size_t sd_ggml_modulate_workspace_size(const ggml_tensor* node);
ggml_ext_op_t sd_ggml_rope_op();
bool sd_ggml_rope_describe(const ggml_tensor* node);
// Apply precomputed rotary matrices, with shapes in ne[] order:
src/core/vulkan/CMakeLists.txt
+2 -1
View File
@@ -48,6 +48,7 @@ endfunction()
# Preserve the native compilation policy: optimization changes RoPE reductions
# and is disabled for KHR cooperative matrix shaders.
sd_vulkan_shader(rope_apply rope_apply ON)
sd_vulkan_shader(modulate modulate ON -O)
sd_vulkan_shader(rms_rope_apply rms_rope_apply ON)
sd_vulkan_shader(rms_rope_cast rms_rope_apply ON -DCAST_F16)
sd_vulkan_shader(quantize_i8_convrot quantize_i8_convrot ${SD_VK_DOT} -O)
@@ -65,5 +66,5 @@ sd_vulkan_shader(mul_mat_i8_swiglu_split_cm2 mul_mat_i8_tensorwise_cm2 ${SD_VK_C
sd_vulkan_shader(mul_mat_i8_residual_cm2 mul_mat_i8_tensorwise_cm2 ${SD_VK_CM2} -O -DRESIDUAL=1)
add_custom_target(sd-vulkan-shaders DEPENDS ${SD_VULKAN_HEADERS})
add_dependencies(${SD_LIB} sd-vulkan-shaders)
target_sources(${SD_LIB} PRIVATE "${CMAKE_CURRENT_SOURCE_DIR}/rope.cpp" "${CMAKE_CURRENT_SOURCE_DIR}/int8.cpp")
target_sources(${SD_LIB} PRIVATE "${CMAKE_CURRENT_SOURCE_DIR}/rope.cpp" "${CMAKE_CURRENT_SOURCE_DIR}/int8.cpp" "${CMAKE_CURRENT_SOURCE_DIR}/modulate.cpp")
target_include_directories(${SD_LIB} PRIVATE "${SD_VULKAN_GENERATED}")
src/core/vulkan/int8.cpp
+26 -30
View File
@@ -165,22 +165,15 @@ namespace sd_vulkan_i8 {
const ggml_tensor* nodes[4]{};
const ggml_tensor* mod = nullptr;
const bool affine = quant->src[0]->op == GGML_OP_ADD;
const int count = affine ? 4 : 3;
const bool affine = sd_ggml_modulate_describe(quant->src[0]);
const int count = affine ? 2 : 3;
if (affine) {
const ggml_tensor* shift = quant->src[0];
const ggml_tensor* add = shift->src[0];
if (add->op != GGML_OP_ADD || add->src[1]->op != GGML_OP_MUL || add->src[1]->src[0] != add->src[0]) {
return 0;
}
nodes[0] = add->src[1];
nodes[1] = add;
nodes[2] = shift;
nodes[3] = quant;
nodes[0] = quant->src[0];
nodes[1] = quant;
mod = nodes[0]->src[1];
const ggml_tensor* vector = shift->src[1];
if (vector->type != GGML_TYPE_F32 || !ggml_is_contiguous(vector) ||
vector->ne[0] != add->ne[0] || ggml_nelements(vector) != add->ne[0]) {
const ggml_tensor* vector = nodes[0]->src[2];
if (vector && (vector->type != GGML_TYPE_F32 || !ggml_is_contiguous(vector) ||
vector->ne[0] != nodes[0]->ne[0] || ggml_nelements(vector) != nodes[0]->ne[0])) {
return 0;
}
} else {
@@ -215,7 +208,7 @@ namespace sd_vulkan_i8 {
}
static int match_rms_modulate_convrot(const Device& device, const ggml_cgraph* graph, int q, int* indices) {
if (!can_modulate(device, 5376, true) || match_modulate_convrot(device, graph, q, indices + 2) != 4) {
if (!can_modulate(device, 5376, true) || match_modulate_convrot(device, graph, q, indices + 2) != 2) {
return 0;
}
const auto weighted = node(graph, indices[2])->src[0];
@@ -236,27 +229,27 @@ namespace sd_vulkan_i8 {
}
}
return next < 0 ? 6 : 0;
return next < 0 ? 4 : 0;
}
static int can_fuse_modulate_convrot(const Device& device, const ggml_cgraph* graph, int i) {
const ggml_op op = node(graph, i)->op;
if (op == GGML_OP_RMS_NORM) {
int indices[6];
if (match_rms_modulate_convrot(device, graph, i + 5, indices) != 6) {
int indices[4];
if (match_rms_modulate_convrot(device, graph, i + 3, indices) != 4) {
return 0;
}
for (int j = 0; j < 6; ++j) {
for (int j = 0; j < 4; ++j) {
if (indices[j] != i + j) {
return 0;
}
}
return 6;
return 4;
}
if (op != GGML_OP_NORM && op != GGML_OP_MUL) {
if (op != GGML_OP_NORM && !sd_ggml_modulate_describe(node(graph, i))) {
return 0;
}
const int count = op == GGML_OP_NORM ? 3 : 4;
const int count = op == GGML_OP_NORM ? 3 : 2;
int indices[4];
if (match_modulate_convrot(device, graph, i + count - 1, indices) != count) {
return 0;
@@ -491,23 +484,24 @@ namespace sd_vulkan_i8 {
static ggml_status quantize(const ggml_ext_execution_context* e, const Device& d, ggml_tensor* const* nodes, int count) {
const auto* output = nodes[count - 1];
const auto* first = nodes[0];
const bool modulate = count > 1 && (first->op == GGML_OP_NORM || first->op == GGML_OP_MUL || first->op == GGML_OP_RMS_NORM);
const bool modulate = count > 1 && (first->op == GGML_OP_NORM || sd_ggml_modulate_describe(first) || first->op == GGML_OP_RMS_NORM);
const bool rms = modulate && first->op == GGML_OP_RMS_NORM;
const bool affine = modulate && (rms || first->op == GGML_OP_MUL);
const bool affine = modulate && (rms || sd_ggml_modulate_describe(first));
const bool sigmoid = count > 1 && first->op == GGML_OP_UNARY;
const auto* x = count == 1 ? output->src[0] : first->src[0];
const auto* up = count == 1 || modulate ? nullptr : sigmoid ? nodes[1]->src[0]
: first->src[1];
const int extra = rms ? 2 : 0;
const auto* mod = modulate ? nodes[extra + (affine ? 0 : 1)]->src[1] : nullptr;
const auto* shift = modulate ? (affine ? nodes[extra + 2]->src[1] : mod) : nullptr;
const auto* shift = modulate ? (affine ? nodes[extra]->src[2] : mod) : nullptr;
const auto* weight = rms ? nodes[1]->src[1] : x;
std::vector<ggml_ext_vulkan_binding> inputs = {{x}};
if (up)
inputs.push_back({up});
if (modulate) {
inputs.push_back({mod});
inputs.push_back({shift});
if (shift)
inputs.push_back({shift});
if (rms)
inputs.push_back({weight});
}
@@ -521,7 +515,7 @@ namespace sd_vulkan_i8 {
const uint32_t groups[] = {std::min(rows, d.info.max_group_count[0]), 1, 1};
ggml_status status;
if (modulate) {
std::vector<ggml_ext_vulkan_binding> bindings = {{x}, {mod}, {shift}, packed, scales};
std::vector<ggml_ext_vulkan_binding> bindings = {{x}, {mod}, {shift ? shift : mod}, packed, scales};
if (rms)
bindings.push_back({weight});
ModulateParameters pc = {rows, 0, 0, 0, affine && !rms ? 0.f : param_f32(first, 0), k, 0};
@@ -534,9 +528,9 @@ namespace sd_vulkan_i8 {
*offsets[i] = info.shader_offset / 4;
}
auto pipeline = rms ? SD_VK_PIPELINE(rms_modulate_convrot, 6, ModulateParameters) : SD_VK_PIPELINE(modulate_convrot, 5, ModulateParameters);
const uint32_t constants[] = {k, uint32_t(affine)};
const uint32_t constants[] = {k, uint32_t(affine), uint32_t(shift != nullptr)};
pipeline.constants = constants;
pipeline.constant_count = 2;
pipeline.constant_count = 3;
status = d.api->dispatch(e->backend, &pipeline, bindings.data(), &pc, groups);
} else {
ggml_ext_vulkan_binding bindings[] = {{x}, packed, scales, {up}};
@@ -788,6 +782,9 @@ void sd_ggml_register_vulkan_i8(ggml_backend_dev_t device) {
d.custom_op = sd_ggml_i8_op();
GGML_ASSERT(ggml_ext_register_kernel(device, &d));
d.optimize = nullptr;
d.name = "sd.modulate_convrot.vulkan.v1";
d.custom_op = sd_ggml_modulate_op();
GGML_ASSERT(ggml_ext_register_kernel(device, &d));
d.custom_op = nullptr;
const struct {
ggml_op op;
@@ -795,7 +792,6 @@ void sd_ggml_register_vulkan_i8(ggml_backend_dev_t device) {
} fusions[] = {
{GGML_OP_NORM, "sd.norm_convrot.vulkan.v1"},
{GGML_OP_RMS_NORM, "sd.rms_modulate_convrot.vulkan.v1"},
{GGML_OP_MUL, "sd.modulate_convrot.vulkan.v1"},
{GGML_OP_GLU, "sd.swiglu_convrot.vulkan.v1"},
{GGML_OP_UNARY, "sd.sigmoid_convrot.vulkan.v1"}};
for (const auto& fusion : fusions) {
src/core/vulkan/modulate.cpp
+134
View File
@@ -0,0 +1,134 @@
#include "modulate.h"
#include <algorithm>
#include <climits>
#include "core/ggml_extensions.h"
#include "ggml-ext-vulkan.h"
namespace sd_vulkan_modulate {
struct Parameters {
uint32_t ne[4], scale_ne[4], shift_ne[4];
uint32_t x_stride[4], scale_stride[4], shift_stride[4];
uint32_t offsets[4];
uint32_t count, has_shift;
};
static_assert(sizeof(Parameters) == 120, "Modulate push constant layout mismatch");
static const ggml_ext_vulkan_api* api(ggml_backend_dev_t device) {
auto query = reinterpret_cast<ggml_ext_vulkan_api_fn>(ggml_backend_reg_get_proc_address(
ggml_backend_dev_backend_reg(device), "ggml_ext_vulkan_v1"));
return query ? query() : nullptr;
}
static bool supports(ggml_backend_dev_t device, const ggml_tensor* node, void*) {
const auto* service = api(device);
ggml_ext_vulkan_device_info info;
if (!service || !service->device_info(device, &info) || !sd_ggml_modulate_describe(node) ||
info.max_group_size[0] < 256 || info.max_group_invocations < 256 || info.max_push_constants < sizeof(Parameters) || ggml_nelements(node) > UINT32_MAX - 255) {
return false;
}
const ggml_tensor* tensors[] = {node, node->src[0], node->src[1], node->src[2]};
for (const auto* t : tensors) {
if (!t) {
continue;
}
size_t offset = 0;
const auto* owner = t;
while (!owner->data && owner->view_src) {
offset += owner->view_offs;
owner = owner->view_src;
}
if ((reinterpret_cast<uintptr_t>(owner->data) + offset) % 4 ||
ggml_nbytes(t) > info.max_buffer_range - (info.buffer_alignment - 1)) {
return false;
}
for (int d = 0; d < 4; ++d) {
if (t->nb[d] / 4 > UINT32_MAX || t->ne[d] > UINT32_MAX) {
return false;
}
}
}
return (uint64_t(ggml_nelements(node)) + 255) / 256 <= uint64_t(info.max_group_count[0]) * info.max_group_count[1];
}
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*) {
auto device = ggml_backend_get_device(e->backend);
const auto* service = api(device);
ggml_ext_vulkan_device_info info;
if (!service || !service->device_info(device, &info)) {
return GGML_STATUS_FAILED;
}
const auto* dst = nodes[0];
const auto* x = dst->src[0];
const auto* scale = dst->src[1];
const auto* shift = dst->src[2];
ggml_ext_vulkan_binding bindings[] = {{x}, {scale}, {shift ? shift : scale}, {dst}};
ggml_ext_vulkan_buffer_info buffers[4];
for (int i = 0; i < 4; ++i) {
if (!service->buffer_info(e->backend, &bindings[i], &buffers[i])) {
return GGML_STATUS_FAILED;
}
}
bool alias = false;
for (int i = 0; i < 3; ++i) {
alias |= buffers[i].storage == buffers[3].storage && buffers[i].offset < buffers[3].offset + buffers[3].size &&
buffers[3].offset < buffers[i].offset + buffers[i].size;
}
const auto destination = bindings[3];
if (alias) {
bindings[3] = {nullptr, e->workspace, 0, ggml_nbytes(dst)};
if (!service->buffer_info(e->backend, &bindings[3], &buffers[3])) {
return GGML_STATUS_FAILED;
}
}
Parameters p = {};
p.count = uint32_t(ggml_nelements(dst));
p.has_shift = shift != nullptr;
for (int d = 0; d < 4; ++d) {
p.ne[d] = uint32_t(x->ne[d]);
p.scale_ne[d] = uint32_t(scale->ne[d]);
p.shift_ne[d] = shift ? uint32_t(shift->ne[d]) : 1;
p.x_stride[d] = uint32_t(x->nb[d] / 4);
p.scale_stride[d] = uint32_t(scale->nb[d] / 4);
p.shift_stride[d] = shift ? uint32_t(shift->nb[d] / 4) : 0;
p.offsets[d] = buffers[d].shader_offset / 4;
}
const uint32_t count = (p.count + 255) / 256;
const uint32_t width = std::min(count, info.max_group_count[0]);
const uint32_t groups[] = {width, (count + width - 1) / width, 1};
ggml_ext_vulkan_pipeline pipeline = {};
pipeline.name = "sd.modulate.vulkan";
pipeline.spirv = sd_vulkan_shaders::modulate;
pipeline.spirv_size = sizeof(sd_vulkan_shaders::modulate);
pipeline.binding_count = 4;
pipeline.push_constant_size = sizeof(Parameters);
auto status = service->dispatch(e->backend, &pipeline, bindings, &p, groups);
if (status == GGML_STATUS_SUCCESS && alias) {
status = service->copy(e->backend, &bindings[3], &destination);
}
return status;
}
}
void sd_ggml_register_vulkan_modulate(ggml_backend_dev_t device) {
if (!sd_vulkan_modulate::api(device)) {
return;
}
ggml_ext_kernel_desc desc = {};
desc.abi_version = GGML_EXT_ABI_VERSION;
desc.name = "sd.modulate.vulkan.v1";
desc.op = GGML_OP_CUSTOM;
desc.custom_op = sd_ggml_modulate_op();
desc.priority = 100;
desc.supports = sd_vulkan_modulate::supports;
desc.prepare = sd_vulkan_modulate::prepare;
desc.destroy = sd_vulkan_modulate::destroy;
desc.execute = sd_vulkan_modulate::execute;
GGML_ASSERT(ggml_ext_register_kernel(device, &desc));
}
src/core/vulkan/shaders/modulate.comp
+39
View File
@@ -0,0 +1,39 @@
#version 450
layout(local_size_x = 256) in;
layout(binding = 0) readonly buffer X { float x[]; };
layout(binding = 1) readonly buffer S { float scale[]; };
layout(binding = 2) readonly buffer H { float shift[]; };
layout(binding = 3) writeonly buffer D { float dst[]; };
layout(push_constant) uniform Parameters {
uint ne[4], scale_ne[4], shift_ne[4];
uint x_stride[4], scale_stride[4], shift_stride[4];
uint offsets[4];
uint count, has_shift;
} p;
void main() {
uint group = gl_WorkGroupID.y * gl_NumWorkGroups.x + gl_WorkGroupID.x;
if (group >= (p.count + 255) / 256) {
return;
}
uint i = group * 256 + gl_LocalInvocationID.x;
if (i >= p.count) {
return;
}
uint index = i;
uint xi = p.offsets[0], si = p.offsets[1], hi = p.offsets[2];
for (uint d = 0; d < 4; ++d) {
uint 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];
hi += (coord % p.shift_ne[d]) * p.shift_stride[d];
}
precise float product = x[xi] * scale[si];
precise float value = x[xi] + product;
if (p.has_shift != 0) {
value += shift[hi];
}
dst[p.offsets[3] + i] = value;
}
src/core/vulkan/shaders/modulate_convrot.comp
+4 -1
View File
@@ -4,6 +4,7 @@
layout(constant_id = 0) const uint K = 4096;
layout(constant_id = 1) const bool AFFINE = false;
layout(constant_id = 2) const bool HAS_SHIFT = true;
layout(local_size_x = 512, local_size_y = 1, local_size_z = 1) in;
layout(binding = 0) readonly buffer X {
float data_x[];
@@ -111,7 +112,9 @@ void main() {
if (AFFINE) {
precise float modulated = values[col] * data_m[p.m_offset + col];
value = values[col] + modulated;
value += data_h[p.h_offset + col];
if (HAS_SHIFT) {
value += data_h[p.h_offset + col];
}
} else {
value = (values[col] - mean) * inv_std;
value *= data_m[p.m_offset + col];