Commit 290148d for stable-diffusion.cpp
commit 290148d2e31ca6ef9757d567905776b5f3646db5
Author: leejet <leejet714@gmail.com>
Date: Mon Nov 2 01:19:28 2026 +0800
feat: add custom modulate operator and ConvRot fusion (#2258)
diff --git a/CMakeLists.txt b/CMakeLists.txt
index 235556f..15bd917 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -378,7 +378,7 @@ if (NOT SD_USE_UPSTREAM_GGML AND (SD_CUDA OR SD_HIPBLAS OR SD_MUSA))
endif()
if (SD_VULKAN AND NOT SD_USE_UPSTREAM_GGML)
add_subdirectory(src/core/vulkan)
- list(APPEND SD_BACKEND_REGISTRATIONS sd_ggml_register_vulkan_rope sd_ggml_register_vulkan_i8)
+ list(APPEND SD_BACKEND_REGISTRATIONS sd_ggml_register_vulkan_rope sd_ggml_register_vulkan_i8 sd_ggml_register_vulkan_modulate)
endif()
if(SD_BACKEND_REGISTRATIONS)
string(REPLACE ";" ",\n " SD_BACKEND_REGISTRATIONS "${SD_BACKEND_REGISTRATIONS}")
diff --git a/cmake/backend_extensions.h.in b/cmake/backend_extensions.h.in
index cf41592..6b42ea1 100644
--- a/cmake/backend_extensions.h.in
+++ b/cmake/backend_extensions.h.in
@@ -5,6 +5,7 @@
void sd_ggml_register_vulkan_rope(ggml_backend_dev_t device);
void sd_ggml_register_vulkan_i8(ggml_backend_dev_t device);
+void sd_ggml_register_vulkan_modulate(ggml_backend_dev_t device);
static void (*const sd_backend_extensions[])(ggml_backend_dev_t) = {
@SD_BACKEND_REGISTRATIONS@
diff --git a/docs/ggml_extensions.md b/docs/ggml_extensions.md
index 6787c44..5c05989 100644
--- a/docs/ggml_extensions.md
+++ b/docs/ggml_extensions.md
@@ -135,6 +135,25 @@ CPU kernels and fusion rules live in `src/core/cpu/`, with registration collecte
in `cpu/registration.cpp`. CUDA-style and Vulkan implementations live in their
respective backend directories.
+`ggml_ext_modulate(ctx, x, shift, scale)` emits `sd.modulate.v1` for F32 inputs
+when the runner context uses CPU, CUDA/HIP/MUSA, or Vulkan.
+Its explicit inputs are `x`, `scale`, and optional `shift`; it computes
+`(x + x * scale) + shift` with separate FP32 multiply/add rounding. The output
+is contiguous and out of place. Inputs may be strided and use ggml's
+dimension-wise repetition rules. Callers prepare token, batch, or frame layouts
+before constructing the operation; normalization and residual gates remain
+separate operations. CPU, CUDA/HIP/MUSA, and Vulkan have standalone kernels.
+Other backends (including Metal, SYCL, and OpenCL), upstream ggml builds, and
+other input types retain the original MUL/ADD graph.
+
+ConvRot fusion matches the Modulate node directly, including scale-only
+modulation, instead of recognizing its former MUL/ADD/ADD decomposition.
+CUDA accepts contiguous per-channel parameters for `Modulate -> ConvRot` and
+`RMSNorm -> Mul(weight) -> Modulate -> ConvRot`. Vulkan limits the first fusion
+to widths 4096/6144/5376 and the RMSNorm fusion to width 5376.
+Other layouts execute the standalone operation followed by ConvRot. The
+gated-residual MUL/ADD fusion remains available for residual branches.
+
`src/core/cpu/rope.cpp` registers the CPU implementation of
`sd.rope_apply.v1` and a
`RMSNorm -> Mul -> RoPE` fusion. The fusion accepts a one-dimensional weight
@@ -269,8 +288,9 @@ unsupported layouts retain ordinary operator execution. Fusion and force-cuBLAS
disable settings are honored.
`src/core/cuda/convrot.cu` registers CUDA fusion for `Scale -> ConvRot`,
-`Norm -> Mul -> ConvRot`, `RMSNorm -> Mul(weight) -> Mul(modulation) -> Add ->
-Add(shift) -> ConvRot`, `SwiGLU -> ConvRot`, and `SiLU/Sigmoid -> Mul -> ConvRot`.
+`Norm -> Mul -> ConvRot`, `Modulate -> ConvRot`,
+`RMSNorm -> Mul(weight) -> Modulate -> ConvRot`, `SwiGLU -> ConvRot`, and
+`SiLU/Sigmoid -> Mul -> ConvRot`.
ConvRot outputs are custom nodes. SwiGLU accepts packed,
swapped, or separate gate/up inputs. Independent modulation and up-branch
producers are moved ahead of their consumers through the safe graph-delay API.
diff --git a/src/core/cpu/modulate.cpp b/src/core/cpu/modulate.cpp
new file mode 100644
index 0000000..84a90ab
--- /dev/null
+++ b/src/core/cpu/modulate.cpp
@@ -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
diff --git a/src/core/cpu/registration.cpp b/src/core/cpu/registration.cpp
index a4f5bd3..d63fd1f 100644
--- a/src/core/cpu/registration.cpp
+++ b/src/core/cpu/registration.cpp
@@ -6,6 +6,7 @@ namespace sd_cpu {
void register_extensions(ggml_backend_dev_t device) {
register_i8(device);
register_rope(device);
+ register_modulate(device);
}
}
diff --git a/src/core/cpu/registration.h b/src/core/cpu/registration.h
index c711f6f..eba629c 100644
--- a/src/core/cpu/registration.h
+++ b/src/core/cpu/registration.h
@@ -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
diff --git a/src/core/cuda/CMakeLists.txt b/src/core/cuda/CMakeLists.txt
index 7bca6aa..f366899 100644
--- a/src/core/cuda/CMakeLists.txt
+++ b/src/core/cuda/CMakeLists.txt
@@ -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")
diff --git a/src/core/cuda/convrot.cu b/src/core/cuda/convrot.cu
index d5d159b..cf78961 100644
--- a/src/core/cuda/convrot.cu
+++ b/src/core/cuda/convrot.cu
@@ -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;
diff --git a/src/core/cuda/convrot_kernel.h b/src/core/cuda/convrot_kernel.h
index 748c08e..04abf97 100644
--- a/src/core/cuda/convrot_kernel.h
+++ b/src/core/cuda/convrot_kernel.h
@@ -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);
diff --git a/src/core/cuda/modulate.cu b/src/core/cuda/modulate.cu
new file mode 100644
index 0000000..c4af6f6
--- /dev/null
+++ b/src/core/cuda/modulate.cu
@@ -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));
+ }
+}
diff --git a/src/core/cuda/registration.cu b/src/core/cuda/registration.cu
index d03b114..291edd6 100644
--- a/src/core/cuda/registration.cu
+++ b/src/core/cuda/registration.cu
@@ -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);
diff --git a/src/core/ggml_extend.cpp b/src/core/ggml_extend.cpp
index 6d72cae..ca1b231 100644
--- a/src/core/ggml_extend.cpp
+++ b/src/core/ggml_extend.cpp
@@ -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,
diff --git a/src/core/ggml_extend.h b/src/core/ggml_extend.h
index 4a97704..9ddfb0a 100644
--- a/src/core/ggml_extend.h
+++ b/src/core/ggml_extend.h
@@ -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);
diff --git a/src/core/ggml_extensions.cpp b/src/core/ggml_extensions.cpp
index 1f0aa1d..27089f6 100644
--- a/src/core/ggml_extensions.cpp
+++ b/src/core/ggml_extensions.cpp
@@ -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;
diff --git a/src/core/ggml_extensions.h b/src/core/ggml_extensions.h
index 49b0e76..86c9121 100644
--- a/src/core/ggml_extensions.h
+++ b/src/core/ggml_extensions.h
@@ -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:
diff --git a/src/core/vulkan/CMakeLists.txt b/src/core/vulkan/CMakeLists.txt
index 8e55d99..42ee931 100644
--- a/src/core/vulkan/CMakeLists.txt
+++ b/src/core/vulkan/CMakeLists.txt
@@ -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}")
diff --git a/src/core/vulkan/int8.cpp b/src/core/vulkan/int8.cpp
index 4caadfb..7e3e67d 100644
--- a/src/core/vulkan/int8.cpp
+++ b/src/core/vulkan/int8.cpp
@@ -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) {
diff --git a/src/core/vulkan/modulate.cpp b/src/core/vulkan/modulate.cpp
new file mode 100644
index 0000000..4561447
--- /dev/null
+++ b/src/core/vulkan/modulate.cpp
@@ -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));
+}
diff --git a/src/core/vulkan/shaders/modulate.comp b/src/core/vulkan/shaders/modulate.comp
new file mode 100644
index 0000000..df07626
--- /dev/null
+++ b/src/core/vulkan/shaders/modulate.comp
@@ -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;
+}
diff --git a/src/core/vulkan/shaders/modulate_convrot.comp b/src/core/vulkan/shaders/modulate_convrot.comp
index 3e4b3a2..3cac110 100644
--- a/src/core/vulkan/shaders/modulate_convrot.comp
+++ b/src/core/vulkan/shaders/modulate_convrot.comp
@@ -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];
diff --git a/src/model/diffusion/anima.hpp b/src/model/diffusion/anima.hpp
index 66e1a0b..e9e33b8 100644
--- a/src/model/diffusion/anima.hpp
+++ b/src/model/diffusion/anima.hpp
@@ -129,7 +129,7 @@ namespace Anima {
auto gate = emb_chunks[2];
auto x = norm->forward(ctx, hidden_states);
- x = Flux::modulate(ctx->ggml_ctx, x, shift, scale);
+ x = Flux::modulate(ctx, x, shift, scale);
return {x, gate};
}
@@ -169,7 +169,7 @@ namespace Anima {
auto scale = emb_chunks[1];
auto x = norm->forward(ctx, hidden_states);
- x = Flux::modulate(ctx->ggml_ctx, x, shift, scale);
+ x = Flux::modulate(ctx, x, shift, scale);
return x;
}
};
diff --git a/src/model/diffusion/boogu.hpp b/src/model/diffusion/boogu.hpp
index ac67ff9..3f652e1 100644
--- a/src/model/diffusion/boogu.hpp
+++ b/src/model/diffusion/boogu.hpp
@@ -128,9 +128,9 @@ namespace Boogu {
}
};
- __STATIC_INLINE__ ggml_tensor* scale_modulate(ggml_context* ctx, ggml_tensor* x, ggml_tensor* scale) {
- scale = ggml_reshape_3d(ctx, scale, scale->ne[0], 1, scale->ne[1]);
- return ggml_add(ctx, x, ggml_mul(ctx, x, scale));
+ __STATIC_INLINE__ ggml_tensor* scale_modulate(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* scale) {
+ scale = ggml_reshape_3d(ctx->ggml_ctx, scale, scale->ne[0], 1, scale->ne[1]);
+ return ggml_ext_modulate(ctx, x, nullptr, scale);
}
__STATIC_INLINE__ ggml_tensor* gate_residual(ggml_context* ctx, ggml_tensor* residual, ggml_tensor* x, ggml_tensor* gate) {
@@ -186,7 +186,7 @@ namespace Boogu {
auto scale_mlp = mods[2];
auto gate_mlp = mods[3];
- x = scale_modulate(ctx->ggml_ctx, norm->forward(ctx, x), scale_msa);
+ x = scale_modulate(ctx, norm->forward(ctx, x), scale_msa);
return {x, gate_msa, scale_mlp, gate_mlp};
}
};
@@ -231,7 +231,7 @@ namespace Boogu {
auto linear_2 = std::dynamic_pointer_cast<Linear>(blocks["linear_2"]);
auto emb = linear_1->forward(ctx, ggml_silu(ctx->ggml_ctx, conditioning_embedding));
- x = scale_modulate(ctx->ggml_ctx, norm->forward(ctx, x), emb);
+ x = scale_modulate(ctx, norm->forward(ctx, x), emb);
x = linear_2->forward(ctx, x);
return x;
}
@@ -334,7 +334,7 @@ namespace Boogu {
auto attn_output = attn->forward(ctx, norm_hidden_states, norm_hidden_states, rotary_emb, attention_mask);
hidden_states = gate_residual(ctx->ggml_ctx, hidden_states, norm2->forward(ctx, attn_output), gate_msa);
- auto mlp_input = scale_modulate(ctx->ggml_ctx, ffn_norm1->forward(ctx, hidden_states), scale_mlp);
+ auto mlp_input = scale_modulate(ctx, ffn_norm1->forward(ctx, hidden_states), scale_mlp);
auto mlp_output = feed_forward->forward(ctx, mlp_input);
hidden_states = gate_residual(ctx->ggml_ctx, hidden_states, ffn_norm2->forward(ctx, mlp_output), gate_mlp);
} else {
@@ -512,14 +512,14 @@ namespace Boogu {
img_hidden_states = gate_residual(ctx->ggml_ctx, img_hidden_states, img_attn_norm->forward(ctx, img_attn_out), img_gate_msa);
img_hidden_states = gate_residual(ctx->ggml_ctx, img_hidden_states, img_self_attn_norm->forward(ctx, img_self_attn_out), img_gate_self);
- auto img_mlp_input = scale_modulate(ctx->ggml_ctx, img_norm2_out, img_scale_mlp);
+ auto img_mlp_input = scale_modulate(ctx, img_norm2_out, img_scale_mlp);
img_shift_mlp = ggml_reshape_3d(ctx->ggml_ctx, img_shift_mlp, img_shift_mlp->ne[0], 1, img_shift_mlp->ne[1]);
img_mlp_input = ggml_add(ctx->ggml_ctx, img_mlp_input, img_shift_mlp);
auto img_mlp_out = img_feed_forward->forward(ctx, img_ffn_norm1->forward(ctx, img_mlp_input));
img_hidden_states = gate_residual(ctx->ggml_ctx, img_hidden_states, img_ffn_norm2->forward(ctx, img_mlp_out), img_gate_mlp);
instruct_hidden_states = gate_residual(ctx->ggml_ctx, instruct_hidden_states, instruct_attn_norm->forward(ctx, instruct_attn_out), instruct_gate_msa);
- auto instruct_mlp_input = scale_modulate(ctx->ggml_ctx, instruct_norm2_out, instruct_scale_mlp);
+ auto instruct_mlp_input = scale_modulate(ctx, instruct_norm2_out, instruct_scale_mlp);
instruct_shift_mlp = ggml_reshape_3d(ctx->ggml_ctx, instruct_shift_mlp, instruct_shift_mlp->ne[0], 1, instruct_shift_mlp->ne[1]);
instruct_mlp_input = ggml_add(ctx->ggml_ctx, instruct_mlp_input, instruct_shift_mlp);
auto instruct_mlp_out = instruct_feed_forward->forward(ctx, instruct_ffn_norm1->forward(ctx, instruct_mlp_input));
diff --git a/src/model/diffusion/ernie_image.hpp b/src/model/diffusion/ernie_image.hpp
index 59b06cc..15badb9 100644
--- a/src/model/diffusion/ernie_image.hpp
+++ b/src/model/diffusion/ernie_image.hpp
@@ -246,13 +246,13 @@ namespace ErnieImage {
auto residual = x;
x = adaLN_sa_ln->forward(ctx, x);
- x = Flux::modulate(ctx->ggml_ctx, x, shift_msa, scale_msa, true);
+ x = Flux::modulate(ctx, x, shift_msa, scale_msa, true);
auto attn_out = self_attention->forward(ctx, x, pe, attention_mask);
x = ggml_add(ctx->ggml_ctx, residual, ggml_mul(ctx->ggml_ctx, attn_out, gate_msa));
residual = x;
x = adaLN_mlp_ln->forward(ctx, x);
- x = Flux::modulate(ctx->ggml_ctx, x, shift_mlp, scale_mlp, true);
+ x = Flux::modulate(ctx, x, shift_mlp, scale_mlp, true);
x = ggml_add(ctx->ggml_ctx, residual, ggml_mul(ctx->ggml_ctx, mlp->forward(ctx, x), gate_mlp));
return x;
}
@@ -274,7 +274,7 @@ namespace ErnieImage {
auto shift = mods[1];
x = norm->forward(ctx, x);
- x = Flux::modulate(ctx->ggml_ctx, x, shift, scale);
+ x = Flux::modulate(ctx, x, shift, scale);
return x;
}
};
diff --git a/src/model/diffusion/fibo2.hpp b/src/model/diffusion/fibo2.hpp
index a642cb4..0c42235 100644
--- a/src/model/diffusion/fibo2.hpp
+++ b/src/model/diffusion/fibo2.hpp
@@ -262,7 +262,7 @@ namespace Fibo2 {
}
return m;
};
- auto h = ZImage::modulate(c, norm("attention_norm1", x), select(0, false));
+ auto h = ZImage::modulate(ctx, norm("attention_norm1", x), select(0, false));
h = std::dynamic_pointer_cast<Attention>(blocks["attention"])->forward(ctx, h, h, pe);
x = ggml_add(c, x, ggml_mul(c, norm("attention_norm2", h), select(1, true)));
if (injection) {
@@ -270,7 +270,7 @@ namespace Fibo2 {
h = ggml_mul(c, norm("cross_attn_norm2", h), ggml_tanh(c, params["cross_attn_gate"]));
x = ggml_add(c, x, h);
}
- h = ZImage::modulate(c, norm("ffn_norm1", x), select(2, false));
+ h = ZImage::modulate(ctx, norm("ffn_norm1", x), select(2, false));
h = std::dynamic_pointer_cast<FeedForward>(blocks["feed_forward"])->forward(ctx, h);
return ggml_add(c, x, ggml_mul(c, norm("ffn_norm2", h), select(3, true)));
}
diff --git a/src/model/diffusion/flux.hpp b/src/model/diffusion/flux.hpp
index c9e3f5f..40edca0 100644
--- a/src/model/diffusion/flux.hpp
+++ b/src/model/diffusion/flux.hpp
@@ -413,7 +413,7 @@ namespace Flux {
}
};
- __STATIC_INLINE__ ggml_tensor* modulate(ggml_context* ctx,
+ __STATIC_INLINE__ ggml_tensor* modulate(GGMLRunnerContext* ctx,
ggml_tensor* x,
ggml_tensor* shift,
ggml_tensor* scale,
@@ -422,12 +422,10 @@ namespace Flux {
// scale: [N, C]
// shift: [N, C]
if (!skip_reshape) {
- scale = ggml_reshape_3d(ctx, scale, scale->ne[0], 1, scale->ne[1]); // [N, 1, C]
- shift = ggml_reshape_3d(ctx, shift, shift->ne[0], 1, shift->ne[1]); // [N, 1, C]
+ scale = ggml_reshape_3d(ctx->ggml_ctx, scale, scale->ne[0], 1, scale->ne[1]); // [N, 1, C]
+ shift = ggml_reshape_3d(ctx->ggml_ctx, shift, shift->ne[0], 1, shift->ne[1]); // [N, 1, C]
}
- x = ggml_add(ctx, x, ggml_mul(ctx, x, scale));
- x = ggml_add(ctx, x, shift);
- return x;
+ return ggml_ext_modulate(ctx, x, shift, scale);
}
struct DoubleStreamBlock : public GGMLBlock {
@@ -540,7 +538,7 @@ namespace Flux {
// prepare image for attention
auto img_modulated = img_norm1->forward(ctx, img);
- img_modulated = Flux::modulate(ctx->ggml_ctx, img_modulated, img_mod1.shift, img_mod1.scale);
+ img_modulated = Flux::modulate(ctx, img_modulated, img_mod1.shift, img_mod1.scale);
auto img_qkv = img_attn->pre_attention(ctx, img_modulated); // q,k,v: [N, n_img_token, n_head, d_head]
auto img_q = img_qkv[0];
auto img_k = img_qkv[1];
@@ -548,7 +546,7 @@ namespace Flux {
// prepare txt for attention
auto txt_modulated = txt_norm1->forward(ctx, txt);
- txt_modulated = Flux::modulate(ctx->ggml_ctx, txt_modulated, txt_mod1.shift, txt_mod1.scale);
+ txt_modulated = Flux::modulate(ctx, txt_modulated, txt_mod1.shift, txt_mod1.scale);
auto txt_qkv = txt_attn->pre_attention(ctx, txt_modulated); // q,k,v: [N, n_txt_token, n_head, d_head]
auto txt_q = txt_qkv[0];
auto txt_k = txt_qkv[1];
@@ -580,14 +578,14 @@ namespace Flux {
// calculate the img bloks
img = ggml_add(ctx->ggml_ctx, img, ggml_mul(ctx->ggml_ctx, img_attn->post_attention(ctx, img_attn_out), img_mod1.gate));
- auto img_mlp_out = img_mlp->forward(ctx, Flux::modulate(ctx->ggml_ctx, img_norm2->forward(ctx, img), img_mod2.shift, img_mod2.scale));
+ auto img_mlp_out = img_mlp->forward(ctx, Flux::modulate(ctx, img_norm2->forward(ctx, img), img_mod2.shift, img_mod2.scale));
img = ggml_add(ctx->ggml_ctx, img, ggml_mul(ctx->ggml_ctx, img_mlp_out, img_mod2.gate));
// calculate the txt bloks
txt = ggml_add(ctx->ggml_ctx, txt, ggml_mul(ctx->ggml_ctx, txt_attn->post_attention(ctx, txt_attn_out), txt_mod1.gate));
- auto txt_mlp_out = txt_mlp->forward(ctx, Flux::modulate(ctx->ggml_ctx, txt_norm2->forward(ctx, txt), txt_mod2.shift, txt_mod2.scale));
+ auto txt_mlp_out = txt_mlp->forward(ctx, Flux::modulate(ctx, txt_norm2->forward(ctx, txt), txt_mod2.shift, txt_mod2.scale));
txt = ggml_add(ctx->ggml_ctx, txt, ggml_mul(ctx->ggml_ctx, txt_mlp_out, txt_mod2.gate));
return {img, txt};
@@ -670,7 +668,7 @@ namespace Flux {
}
}
- auto x_mod = Flux::modulate(ctx->ggml_ctx, pre_norm->forward(ctx, x), mod.shift, mod.scale);
+ auto x_mod = Flux::modulate(ctx, pre_norm->forward(ctx, x), mod.shift, mod.scale);
auto qkv_mlp = linear1->forward(ctx, x_mod); // [N, n_token, hidden_size * 3 + mlp_hidden_dim*mlp_mult_factor]
int64_t head_dim = hidden_size / num_heads;
@@ -752,7 +750,7 @@ namespace Flux {
scale = m_vec[1];
}
- x = Flux::modulate(ctx->ggml_ctx, norm_final->forward(ctx, x), shift, scale);
+ x = Flux::modulate(ctx, norm_final->forward(ctx, x), shift, scale);
x = linear->forward(ctx, x);
return x;
diff --git a/src/model/diffusion/flux3.hpp b/src/model/diffusion/flux3.hpp
index 6e4cd64..fac1b1a 100644
--- a/src/model/diffusion/flux3.hpp
+++ b/src/model/diffusion/flux3.hpp
@@ -83,7 +83,7 @@ namespace Flux3 {
int64_t offset = 0;
for (size_t i = 0; i < lengths.size(); ++i) {
auto part = token_slice(gc, x, offset, lengths[i]);
- part = Flux::modulate(gc, ggml_norm(gc, part, 1e-6f), mods[i].shift, mods[i].scale);
+ part = Flux::modulate(ctx, ggml_norm(gc, part, 1e-6f), mods[i].shift, mods[i].scale);
x_mod = x_mod ? ggml_concat(gc, x_mod, part, 1) : part;
offset += lengths[i];
}
diff --git a/src/model/diffusion/ideogram4.hpp b/src/model/diffusion/ideogram4.hpp
index 5a5cb00..270344d 100644
--- a/src/model/diffusion/ideogram4.hpp
+++ b/src/model/diffusion/ideogram4.hpp
@@ -103,9 +103,9 @@ namespace Ideogram4 {
return ggml_reshape_3d(ctx, x, hidden_size * LLM_HIDDEN_STATE_LAYERS, token_count, batch_count);
}
- __STATIC_INLINE__ ggml_tensor* modulate(ggml_context* ctx, ggml_tensor* x, ggml_tensor* scale) {
- scale = to_token_modulation(ctx, scale);
- return ggml_add(ctx, x, ggml_mul(ctx, x, scale));
+ __STATIC_INLINE__ ggml_tensor* modulate(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* scale) {
+ scale = to_token_modulation(ctx->ggml_ctx, scale);
+ return ggml_ext_modulate(ctx, x, nullptr, scale);
}
__STATIC_INLINE__ ggml_tensor* patchify(ggml_context* ctx, ggml_tensor* x, const Ideogram4Config& config) {
@@ -285,13 +285,13 @@ namespace Ideogram4 {
auto gate_mlp = to_token_modulation(ctx->ggml_ctx, ggml_tanh(ctx->ggml_ctx, mods[3]));
auto attn_out = attention_norm1->forward(ctx, x);
- attn_out = modulate(ctx->ggml_ctx, attn_out, scale_msa);
+ attn_out = modulate(ctx, attn_out, scale_msa);
attn_out = attention->forward(ctx, attn_out, pe, mask);
attn_out = attention_norm2->forward(ctx, attn_out);
x = ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, attn_out, gate_msa));
auto ffn_out = ffn_norm1->forward(ctx, x);
- ffn_out = modulate(ctx->ggml_ctx, ffn_out, scale_mlp);
+ ffn_out = modulate(ctx, ffn_out, scale_mlp);
ffn_out = feed_forward->forward(ctx, ffn_out);
ffn_out = ffn_norm2->forward(ctx, ffn_out);
x = ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, ffn_out, gate_mlp));
@@ -337,7 +337,7 @@ namespace Ideogram4 {
auto scale = adaln_modulation->forward(ctx, ggml_silu(ctx->ggml_ctx, c));
x = norm_final->forward(ctx, x);
- x = modulate(ctx->ggml_ctx, x, scale);
+ x = modulate(ctx, x, scale);
x = linear->forward(ctx, x);
return x;
}
diff --git a/src/model/diffusion/krea2.hpp b/src/model/diffusion/krea2.hpp
index d21574b..fba817a 100644
--- a/src/model/diffusion/krea2.hpp
+++ b/src/model/diffusion/krea2.hpp
@@ -465,8 +465,8 @@ namespace Krea2 {
auto pre_x_main = ggml_view_3d(ctx->ggml_ctx, pre_x, D, len_main, B, nb1, nb2, 0);
auto pre_x_refs = ggml_view_3d(ctx->ggml_ctx, pre_x, D, len_refs, B, nb1, nb2, len_main * nb1);
- auto attn_in_main = Flux::modulate(ctx->ggml_ctx, pre_x_main, mods_main[1], mods_main[0], true);
- auto attn_in_refs = Flux::modulate(ctx->ggml_ctx, pre_x_refs, mods_refs[1], mods_refs[0], true);
+ auto attn_in_main = Flux::modulate(ctx, pre_x_main, mods_main[1], mods_main[0], true);
+ auto attn_in_refs = Flux::modulate(ctx, pre_x_refs, mods_refs[1], mods_refs[0], true);
auto attn_input = ggml_concat(ctx->ggml_ctx, attn_in_main, attn_in_refs, 1);
@@ -487,8 +487,8 @@ namespace Krea2 {
auto post_x_main = ggml_view_3d(ctx->ggml_ctx, post_x, D, len_main, B, post_x->nb[1], post_x->nb[2], 0);
auto post_x_refs = ggml_view_3d(ctx->ggml_ctx, post_x, D, len_refs, B, post_x->nb[1], post_x->nb[2], len_main * post_x->nb[1]);
- auto mlp_in_main = Flux::modulate(ctx->ggml_ctx, post_x_main, mods_main[4], mods_main[3], true);
- auto mlp_in_refs = Flux::modulate(ctx->ggml_ctx, post_x_refs, mods_refs[4], mods_refs[3], true);
+ auto mlp_in_main = Flux::modulate(ctx, post_x_main, mods_main[4], mods_main[3], true);
+ auto mlp_in_refs = Flux::modulate(ctx, post_x_refs, mods_refs[4], mods_refs[3], true);
auto mlp_input = ggml_concat(ctx->ggml_ctx, mlp_in_main, mlp_in_refs, 1);
auto mlp_out = mlp->forward(ctx, mlp_input);
@@ -503,7 +503,7 @@ namespace Krea2 {
x = ggml_add(ctx->ggml_ctx, x, mlp_res);
} else {
auto mods = mod->forward(ctx, vec);
- auto attn_input = Flux::modulate(ctx->ggml_ctx,
+ auto attn_input = Flux::modulate(ctx,
prenorm->forward(ctx, x),
mods[1],
mods[0],
@@ -511,7 +511,7 @@ namespace Krea2 {
auto attn_out = attn->forward(ctx, attn_input, pe);
x = ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, attn_out, mods[2]));
- auto mlp_input = Flux::modulate(ctx->ggml_ctx,
+ auto mlp_input = Flux::modulate(ctx,
postnorm->forward(ctx, x),
mods[4],
mods[3],
@@ -595,7 +595,7 @@ namespace Krea2 {
auto modulation = std::dynamic_pointer_cast<KreaFinalModulation>(blocks["modulation"]);
auto mods = modulation->forward(ctx, vec);
- x = Flux::modulate(ctx->ggml_ctx,
+ x = Flux::modulate(ctx,
norm->forward(ctx, x),
mods[1],
mods[0],
diff --git a/src/model/diffusion/lens.hpp b/src/model/diffusion/lens.hpp
index 62ad6d3..bcfefd9 100644
--- a/src/model/diffusion/lens.hpp
+++ b/src/model/diffusion/lens.hpp
@@ -237,9 +237,9 @@ namespace Lens {
auto txt_mod_vec = ggml_ext_chunk(ctx->ggml_ctx, txt_mod_params, 6, 0);
auto img_normed = img_norm1->forward(ctx, img);
- auto img_modulated = Flux::modulate(ctx->ggml_ctx, img_normed, img_mod_vec[0], img_mod_vec[1]);
+ auto img_modulated = Flux::modulate(ctx, img_normed, img_mod_vec[0], img_mod_vec[1]);
auto txt_normed = txt_norm1->forward(ctx, txt);
- auto txt_modulated = Flux::modulate(ctx->ggml_ctx, txt_normed, txt_mod_vec[0], txt_mod_vec[1]);
+ auto txt_modulated = Flux::modulate(ctx, txt_normed, txt_mod_vec[0], txt_mod_vec[1]);
auto [img_attn_output, txt_attn_output] = attn->forward(ctx, img_modulated, txt_modulated, pe);
@@ -247,9 +247,9 @@ namespace Lens {
txt = ggml_add(ctx->ggml_ctx, txt, ggml_mul(ctx->ggml_ctx, txt_attn_output, txt_mod_vec[2]));
auto img_normed2 = img_norm2->forward(ctx, img);
- auto img_modulated2 = Flux::modulate(ctx->ggml_ctx, img_normed2, img_mod_vec[3], img_mod_vec[4]);
+ auto img_modulated2 = Flux::modulate(ctx, img_normed2, img_mod_vec[3], img_mod_vec[4]);
auto txt_normed2 = txt_norm2->forward(ctx, txt);
- auto txt_modulated2 = Flux::modulate(ctx->ggml_ctx, txt_normed2, txt_mod_vec[3], txt_mod_vec[4]);
+ auto txt_modulated2 = Flux::modulate(ctx, txt_normed2, txt_mod_vec[3], txt_mod_vec[4]);
img = ggml_add(ctx->ggml_ctx, img, ggml_mul(ctx->ggml_ctx, img_mlp->forward(ctx, img_modulated2), img_mod_vec[5]));
txt = ggml_add(ctx->ggml_ctx, txt, ggml_mul(ctx->ggml_ctx, txt_mlp->forward(ctx, txt_modulated2), txt_mod_vec[5]));
@@ -272,7 +272,7 @@ namespace Lens {
auto scale = mods[0];
auto shift = mods[1];
x = ggml_norm(ctx->ggml_ctx, x, eps);
- return Flux::modulate(ctx->ggml_ctx, x, shift, scale);
+ return Flux::modulate(ctx, x, shift, scale);
}
};
diff --git a/src/model/diffusion/lingbot_video.hpp b/src/model/diffusion/lingbot_video.hpp
index 1ea7b1c..d15d754 100644
--- a/src/model/diffusion/lingbot_video.hpp
+++ b/src/model/diffusion/lingbot_video.hpp
@@ -485,12 +485,12 @@ namespace LingBotVideo {
auto scale_mlp = mods[4];
auto gate_mlp = ggml_tanh(ctx->ggml_ctx, mods[5]);
- auto attn_in = Flux::modulate(ctx->ggml_ctx, norm1->forward(ctx, x), shift_msa, scale_msa, true);
+ auto attn_in = Flux::modulate(ctx, norm1->forward(ctx, x), shift_msa, scale_msa, true);
auto attn_out = attn->forward(ctx, attn_in, pe, attention_mask);
attn_out = norm_post_attn->forward(ctx, attn_out);
x = ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, gate_msa, attn_out));
- auto ffn_in = Flux::modulate(ctx->ggml_ctx, norm2->forward(ctx, x), shift_mlp, scale_mlp, true);
+ auto ffn_in = Flux::modulate(ctx, norm2->forward(ctx, x), shift_mlp, scale_mlp, true);
auto ffn_out = ffn->forward(ctx, ffn_in);
ffn_out = norm_post_ffn->forward(ctx, ffn_out);
x = ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, gate_mlp, ffn_out));
@@ -587,7 +587,7 @@ namespace LingBotVideo {
2,
0);
hidden_states = norm_out->forward(ctx, hidden_states);
- hidden_states = Flux::modulate(ctx->ggml_ctx, hidden_states, final_mods[0], final_mods[1], true);
+ hidden_states = Flux::modulate(ctx, hidden_states, final_mods[0], final_mods[1], true);
hidden_states = proj_out->forward(ctx, hidden_states);
hidden_states = ggml_ext_slice(ctx->ggml_ctx, hidden_states, 1, 0, n_img);
diff --git a/src/model/diffusion/ltxv.hpp b/src/model/diffusion/ltxv.hpp
index 23c5f5f..32e3801 100644
--- a/src/model/diffusion/ltxv.hpp
+++ b/src/model/diffusion/ltxv.hpp
@@ -38,12 +38,12 @@ namespace LTXV {
return mod;
}
- __STATIC_INLINE__ ggml_tensor* modulate(ggml_context* ctx,
+ __STATIC_INLINE__ ggml_tensor* modulate(GGMLRunnerContext* ctx,
ggml_tensor* x,
ggml_tensor* shift,
ggml_tensor* scale) {
- shift = align_token_modulation(ctx, x, shift);
- scale = align_token_modulation(ctx, x, scale);
+ shift = align_token_modulation(ctx->ggml_ctx, x, shift);
+ scale = align_token_modulation(ctx->ggml_ctx, x, scale);
return Flux::modulate(ctx, x, shift, scale, true);
}
@@ -821,7 +821,7 @@ namespace LTXV {
auto gate_mlp = mods[5];
auto x_norm = rms_norm(ctx->ggml_ctx, x);
- x_norm = LTXV::modulate(ctx->ggml_ctx, x_norm, shift_msa, scale_msa);
+ x_norm = LTXV::modulate(ctx, x_norm, shift_msa, scale_msa);
auto msa = attn1->forward(ctx, x_norm, nullptr, self_attention_mask, pe);
x = ggml_add(ctx->ggml_ctx, x, apply_gate(ctx->ggml_ctx, msa, gate_msa));
@@ -831,12 +831,12 @@ namespace LTXV {
auto gate_q = mods[8];
auto q = rms_norm(ctx->ggml_ctx, x);
- q = LTXV::modulate(ctx->ggml_ctx, q, shift_q, scale_q);
+ q = LTXV::modulate(ctx, q, shift_q, scale_q);
auto context_mod = context;
if (prompt_timestep != nullptr) {
auto prompt_mods = get_prompt_scale_shift_values(ctx, prompt_timestep);
- context_mod = LTXV::modulate(ctx->ggml_ctx, context_mod, prompt_mods[0], prompt_mods[1]);
+ context_mod = LTXV::modulate(ctx, context_mod, prompt_mods[0], prompt_mods[1]);
}
auto mca = attn2->forward(ctx, q, context_mod, attention_mask, nullptr, nullptr);
@@ -847,7 +847,7 @@ namespace LTXV {
}
auto y = rms_norm(ctx->ggml_ctx, x);
- y = LTXV::modulate(ctx->ggml_ctx, y, shift_mlp, scale_mlp);
+ y = LTXV::modulate(ctx, y, shift_mlp, scale_mlp);
auto mlp_out = ff->forward(ctx, y);
x = ggml_add(ctx->ggml_ctx, x, apply_gate(ctx->ggml_ctx, mlp_out, gate_mlp));
return x;
@@ -1199,11 +1199,11 @@ namespace LTXV {
if (cross_attention_adaln) {
auto q_mods = get_ada_values(ctx, table, timestep, dim, 9, 6, 3);
auto q = rms_norm(ctx->ggml_ctx, x);
- q = LTXV::modulate(ctx->ggml_ctx, q, q_mods[0], q_mods[1]);
+ q = LTXV::modulate(ctx, q, q_mods[0], q_mods[1]);
auto context_mod = context;
if (prompt_timestep != nullptr && prompt_table != nullptr) {
auto p_mods = get_ada_values(ctx, prompt_table, prompt_timestep, dim, 2);
- context_mod = LTXV::modulate(ctx->ggml_ctx, context_mod, p_mods[0], p_mods[1]);
+ context_mod = LTXV::modulate(ctx, context_mod, p_mods[0], p_mods[1]);
}
auto out = attn->forward(ctx, q, context_mod, attention_mask, nullptr, nullptr);
return apply_gate(ctx->ggml_ctx, out, q_mods[2]);
@@ -1250,7 +1250,7 @@ namespace LTXV {
auto v_mods = get_ada_values(ctx, v_table, v_timestep, v_dim, cross_attention_adaln ? 9 : 6);
auto v_norm = rms_norm(ctx->ggml_ctx, vx);
- v_norm = LTXV::modulate(ctx->ggml_ctx, v_norm, v_mods[0], v_mods[1]);
+ v_norm = LTXV::modulate(ctx, v_norm, v_mods[0], v_mods[1]);
auto v_sa = attn1->forward(ctx, v_norm, nullptr, self_attention_mask, v_pe);
vx = ggml_add(ctx->ggml_ctx, vx, apply_gate(ctx->ggml_ctx, v_sa, v_mods[2]));
auto v_txt = apply_text_cross_attention(ctx,
@@ -1268,7 +1268,7 @@ namespace LTXV {
if (run_ax) {
auto a_mods = get_ada_values(ctx, a_table, a_timestep, a_dim, cross_attention_adaln ? 9 : 6);
auto a_norm = rms_norm(ctx->ggml_ctx, ax);
- a_norm = LTXV::modulate(ctx->ggml_ctx, a_norm, a_mods[0], a_mods[1]);
+ a_norm = LTXV::modulate(ctx, a_norm, a_mods[0], a_mods[1]);
auto a_sa = audio_attn1->forward(ctx, a_norm, nullptr, nullptr, a_pe);
ax = ggml_add(ctx->ggml_ctx, ax, apply_gate(ctx->ggml_ctx, a_sa, a_mods[2]));
auto a_txt = apply_text_cross_attention(ctx,
@@ -1291,8 +1291,8 @@ namespace LTXV {
auto a2v_video_table = ggml_ext_slice(ctx->ggml_ctx, params["scale_shift_table_a2v_ca_video"], 1, 0, 4);
auto a2v_audio = get_ada_values(ctx, a2v_audio_table, a_cross_scale_shift_timestep, a_dim, 4);
auto a2v_video = get_ada_values(ctx, a2v_video_table, v_cross_scale_shift_timestep, v_dim, 4);
- auto vx_scaled = LTXV::modulate(ctx->ggml_ctx, vx_norm3, a2v_video[1], a2v_video[0]);
- auto ax_scaled = LTXV::modulate(ctx->ggml_ctx, ax_norm3, a2v_audio[1], a2v_audio[0]);
+ auto vx_scaled = LTXV::modulate(ctx, vx_norm3, a2v_video[1], a2v_video[0]);
+ auto ax_scaled = LTXV::modulate(ctx, ax_norm3, a2v_audio[1], a2v_audio[0]);
auto a2v_out = audio_to_video_attn->forward(ctx, vx_scaled, ax_scaled, nullptr, v_cross_pe, a_cross_pe);
auto a2v_gate_table = ggml_ext_slice(ctx->ggml_ctx, params["scale_shift_table_a2v_ca_video"], 1, 4, 5);
auto a2v_gate = get_ada_values(ctx, a2v_gate_table, v_cross_gate_timestep, v_dim, 1)[0];
@@ -1304,8 +1304,8 @@ namespace LTXV {
auto v2a_video_table = ggml_ext_slice(ctx->ggml_ctx, params["scale_shift_table_a2v_ca_video"], 1, 0, 4);
auto v2a_audio = get_ada_values(ctx, v2a_audio_table, a_cross_scale_shift_timestep, a_dim, 4);
auto v2a_video = get_ada_values(ctx, v2a_video_table, v_cross_scale_shift_timestep, v_dim, 4);
- auto ax_scaled = LTXV::modulate(ctx->ggml_ctx, ax_norm3, v2a_audio[3], v2a_audio[2]);
- auto vx_scaled = LTXV::modulate(ctx->ggml_ctx, vx_norm3, v2a_video[3], v2a_video[2]);
+ auto ax_scaled = LTXV::modulate(ctx, ax_norm3, v2a_audio[3], v2a_audio[2]);
+ auto vx_scaled = LTXV::modulate(ctx, vx_norm3, v2a_video[3], v2a_video[2]);
auto v2a_out = video_to_audio_attn->forward(ctx, ax_scaled, vx_scaled, nullptr, a_cross_pe, v_cross_pe);
auto v2a_gate_table = ggml_ext_slice(ctx->ggml_ctx, params["scale_shift_table_a2v_ca_audio"], 1, 4, 5);
auto v2a_gate = get_ada_values(ctx, v2a_gate_table, a_cross_gate_timestep, a_dim, 1)[0];
@@ -1313,14 +1313,14 @@ namespace LTXV {
}
auto a_ff_mods = get_ada_values(ctx, a_table, a_timestep, a_dim, cross_attention_adaln ? 9 : 6, 3, 3);
auto ax_scaled = rms_norm(ctx->ggml_ctx, ax);
- ax_scaled = LTXV::modulate(ctx->ggml_ctx, ax_scaled, a_ff_mods[0], a_ff_mods[1]);
+ ax_scaled = LTXV::modulate(ctx, ax_scaled, a_ff_mods[0], a_ff_mods[1]);
auto a_ff_out = audio_ff->forward(ctx, ax_scaled);
ax = ggml_add(ctx->ggml_ctx, ax, apply_gate(ctx->ggml_ctx, a_ff_out, a_ff_mods[2]));
}
auto v_ff_mods = get_ada_values(ctx, v_table, v_timestep, v_dim, cross_attention_adaln ? 9 : 6, 3, 3);
auto vx_scaled = rms_norm(ctx->ggml_ctx, vx);
- vx_scaled = LTXV::modulate(ctx->ggml_ctx, vx_scaled, v_ff_mods[0], v_ff_mods[1]);
+ vx_scaled = LTXV::modulate(ctx, vx_scaled, v_ff_mods[0], v_ff_mods[1]);
auto v_ff_out = ff->forward(ctx, vx_scaled);
vx = ggml_add(ctx->ggml_ctx, vx, apply_gate(ctx->ggml_ctx, v_ff_out, v_ff_mods[2]));
@@ -1720,14 +1720,14 @@ namespace LTXV {
auto v_shift_scale = get_output_scale_shift(ctx, params["scale_shift_table"], v_embedded_time, config.hidden_size);
vx = norm_out->forward(ctx, vx);
- vx = LTXV::modulate(ctx->ggml_ctx, vx, v_shift_scale[0], v_shift_scale[1]);
+ vx = LTXV::modulate(ctx, vx, v_shift_scale[0], v_shift_scale[1]);
vx = proj_out->forward(ctx, vx);
vx = unpatchify_video(ctx, vx, width, height, frames);
if (ax != nullptr && audio_time > 0) {
auto a_shift_scale = get_output_scale_shift(ctx, params["audio_scale_shift_table"], a_embedded_time, config.audio_hidden_size);
ax = audio_norm_out->forward(ctx, ax);
- ax = LTXV::modulate(ctx->ggml_ctx, ax, a_shift_scale[0], a_shift_scale[1]);
+ ax = LTXV::modulate(ctx, ax, a_shift_scale[0], a_shift_scale[1]);
ax = audio_proj_out->forward(ctx, ax);
ax = unpatchify_audio(ctx, ax, audio_time);
}
diff --git a/src/model/diffusion/minimax_h3.hpp b/src/model/diffusion/minimax_h3.hpp
index 0915b52..a7c2f1f 100644
--- a/src/model/diffusion/minimax_h3.hpp
+++ b/src/model/diffusion/minimax_h3.hpp
@@ -445,9 +445,7 @@ namespace MiniMaxH3 {
inputs.push_back(ggml_ext_slice(ctx->ggml_ctx, x, 1, segment.start, segment.end, false));
}
auto modulate = [&](ggml_tensor* value, const std::vector<ggml_tensor*>& row, int shift, int scale) {
- return ggml_add(ctx->ggml_ctx,
- ggml_add(ctx->ggml_ctx, value, ggml_mul(ctx->ggml_ctx, value, row[scale])),
- row[shift]);
+ return ggml_ext_modulate(ctx, value, row[shift], row[scale]);
};
ggml_tensor* qkv = nullptr;
for (size_t i = 0; i < segments.size(); ++i) {
diff --git a/src/model/diffusion/mmdit.hpp b/src/model/diffusion/mmdit.hpp
index cdf65f3..613cfc1 100644
--- a/src/model/diffusion/mmdit.hpp
+++ b/src/model/diffusion/mmdit.hpp
@@ -372,18 +372,16 @@ public:
}
};
-__STATIC_INLINE__ ggml_tensor* modulate(ggml_context* ctx,
+__STATIC_INLINE__ ggml_tensor* modulate(GGMLRunnerContext* ctx,
ggml_tensor* x,
ggml_tensor* shift,
ggml_tensor* scale) {
// x: [N, L, C]
// scale: [N, C]
// shift: [N, C]
- scale = ggml_reshape_3d(ctx, scale, scale->ne[0], 1, scale->ne[1]); // [N, 1, C]
- shift = ggml_reshape_3d(ctx, shift, shift->ne[0], 1, shift->ne[1]); // [N, 1, C]
- x = ggml_add(ctx, x, ggml_mul(ctx, x, scale));
- x = ggml_add(ctx, x, shift);
- return x;
+ scale = ggml_reshape_3d(ctx->ggml_ctx, scale, scale->ne[0], 1, scale->ne[1]); // [N, 1, C]
+ shift = ggml_reshape_3d(ctx->ggml_ctx, shift, shift->ne[0], 1, shift->ne[1]); // [N, 1, C]
+ return ggml_ext_modulate(ctx, x, shift, scale);
}
struct DismantledBlock : public GGMLBlock {
@@ -455,10 +453,10 @@ public:
auto x_norm = norm1->forward(ctx, x);
- auto attn_in = modulate(ctx->ggml_ctx, x_norm, shift_msa, scale_msa);
+ auto attn_in = modulate(ctx, x_norm, shift_msa, scale_msa);
auto qkv = attn->pre_attention(ctx, attn_in);
- auto attn2_in = modulate(ctx->ggml_ctx, x_norm, shift_msa2, scale_msa2);
+ auto attn2_in = modulate(ctx, x_norm, shift_msa2, scale_msa2);
auto qkv2 = attn2->pre_attention(ctx, attn2_in);
return {qkv, qkv2, {x, gate_msa, shift_mlp, scale_mlp, gate_mlp, gate_msa2}};
@@ -488,13 +486,13 @@ public:
auto scale_mlp = m_vec[4]; // [N, hidden_size]
auto gate_mlp = m_vec[5]; // [N, hidden_size]
- auto attn_in = modulate(ctx->ggml_ctx, norm1->forward(ctx, x), shift_msa, scale_msa);
+ auto attn_in = modulate(ctx, norm1->forward(ctx, x), shift_msa, scale_msa);
auto qkv = attn->pre_attention(ctx, attn_in);
return {qkv, {x, gate_msa, shift_mlp, scale_mlp, gate_mlp}};
} else {
- auto attn_in = modulate(ctx->ggml_ctx, norm1->forward(ctx, x), shift_msa, scale_msa);
+ auto attn_in = modulate(ctx, norm1->forward(ctx, x), shift_msa, scale_msa);
auto qkv = attn->pre_attention(ctx, attn_in);
return {qkv, {nullptr, nullptr, nullptr, nullptr, nullptr}};
@@ -533,7 +531,7 @@ public:
x = ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, attn_out, gate_msa));
x = ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, attn2_out, gate_msa2));
- auto mlp_out = mlp->forward(ctx, modulate(ctx->ggml_ctx, norm2->forward(ctx, x), shift_mlp, scale_mlp));
+ auto mlp_out = mlp->forward(ctx, modulate(ctx, norm2->forward(ctx, x), shift_mlp, scale_mlp));
x = ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, mlp_out, gate_mlp));
return x;
@@ -565,7 +563,7 @@ public:
attn_out = attn->post_attention(ctx, attn_out);
x = ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, attn_out, gate_msa));
- auto mlp_out = mlp->forward(ctx, modulate(ctx->ggml_ctx, norm2->forward(ctx, x), shift_mlp, scale_mlp));
+ auto mlp_out = mlp->forward(ctx, modulate(ctx, norm2->forward(ctx, x), shift_mlp, scale_mlp));
x = ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, mlp_out, gate_mlp));
return x;
@@ -756,7 +754,7 @@ public:
auto shift = m_vec[0]; // [N, hidden_size]
auto scale = m_vec[1]; // [N, hidden_size]
- x = modulate(ctx->ggml_ctx, norm_final->forward(ctx, x), shift, scale);
+ x = modulate(ctx, norm_final->forward(ctx, x), shift, scale);
x = linear->forward(ctx, x);
return x;
diff --git a/src/model/diffusion/pixart.hpp b/src/model/diffusion/pixart.hpp
index 45f9e43..dc6e2c2 100644
--- a/src/model/diffusion/pixart.hpp
+++ b/src/model/diffusion/pixart.hpp
@@ -198,11 +198,11 @@ namespace PixArt {
return ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, y, g));
};
- auto h = modulate(ctx->ggml_ctx, norm(ctx->ggml_ctx, x), mv[0], mv[1]);
+ auto h = modulate(ctx, norm(ctx->ggml_ctx, x), mv[0], mv[1]);
x = gate(attn1->forward(ctx, h, h), mv[2]);
// ada_norm_single: no norm before cross-attention (PixArtMS.py)
x = ggml_add(ctx->ggml_ctx, x, attn2->forward(ctx, x, context, context_mask));
- h = modulate(ctx->ggml_ctx, norm(ctx->ggml_ctx, x), mv[3], mv[4]);
+ h = modulate(ctx, norm(ctx->ggml_ctx, x), mv[3], mv[4]);
h = proj->forward(ctx, h);
h = ggml_ext_gelu(ctx->ggml_ctx, h, true);
h = fc2->forward(ctx, h);
@@ -288,7 +288,7 @@ namespace PixArt {
ggml_reshape_2d(ctx->ggml_ctx, ggml_ext_cont(ctx->ggml_ctx, ss), config.hidden_size * 2, N),
2, 0);
h = std::dynamic_pointer_cast<LayerNorm>(blocks["final_layer.norm_final"])->forward(ctx, h);
- h = modulate(ctx->ggml_ctx, h, parts[0], parts[1]);
+ h = modulate(ctx, h, parts[0], parts[1]);
h = std::dynamic_pointer_cast<Linear>(blocks["final_layer.linear"])->forward(ctx, h); // [N, hp*wp, p*p*out_ch]
h = DiT::unpatchify(ctx->ggml_ctx, h, hp, wp, static_cast<int>(p), static_cast<int>(p), false);
return h; // [N, out_channels, H, W]
diff --git a/src/model/diffusion/qwen_image.hpp b/src/model/diffusion/qwen_image.hpp
index 4c473da..6435222 100644
--- a/src/model/diffusion/qwen_image.hpp
+++ b/src/model/diffusion/qwen_image.hpp
@@ -350,11 +350,11 @@ namespace Qwen {
auto txt_mod_param_vec = get_mod_params_vec(ctx->ggml_ctx, txt_mod_params);
auto img_normed = img_norm1->forward(ctx, img);
- auto img_modulated = Flux::modulate(ctx->ggml_ctx, img_normed, img_mod_param_vec[0], img_mod_param_vec[1], modulate_index != nullptr);
+ auto img_modulated = Flux::modulate(ctx, img_normed, img_mod_param_vec[0], img_mod_param_vec[1], modulate_index != nullptr);
auto img_gate1 = img_mod_param_vec[2];
auto txt_normed = txt_norm1->forward(ctx, txt);
- auto txt_modulated = Flux::modulate(ctx->ggml_ctx, txt_normed, txt_mod_param_vec[0], txt_mod_param_vec[1]);
+ auto txt_modulated = Flux::modulate(ctx, txt_normed, txt_mod_param_vec[0], txt_mod_param_vec[1]);
auto txt_gate1 = txt_mod_param_vec[2];
auto [img_attn_output, txt_attn_output] = attn->forward(ctx, img_modulated, txt_modulated, pe);
@@ -363,11 +363,11 @@ namespace Qwen {
txt = ggml_add(ctx->ggml_ctx, txt, ggml_mul(ctx->ggml_ctx, txt_attn_output, txt_gate1));
auto img_normed2 = img_norm2->forward(ctx, img);
- auto img_modulated2 = Flux::modulate(ctx->ggml_ctx, img_normed2, img_mod_param_vec[3], img_mod_param_vec[4], modulate_index != nullptr);
+ auto img_modulated2 = Flux::modulate(ctx, img_normed2, img_mod_param_vec[3], img_mod_param_vec[4], modulate_index != nullptr);
auto img_gate2 = img_mod_param_vec[5];
auto txt_normed2 = txt_norm2->forward(ctx, txt);
- auto txt_modulated2 = Flux::modulate(ctx->ggml_ctx, txt_normed2, txt_mod_param_vec[3], txt_mod_param_vec[4]);
+ auto txt_modulated2 = Flux::modulate(ctx, txt_normed2, txt_mod_param_vec[3], txt_mod_param_vec[4]);
auto txt_gate2 = txt_mod_param_vec[5];
auto img_mlp_out = img_mlp->forward(ctx, img_modulated2);
@@ -407,7 +407,7 @@ namespace Qwen {
auto shift = mods[1];
x = norm->forward(ctx, x);
- x = Flux::modulate(ctx->ggml_ctx, x, shift, scale);
+ x = Flux::modulate(ctx, x, shift, scale);
return x;
}
diff --git a/src/model/diffusion/z_image.hpp b/src/model/diffusion/z_image.hpp
index 229c6aa..1c66c96 100644
--- a/src/model/diffusion/z_image.hpp
+++ b/src/model/diffusion/z_image.hpp
@@ -294,17 +294,16 @@ namespace ZImage {
}
};
- __STATIC_INLINE__ ggml_tensor* modulate(ggml_context* ctx,
+ __STATIC_INLINE__ ggml_tensor* modulate(GGMLRunnerContext* ctx,
ggml_tensor* x,
ggml_tensor* scale) {
// x: [N, L, C]
// scale: [N, C], or [N, L, C] when the caller modulates per token (LLaDA-Image editing
// feeds a per-token timestep embedding so each segment carries its own modulation).
if (scale->ne[1] != x->ne[1]) {
- scale = ggml_reshape_3d(ctx, scale, scale->ne[0], 1, scale->ne[1]); // [N, 1, C]
+ scale = ggml_reshape_3d(ctx->ggml_ctx, scale, scale->ne[0], 1, scale->ne[1]); // [N, 1, C]
}
- x = ggml_add(ctx, x, ggml_mul(ctx, x, scale));
- return x;
+ return ggml_ext_modulate(ctx, x, nullptr, scale);
}
struct JointTransformerBlock : public GGMLBlock {
@@ -361,14 +360,14 @@ namespace ZImage {
auto gate_mlp = mods[3];
auto residual = x;
- x = modulate(ctx->ggml_ctx, attention_norm1->forward(ctx, x), scale_msa);
+ x = modulate(ctx, attention_norm1->forward(ctx, x), scale_msa);
x = attention->forward(ctx, x, pe, mask);
x = attention_norm2->forward(ctx, x);
x = ggml_mul(ctx->ggml_ctx, x, ggml_tanh(ctx->ggml_ctx, gate_msa));
x = ggml_add(ctx->ggml_ctx, x, residual);
residual = x;
- x = modulate(ctx->ggml_ctx, ffn_norm1->forward(ctx, x), scale_mlp);
+ x = modulate(ctx, ffn_norm1->forward(ctx, x), scale_mlp);
x = feed_forward->forward(ctx, x);
x = ffn_norm2->forward(ctx, x);
x = ggml_mul(ctx->ggml_ctx, x, ggml_tanh(ctx->ggml_ctx, gate_mlp));
@@ -505,14 +504,14 @@ namespace ZImage {
auto gate_mlp = mods[3];
auto residual = c;
- c = modulate(ctx->ggml_ctx, attention_norm1->forward(ctx, c), scale_msa);
+ c = modulate(ctx, attention_norm1->forward(ctx, c), scale_msa);
c = forward_attention(ctx, c, pe);
c = attention_norm2->forward(ctx, c);
c = ggml_mul(ctx->ggml_ctx, c, ggml_tanh(ctx->ggml_ctx, gate_msa));
c = ggml_add(ctx->ggml_ctx, c, residual);
residual = c;
- c = modulate(ctx->ggml_ctx, ffn_norm1->forward(ctx, c), scale_mlp);
+ c = modulate(ctx, ffn_norm1->forward(ctx, c), scale_mlp);
c = feed_forward->forward(ctx, c);
c = ffn_norm2->forward(ctx, c);
c = ggml_mul(ctx->ggml_ctx, c, ggml_tanh(ctx->ggml_ctx, gate_mlp));
@@ -544,7 +543,7 @@ namespace ZImage {
auto scale = adaLN_modulation_1->forward(ctx, ggml_silu(ctx->ggml_ctx, c)); // [N, hidden_size]
x = norm_final->forward(ctx, x);
- x = modulate(ctx->ggml_ctx, x, scale);
+ x = modulate(ctx, x, scale);
x = linear->forward(ctx, x);
return x;