diff --git a/CMakeLists.txt b/CMakeLists.txt index 0dd9f228..7461b49c 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -348,15 +348,15 @@ if (SD_CUDA) target_sources(${SD_LIB} PRIVATE src/core/cuda/rope.cu src/core/cuda/int8.cu) set_source_files_properties(src/core/cuda/rope.cu src/core/cuda/int8.cu PROPERTIES COMPILE_OPTIONS "--use_fast_math") if (GGML_CUDA_CUTLASS_PATH) - target_sources(${SD_LIB} PRIVATE src/core/cuda/cute_attention.cu src/core/cuda/i8_projection.cu) - set_source_files_properties(src/core/cuda/cute_attention.cu src/core/cuda/i8_projection.cu PROPERTIES + target_sources(${SD_LIB} PRIVATE src/core/cuda/cute_attention.cu src/core/cuda/i8_projection.cu src/core/cuda/i8_gated_residual.cu) + set_source_files_properties(src/core/cuda/cute_attention.cu src/core/cuda/i8_projection.cu src/core/cuda/i8_gated_residual.cu PROPERTIES INCLUDE_DIRECTORIES "${GGML_CUDA_CUTLASS_PATH}/include" COMPILE_OPTIONS "--use_fast_math;--extended-lambda;--expt-relaxed-constexpr") if (MSVC) - set_property(SOURCE src/core/cuda/cute_attention.cu src/core/cuda/i8_projection.cu APPEND PROPERTY COMPILE_OPTIONS -Xcompiler=/Zc:preprocessor,/Zc:__cplusplus) + set_property(SOURCE src/core/cuda/cute_attention.cu src/core/cuda/i8_projection.cu src/core/cuda/i8_gated_residual.cu APPEND PROPERTY COMPILE_OPTIONS -Xcompiler=/Zc:preprocessor,/Zc:__cplusplus) endif() if (GGML_CUDA_FORCE_CUBLAS) - set_property(SOURCE src/core/cuda/i8_projection.cu APPEND PROPERTY COMPILE_DEFINITIONS SD_CUDA_FORCE_CUBLAS) + set_property(SOURCE src/core/cuda/i8_projection.cu src/core/cuda/i8_gated_residual.cu APPEND PROPERTY COMPILE_DEFINITIONS SD_CUDA_FORCE_CUBLAS) endif() set_property(SOURCE src/core/ggml_extensions.cpp APPEND PROPERTY COMPILE_DEFINITIONS SD_CUDA_CUTE_ATTENTION) endif() diff --git a/docs/ggml_extensions.md b/docs/ggml_extensions.md index 97e8efc2..565dfebc 100644 --- a/docs/ggml_extensions.md +++ b/docs/ggml_extensions.md @@ -196,7 +196,25 @@ CUDA GEMM requires Turing or newer, contiguous tensors, K/N divisible by four, float-aligned inputs, and at most 65535 rows; CUDA ConvRot uses groups of 256. Unsupported custom nodes can fall back to CPU through the scheduler. -The builtin CUDA/Vulkan INT8 GEMM, ConvRot quantization, and their other fusion rules +`src/core/cuda/i8_gated_residual.cu` implements +`INT8 GEMM -> [Scale] -> Mul(gate) -> Add(residual) -> [Concat]` as a registered +CUDA fusion, for both builtin GEMM and custom I8 nodes with separate activation +scales. The corresponding native CUDA fusion rules and CUTLASS residual epilogue +have been removed from ggml. Its optimizer moves independent gate producers ahead +of GEMM and optional scaling through `ggml_ext_graph_delay`, before fusion matching. +The former native INT8 affine reorder rule has also moved into this extension. +It preserves the original tile selection, small-row +threshold, Stream-K scheduling, Ada column bands, and separate FP32 rounding of +gate multiplication and residual addition. The plan bounds Stream-K workspace +over possible occupancy values using host metadata; execution verifies the +actual requirement against that bound. Device occupancy is cached on first +execution; each launch rebuilds parameters for the current tensor addresses and +initializes its Stream-K barriers. Overlapping outputs use scratch, except +for an exact in-place residual without concatenation. External consumers and +unsupported layouts retain ordinary operator execution. Fusion and force-cuBLAS +disable settings are honored. + +The builtin CUDA/Vulkan INT8 GEMM, ConvRot quantization, and remaining fusion rules still need to be externalized. Vulkan execution services and kernels also remain pending. This implementation does not yet externalize all model-specific kernels or add dynamic dtypes. diff --git a/ggml b/ggml index c4e98f10..a6903f4a 160000 --- a/ggml +++ b/ggml @@ -1 +1 @@ -Subproject commit c4e98f1002427f3fef63cd6ecfbe4781ee67a808 +Subproject commit a6903f4a06c5d3b980964fedece61587e86c5678 diff --git a/src/core/cuda/i8_gated_residual.cu b/src/core/cuda/i8_gated_residual.cu new file mode 100644 index 00000000..37d67792 --- /dev/null +++ b/src/core/cuda/i8_gated_residual.cu @@ -0,0 +1,465 @@ +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include "core/int8_ops.h" + +namespace sd_i8_gated { + struct BandedStreamK : cutlass::gemm::threadblock::ThreadblockSwizzleStreamK { + using Base = cutlass::gemm::threadblock::ThreadblockSwizzleStreamK; + using Base::Base; + + CUTLASS_DEVICE cutlass::gemm::GemmCoord get_tile_offset(int index) const { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == 890 + const auto shape = tiled_shape(); + if (problem_size.m() >= 4096 && problem_size.n() >= 16384 && shape.m() >= shape.n()) { + // Column bands retain weight tiles in Ada's L2 across successive row tiles. + constexpr int band_width = 16; + if (index >= shape.m() * shape.n()) { + return {shape.m(), shape.n(), get_batch_idx()}; + } + const int band = index / (shape.m() * band_width); + const int width = min(band_width, shape.n() - band * band_width); + const int within = index - band * shape.m() * band_width; + return {within / width, band * band_width + within % width, get_batch_idx()}; + } +#endif + return Base::get_tile_offset(index); + } + }; + + template + struct TileConfig { + using Tile = cutlass::gemm::GemmShape; + using Warp = cutlass::gemm::GemmShape; + }; + + enum Tile { + TILE_128_256, + TILE_128_128, + TILE_64_128, + TILE_16_128, + }; + + static bool small_rows_enabled(int64_t rows, int64_t n, int64_t k) { + static const char* setting = std::getenv("GGML_CUDA_I8_SMALL_ROWS"); + return setting ? std::atoi(setting) != 0 : rows < 32 && n >= 4096 && k >= 4096; + } + + static Tile select_tile(int64_t rows, int64_t n, int nsm) { + if (rows < 32) { + return TILE_16_128; + } + const int64_t columns = (n + 127) / 128; + if (((rows + 15) / 16) * columns <= nsm) { + return TILE_16_128; + } + if (((rows + 63) / 64) * columns <= nsm) { + return TILE_64_128; + } + const int64_t wide_tiles = ((rows + 127) / 128) * ((n + 255) / 256); + const int64_t square_tiles = ((rows + 127) / 128) * columns; + // A narrower tile costs more per output element, but can waste less of the last wave. + const double wide_cost = double((wide_tiles + nsm - 1) / nsm) / wide_tiles; + const double square_cost = 1.08 * double((square_tiles + nsm - 1) / nsm) / square_tiles; + return wide_cost <= square_cost ? TILE_128_256 : TILE_128_128; + } + + template + struct MulRound { + CUTLASS_DEVICE T operator()(const T& a, const T& b) const { + T result; + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < T::kElements; ++i) { + result[i] = __fmul_rn(a[i], b[i]); + } + return result; + } + }; + + template + struct AddRound { + CUTLASS_DEVICE T operator()(const T& a, const T& b) const { + T result; + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < T::kElements; ++i) { + result[i] = __fadd_rn(a[i], b[i]); + } + return result; + } + }; + + template + struct KernelTypes { + using Tile = typename Config::Tile; + using Warp = typename Config::Warp; + using Swizzle = std::conditional_t>; + using Map = cutlass::epilogue::threadblock::OutputTileThreadLayout; + using Col = cutlass::epilogue::threadblock::VisitorColBroadcast>; + using Row = cutlass::epilogue::threadblock::VisitorRowBroadcast>; + using Mul = cutlass::epilogue::threadblock::VisitorCompute; + using ScaleX = cutlass::epilogue::threadblock::Sm80EVT; + using ScaleW = cutlass::epilogue::threadblock::Sm80EVT; + using Residual = cutlass::epilogue::threadblock::VisitorAuxLoad>; + using GateMul = cutlass::epilogue::threadblock::VisitorCompute; + using Add = cutlass::epilogue::threadblock::VisitorCompute; + using Scalar = cutlass::epilogue::threadblock::VisitorScalarBroadcast; + using ScaleOut = cutlass::epilogue::threadblock::Sm80EVT; + using Scaled = std::conditional_t; + using Gated = cutlass::epilogue::threadblock::Sm80EVT; + using WithResidual = cutlass::epilogue::threadblock::Sm80EVT; + using Output = WithResidual; + using Store = cutlass::epilogue::threadblock::VisitorAuxStore>; + using Epilogue = cutlass::epilogue::threadblock::Sm80EVT; + using Kernel = typename cutlass::gemm::kernel::DefaultGemmWithVisitor< + int8_t, + cutlass::layout::RowMajor, + cutlass::ComplexTransform::kNone, + 16, + int8_t, + cutlass::layout::ColumnMajor, + cutlass::ComplexTransform::kNone, + 16, + float, + cutlass::layout::RowMajor, + 4, + int, + float, + cutlass::arch::OpClassTensorOp, + cutlass::arch::Sm80, + Tile, + Warp, + cutlass::gemm::GemmShape<16, 8, 32>, + Epilogue, + Swizzle, + 4, + cutlass::arch::OpMultiplyAddSaturate>::GemmKernel; + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + }; + + struct DeviceInfo { + int ordinal = -1; + int major = 0; + int sms = 0; + int threads_per_sm = 0; + size_t shared = 0; + }; + static DeviceInfo device_info(ggml_backend_dev_t device) { + static std::mutex mutex; + static std::unordered_map cache; + std::lock_guard lock(mutex); + auto found = cache.find(device); + if (found != cache.end()) + return found->second; + DeviceInfo info; + auto reg = ggml_backend_dev_backend_reg(device); + auto query = reinterpret_cast(ggml_backend_reg_get_proc_address(reg, "ggml_ext_cuda_device_v1")); + int ordinal = query ? query(device) : -1; + cudaDeviceProp props = {}; + if (ordinal >= 0 && cudaGetDeviceProperties(&props, ordinal) == cudaSuccess) + info = {ordinal, props.major, props.multiProcessorCount, props.maxThreadsPerMultiProcessor, size_t(props.sharedMemPerBlockOptin)}; + cache.emplace(device, info); + return info; + } + static float param_f32(const ggml_tensor* tensor, int index) { + float value; + std::memcpy(&value, tensor->op_params + index, sizeof(value)); + return value; + } + static bool aligned(const ggml_tensor* tensor) { + if (!tensor) + return false; + size_t offset = 0; + while (tensor->view_src) { + offset += tensor->view_offs; + tensor = tensor->view_src; + } + return (reinterpret_cast(tensor->data) + offset) % 16 == 0; + } + struct Operation { + SDI8Operation mat; + const ggml_tensor* gate; + const ggml_tensor* residual; + const ggml_tensor* prefix; + float scale; + int64_t rows, n, k; + }; + static Operation describe(ggml_tensor* const* nodes, int count) { + Operation op = {}; + sd_ggml_i8_describe(nodes[0], &op.mat); + const bool concat = nodes[count - 1]->op == GGML_OP_CONCAT; + const auto* add = nodes[count - (concat ? 2 : 1)]; + const auto* mul = nodes[count - (concat ? 3 : 2)]; + op.gate = mul->src[1]; + op.residual = add->src[0] == mul ? add->src[1] : add->src[0]; + op.prefix = concat ? nodes[count - 1]->src[0] : nullptr; + op.scale = nodes[1]->op == GGML_OP_SCALE ? param_f32(nodes[1], 0) : 1.f; + op.rows = ggml_nrows(nodes[0]); + op.n = nodes[0]->ne[0]; + op.k = op.mat.weight->ne[0]; + return op; + } + template + static typename Types::Gemm::Arguments arguments(const Operation& op, float* output, bool metadata) { + const auto* x = metadata ? nullptr : static_cast(op.mat.input->data); + const auto* w = metadata ? nullptr : static_cast(op.mat.weight->data); + const auto* xs = metadata ? nullptr : op.mat.packed ? reinterpret_cast(x + op.k * ((op.rows + 3) / 4 * 4)) + : static_cast(op.mat.input_scale->data); + const auto* ws = metadata ? nullptr : static_cast(op.mat.weight_scale->data); + const auto* gate = metadata ? nullptr : static_cast(op.gate->data); + auto* residual = metadata ? nullptr : static_cast(op.residual->data); + typename Types::ScaleW::Arguments scales = {{{}, {xs, 0.f, {}}, {}}, {ws, 0.f, {}}, {}}; + typename Types::Scaled::Arguments scaled; + if constexpr (std::is_same::value) + scaled = scales; + else + scaled = {scales, {{op.scale}}, {}}; + typename Types::Epilogue::Arguments epilogue = {{{scaled, {gate, 0.f, {}}, {}}, {residual, 0.f, {op.n, cute::_1{}, op.rows * op.n}}, {}}, {output, {op.n, cute::_1{}, op.rows * op.n}}}; + return {cutlass::gemm::GemmUniversalMode::kGemm, {int(op.rows), int(op.n), int(op.k)}, 1, epilogue, x, w, nullptr, nullptr, op.rows * op.k, op.n * op.k, 0, 0, op.k, op.k, 0, 0}; + } + struct Plan { + DeviceInfo device; + size_t kernel_workspace = 0; + size_t output_offset = 0; + int occupancy = 0; + ggml_status (*launch)(Plan&, const Operation&, float*, void*, cudaStream_t) = nullptr; + }; + template + static ggml_status launch(Plan& plan, const Operation& op, float* output, void* scratch, cudaStream_t stream) { + using Kernel = typename Types::Kernel; + auto args = arguments(op, output, false); + if (Kernel::can_implement(args) != cutlass::Status::kSuccess) + return GGML_STATUS_FAILED; + if (plan.occupancy == 0) { + if (sizeof(typename Kernel::SharedStorage) >= (48 << 10) && + cudaFuncSetAttribute(cutlass::Kernel2, cudaFuncAttributeMaxDynamicSharedMemorySize, sizeof(typename Kernel::SharedStorage)) != cudaSuccess) + return GGML_STATUS_FAILED; + int occupancy = 0; + if (cudaOccupancyMaxActiveBlocksPerMultiprocessorWithFlags(&occupancy, cutlass::Kernel2, + Kernel::kThreadCount, sizeof(typename Kernel::SharedStorage), cudaOccupancyDisableCachingOverride) != cudaSuccess || + occupancy <= 0) + return GGML_STATUS_FAILED; + plan.occupancy = occupancy; + } + typename Kernel::Params params(args, plan.device.sms, plan.occupancy); + if (params.get_workspace_size() > plan.kernel_workspace || params.init_workspace(scratch, stream) != cutlass::Status::kSuccess) + return GGML_STATUS_FAILED; + cutlass::arch::synclog_setup(); + cutlass::Kernel2<<>>(params); + return cudaGetLastError() == cudaSuccess ? GGML_STATUS_SUCCESS : GGML_STATUS_FAILED; + } + template + static bool configure(const Operation& op, Plan& plan) { + using Types = KernelTypes; + using Kernel = typename Types::Kernel; + if (sizeof(typename Kernel::SharedStorage) > plan.device.shared) + return false; + auto args = arguments(op, nullptr, true); + if (Kernel::can_implement(args) != cutlass::Status::kSuccess) + return false; + size_t workspace = 0; + // Bound every possible occupancy without setting a CUDA device or initializing a kernel during graph measurement. + for (int occupancy = 1; occupancy <= plan.device.threads_per_sm / Kernel::kThreadCount; ++occupancy) { + typename Kernel::Params params(args, plan.device.sms, occupancy); + const auto grid = params.get_grid_dims(); + if (grid.y > 65535 || grid.z > 65535) + return false; + workspace = std::max(workspace, params.get_workspace_size()); + } + plan.kernel_workspace = workspace; + plan.output_offset = (workspace + 255) / 256 * 256; + plan.launch = launch; + return true; + } + template + static bool configure(const Operation& op, Plan& plan) { + if (op.scale != 1.f) + return op.rows >= 8192 ? configure(op, plan) : configure(op, plan); + return op.rows >= 8192 ? configure(op, plan) : configure(op, plan); + } + static bool configure(const Operation& op, Plan& plan) { + auto try_tile = [&](Tile tile) { + switch (tile) { + case TILE_128_256: + return configure>(op, plan); + case TILE_128_128: + return configure>(op, plan); + case TILE_64_128: + return configure>(op, plan); + case TILE_16_128: + return configure>(op, plan); + } + return false; + }; + const auto tile = select_tile(op.rows, op.n, plan.device.sms); + if (try_tile(tile)) + return true; + for (auto fallback : {TILE_128_256, TILE_128_128, TILE_64_128, TILE_16_128}) + if (fallback != tile && try_tile(fallback)) + return true; + return false; + } + static bool fusion_disabled() { + static const bool disabled = [] { const char* v = std::getenv("GGML_CUDA_DISABLE_FUSION"); return v && std::atoi(v) != 0; }(); + return disabled; + } + static void optimize(ggml_backend_dev_t, ggml_cgraph* graph, void*) { +#ifndef SD_CUDA_FORCE_CUBLAS + if (fusion_disabled()) + return; + const int total = ggml_graph_n_nodes(graph); + for (int i = 2; i + 1 < total; ++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; + const auto* producer = mul->src[0]; + const bool scaled = producer->op == GGML_OP_SCALE; + const auto* mat = scaled ? producer->src[0] : producer; + SDI8Operation op; + if (!sd_ggml_i8_describe(mat, &op) || op.input->type != GGML_TYPE_I8 || !op.weight_scale || op.bias) + continue; + int scale_index = -1; + for (int j = i - 1; j >= 0 && j >= i - 64; --j) { + if (scaled && ggml_graph_node(graph, j) == producer) + scale_index = j; + if (ggml_graph_node(graph, j) != mat) + continue; + if (scaled) { + if (scale_index < 0 || (j + 1 < scale_index && !ggml_ext_graph_delay(graph, j, 1, scale_index))) + break; + j = scale_index - 1; + } + ggml_ext_graph_delay(graph, j, scaled ? 2 : 1, i); + break; + } + } +#endif + } + static int match(ggml_backend_dev_t device, const ggml_cgraph* graph, int first, void*) { +#ifdef SD_CUDA_FORCE_CUBLAS + return 0; +#else + const int available = ggml_graph_n_nodes(const_cast(graph)) - first; + if (fusion_disabled() || available < 3) + return 0; + ggml_tensor* nodes[5] = {}; + for (int i = 0; i < std::min(available, 5); ++i) + nodes[i] = ggml_graph_node(const_cast(graph), first + i); + SDI8Operation mat; + if (!sd_ggml_i8_describe(nodes[0], &mat) || mat.input->type != GGML_TYPE_I8 || !mat.weight_scale || mat.bias) + return 0; + const int64_t rows = ggml_nrows(nodes[0]), n = nodes[0]->ne[0], k = mat.weight->ne[0]; + if (rows < (small_rows_enabled(rows, n, k) ? 1 : 128) || rows > INT_MAX || n > INT_MAX || k > INT_MAX || n % 128 || k % 64 || + ((rows + 15) / 16) > INT_MAX / ((n + 127) / 128)) + return 0; + for (const auto* t : {mat.weight, mat.input, mat.weight_scale, mat.packed ? mat.weight_scale : mat.input_scale}) + if (!aligned(t)) + return 0; + const bool scaled = nodes[1]->op == GGML_OP_SCALE; + int count = scaled ? 4 : 3; + if (available < count || (scaled && (nodes[1]->src[0] != nodes[0] || param_f32(nodes[1], 1) != 0.f))) + return 0; + const auto* mul = nodes[count - 2]; + const auto* add = nodes[count - 1]; + if (mul->op != GGML_OP_MUL || add->op != GGML_OP_ADD || mul->src[0] != nodes[count - 3] || + (add->src[0] != mul && add->src[1] != mul)) + return 0; + const auto* gate = mul->src[1]; + const auto* residual = add->src[0] == mul ? add->src[1] : add->src[0]; + if (!gate || !residual || gate->ne[0] != n || ggml_nelements(gate) != n || !aligned(gate) || !aligned(residual)) + return 0; + for (const auto* t : {gate, residual, mul, add}) + if (t->type != GGML_TYPE_F32 || !ggml_is_contiguous(t) || (t != gate && !ggml_are_same_shape(nodes[0], t))) + return 0; + if (available > count) { + const auto* concat = nodes[count]; + if (concat->op == GGML_OP_CONCAT && concat->src[1] == add && concat->op_params[0] == 1 && + concat->type == GGML_TYPE_F32 && concat->ne[2] == 1 && concat->ne[3] == 1 && ggml_is_contiguous(concat) && + concat->src[0]->type == GGML_TYPE_F32 && ggml_is_contiguous(concat->src[0])) + ++count; + } + const auto* prefix = nodes[count - 1]->op == GGML_OP_CONCAT ? nodes[count - 1]->src[0] : nullptr; + for (int i = 0; i < count; ++i) + if (nodes[i] == gate || nodes[i] == residual || nodes[i] == prefix) + return 0; + Plan plan; + plan.device = device_info(device); + if (plan.device.major < 8 || plan.device.sms <= 0) + return 0; + return configure(describe(nodes, count), plan) ? count : 0; +#endif + } + static ggml_status prepare(ggml_backend_t backend, ggml_tensor* const* nodes, int count, void** state, size_t* workspace, void*) { + auto* plan = new Plan; + *state = plan; + plan->device = device_info(ggml_backend_get_device(backend)); + if (!configure(describe(nodes, count), *plan)) + return GGML_STATUS_FAILED; + *workspace = plan->output_offset + ggml_nbytes(nodes[count - 1]); + return GGML_STATUS_SUCCESS; + } + static void destroy(void* state, void*) { + delete static_cast(state); + } + 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 execute(const ggml_ext_execution_context* execution, ggml_tensor* const* nodes, int count, void* state, void*) { + auto& plan = *static_cast(state); + const auto op = describe(nodes, count); + auto* dst = nodes[count - 1]; + auto reg = ggml_backend_dev_backend_reg(ggml_backend_get_device(execution->backend)); + auto query = reinterpret_cast(ggml_backend_reg_get_proc_address(reg, "ggml_ext_cuda_context_v1")); + ggml_ext_cuda_context context; + if (!query || !query(execution->backend, &context) || context.device != plan.device.ordinal) + return GGML_STATUS_FAILED; + auto stream = static_cast(context.stream); + auto* scratch = static_cast(ggml_backend_buffer_get_base(execution->workspace)); + bool alias = reinterpret_cast(dst->data) % 16 != 0; + for (const auto* input : {op.mat.weight, op.mat.input, op.mat.input_scale, op.mat.weight_scale, op.gate, op.residual, op.prefix}) { + if (input == op.residual && !op.prefix && input->data == dst->data) + continue; + alias |= overlaps(dst, input); + } + auto* output = alias ? scratch + plan.output_offset : static_cast(dst->data); + const size_t prefix_bytes = op.prefix ? ggml_nbytes(op.prefix) : 0; + auto status = plan.launch(plan, op, reinterpret_cast(output + prefix_bytes), scratch, stream); + if (status != GGML_STATUS_SUCCESS) + return status; + if (prefix_bytes && cudaMemcpyAsync(output, op.prefix->data, prefix_bytes, cudaMemcpyDeviceToDevice, stream) != cudaSuccess) + return GGML_STATUS_FAILED; + if (alias && cudaMemcpyAsync(dst->data, output, ggml_nbytes(dst), cudaMemcpyDeviceToDevice, stream) != cudaSuccess) + return GGML_STATUS_FAILED; + return GGML_STATUS_SUCCESS; + } +} + +void sd_ggml_register_cuda_i8_gated_residual(ggml_backend_dev_t device) { + ggml_ext_kernel_desc desc = {}; + desc.abi_version = GGML_EXT_ABI_VERSION; + desc.name = "sd.i8_gated_residual.cuda.v1"; + desc.op = GGML_OP_MUL_MAT; + desc.priority = 325; + desc.optimize = sd_i8_gated::optimize; + desc.match = sd_i8_gated::match; + desc.prepare = sd_i8_gated::prepare; + desc.execute = sd_i8_gated::execute; + desc.destroy = sd_i8_gated::destroy; + GGML_ASSERT(ggml_ext_register_kernel(device, &desc)); + desc.name = "sd.custom_i8_gated_residual.cuda.v1"; + desc.op = GGML_OP_CUSTOM; + desc.custom_op = sd_ggml_i8_op(); + desc.optimize = nullptr; + GGML_ASSERT(ggml_ext_register_kernel(device, &desc)); +} diff --git a/src/core/ggml_extensions.cpp b/src/core/ggml_extensions.cpp index 8b2fe349..d48757d4 100644 --- a/src/core/ggml_extensions.cpp +++ b/src/core/ggml_extensions.cpp @@ -14,6 +14,7 @@ void sd_ggml_register_cuda_i8(ggml_backend_dev_t device); void sd_ggml_register_cuda_rope(ggml_backend_dev_t device); #endif #ifdef SD_CUDA_CUTE_ATTENTION +void sd_ggml_register_cuda_i8_gated_residual(ggml_backend_dev_t device); void sd_ggml_register_cuda_i8_projection(ggml_backend_dev_t device); void sd_ggml_register_cuda_attention(ggml_backend_dev_t device); #endif @@ -185,6 +186,7 @@ void sd_ggml_register_extensions(ggml_backend_t backend) { #endif #ifdef SD_CUDA_CUTE_ATTENTION sd_ggml_register_cuda_i8_projection(device); + sd_ggml_register_cuda_i8_gated_residual(device); sd_ggml_register_cuda_attention(device); #endif return;