diff --git a/CMakeLists.txt b/CMakeLists.txt index 7a927fc6..fe479b6b 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -348,12 +348,15 @@ if (SD_CUDA) target_sources(${SD_LIB} PRIVATE src/core/cuda/rope.cu) set_source_files_properties(src/core/cuda/rope.cu PROPERTIES COMPILE_OPTIONS "--use_fast_math") if (GGML_CUDA_CUTLASS_PATH) - target_sources(${SD_LIB} PRIVATE src/core/cuda/cute_attention.cu) - set_source_files_properties(src/core/cuda/cute_attention.cu PROPERTIES + 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 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 APPEND PROPERTY COMPILE_OPTIONS -Xcompiler=/Zc:preprocessor,/Zc:__cplusplus) + set_property(SOURCE src/core/cuda/cute_attention.cu src/core/cuda/i8_projection.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) 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 cc8c52c9..e1bae6fb 100644 --- a/docs/ggml_extensions.md +++ b/docs/ggml_extensions.md @@ -67,16 +67,20 @@ compute flags, external uses, output flags, and views; the initial implementatio accepts a conservative set of pure operators with only the final node exposed as the fused output. The rule remains responsible for mathematical equivalence, dtype, shape, layout, and hardware constraints. Internal views are accepted when -their storage belongs to the fused subgraph. The final output must own storage, -except for a terminal `CPY`, which retains its original destination and write +their storage belongs to the fused subgraph. In-place intermediates are accepted +only when their storage owner is an earlier node in the same match. External +consumers and output flags still prevent eliminating those results. The final +output must own storage, except for a terminal `CPY`, which retains its original destination and write semantics. The extension must handle input/output storage overlap, including allocator reuse, with scratch space when necessary. Other in-place writes are -not eligible for extension fusion. +not eligible for extension fusion when they target storage outside the match. An optional `optimize` callback runs before matching and allocation. Use `ggml_ext_graph_delay` to move a producer range after independent nodes. The helper rejects dependencies, unknown side effects, and writes that alias moved -inputs or outputs. It can cross independent copies, including dtype conversions. +inputs or outputs. It can cross independent copies and known in-place operators +when their storage does not alias the moved inputs or outputs. Owning dtype +conversions can also be moved or crossed. Reordering must preserve edges and tensor metadata. Fusion legality is checked again after reordering. @@ -140,12 +144,21 @@ CUDA stream, and scheduler-owned workspace; they do not use ggml's CUDA pool or private conversion helpers. The existing CUDA fusion and CuTe disable environment variables are honored. -The combined INT8 GEMM epilogue feeding FP16 Q directly into attention still needs -to be externalized. The current extension can fuse RMSNorm/RoPE with attention, -while other INT8 GEMM and standalone projection fusions retain their native -implementations. Native INT8 projection matchers do not recognize custom RoPE -yet, so those graphs currently run GEMM followed by the external RoPE/attention -path. Restoring their combined epilogues is part of the remaining migration. -Vulkan execution services and kernels also remain pending. +`src/core/cuda/i8_projection.cu` supplies the INT8 GEMM projection epilogues: +`GEMM -> Reshape -> RMSNorm -> Mul -> RoPE`, optional scaling and FP16 copy, +and K/V prefix concatenation with FP16 output. CuTe can consume a directly +produced FP16 Q without materializing the FP32 GEMM or RoPE intermediates. +Independent affine operands, including broadcasted constants, are computed before +their normalization/projection +chain when graph dependencies allow it. The attention optimizer can then move +the complete five-node projection when safe. +These rules accept both builtin and custom RoPE nodes. They preserve the original +large-matrix selection thresholds, full interleaved head dimension 128, packed +activation layout, and separate weight scales. External consumers prevent fusion; +aliasing terminal copies use scheduler workspace. The CUDA projection kernels +and matchers now live in sd.cpp. `GGML_CUDA_FORCE_CUBLAS` disables these epilogues. + +Ordinary INT8 GEMM, ConvRot quantization, and their other 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 7ffa0f0f..1223c158 160000 --- a/ggml +++ b/ggml @@ -1 +1 @@ -Subproject commit 7ffa0f0fe1cdcf6a60eeba1d92f01d53ccd19a55 +Subproject commit 1223c158c84475a5431b34bf148e3223c8364b4a diff --git a/src/core/cuda/cute_attention.cu b/src/core/cuda/cute_attention.cu index 53ed961e..15f210a8 100644 --- a/src/core/cuda/cute_attention.cu +++ b/src/core/cuda/cute_attention.cu @@ -44,6 +44,7 @@ #include #include #include +#include "core/cuda/i8_projection.h" #include "core/cuda/rope.h" #include "ggml-ext.h" @@ -578,7 +579,14 @@ static int attention_match(ggml_backend_dev_t device, const ggml_cgraph* graph, if (!fusion_enabled()) return 0; int offset = 0; - if (graph_node(graph, first)->op == GGML_OP_RMS_NORM) { + if (graph_node(graph, first)->op == GGML_OP_MUL_MAT) { + if (first + 5 >= ggml_graph_n_nodes(const_cast(graph)) || + !sd_cuda_i8_rope_supports(device, graph_node(graph, first), graph_node(graph, first + 1), + graph_node(graph, first + 2), graph_node(graph, first + 3), graph_node(graph, first + 4)) || + graph_node(graph, first + 5)->src[0] != graph_node(graph, first + 4)) + return 0; + offset = 5; + } else if (graph_node(graph, first)->op == GGML_OP_RMS_NORM) { if (first + 3 >= ggml_graph_n_nodes(const_cast(graph)) || !sd_cuda_rope_fusion_supports(graph_node(graph, first), graph_node(graph, first + 1), graph_node(graph, first + 2)) || graph_node(graph, first + 3)->src[0] != graph_node(graph, first + 2)) @@ -601,6 +609,9 @@ static void attention_optimize(ggml_backend_dev_t device, ggml_cgraph* graph, vo for (int j = i - 3; j >= 0 && j >= i - 32; --j) { if (attn->src[0] == graph_node(graph, j + 2) && sd_cuda_rope_fusion_supports(graph_node(graph, j), graph_node(graph, j + 1), graph_node(graph, j + 2))) { + if (j >= 2 && sd_cuda_i8_rope_supports(device, graph_node(graph, j - 2), graph_node(graph, j - 1), graph_node(graph, j), graph_node(graph, j + 1), graph_node(graph, j + 2)) && + ggml_ext_graph_delay(graph, j - 2, 5, i)) + break; ggml_ext_graph_delay(graph, j, 3, i); break; } @@ -616,7 +627,8 @@ struct AttentionPlan { static ggml_status attention_prepare(ggml_backend_t, ggml_tensor* const* nodes, int count, void** state, size_t* workspace_size, void*) { auto* plan = new AttentionPlan{}; *state = plan; - plan->index = nodes[0]->op == GGML_OP_RMS_NORM ? 3 : 0; + plan->index = nodes[0]->op == GGML_OP_MUL_MAT ? 5 : nodes[0]->op == GGML_OP_RMS_NORM ? 3 + : 0; GGML_ASSERT(plan->index < count); const auto* attn = nodes[plan->index]; size_t size = 0; @@ -671,7 +683,10 @@ static ggml_status attention_execute(const ggml_ext_execution_context* execution for (int i = 0; i < 3; ++i) { const auto* src = attn->src[i]; const float scale = i == 0 ? param_f32(attn, 0) : 1.0f; - if (i == 0 && plan.index == 3) { + if (i == 0 && plan.index == 5) { + if (sd_cuda_i8_rope_launch(stream, nodes, work[0], scale) != cudaSuccess) + return GGML_STATUS_FAILED; + } else if (i == 0 && plan.index == 3) { if (sd_cuda_rms_rope_launch(stream, nodes, work[0], true, scale) != cudaSuccess) return GGML_STATUS_FAILED; } else if (i != 0 && src->type == GGML_TYPE_F16 && ggml_is_contiguous(src) && uintptr_t(src->data) % 16 == 0) { @@ -766,4 +781,8 @@ void sd_ggml_register_cuda_attention(ggml_backend_dev_t device) { desc.preferred = nullptr; desc.optimize = nullptr; GGML_ASSERT(ggml_ext_register_kernel(device, &desc)); + desc.name = "sd.i8_rope_cute_attention.cuda.v1"; + desc.op = GGML_OP_MUL_MAT; + desc.priority = 400; + GGML_ASSERT(ggml_ext_register_kernel(device, &desc)); } diff --git a/src/core/cuda/i8_projection.cu b/src/core/cuda/i8_projection.cu new file mode 100644 index 00000000..819ea751 --- /dev/null +++ b/src/core/cuda/i8_projection.cu @@ -0,0 +1,389 @@ +#include "core/cuda/i8_projection.h" +#include "core/cuda/rope.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +static __device__ __forceinline__ float warp_sum8(float value) { +#pragma unroll + for (int offset = 4; offset > 0; offset >>= 1) + value = __fadd_rn(value, __shfl_xor_sync(0xffffffff, value, offset, 8)); + return value; +} + +template +struct I8RopeStore : cutlass::epilogue::threadblock::VisitorAuxStore> { + using Base = cutlass::epilogue::threadblock::VisitorAuxStore>; + struct Arguments : Base::Arguments { + const float* weight; + const float* pe; + float eps; + float scale; + int64_t tokens; + half* half_output; + int64_t prefix_tokens; + }; + using Params = Arguments; + using SharedStorage = typename Base::SharedStorage; + const Params* extra; + template + static constexpr Params to_underlying_arguments(Shape const&, Arguments const& args, void*) { return args; } + CUTLASS_HOST_DEVICE I8RopeStore(Params const& p, SharedStorage const& s) + : Base(p, s), extra(&p) {} + template + struct Callbacks : CB { + const Params* extra; + CUTLASS_DEVICE Callbacks(CB const& base, const Params* p) + : CB(base), extra(p) {} + CUTLASS_DEVICE float sum32(float4 v) { + float a = warp_sum8(__fmul_rn(v.x, v.x)); + float b = warp_sum8(__fmul_rn(v.y, v.y)); + float c = warp_sum8(__fmul_rn(v.z, v.z)); + float d = warp_sum8(__fmul_rn(v.w, v.w)); + return __fadd_rn(__fadd_rn(a, c), __fadd_rn(b, d)); + } + CUTLASS_DEVICE void end_step(int step) { + static_assert(Map::Detail::kAccessWidth == 16); + auto src = cute::filter(this->tC_rAux); + auto coords = cute::filter(this->tC_cAux(cute::_, cute::_, cute::_, step)); + static_assert(cute::size(src) == 4); + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < cute::size(src); i += 2) { + float4 va = *reinterpret_cast(&src(i)), vb = *reinterpret_cast(&src(i + 1)); + float norm = 1.0f; + if constexpr (rms) { + // Preserve the standalone RMSNorm reduction tree across four 32-element groups. + float sum = __fadd_rn(sum32(va), sum32(vb)); + sum = __fadd_rn(sum, __shfl_xor_sync(0xffffffff, sum, 8)); + norm = rsqrtf(sum / 128.f + extra->eps); + } + CUTLASS_PRAGMA_UNROLL + for (int j = 0; j < 2; ++j) { + auto coord = coords(i + j); + int row = cute::get<0>(coord), col = cute::get<1>(coord); + int heads = cute::get<1>(this->problem_shape) / 128; + if (cute::elem_less(coord, this->problem_shape)) { + float4 value = j ? vb : va; + int local = col % 128; + int64_t token = row % extra->tokens, batch = row / extra->tokens, head = col / 128; + float4 r = value; + if constexpr (rms) { + float4 weight = ((const float4*)extra->weight)[local / 4]; + value.x = __fmul_rn(__fmul_rn(norm, value.x), weight.x); + value.y = __fmul_rn(__fmul_rn(norm, value.y), weight.y); + value.z = __fmul_rn(__fmul_rn(norm, value.z), weight.z); + value.w = __fmul_rn(__fmul_rn(norm, value.w), weight.w); + float4 p0 = ((const float4*)extra->pe)[token * 64 + local / 2], p1 = ((const float4*)extra->pe)[token * 64 + local / 2 + 1]; + r.x = __fadd_rn(__fmul_rn(value.x, p0.x), __fmul_rn(value.y, p0.y)); + r.y = __fadd_rn(__fmul_rn(value.x, p0.z), __fmul_rn(value.y, p0.w)); + r.z = __fadd_rn(__fmul_rn(value.z, p1.x), __fmul_rn(value.w, p1.y)); + r.w = __fadd_rn(__fmul_rn(value.z, p1.z), __fmul_rn(value.w, p1.w)); + } + int64_t offset = ((batch * heads + head) * (extra->tokens + extra->prefix_tokens) + token + extra->prefix_tokens) * 128 + local; + if constexpr (std::is_same::value) { + half2 a = __floats2half2_rn(__fmul_rn(r.x, extra->scale), __fmul_rn(r.y, extra->scale)); + half2 b = __floats2half2_rn(__fmul_rn(r.z, extra->scale), __fmul_rn(r.w, extra->scale)); + half2* dst = (half2*)(extra->half_output + offset); + dst[0] = a; + dst[1] = b; + } else { + *((float4*)(this->params_ptr->ptr_aux + offset)) = r; + } + } + } + } + } + }; + template + CUTLASS_DEVICE auto get_callbacks(cutlass::gemm::GemmCoord tile, int tid, Shape shape) { + auto cb = Base::get_callbacks(tile, tid, shape); + return Callbacks(cb, extra); + } +}; + +template +struct ProjectionKernel { + using Tile = cutlass::gemm::GemmShape<128, 256, 64>; + using Warp = cutlass::gemm::GemmShape<64, 64, 64>; + using Swizzle = cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>; + 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 Store = I8RopeStore; + using EVT = 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>, + EVT, + Swizzle, + 4, + cutlass::arch::OpMultiplyAddSaturate>::GemmKernel; + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; +}; + +static float param_f32(const ggml_tensor* tensor, int index) { + float result; + std::memcpy(&result, tensor->op_params + index, sizeof(result)); + return result; +} + +static bool device_supports(ggml_backend_dev_t device) { +#ifdef SD_CUDA_FORCE_CUBLAS + (void)device; + return false; +#else + 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; + 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 = {}; + size_t shared = std::max({sizeof(typename ProjectionKernel::Kernel::SharedStorage), + sizeof(typename ProjectionKernel::Kernel::SharedStorage), sizeof(typename ProjectionKernel::Kernel::SharedStorage)}); + bool supported = ordinal >= 0 && cudaGetDeviceProperties(&props, ordinal) == cudaSuccess && props.major >= 8 && props.sharedMemPerBlockOptin >= shared; + cache.emplace(device, supported); + return supported; +#endif +} + +static bool aligned(const ggml_tensor* tensor) { + size_t offset = 0; + while (tensor->view_src) { + offset += tensor->view_offs; + tensor = tensor->view_src; + } + return (reinterpret_cast(tensor->data) + offset) % 16 == 0; +} + +static bool projection_supports(ggml_backend_dev_t device, const ggml_tensor* mat, const ggml_tensor* view) { + const auto* w = mat->src[0]; + const auto* x = mat->src[1]; + const auto* ws = mat->src[2]; + if (mat->op != GGML_OP_MUL_MAT || mat->type != GGML_TYPE_F32 || mat->op_params[4] <= 0 || !w || !x || !ws || mat->src[3] || + w->type != GGML_TYPE_I8 || x->type != GGML_TYPE_I8 || ws->type != GGML_TYPE_F32 || + w->ne[2] != 1 || w->ne[3] != 1 || mat->ne[0] != w->ne[1] || ggml_nelements(ws) != mat->ne[0] || + !ggml_is_contiguous(mat) || !ggml_is_contiguous(w) || !ggml_is_contiguous(x) || !ggml_is_contiguous(ws) || + !aligned(w) || !aligned(x) || !aligned(ws) || !aligned(mat) || + view->op != GGML_OP_RESHAPE || view->src[0] != mat || view->view_offs != 0 || !ggml_is_contiguous(view) || + view->ne[0] != 128 || view->ne[0] * view->ne[1] != mat->ne[0] || ggml_nelements(view) != ggml_nelements(mat)) + return false; + const int64_t rows = ggml_nrows(mat), n = mat->ne[0], k = w->ne[0]; + if (rows < 4096 || rows > INT_MAX || n < 4096 || n > INT_MAX || k < 4096 || k > INT_MAX || n % 256 || k % 64 || n / 256 > 65535 || + x->ne[0] != k || x->ne[1] != (rows + 3) / 4 * 4 + (rows * 4 + k - 1) / k || x->ne[2] != 1 || x->ne[3] != 1) + return false; + return device_supports(device); +} + +bool sd_cuda_i8_rope_supports(ggml_backend_dev_t device, const ggml_tensor* mat, const ggml_tensor* view, const ggml_tensor* norm, const ggml_tensor* mul, const ggml_tensor* rope) { + if (!projection_supports(device, mat, view) || norm->src[0] != view || + !sd_cuda_rope_fusion_supports(norm, mul, rope) || rope->src[1]->ne[2] != 64 || !aligned(mul->src[1]) || !aligned(rope->src[1])) + return false; + int32_t interleaved; + std::memcpy(&interleaved, ggml_ext_op_get(rope) ? ggml_ext_op_attributes(rope) : rope->op_params, sizeof(interleaved)); + return interleaved != 0; +} + +template +static cudaError_t projection_launch(cudaStream_t stream, const ggml_tensor* mat, const ggml_tensor* view, const ggml_tensor* norm, const ggml_tensor* mul, const ggml_tensor* rope, void* output, float scale, int64_t prefix_tokens) { + using P = ProjectionKernel; + using Gemm = typename P::Gemm; + const int rows = int(ggml_nrows(mat)), n = int(mat->ne[0]), k = int(mat->src[0]->ne[0]); + const auto* x = static_cast(mat->src[1]->data); + const auto* w = static_cast(mat->src[0]->data); + const auto* xs = reinterpret_cast(x + int64_t(k) * ((int64_t(rows) + 3) / 4 * 4)); + const auto* ws = static_cast(mat->src[2]->data); + const auto* nw = RMS ? static_cast(mul->src[1]->data) : nullptr; + const auto* pe = RMS ? static_cast(rope->src[1]->data) : nullptr; + // The visitor forms FP32 addresses even when its only writes target the FP16 output. + float* address_space = std::is_same::value ? static_cast(mat->data) : static_cast(output); + if ((reinterpret_cast(x) | reinterpret_cast(w) | reinterpret_cast(xs) | + reinterpret_cast(ws) | reinterpret_cast(nw) | reinterpret_cast(pe) | + reinterpret_cast(address_space)) % + 16 || + reinterpret_cast(output) % 8) + return cudaErrorInvalidValue; + typename P::EVT::Arguments evt = {{{{}, {xs, 0.f, {}}, {}}, {ws, 0.f, {}}, {}}, + {{address_space, {n, cute::_1{}, int64_t(rows) * n}}, nw, pe, RMS ? param_f32(norm, 0) : 0.f, scale, view->ne[2], static_cast(output), prefix_tokens}}; + typename Gemm::Arguments args(cutlass::gemm::GemmUniversalMode::kGemm, {rows, n, k}, 1, evt, x, w, nullptr, nullptr, + int64_t(rows) * k, int64_t(n) * k, 0, 0, k, k, 0, 0); + Gemm gemm; + if (gemm.can_implement(args) != cutlass::Status::kSuccess) + return cudaErrorInvalidValue; + if (gemm.initialize(args, nullptr, stream) != cutlass::Status::kSuccess || gemm(stream) != cutlass::Status::kSuccess) + return cudaErrorUnknown; + return cudaGetLastError(); +} + +cudaError_t sd_cuda_i8_rope_launch(cudaStream_t stream, ggml_tensor* const* nodes, void* output, float scale) { + return projection_launch(stream, nodes[0], nodes[1], nodes[2], nodes[3], nodes[4], output, scale, 0); +} + +static const ggml_tensor* node_at(const ggml_cgraph* graph, int index) { + return ggml_graph_node(const_cast(graph), index); +} + +static int concat_match(const ggml_cgraph* graph, int first, bool key) { + const int offset = key ? 5 : 2, count = key ? 7 : 5; + if (first + count > ggml_graph_n_nodes(const_cast(graph))) + return 0; + const auto* concat = node_at(graph, first + offset); + const auto* copy = node_at(graph, first + count - 1); + if (concat->op != GGML_OP_CONCAT || concat->op_params[0] != (key ? 1 : 2) || concat->src[1] != node_at(graph, first + offset - 1) || + copy->op != GGML_OP_CPY || copy->src[0] != node_at(graph, first + count - 2) || copy->type != GGML_TYPE_F16 || + !ggml_is_contiguous(copy) || !ggml_are_same_shape(copy, copy->src[0])) + return 0; + for (int i = 0; i < 2; ++i) + if (!concat->src[i] || concat->src[i]->type != GGML_TYPE_F32 || !ggml_is_contiguous(concat->src[i])) + return 0; + const auto* view = node_at(graph, first + 1); + const auto* prefix = concat->src[0]; + if (prefix->ne[key ? 1 : 2] > INT_MAX - view->ne[2] || (ggml_nelements(prefix) + 255) / 256 > INT_MAX) + return 0; + if (!key) { + const auto* permute = node_at(graph, first + 3); + if (permute->op != GGML_OP_PERMUTE || permute->src[0] != concat || permute->op_params[0] != 0 || permute->op_params[1] != 2 || + permute->op_params[2] != 1 || permute->op_params[3] != 3) + return 0; + } + return count; +} + +static int projection_match(ggml_backend_dev_t device, const ggml_cgraph* graph, int first, void*) { + static const bool disabled = [] {const char* value=std::getenv("GGML_CUDA_DISABLE_FUSION");return value&&std::atoi(value)!=0; }(); + if (disabled || first + 4 >= ggml_graph_n_nodes(const_cast(graph))) + return 0; + const auto* mat = node_at(graph, first); + const auto* view = node_at(graph, first + 1); + const auto* norm = node_at(graph, first + 2); + const bool key = norm->op == GGML_OP_RMS_NORM; + if (key) { + if (!sd_cuda_i8_rope_supports(device, mat, view, norm, node_at(graph, first + 3), node_at(graph, first + 4))) + return 0; + if (int count = concat_match(graph, first, true)) + return count; + if (first + 6 < ggml_graph_n_nodes(const_cast(graph))) { + const auto* rope = node_at(graph, first + 4); + const auto* scale = node_at(graph, first + 5); + const auto* copy = node_at(graph, first + 6); + if (scale->op == GGML_OP_SCALE && scale->src[0] == rope && param_f32(scale, 1) == 0.f && copy->op == GGML_OP_CPY && copy->src[0] == scale && + copy->type == GGML_TYPE_F16 && ggml_are_same_shape(rope, scale) && ggml_are_same_shape(rope, copy) && ggml_is_contiguous(copy)) + return 7; + } + return 5; + } + return projection_supports(device, mat, view) ? concat_match(graph, first, false) : 0; +} + +template +static __global__ void prefix_copy(const float* input, half2* output, int64_t count4, int64_t prefix, int64_t tokens, int64_t heads, bool permuted) { + const int64_t i = int64_t(blockIdx.x) * blockDim.x + threadIdx.x; + if (i >= count4) + return; + const int64_t col = i % 32, row = i / 32; + const int64_t token = row % prefix, head = (row / prefix) % heads, batch = row / (prefix * heads); + const int64_t source_row = permuted ? (batch * prefix + token) * heads + head : (batch * heads + head) * prefix + token; + const float* source = input + source_row * 128 + col * 4; + float4 value; + if constexpr (Aligned) + value = *reinterpret_cast(source); + else + value = make_float4(source[0], source[1], source[2], source[3]); + const int64_t target = ((batch * heads + head) * (prefix + tokens) + token) * 32 + col; + output[2 * target] = __floats2half2_rn(value.x, value.y); + output[2 * target + 1] = __floats2half2_rn(value.z, value.w); +} + +static ggml_status projection_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 projection_destroy(void*, void*) {} + +static bool overlaps(const ggml_tensor* a, const ggml_tensor* b) { + 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 projection_execute(const ggml_ext_execution_context* execution, ggml_tensor* const* nodes, int count, void*, void*) { + const auto* mat = nodes[0]; + const auto* view = nodes[1]; + const bool key = nodes[2]->op == GGML_OP_RMS_NORM; + const auto* concat = nodes[key ? count == 7 ? 5 : 4 : 2]; + const bool joined = concat->op == GGML_OP_CONCAT; + 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)) + return GGML_STATUS_FAILED; + auto stream = static_cast(context.stream); + bool alias = reinterpret_cast(dst->data) % (dst->type == GGML_TYPE_F32 ? 16 : 8) != 0; + for (const auto* input : {mat->src[0], mat->src[1], mat->src[2], key ? nodes[3]->src[1] : nullptr, key ? nodes[4]->src[1] : nullptr, joined ? concat->src[0] : nullptr}) + if (input) + alias |= overlaps(dst, input); + void* output = alias ? ggml_backend_buffer_get_base(execution->workspace) : dst->data; + int64_t prefix = joined ? concat->src[0]->ne[key ? 1 : 2] : 0; + float scale = key && count == 7 && !joined ? param_f32(nodes[5], 0) : 1.f; + cudaError_t error; + if (dst->type == GGML_TYPE_F32) + error = projection_launch(stream, mat, view, nodes[2], nodes[3], nodes[4], output, scale, 0); + else if (key) + error = projection_launch(stream, mat, view, nodes[2], nodes[3], nodes[4], output, scale, prefix); + else + error = projection_launch(stream, mat, view, nullptr, nullptr, nullptr, output, scale, prefix); + if (error == cudaSuccess && joined && prefix > 0) { + const auto* input = concat->src[0]; + const int64_t count4 = ggml_nelements(input) / 4; + auto kernel = reinterpret_cast(input->data) % 16 == 0 ? prefix_copy : prefix_copy; + kernel<<((count4 + 255) / 256), 256, 0, stream>>>(static_cast(input->data), static_cast(output), count4, prefix, view->ne[2], view->ne[1], !key); + error = cudaGetLastError(); + } + if (error == cudaSuccess && alias) + error = cudaMemcpyAsync(dst->data, output, ggml_nbytes(dst), cudaMemcpyDeviceToDevice, stream); + return error == cudaSuccess ? GGML_STATUS_SUCCESS : GGML_STATUS_FAILED; +} + +void sd_ggml_register_cuda_i8_projection(ggml_backend_dev_t device) { + ggml_ext_kernel_desc desc = {}; + desc.abi_version = GGML_EXT_ABI_VERSION; + desc.name = "sd.i8_projection.cuda.v1"; + desc.op = GGML_OP_MUL_MAT; + desc.priority = 350; + desc.match = projection_match; + desc.prepare = projection_prepare; + desc.execute = projection_execute; + desc.destroy = projection_destroy; + GGML_ASSERT(ggml_ext_register_kernel(device, &desc)); +} diff --git a/src/core/cuda/i8_projection.h b/src/core/cuda/i8_projection.h new file mode 100644 index 00000000..fb815d08 --- /dev/null +++ b/src/core/cuda/i8_projection.h @@ -0,0 +1,11 @@ +#ifndef __SD_CORE_CUDA_I8_PROJECTION_H__ +#define __SD_CORE_CUDA_I8_PROJECTION_H__ + +#include +#include "ggml-ext.h" + +bool sd_cuda_i8_rope_supports(ggml_backend_dev_t device, const ggml_tensor* mat, const ggml_tensor* view, const ggml_tensor* norm, const ggml_tensor* mul, const ggml_tensor* rope); +cudaError_t sd_cuda_i8_rope_launch(cudaStream_t stream, ggml_tensor* const* nodes, void* output, float scale); +void sd_ggml_register_cuda_i8_projection(ggml_backend_dev_t device); + +#endif diff --git a/src/core/cuda/rope.cu b/src/core/cuda/rope.cu index f42b9c82..2217f19f 100644 --- a/src/core/cuda/rope.cu +++ b/src/core/cuda/rope.cu @@ -235,6 +235,40 @@ static int rope_match(ggml_backend_dev_t, const ggml_cgraph* graph, int first, v return 3; } +static void rope_optimize(ggml_backend_dev_t, ggml_cgraph* graph, void*) { + static const bool disabled = [] { + const char* value = std::getenv("GGML_CUDA_DISABLE_FUSION"); + return value && std::atoi(value) != 0; + }(); + if (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* rope = ggml_graph_node(graph, i + 1); + if (mul->op != GGML_OP_MUL || !sd_cuda_rope_is_node(rope) || rope->src[0] != mul || + !mul->src[0] || mul->src[0]->op != GGML_OP_RMS_NORM) + continue; + const auto* norm = mul->src[0]; + for (int j = i - 1; j >= 0 && j >= i - 64; --j) { + if (ggml_graph_node(graph, j) != norm) + continue; + if (j + 1 < i) { + if (j >= 2) { + const auto* mat = ggml_graph_node(graph, j - 2); + const auto* view = ggml_graph_node(graph, j - 1); + if (mat->op == GGML_OP_MUL_MAT && mat->src[0]->type == GGML_TYPE_I8 && + view->op == GGML_OP_RESHAPE && view->src[0] == mat && norm->src[0] == view && + ggml_ext_graph_delay(graph, j - 2, 3, i)) + break; + } + ggml_ext_graph_delay(graph, j, 1, i); + } + break; + } + } +} + static ggml_status rope_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; @@ -298,8 +332,9 @@ void sd_ggml_register_cuda_rope(ggml_backend_dev_t device) { GGML_ASSERT(ggml_ext_register_kernel(device, &desc)); desc.name = "sd.rms_norm_mul_rope.cuda.v1"; desc.op = GGML_OP_RMS_NORM; - desc.priority = 200; + desc.priority = 250; desc.supports = nullptr; desc.match = rope_match; + desc.optimize = rope_optimize; GGML_ASSERT(ggml_ext_register_kernel(device, &desc)); } diff --git a/src/core/ggml_extensions.cpp b/src/core/ggml_extensions.cpp index b3c2ec23..a5dbffff 100644 --- a/src/core/ggml_extensions.cpp +++ b/src/core/ggml_extensions.cpp @@ -12,6 +12,7 @@ void sd_ggml_register_cuda_rope(ggml_backend_dev_t device); #endif #ifdef SD_CUDA_CUTE_ATTENTION +void sd_ggml_register_cuda_i8_projection(ggml_backend_dev_t device); void sd_ggml_register_cuda_attention(ggml_backend_dev_t device); #endif @@ -180,6 +181,7 @@ void sd_ggml_register_extensions(ggml_backend_t backend) { sd_ggml_register_cuda_rope(device); #endif #ifdef SD_CUDA_CUTE_ATTENTION + sd_ggml_register_cuda_i8_projection(device); sd_ggml_register_cuda_attention(device); #endif return;