diff --git a/CMakeLists.txt b/CMakeLists.txt index 123307d5..b2c13a1e 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -345,8 +345,7 @@ if (SD_CUDA) set_property(SOURCE src/core/ggml_extend_backend.cpp APPEND PROPERTY COMPILE_DEFINITIONS SD_USE_CUDA) if (NOT SD_USE_UPSTREAM_GGML) enable_language(CUDA) - target_sources(${SD_LIB} PRIVATE src/core/cuda/rope.cu src/core/cuda/int8.cu src/core/cuda/convrot.cu src/core/cuda/affine.cu) - set_source_files_properties(src/core/cuda/rope.cu src/core/cuda/int8.cu src/core/cuda/convrot.cu src/core/cuda/affine.cu PROPERTIES COMPILE_OPTIONS "--use_fast_math") + set_target_properties(${SD_LIB} PROPERTIES CUDA_STANDARD 17 CUDA_STANDARD_REQUIRED ON) if (GGML_CUDA_CUTLASS_PATH) target_sources(${SD_LIB} PRIVATE src/core/cuda/cute_attention.cu src/core/cuda/i8_projection.cu src/core/cuda/i8_gated_residual.cu src/core/cuda/i8_gemm.cu) set_source_files_properties(src/core/cuda/cute_attention.cu src/core/cuda/i8_projection.cu src/core/cuda/i8_gated_residual.cu src/core/cuda/i8_gemm.cu PROPERTIES @@ -358,17 +357,28 @@ if (SD_CUDA) if (GGML_CUDA_FORCE_CUBLAS) set_property(SOURCE src/core/cuda/i8_projection.cu src/core/cuda/i8_gated_residual.cu src/core/cuda/i8_gemm.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) - set_property(SOURCE src/core/cuda/int8.cu APPEND PROPERTY COMPILE_DEFINITIONS SD_CUDA_I8_CUTLASS) endif() - set_target_properties(${SD_LIB} PROPERTIES CUDA_STANDARD 17 CUDA_STANDARD_REQUIRED ON) - target_link_libraries(${SD_LIB} PRIVATE CUDA::cudart CUDA::cublas) - set_property(SOURCE src/core/ggml_extensions.cpp APPEND PROPERTY COMPILE_DEFINITIONS SD_CUDA_EXTENSIONS) endif() endif() +set(SD_BACKEND_REGISTRATIONS) +if (NOT SD_USE_UPSTREAM_GGML AND (SD_CUDA OR SD_HIPBLAS OR SD_MUSA)) + add_subdirectory(src/core/cuda) + foreach(VENDOR CUDA HIP MUSA) + if(SD_${VENDOR} OR (VENDOR STREQUAL "HIP" AND SD_HIPBLAS)) + string(TOLOWER "${VENDOR}" NAMESPACE) + list(APPEND SD_BACKEND_REGISTRATIONS "sd_${NAMESPACE}::register_extensions") + endif() + endforeach() +endif() if (SD_VULKAN AND NOT SD_USE_UPSTREAM_GGML) add_subdirectory(src/core/vulkan) - set_property(SOURCE src/core/ggml_extensions.cpp APPEND PROPERTY COMPILE_DEFINITIONS SD_VULKAN_EXTENSIONS) + list(APPEND SD_BACKEND_REGISTRATIONS sd_ggml_register_vulkan_rope) +endif() +if(SD_BACKEND_REGISTRATIONS) + string(REPLACE ";" ",\n " SD_BACKEND_REGISTRATIONS "${SD_BACKEND_REGISTRATIONS}") + configure_file(cmake/backend_extensions.h.in generated/sd_backend_extensions.h @ONLY) + set_property(SOURCE src/core/ggml_extensions.cpp APPEND PROPERTY INCLUDE_DIRECTORIES "${CMAKE_CURRENT_BINARY_DIR}/generated") + set_property(SOURCE src/core/ggml_extensions.cpp APPEND PROPERTY COMPILE_DEFINITIONS SD_BACKEND_EXTENSIONS) endif() if (SD_HIPBLAS) find_package(hip REQUIRED CONFIG HINTS "$ENV{ROCM_PATH}" /opt/rocm) diff --git a/cmake/backend_extensions.h.in b/cmake/backend_extensions.h.in new file mode 100644 index 00000000..ec28c6bc --- /dev/null +++ b/cmake/backend_extensions.h.in @@ -0,0 +1,12 @@ +#ifndef __SD_BACKEND_EXTENSIONS_H__ +#define __SD_BACKEND_EXTENSIONS_H__ + +#include "core/cuda/registration.h" + +void sd_ggml_register_vulkan_rope(ggml_backend_dev_t device); + +static void (*const sd_backend_extensions[])(ggml_backend_dev_t) = { + @SD_BACKEND_REGISTRATIONS@ +}; + +#endif diff --git a/docs/ggml_extensions.md b/docs/ggml_extensions.md index 1ff829ea..aeece307 100644 --- a/docs/ggml_extensions.md +++ b/docs/ggml_extensions.md @@ -112,14 +112,20 @@ Services are discovered through `ggml_backend_reg_get_proc_address`: | `ggml_ext_cuda_context_v1` | Borrow the current CUDA device and stream, ordered with native work. | | `ggml_ext_cuda_device_v1` | Get the physical CUDA ordinal for capability checks before backend creation, including virtual ggml devices. | | `ggml_ext_cuda_blas_v1` | Borrow the current stream's cuBLAS handle during execution, with host scalar pointers. | +| `ggml_ext_gpu_device_v1` | Query CUDA/HIP/MUSA runtime kind, physical ordinal, shared-memory limit, warp size, and INT8 BLAS capability. | +| `ggml_ext_gpu_context_v1` | Borrow the matching runtime's device and execution stream. | +| `ggml_ext_gpu_blas_v1` | Borrow the matching cuBLAS, hipBLAS, or muBLAS handle on that stream. | The borrowed cuBLAS handle remains owned by the backend. Extensions must not retain or destroy it, or change its stream, pointer mode, math mode, or workspace. This -service is exposed by the CUDA backend, not HIP/MUSA. +CUDA-specific BLAS service is retained for existing callers; the runtime-aware +services support CUDA, HIP, and MUSA with the same ownership rules. Extensions +call vendor runtime and BLAS APIs directly. These services do not execute GEMM, +quantization, or model fusion on the application's behalf. Workspace is a backend buffer. Its base is a CUDA device pointer on CUDA; callers -must use the representation appropriate to their backend. A Vulkan command -recording service is not implemented yet. +must use the representation appropriate to their backend. Vulkan extensions use +the command recording services described below. ## sd.cpp integration @@ -127,8 +133,8 @@ recording service is not implemented yet. `sd.rope_apply.v1`, an override for existing `GGML_OP_ROPE_APPLY` nodes, and a `RMSNorm -> Mul -> RoPE` fusion. The fusion accepts a one-dimensional weight vector. Both interleaved and split-half RoPE, partial rotation, and strided input -rows are supported. Model construction uses the custom operator on CPU and CUDA -and retains existing backend implementations elsewhere. +rows are supported. Model construction uses the custom operator on CPU, CUDA, +HIP, MUSA, and Vulkan, and retains existing backend implementations elsewhere. `src/core/cuda/rope.cu` supplies standalone CUDA RoPE, the same builtin override, and `RMSNorm -> Mul -> RoPE [-> Scale -> FP16 Copy]` fusion. CuTe attention shares @@ -136,8 +142,16 @@ its RMSNorm/RoPE kernel. Partial rotation and arbitrary float-aligned input and position-encoding strides are supported by the standalone kernel; RMSNorm fusion requires contiguous position encoding and weights. Copy fusion uses workspace when the destination overlaps an input. These registrations do not require -CUTLASS. The old CUDA RoPE dispatch is removed; HIP and MUSA retain their existing -implementations. +CUTLASS. The old CUDA/HIP/MUSA RoPE dispatch is removed. + +Shared CUDA-style kernels live in `src/core/cuda/` and are compiled separately +for each enabled runtime. `SD_CUDA_NAMESPACE` expands to `sd_cuda`, `sd_hip`, or +`sd_musa`, so implementations can coexist. `runtime.h` handles the small runtime +differences, including logical 32-lane shuffles on HIP. The application defines +and launches its kernels directly; it includes only public ggml headers. +CMake generates the registration table from enabled backend options. The shared +entry point uses one `SD_BACKEND_EXTENSIONS` macro; backend registrars verify the +device kind, and disabled backends add no symbol references. `SD_USE_UPSTREAM_GGML` excludes these registrations. The stable-diffusion public C API is unchanged. INT8 ConvRot retains its existing packed activation layout and @@ -216,10 +230,12 @@ Custom I8 nodes also retain cuBLAS when the reduction could overflow I32, to preserve their previous accumulation behavior. Fusion disable and force-cuBLAS settings are honored. No CUTLASS headers are needed for the cuBLAS path. -The native CUDA CUTLASS GEMM implementation and its SwiGLU/row-concat matchers -have been removed from ggml. Native cuBLAS GEMM and standalone ConvRot remain -compatibility implementations for callers without these registrations. HIP -retains its native INT8/activation fusion paths. +The native CUDA/HIP INT8 GEMM, standalone ConvRot, and associated fusion paths +have been removed from ggml. The shared application implementation calls cuBLAS +or hipBLAS directly. HIP retains the existing architecture eligibility policy; +MUSA does not advertise INT8 BLAS support and uses scheduler fallback for GEMM. +Standalone ConvRot kernels are compiled for all three runtimes. CUTLASS +implementations remain CUDA-specific. `src/core/cuda/i8_gated_residual.cu` implements `INT8 GEMM -> [Scale] -> Mul(gate) -> Add(residual) -> [Concat]` as a registered @@ -255,8 +271,8 @@ overlapping outputs. Long rows or external intermediate consumers retain ordinary execution. These fusions do not require CUTLASS and honor `GGML_CUDA_DISABLE_FUSION`. -The corresponding scale, normalization, and unary-activation fusion rules in -ggml now only serve HIP. CUDA GEMM/SwiGLU/ConvRot chains are selected by the +The corresponding scale, normalization, and unary-activation fusion rules have +been removed from ggml for CUDA and HIP. GEMM/SwiGLU/ConvRot chains are selected by the application INT8 extension before the standalone activation fusion. `src/core/cuda/affine.cu` registers FP32 `Mul(gate) -> Add(residual)`, @@ -268,8 +284,8 @@ Add inputs, and fusion disabling retain ordinary operators. Internal in-place Add outputs are accepted when they overwrite the complete fused Mul allocation. Exact in-place residual or input addresses are safe for the pointwise kernel; partial overlap, broadcast gate overlap, and FP16 copy overlap use scheduler workspace. These registrations -require no CUTLASS and support reusable CUDA Graph replay. Their former native -fusion implementations now only serve HIP/MUSA. +require no CUTLASS, compile for CUDA/HIP/MUSA, and support reusable CUDA Graph +replay. Their former native fusion implementations have been removed from ggml. CUDA kernels can opt into device replay with `ggml_ext_register_kernel_ex` and `GGML_EXT_KERNEL_CAPTURE_SAFE`. The original registration function retains its diff --git a/ggml b/ggml index 64edc3b7..a2940033 160000 --- a/ggml +++ b/ggml @@ -1 +1 @@ -Subproject commit 64edc3b78462898ac6c7a0261945240ad2891706 +Subproject commit a29400335191f983544d14df03773a50420224d7 diff --git a/src/core/compute_workspace.cpp b/src/core/compute_workspace.cpp index de1cea86..1c11c12d 100644 --- a/src/core/compute_workspace.cpp +++ b/src/core/compute_workspace.cpp @@ -147,6 +147,9 @@ namespace sd { append(usage); append(type); } + } else { + const auto alignment_offset = reinterpret_cast(tensor->data) % ggml_backend_get_alignment(external); + append(alignment_offset); } const size_t index = ggml_hash_find(&graph->visited_hash_set, tensor); const int32_t uses = index != GGML_HASHSET_FULL && ggml_bitset_get(graph->visited_hash_set.used, index) @@ -187,7 +190,8 @@ namespace sd { ggml_backend_buffer_set_usage(buffer, GGML_BACKEND_BUFFER_USAGE_WEIGHTS); } copy->buffer = buffer; - copy->data = reinterpret_cast(static_cast(1)); + // Capability checks inspect alignment even before weights are resident. + copy->data = source->data ? source->data : reinterpret_cast(ggml_backend_get_alignment(external)); copy->extra = nullptr; } } diff --git a/src/core/cuda/CMakeLists.txt b/src/core/cuda/CMakeLists.txt new file mode 100644 index 00000000..3d4509d6 --- /dev/null +++ b/src/core/cuda/CMakeLists.txt @@ -0,0 +1,90 @@ +cmake_policy(SET CMP0079 NEW) + +function(sd_cuda_sources VENDOR LANGUAGE) + set(SOURCES) + foreach(OP rope int8 convrot affine 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") + set_source_files_properties("${SOURCE}" PROPERTIES LANGUAGE ${LANGUAGE}) + list(APPEND SOURCES "${SOURCE}") + endforeach() + set(TARGET "sd-${VENDOR}") + add_library(${TARGET} OBJECT ${SOURCES}) + target_include_directories(${TARGET} PRIVATE "${PROJECT_SOURCE_DIR}/src") + target_link_libraries(${TARGET} PRIVATE ggml) + set_target_properties(${TARGET} PROPERTIES POSITION_INDEPENDENT_CODE ON) + target_compile_features(${TARGET} PRIVATE cxx_std_17) + target_link_libraries(${SD_LIB} PRIVATE ${TARGET}) + set(SD_CUDA_TARGET ${TARGET} PARENT_SCOPE) + set(SD_CUDA_SOURCES ${SOURCES} PARENT_SCOPE) +endfunction() + +if(SD_CUDA) + enable_language(CUDA) + sd_cuda_sources(cuda CUDA) + set_source_files_properties(${SD_CUDA_SOURCES} PROPERTIES COMPILE_OPTIONS "--use_fast_math") + if(GGML_CUDA_CUTLASS_PATH) + set_property(SOURCE "${CMAKE_CURRENT_BINARY_DIR}/int8_cuda.cu" APPEND PROPERTY COMPILE_DEFINITIONS SD_CUDA_I8_CUTLASS) + set_property(SOURCE "${CMAKE_CURRENT_BINARY_DIR}/registration_cuda.cu" APPEND PROPERTY COMPILE_DEFINITIONS SD_CUDA_CUTE_ATTENTION) + set_property(SOURCE "${CMAKE_CURRENT_BINARY_DIR}/int8_cuda.cu" APPEND PROPERTY INCLUDE_DIRECTORIES "${PROJECT_SOURCE_DIR}/src/core/cuda") + endif() + set_target_properties(${SD_CUDA_TARGET} PROPERTIES CUDA_STANDARD 17 CUDA_STANDARD_REQUIRED ON) + target_link_libraries(${SD_CUDA_TARGET} PRIVATE CUDA::cudart CUDA::cublas) + target_link_libraries(${SD_LIB} PRIVATE CUDA::cudart CUDA::cublas) +endif() + +if(SD_HIPBLAS) + list(APPEND CMAKE_PREFIX_PATH "$ENV{ROCM_PATH}" /opt/rocm /usr) + find_package(hip REQUIRED) + find_package(hipblas REQUIRED) + if(WIN32 OR CMAKE_CXX_COMPILER MATCHES "hipcc") + sd_cuda_sources(hip CXX) + target_link_libraries(${SD_CUDA_TARGET} PRIVATE hip::device) + else() + if(NOT CMAKE_HIP_ARCHITECTURES) + if(GPU_TARGETS) + set(CMAKE_HIP_ARCHITECTURES ${GPU_TARGETS}) + elseif(AMDGPU_TARGETS) + set(CMAKE_HIP_ARCHITECTURES ${AMDGPU_TARGETS}) + endif() + endif() + enable_language(HIP) + sd_cuda_sources(hip HIP) + set_target_properties(${SD_CUDA_TARGET} PROPERTIES HIP_STANDARD 17 HIP_STANDARD_REQUIRED ON) + endif() + set_source_files_properties(${SD_CUDA_SOURCES} PROPERTIES COMPILE_OPTIONS "-ffast-math") + target_link_libraries(${SD_CUDA_TARGET} PRIVATE hip::host roc::hipblas) + target_link_libraries(${SD_LIB} PRIVATE hip::host roc::hipblas) +endif() + +if(SD_MUSA) + if(NOT MUSA_PATH) + if(DEFINED ENV{MUSA_PATH}) + set(MUSA_PATH "$ENV{MUSA_PATH}") + elseif(EXISTS /opt/musa) + set(MUSA_PATH /opt/musa) + else() + set(MUSA_PATH /usr/local/musa) + endif() + endif() + list(APPEND CMAKE_MODULE_PATH "${MUSA_PATH}/cmake") + find_package(MUSAToolkit REQUIRED) + set(CMAKE_CXX_COMPILER "${MUSA_PATH}/bin/clang++") + sd_cuda_sources(musa CXX) + set(FLAGS -Od3 -fno-strict-aliasing -ffast-math -fsigned-char -x musa -mtgpu -fmusa-flush-denormals-to-zero) + if(NOT DEFINED MUSA_ARCHITECTURES) + set(MUSA_ARCHITECTURES "22;31") + endif() + foreach(ARCH ${MUSA_ARCHITECTURES}) + list(APPEND FLAGS "--cuda-gpu-arch=mp_${ARCH}") + endforeach() + set_source_files_properties(${SD_CUDA_SOURCES} PROPERTIES COMPILE_OPTIONS "${FLAGS}") + if(GGML_STATIC) + target_link_libraries(${SD_CUDA_TARGET} PRIVATE MUSA::musart_static MUSA::mublas_static) + target_link_libraries(${SD_LIB} PRIVATE MUSA::musart_static MUSA::mublas_static) + else() + target_link_libraries(${SD_CUDA_TARGET} PRIVATE MUSA::musart MUSA::mublas) + target_link_libraries(${SD_LIB} PRIVATE MUSA::musart MUSA::mublas) + endif() +endif() diff --git a/src/core/cuda/affine.cu b/src/core/cuda/affine.cu index d96f6d1c..e410d40b 100644 --- a/src/core/cuda/affine.cu +++ b/src/core/cuda/affine.cu @@ -1,183 +1,183 @@ -#include -#include #include #include #include #include "ggml-ext.h" +#include "runtime.h" -namespace sd_affine_cuda { - static bool disabled() { - static const bool value = [] { const char* v = std::getenv("GGML_CUDA_DISABLE_FUSION"); return v && std::atoi(v) != 0; }(); - return value; - } - static bool aligned(const ggml_tensor* tensor, size_t alignment) { - size_t offset = 0; - while (!tensor->data && tensor->view_src) { - offset += tensor->view_offs; - tensor = tensor->view_src; +namespace SD_CUDA_NAMESPACE { + namespace sd_affine_cuda { + static bool disabled() { + static const bool value = [] { const char* v = std::getenv("GGML_CUDA_DISABLE_FUSION"); return v && std::atoi(v) != 0; }(); + return value; } - return (reinterpret_cast(tensor->data) + offset) % alignment == 0; - } - static float param(const ggml_tensor* tensor, int index) { - float value; - std::memcpy(&value, tensor->op_params + index, sizeof(value)); - return value; - } - static bool contiguous_f32(const ggml_tensor* tensor) { - return tensor && tensor->type == GGML_TYPE_F32 && ggml_is_contiguous(tensor) && aligned(tensor, 16); - } - static bool bounded(const ggml_tensor* tensor) { - return ggml_nelements(tensor) > 0 && ggml_nelements(tensor) <= int64_t(INT_MAX) * 1024; - } - static int match(ggml_backend_dev_t, const ggml_cgraph* graph, int first, void*) { - auto* g = const_cast(graph); - if (disabled() || first + 1 >= ggml_graph_n_nodes(g)) + static bool aligned(const ggml_tensor* tensor, size_t alignment) { + size_t offset = 0; + while (!tensor->data && tensor->view_src) { + offset += tensor->view_offs; + tensor = tensor->view_src; + } + return (reinterpret_cast(tensor->data) + offset) % alignment == 0; + } + static float param(const ggml_tensor* tensor, int index) { + float value; + std::memcpy(&value, tensor->op_params + index, sizeof(value)); + return value; + } + static bool contiguous_f32(const ggml_tensor* tensor) { + return tensor && tensor->type == GGML_TYPE_F32 && ggml_is_contiguous(tensor) && aligned(tensor, 16); + } + static bool bounded(const ggml_tensor* tensor) { + return ggml_nelements(tensor) > 0 && ggml_nelements(tensor) <= int64_t(INT_MAX) * 1024; + } + static int match(ggml_backend_dev_t, const ggml_cgraph* graph, int first, void*) { + auto* g = const_cast(graph); + if (disabled() || first + 1 >= ggml_graph_n_nodes(g)) + return 0; + const auto* node = ggml_graph_node(g, first); + const auto* next = ggml_graph_node(g, first + 1); + if (node->op == GGML_OP_MUL) { + if (next->op != GGML_OP_ADD || (next->src[0] != node && next->src[1] != node)) + return 0; + const auto* x = node->src[0]; + const auto* gate = node->src[1]; + const auto* residual = next->src[0] == node ? next->src[1] : next->src[0]; + return residual != node && contiguous_f32(x) && contiguous_f32(gate) && contiguous_f32(residual) && contiguous_f32(next) && + node->type == GGML_TYPE_F32 && ggml_are_same_shape(node, next) && ggml_are_same_shape(x, next) && + ggml_are_same_shape(residual, next) && gate->ne[0] == x->ne[0] && ggml_nelements(gate) == x->ne[0] && + x->ne[0] % 4 == 0 && bounded(next) + ? 2 + : 0; + } + if (node->op == GGML_OP_CONCAT) { + const bool permuted = next->op == GGML_OP_PERMUTE; + const int count = permuted ? 3 : 2; + if (first + count > ggml_graph_n_nodes(g) || node->type != GGML_TYPE_F32 || node->ne[0] % 4 || + node->op_params[0] != (permuted ? 2 : 1) || !ggml_is_contiguous(node) || + !contiguous_f32(node->src[0]) || !contiguous_f32(node->src[1])) + return 0; + if (permuted && (next->src[0] != node || next->ne[0] != node->ne[0] || next->ne[1] != node->ne[2] || + next->ne[2] != node->ne[1] || next->ne[3] != node->ne[3] || next->nb[0] != node->nb[0] || + next->nb[1] != node->nb[2] || next->nb[2] != node->nb[1] || next->nb[3] != node->nb[3])) + return 0; + const auto* dst = ggml_graph_node(g, first + count - 1); + return dst->op == GGML_OP_CPY && dst->src[0] == (permuted ? next : node) && dst->type == GGML_TYPE_F16 && + ggml_is_contiguous(dst) && ggml_are_same_shape(dst, dst->src[0]) && aligned(dst, 4) && bounded(dst) + ? count + : 0; + } + if (node->op == GGML_OP_CONT && first + 3 < ggml_graph_n_nodes(g)) { + const auto* src = node->src[0]; + const auto* permute = ggml_graph_node(g, first + 2); + const auto* dst = ggml_graph_node(g, first + 3); + return ggml_are_same_shape(src, node) && src->type == GGML_TYPE_F32 && src->nb[0] == 4 && src->ne[0] % 4 == 0 && + aligned(src, 16) && (src->nb[1] | src->nb[2] | src->nb[3]) % 16 == 0 && + next->op == GGML_OP_SCALE && next->src[0] == node && next->type == GGML_TYPE_F32 && param(next, 1) == 0.f && + permute->op == GGML_OP_PERMUTE && permute->src[0] == next && permute->op_params[0] == 0 && + permute->op_params[1] == 2 && permute->op_params[2] == 1 && permute->op_params[3] == 3 && + dst->op == GGML_OP_CPY && dst->src[0] == permute && dst->type == GGML_TYPE_F16 && + ggml_are_same_shape(dst, permute) && ggml_is_contiguous(dst) && aligned(dst, 4) && bounded(dst) + ? 4 + : 0; + } return 0; - const auto* node = ggml_graph_node(g, first); - const auto* next = ggml_graph_node(g, first + 1); - if (node->op == GGML_OP_MUL) { - if (next->op != GGML_OP_ADD || (next->src[0] != node && next->src[1] != node)) - return 0; - const auto* x = node->src[0]; - const auto* gate = node->src[1]; - const auto* residual = next->src[0] == node ? next->src[1] : next->src[0]; - return residual != node && contiguous_f32(x) && contiguous_f32(gate) && contiguous_f32(residual) && contiguous_f32(next) && - node->type == GGML_TYPE_F32 && ggml_are_same_shape(node, next) && ggml_are_same_shape(x, next) && - ggml_are_same_shape(residual, next) && gate->ne[0] == x->ne[0] && ggml_nelements(gate) == x->ne[0] && - x->ne[0] % 4 == 0 && bounded(next) - ? 2 - : 0; } - if (node->op == GGML_OP_CONCAT) { - const bool permuted = next->op == GGML_OP_PERMUTE; - const int count = permuted ? 3 : 2; - if (first + count > ggml_graph_n_nodes(g) || node->type != GGML_TYPE_F32 || node->ne[0] % 4 || - node->op_params[0] != (permuted ? 2 : 1) || !ggml_is_contiguous(node) || - !contiguous_f32(node->src[0]) || !contiguous_f32(node->src[1])) - return 0; - if (permuted && (next->src[0] != node || next->ne[0] != node->ne[0] || next->ne[1] != node->ne[2] || - next->ne[2] != node->ne[1] || next->ne[3] != node->ne[3] || next->nb[0] != node->nb[0] || - next->nb[1] != node->nb[2] || next->nb[2] != node->nb[1] || next->nb[3] != node->nb[3])) - return 0; - const auto* dst = ggml_graph_node(g, first + count - 1); - return dst->op == GGML_OP_CPY && dst->src[0] == (permuted ? next : node) && dst->type == GGML_TYPE_F16 && - ggml_is_contiguous(dst) && ggml_are_same_shape(dst, dst->src[0]) && aligned(dst, 4) && bounded(dst) - ? count - : 0; + static __global__ void gated_residual(const float4* x, const float4* gate, const float4* residual, float4* dst, int64_t n4, int64_t count4) { + const int64_t i = int64_t(blockIdx.x) * blockDim.x + threadIdx.x; + if (i >= count4) + return; + const float4 a = x[i], g = gate[i % n4], r = residual[i]; + // The graph performs two FP32 operations; contracting them into FMA changes rounding. + dst[i] = make_float4(__fadd_rn(__fmul_rn(a.x, g.x), r.x), __fadd_rn(__fmul_rn(a.y, g.y), r.y), + __fadd_rn(__fmul_rn(a.z, g.z), r.z), __fadd_rn(__fmul_rn(a.w, g.w), r.w)); } - if (node->op == GGML_OP_CONT && first + 3 < ggml_graph_n_nodes(g)) { - const auto* src = node->src[0]; - const auto* permute = ggml_graph_node(g, first + 2); - const auto* dst = ggml_graph_node(g, first + 3); - return ggml_are_same_shape(src, node) && src->type == GGML_TYPE_F32 && src->nb[0] == 4 && src->ne[0] % 4 == 0 && - aligned(src, 16) && (src->nb[1] | src->nb[2] | src->nb[3]) % 16 == 0 && - next->op == GGML_OP_SCALE && next->src[0] == node && next->type == GGML_TYPE_F32 && param(next, 1) == 0.f && - permute->op == GGML_OP_PERMUTE && permute->src[0] == next && permute->op_params[0] == 0 && - permute->op_params[1] == 2 && permute->op_params[2] == 1 && permute->op_params[3] == 3 && - dst->op == GGML_OP_CPY && dst->src[0] == permute && dst->type == GGML_TYPE_F16 && - ggml_are_same_shape(dst, permute) && ggml_is_contiguous(dst) && aligned(dst, 4) && bounded(dst) - ? 4 - : 0; + template + static __global__ void concat_cast(const float4* prefix, const float4* current, half2* dst, int64_t columns4, int64_t prefix_tokens, int64_t tokens, int64_t heads, int64_t count4) { + const int64_t i = int64_t(blockIdx.x) * blockDim.x + threadIdx.x; + if (i >= count4) + return; + const int64_t col = i % columns4, row = i / columns4; + const int64_t token = row % tokens, head = row / tokens % heads, batch = row / (tokens * heads); + const bool from_prefix = token < prefix_tokens; + const int64_t src_tokens = from_prefix ? prefix_tokens : tokens - prefix_tokens; + const int64_t src_token = from_prefix ? token : token - prefix_tokens; + const int64_t src_row = Permuted ? (batch * src_tokens + src_token) * heads + head + : (batch * heads + head) * src_tokens + src_token; + const float4 value = (from_prefix ? prefix : current)[src_row * columns4 + col]; + dst[2 * i] = __floats2half2_rn(value.x, value.y); + dst[2 * i + 1] = __floats2half2_rn(value.z, value.w); } - return 0; - } - static __global__ void gated_residual(const float4* x, const float4* gate, const float4* residual, float4* dst, int64_t n4, int64_t count4) { - const int64_t i = int64_t(blockIdx.x) * blockDim.x + threadIdx.x; - if (i >= count4) - return; - const float4 a = x[i], g = gate[i % n4], r = residual[i]; - // The graph performs two FP32 operations; contracting them into FMA changes rounding. - dst[i] = make_float4(__fadd_rn(__fmul_rn(a.x, g.x), r.x), __fadd_rn(__fmul_rn(a.y, g.y), r.y), - __fadd_rn(__fmul_rn(a.z, g.z), r.z), __fadd_rn(__fmul_rn(a.w, g.w), r.w)); - } - template - static __global__ void concat_cast(const float4* prefix, const float4* current, half2* dst, int64_t columns4, int64_t prefix_tokens, int64_t tokens, int64_t heads, int64_t count4) { - const int64_t i = int64_t(blockIdx.x) * blockDim.x + threadIdx.x; - if (i >= count4) - return; - const int64_t col = i % columns4, row = i / columns4; - const int64_t token = row % tokens, head = row / tokens % heads, batch = row / (tokens * heads); - const bool from_prefix = token < prefix_tokens; - const int64_t src_tokens = from_prefix ? prefix_tokens : tokens - prefix_tokens; - const int64_t src_token = from_prefix ? token : token - prefix_tokens; - const int64_t src_row = Permuted ? (batch * src_tokens + src_token) * heads + head - : (batch * heads + head) * src_tokens + src_token; - const float4 value = (from_prefix ? prefix : current)[src_row * columns4 + col]; - dst[2 * i] = __floats2half2_rn(value.x, value.y); - dst[2 * i + 1] = __floats2half2_rn(value.z, value.w); - } - static __global__ void scale_permute(const float* src, half2* dst, int64_t dim, int64_t tokens, int64_t heads, int64_t count, int64_t nb1, int64_t nb2, int64_t nb3, float scale) { - const int64_t index = (int64_t(blockIdx.x) * blockDim.x + threadIdx.x) * 4; - if (index >= count) - return; - const int64_t col = index % dim, row = index / dim; - const int64_t token = row % tokens, head = row / tokens % heads, batch = row / (tokens * heads); - const float4 value = *reinterpret_cast(src + col + head * nb1 + token * nb2 + batch * nb3); - dst[index / 2] = __floats2half2_rn(__fmul_rn(value.x, scale), __fmul_rn(value.y, scale)); - dst[index / 2 + 1] = __floats2half2_rn(__fmul_rn(value.z, scale), __fmul_rn(value.w, scale)); - } - 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 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 count, void*, void*) { - 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); - const auto* node = nodes[0]; - const auto* a = node->src[0]; - const auto* b = node->op == GGML_OP_CONT ? nullptr : node->src[1]; - auto* dst = nodes[count - 1]; - const auto* residual = node->op == GGML_OP_MUL ? (dst->src[0] == node ? dst->src[1] : dst->src[0]) : nullptr; - const bool gate = residual != nullptr; - const bool alias = (overlaps(dst, a) && !(gate && dst->data == a->data)) || (b && overlaps(dst, b)) || - (residual && overlaps(dst, residual) && dst->data != residual->data); - void* output = alias ? ggml_backend_buffer_get_base(execution->workspace) : dst->data; - const int64_t count4 = ggml_nelements(dst) / 4; - const unsigned blocks = unsigned((count4 + 255) / 256); - if (gate) { - gated_residual<<>>(static_cast(a->data), static_cast(b->data), - static_cast(residual->data), static_cast(output), a->ne[0] / 4, count4); - } else if (node->op == GGML_OP_CONCAT) { - const bool permuted = count == 3; - auto kernel = permuted ? concat_cast : concat_cast; - kernel<<>>(static_cast(a->data), static_cast(b->data), - static_cast(output), node->ne[0] / 4, a->ne[permuted ? 2 : 1], node->ne[permuted ? 2 : 1], - node->ne[permuted ? 1 : 2], count4); - } else { - scale_permute<<>>(static_cast(a->data), static_cast(output), - a->ne[0], a->ne[2], a->ne[1], ggml_nelements(dst), a->nb[1] / 4, a->nb[2] / 4, a->nb[3] / 4, param(nodes[1], 0)); + static __global__ void scale_permute(const float* src, half2* dst, int64_t dim, int64_t tokens, int64_t heads, int64_t count, int64_t nb1, int64_t nb2, int64_t nb3, float scale) { + const int64_t index = (int64_t(blockIdx.x) * blockDim.x + threadIdx.x) * 4; + if (index >= count) + return; + const int64_t col = index % dim, row = index / dim; + const int64_t token = row % tokens, head = row / tokens % heads, batch = row / (tokens * heads); + const float4 value = *reinterpret_cast(src + col + head * nb1 + token * nb2 + batch * nb3); + dst[index / 2] = __floats2half2_rn(__fmul_rn(value.x, scale), __fmul_rn(value.y, scale)); + dst[index / 2 + 1] = __floats2half2_rn(__fmul_rn(value.z, scale), __fmul_rn(value.w, scale)); + } + 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 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 count, void*, void*) { + ggml_ext_gpu_context context; + if (!cuda_context(execution->backend, &context)) + return GGML_STATUS_FAILED; + auto stream = static_cast(context.stream); + const auto* node = nodes[0]; + const auto* a = node->src[0]; + const auto* b = node->op == GGML_OP_CONT ? nullptr : node->src[1]; + auto* dst = nodes[count - 1]; + const auto* residual = node->op == GGML_OP_MUL ? (dst->src[0] == node ? dst->src[1] : dst->src[0]) : nullptr; + const bool gate = residual != nullptr; + const bool alias = (overlaps(dst, a) && !(gate && dst->data == a->data)) || (b && overlaps(dst, b)) || + (residual && overlaps(dst, residual) && dst->data != residual->data); + void* output = alias ? ggml_backend_buffer_get_base(execution->workspace) : dst->data; + const int64_t count4 = ggml_nelements(dst) / 4; + const unsigned blocks = unsigned((count4 + 255) / 256); + if (gate) { + gated_residual<<>>(static_cast(a->data), static_cast(b->data), + static_cast(residual->data), static_cast(output), a->ne[0] / 4, count4); + } else if (node->op == GGML_OP_CONCAT) { + const bool permuted = count == 3; + auto kernel = permuted ? concat_cast : concat_cast; + kernel<<>>(static_cast(a->data), static_cast(b->data), + static_cast(output), node->ne[0] / 4, a->ne[permuted ? 2 : 1], node->ne[permuted ? 2 : 1], + node->ne[permuted ? 1 : 2], count4); + } else { + scale_permute<<>>(static_cast(a->data), static_cast(output), + a->ne[0], a->ne[2], a->ne[1], ggml_nelements(dst), a->nb[1] / 4, a->nb[2] / 4, a->nb[3] / 4, param(nodes[1], 0)); + } + auto error = cuda_last_error(); + 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; } - auto 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_affine(ggml_backend_dev_t device) { - ggml_ext_kernel_desc desc = {}; - desc.abi_version = GGML_EXT_ABI_VERSION; - desc.priority = 200; - desc.match = sd_affine_cuda::match; - desc.prepare = sd_affine_cuda::prepare; - desc.destroy = sd_affine_cuda::destroy; - desc.execute = sd_affine_cuda::execute; - desc.name = "sd.gated_residual.cuda.v1"; - desc.op = GGML_OP_MUL; - GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE)); - desc.name = "sd.concat_cast.cuda.v1"; - desc.op = GGML_OP_CONCAT; - GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE)); - desc.name = "sd.cont_scale_permute_cast.cuda.v1"; - desc.op = GGML_OP_CONT; - GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE)); + void register_affine(ggml_backend_dev_t device) { + ggml_ext_kernel_desc desc = {}; + desc.abi_version = GGML_EXT_ABI_VERSION; + desc.priority = 200; + desc.match = sd_affine_cuda::match; + desc.prepare = sd_affine_cuda::prepare; + desc.destroy = sd_affine_cuda::destroy; + desc.execute = sd_affine_cuda::execute; + desc.name = "sd.gated_residual." SD_CUDA_NAME ".v1"; + desc.op = GGML_OP_MUL; + GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE)); + desc.name = "sd.concat_cast." SD_CUDA_NAME ".v1"; + desc.op = GGML_OP_CONCAT; + GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE)); + desc.name = "sd.cont_scale_permute_cast." SD_CUDA_NAME ".v1"; + desc.op = GGML_OP_CONT; + 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/convrot.cu b/src/core/cuda/convrot.cu index eba72cc1..a48ffe56 100644 --- a/src/core/cuda/convrot.cu +++ b/src/core/cuda/convrot.cu @@ -2,314 +2,295 @@ #include #include #include -#include #include -#include #include "convrot_kernel.h" #include "core/int8_ops.h" +#include "runtime.h" -namespace sd_convrot { - struct DeviceInfo { - int ordinal = -1; - 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 query = reinterpret_cast(ggml_backend_reg_get_proc_address(ggml_backend_dev_backend_reg(device), "ggml_ext_cuda_device_v1")); - int ordinal = query ? query(device) : -1; - cudaDeviceProp props = {}; - if (ordinal >= 0 && cudaGetDeviceProperties(&props, ordinal) == cudaSuccess) - info = {ordinal, size_t(props.sharedMemPerBlockOptin)}; - cache.emplace(device, info); - return info; - } - 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 float param_f32(const ggml_tensor* tensor, int index) { - float value; - std::memcpy(&value, tensor->op_params + index, sizeof(value)); - return value; - } - static bool f32(const ggml_tensor* tensor) { - return tensor && tensor->type == GGML_TYPE_F32 && ggml_is_contiguous(tensor); - } - static bool vector(const ggml_tensor* tensor, int64_t k) { - return f32(tensor) && tensor->ne[0] == k && ggml_nelements(tensor) == k; - } - enum Kind { SCALE, - NORM, - RMS_MOD, - SWIGLU, - SILU, - SIGMOID }; - struct Operation { - Kind kind; - const ggml_tensor* input = nullptr; - const ggml_tensor* up = nullptr; - const ggml_tensor* modulation = nullptr; - const ggml_tensor* weight = nullptr; - const ggml_tensor* shift = nullptr; - int64_t k, rows, stride, up_stride, gate_offset = 0, up_offset = 0; - float eps = 0.f, scale = 1.f, bias = 0.f; - }; - static bool describe(ggml_tensor* const* nodes, int count, Operation& op) { - int group; - const auto* first = nodes[0]; - const auto* quant = nodes[count - 1]; - if (!sd_ggml_convrot_describe(quant, &group) || group != 256) - return false; - const auto* value = quant->src[0]; - op.k = value->ne[0]; - op.rows = ggml_nrows(value); - op.input = first->src[0]; - if (!f32(op.input) || !f32(first) || op.k > INT_MAX || op.rows > INT_MAX) - return false; - op.stride = op.input->ne[0]; - op.up_stride = op.k; - if (count == 2 && first->op == GGML_OP_SCALE && value == first) { - op.kind = SCALE; - op.scale = param_f32(first, 0); - op.bias = param_f32(first, 1); - } else if (count == 3 && first->op == GGML_OP_NORM && nodes[1]->op == GGML_OP_MUL && nodes[1]->src[0] == first && value == nodes[1]) { - op.kind = NORM; - op.modulation = nodes[1]->src[1]; - 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]) { - op.kind = RMS_MOD; - op.weight = nodes[1]->src[1]; - op.modulation = nodes[2]->src[1]; - op.shift = nodes[4]->src[1]; - 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)) - 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; - op.up = first->src[1] ? first->src[1] : op.input; - if (!f32(op.up)) - return false; - op.up_stride = op.up->ne[0]; - if (!first->src[1]) { - if (first->op_params[1]) - op.gate_offset = op.k; - else - op.up_offset = op.k; - } - } else if (count == 3 && first->op == GGML_OP_UNARY && nodes[1]->op == GGML_OP_MUL && value == nodes[1] && - (nodes[1]->src[0] == first || nodes[1]->src[1] == first)) { - auto unary = ggml_get_unary_op(first); - if (unary != GGML_UNARY_OP_SILU && unary != GGML_UNARY_OP_SIGMOID) - return false; - op.kind = unary == GGML_UNARY_OP_SILU ? SILU : SIGMOID; - op.up = nodes[1]->src[0] == first ? nodes[1]->src[1] : nodes[1]->src[0]; - if (!f32(op.up) || !ggml_are_same_shape(op.input, op.up)) - return false; - } else - return false; - for (int i = 0; i < count - 1; ++i) { - if (!f32(nodes[i]) || !ggml_are_same_shape(nodes[i], value)) - return false; - for (const auto* input : {op.input, op.up, op.modulation, op.weight, op.shift}) - if (input == nodes[i]) - return false; +namespace SD_CUDA_NAMESPACE { + namespace sd_convrot { + + 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; } - return true; - } - static bool supported(const Operation& op, const DeviceInfo& device) { - const size_t extra = op.kind == NORM ? 384 : op.kind == RMS_MOD ? 256 - : 128; - return device.ordinal >= 0 && size_t(op.k) * sizeof(float) + extra <= device.shared; - } - static int match(ggml_backend_dev_t device, const ggml_cgraph* graph, int first, void*) { - if (fusion_disabled()) - return 0; - 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 - : 2; - if (available < count) - return 0; - ggml_tensor* nodes[6]; - for (int j = 0; j < count; ++j) - nodes[j] = ggml_graph_node(g, first + j); - Operation op; - return describe(nodes, count, op) && supported(op, device_info(device)) ? count : 0; - } - 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) { - auto* quant = ggml_graph_node(graph, i); + 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 f32(const ggml_tensor* tensor) { + return tensor && tensor->type == GGML_TYPE_F32 && ggml_is_contiguous(tensor); + } + static bool vector(const ggml_tensor* tensor, int64_t k) { + return f32(tensor) && tensor->ne[0] == k && ggml_nelements(tensor) == k; + } + enum Kind { SCALE, + NORM, + RMS_MOD, + SWIGLU, + SILU, + SIGMOID }; + struct Operation { + Kind kind; + const ggml_tensor* input = nullptr; + const ggml_tensor* up = nullptr; + const ggml_tensor* modulation = nullptr; + const ggml_tensor* weight = nullptr; + const ggml_tensor* shift = nullptr; + int64_t k, rows, stride, up_stride, gate_offset = 0, up_offset = 0; + float eps = 0.f, scale = 1.f, bias = 0.f; + }; + static bool describe(ggml_tensor* const* nodes, int count, Operation& op) { int group; + const auto* first = nodes[0]; + const auto* quant = nodes[count - 1]; if (!sd_ggml_convrot_describe(quant, &group) || group != 256) - continue; - 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) { - chain[0] = weight->src[0]; - chain[1] = weight; - chain[2] = add->src[1]; - chain[3] = add; - chain[4] = value; - chain[5] = quant; - count = 6; + return false; + const auto* value = quant->src[0]; + op.k = value->ne[0]; + op.rows = ggml_nrows(value); + op.input = first->src[0]; + if (!f32(op.input) || !f32(first) || op.k > INT_MAX || op.rows > INT_MAX) + return false; + op.stride = op.input->ne[0]; + op.up_stride = op.k; + if (count == 2 && first->op == GGML_OP_SCALE && value == first) { + op.kind = SCALE; + op.scale = param_f32(first, 0); + op.bias = param_f32(first, 1); + } else if (count == 3 && first->op == GGML_OP_NORM && nodes[1]->op == GGML_OP_MUL && nodes[1]->src[0] == first && value == nodes[1]) { + op.kind = NORM; + op.modulation = nodes[1]->src[1]; + 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]) { + op.kind = RMS_MOD; + op.weight = nodes[1]->src[1]; + op.modulation = nodes[2]->src[1]; + op.shift = nodes[4]->src[1]; + 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)) + 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; + op.up = first->src[1] ? first->src[1] : op.input; + if (!f32(op.up)) + return false; + op.up_stride = op.up->ne[0]; + if (!first->src[1]) { + if (first->op_params[1]) + op.gate_offset = op.k; + else + op.up_offset = op.k; } - } else if (value->op == GGML_OP_MUL) { - for (int src = 0; src < 2; ++src) - if (value->src[src]->op == GGML_OP_UNARY) { - chain[0] = value->src[src]; - chain[1] = value; - chain[2] = quant; - count = 3; - break; - } + } else if (count == 3 && first->op == GGML_OP_UNARY && nodes[1]->op == GGML_OP_MUL && value == nodes[1] && + (nodes[1]->src[0] == first || nodes[1]->src[1] == first)) { + auto unary = ggml_get_unary_op(first); + if (unary != GGML_UNARY_OP_SILU && unary != GGML_UNARY_OP_SIGMOID) + return false; + op.kind = unary == GGML_UNARY_OP_SILU ? SILU : SIGMOID; + op.up = nodes[1]->src[0] == first ? nodes[1]->src[1] : nodes[1]->src[0]; + if (!f32(op.up) || !ggml_are_same_shape(op.input, op.up)) + return false; + } else + return false; + for (int i = 0; i < count - 1; ++i) { + if (!f32(nodes[i]) || !ggml_are_same_shape(nodes[i], value)) + return false; + for (const auto* input : {op.input, op.up, op.modulation, op.weight, op.shift}) + if (input == nodes[i]) + return false; } + return true; + } + static bool supported(const Operation& op, const ggml_ext_gpu_device_info& device) { + const size_t extra = op.kind == NORM ? 384 : op.kind == RMS_MOD ? 256 + : 128; + return device.device >= 0 && size_t(op.k) * sizeof(float) + extra <= device.shared_memory; + } + static int match(ggml_backend_dev_t device, const ggml_cgraph* graph, int first, void*) { + if (fusion_disabled()) + return 0; + 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 + : 2; + if (available < count) + return 0; + ggml_tensor* nodes[6]; + for (int j = 0; j < count; ++j) + nodes[j] = ggml_graph_node(g, first + j); Operation op; - if (!count || !describe(chain, count, op) || !supported(op, device_info(device))) - continue; - int begin = -1; - for (int j = std::max(0, i - 64); j < i; ++j) - if (ggml_graph_node(graph, j) == chain[0]) { - begin = j; - break; + return describe(nodes, count, op) && supported(op, cuda_device_info(device)) ? count : 0; + } + 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) { + auto* quant = ggml_graph_node(graph, i); + int group; + if (!sd_ggml_convrot_describe(quant, &group) || group != 256) + continue; + 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) { + chain[0] = weight->src[0]; + chain[1] = weight; + chain[2] = add->src[1]; + chain[3] = add; + chain[4] = value; + chain[5] = quant; + count = 6; + } + } else if (value->op == GGML_OP_MUL) { + for (int src = 0; src < 2; ++src) + if (value->src[src]->op == GGML_OP_UNARY) { + chain[0] = value->src[src]; + chain[1] = value; + chain[2] = quant; + count = 3; + break; + } } - if (begin < 0) - continue; - // Move the owning producer with its in-place consumers so alias checks retain their storage relationship. - for (int part = 1; part < count; ++part) { - int next = -1; - for (int j = begin + part; j <= i; ++j) - if (ggml_graph_node(graph, j) == chain[part]) { - next = j; + Operation op; + if (!count || !describe(chain, count, op) || !supported(op, cuda_device_info(device))) + continue; + int begin = -1; + for (int j = std::max(0, i - 64); j < i; ++j) + if (ggml_graph_node(graph, j) == chain[0]) { + begin = j; break; } - if (next < 0 || (next > begin + part && !ggml_ext_graph_delay(graph, begin, part, next))) - break; - begin = next - part; + if (begin < 0) + continue; + // Move the owning producer with its in-place consumers so alias checks retain their storage relationship. + for (int part = 1; part < count; ++part) { + int next = -1; + for (int j = begin + part; j <= i; ++j) + if (ggml_graph_node(graph, j) == chain[part]) { + next = j; + break; + } + if (next < 0 || (next > begin + part && !ggml_ext_graph_delay(graph, begin, part, next))) + break; + begin = next - part; + } } } - } - struct Plan { - DeviceInfo device; - Operation op; - bool configured = false; - }; - 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 (!describe(nodes, count, plan->op) || !supported(plan->op, plan->device)) - return GGML_STATUS_FAILED; - *workspace = 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); - } - template - static cudaError_t launch(Plan& p, const Operation& op, int8_t* output, cudaStream_t stream) { - constexpr bool warp = Scaled || RMS; - constexpr int extra = Normalize ? 384 : RMS ? 256 - : 128; - if (!p.configured) { - auto status = cudaFuncSetAttribute(quantize, - cudaFuncAttributeMaxDynamicSharedMemorySize, int(p.device.shared - extra)); - if (status != cudaSuccess) - return status; - p.configured = true; - } - const int threads = Normalize ? 1024 : RMS ? (op.k <= 1024 ? 128 : op.k <= 6144 ? 256 - : op.k <= 8192 ? 512 - : 1024) - : Swiglu && !Sigmoid && op.k >= 16384 ? 512 - : 256; - const int64_t padded = (op.rows + 3) / 4 * 4; - auto data = [](const ggml_tensor* t) { return t ? static_cast(t->data) : nullptr; }; - 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 cudaGetLastError(); - } - 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 = plan.op; - 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); - bool alias = false; - for (const auto* input : {op.input, op.up, op.modulation, op.weight, op.shift}) - alias |= overlaps(dst, input); - auto* output = static_cast(alias ? ggml_backend_buffer_get_base(execution->workspace) : dst->data); - const int64_t padded = (op.rows + 3) / 4 * 4; - if (padded > op.rows && cudaMemsetAsync(output + op.k * op.rows, 0, size_t(op.k) * (padded - op.rows), stream) != cudaSuccess) - return GGML_STATUS_FAILED; - cudaError_t error; - switch (op.kind) { - case SCALE: - error = launch(plan, op, output, stream); - break; - case NORM: - error = launch(plan, op, output, stream); - break; - case RMS_MOD: - error = launch(plan, op, output, stream); - break; - case SWIGLU: - case SILU: - error = launch(plan, op, output, stream); - break; - case SIGMOID: - error = launch(plan, op, output, stream); - break; - default: + struct Plan { + ggml_ext_gpu_device_info device; + Operation op; + bool configured = false; + }; + 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 = cuda_device_info(ggml_backend_get_device(backend)); + if (!describe(nodes, count, plan->op) || !supported(plan->op, plan->device)) return GGML_STATUS_FAILED; + *workspace = 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); + } + template + static CUDAError launch(Plan& p, const Operation& op, int8_t* output, CUDAStream stream) { + constexpr bool warp = Scaled || RMS; + constexpr int extra = Normalize ? 384 : RMS ? 256 + : 128; + if (!p.configured) { + auto status = cuda_set_shared_memory(quantize, + int(p.device.shared_memory - extra)); + if (status != cuda_success) + return status; + p.configured = true; + } + const int threads = Normalize ? 1024 : RMS ? (op.k <= 1024 ? 128 : op.k <= 6144 ? 256 + : op.k <= 8192 ? 512 + : 1024) + : Swiglu && !Sigmoid && op.k >= 16384 ? 512 + : 256; + const int64_t padded = (op.rows + 3) / 4 * 4; + auto data = [](const ggml_tensor* t) { return t ? static_cast(t->data) : nullptr; }; + 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(); + } + 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 = plan.op; + auto* dst = nodes[count - 1]; + ggml_ext_gpu_context context; + if (!cuda_context(execution->backend, &context) || context.device != plan.device.device) + return GGML_STATUS_FAILED; + auto stream = static_cast(context.stream); + bool alias = false; + for (const auto* input : {op.input, op.up, op.modulation, op.weight, op.shift}) + alias |= overlaps(dst, input); + auto* output = static_cast(alias ? ggml_backend_buffer_get_base(execution->workspace) : dst->data); + const int64_t padded = (op.rows + 3) / 4 * 4; + if (padded > op.rows && cuda_memset_async(output + op.k * op.rows, 0, size_t(op.k) * (padded - op.rows), stream) != cuda_success) + return GGML_STATUS_FAILED; + CUDAError error; + switch (op.kind) { + case SCALE: + error = launch(plan, op, output, stream); + break; + case NORM: + error = launch(plan, op, output, stream); + break; + case RMS_MOD: + error = launch(plan, op, output, stream); + break; + case SWIGLU: + case SILU: + error = launch(plan, op, output, stream); + break; + case SIGMOID: + error = launch(plan, op, output, stream); + break; + default: + return GGML_STATUS_FAILED; + } + if (error != cuda_success || (alias && cuda_copy_async(dst->data, output, size_t(op.k) * padded + op.rows * sizeof(float), stream) != cuda_success)) + return GGML_STATUS_FAILED; + return GGML_STATUS_SUCCESS; } - if (error != cudaSuccess || (alias && cudaMemcpyAsync(dst->data, output, size_t(op.k) * padded + op.rows * sizeof(float), cudaMemcpyDeviceToDevice, stream) != cudaSuccess)) - return GGML_STATUS_FAILED; - return GGML_STATUS_SUCCESS; } -} -void sd_ggml_register_cuda_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}) { - 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) + ".cuda.v1"; - desc.name = name.c_str(); - desc.op = op; - desc.priority = 225; - desc.optimize = op == GGML_OP_SCALE ? sd_convrot::optimize : nullptr; - desc.match = sd_convrot::match; - desc.prepare = sd_convrot::prepare; - desc.execute = sd_convrot::execute; - desc.destroy = sd_convrot::destroy; - GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE)); + 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}) { + 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.priority = 225; + desc.optimize = op == GGML_OP_SCALE ? sd_convrot::optimize : nullptr; + desc.match = sd_convrot::match; + desc.prepare = sd_convrot::prepare; + desc.execute = sd_convrot::execute; + desc.destroy = sd_convrot::destroy; + 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/convrot_kernel.h b/src/core/cuda/convrot_kernel.h index 970bf88b..748c08e7 100644 --- a/src/core/cuda/convrot_kernel.h +++ b/src/core/cuda/convrot_kernel.h @@ -1,249 +1,251 @@ #ifndef __SD_CORE_CUDA_CONVROT_KERNEL_H__ #define __SD_CORE_CUDA_CONVROT_KERNEL_H__ -#include #include -namespace sd_convrot { - static __device__ __forceinline__ float warp_reduce_sum(float x) { +#include "runtime.h" +namespace SD_CUDA_NAMESPACE { + namespace sd_convrot { + static __device__ __forceinline__ float warp_reduce_sum(float x) { #pragma unroll - for (int offset = 16; offset > 0; offset >>= 1) - x += __shfl_xor_sync(0xffffffff, x, offset, 32); - return x; - } - static __device__ __forceinline__ float2 warp_reduce_sum(float2 x) { -#pragma unroll - for (int offset = 16; offset > 0; offset >>= 1) { - x.x += __shfl_xor_sync(0xffffffff, x.x, offset, 32); - x.y += __shfl_xor_sync(0xffffffff, x.y, offset, 32); + for (int offset = 16; offset > 0; offset >>= 1) + x += cuda_shuffle_xor(x, offset, 32); + return x; } - return x; - } - template - static __device__ __forceinline__ float warp_reduce_max(float x) { + static __device__ __forceinline__ float2 warp_reduce_sum(float2 x) { #pragma unroll - for (int offset = Width / 2; offset > 0; offset >>= 1) - x = fmaxf(x, __shfl_xor_sync(0xffffffff, x, offset, Width)); - return x; - } - static __device__ float2 block_sum(float2 x, float2* shared) { - x = warp_reduce_sum(x); - if (threadIdx.x % 32 == 0) - shared[threadIdx.x / 32] = x; - __syncthreads(); - return warp_reduce_sum(shared[threadIdx.x % 32]); - } - static __device__ __forceinline__ float convrot_h4_cuda(float a, float b, float c, float d, int index) { - a = index == 3 ? -a : a; - b = index == 2 ? -b : b; - c = index == 1 ? -c : c; - d = index == 0 ? -d : d; - return __fadd_rn(__fadd_rn(__fadd_rn(a, b), c), d); - } - - template - __global__ void quantize( - const T* src, - int8_t* dst, - float* scales, - int64_t k, - int64_t rows, - int64_t dst_stride, - const float* up, - int64_t src_stride, - int64_t up_stride, - const float* activation_scales = nullptr, - const float* weight_scales = nullptr, - const float* modulation = nullptr, - float eps = 0.0f, - float input_scale = 1.0f, - float input_bias = 0.0f, - const float* rms_weight = nullptr, - const float* rms_shift = nullptr) { - constexpr int group_size = 256; - - const int64_t row = blockIdx.x; - const int tid = threadIdx.x; - if (row >= rows) { - return; - } - - extern __shared__ float values[]; - __shared__ float maxima[32]; - const T* src_row = src + row * src_stride; - int8_t* dst_row = dst + row * dst_stride; - - [[maybe_unused]] float rms_scale = 1.0f; - if constexpr (rms_modulation) { - __shared__ float sums[32]; - // Preserve the unfused RMSNorm's 1024-thread reduction order with fewer physical warps. - for (int virtual_tid = tid; virtual_tid < 1024; virtual_tid += blockDim.x) { - float partial = 0.0f; - for (int64_t i = virtual_tid; i < k; i += 1024) { - const float value = src_row[i]; - partial += value * value; - } - partial = warp_reduce_sum(partial); - if (tid % 32 == 0) { - sums[virtual_tid / 32] = partial; - } + for (int offset = 16; offset > 0; offset >>= 1) { + x.x += cuda_shuffle_xor(x.x, offset, 32); + x.y += cuda_shuffle_xor(x.y, offset, 32); } + return x; + } + template + static __device__ __forceinline__ float warp_reduce_max(float x) { +#pragma unroll + for (int offset = Width / 2; offset > 0; offset >>= 1) + x = fmaxf(x, cuda_shuffle_xor(x, offset, Width)); + return x; + } + static __device__ float2 block_sum(float2 x, float2* shared) { + x = warp_reduce_sum(x); + if (threadIdx.x % 32 == 0) + shared[threadIdx.x / 32] = x; __syncthreads(); - const float sum = warp_reduce_sum(sums[tid % 32]); - rms_scale = rsqrtf(sum / k + eps); + return warp_reduce_sum(shared[threadIdx.x % 32]); + } + static __device__ __forceinline__ float convrot_h4_cuda(float a, float b, float c, float d, int index) { + a = index == 3 ? -a : a; + b = index == 2 ? -b : b; + c = index == 1 ? -c : c; + d = index == 0 ? -d : d; + return __fadd_rn(__fadd_rn(__fadd_rn(a, b), c), d); } - if constexpr (swiglu || warp_rotation) { - const int lane = tid % 32; - for (int64_t group = tid / 32; group < k / group_size; group += blockDim.x / 32) { - float v[8]; -#pragma unroll - for (int j = 0; j < 8; ++j) { - const int64_t i = group * group_size + j * 32 + lane; - if constexpr (!swiglu) { - 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 (scaled) { - value = __fmaf_rn(value, input_scale, input_bias); - } - v[j] = value * (1.0f / 16.0f); - } else { - float gate_value, up_value; - if constexpr (dequantize) { - gate_value = (float)((const int32_t*)src_row)[i] * activation_scales[row]; - up_value = (float)((const int32_t*)up)[row * up_stride + i] * activation_scales[row]; - gate_value *= weight_scales[i]; - up_value *= weight_scales[k + i]; - } else { - gate_value = src_row[i]; - up_value = up[row * up_stride + i]; - } - const float activated = sigmoid_gate ? 1.0f / (1.0f + expf(-gate_value)) : (gate_value / (1.0f + expf(-gate_value))); - v[j] = (activated * up_value) * (1.0f / 16.0f); + template + __global__ void quantize( + const T* src, + int8_t* dst, + float* scales, + int64_t k, + int64_t rows, + int64_t dst_stride, + const float* up, + int64_t src_stride, + int64_t up_stride, + const float* activation_scales = nullptr, + const float* weight_scales = nullptr, + const float* modulation = nullptr, + float eps = 0.0f, + float input_scale = 1.0f, + float input_bias = 0.0f, + const float* rms_weight = nullptr, + const float* rms_shift = nullptr) { + constexpr int group_size = 256; + + const int64_t row = blockIdx.x; + const int tid = threadIdx.x; + if (row >= rows) { + return; + } + + extern __shared__ float values[]; + __shared__ float maxima[32]; + const T* src_row = src + row * src_stride; + int8_t* dst_row = dst + row * dst_stride; + + [[maybe_unused]] float rms_scale = 1.0f; + if constexpr (rms_modulation) { + __shared__ float sums[32]; + // Preserve the unfused RMSNorm's 1024-thread reduction order with fewer physical warps. + for (int virtual_tid = tid; virtual_tid < 1024; virtual_tid += blockDim.x) { + float partial = 0.0f; + for (int64_t i = virtual_tid; i < k; i += 1024) { + const float value = src_row[i]; + partial += value * value; + } + partial = warp_reduce_sum(partial); + if (tid % 32 == 0) { + sums[virtual_tid / 32] = partial; } } - // One logical warp owns H256, keeping every butterfly in registers. -#pragma unroll - for (int stride = 1; stride <= 4; stride *= 4) { - const int base = lane & ~(3 * stride); + __syncthreads(); + const float sum = warp_reduce_sum(sums[tid % 32]); + rms_scale = rsqrtf(sum / k + eps); + } + + if constexpr (swiglu || warp_rotation) { + const int lane = tid % 32; + for (int64_t group = tid / 32; group < k / group_size; group += blockDim.x / 32) { + float v[8]; #pragma unroll for (int j = 0; j < 8; ++j) { - const float a = __shfl_sync(0xffffffff, v[j], base, 32); - const float b = __shfl_sync(0xffffffff, v[j], base + stride, 32); - const float c = __shfl_sync(0xffffffff, v[j], base + 2 * stride, 32); - const float d = __shfl_sync(0xffffffff, v[j], base + 3 * stride, 32); - v[j] = convrot_h4_cuda(a, b, c, d, (lane / stride) % 4); + const int64_t i = group * group_size + j * 32 + lane; + if constexpr (!swiglu) { + 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 (scaled) { + value = __fmaf_rn(value, input_scale, input_bias); + } + v[j] = value * (1.0f / 16.0f); + } else { + float gate_value, up_value; + if constexpr (dequantize) { + gate_value = (float)((const int32_t*)src_row)[i] * activation_scales[row]; + up_value = (float)((const int32_t*)up)[row * up_stride + i] * activation_scales[row]; + gate_value *= weight_scales[i]; + up_value *= weight_scales[k + i]; + } else { + gate_value = src_row[i]; + up_value = up[row * up_stride + i]; + } + const float activated = sigmoid_gate ? 1.0f / (1.0f + expf(-gate_value)) : (gate_value / (1.0f + expf(-gate_value))); + v[j] = (activated * up_value) * (1.0f / 16.0f); + } + } + // One logical warp owns H256, keeping every butterfly in registers. +#pragma unroll + for (int stride = 1; stride <= 4; stride *= 4) { + const int base = lane & ~(3 * stride); +#pragma unroll + for (int j = 0; j < 8; ++j) { + const float a = cuda_shuffle(v[j], base, 32); + const float b = cuda_shuffle(v[j], base + stride, 32); + const float c = cuda_shuffle(v[j], base + 2 * stride, 32); + const float d = cuda_shuffle(v[j], base + 3 * stride, 32); + v[j] = convrot_h4_cuda(a, b, c, d, (lane / stride) % 4); + } + } +#pragma unroll + for (int j = 0; j < 8; j += 2) { + const float a = cuda_shuffle(v[j], lane % 16, 32); + const float b = cuda_shuffle(v[j], lane % 16 + 16, 32); + const float c = cuda_shuffle(v[j + 1], lane % 16, 32); + const float d = cuda_shuffle(v[j + 1], lane % 16 + 16, 32); + v[j] = convrot_h4_cuda(a, b, c, d, lane / 16); + v[j + 1] = convrot_h4_cuda(a, b, c, d, 2 + lane / 16); + } +#pragma unroll + for (int j = 0; j < 2; ++j) { + const float a = v[j], b = v[j + 2], c = v[j + 4], d = v[j + 6]; +#pragma unroll + for (int t = 0; t < 4; ++t) { + values[group * group_size + (j + 2 * t) * 32 + lane] = convrot_h4_cuda(a, b, c, d, t); + } } } -#pragma unroll - for (int j = 0; j < 8; j += 2) { - const float a = __shfl_sync(0xffffffff, v[j], lane % 16, 32); - const float b = __shfl_sync(0xffffffff, v[j], lane % 16 + 16, 32); - const float c = __shfl_sync(0xffffffff, v[j + 1], lane % 16, 32); - const float d = __shfl_sync(0xffffffff, v[j + 1], lane % 16 + 16, 32); - v[j] = convrot_h4_cuda(a, b, c, d, lane / 16); - v[j + 1] = convrot_h4_cuda(a, b, c, d, 2 + lane / 16); - } -#pragma unroll - for (int j = 0; j < 2; ++j) { - const float a = v[j], b = v[j + 2], c = v[j + 4], d = v[j + 6]; -#pragma unroll - for (int t = 0; t < 4; ++t) { - values[group * group_size + (j + 2 * t) * 32 + lane] = convrot_h4_cuda(a, b, c, d, t); - } - } - } - __syncthreads(); - } else { - constexpr int group_threads = group_size / 4; - constexpr int groups_per_wave = 1024 / group_threads; - if constexpr (normalize) { - __shared__ float2 sums[32]; - float2 mean_var = make_float2(0.0f, 0.0f); - for (int64_t i = tid; i < k; i += 1024) { - const float value = src_row[i]; - mean_var.x += value; - mean_var.y += value * value; - } - mean_var = block_sum(mean_var, sums); - const float mean = mean_var.x / k; - const float variance = mean_var.y / k - mean * mean; - const float inv_std = rsqrtf(variance + eps); - for (int64_t i = tid; i < k; i += 1024) { - const float normalized = __fmul_rn(__fsub_rn(src_row[i], mean), inv_std); - values[i] = __fmul_rn(__fmul_rn(normalized, modulation[i]), 1.0f / 16.0f); - } + __syncthreads(); } else { - for (int64_t i = tid; i < k; i += blockDim.x) { - const float value = scaled ? __fmaf_rn(src_row[i], input_scale, input_bias) : src_row[i]; - values[i] = value * (1.0f / 16.0f); + constexpr int group_threads = group_size / 4; + constexpr int groups_per_wave = 1024 / group_threads; + if constexpr (normalize) { + __shared__ float2 sums[32]; + float2 mean_var = make_float2(0.0f, 0.0f); + for (int64_t i = tid; i < k; i += 1024) { + const float value = src_row[i]; + mean_var.x += value; + mean_var.y += value * value; + } + mean_var = block_sum(mean_var, sums); + const float mean = mean_var.x / k; + const float variance = mean_var.y / k - mean * mean; + const float inv_std = rsqrtf(variance + eps); + for (int64_t i = tid; i < k; i += 1024) { + const float normalized = __fmul_rn(__fsub_rn(src_row[i], mean), inv_std); + values[i] = __fmul_rn(__fmul_rn(normalized, modulation[i]), 1.0f / 16.0f); + } + } else { + for (int64_t i = tid; i < k; i += blockDim.x) { + const float value = scaled ? __fmaf_rn(src_row[i], input_scale, input_bias) : src_row[i]; + values[i] = value * (1.0f / 16.0f); + } + } + __syncthreads(); + + const int group_lane = tid % group_threads; + const int wave_group = tid / group_threads; + const int64_t groups = k / group_size; + for (int64_t group_base = 0; group_base < groups; group_base += groups_per_wave) { + const int64_t group = group_base + wave_group; +#pragma unroll + for (int stride = 1; stride < group_size; stride *= 4) { + if (group < groups) { + float* group_values = values + group * group_size; + const int base = (group_lane / stride) * 4 * stride + group_lane % stride; + const int i0 = base; + const int i1 = i0 + stride; + const int i2 = i1 + stride; + const int i3 = i2 + stride; + const float a = group_values[i0]; + const float b = group_values[i1]; + const float c = group_values[i2]; + const float d = group_values[i3]; + group_values[i0] = a + b + c - d; + group_values[i1] = a + b - c + d; + group_values[i2] = a - b + c + d; + group_values[i3] = -a + b + c + d; + } + __syncthreads(); + } + } + } + + float local_max = 0.0f; + for (int64_t i = tid; i < k; i += blockDim.x) { + local_max = fmaxf(local_max, fabsf(values[i])); + } + local_max = warp_reduce_max<32>(local_max); + const int lane = tid & 31; + const int warp = tid >> 5; + if (lane == 0) { + maxima[warp] = local_max; + } + __syncthreads(); + + if (warp == 0) { + local_max = lane < blockDim.x / 32 ? maxima[lane] : 0.0f; + local_max = warp_reduce_max<32>(local_max); + if (lane == 0) { + maxima[0] = local_max / 127.0f; + if (scales != nullptr) { + scales[row] = maxima[0]; + } else { + *((float*)(dst_row + k)) = maxima[0]; + } } } __syncthreads(); - const int group_lane = tid % group_threads; - const int wave_group = tid / group_threads; - const int64_t groups = k / group_size; - for (int64_t group_base = 0; group_base < groups; group_base += groups_per_wave) { - const int64_t group = group_base + wave_group; -#pragma unroll - for (int stride = 1; stride < group_size; stride *= 4) { - if (group < groups) { - float* group_values = values + group * group_size; - const int base = (group_lane / stride) * 4 * stride + group_lane % stride; - const int i0 = base; - const int i1 = i0 + stride; - const int i2 = i1 + stride; - const int i3 = i2 + stride; - const float a = group_values[i0]; - const float b = group_values[i1]; - const float c = group_values[i2]; - const float d = group_values[i3]; - group_values[i0] = a + b + c - d; - group_values[i1] = a + b - c + d; - group_values[i2] = a - b + c + d; - group_values[i3] = -a + b + c + d; - } - __syncthreads(); - } + const float row_scale = maxima[0]; + const float inv_scale = row_scale == 0.0f ? 0.0f : 1.0f / row_scale; + for (int64_t i = tid; i < k; i += blockDim.x) { + int value = row_scale == 0.0f ? 0 : __float2int_rn(values[i] * inv_scale); + value = max(-127, min(127, value)); + dst_row[i] = (int8_t)value; } } - float local_max = 0.0f; - for (int64_t i = tid; i < k; i += blockDim.x) { - local_max = fmaxf(local_max, fabsf(values[i])); - } - local_max = warp_reduce_max<32>(local_max); - const int lane = tid & 31; - const int warp = tid >> 5; - if (lane == 0) { - maxima[warp] = local_max; - } - __syncthreads(); - - if (warp == 0) { - local_max = lane < blockDim.x / 32 ? maxima[lane] : 0.0f; - local_max = warp_reduce_max<32>(local_max); - if (lane == 0) { - maxima[0] = local_max / 127.0f; - if (scales != nullptr) { - scales[row] = maxima[0]; - } else { - *((float*)(dst_row + k)) = maxima[0]; - } - } - } - __syncthreads(); - - const float row_scale = maxima[0]; - const float inv_scale = row_scale == 0.0f ? 0.0f : 1.0f / row_scale; - for (int64_t i = tid; i < k; i += blockDim.x) { - int value = row_scale == 0.0f ? 0 : __float2int_rn(values[i] * inv_scale); - value = max(-127, min(127, value)); - dst_row[i] = (int8_t)value; - } } - } #endif diff --git a/src/core/cuda/cute_attention.cu b/src/core/cuda/cute_attention.cu index db7924db..80ab09dd 100644 --- a/src/core/cuda/cute_attention.cu +++ b/src/core/cuda/cute_attention.cu @@ -588,7 +588,7 @@ static int attention_match(ggml_backend_dev_t device, const ggml_cgraph* graph, 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)) || + !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)) return 0; offset = 3; @@ -608,7 +608,7 @@ static void attention_optimize(ggml_backend_dev_t device, ggml_cgraph* graph, vo continue; 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))) { + 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; @@ -687,7 +687,7 @@ static ggml_status attention_execute(const ggml_ext_execution_context* execution 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) + 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) { inputs[i] = static_cast(src->data); diff --git a/src/core/cuda/i8_projection.cu b/src/core/cuda/i8_projection.cu index cc2d9d3d..4a782bf6 100644 --- a/src/core/cuda/i8_projection.cu +++ b/src/core/cuda/i8_projection.cu @@ -208,7 +208,7 @@ static bool projection_supports(ggml_backend_dev_t device, const ggml_tensor* ma 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])) + !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)); diff --git a/src/core/cuda/int8.cu b/src/core/cuda/int8.cu index b0abd860..8931300a 100644 --- a/src/core/cuda/int8.cu +++ b/src/core/cuda/int8.cu @@ -1,745 +1,751 @@ #include #include "convrot_kernel.h" #include "core/int8_ops.h" -#include "i8_gemm.h" +#ifdef SD_CUDA_I8_CUTLASS +#include "core/cuda/i8_gemm.h" +#else +struct SDI8CudaGemmPlan; +#endif +#include "runtime.h" +#if defined(SD_CUDA_VENDOR_HIP) +#include +#elif !defined(SD_CUDA_VENDOR_MUSA) #include -#include +#endif #include #include #include #include -namespace sd_i8_cuda { - template - static __device__ __forceinline__ float warp_reduce_max(float value) { +namespace SD_CUDA_NAMESPACE { + namespace sd_i8_cuda { + template + static __device__ __forceinline__ float warp_reduce_max(float value) { #pragma unroll - for (int offset = Width / 2; offset > 0; offset >>= 1) - value = fmaxf(value, __shfl_xor_sync(0xffffffff, value, offset, Width)); - return value; - } - static __global__ void quantize_rowwise_i8_cuda(const float* src, int8_t* dst, float* scales, int64_t k, int64_t rows) { - const int64_t row = blockIdx.x; - const int tid = threadIdx.x; - if (row >= rows) { - return; + for (int offset = Width / 2; offset > 0; offset >>= 1) + value = fmaxf(value, cuda_shuffle_xor(value, offset, Width)); + return value; } - - __shared__ float maxima[256]; - const float* src_row = src + row * k; - int8_t* dst_row = dst + row * k; - float local_max = 0.0f; - for (int64_t i = tid; i < k; i += blockDim.x) { - local_max = fmaxf(local_max, fabsf(src_row[i])); - } - maxima[tid] = local_max; - __syncthreads(); - - for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) { - if (tid < stride) { - maxima[tid] = fmaxf(maxima[tid], maxima[tid + stride]); + static __global__ void quantize_rowwise_i8_cuda(const float* src, int8_t* dst, float* scales, int64_t k, int64_t rows) { + const int64_t row = blockIdx.x; + const int tid = threadIdx.x; + if (row >= rows) { + return; } - __syncthreads(); - } - const float scale = maxima[0] / 127.0f; - if (tid == 0) { - scales[row] = scale; - } - if (scale == 0.0f) { + __shared__ float maxima[256]; + const float* src_row = src + row * k; + int8_t* dst_row = dst + row * k; + float local_max = 0.0f; for (int64_t i = tid; i < k; i += blockDim.x) { - dst_row[i] = 0; + local_max = fmaxf(local_max, fabsf(src_row[i])); } - return; - } + maxima[tid] = local_max; + __syncthreads(); - const float inv_scale = 1.0f / scale; - for (int64_t i = tid; i < k; i += blockDim.x) { - int value = __float2int_rn(src_row[i] * inv_scale); - value = max(-127, min(127, value)); - dst_row[i] = (int8_t)value; - } - } - static __device__ __forceinline__ float convrot_h4_cuda(float a, float b, float c, float d, int index) { - a = index == 3 ? -a : a; - b = index == 2 ? -b : b; - c = index == 1 ? -c : c; - d = index == 0 ? -d : d; - return __fadd_rn(__fadd_rn(__fadd_rn(a, b), c), d); - } - static __global__ void quantize_convrot(const float* src, int8_t* dst, float* scales, int64_t k, int64_t rows, int64_t dst_stride) { - constexpr int group_size = 256; - const int64_t row = blockIdx.x; - const int tid = threadIdx.x; - if (row >= rows) - return; - extern __shared__ float values[]; - __shared__ float maxima[32]; - const float* src_row = src + row * k; - int8_t* dst_row = dst + row * dst_stride; - { - const int lane = tid % 32; - for (int64_t group = tid / 32; group < k / group_size; group += blockDim.x / 32) { - float v[8]; -#pragma unroll - for (int j = 0; j < 8; ++j) { - const int64_t i = group * group_size + j * 32 + lane; - const float value = src_row[i]; - v[j] = value * (1.0f / 16.0f); + for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) { + if (tid < stride) { + maxima[tid] = fmaxf(maxima[tid], maxima[tid + stride]); } - // One logical warp owns H256, keeping every butterfly in registers. -#pragma unroll - for (int stride = 1; stride <= 4; stride *= 4) { - const int base = lane & ~(3 * stride); + __syncthreads(); + } + + const float scale = maxima[0] / 127.0f; + if (tid == 0) { + scales[row] = scale; + } + if (scale == 0.0f) { + for (int64_t i = tid; i < k; i += blockDim.x) { + dst_row[i] = 0; + } + return; + } + + const float inv_scale = 1.0f / scale; + for (int64_t i = tid; i < k; i += blockDim.x) { + int value = __float2int_rn(src_row[i] * inv_scale); + value = max(-127, min(127, value)); + dst_row[i] = (int8_t)value; + } + } + static __device__ __forceinline__ float convrot_h4_cuda(float a, float b, float c, float d, int index) { + a = index == 3 ? -a : a; + b = index == 2 ? -b : b; + c = index == 1 ? -c : c; + d = index == 0 ? -d : d; + return __fadd_rn(__fadd_rn(__fadd_rn(a, b), c), d); + } + static __global__ void quantize_convrot(const float* src, int8_t* dst, float* scales, int64_t k, int64_t rows, int64_t dst_stride) { + constexpr int group_size = 256; + const int64_t row = blockIdx.x; + const int tid = threadIdx.x; + if (row >= rows) + return; + extern __shared__ float values[]; + __shared__ float maxima[32]; + const float* src_row = src + row * k; + int8_t* dst_row = dst + row * dst_stride; + { + const int lane = tid % 32; + for (int64_t group = tid / 32; group < k / group_size; group += blockDim.x / 32) { + float v[8]; #pragma unroll for (int j = 0; j < 8; ++j) { - const float a = __shfl_sync(0xffffffff, v[j], base, 32); - const float b = __shfl_sync(0xffffffff, v[j], base + stride, 32); - const float c = __shfl_sync(0xffffffff, v[j], base + 2 * stride, 32); - const float d = __shfl_sync(0xffffffff, v[j], base + 3 * stride, 32); - v[j] = convrot_h4_cuda(a, b, c, d, (lane / stride) % 4); + const int64_t i = group * group_size + j * 32 + lane; + const float value = src_row[i]; + v[j] = value * (1.0f / 16.0f); } - } + // One logical warp owns H256, keeping every butterfly in registers. #pragma unroll - for (int j = 0; j < 8; j += 2) { - const float a = __shfl_sync(0xffffffff, v[j], lane % 16, 32); - const float b = __shfl_sync(0xffffffff, v[j], lane % 16 + 16, 32); - const float c = __shfl_sync(0xffffffff, v[j + 1], lane % 16, 32); - const float d = __shfl_sync(0xffffffff, v[j + 1], lane % 16 + 16, 32); - v[j] = convrot_h4_cuda(a, b, c, d, lane / 16); - v[j + 1] = convrot_h4_cuda(a, b, c, d, 2 + lane / 16); - } + for (int stride = 1; stride <= 4; stride *= 4) { + const int base = lane & ~(3 * stride); #pragma unroll - for (int j = 0; j < 2; ++j) { - const float a = v[j], b = v[j + 2], c = v[j + 4], d = v[j + 6]; + for (int j = 0; j < 8; ++j) { + const float a = cuda_shuffle(v[j], base, 32); + const float b = cuda_shuffle(v[j], base + stride, 32); + const float c = cuda_shuffle(v[j], base + 2 * stride, 32); + const float d = cuda_shuffle(v[j], base + 3 * stride, 32); + v[j] = convrot_h4_cuda(a, b, c, d, (lane / stride) % 4); + } + } +#pragma unroll + for (int j = 0; j < 8; j += 2) { + const float a = cuda_shuffle(v[j], lane % 16, 32); + const float b = cuda_shuffle(v[j], lane % 16 + 16, 32); + const float c = cuda_shuffle(v[j + 1], lane % 16, 32); + const float d = cuda_shuffle(v[j + 1], lane % 16 + 16, 32); + v[j] = convrot_h4_cuda(a, b, c, d, lane / 16); + v[j + 1] = convrot_h4_cuda(a, b, c, d, 2 + lane / 16); + } +#pragma unroll + for (int j = 0; j < 2; ++j) { + const float a = v[j], b = v[j + 2], c = v[j + 4], d = v[j + 6]; #pragma unroll - for (int t = 0; t < 4; ++t) { - values[group * group_size + (j + 2 * t) * 32 + lane] = convrot_h4_cuda(a, b, c, d, t); + for (int t = 0; t < 4; ++t) { + values[group * group_size + (j + 2 * t) * 32 + lane] = convrot_h4_cuda(a, b, c, d, t); + } } } + __syncthreads(); } - __syncthreads(); - } - float local_max = 0.0f; - for (int64_t i = tid; i < k; i += blockDim.x) { - local_max = fmaxf(local_max, fabsf(values[i])); - } - local_max = warp_reduce_max<32>(local_max); - const int lane = tid & 31; - const int warp = tid >> 5; - if (lane == 0) { - maxima[warp] = local_max; - } - __syncthreads(); - - if (warp == 0) { - local_max = lane < blockDim.x / 32 ? maxima[lane] : 0.0f; - local_max = warp_reduce_max<32>(local_max); + float local_max = 0.0f; + for (int64_t i = tid; i < k; i += blockDim.x) { + local_max = fmaxf(local_max, fabsf(values[i])); + } + local_max = warp_reduce_max<32>(local_max); + const int lane = tid & 31; + const int warp = tid >> 5; if (lane == 0) { - maxima[0] = local_max / 127.0f; - if (scales != nullptr) { - scales[row] = maxima[0]; - } else { - *((float*)(dst_row + k)) = maxima[0]; + maxima[warp] = local_max; + } + __syncthreads(); + + if (warp == 0) { + local_max = lane < blockDim.x / 32 ? maxima[lane] : 0.0f; + local_max = warp_reduce_max<32>(local_max); + if (lane == 0) { + maxima[0] = local_max / 127.0f; + if (scales != nullptr) { + scales[row] = maxima[0]; + } else { + *((float*)(dst_row + k)) = maxima[0]; + } } } - } - __syncthreads(); + __syncthreads(); - const float row_scale = maxima[0]; - const float inv_scale = row_scale == 0.0f ? 0.0f : 1.0f / row_scale; - for (int64_t i = tid; i < k; i += blockDim.x) { - int value = row_scale == 0.0f ? 0 : __float2int_rn(values[i] * inv_scale); - value = max(-127, min(127, value)); - dst_row[i] = (int8_t)value; - } - } - - static __global__ void convrot_group_amax_i8_cuda( - const float* src, - float* partial_max, - int64_t k, - int64_t groups, - int64_t total_groups) { - const int64_t index = blockIdx.x; - const int tid = threadIdx.x; - if (index >= total_groups) { - return; + const float row_scale = maxima[0]; + const float inv_scale = row_scale == 0.0f ? 0.0f : 1.0f / row_scale; + for (int64_t i = tid; i < k; i += blockDim.x) { + int value = row_scale == 0.0f ? 0 : __float2int_rn(values[i] * inv_scale); + value = max(-127, min(127, value)); + dst_row[i] = (int8_t)value; + } } - __shared__ float values[256]; - __shared__ float maxima[2]; - const int64_t row = index / groups; - const int64_t group = index - row * groups; - const float* input = src + row * k + group * 256; - values[tid] = input[tid] * (1.0f / 16.0f); - values[tid + 64] = input[tid + 64] * (1.0f / 16.0f); - values[tid + 128] = input[tid + 128] * (1.0f / 16.0f); - values[tid + 192] = input[tid + 192] * (1.0f / 16.0f); - __syncthreads(); + static __global__ void convrot_group_amax_i8_cuda( + const float* src, + float* partial_max, + int64_t k, + int64_t groups, + int64_t total_groups) { + const int64_t index = blockIdx.x; + const int tid = threadIdx.x; + if (index >= total_groups) { + return; + } + + __shared__ float values[256]; + __shared__ float maxima[2]; + const int64_t row = index / groups; + const int64_t group = index - row * groups; + const float* input = src + row * k + group * 256; + values[tid] = input[tid] * (1.0f / 16.0f); + values[tid + 64] = input[tid + 64] * (1.0f / 16.0f); + values[tid + 128] = input[tid + 128] * (1.0f / 16.0f); + values[tid + 192] = input[tid + 192] * (1.0f / 16.0f); + __syncthreads(); #pragma unroll - for (int stride = 1; stride < 256; stride *= 4) { - const int base = (tid / stride) * 4 * stride + tid % stride; - const int i0 = base; - const int i1 = i0 + stride; - const int i2 = i1 + stride; - const int i3 = i2 + stride; - const float a = values[i0]; - const float b = values[i1]; - const float c = values[i2]; - const float d = values[i3]; - values[i0] = a + b + c - d; - values[i1] = a + b - c + d; - values[i2] = a - b + c + d; - values[i3] = -a + b + c + d; - __syncthreads(); - } + for (int stride = 1; stride < 256; stride *= 4) { + const int base = (tid / stride) * 4 * stride + tid % stride; + const int i0 = base; + const int i1 = i0 + stride; + const int i2 = i1 + stride; + const int i3 = i2 + stride; + const float a = values[i0]; + const float b = values[i1]; + const float c = values[i2]; + const float d = values[i3]; + values[i0] = a + b + c - d; + values[i1] = a + b - c + d; + values[i2] = a - b + c + d; + values[i3] = -a + b + c + d; + __syncthreads(); + } - float value = fmaxf(fabsf(values[tid]), fabsf(values[tid + 64])); - value = fmaxf(value, fabsf(values[tid + 128])); - value = fmaxf(value, fabsf(values[tid + 192])); - value = warp_reduce_max<32>(value); - if ((tid & 31) == 0) { - maxima[tid >> 5] = value; - } - __syncthreads(); - if (tid == 0) { - partial_max[index] = fmaxf(maxima[0], maxima[1]); - } - } - - static __global__ void reduce_convrot_row_amax_i8_cuda( - const float* partial_max, - float* scales, - int64_t groups, - int64_t rows) { - const int64_t row = blockIdx.x; - const int tid = threadIdx.x; - if (row >= rows) { - return; - } - - __shared__ float maxima[256]; - float value = 0.0f; - for (int64_t group = tid; group < groups; group += blockDim.x) { - value = fmaxf(value, partial_max[row * groups + group]); - } - maxima[tid] = value; - __syncthreads(); - for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) { - if (tid < stride) { - maxima[tid] = fmaxf(maxima[tid], maxima[tid + stride]); + float value = fmaxf(fabsf(values[tid]), fabsf(values[tid + 64])); + value = fmaxf(value, fabsf(values[tid + 128])); + value = fmaxf(value, fabsf(values[tid + 192])); + value = warp_reduce_max<32>(value); + if ((tid & 31) == 0) { + maxima[tid >> 5] = value; } __syncthreads(); - } - if (tid == 0) { - scales[row] = maxima[0] / 127.0f; - } - } - - static __global__ void convrot_group_quantize_i8_cuda( - const float* src, - int8_t* dst, - const float* scales, - int64_t k, - int64_t groups, - int64_t total_groups, - int64_t dst_stride) { - const int64_t index = blockIdx.x; - const int tid = threadIdx.x; - if (index >= total_groups) { - return; + if (tid == 0) { + partial_max[index] = fmaxf(maxima[0], maxima[1]); + } } - __shared__ float values[256]; - const int64_t row = index / groups; - const int64_t group = index - row * groups; - const float* input = src + row * k + group * 256; - int8_t* output = dst + row * dst_stride + group * 256; - values[tid] = input[tid] * (1.0f / 16.0f); - values[tid + 64] = input[tid + 64] * (1.0f / 16.0f); - values[tid + 128] = input[tid + 128] * (1.0f / 16.0f); - values[tid + 192] = input[tid + 192] * (1.0f / 16.0f); - __syncthreads(); + static __global__ void reduce_convrot_row_amax_i8_cuda( + const float* partial_max, + float* scales, + int64_t groups, + int64_t rows) { + const int64_t row = blockIdx.x; + const int tid = threadIdx.x; + if (row >= rows) { + return; + } -#pragma unroll - for (int stride = 1; stride < 256; stride *= 4) { - const int base = (tid / stride) * 4 * stride + tid % stride; - const int i0 = base; - const int i1 = i0 + stride; - const int i2 = i1 + stride; - const int i3 = i2 + stride; - const float a = values[i0]; - const float b = values[i1]; - const float c = values[i2]; - const float d = values[i3]; - values[i0] = a + b + c - d; - values[i1] = a + b - c + d; - values[i2] = a - b + c + d; - values[i3] = -a + b + c + d; + __shared__ float maxima[256]; + float value = 0.0f; + for (int64_t group = tid; group < groups; group += blockDim.x) { + value = fmaxf(value, partial_max[row * groups + group]); + } + maxima[tid] = value; __syncthreads(); + for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) { + if (tid < stride) { + maxima[tid] = fmaxf(maxima[tid], maxima[tid + stride]); + } + __syncthreads(); + } + if (tid == 0) { + scales[row] = maxima[0] / 127.0f; + } } - const float row_scale = scales[row]; - const float inv_scale = row_scale == 0.0f ? 0.0f : 1.0f / row_scale; - const int indices[4] = {tid, tid + 64, tid + 128, tid + 192}; + static __global__ void convrot_group_quantize_i8_cuda( + const float* src, + int8_t* dst, + const float* scales, + int64_t k, + int64_t groups, + int64_t total_groups, + int64_t dst_stride) { + const int64_t index = blockIdx.x; + const int tid = threadIdx.x; + if (index >= total_groups) { + return; + } + + __shared__ float values[256]; + const int64_t row = index / groups; + const int64_t group = index - row * groups; + const float* input = src + row * k + group * 256; + int8_t* output = dst + row * dst_stride + group * 256; + values[tid] = input[tid] * (1.0f / 16.0f); + values[tid + 64] = input[tid + 64] * (1.0f / 16.0f); + values[tid + 128] = input[tid + 128] * (1.0f / 16.0f); + values[tid + 192] = input[tid + 192] * (1.0f / 16.0f); + __syncthreads(); + #pragma unroll - for (int i = 0; i < 4; ++i) { - int value = row_scale == 0.0f ? 0 : __float2int_rn(values[indices[i]] * inv_scale); - value = max(-127, min(127, value)); - output[indices[i]] = (int8_t)value; - } - if (group == 0 && tid == 0 && dst_stride != k) { - *((float*)(dst + row * dst_stride + k)) = row_scale; - } - } + for (int stride = 1; stride < 256; stride *= 4) { + const int base = (tid / stride) * 4 * stride + tid % stride; + const int i0 = base; + const int i1 = i0 + stride; + const int i2 = i1 + stride; + const int i3 = i2 + stride; + const float a = values[i0]; + const float b = values[i1]; + const float c = values[i2]; + const float d = values[i3]; + values[i0] = a + b + c - d; + values[i1] = a + b - c + d; + values[i2] = a - b + c + d; + values[i3] = -a + b + c + d; + __syncthreads(); + } - static __global__ void dequantize_i32_rows_cuda( - const int32_t* src, - const float* scales, - const float* weight_scales, - const float* bias, - float* dst, - int64_t n, - int64_t rows) { - const int64_t output = (int64_t)blockIdx.x * blockDim.x + threadIdx.x; - const int64_t row = blockIdx.y + int64_t(blockIdx.z) * 65535; - if (output < n && row < rows) { - const int64_t index = row * n + output; - const float activation_scale = scales[row]; - float value = (float)src[index] * activation_scale; + const float row_scale = scales[row]; + const float inv_scale = row_scale == 0.0f ? 0.0f : 1.0f / row_scale; + const int indices[4] = {tid, tid + 64, tid + 128, tid + 192}; +#pragma unroll + for (int i = 0; i < 4; ++i) { + int value = row_scale == 0.0f ? 0 : __float2int_rn(values[indices[i]] * inv_scale); + value = max(-127, min(127, value)); + output[indices[i]] = (int8_t)value; + } + if (group == 0 && tid == 0 && dst_stride != k) { + *((float*)(dst + row * dst_stride + k)) = row_scale; + } + } + + static __global__ void dequantize_i32_rows_cuda( + const int32_t* src, + const float* scales, + const float* weight_scales, + const float* bias, + float* dst, + int64_t n, + int64_t rows) { + const int64_t output = (int64_t)blockIdx.x * blockDim.x + threadIdx.x; + const int64_t row = blockIdx.y + int64_t(blockIdx.z) * 65535; + if (output < n && row < rows) { + const int64_t index = row * n + output; + const float activation_scale = scales[row]; + float value = (float)src[index] * activation_scale; + if (weight_scales != nullptr) { + value *= weight_scales[output]; + } + if (bias != nullptr) { + value += bias[output]; + } + dst[index] = value; + } + } + + static __global__ void dequantize_i32_rows_vec4_cuda( + const int4* src, + const float* scales, + const float4* weight_scales, + const float4* bias, + float4* dst, + int64_t n4, + int64_t rows) { + const int64_t output = (int64_t)blockIdx.x * blockDim.x + threadIdx.x; + const int64_t row = blockIdx.y + int64_t(blockIdx.z) * 65535; + if (output >= n4 || row >= rows) { + return; + } + const int64_t index = row * n4 + output; + const int4 v = src[index]; + const float scale = scales[row]; + float4 value = make_float4(v.x * scale, v.y * scale, v.z * scale, v.w * scale); if (weight_scales != nullptr) { - value *= weight_scales[output]; + const float4 w = weight_scales[output]; + value.x *= w.x; + value.y *= w.y; + value.z *= w.z; + value.w *= w.w; } if (bias != nullptr) { - value += bias[output]; + const float4 b = bias[output]; + value.x += b.x; + value.y += b.y; + value.z += b.z; + value.w += b.w; } dst[index] = value; } - } - static __global__ void dequantize_i32_rows_vec4_cuda( - const int4* src, - const float* scales, - const float4* weight_scales, - const float4* bias, - float4* dst, - int64_t n4, - int64_t rows) { - const int64_t output = (int64_t)blockIdx.x * blockDim.x + threadIdx.x; - const int64_t row = blockIdx.y + int64_t(blockIdx.z) * 65535; - if (output >= n4 || row >= rows) { - return; - } - const int64_t index = row * n4 + output; - const int4 v = src[index]; - const float scale = scales[row]; - float4 value = make_float4(v.x * scale, v.y * scale, v.z * scale, v.w * scale); - if (weight_scales != nullptr) { - const float4 w = weight_scales[output]; - value.x *= w.x; - value.y *= w.y; - value.z *= w.z; - value.w *= w.w; - } - if (bias != nullptr) { - const float4 b = bias[output]; - value.x += b.x; - value.y += b.y; - value.z += b.z; - value.w += b.w; - } - dst[index] = value; - } - - struct DeviceInfo { - int ordinal = -1; - int major = 0; - int minor = 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.minor, size_t(props.sharedMemPerBlockOptin)}; - cache.emplace(device, info); - return info; - } - static bool aligned(const ggml_tensor* tensor, size_t alignment = 4) { - if (!tensor) - return true; - size_t offset = 0; - while (tensor->view_src) { - offset += tensor->view_offs; - tensor = tensor->view_src; - } - return (reinterpret_cast(tensor->data) + offset) % alignment == 0; - } - static bool supports(ggml_backend_dev_t device, const ggml_tensor* node, void*) { - SDI8Operation op; - if (!sd_ggml_i8_describe(node, &op) || !aligned(node)) - return false; - const int64_t k = op.weight->ne[0], n = op.weight->ne[1], rows = ggml_nrows(node); - if (k > INT_MAX || n > INT_MAX || rows > INT_MAX - 3 || k % 4 || n % 4 || (op.input->type == GGML_TYPE_F32 && op.group_size && (op.group_size != 256 || rows > INT_MAX / (k / 256)))) - return false; - for (auto* input : node->src) - if (!aligned(input)) - return false; - const auto info = device_info(device); - return (info.major > 7 || (info.major == 7 && info.minor >= 5)) && - ggml_backend_reg_get_proc_address(ggml_backend_dev_backend_reg(device), "ggml_ext_cuda_blas_v1"); - } - static bool convrot_supports(ggml_backend_dev_t device, const ggml_tensor* node, void*) { - int group; - if (!sd_ggml_convrot_describe(node, &group) || group != 256 || !aligned(node) || !aligned(node->src[0])) - return false; - const int64_t k = node->ne[0], rows = ggml_nrows(node->src[0]); - return k <= INT_MAX && rows <= INT_MAX && rows <= INT_MAX / (k / 256) && device_info(device).ordinal >= 0; - } - enum Kind { MAT, - QUANTIZE, - SWIGLU, - SWIGLU_QUANTIZE, - CONCAT }; - 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 int match(ggml_backend_dev_t device, const ggml_cgraph* graph, int first, void*) { - auto* g = const_cast(graph); - if (fusion_disabled() || first + 1 >= ggml_graph_n_nodes(g)) - return 0; - const auto* mat = ggml_graph_node(g, first); - SDI8Operation op; - if (!supports(device, mat, nullptr) || !sd_ggml_i8_describe(mat, &op) || !op.weight_scale || op.bias) - return 0; - const auto* next = ggml_graph_node(g, first + 1); - if (next->op == GGML_OP_CONCAT && next->src[1] == mat && next->src[0] != mat && next->op_params[0] == 1 && - next->type == GGML_TYPE_F32 && mat->ne[2] == 1 && mat->ne[3] == 1 && - next->ne[2] == 1 && next->ne[3] == 1 && ggml_is_contiguous(next) && ggml_is_contiguous(next->src[0])) - return 2; - if (next->op != GGML_OP_GLU || ggml_get_glu_op(next) != GGML_GLU_OP_SWIGLU || next->op_params[1] != 0 || - next->src[0] != mat || next->src[1] || next->type != GGML_TYPE_F32 || !ggml_is_contiguous(next) || - ggml_nelements(next) > int64_t(INT_MAX) * 256) - return 0; - if (first + 2 < ggml_graph_n_nodes(g)) { - const auto* quant = ggml_graph_node(g, first + 2); - if (quant->src[0] == next && convrot_supports(device, quant, nullptr) && - size_t(next->ne[0]) * sizeof(float) + 128 <= device_info(device).shared) - return 3; - } - return 2; - } - struct Plan { - DeviceInfo device; - Kind kind; - int64_t k, n, rows, padded; - size_t quant, scales, accum, partial, output, gemm_offset; - SDI8CudaGemmPlan* gemm = nullptr; - bool fast_swiglu = false; - }; - static size_t reserve(size_t& total, size_t bytes) { - const size_t offset = total; - total += (bytes + 255) / 256 * 256; - return offset; - } - static ggml_status prepare(ggml_backend_t backend, ggml_tensor* const* nodes, int count, void** state, size_t* workspace, void*) { - const auto* dst = nodes[count - 1]; - SDI8Operation op; - const bool mat = sd_ggml_i8_describe(nodes[0], &op); - auto* p = new Plan{}; - *state = p; - auto device = ggml_backend_get_device(backend); - p->device = device_info(device); - p->kind = !mat ? QUANTIZE : count == 3 ? SWIGLU_QUANTIZE - : count == 1 ? MAT - : dst->op == GGML_OP_CONCAT ? CONCAT - : SWIGLU; - p->k = mat ? op.weight->ne[0] : dst->ne[0]; - p->n = mat ? nodes[0]->ne[0] : 0; - p->rows = mat ? ggml_nrows(nodes[0]) : ggml_nrows(dst->src[0]); - p->padded = (p->rows + 3) / 4 * 4; - size_t total = 0; - if (mat) { - p->quant = reserve(total, op.input->type == GGML_TYPE_F32 || (!op.packed && p->rows != p->padded) ? size_t(p->k) * p->padded : 0); - p->scales = reserve(total, op.input->type == GGML_TYPE_F32 ? size_t(p->rows) * sizeof(float) : 0); -#ifdef SD_CUDA_I8_CUTLASS - size_t gemm_bytes = 0; - // Custom GEMM previously used cuBLAS; retain that path when I32 accumulation can overflow. - if (op.weight_scale && !op.bias && (nodes[0]->op == GGML_OP_MUL_MAT || p->k <= INT_MAX / 16384) && - aligned(op.weight, 16) && aligned(op.weight_scale, 16) && aligned(dst, 16) && - (op.input->type == GGML_TYPE_F32 || (aligned(op.input, 16) && aligned(op.input_scale, 16))) && - (p->kind != CONCAT || ggml_nbytes(dst->src[0]) % 16 == 0)) { - if (p->kind == SWIGLU || p->kind == SWIGLU_QUANTIZE) { - p->gemm = sd_cuda_i8_gemm_prepare(device, p->rows, p->n, p->k, true, &gemm_bytes); - p->fast_swiglu = p->gemm != nullptr; - } - if (!p->gemm) - p->gemm = sd_cuda_i8_gemm_prepare(device, p->rows, p->n, p->k, false, &gemm_bytes); + static bool aligned(const ggml_tensor* tensor, size_t alignment = 4) { + if (!tensor) + return true; + size_t offset = 0; + while (tensor->view_src) { + offset += tensor->view_offs; + tensor = tensor->view_src; } - p->gemm_offset = reserve(total, gemm_bytes); -#endif - const bool intermediate = !p->gemm || p->kind == SWIGLU_QUANTIZE || (p->kind == SWIGLU && !p->fast_swiglu); - p->accum = reserve(total, intermediate ? size_t(p->fast_swiglu ? p->n / 2 : p->n) * p->padded * sizeof(float) : 0); + return (reinterpret_cast(tensor->data) + offset) % alignment == 0; } - const bool large_rotation = (!mat || (op.input->type == GGML_TYPE_F32 && op.group_size == 256)) && size_t(p->k) * sizeof(float) + 128 > p->device.shared; - p->partial = reserve(total, large_rotation ? size_t(p->rows) * (p->k / 256) * sizeof(float) : 0); - p->output = reserve(total, p->gemm && p->kind == SWIGLU_QUANTIZE ? 0 : ggml_nbytes(dst)); - *workspace = total; - return GGML_STATUS_SUCCESS; - } - static void destroy(void* state, void*) { - auto* p = static_cast(state); + static bool supports(ggml_backend_dev_t device, const ggml_tensor* node, void*) { + SDI8Operation op; + if (!sd_ggml_i8_describe(node, &op) || !aligned(node)) + return false; + const int64_t k = op.weight->ne[0], n = op.weight->ne[1], rows = ggml_nrows(node); + if (k > INT_MAX || n > INT_MAX || rows > INT_MAX - 3 || k % 4 || n % 4 || (op.input->type == GGML_TYPE_F32 && op.group_size && (op.group_size != 256 || rows > INT_MAX / (k / 256)))) + return false; + for (auto* input : node->src) + if (!aligned(input)) + return false; + const auto info = cuda_device_info(device); + return info.int8_blas && + ggml_backend_reg_get_proc_address(ggml_backend_dev_backend_reg(device), "ggml_ext_gpu_blas_v1"); + } + static bool convrot_supports(ggml_backend_dev_t device, const ggml_tensor* node, void*) { + int group; + if (!sd_ggml_convrot_describe(node, &group) || group != 256 || !aligned(node) || !aligned(node->src[0])) + return false; + const int64_t k = node->ne[0], rows = ggml_nrows(node->src[0]); + return k <= INT_MAX && rows <= INT_MAX && rows <= INT_MAX / (k / 256) && cuda_device_info(device).device >= 0; + } + enum Kind { MAT, + QUANTIZE, + SWIGLU, + SWIGLU_QUANTIZE, + CONCAT }; + 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 int match(ggml_backend_dev_t device, const ggml_cgraph* graph, int first, void*) { + auto* g = const_cast(graph); + if (fusion_disabled() || first + 1 >= ggml_graph_n_nodes(g)) + return 0; + const auto* mat = ggml_graph_node(g, first); + SDI8Operation op; + if (!supports(device, mat, nullptr) || !sd_ggml_i8_describe(mat, &op) || !op.weight_scale || op.bias) + return 0; + const auto* next = ggml_graph_node(g, first + 1); + if (next->op == GGML_OP_CONCAT && next->src[1] == mat && next->src[0] != mat && next->op_params[0] == 1 && + next->type == GGML_TYPE_F32 && mat->ne[2] == 1 && mat->ne[3] == 1 && + next->ne[2] == 1 && next->ne[3] == 1 && ggml_is_contiguous(next) && ggml_is_contiguous(next->src[0])) + return 2; + if (next->op != GGML_OP_GLU || ggml_get_glu_op(next) != GGML_GLU_OP_SWIGLU || next->op_params[1] != 0 || + next->src[0] != mat || next->src[1] || next->type != GGML_TYPE_F32 || !ggml_is_contiguous(next) || + ggml_nelements(next) > int64_t(INT_MAX) * 256) + return 0; + if (first + 2 < ggml_graph_n_nodes(g)) { + const auto* quant = ggml_graph_node(g, first + 2); + if (quant->src[0] == next && convrot_supports(device, quant, nullptr) && + size_t(next->ne[0]) * sizeof(float) + 128 <= cuda_device_info(device).shared_memory) + return 3; + } + return 2; + } + struct Plan { + ggml_ext_gpu_device_info device; + Kind kind; + int64_t k, n, rows, padded; + size_t quant, scales, accum, partial, output, gemm_offset; + SDI8CudaGemmPlan* gemm = nullptr; + bool fast_swiglu = false; + }; + static size_t reserve(size_t& total, size_t bytes) { + const size_t offset = total; + total += (bytes + 255) / 256 * 256; + return offset; + } + static ggml_status prepare(ggml_backend_t backend, ggml_tensor* const* nodes, int count, void** state, size_t* workspace, void*) { + const auto* dst = nodes[count - 1]; + SDI8Operation op; + const bool mat = sd_ggml_i8_describe(nodes[0], &op); + auto* p = new Plan{}; + *state = p; + auto device = ggml_backend_get_device(backend); + p->device = cuda_device_info(device); + p->kind = !mat ? QUANTIZE : count == 3 ? SWIGLU_QUANTIZE + : count == 1 ? MAT + : dst->op == GGML_OP_CONCAT ? CONCAT + : SWIGLU; + p->k = mat ? op.weight->ne[0] : dst->ne[0]; + p->n = mat ? nodes[0]->ne[0] : 0; + p->rows = mat ? ggml_nrows(nodes[0]) : ggml_nrows(dst->src[0]); + p->padded = (p->rows + 3) / 4 * 4; + size_t total = 0; + if (mat) { + p->quant = reserve(total, op.input->type == GGML_TYPE_F32 || (!op.packed && p->rows != p->padded) ? size_t(p->k) * p->padded : 0); + p->scales = reserve(total, op.input->type == GGML_TYPE_F32 ? size_t(p->rows) * sizeof(float) : 0); #ifdef SD_CUDA_I8_CUTLASS - if (p) - sd_cuda_i8_gemm_free(p->gemm); + size_t gemm_bytes = 0; + // Custom GEMM previously used cuBLAS; retain that path when I32 accumulation can overflow. + if (op.weight_scale && !op.bias && (nodes[0]->op == GGML_OP_MUL_MAT || p->k <= INT_MAX / 16384) && + aligned(op.weight, 16) && aligned(op.weight_scale, 16) && aligned(dst, 16) && + (op.input->type == GGML_TYPE_F32 || (aligned(op.input, 16) && aligned(op.input_scale, 16))) && + (p->kind != CONCAT || ggml_nbytes(dst->src[0]) % 16 == 0)) { + if (p->kind == SWIGLU || p->kind == SWIGLU_QUANTIZE) { + p->gemm = sd_cuda_i8_gemm_prepare(device, p->rows, p->n, p->k, true, &gemm_bytes); + p->fast_swiglu = p->gemm != nullptr; + } + if (!p->gemm) + p->gemm = sd_cuda_i8_gemm_prepare(device, p->rows, p->n, p->k, false, &gemm_bytes); + } + p->gemm_offset = reserve(total, gemm_bytes); #endif - delete p; - } - 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 cudaError_t convrot_launch(const Plan& p, cudaStream_t stream, const float* src, int8_t* dst, float* scales, float* partial) { - if (p.padded > p.rows) { - auto error = cudaMemsetAsync(dst + p.k * p.rows, 0, size_t(p.k) * (p.padded - p.rows), stream); - if (error != cudaSuccess) - return error; + const bool intermediate = !p->gemm || p->kind == SWIGLU_QUANTIZE || (p->kind == SWIGLU && !p->fast_swiglu); + p->accum = reserve(total, intermediate ? size_t(p->fast_swiglu ? p->n / 2 : p->n) * p->padded * sizeof(float) : 0); + } + const bool large_rotation = (!mat || (op.input->type == GGML_TYPE_F32 && op.group_size == 256)) && size_t(p->k) * sizeof(float) + 128 > p->device.shared_memory; + p->partial = reserve(total, large_rotation ? size_t(p->rows) * (p->k / 256) * sizeof(float) : 0); + p->output = reserve(total, p->gemm && p->kind == SWIGLU_QUANTIZE ? 0 : ggml_nbytes(dst)); + *workspace = total; + return GGML_STATUS_SUCCESS; } - const size_t shared = size_t(p.k) * sizeof(float); - if (shared + 128 <= p.device.shared) { + static void destroy(void* state, void*) { + auto* p = static_cast(state); +#ifdef SD_CUDA_I8_CUTLASS + if (p) + sd_cuda_i8_gemm_free(p->gemm); +#endif + delete p; + } + 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 CUDAError convrot_launch(const Plan& p, CUDAStream stream, const float* src, int8_t* dst, float* scales, float* partial) { + if (p.padded > p.rows) { + auto error = cuda_memset_async(dst + p.k * p.rows, 0, size_t(p.k) * (p.padded - p.rows), stream); + if (error != cuda_success) + return error; + } + const size_t shared = size_t(p.k) * sizeof(float); + if (shared + 128 <= p.device.shared_memory) { + static std::mutex mutex; + static std::unordered_map configured; + { + std::lock_guard lock(mutex); + if (!configured[p.device.device]) { + auto error = cuda_set_shared_memory(quantize_convrot, int(p.device.shared_memory - 128)); + if (error != cuda_success) + return error; + configured[p.device.device] = true; + } + } + int threads = p.k <= 1024 ? 128 : p.k <= 8192 ? 256 + : 1024; + quantize_convrot<<>>(src, dst, scales, p.k, p.rows, p.k); + } else { + const int64_t groups = p.k / 256, total_groups = p.rows * groups; + convrot_group_amax_i8_cuda<<>>(src, partial, p.k, groups, total_groups); + reduce_convrot_row_amax_i8_cuda<<>>(partial, scales, groups, p.rows); + convrot_group_quantize_i8_cuda<<>>(src, dst, scales, p.k, groups, total_groups, p.k); + } + return cuda_last_error(); + } + template + static __global__ void swiglu(const float* src, const float* xs, const float* ws, float* dst, int64_t n, int64_t rows) { + const int64_t i = int64_t(blockIdx.x) * blockDim.x + threadIdx.x; + const int64_t half = n / 2; + if (i >= rows * half) + return; + const int64_t row = i / half, col = i % half; + float gate, up; + if constexpr (Dequantize) { + gate = float(reinterpret_cast(src)[row * n + col]) * xs[row]; + up = float(reinterpret_cast(src)[row * n + half + col]) * xs[row]; + gate *= ws[col]; + up *= ws[half + col]; + } else { + gate = src[row * n + col]; + up = src[row * n + half + col]; + } + dst[i] = __fmul_rn(gate / (1.f + expf(-gate)), up); + } + template + static CUDAError quantize_epilogue(const Plan& p, const float* src, const float* xs, const float* ws, int8_t* output, CUDAStream stream) { + const int64_t hidden = p.n / 2; + auto kernel = sd_convrot::quantize; static std::mutex mutex; - static std::unordered_map configured; + static std::unordered_map initialized; { std::lock_guard lock(mutex); - if (!configured[p.device.ordinal]) { - auto error = cudaFuncSetAttribute(quantize_convrot, cudaFuncAttributeMaxDynamicSharedMemorySize, int(p.device.shared - 128)); - if (error != cudaSuccess) + if (!initialized[p.device.device]) { + auto error = cuda_set_shared_memory(kernel, int(p.device.shared_memory - 128)); + if (error != cuda_success) return error; - configured[p.device.ordinal] = true; + initialized[p.device.device] = true; } } - int threads = p.k <= 1024 ? 128 : p.k <= 8192 ? 256 - : 1024; - quantize_convrot<<>>(src, dst, scales, p.k, p.rows, p.k); - } else { - const int64_t groups = p.k / 256, total_groups = p.rows * groups; - convrot_group_amax_i8_cuda<<>>(src, partial, p.k, groups, total_groups); - reduce_convrot_row_amax_i8_cuda<<>>(partial, scales, groups, p.rows); - convrot_group_quantize_i8_cuda<<>>(src, dst, scales, p.k, groups, total_groups, p.k); - } - return cudaGetLastError(); - } - template - static __global__ void swiglu(const float* src, const float* xs, const float* ws, float* dst, int64_t n, int64_t rows) { - const int64_t i = int64_t(blockIdx.x) * blockDim.x + threadIdx.x; - const int64_t half = n / 2; - if (i >= rows * half) - return; - const int64_t row = i / half, col = i % half; - float gate, up; - if constexpr (Dequantize) { - gate = float(reinterpret_cast(src)[row * n + col]) * xs[row]; - up = float(reinterpret_cast(src)[row * n + half + col]) * xs[row]; - gate *= ws[col]; - up *= ws[half + col]; - } else { - gate = src[row * n + col]; - up = src[row * n + half + col]; - } - dst[i] = __fmul_rn(gate / (1.f + expf(-gate)), up); - } - template - static cudaError_t quantize_epilogue(const Plan& p, const float* src, const float* xs, const float* ws, int8_t* output, cudaStream_t stream) { - const int64_t hidden = p.n / 2; - auto kernel = sd_convrot::quantize; - static std::mutex mutex; - static std::unordered_map initialized; - { - std::lock_guard lock(mutex); - if (!initialized[p.device.ordinal]) { - auto error = cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, int(p.device.shared - 128)); - if (error != cudaSuccess) + if (p.padded > p.rows) { + auto error = cuda_memset_async(output + hidden * p.rows, 0, size_t(hidden) * (p.padded - p.rows), stream); + if (error != cuda_success) return error; - initialized[p.device.ordinal] = true; } + const int threads = Rotated && !Swiglu && !Dequantize ? (hidden <= 1024 ? 128 : hidden <= 8192 ? 256 + : 1024) + : 256; + kernel<<>>(src, output, + reinterpret_cast(output + hidden * p.padded), hidden, p.rows, hidden, + Swiglu ? src + hidden : nullptr, Swiglu ? p.n : hidden, Swiglu ? p.n : hidden, xs, ws, nullptr, 0.f, 1.f, 0.f, nullptr, nullptr); + return cuda_last_error(); } - if (p.padded > p.rows) { - auto error = cudaMemsetAsync(output + hidden * p.rows, 0, size_t(hidden) * (p.padded - p.rows), stream); - if (error != cudaSuccess) - return error; - } - const int threads = Rotated && !Swiglu && !Dequantize ? (hidden <= 1024 ? 128 : hidden <= 8192 ? 256 - : 1024) - : 256; - kernel<<>>(src, output, - reinterpret_cast(output + hidden * p.padded), hidden, p.rows, hidden, - Swiglu ? src + hidden : nullptr, Swiglu ? p.n : hidden, Swiglu ? p.n : hidden, xs, ws, nullptr, 0.f, 1.f, 0.f, nullptr, nullptr); - return cudaGetLastError(); - } - static ggml_status execute(const ggml_ext_execution_context* execution, ggml_tensor* const* nodes, int count, void* state, void*) { - const auto& p = *static_cast(state); - 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 != p.device.ordinal) - return GGML_STATUS_FAILED; - auto stream = static_cast(context.stream); - auto* scratch = static_cast(ggml_backend_buffer_get_base(execution->workspace)); - auto* dst = nodes[count - 1]; - SDI8Operation op; - const bool mat = sd_ggml_i8_describe(nodes[0], &op); - bool alias = false; - for (int i = 0; i < count; ++i) - for (auto* input : nodes[i]->src) { - bool internal = false; - for (int j = 0; j < count; ++j) - internal |= input == nodes[j]; - if (!internal) - alias |= overlaps(dst, input); - } - // CUTLASS quantization epilogues read only the workspace intermediate after GEMM completes. - if (p.gemm && p.kind == SWIGLU_QUANTIZE) - alias = false; - void* output = alias ? scratch + p.output : dst->data; - cudaError_t error = cudaSuccess; - if (!mat) { - auto* quant = static_cast(output); - auto* scales = reinterpret_cast(quant + p.k * p.padded); - error = convrot_launch(p, stream, static_cast(dst->src[0]->data), quant, scales, reinterpret_cast(scratch + p.partial)); - } else { - auto* quant = reinterpret_cast(scratch + p.quant); - const float* scales; - if (op.input->type == GGML_TYPE_F32) { - auto* temporary_scales = reinterpret_cast(scratch + p.scales); - scales = temporary_scales; - if (op.group_size) - error = convrot_launch(p, stream, static_cast(op.input->data), quant, temporary_scales, reinterpret_cast(scratch + p.partial)); - else { - if (p.padded > p.rows) - error = cudaMemsetAsync(quant + p.k * p.rows, 0, size_t(p.k) * (p.padded - p.rows), stream); - if (error == cudaSuccess) - quantize_rowwise_i8_cuda<<>>(static_cast(op.input->data), quant, temporary_scales, p.k, p.rows); - } - } else { - auto* data = static_cast(op.input->data); - scales = op.packed ? reinterpret_cast(data + p.k * p.padded) : static_cast(op.input_scale->data); - if (op.packed || p.rows == p.padded) - quant = data; - else { - error = cudaMemcpyAsync(quant, data, size_t(p.k) * p.rows, cudaMemcpyDeviceToDevice, stream); - if (error == cudaSuccess) - error = cudaMemsetAsync(quant + p.k * p.rows, 0, size_t(p.k) * (p.padded - p.rows), stream); - } - } - if (error != cudaSuccess || cudaGetLastError() != cudaSuccess) + static ggml_status execute(const ggml_ext_execution_context* execution, ggml_tensor* const* nodes, int count, void* state, void*) { + const auto& p = *static_cast(state); + auto reg = ggml_backend_dev_backend_reg(ggml_backend_get_device(execution->backend)); + ggml_ext_gpu_context context; + if (!cuda_context(execution->backend, &context) || context.device != p.device.device) return GGML_STATUS_FAILED; - const auto* ws = op.weight_scale ? static_cast(op.weight_scale->data) : nullptr; - const auto* bias = op.bias ? static_cast(op.bias->data) : nullptr; - const size_t prefix = p.kind == CONCAT ? ggml_nbytes(dst->src[0]) : 0; - auto* gemm_output = reinterpret_cast(static_cast(output) + prefix); - auto* accum = reinterpret_cast(scratch + p.accum); - const bool intermediate = p.kind == SWIGLU_QUANTIZE || (p.kind == SWIGLU && !p.fast_swiglu); -#ifdef SD_CUDA_I8_CUTLASS - if (p.gemm) { - auto status = sd_cuda_i8_gemm_execute(p.gemm, quant, static_cast(op.weight->data), scales, ws, - intermediate ? reinterpret_cast(accum) : gemm_output, scratch + p.gemm_offset, stream); - if (status != GGML_STATUS_SUCCESS) - return status; - } else -#endif - { - auto blas_query = reinterpret_cast(ggml_backend_reg_get_proc_address(reg, "ggml_ext_cuda_blas_v1")); - auto handle = blas_query ? static_cast(blas_query(execution->backend)) : nullptr; - if (!handle) - return GGML_STATUS_FAILED; - const int32_t alpha = 1, beta = 0; - if (cublasGemmEx(handle, CUBLAS_OP_T, CUBLAS_OP_N, int(p.n), int(p.padded), int(p.k), &alpha, - op.weight->data, CUDA_R_8I, int(p.k), quant, CUDA_R_8I, int(p.k), &beta, accum, CUDA_R_32I, - int(p.n), CUBLAS_COMPUTE_32I, CUBLAS_GEMM_DEFAULT_TENSOR_OP) != CUBLAS_STATUS_SUCCESS) - return GGML_STATUS_FAILED; - } - const auto* values = reinterpret_cast(accum); - if (p.kind == SWIGLU_QUANTIZE) { - if (p.fast_swiglu) - error = quantize_epilogue(p, values, nullptr, nullptr, static_cast(output), stream); - else if (p.gemm) - error = quantize_epilogue(p, values, nullptr, nullptr, static_cast(output), stream); - else - error = quantize_epilogue(p, values, scales, ws, static_cast(output), stream); - } else if (p.kind == SWIGLU && !p.fast_swiglu) { - const unsigned blocks = unsigned((p.rows * (p.n / 2) + 255) / 256); - if (p.gemm) - swiglu<<>>(values, scales, ws, static_cast(output), p.n, p.rows); - else - swiglu<<>>(values, scales, ws, static_cast(output), p.n, p.rows); - error = cudaGetLastError(); - } else if (!p.gemm) { - const unsigned gy = unsigned(std::min(p.rows, 65535)), gz = unsigned((p.rows + 65534) / 65535); - if ((reinterpret_cast(ws) | reinterpret_cast(bias) | reinterpret_cast(gemm_output)) % 16 == 0) { - dequantize_i32_rows_vec4_cuda<<>>( - reinterpret_cast(accum), scales, reinterpret_cast(ws), reinterpret_cast(bias), reinterpret_cast(gemm_output), p.n / 4, p.rows); - } else { - dequantize_i32_rows_cuda<<>>(accum, scales, ws, bias, gemm_output, p.n, p.rows); + auto stream = static_cast(context.stream); + auto* scratch = static_cast(ggml_backend_buffer_get_base(execution->workspace)); + auto* dst = nodes[count - 1]; + SDI8Operation op; + const bool mat = sd_ggml_i8_describe(nodes[0], &op); + bool alias = false; + for (int i = 0; i < count; ++i) + for (auto* input : nodes[i]->src) { + bool internal = false; + for (int j = 0; j < count; ++j) + internal |= input == nodes[j]; + if (!internal) + alias |= overlaps(dst, input); } - error = cudaGetLastError(); + // CUTLASS quantization epilogues read only the workspace intermediate after GEMM completes. + if (p.gemm && p.kind == SWIGLU_QUANTIZE) + alias = false; + void* output = alias ? scratch + p.output : dst->data; + CUDAError error = cuda_success; + if (!mat) { + auto* quant = static_cast(output); + auto* scales = reinterpret_cast(quant + p.k * p.padded); + error = convrot_launch(p, stream, static_cast(dst->src[0]->data), quant, scales, reinterpret_cast(scratch + p.partial)); + } else { + auto* quant = reinterpret_cast(scratch + p.quant); + const float* scales; + if (op.input->type == GGML_TYPE_F32) { + auto* temporary_scales = reinterpret_cast(scratch + p.scales); + scales = temporary_scales; + if (op.group_size) + error = convrot_launch(p, stream, static_cast(op.input->data), quant, temporary_scales, reinterpret_cast(scratch + p.partial)); + else { + if (p.padded > p.rows) + error = cuda_memset_async(quant + p.k * p.rows, 0, size_t(p.k) * (p.padded - p.rows), stream); + if (error == cuda_success) + quantize_rowwise_i8_cuda<<>>(static_cast(op.input->data), quant, temporary_scales, p.k, p.rows); + } + } else { + auto* data = static_cast(op.input->data); + scales = op.packed ? reinterpret_cast(data + p.k * p.padded) : static_cast(op.input_scale->data); + if (op.packed || p.rows == p.padded) + quant = data; + else { + error = cuda_copy_async(quant, data, size_t(p.k) * p.rows, stream); + if (error == cuda_success) + error = cuda_memset_async(quant + p.k * p.rows, 0, size_t(p.k) * (p.padded - p.rows), stream); + } + } + if (error != cuda_success || cuda_last_error() != cuda_success) + return GGML_STATUS_FAILED; + const auto* ws = op.weight_scale ? static_cast(op.weight_scale->data) : nullptr; + const auto* bias = op.bias ? static_cast(op.bias->data) : nullptr; + const size_t prefix = p.kind == CONCAT ? ggml_nbytes(dst->src[0]) : 0; + auto* gemm_output = reinterpret_cast(static_cast(output) + prefix); + auto* accum = reinterpret_cast(scratch + p.accum); + const bool intermediate = p.kind == SWIGLU_QUANTIZE || (p.kind == SWIGLU && !p.fast_swiglu); +#ifdef SD_CUDA_I8_CUTLASS + if (p.gemm) { + auto status = sd_cuda_i8_gemm_execute(p.gemm, quant, static_cast(op.weight->data), scales, ws, + intermediate ? reinterpret_cast(accum) : gemm_output, scratch + p.gemm_offset, stream); + if (status != GGML_STATUS_SUCCESS) + return status; + } else +#endif + { + auto blas_query = reinterpret_cast(ggml_backend_reg_get_proc_address(reg, "ggml_ext_gpu_blas_v1")); + auto handle = blas_query ? blas_query(execution->backend) : nullptr; + if (!handle) + return GGML_STATUS_FAILED; +#if defined(SD_CUDA_VENDOR_MUSA) + return GGML_STATUS_FAILED; +#else + const int32_t alpha = 1, beta = 0; +#if defined(SD_CUDA_VENDOR_HIP) +#if HIP_VERSION >= 60500000 + constexpr auto compute_type = HIPBLAS_COMPUTE_32I; +#else + constexpr auto compute_type = HIPBLAS_R_32I; +#endif + const auto status = hipblasGemmEx(static_cast(handle), HIPBLAS_OP_T, HIPBLAS_OP_N, + int(p.n), int(p.padded), int(p.k), &alpha, + op.weight->data, HIPBLAS_R_8I, int(p.k), quant, HIPBLAS_R_8I, int(p.k), + &beta, accum, HIPBLAS_R_32I, int(p.n), compute_type, HIPBLAS_GEMM_DEFAULT); + if (status != HIPBLAS_STATUS_SUCCESS) +#else + const auto status = cublasGemmEx(static_cast(handle), CUBLAS_OP_T, CUBLAS_OP_N, + int(p.n), int(p.padded), int(p.k), &alpha, + op.weight->data, CUDA_R_8I, int(p.k), quant, CUDA_R_8I, int(p.k), + &beta, accum, CUDA_R_32I, int(p.n), CUBLAS_COMPUTE_32I, CUBLAS_GEMM_DEFAULT_TENSOR_OP); + if (status != CUBLAS_STATUS_SUCCESS) +#endif + return GGML_STATUS_FAILED; +#endif + } + const auto* values = reinterpret_cast(accum); + if (p.kind == SWIGLU_QUANTIZE) { + if (p.fast_swiglu) + error = quantize_epilogue(p, values, nullptr, nullptr, static_cast(output), stream); + else if (p.gemm) + error = quantize_epilogue(p, values, nullptr, nullptr, static_cast(output), stream); + else + error = quantize_epilogue(p, values, scales, ws, static_cast(output), stream); + } else if (p.kind == SWIGLU && !p.fast_swiglu) { + const unsigned blocks = unsigned((p.rows * (p.n / 2) + 255) / 256); + if (p.gemm) + swiglu<<>>(values, scales, ws, static_cast(output), p.n, p.rows); + else + swiglu<<>>(values, scales, ws, static_cast(output), p.n, p.rows); + error = cuda_last_error(); + } else if (!p.gemm) { + const unsigned gy = unsigned(std::min(p.rows, 65535)), gz = unsigned((p.rows + 65534) / 65535); + if ((reinterpret_cast(ws) | reinterpret_cast(bias) | reinterpret_cast(gemm_output)) % 16 == 0) { + dequantize_i32_rows_vec4_cuda<<>>( + reinterpret_cast(accum), scales, reinterpret_cast(ws), reinterpret_cast(bias), reinterpret_cast(gemm_output), p.n / 4, p.rows); + } else { + dequantize_i32_rows_cuda<<>>(accum, scales, ws, bias, gemm_output, p.n, p.rows); + } + error = cuda_last_error(); + } + if (error == cuda_success && prefix) + error = cuda_copy_async(output, dst->src[0]->data, prefix, stream); } - if (error == cudaSuccess && prefix) - error = cudaMemcpyAsync(output, dst->src[0]->data, prefix, cudaMemcpyDeviceToDevice, stream); + 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; } - if (error == cudaSuccess && alias) - error = cudaMemcpyAsync(dst->data, output, ggml_nbytes(dst), cudaMemcpyDeviceToDevice, stream); - return error == cudaSuccess ? GGML_STATUS_SUCCESS : GGML_STATUS_FAILED; + + } + + void register_i8(ggml_backend_dev_t device) { + ggml_ext_kernel_desc desc = {}; + desc.abi_version = GGML_EXT_ABI_VERSION; + desc.name = "sd.linear_i8." SD_CUDA_NAME ".v1"; + desc.op = GGML_OP_CUSTOM; + desc.custom_op = sd_ggml_i8_op(); + desc.priority = 100; + desc.supports = sd_i8_cuda::supports; + desc.prepare = sd_i8_cuda::prepare; + desc.execute = sd_i8_cuda::execute; + desc.destroy = sd_i8_cuda::destroy; + GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE)); + desc.name = "sd.linear_i8." SD_CUDA_NAME ".override.v1"; + desc.op = GGML_OP_MUL_MAT; + desc.custom_op = nullptr; + GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE)); + desc.name = "sd.linear_i8.epilogue." SD_CUDA_NAME ".v1"; + desc.priority = 240; + desc.supports = nullptr; + desc.match = sd_i8_cuda::match; + GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE)); + desc.name = "sd.linear_i8.epilogue." SD_CUDA_NAME ".custom.v1"; + desc.op = GGML_OP_CUSTOM; + desc.custom_op = sd_ggml_i8_op(); + GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE)); + desc.priority = 100; + desc.match = nullptr; + desc.name = "sd.quantize_i8_convrot." SD_CUDA_NAME ".v1"; + desc.custom_op = sd_ggml_convrot_op(); + desc.supports = sd_i8_cuda::convrot_supports; + GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE)); + desc.name = "sd.quantize_i8_convrot." SD_CUDA_NAME ".override.v1"; + desc.op = GGML_OP_QUANTIZE_I8_CONVROT; + desc.custom_op = nullptr; + GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE)); } } - -void sd_ggml_register_cuda_i8(ggml_backend_dev_t device) { - ggml_ext_kernel_desc desc = {}; - desc.abi_version = GGML_EXT_ABI_VERSION; - desc.name = "sd.linear_i8.cuda.v1"; - desc.op = GGML_OP_CUSTOM; - desc.custom_op = sd_ggml_i8_op(); - desc.priority = 100; - desc.supports = sd_i8_cuda::supports; - desc.prepare = sd_i8_cuda::prepare; - desc.execute = sd_i8_cuda::execute; - desc.destroy = sd_i8_cuda::destroy; - GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE)); - desc.name = "sd.linear_i8.cuda.override.v1"; - desc.op = GGML_OP_MUL_MAT; - desc.custom_op = nullptr; - GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE)); - desc.name = "sd.linear_i8.epilogue.cuda.v1"; - desc.priority = 240; - desc.supports = nullptr; - desc.match = sd_i8_cuda::match; - GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE)); - desc.name = "sd.linear_i8.epilogue.cuda.custom.v1"; - desc.op = GGML_OP_CUSTOM; - desc.custom_op = sd_ggml_i8_op(); - GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE)); - desc.priority = 100; - desc.match = nullptr; - desc.name = "sd.quantize_i8_convrot.cuda.v1"; - desc.custom_op = sd_ggml_convrot_op(); - desc.supports = sd_i8_cuda::convrot_supports; - GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE)); - desc.name = "sd.quantize_i8_convrot.cuda.override.v1"; - desc.op = GGML_OP_QUANTIZE_I8_CONVROT; - desc.custom_op = 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/registration.cu b/src/core/cuda/registration.cu new file mode 100644 index 00000000..601890c7 --- /dev/null +++ b/src/core/cuda/registration.cu @@ -0,0 +1,28 @@ +#include "runtime.h" + +#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 + +namespace SD_CUDA_NAMESPACE { + void register_rope(ggml_backend_dev_t device); + 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_extensions(ggml_backend_dev_t device) { + if (cuda_device_info(device).device < 0) + return; + register_affine(device); + register_i8(device); + register_convrot(device); + register_rope(device); +#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 + } +} diff --git a/src/core/cuda/registration.h b/src/core/cuda/registration.h new file mode 100644 index 00000000..f98c36fb --- /dev/null +++ b/src/core/cuda/registration.h @@ -0,0 +1,13 @@ +#ifndef __SD_CORE_CUDA_REGISTRATION_H__ +#define __SD_CORE_CUDA_REGISTRATION_H__ +#include "ggml-backend.h" +namespace sd_cuda { + void register_extensions(ggml_backend_dev_t device); +} +namespace sd_hip { + void register_extensions(ggml_backend_dev_t device); +} +namespace sd_musa { + void register_extensions(ggml_backend_dev_t device); +} +#endif diff --git a/src/core/cuda/rope.cu b/src/core/cuda/rope.cu index a0aa7f08..ffb7fa31 100644 --- a/src/core/cuda/rope.cu +++ b/src/core/cuda/rope.cu @@ -1,188 +1,188 @@ #include "core/cuda/rope.h" #include "core/ggml_extensions.h" -#include #include #include #include #include -static __device__ __forceinline__ float sd_warp_sum(float v) { - for (int mask = 16; mask > 0; mask >>= 1) - v = __fadd_rn(v, __shfl_xor_sync(0xffffffff, v, mask)); - return v; -} -template -static __global__ void rms_norm_mul_rope_apply_fused( - const float* x, - const float* weight, - const float* pe, - T* dst, - int64_t heads, - int64_t tokens, - int64_t rows, - int64_t nb1, - int64_t nb2, - int64_t nb3, - int rot_half, - float eps, - float output_scale) { - const int lane = threadIdx.x; - const int64_t row = int64_t(blockIdx.x) * blockDim.y + threadIdx.y; - if (row >= rows) { - return; +namespace SD_CUDA_NAMESPACE { + static __device__ __forceinline__ float sd_warp_sum(float v) { + for (int mask = 16; mask > 0; mask >>= 1) + v = __fadd_rn(v, cuda_shuffle_xor(v, mask)); + return v; } - const int64_t head = row % heads; - const int64_t token = (row / heads) % tokens; - const int64_t batch = row / (heads * tokens); - const float* x_row = x + head * nb1 + token * nb2 + batch * nb3; - float v[D / 32], sums[8] = {}; + template + static __global__ void rms_norm_mul_rope_apply_fused( + const float* x, + const float* weight, + const float* pe, + T* dst, + int64_t heads, + int64_t tokens, + int64_t rows, + int64_t nb1, + int64_t nb2, + int64_t nb3, + int rot_half, + float eps, + float output_scale) { + const int lane = threadIdx.x; + const int64_t row = int64_t(blockIdx.x) * blockDim.y + threadIdx.y; + if (row >= rows) { + return; + } + const int64_t head = row % heads; + const int64_t token = (row / heads) % tokens; + const int64_t batch = row / (heads * tokens); + const float* x_row = x + head * nb1 + token * nb2 + batch * nb3; + float v[D / 32], sums[8] = {}; #pragma unroll - for (int j = 0; j < D / 32; ++j) { - v[j] = x_row[j * 32 + lane]; - sums[j] = sd_warp_sum(__fmul_rn(v[j], v[j])); - } - // Preserve the 256-thread RMSNorm kernel's reduction tree, including inactive warps. - const float sum = __fadd_rn( - __fadd_rn(__fadd_rn(sums[0], sums[4]), __fadd_rn(sums[2], sums[6])), - __fadd_rn(__fadd_rn(sums[1], sums[5]), __fadd_rn(sums[3], sums[7]))); - const float scale = rsqrtf(sum / float(D) + eps); - T* out = dst + ((batch * heads + head) * tokens + token) * D; + for (int j = 0; j < D / 32; ++j) { + v[j] = x_row[j * 32 + lane]; + sums[j] = sd_warp_sum(__fmul_rn(v[j], v[j])); + } + // Preserve the 256-thread RMSNorm kernel's reduction tree, including inactive warps. + const float sum = __fadd_rn( + __fadd_rn(__fadd_rn(sums[0], sums[4]), __fadd_rn(sums[2], sums[6])), + __fadd_rn(__fadd_rn(sums[1], sums[5]), __fadd_rn(sums[3], sums[7]))); + const float scale = rsqrtf(sum / float(D) + eps); + T* out = dst + ((batch * heads + head) * tokens + token) * D; #pragma unroll - for (int j = 0; j < D / 32; ++j) { - const int col = j * 32 + lane; - const float value = __fmul_rn(__fmul_rn(scale, v[j]), weight[col]); - float rotated = value; - if constexpr (interleaved) { - const float other = __shfl_xor_sync(0xffffffff, value, 1, 32); - if (col < 2 * rot_half) { - const float* m = pe + 4 * (token * rot_half + col / 2); - const float4 matrix = make_float4(m[0], m[1], m[2], m[3]); - rotated = lane % 2 == 0 - ? __fadd_rn(__fmul_rn(value, matrix.x), __fmul_rn(other, matrix.y)) - : __fadd_rn(__fmul_rn(other, matrix.z), __fmul_rn(value, matrix.w)); + for (int j = 0; j < D / 32; ++j) { + const int col = j * 32 + lane; + const float value = __fmul_rn(__fmul_rn(scale, v[j]), weight[col]); + float rotated = value; + if constexpr (interleaved) { + const float other = cuda_shuffle_xor(value, 1, 32); + if (col < 2 * rot_half) { + const float* m = pe + 4 * (token * rot_half + col / 2); + const float4 matrix = make_float4(m[0], m[1], m[2], m[3]); + rotated = lane % 2 == 0 + ? __fadd_rn(__fmul_rn(value, matrix.x), __fmul_rn(other, matrix.y)) + : __fadd_rn(__fmul_rn(other, matrix.z), __fmul_rn(value, matrix.w)); + } + } else { + if (col < 2 * rot_half) { + const int partner = col < rot_half ? col + rot_half : col - rot_half; + const float other = __fmul_rn(__fmul_rn(scale, x_row[partner]), weight[partner]); + const float* m = pe + 4 * (token * rot_half + col % rot_half); + const float4 matrix = make_float4(m[0], m[1], m[2], m[3]); + rotated = col < rot_half + ? __fsub_rn(__fmul_rn(value, matrix.x), __fmul_rn(other, matrix.z)) + : __fadd_rn(__fmul_rn(other, matrix.z), __fmul_rn(value, matrix.x)); + } } - } else { - if (col < 2 * rot_half) { - const int partner = col < rot_half ? col + rot_half : col - rot_half; - const float other = __fmul_rn(__fmul_rn(scale, x_row[partner]), weight[partner]); - const float* m = pe + 4 * (token * rot_half + col % rot_half); - const float4 matrix = make_float4(m[0], m[1], m[2], m[3]); - rotated = col < rot_half - ? __fsub_rn(__fmul_rn(value, matrix.x), __fmul_rn(other, matrix.z)) - : __fadd_rn(__fmul_rn(other, matrix.z), __fmul_rn(value, matrix.x)); + if constexpr (std::is_same::value) { + out[col] = __float2half_rn(__fmul_rn(rotated, output_scale)); + } else { + out[col] = rotated; } } - if constexpr (std::is_same::value) { - out[col] = __float2half_rn(__fmul_rn(rotated, output_scale)); + } + + struct RopeApplyStrides { + int64_t x[4]; + int64_t pe[4]; + }; + + template + static __global__ void rope_apply_f32( + const float* x, + const float* pe, + float* dst, + int64_t half_dim, + int64_t rot_half, + int64_t heads, + int64_t tokens, + int64_t pairs, + RopeApplyStrides strides) { + const int64_t index = (int64_t)blockIdx.x * blockDim.x + threadIdx.x; + if (index >= pairs) { + return; + } + const int64_t j = index % half_dim; + const int64_t row = index / half_dim; + const int64_t token = row % tokens; + const int64_t head = (row / tokens) % heads; + const int64_t batch = row / (tokens * heads); + const int64_t i0 = interleaved || j >= rot_half ? 2 * j : j; + const int64_t i1 = interleaved || j >= rot_half ? 2 * j + 1 : j + rot_half; + const float* x_row = x + head * strides.x[1] + token * strides.x[2] + batch * strides.x[3]; + const float a = x_row[i0 * strides.x[0]]; + const float b = x_row[i1 * strides.x[0]]; + float* out = dst + row * half_dim * 2; + if (j >= rot_half) { + out[i0] = a; + out[i1] = b; + return; + } + const float* matrix = pe + j * strides.pe[2] + token * strides.pe[3]; + const float c = matrix[0]; + const float s = matrix[strides.pe[1]]; + // Match the separate mul/add rounding of the unfused graph. + if (interleaved) { + out[i0] = __fadd_rn(__fmul_rn(a, c), __fmul_rn(b, matrix[strides.pe[0]])); + out[i1] = __fadd_rn(__fmul_rn(a, s), __fmul_rn(b, matrix[strides.pe[0] + strides.pe[1]])); } else { - out[col] = rotated; + out[i0] = __fsub_rn(__fmul_rn(a, c), __fmul_rn(b, s)); + out[i1] = __fadd_rn(__fmul_rn(a, s), __fmul_rn(b, c)); } } -} -struct RopeApplyStrides { - int64_t x[4]; - int64_t pe[4]; -}; - -template -static __global__ void rope_apply_f32( - const float* x, - const float* pe, - float* dst, - int64_t half_dim, - int64_t rot_half, - int64_t heads, - int64_t tokens, - int64_t pairs, - RopeApplyStrides strides) { - const int64_t index = (int64_t)blockIdx.x * blockDim.x + threadIdx.x; - if (index >= pairs) { - return; + static bool rope_interleaved(const ggml_tensor* rope) { + int32_t value; + std::memcpy(&value, ggml_ext_op_get(rope) ? ggml_ext_op_attributes(rope) : rope->op_params, sizeof(value)); + return value != 0; } - const int64_t j = index % half_dim; - const int64_t row = index / half_dim; - const int64_t token = row % tokens; - const int64_t head = (row / tokens) % heads; - const int64_t batch = row / (tokens * heads); - const int64_t i0 = interleaved || j >= rot_half ? 2 * j : j; - const int64_t i1 = interleaved || j >= rot_half ? 2 * j + 1 : j + rot_half; - const float* x_row = x + head * strides.x[1] + token * strides.x[2] + batch * strides.x[3]; - const float a = x_row[i0 * strides.x[0]]; - const float b = x_row[i1 * strides.x[0]]; - float* out = dst + row * half_dim * 2; - if (j >= rot_half) { - out[i0] = a; - out[i1] = b; - return; + + static float param_f32(const ggml_tensor* node, int index) { + float value; + std::memcpy(&value, node->op_params + index, sizeof(value)); + return value; } - const float* matrix = pe + j * strides.pe[2] + token * strides.pe[3]; - const float c = matrix[0]; - const float s = matrix[strides.pe[1]]; - // Match the separate mul/add rounding of the unfused graph. - if (interleaved) { - out[i0] = __fadd_rn(__fmul_rn(a, c), __fmul_rn(b, matrix[strides.pe[0]])); - out[i1] = __fadd_rn(__fmul_rn(a, s), __fmul_rn(b, matrix[strides.pe[0] + strides.pe[1]])); - } else { - out[i0] = __fsub_rn(__fmul_rn(a, c), __fmul_rn(b, s)); - out[i1] = __fadd_rn(__fmul_rn(a, s), __fmul_rn(b, c)); + + static bool rope_is_node(const ggml_tensor* node) { + return node && (node->op == GGML_OP_ROPE_APPLY || ggml_ext_op_get(node) == sd_ggml_rope_op()); } -} -static bool rope_interleaved(const ggml_tensor* rope) { - int32_t value; - std::memcpy(&value, ggml_ext_op_get(rope) ? ggml_ext_op_attributes(rope) : rope->op_params, sizeof(value)); - return value != 0; -} - -static float param_f32(const ggml_tensor* node, int index) { - float value; - std::memcpy(&value, node->op_params + index, sizeof(value)); - return value; -} - -bool sd_cuda_rope_is_node(const ggml_tensor* node) { - return node && (node->op == GGML_OP_ROPE_APPLY || ggml_ext_op_get(node) == sd_ggml_rope_op()); -} - -static bool rope_supports(ggml_backend_dev_t, const ggml_tensor* rope, void*) { - const auto* x = rope->src[0]; - const auto* pe = rope->src[1]; - if (!x || !pe || x->type != GGML_TYPE_F32 || pe->type != GGML_TYPE_F32 || rope->type != GGML_TYPE_F32 || - x->ne[0] % 2 || pe->ne[0] != 2 || pe->ne[1] != 2 || pe->ne[2] <= 0 || pe->ne[2] > x->ne[0] / 2 || - pe->ne[3] != x->ne[2] || !ggml_is_contiguous(rope) || rope->ne[0] != x->ne[0] || - rope->ne[1] != x->ne[2] || rope->ne[2] != x->ne[1] * x->ne[3] || rope->ne[3] != 1 || - (ggml_nelements(x) / 2 + 255) / 256 > INT_MAX) { - return false; - } - for (int i = 0; i < 4; ++i) { - if (x->nb[i] % sizeof(float) || pe->nb[i] % sizeof(float)) + static bool rope_supports(ggml_backend_dev_t, const ggml_tensor* rope, void*) { + const auto* x = rope->src[0]; + const auto* pe = rope->src[1]; + if (!x || !pe || x->type != GGML_TYPE_F32 || pe->type != GGML_TYPE_F32 || rope->type != GGML_TYPE_F32 || + x->ne[0] % 2 || pe->ne[0] != 2 || pe->ne[1] != 2 || pe->ne[2] <= 0 || pe->ne[2] > x->ne[0] / 2 || + pe->ne[3] != x->ne[2] || !ggml_is_contiguous(rope) || rope->ne[0] != x->ne[0] || + rope->ne[1] != x->ne[2] || rope->ne[2] != x->ne[1] * x->ne[3] || rope->ne[3] != 1 || + (ggml_nelements(x) / 2 + 255) / 256 > INT_MAX) { return false; + } + for (int i = 0; i < 4; ++i) { + if (x->nb[i] % sizeof(float) || pe->nb[i] % sizeof(float)) + return false; + } + return true; } - return true; -} -bool sd_cuda_rope_fusion_supports(const ggml_tensor* norm, const ggml_tensor* mul, const ggml_tensor* rope) { - if (norm->op != GGML_OP_RMS_NORM || mul->op != GGML_OP_MUL || !sd_cuda_rope_is_node(rope) || - mul->src[0] != norm || rope->src[0] != mul || !rope_supports(nullptr, rope, nullptr)) - 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 && - 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] && - x->nb[0] == sizeof(float) && x->nb[1] % sizeof(float) == 0 && x->nb[2] % sizeof(float) == 0 && - x->nb[3] % sizeof(float) == 0 && ggml_is_contiguous(weight) && ggml_is_contiguous(rope->src[1]) && - (ggml_nrows(x) + 3) / 4 <= INT_MAX; -} + bool rope_fusion_supports(const ggml_tensor* norm, const ggml_tensor* mul, const ggml_tensor* rope) { + if (norm->op != GGML_OP_RMS_NORM || mul->op != GGML_OP_MUL || !rope_is_node(rope) || + mul->src[0] != norm || rope->src[0] != mul || !rope_supports(nullptr, rope, nullptr)) + 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 && + 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] && + x->nb[0] == sizeof(float) && x->nb[1] % sizeof(float) == 0 && x->nb[2] % sizeof(float) == 0 && + x->nb[3] % sizeof(float) == 0 && ggml_is_contiguous(weight) && ggml_is_contiguous(rope->src[1]) && + (ggml_nrows(x) + 3) / 4 <= INT_MAX; + } -template -static cudaError_t rms_rope_launch(cudaStream_t stream, ggml_tensor* const* nodes, T* output, float scale) { - const auto* x = nodes[0]->src[0]; - const auto* pe = nodes[2]->src[1]; - const int64_t rows = ggml_nrows(x); + template + static CUDAError rms_rope_launch(CUDAStream stream, ggml_tensor* const* nodes, T* output, float scale) { + const auto* x = nodes[0]->src[0]; + const auto* pe = nodes[2]->src[1]; + const int64_t rows = ggml_nrows(x); #define SD_ROPE(D) \ case D: \ rms_norm_mul_rope_apply_fused<<((rows + 3) / 4), dim3(32, 4), 0, stream>>>( \ @@ -190,151 +190,151 @@ static cudaError_t rms_rope_launch(cudaStream_t stream, ggml_tensor* const* node static_cast(pe->data), output, x->ne[1], x->ne[2], rows, \ x->nb[1] / 4, x->nb[2] / 4, x->nb[3] / 4, int(pe->ne[2]), param_f32(nodes[0], 0), scale); \ break - switch (x->ne[0]) { - SD_ROPE(32); - SD_ROPE(64); - SD_ROPE(96); - SD_ROPE(128); - SD_ROPE(160); - SD_ROPE(192); - SD_ROPE(224); - SD_ROPE(256); - default: - return cudaErrorInvalidValue; - } + switch (x->ne[0]) { + SD_ROPE(32); + SD_ROPE(64); + SD_ROPE(96); + SD_ROPE(128); + SD_ROPE(160); + SD_ROPE(192); + SD_ROPE(224); + SD_ROPE(256); + default: + return cuda_invalid_value; + } #undef SD_ROPE - return cudaGetLastError(); -} - -cudaError_t sd_cuda_rms_rope_launch(cudaStream_t stream, ggml_tensor* const* nodes, void* output, bool fp16, float scale) { - if (fp16) { - return rope_interleaved(nodes[2]) ? rms_rope_launch(stream, nodes, static_cast(output), scale) - : rms_rope_launch(stream, nodes, static_cast(output), scale); + return cuda_last_error(); } - return rope_interleaved(nodes[2]) ? rms_rope_launch(stream, nodes, static_cast(output), scale) - : rms_rope_launch(stream, nodes, static_cast(output), scale); -} -static int rope_match(ggml_backend_dev_t, 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) - return 0; - auto* g = const_cast(graph); - if (first + 2 >= ggml_graph_n_nodes(g) || !sd_cuda_rope_fusion_supports( - ggml_graph_node(g, first), ggml_graph_node(g, first + 1), ggml_graph_node(g, first + 2))) - return 0; - if (first + 4 < ggml_graph_n_nodes(g)) { - const auto* rope = ggml_graph_node(g, first + 2); - const auto* scale = ggml_graph_node(g, first + 3); - const auto* copy = ggml_graph_node(g, first + 4); - if (scale->op == GGML_OP_SCALE && copy->op == GGML_OP_CPY && scale->src[0] == rope && copy->src[0] == scale && - scale->type == GGML_TYPE_F32 && copy->type == GGML_TYPE_F16 && param_f32(scale, 1) == 0.f && - ggml_are_same_shape(rope, scale) && ggml_are_same_shape(rope, copy) && ggml_is_contiguous(copy)) - return 5; + CUDAError rms_rope_launch(CUDAStream stream, ggml_tensor* const* nodes, void* output, bool fp16, float scale) { + if (fp16) { + return rope_interleaved(nodes[2]) ? rms_rope_launch(stream, nodes, static_cast(output), scale) + : rms_rope_launch(stream, nodes, static_cast(output), scale); + } + return rope_interleaved(nodes[2]) ? rms_rope_launch(stream, nodes, static_cast(output), scale) + : rms_rope_launch(stream, nodes, static_cast(output), scale); } - 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) + static int rope_match(ggml_backend_dev_t, 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) + return 0; + auto* g = const_cast(graph); + if (first + 2 >= ggml_graph_n_nodes(g) || !rope_fusion_supports( + ggml_graph_node(g, first), ggml_graph_node(g, first + 1), ggml_graph_node(g, first + 2))) + return 0; + if (first + 4 < ggml_graph_n_nodes(g)) { + const auto* rope = ggml_graph_node(g, first + 2); + const auto* scale = ggml_graph_node(g, first + 3); + const auto* copy = ggml_graph_node(g, first + 4); + if (scale->op == GGML_OP_SCALE && copy->op == GGML_OP_CPY && scale->src[0] == rope && copy->src[0] == scale && + scale->type == GGML_TYPE_F32 && copy->type == GGML_TYPE_F16 && param_f32(scale, 1) == 0.f && + ggml_are_same_shape(rope, scale) && ggml_are_same_shape(rope, copy) && ggml_is_contiguous(copy)) + return 5; + } + 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 || !rope_is_node(rope) || rope->src[0] != mul || + !mul->src[0] || mul->src[0]->op != GGML_OP_RMS_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; + 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); } - ggml_ext_graph_delay(graph, j, 1, i); + break; } - 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; -} + 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; + } -static void rope_destroy(void*, void*) {} + static void rope_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 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 rope_execute(const ggml_ext_execution_context* execution, ggml_tensor* const* nodes, int count, void*, void*) { - 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); - const auto* rope = nodes[count == 1 ? 0 : 2]; - const auto* x = count == 1 ? rope->src[0] : nodes[0]->src[0]; - auto* dst = nodes[count - 1]; - const bool alias = overlaps(dst, x) || overlaps(dst, rope->src[1]) || (count > 1 && overlaps(dst, nodes[1]->src[1])); - void* output = alias ? ggml_backend_buffer_get_base(execution->workspace) : dst->data; - cudaError_t error; - if (count > 1) { - error = sd_cuda_rms_rope_launch(stream, nodes, output, count == 5, count == 5 ? param_f32(nodes[3], 0) : 1.f); - } else { - RopeApplyStrides strides; - const auto* pe = rope->src[1]; - for (int i = 0; i < 4; ++i) { - strides.x[i] = x->nb[i] / 4; - strides.pe[i] = pe->nb[i] / 4; + static ggml_status rope_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; + auto stream = static_cast(context.stream); + const auto* rope = nodes[count == 1 ? 0 : 2]; + const auto* x = count == 1 ? rope->src[0] : nodes[0]->src[0]; + auto* dst = nodes[count - 1]; + const bool alias = overlaps(dst, x) || overlaps(dst, rope->src[1]) || (count > 1 && overlaps(dst, nodes[1]->src[1])); + 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); + } else { + RopeApplyStrides strides; + const auto* pe = rope->src[1]; + for (int i = 0; i < 4; ++i) { + strides.x[i] = x->nb[i] / 4; + strides.pe[i] = pe->nb[i] / 4; + } + const int64_t pairs = ggml_nelements(x) / 2; + auto kernel = rope_interleaved(rope) ? rope_apply_f32 : rope_apply_f32; + kernel<<((pairs + 255) / 256), 256, 0, stream>>>(static_cast(x->data), static_cast(pe->data), + static_cast(output), x->ne[0] / 2, pe->ne[2], x->ne[1], x->ne[2], pairs, strides); + error = cuda_last_error(); } - const int64_t pairs = ggml_nelements(x) / 2; - auto kernel = rope_interleaved(rope) ? rope_apply_f32 : rope_apply_f32; - kernel<<((pairs + 255) / 256), 256, 0, stream>>>(static_cast(x->data), static_cast(pe->data), - static_cast(output), x->ne[0] / 2, pe->ne[2], x->ne[1], x->ne[2], pairs, strides); - error = cudaGetLastError(); + 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; + } + + void register_rope(ggml_backend_dev_t device) { + ggml_ext_kernel_desc desc = {}; + desc.abi_version = GGML_EXT_ABI_VERSION; + desc.name = "sd.rope_apply." SD_CUDA_NAME ".v1"; + desc.op = GGML_OP_CUSTOM; + desc.custom_op = sd_ggml_rope_op(); + desc.priority = 100; + desc.supports = rope_supports; + desc.prepare = rope_prepare; + desc.destroy = rope_destroy; + desc.execute = rope_execute; + GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE)); + desc.name = "sd.rope_apply." SD_CUDA_NAME ".override.v1"; + desc.op = GGML_OP_ROPE_APPLY; + desc.custom_op = nullptr; + GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE)); + desc.name = "sd.rms_norm_mul_rope." SD_CUDA_NAME ".v1"; + desc.op = GGML_OP_RMS_NORM; + desc.priority = 250; + desc.supports = nullptr; + desc.match = rope_match; + desc.optimize = rope_optimize; + GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE)); } - 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_rope(ggml_backend_dev_t device) { - ggml_ext_kernel_desc desc = {}; - desc.abi_version = GGML_EXT_ABI_VERSION; - desc.name = "sd.rope_apply.cuda.v1"; - desc.op = GGML_OP_CUSTOM; - desc.custom_op = sd_ggml_rope_op(); - desc.priority = 100; - desc.supports = rope_supports; - desc.prepare = rope_prepare; - desc.destroy = rope_destroy; - desc.execute = rope_execute; - GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE)); - desc.name = "sd.rope_apply.cuda.override.v1"; - desc.op = GGML_OP_ROPE_APPLY; - desc.custom_op = nullptr; - GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE)); - desc.name = "sd.rms_norm_mul_rope.cuda.v1"; - desc.op = GGML_OP_RMS_NORM; - desc.priority = 250; - desc.supports = nullptr; - desc.match = rope_match; - desc.optimize = rope_optimize; - 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.h b/src/core/cuda/rope.h index c017718c..0d45f689 100644 --- a/src/core/cuda/rope.h +++ b/src/core/cuda/rope.h @@ -1,12 +1,11 @@ #ifndef __SD_CORE_CUDA_ROPE_H__ #define __SD_CORE_CUDA_ROPE_H__ -#include -#include "ggml-ext.h" +#include "runtime.h" -bool sd_cuda_rope_is_node(const ggml_tensor* node); -bool sd_cuda_rope_fusion_supports(const ggml_tensor* norm, const ggml_tensor* mul, const ggml_tensor* rope); -cudaError_t sd_cuda_rms_rope_launch(cudaStream_t stream, ggml_tensor* const* nodes, void* output, bool fp16, float scale); -void sd_ggml_register_cuda_rope(ggml_backend_dev_t device); +namespace SD_CUDA_NAMESPACE { + bool rope_fusion_supports(const ggml_tensor* norm, const ggml_tensor* mul, const ggml_tensor* rope); + CUDAError rms_rope_launch(CUDAStream stream, ggml_tensor* const* nodes, void* output, bool fp16, float scale); +} #endif diff --git a/src/core/cuda/runtime.h b/src/core/cuda/runtime.h new file mode 100644 index 00000000..f0db1dcb --- /dev/null +++ b/src/core/cuda/runtime.h @@ -0,0 +1,111 @@ +#ifndef __SD_CORE_CUDA_RUNTIME_H__ +#define __SD_CORE_CUDA_RUNTIME_H__ + +#include "ggml-ext.h" +#if defined(SD_CUDA_VENDOR_HIP) +#include +#include +#define SD_CUDA_NAMESPACE sd_hip +#define SD_CUDA_NAME "hip" +#elif defined(SD_CUDA_VENDOR_MUSA) +#include +#include +#define SD_CUDA_NAMESPACE sd_musa +#define SD_CUDA_NAME "musa" +#else +#include +#include +#define SD_CUDA_NAMESPACE sd_cuda +#define SD_CUDA_NAME "cuda" +#endif + +namespace SD_CUDA_NAMESPACE { +#if defined(SD_CUDA_VENDOR_HIP) + using CUDAError = hipError_t; + using CUDAStream = hipStream_t; + static constexpr auto cuda_success = hipSuccess; + static constexpr auto cuda_invalid_value = hipErrorInvalidValue; + static constexpr auto cuda_kind = GGML_EXT_GPU_HIP; + static CUDAError cuda_last_error() { + return hipGetLastError(); + } + static CUDAError cuda_memset_async(void* dst, int value, size_t bytes, CUDAStream stream) { + return hipMemsetAsync(dst, value, bytes, stream); + } + static CUDAError cuda_copy_async(void* dst, const void* src, size_t bytes, CUDAStream stream) { + return hipMemcpyAsync(dst, src, bytes, hipMemcpyDeviceToDevice, stream); + } + template + static CUDAError cuda_set_shared_memory(Kernel, int) { + // HIP uses the device's shared-memory limit without CUDA's opt-in attribute. + return cuda_success; + } +#elif defined(SD_CUDA_VENDOR_MUSA) + using CUDAError = musaError_t; + using CUDAStream = musaStream_t; + static constexpr auto cuda_success = musaSuccess; + static constexpr auto cuda_invalid_value = musaErrorInvalidValue; + static constexpr auto cuda_kind = GGML_EXT_GPU_MUSA; + static CUDAError cuda_last_error() { + return musaGetLastError(); + } + static CUDAError cuda_memset_async(void* dst, int value, size_t bytes, CUDAStream stream) { + return musaMemsetAsync(dst, value, bytes, stream); + } + static CUDAError cuda_copy_async(void* dst, const void* src, size_t bytes, CUDAStream stream) { + return musaMemcpyAsync(dst, src, bytes, musaMemcpyDeviceToDevice, stream); + } + template + static CUDAError cuda_set_shared_memory(Kernel kernel, int bytes) { + return musaFuncSetAttribute(kernel, musaFuncAttributeMaxDynamicSharedMemorySize, bytes); + } +#else + using CUDAError = cudaError_t; + using CUDAStream = cudaStream_t; + static constexpr auto cuda_success = cudaSuccess; + static constexpr auto cuda_invalid_value = cudaErrorInvalidValue; + static constexpr auto cuda_kind = GGML_EXT_GPU_CUDA; + static CUDAError cuda_last_error() { + return cudaGetLastError(); + } + static CUDAError cuda_memset_async(void* dst, int value, size_t bytes, CUDAStream stream) { + return cudaMemsetAsync(dst, value, bytes, stream); + } + static CUDAError cuda_copy_async(void* dst, const void* src, size_t bytes, CUDAStream stream) { + return cudaMemcpyAsync(dst, src, bytes, cudaMemcpyDeviceToDevice, stream); + } + template + static CUDAError cuda_set_shared_memory(Kernel kernel, int bytes) { + return cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, bytes); + } +#endif + template + static __device__ __forceinline__ T cuda_shuffle(T value, int lane, int width = 32) { +#if defined(SD_CUDA_VENDOR_HIP) + return __shfl(value, lane, width); +#else + return __shfl_sync(0xffffffff, value, lane, width); +#endif + } + template + static __device__ __forceinline__ T cuda_shuffle_xor(T value, int lane, int width = 32) { +#if defined(SD_CUDA_VENDOR_HIP) + return __shfl_xor(value, lane, width); +#else + return __shfl_xor_sync(0xffffffff, value, lane, width); +#endif + } + static bool cuda_context(ggml_backend_t backend, ggml_ext_gpu_context* context) { + auto reg = ggml_backend_dev_backend_reg(ggml_backend_get_device(backend)); + auto query = reinterpret_cast(ggml_backend_reg_get_proc_address(reg, "ggml_ext_gpu_context_v1")); + return query && query(backend, context) && context->kind == cuda_kind; + } + static ggml_ext_gpu_device_info cuda_device_info(ggml_backend_dev_t device) { + ggml_ext_gpu_device_info info = {}; + auto query = reinterpret_cast(ggml_backend_reg_get_proc_address(ggml_backend_dev_backend_reg(device), "ggml_ext_gpu_device_v1")); + if (!query || !query(device, &info) || info.kind != cuda_kind) + return {cuda_kind, -1, 0, 0, false}; + return info; + } +} +#endif diff --git a/src/core/ggml_extensions.cpp b/src/core/ggml_extensions.cpp index 8fe44279..c63ff8ac 100644 --- a/src/core/ggml_extensions.cpp +++ b/src/core/ggml_extensions.cpp @@ -9,19 +9,8 @@ #include #include -#ifdef SD_VULKAN_EXTENSIONS -void sd_ggml_register_vulkan_rope(ggml_backend_dev_t device); -#endif -#ifdef SD_CUDA_EXTENSIONS -void sd_ggml_register_cuda_affine(ggml_backend_dev_t device); -void sd_ggml_register_cuda_i8(ggml_backend_dev_t device); -void sd_ggml_register_cuda_convrot(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); +#ifdef SD_BACKEND_EXTENSIONS +#include "sd_backend_extensions.h" #endif static bool rope_infer(ggml_tensor* const* inputs, int count, const void*, ggml_ext_tensor_desc* out) { @@ -171,37 +160,19 @@ void sd_ggml_register_extensions(ggml_backend_t backend) { if (!backend) { return; } - auto device = ggml_backend_get_device(backend); - auto reg = ggml_backend_dev_backend_reg(device); - const bool cpu = ggml_backend_reg_get_proc_address(reg, "ggml_ext_cpu_parallel_v1") != nullptr; - const bool cuda = ggml_backend_reg_get_proc_address(reg, "ggml_ext_cuda_context_v1") != nullptr; - const bool vulkan = ggml_backend_reg_get_proc_address(reg, "ggml_ext_vulkan_v1") != nullptr; - if (!cpu && !cuda && !vulkan) { - return; - } + auto device = ggml_backend_get_device(backend); + auto reg = ggml_backend_dev_backend_reg(device); + const bool cpu = ggml_backend_reg_get_proc_address(reg, "ggml_ext_cpu_parallel_v1") != nullptr; static std::mutex mutex; static std::unordered_set registered; std::lock_guard lock(mutex); if (!registered.insert(device).second) { return; } - if (vulkan) { -#ifdef SD_VULKAN_EXTENSIONS - sd_ggml_register_vulkan_rope(device); -#endif - return; - } - if (cuda) { -#ifdef SD_CUDA_EXTENSIONS - sd_ggml_register_cuda_affine(device); - sd_ggml_register_cuda_i8(device); - sd_ggml_register_cuda_convrot(device); - sd_ggml_register_cuda_rope(device); -#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); + if (!cpu) { +#ifdef SD_BACKEND_EXTENSIONS + for (auto register_backend : sd_backend_extensions) + register_backend(device); #endif return; } @@ -234,11 +205,8 @@ ggml_tensor* sd_ggml_rope_apply(ggml_context* ctx, ggml_backend_t backend, ggml_ if (backend) { auto reg = ggml_backend_dev_backend_reg(ggml_backend_get_device(backend)); if (ggml_backend_reg_get_proc_address(reg, "ggml_ext_cpu_parallel_v1") -#ifdef SD_CUDA_EXTENSIONS - || ggml_backend_reg_get_proc_address(reg, "ggml_ext_cuda_context_v1") -#endif -#ifdef SD_VULKAN_EXTENSIONS - || ggml_backend_reg_get_proc_address(reg, "ggml_ext_vulkan_v1") +#ifdef SD_BACKEND_EXTENSIONS + || ggml_backend_reg_get_proc_address(reg, "ggml_ext_gpu_context_v1") || ggml_backend_reg_get_proc_address(reg, "ggml_ext_vulkan_v1") #endif ) { ggml_tensor* inputs[] = {x, pe}; diff --git a/src/core/vulkan/rope.cpp b/src/core/vulkan/rope.cpp index 8a485139..4ed0811d 100644 --- a/src/core/vulkan/rope.cpp +++ b/src/core/vulkan/rope.cpp @@ -200,6 +200,8 @@ namespace sd_vulkan_rope { } void sd_ggml_register_vulkan_rope(ggml_backend_dev_t device) { + if (!sd_vulkan_rope::api(device)) + return; ggml_ext_kernel_desc desc = {}; desc.abi_version = GGML_EXT_ABI_VERSION; desc.name = "sd.rope_apply.vulkan.v1";