diff --git a/src/core/cuda/rms_affine.cu b/src/core/cuda/rms_affine.cu index d9ac351ed..9f4192b25 100644 --- a/src/core/cuda/rms_affine.cu +++ b/src/core/cuda/rms_affine.cu @@ -2,6 +2,7 @@ #include #include #include "bf16.h" +#include "core/ggml_extensions.h" #include "rms_affine.h" namespace SD_CUDA_NAMESPACE { @@ -18,6 +19,28 @@ namespace SD_CUDA_NAMESPACE { return t && t->type == GGML_TYPE_F32 && ggml_is_contiguous(t); } + static int modulate_match(ggml_backend_dev_t, const ggml_cgraph* graph, int first, void*) { + auto* g = const_cast(graph); + if (disabled() || first + 2 >= ggml_graph_n_nodes(g)) + return 0; + const auto* norm = ggml_graph_node(g, first); + const auto* affine = ggml_graph_node(g, first + 1); + const auto* modulate = ggml_graph_node(g, first + 2); + if (norm->op != GGML_OP_RMS_NORM || affine->op != GGML_OP_MUL || affine->src[0] != norm || + !sd_ggml_modulate_describe(modulate) || modulate->src[0] != affine) + return 0; + const auto* x = norm->src[0]; + if (!f32(x) || !f32(norm) || !f32(affine) || x->ne[0] < 32 || x->ne[0] > 8192 || + ggml_nrows(x) > INT_MAX || !ggml_are_same_shape(x, norm) || !ggml_are_same_shape(x, affine)) + return 0; + for (const auto* parameter : {affine->src[1], modulate->src[1], modulate->src[2]}) { + if (parameter && (!f32(parameter) || parameter == norm || parameter == affine || + parameter->ne[0] != x->ne[0] || ggml_nelements(parameter) != x->ne[0])) + return 0; + } + return affine->src[1] ? 3 : 0; + } + static int match(ggml_backend_dev_t, const ggml_cgraph* graph, int first, void*) { auto* g = const_cast(graph); if (disabled() || first + 3 >= ggml_graph_n_nodes(g)) @@ -44,13 +67,17 @@ namespace SD_CUDA_NAMESPACE { static void optimize(ggml_backend_dev_t, ggml_cgraph* graph, void*) { if (disabled()) return; - for (int i = 2; i + 1 < ggml_graph_n_nodes(graph); ++i) { + for (int i = 2; i < ggml_graph_n_nodes(graph); ++i) { const auto* mul = ggml_graph_node(graph, i); - const auto* add = ggml_graph_node(graph, i + 1); - if (mul->op != GGML_OP_MUL || add->op != GGML_OP_ADD || (add->src[0] != mul && add->src[1] != mul)) - continue; + if (!sd_ggml_modulate_describe(mul)) { + if (i + 1 >= ggml_graph_n_nodes(graph)) + continue; + const auto* add = ggml_graph_node(graph, i + 1); + if (mul->op != GGML_OP_MUL || add->op != GGML_OP_ADD || (add->src[0] != mul && add->src[1] != mul)) + continue; + } const auto* affine = mul->src[0]; - if (affine->op != GGML_OP_MUL || affine->src[0]->op != GGML_OP_RMS_NORM) + if (!affine || affine->op != GGML_OP_MUL || !affine->src[0] || affine->src[0]->op != GGML_OP_RMS_NORM) continue; for (int j = i - 2; j >= 0 && j >= i - 32; --j) { if (ggml_graph_node(graph, j) == affine->src[0] && ggml_graph_node(graph, j + 1) == affine) { @@ -68,64 +95,136 @@ namespace SD_CUDA_NAMESPACE { return value; } - template - static __global__ void kernel(const float* x, const float* weight, const float* gate, const float* residual, T* dst, int width, float eps) { + enum class Kind { Residual, + SelfResidual, + Modulate }; + + struct Operation { + Kind kind = Kind::Residual; + const ggml_tensor* input; + const ggml_tensor* weight; + const ggml_tensor* gate; + const ggml_tensor* residual = nullptr; + const ggml_tensor* shift = nullptr; + ggml_tensor* output; + float eps; + }; + + static Operation describe(ggml_tensor* const* nodes, int count) { + Operation op; + op.input = nodes[0]->src[0]; + op.weight = nodes[1]->src[1]; + op.gate = nodes[2]->src[1]; + op.output = nodes[count - 1]; + std::memcpy(&op.eps, nodes[0]->op_params, sizeof(op.eps)); + if (count == 3) { + op.kind = Kind::Modulate; + op.shift = op.output->src[2]; + } else { + op.residual = op.output->src[0] == nodes[2] ? op.output->src[1] : op.output->src[0]; + if (op.residual == nodes[1]) { + op.kind = Kind::SelfResidual; + op.residual = nullptr; + } + } + return op; + } + + template + static __global__ void kernel(const float* x, const float* weight, const float* gate, const float* residual, const float* shift, T* dst, int width, float eps) { const int col0 = threadIdx.x; const int64_t base = int64_t(blockIdx.x) * width; - float values[8], sum = 0.f; - for (int j = 0, col = col0; col < width; ++j, col += Block) { - values[j] = x[base + col]; - sum = __fmaf_rn(values[j], values[j], sum); - } - sum = warp_sum(sum); + float values[Groups][8]; __shared__ float sums[32]; const int lane = col0 % 32; - if (lane == 0) - sums[col0 / 32] = sum; +#pragma unroll + for (int group = 0; group < Groups; ++group) { + float sum = 0.f; + for (int j = 0, col = col0 + group * Block; col < width; ++j, col += Block * Groups) { + values[group][j] = x[base + col]; + sum = __fmaf_rn(values[group][j], values[group][j], sum); + } + sum = warp_sum(sum); + if (lane == 0) + sums[(col0 + group * Block) / 32] = sum; + } + // Groups preserve the native reduction tree when using fewer physical threads. __syncthreads(); - sum = warp_sum(lane < Block / 32 ? sums[lane] : 0.f); + const float sum = warp_sum(lane < Block * Groups / 32 ? sums[lane] : 0.f); const float scale = rsqrtf(sum / width + eps); - for (int j = 0, col = col0; col < width; ++j, col += Block) { - const float value = __fmul_rn(__fmul_rn(scale, values[j]), weight[col]); - const float r = Modulate ? value : residual[base + col]; - dst[base + col] = T(__fadd_rn(__fmul_rn(value, gate[col]), r)); +#pragma unroll + for (int group = 0; group < Groups; ++group) { + for (int j = 0, col = col0 + group * Block; col < width; ++j, col += Block * Groups) { + float value = __fmul_rn(__fmul_rn(scale, values[group][j]), weight[col]); + if constexpr (Op == Kind::Modulate) { + value = __fadd_rn(value, __fmul_rn(value, gate[col])); + if (shift) + value = __fadd_rn(value, shift[col]); + } else { + const float r = Op == Kind::SelfResidual ? value : residual[base + col]; + value = __fadd_rn(__fmul_rn(value, gate[col]), r); + } + dst[base + col] = T(value); + } } } static bool overlaps(const ggml_tensor* a, const ggml_tensor* b) { + if (!b) + return false; const auto x = reinterpret_cast(a->data), y = reinterpret_cast(b->data); return x <= y ? y - x < ggml_nbytes(a) : x - y < ggml_nbytes(b); } - static ggml_status prepare(ggml_backend_t, ggml_tensor* const* nodes, int, void**, size_t* workspace, void*) { - *workspace = ggml_nbytes(nodes[3]); + static ggml_status prepare(ggml_backend_t, ggml_tensor* const* nodes, int count, void**, size_t* workspace, void*) { + *workspace = ggml_nbytes(nodes[count - 1]); return GGML_STATUS_SUCCESS; } static void destroy(void*, void*) {} - static ggml_status execute(const ggml_ext_execution_context* execution, ggml_tensor* const* nodes, int, void*, void*) { + template + static CUDAError launch(CUDAStream stream, const Operation& op, T* output) { + auto compute = kernel<256, 1, Op, T>; + int threads = 256; + if (op.input->ne[0] >= 1024) { + if constexpr (Op == Kind::Modulate) { + compute = kernel<512, 2, Op, T>; + threads = 512; + } else { + compute = kernel<1024, 1, Op, T>; + threads = 1024; + } + } + compute<<>>( + static_cast(op.input->data), static_cast(op.weight->data), + static_cast(op.gate->data), op.residual ? static_cast(op.residual->data) : nullptr, + op.shift ? static_cast(op.shift->data) : nullptr, output, int(op.input->ne[0]), op.eps); + return cuda_last_error(); + } + + static ggml_status execute(const ggml_ext_execution_context* execution, ggml_tensor* const* nodes, int count, void*, void*) { ggml_ext_gpu_context context; if (!cuda_context(execution->backend, &context)) return GGML_STATUS_FAILED; - const auto* x = nodes[0]->src[0]; - const auto* w = nodes[1]->src[1]; - const auto* gate = nodes[2]->src[1]; - auto* dst = nodes[3]; - const auto* residual = dst->src[0] == nodes[2] ? dst->src[1] : dst->src[0]; - const bool modulation = residual == nodes[1]; - const bool alias = (overlaps(dst, x) && dst->data != x->data) || overlaps(dst, w) || overlaps(dst, gate) || - (!modulation && overlaps(dst, residual) && dst->data != residual->data); - auto* output = static_cast(alias ? ggml_backend_buffer_get_base(execution->workspace) : dst->data); - float eps; - std::memcpy(&eps, nodes[0]->op_params, sizeof(eps)); - auto launch = x->ne[0] < 1024 ? (modulation ? kernel<256, true> : kernel<256, false>) - : (modulation ? kernel<1024, true> : kernel<1024, false>); + const auto op = describe(nodes, count); + auto* dst = op.output; + const bool alias = (overlaps(dst, op.input) && dst->data != op.input->data) || + (overlaps(dst, op.residual) && dst->data != op.residual->data) || + overlaps(dst, op.weight) || overlaps(dst, op.gate) || overlaps(dst, op.shift); + auto* output = static_cast(alias ? ggml_backend_buffer_get_base(execution->workspace) : dst->data); const auto stream = static_cast(context.stream); - const int threads = x->ne[0] < 1024 ? 256 : 1024; - launch<<>>( - static_cast(x->data), static_cast(w->data), static_cast(gate->data), - modulation ? nullptr : static_cast(residual->data), output, int(x->ne[0]), eps); - auto error = cuda_last_error(); + CUDAError error = cuda_invalid_value; + switch (op.kind) { + case Kind::Residual: + error = launch(stream, op, output); + break; + case Kind::SelfResidual: + error = launch(stream, op, output); + break; + case Kind::Modulate: + error = launch(stream, op, output); + break; + } if (error == cuda_success && alias) error = cuda_copy_async(dst->data, output, ggml_nbytes(dst), stream); return error == cuda_success ? GGML_STATUS_SUCCESS : GGML_STATUS_FAILED; @@ -137,19 +236,12 @@ namespace SD_CUDA_NAMESPACE { } ggml_status rms_affine_pack(CUDAStream stream, ggml_tensor* const* nodes, void* output) { - const auto* x = nodes[0]->src[0]; - const auto* residual = nodes[3]->src[0] == nodes[2] ? nodes[3]->src[1] : nodes[3]->src[0]; - const bool modulation = residual == nodes[1]; - auto launch = x->ne[0] < 1024 ? (modulation ? sd_rms_affine::kernel<256, true, CUDABFloat16> : sd_rms_affine::kernel<256, false, CUDABFloat16>) - : (modulation ? sd_rms_affine::kernel<1024, true, CUDABFloat16> : sd_rms_affine::kernel<1024, false, CUDABFloat16>); - float eps; - std::memcpy(&eps, nodes[0]->op_params, sizeof(eps)); - const int threads = x->ne[0] < 1024 ? 256 : 1024; - launch<<>>( - static_cast(x->data), static_cast(nodes[1]->src[1]->data), - static_cast(nodes[2]->src[1]->data), modulation ? nullptr : static_cast(residual->data), - static_cast(output), int(x->ne[0]), eps); - return cuda_last_error() == cuda_success ? GGML_STATUS_SUCCESS : GGML_STATUS_FAILED; + const auto op = sd_rms_affine::describe(nodes, 4); + auto* dst = static_cast(output); + const auto error = op.kind == sd_rms_affine::Kind::SelfResidual + ? sd_rms_affine::launch(stream, op, dst) + : sd_rms_affine::launch(stream, op, dst); + return error == cuda_success ? GGML_STATUS_SUCCESS : GGML_STATUS_FAILED; } void register_rms_affine(ggml_backend_dev_t device) { @@ -164,5 +256,11 @@ namespace SD_CUDA_NAMESPACE { desc.destroy = sd_rms_affine::destroy; desc.execute = sd_rms_affine::execute; GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE)); + desc.name = "sd.rms_affine_modulate." SD_CUDA_NAME ".v1"; + // Prefer the existing fusion that also produces ConvRot input when available. + desc.priority = 200; + desc.match = sd_rms_affine::modulate_match; + desc.optimize = nullptr; + GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE)); } } diff --git a/src/core/cuda/rope.cu b/src/core/cuda/rope.cu index 1b724b0a7..99cc824f1 100644 --- a/src/core/cuda/rope.cu +++ b/src/core/cuda/rope.cu @@ -169,7 +169,7 @@ namespace SD_CUDA_NAMESPACE { return false; const auto* x = norm->src[0]; const auto* weight = mul->src[1]; - return x && weight && x->ne[0] >= 32 && x->ne[0] <= 256 && x->ne[0] % 32 == 0 && + return x && weight && weight != norm && x->ne[0] >= 32 && x->ne[0] <= 256 && x->ne[0] % 32 == 0 && x->type == GGML_TYPE_F32 && weight->type == GGML_TYPE_F32 && norm->type == GGML_TYPE_F32 && ggml_are_same_shape(norm, mul) && ggml_are_same_shape(x, norm) && weight->ne[0] == x->ne[0] && ggml_nelements(weight) == x->ne[0] && @@ -232,6 +232,13 @@ namespace SD_CUDA_NAMESPACE { ggml_are_same_shape(rope, scale) && ggml_are_same_shape(rope, copy) && ggml_is_contiguous(copy)) return 5; } + if (first + 3 < ggml_graph_n_nodes(g)) { + const auto* rope = ggml_graph_node(g, first + 2); + const auto* copy = ggml_graph_node(g, first + 3); + if (copy->op == GGML_OP_CPY && copy->src[0] == rope && copy->type == GGML_TYPE_F16 && + ggml_are_same_shape(rope, copy) && ggml_is_contiguous(copy)) + return 4; + } return 3; } @@ -293,7 +300,7 @@ namespace SD_CUDA_NAMESPACE { void* output = alias ? ggml_backend_buffer_get_base(execution->workspace) : dst->data; CUDAError error; if (count > 1) { - error = rms_rope_launch(stream, nodes, output, count == 5, count == 5 ? param_f32(nodes[3], 0) : 1.f); + error = rms_rope_launch(stream, nodes, output, count > 3, count == 5 ? param_f32(nodes[3], 0) : 1.f); } else { RopeApplyStrides strides; const auto* pe = rope->src[1]; diff --git a/src/model/diffusion/moss_image.hpp b/src/model/diffusion/moss_image.hpp index 1f5dc1668..a0fae7448 100644 --- a/src/model/diffusion/moss_image.hpp +++ b/src/model/diffusion/moss_image.hpp @@ -116,7 +116,7 @@ namespace MOSSImage { auto norm = [&](const char* name, ggml_tensor* h) { return std::dynamic_pointer_cast(blocks[name])->forward(ctx, h); }; auto part = [&](int i) { return mod[i]; }; auto y = norm("norm1", x); - y = ggml_add(c, ggml_add(c, y, ggml_mul(c, y, part(1))), part(0)); + y = ggml_ext_modulate(ctx, y, part(0), part(1)); auto qkv = linear("proj.qkv", y); auto split = [&](int i) { return ggml_view_4d(c, qkv, config.head_dim, config.dim / config.head_dim, x->ne[1], x->ne[2], @@ -130,7 +130,7 @@ namespace MOSSImage { y = ggml_mul(c, y, part(2)); x = ggml_add(c, x, y); y = norm("norm2", x); - y = ggml_add(c, ggml_add(c, y, ggml_mul(c, y, part(4))), part(3)); + y = ggml_ext_modulate(ctx, y, part(3), part(4)); auto up_gate = linear("proj.ffn_up_gate", y); y = ggml_ext_situ(ctx, up_gate); y = linear("proj.ffn_down", y); @@ -191,7 +191,7 @@ namespace MOSSImage { x = std::dynamic_pointer_cast(blocks["final_norm_video"])->forward(ctx, x); auto shift = ggml_ext_slice(c, head, 0, 0, config.dim); auto scale = ggml_ext_slice(c, head, 0, config.dim, 2 * config.dim); - x = ggml_add(c, ggml_add(c, x, ggml_mul(c, x, scale)), shift); + x = ggml_ext_modulate(ctx, x, shift, scale); x = linear("final_linear_video", x); return DiT::unpatchify(c, x, image->ne[1], image->ne[0], 1, 1); }