diff --git a/CMakeLists.txt b/CMakeLists.txt index 235556fac..15bd91742 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 cf41592fb..6b42ea117 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 6787c443c..5c059892e 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 000000000..84a90ab68 --- /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 + +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(static_cast(t->data) + offset); + } + + static void modulate_compute(ggml_tensor* dst, int ith, int nth, void*) { + auto* output = static_cast(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_backend_reg_get_proc_address(reg, "ggml_ext_cpu_parallel_v1")); + GGML_ASSERT(run); + auto* dst = nodes[0]; + const auto address = reinterpret_cast(dst->data); + bool alias = false; + for (const auto* input : dst->src) { + if (input) { + const auto src = reinterpret_cast(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 a4f5bd32b..d63fd1f09 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 c711f6fac..eba629c2a 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 7bca6aa33..f366899fe 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 d5d159be1..cf7896114 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(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(a->data), y = reinterpret_cast(b->data); return x <= y ? y - x < ggml_nbytes(a) : x - y < ggml_nbytes(b); } - template + template 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, + auto status = cuda_set_shared_memory(quantize, 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(t->data) : nullptr; }; - quantize<<>>( + quantize<<>>( data(op.input) + op.gate_offset, output, reinterpret_cast(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(plan, op, output, stream); break; + case MOD: + error = launch(plan, op, output, stream); + break; case SWIGLU: case SILU: error = launch(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 748c08e7d..04abf97ac 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 + template __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 000000000..c4af6f6e0 --- /dev/null +++ b/src/core/cuda/modulate.cu @@ -0,0 +1,95 @@ +#include +#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(dst->data); + for (const auto* input : {x, scale, shift}) { + if (input) { + const auto src = reinterpret_cast(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(context.stream); + auto* output = static_cast(alias ? ggml_backend_buffer_get_base(e->workspace) : dst->data); + const unsigned blocks = unsigned(std::min((p.count + 255) / 256, 65535)); + modulate<<>>(static_cast(x->data), static_cast(scale->data), + shift ? static_cast(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 d03b114cc..291edd62a 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 6d72caedf..ca1b231c8 100644 --- a/src/core/ggml_extend.cpp +++ b/src/core/ggml_extend.cpp @@ -6,6 +6,7 @@ #include #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 4a97704d3..9ddfb0ad6 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 1f0aa1de5..27089f6c3 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(node->data); + const auto src = reinterpret_cast(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 49b0e7605..86c9121a7 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 8e55d9997..42ee931a2 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 4caadfbef..7e3e67d48 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 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 bindings = {{x}, {mod}, {shift}, packed, scales}; + std::vector 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 000000000..456144702 --- /dev/null +++ b/src/core/vulkan/modulate.cpp @@ -0,0 +1,134 @@ +#include "modulate.h" +#include +#include +#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_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(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 000000000..df07626d2 --- /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 3e4b3a217..3cac110c9 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 66e1a0ba3..e9e33b88a 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 ac67ff9cd..3f652e158 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(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 59b06cc2f..15badb983 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 a642cb4df..0c42235f8 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(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(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 c9e3f5faa..40edca0ed 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 6e4cd6408..fac1b1a56 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 5a5cb000c..270344dc9 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 d21574b63..fba817aa6 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(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 62ad6d373..bcfefd97b 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 1ea7b1c0c..d15d75408 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 23c5f5f29..32e38012b 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 3dd720cad..2596a4e95 100644 --- a/src/model/diffusion/minimax_h3.hpp +++ b/src/model/diffusion/minimax_h3.hpp @@ -415,9 +415,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& 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 cdf65f300..613cfc11a 100644 --- a/src/model/diffusion/mmdit.hpp +++ b/src/model/diffusion/mmdit.hpp @@ -372,18 +372,16 @@ class SelfAttention : public GGMLBlock { } }; -__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 @@ struct DismantledBlock : public GGMLBlock { 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 @@ struct DismantledBlock : public GGMLBlock { 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 @@ struct DismantledBlock : public GGMLBlock { 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 @@ struct DismantledBlock : public GGMLBlock { 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 @@ struct FinalLayer : public GGMLBlock { 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 45f9e4317..dc6e2c2e7 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(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(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(p), static_cast(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 4c473da35..6435222f5 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 229c6aaee..1c66c96b7 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;