diff --git a/CMakeLists.txt b/CMakeLists.txt index 15bd91742..048f273de 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -350,15 +350,15 @@ if (SD_CUDA) enable_language(CUDA) set_target_properties(${SD_LIB} PROPERTIES CUDA_STANDARD 17 CUDA_STANDARD_REQUIRED ON) if (SD_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 + 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 src/core/cuda/bf16_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 src/core/cuda/bf16_gemm.cu PROPERTIES INCLUDE_DIRECTORIES "${SD_CUDA_CUTLASS_PATH}/include" COMPILE_OPTIONS "--use_fast_math;--extended-lambda;--expt-relaxed-constexpr") if (MSVC) - set_property(SOURCE src/core/cuda/cute_attention.cu src/core/cuda/i8_projection.cu src/core/cuda/i8_gated_residual.cu src/core/cuda/i8_gemm.cu APPEND PROPERTY COMPILE_OPTIONS -Xcompiler=/Zc:preprocessor,/Zc:__cplusplus) + set_property(SOURCE src/core/cuda/cute_attention.cu src/core/cuda/i8_projection.cu src/core/cuda/i8_gated_residual.cu src/core/cuda/i8_gemm.cu src/core/cuda/bf16_gemm.cu APPEND PROPERTY COMPILE_OPTIONS -Xcompiler=/Zc:preprocessor,/Zc:__cplusplus) endif() if (GGML_CUDA_FORCE_CUBLAS) - set_property(SOURCE src/core/cuda/i8_projection.cu src/core/cuda/i8_gated_residual.cu src/core/cuda/i8_gemm.cu APPEND PROPERTY COMPILE_DEFINITIONS SD_CUDA_FORCE_CUBLAS) + set_property(SOURCE src/core/cuda/i8_projection.cu src/core/cuda/i8_gated_residual.cu src/core/cuda/i8_gemm.cu src/core/cuda/bf16_gemm.cu APPEND PROPERTY COMPILE_DEFINITIONS SD_CUDA_FORCE_CUBLAS) endif() endif() endif() diff --git a/src/core/cuda/CMakeLists.txt b/src/core/cuda/CMakeLists.txt index f366899fe..0a95d6cf7 100644 --- a/src/core/cuda/CMakeLists.txt +++ b/src/core/cuda/CMakeLists.txt @@ -2,7 +2,7 @@ cmake_policy(SET CMP0079 NEW) function(sd_cuda_sources VENDOR LANGUAGE) set(SOURCES) - foreach(OP rope int8 convrot affine modulate registration) + foreach(OP rope int8 convrot affine modulate bf16 rms_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") @@ -27,6 +27,7 @@ if(SD_CUDA) set_source_files_properties(sage_attention.cu sol_attention.cu PROPERTIES COMPILE_OPTIONS "--use_fast_math") set_source_files_properties(${SD_CUDA_SOURCES} PROPERTIES COMPILE_OPTIONS "--use_fast_math") if(SD_CUDA_CUTLASS_PATH) + set_property(SOURCE "${CMAKE_CURRENT_BINARY_DIR}/bf16_cuda.cu" APPEND PROPERTY COMPILE_DEFINITIONS SD_CUDA_BF16_CUTLASS) 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") diff --git a/src/core/cuda/bf16.cu b/src/core/cuda/bf16.cu new file mode 100644 index 000000000..59f507394 --- /dev/null +++ b/src/core/cuda/bf16.cu @@ -0,0 +1,277 @@ +#include "bf16.h" +#if defined(SD_CUDA_VENDOR_HIP) +#include +#elif defined(SD_CUDA_VENDOR_MUSA) +#include +#else +#include +#endif +#include +#include +#include +#include "rms_affine.h" +#include "runtime.h" +#ifdef SD_CUDA_BF16_CUTLASS +#include "bf16_gemm.h" +#endif + +namespace SD_CUDA_NAMESPACE { + namespace sd_bf16 { + static float scale_value(const ggml_tensor* node) { + float value; + std::memcpy(&value, node->op_params, sizeof(value)); + return value; + } + + static bool matrix(const ggml_tensor* node, ggml_type type) { + return node && node->type == type && ggml_is_contiguous(node) && + node->ne[2] == 1 && node->ne[3] == 1 && node->ne[0] <= INT_MAX && node->ne[1] <= INT_MAX; + } + + static bool projection(const ggml_tensor* mat, const ggml_tensor* input) { + return mat->op == GGML_OP_MUL_MAT && mat->src[1] == input && mat->op_params[0] == GGML_PREC_DEFAULT && + matrix(input, GGML_TYPE_F32) && input->ne[1] >= 32 && ggml_nelements(input) <= INT_MAX && + matrix(mat, GGML_TYPE_F32) && matrix(mat->src[0], GGML_TYPE_BF16) && !mat->view_src; + } + + static int match(ggml_backend_dev_t, const ggml_cgraph* graph, int first, void*) { + // An explicit cuBLAS compute override must retain the native conversion path. + static const bool disabled = [] { + const char* value = std::getenv("GGML_CUDA_DISABLE_FUSION"); + return (value && std::atoi(value)) || std::getenv("GGML_CUDA_CUBLAS_COMPUTE_TYPE"); + }(); + if (disabled) + return 0; + auto* g = const_cast(graph); + if (first + 1 >= ggml_graph_n_nodes(g)) + return 0; + const auto* glu = ggml_graph_node(g, first); + int offset = 0; + if (glu->op == GGML_OP_RMS_NORM) { + if (!rms_affine_match(graph, first) || first + 4 >= ggml_graph_n_nodes(g)) + return 0; + const auto* input = ggml_graph_node(g, first + 3); + const auto* gate = ggml_graph_node(g, first + 4); + if (!projection(gate, input)) + return 0; + if (first + 6 >= ggml_graph_n_nodes(g)) + return 5; + const auto* up = ggml_graph_node(g, first + 5); + glu = ggml_graph_node(g, first + 6); + if (!projection(up, input) || glu->op != GGML_OP_GLU || glu->src[0] != gate || glu->src[1] != up) + return 5; + offset = 6; + } + if (first + offset + 1 >= ggml_graph_n_nodes(g)) + return 0; + if (glu->op != GGML_OP_GLU || ggml_get_glu_op(glu) != GGML_GLU_OP_SWIGLU || + !matrix(glu, GGML_TYPE_F32) || !matrix(glu->src[0], GGML_TYPE_F32) || + !matrix(glu->src[1], GGML_TYPE_F32) || !ggml_are_same_shape(glu, glu->src[0]) || + !ggml_are_same_shape(glu, glu->src[1]) || glu->ne[1] < 32 || ggml_nelements(glu) > INT_MAX) + return 0; + int count = 2; + const auto* input = glu; + const auto* mat = ggml_graph_node(g, first + offset + 1); + if (mat->op == GGML_OP_SCALE) { + float bias; + std::memcpy(&bias, mat->op_params + 1, sizeof(bias)); + if (mat->src[0] != glu || bias != 0.f || first + offset + 2 >= ggml_graph_n_nodes(g)) + return 0; + input = mat; + mat = ggml_graph_node(g, first + offset + 2); + count = 3; + } + if (!projection(mat, input)) + return 0; + return offset + count; + } + + struct Workspace { + size_t norm = 0, gate = 0, up = 0, packed = 0, output = 0, size = 0; + Workspace(ggml_tensor* const* nodes, int count) { + if (nodes[0]->op == GGML_OP_RMS_NORM) { + norm = reserve(ggml_nelements(nodes[3]) * sizeof(CUDABFloat16)); + if (count > 5) { + gate = reserve(ggml_nbytes(nodes[4])); + up = reserve(ggml_nbytes(nodes[5])); + packed = reserve(ggml_nelements(nodes[6]) * sizeof(CUDABFloat16)); + } + } else { + packed = reserve(ggml_nelements(nodes[0]) * sizeof(CUDABFloat16)); + } + output = reserve(ggml_nbytes(nodes[count - 1])); + } + size_t reserve(size_t bytes) { + const size_t offset = size; + size += GGML_PAD(bytes, 256); + return offset; + } + }; + + static ggml_status prepare(ggml_backend_t, ggml_tensor* const* nodes, int count, void**, size_t* workspace, void*) { + *workspace = Workspace(nodes, count).size; + return GGML_STATUS_SUCCESS; + } + + static void destroy(void*, void*) {} + + static __device__ __forceinline__ float swiglu(float x, float up, float scale) { + return __fmul_rn(__fmul_rn(x / (1.f + expf(-x)), up), scale); + } + + static __global__ void pack_swiglu(const float* gate, const float* up, CUDABFloat16* dst, int64_t size, float scale) { + const int64_t i = int64_t(blockIdx.x) * blockDim.x + threadIdx.x; + if (i < size) { + dst[i] = __float2bfloat16(swiglu(gate[i], up[i], scale)); + } + } + + static __global__ void pack_swiglu4(const float4* gate, const float4* up, CUDABFloat162* dst, int64_t size, float scale) { + const int64_t i = int64_t(blockIdx.x) * blockDim.x + threadIdx.x; + if (i < size) { + const float4 x = gate[i], y = up[i]; + dst[2 * i] = __float22bfloat162_rn(make_float2(swiglu(x.x, y.x, scale), swiglu(x.y, y.y, scale))); + dst[2 * i + 1] = __float22bfloat162_rn(make_float2(swiglu(x.z, y.z, scale), swiglu(x.w, y.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 gemm(void* handle, const ggml_tensor* w, const void* input, int rows, void* output, CUDAStream stream, bool tiled) { +#ifdef SD_CUDA_BF16_CUTLASS + if (tiled) { + const auto result = sd_cuda_bf16_gemm(w->data, input, static_cast(output), rows, w->ne[1], w->ne[0], stream); + if (result == SDBF16GemmResult::Success) + return GGML_STATUS_SUCCESS; + if (result == SDBF16GemmResult::Failed) + return GGML_STATUS_FAILED; + } +#else + (void)tiled; +#endif + const float alpha = 1.f, beta = 0.f; +#if defined(SD_CUDA_VENDOR_HIP) +#if HIP_VERSION >= 60500000 || defined(HIPBLAS_V2) + constexpr auto compute_type = HIPBLAS_COMPUTE_32F; +#else + constexpr auto compute_type = HIPBLAS_R_32F; +#endif +#if HIP_VERSION >= 70000000 || defined(HIPBLAS_V2) + constexpr auto bf16_type = HIP_R_16BF, f32_type = HIP_R_32F; +#else + constexpr auto bf16_type = HIPBLAS_R_16B, f32_type = HIPBLAS_R_32F; +#endif + return hipblasGemmEx(static_cast(handle), HIPBLAS_OP_T, HIPBLAS_OP_N, + int(w->ne[1]), rows, int(w->ne[0]), &alpha, + w->data, bf16_type, int(w->ne[0]), input, bf16_type, int(w->ne[0]), + &beta, output, f32_type, int(w->ne[1]), compute_type, HIPBLAS_GEMM_DEFAULT) == HIPBLAS_STATUS_SUCCESS + ? GGML_STATUS_SUCCESS + : GGML_STATUS_FAILED; +#elif defined(SD_CUDA_VENDOR_MUSA) + return mublasGemmEx(static_cast(handle), MUBLAS_OP_T, MUBLAS_OP_N, + int(w->ne[1]), rows, int(w->ne[0]), &alpha, + w->data, MUSA_R_16BF, int(w->ne[0]), input, MUSA_R_16BF, int(w->ne[0]), + &beta, output, MUSA_R_32F, int(w->ne[1]), MUSA_R_32F, MUBLAS_GEMM_DEFAULT) == MUBLAS_STATUS_SUCCESS + ? GGML_STATUS_SUCCESS + : GGML_STATUS_FAILED; +#else + return cublasGemmEx(static_cast(handle), CUBLAS_OP_T, CUBLAS_OP_N, int(w->ne[1]), rows, int(w->ne[0]), + &alpha, w->data, CUDA_R_16BF, int(w->ne[0]), input, CUDA_R_16BF, int(w->ne[0]), + &beta, output, CUDA_R_32F, int(w->ne[1]), CUBLAS_COMPUTE_32F, CUBLAS_GEMM_DEFAULT_TENSOR_OP) == CUBLAS_STATUS_SUCCESS + ? GGML_STATUS_SUCCESS + : GGML_STATUS_FAILED; +#endif + } + + static ggml_status execute(const ggml_ext_execution_context* execution, ggml_tensor* const* nodes, int count, void*, void* user_data) { + const bool tiled = *static_cast(user_data); + ggml_ext_gpu_context context; + if (!cuda_context(execution->backend, &context)) + return GGML_STATUS_FAILED; + const auto reg = ggml_backend_dev_backend_reg(ggml_backend_get_device(execution->backend)); + const auto query = reinterpret_cast(ggml_backend_reg_get_proc_address(reg, "ggml_ext_gpu_blas_v1")); + const auto handle = query ? query(execution->backend) : nullptr; + if (!handle) + return GGML_STATUS_FAILED; + const auto stream = static_cast(context.stream); + const bool norm = nodes[0]->op == GGML_OP_RMS_NORM; + const int offset = norm ? 6 : 0; + const auto* glu = norm && count == 5 ? nullptr : nodes[offset]; + auto* dst = nodes[count - 1]; + const auto* w = dst->src[0]; + auto* base = static_cast(ggml_backend_buffer_get_base(execution->workspace)); + const Workspace work(nodes, count); + auto* packed = reinterpret_cast(base + work.packed); + const bool alias = overlaps(dst, w); + void* output = alias ? base + work.output : dst->data; + if (norm) { + if (rms_affine_pack(stream, nodes, base + work.norm) != GGML_STATUS_SUCCESS) + return GGML_STATUS_FAILED; + if (gemm(handle, nodes[4]->src[0], base + work.norm, int(nodes[3]->ne[1]), count == 5 ? output : base + work.gate, stream, tiled) != GGML_STATUS_SUCCESS) + return GGML_STATUS_FAILED; + if (count == 5) + return !alias || cuda_copy_async(dst->data, output, ggml_nbytes(dst), stream) == cuda_success ? GGML_STATUS_SUCCESS : GGML_STATUS_FAILED; + if (gemm(handle, nodes[5]->src[0], base + work.norm, int(nodes[3]->ne[1]), base + work.up, stream, tiled) != GGML_STATUS_SUCCESS) + return GGML_STATUS_FAILED; + } + const auto size = ggml_nelements(glu); + const void* gate = norm ? base + work.gate : glu->src[0]->data; + const void* up = norm ? base + work.up : glu->src[1]->data; + const float scale = count - offset == 3 ? scale_value(nodes[offset + 1]) : 1.f; + if (size % 4 == 0 && (reinterpret_cast(gate) | reinterpret_cast(up)) % 16 == 0) { + pack_swiglu4<<>>(static_cast(gate), + static_cast(up), + reinterpret_cast(packed), size / 4, scale); + } else { + pack_swiglu<<>>(static_cast(gate), + static_cast(up), packed, size, scale); + } + if (cuda_last_error() != cuda_success) + return GGML_STATUS_FAILED; + const auto status = gemm(handle, w, packed, int(glu->ne[1]), output, stream, tiled); + if (status != GGML_STATUS_SUCCESS) + return GGML_STATUS_FAILED; + return !alias || cuda_copy_async(dst->data, output, ggml_nbytes(dst), stream) == cuda_success ? GGML_STATUS_SUCCESS : GGML_STATUS_FAILED; + } + } + + void register_bf16(ggml_backend_dev_t device) { + const auto info = cuda_device_info(device); + if (info.device < 0 || !ggml_backend_reg_get_proc_address(ggml_backend_dev_backend_reg(device), "ggml_ext_gpu_blas_v1")) + return; + bool tiled = false; +#if defined(SD_CUDA_VENDOR_HIP) + hipDeviceProp_t props; + if (hipGetDeviceProperties(&props, info.device) != hipSuccess || std::strncmp(props.gcnArchName, "gfx", 3)) + return; + const auto arch = std::strtoul(props.gcnArchName + 3, nullptr, 16); + // Match the native backend's BF16 operand policy; older GPUs use FP32 GEMM. + if (arch < 0x1100 && !(arch >= 0x908 && arch < 0x1010 && arch != 0x909 && arch != 0x90c)) + return; +#elif !defined(SD_CUDA_VENDOR_MUSA) + cudaDeviceProp props; + if (cudaGetDeviceProperties(&props, info.device) != cuda_success || props.major < 8) + return; + tiled = props.major >= 8; +#endif + static bool tiled_gemm[] = {false, true}; + ggml_ext_kernel_desc desc = {}; + desc.user_data = &tiled_gemm[tiled]; + desc.abi_version = GGML_EXT_ABI_VERSION; + desc.name = "sd.swiglu_bf16_matmul." SD_CUDA_NAME ".v1"; + desc.op = GGML_OP_GLU; + desc.priority = 200; + desc.match = sd_bf16::match; + desc.prepare = sd_bf16::prepare; + desc.destroy = sd_bf16::destroy; + desc.execute = sd_bf16::execute; + GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE)); + desc.name = "sd.rms_affine_bf16_matmul." SD_CUDA_NAME ".v1"; + desc.op = GGML_OP_RMS_NORM; + desc.priority = 400; + 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/bf16.h b/src/core/cuda/bf16.h new file mode 100644 index 000000000..b5da3704f --- /dev/null +++ b/src/core/cuda/bf16.h @@ -0,0 +1,26 @@ +#ifndef __SD_CORE_CUDA_BF16_H__ +#define __SD_CORE_CUDA_BF16_H__ + +#include "runtime.h" +#if defined(SD_CUDA_VENDOR_HIP) +#include +#elif defined(SD_CUDA_VENDOR_MUSA) +#include +#else +#include +#endif + +namespace SD_CUDA_NAMESPACE { +#if defined(SD_CUDA_VENDOR_HIP) + using CUDABFloat16 = __hip_bfloat16; + using CUDABFloat162 = __hip_bfloat162; +#elif defined(SD_CUDA_VENDOR_MUSA) + using CUDABFloat16 = __mt_bfloat16; + using CUDABFloat162 = __mt_bfloat162; +#else + using CUDABFloat16 = nv_bfloat16; + using CUDABFloat162 = nv_bfloat162; +#endif +} + +#endif diff --git a/src/core/cuda/bf16_gemm.cu b/src/core/cuda/bf16_gemm.cu new file mode 100644 index 000000000..41d75706c --- /dev/null +++ b/src/core/cuda/bf16_gemm.cu @@ -0,0 +1,25 @@ +#include +#include +#include "bf16_gemm.h" + +SDBF16GemmResult sd_cuda_bf16_gemm(const void* weight, const void* input, float* output, int64_t rows, int64_t n, int64_t k, cudaStream_t stream) { +#ifdef SD_CUDA_FORCE_CUBLAS + return SDBF16GemmResult::Unsupported; +#else + if (rows < 512 || rows > 2048 || (n < 4096 && (n < 2048 || k < 2 * n)) || k % 8 || n % 4 || + (reinterpret_cast(weight) | reinterpret_cast(input) | reinterpret_cast(output)) % 16) + return SDBF16GemmResult::Unsupported; + using Gemm = cutlass::gemm::device::Gemm< + cutlass::bfloat16_t, cutlass::layout::RowMajor, cutlass::bfloat16_t, cutlass::layout::ColumnMajor, + float, cutlass::layout::RowMajor, float, cutlass::arch::OpClassTensorOp, cutlass::arch::Sm80, + cutlass::gemm::GemmShape<64, 128, 32>, cutlass::gemm::GemmShape<32, 64, 32>, cutlass::gemm::GemmShape<16, 8, 16>, + cutlass::epilogue::thread::LinearCombination, cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>, 3>; + Gemm gemm; + const Gemm::Arguments args({int(rows), int(n), int(k)}, + {static_cast(input), k}, + {static_cast(weight), k}, {output, n}, {output, n}, {1.f, 0.f}); + if (gemm.can_implement(args) != cutlass::Status::kSuccess) + return SDBF16GemmResult::Unsupported; + return gemm(args, nullptr, stream) == cutlass::Status::kSuccess ? SDBF16GemmResult::Success : SDBF16GemmResult::Failed; +#endif +} diff --git a/src/core/cuda/bf16_gemm.h b/src/core/cuda/bf16_gemm.h new file mode 100644 index 000000000..6cb36501f --- /dev/null +++ b/src/core/cuda/bf16_gemm.h @@ -0,0 +1,13 @@ +#ifndef __SD_CORE_CUDA_BF16_GEMM_H__ +#define __SD_CORE_CUDA_BF16_GEMM_H__ + +#include +#include + +enum class SDBF16GemmResult { Unsupported, + Success, + Failed }; + +SDBF16GemmResult sd_cuda_bf16_gemm(const void* weight, const void* input, float* output, int64_t rows, int64_t n, int64_t k, cudaStream_t stream); + +#endif diff --git a/src/core/cuda/registration.cu b/src/core/cuda/registration.cu index 291edd62a..d75728021 100644 --- a/src/core/cuda/registration.cu +++ b/src/core/cuda/registration.cu @@ -12,6 +12,8 @@ namespace SD_CUDA_NAMESPACE { void register_convrot(ggml_backend_dev_t device); void register_affine(ggml_backend_dev_t device); void register_modulate(ggml_backend_dev_t device); + void register_bf16(ggml_backend_dev_t device); + void register_rms_affine(ggml_backend_dev_t device); #if !defined(SD_CUDA_VENDOR_HIP) && !defined(SD_CUDA_VENDOR_MUSA) void register_sage_attention(ggml_backend_dev_t device); void register_sol_attention(ggml_backend_dev_t device); @@ -25,6 +27,8 @@ namespace SD_CUDA_NAMESPACE { register_i8(device); register_convrot(device); register_rope(device); + register_bf16(device); + register_rms_affine(device); #if !defined(SD_CUDA_VENDOR_HIP) && !defined(SD_CUDA_VENDOR_MUSA) register_sage_attention(device); register_sol_attention(device); diff --git a/src/core/cuda/rms_affine.cu b/src/core/cuda/rms_affine.cu new file mode 100644 index 000000000..d9ac351ed --- /dev/null +++ b/src/core/cuda/rms_affine.cu @@ -0,0 +1,168 @@ +#include +#include +#include +#include "bf16.h" +#include "rms_affine.h" + +namespace SD_CUDA_NAMESPACE { + namespace sd_rms_affine { + static bool disabled() { + static const bool value = [] { + const char* fusion = std::getenv("GGML_CUDA_DISABLE_FUSION"); + return fusion && std::atoi(fusion); + }(); + return value; + } + + static bool f32(const ggml_tensor* t) { + return t && t->type == GGML_TYPE_F32 && ggml_is_contiguous(t); + } + + static int match(ggml_backend_dev_t, const ggml_cgraph* graph, int first, void*) { + auto* g = const_cast(graph); + if (disabled() || first + 3 >= ggml_graph_n_nodes(g)) + return 0; + const auto* norm = ggml_graph_node(g, first); + const auto* affine = ggml_graph_node(g, first + 1); + const auto* mul = ggml_graph_node(g, first + 2); + const auto* add = ggml_graph_node(g, first + 3); + if (norm->op != GGML_OP_RMS_NORM || affine->op != GGML_OP_MUL || affine->src[0] != norm || + mul->op != GGML_OP_MUL || mul->src[0] != affine || add->op != GGML_OP_ADD || + (add->src[0] != mul && add->src[1] != mul)) + return 0; + const auto* x = norm->src[0]; + const auto* residual = add->src[0] == mul ? add->src[1] : add->src[0]; + return f32(x) && f32(affine->src[1]) && f32(mul->src[1]) && f32(residual) && f32(add) && + x->ne[0] >= 32 && x->ne[0] <= 8192 && ggml_nrows(x) <= INT_MAX && + ggml_nelements(affine->src[1]) == x->ne[0] && ggml_nelements(mul->src[1]) == x->ne[0] && + affine->src[1]->ne[0] == x->ne[0] && mul->src[1]->ne[0] == x->ne[0] && + residual != mul && ggml_are_same_shape(residual, x) && ggml_are_same_shape(add, x) + ? 4 + : 0; + } + + static void optimize(ggml_backend_dev_t, ggml_cgraph* graph, void*) { + if (disabled()) + return; + for (int i = 2; i + 1 < ggml_graph_n_nodes(graph); ++i) { + const auto* mul = ggml_graph_node(graph, i); + const auto* add = ggml_graph_node(graph, i + 1); + if (mul->op != GGML_OP_MUL || add->op != GGML_OP_ADD || (add->src[0] != mul && add->src[1] != mul)) + continue; + const auto* affine = mul->src[0]; + if (affine->op != GGML_OP_MUL || affine->src[0]->op != GGML_OP_RMS_NORM) + continue; + for (int j = i - 2; j >= 0 && j >= i - 32; --j) { + if (ggml_graph_node(graph, j) == affine->src[0] && ggml_graph_node(graph, j + 1) == affine) { + if (j + 2 < i) + ggml_ext_graph_delay(graph, j, 2, i); + break; + } + } + } + } + + static __device__ __forceinline__ float warp_sum(float value) { + for (int offset = 16; offset > 0; offset >>= 1) + value = __fadd_rn(value, cuda_shuffle_xor(value, offset, 32)); + return value; + } + + template + static __global__ void kernel(const float* x, const float* weight, const float* gate, const float* residual, T* dst, int width, float eps) { + const int col0 = threadIdx.x; + const int64_t base = int64_t(blockIdx.x) * width; + float values[8], sum = 0.f; + for (int j = 0, col = col0; col < width; ++j, col += Block) { + values[j] = x[base + col]; + sum = __fmaf_rn(values[j], values[j], sum); + } + sum = warp_sum(sum); + __shared__ float sums[32]; + const int lane = col0 % 32; + if (lane == 0) + sums[col0 / 32] = sum; + __syncthreads(); + sum = warp_sum(lane < Block / 32 ? sums[lane] : 0.f); + const float scale = rsqrtf(sum / width + eps); + for (int j = 0, col = col0; col < width; ++j, col += Block) { + const float value = __fmul_rn(__fmul_rn(scale, values[j]), weight[col]); + const float r = Modulate ? value : residual[base + col]; + dst[base + col] = T(__fadd_rn(__fmul_rn(value, gate[col]), r)); + } + } + + 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, void**, size_t* workspace, void*) { + *workspace = ggml_nbytes(nodes[3]); + return GGML_STATUS_SUCCESS; + } + static void destroy(void*, void*) {} + + static ggml_status execute(const ggml_ext_execution_context* execution, ggml_tensor* const* nodes, int, void*, void*) { + ggml_ext_gpu_context context; + if (!cuda_context(execution->backend, &context)) + return GGML_STATUS_FAILED; + const auto* x = nodes[0]->src[0]; + const auto* w = nodes[1]->src[1]; + const auto* gate = nodes[2]->src[1]; + auto* dst = nodes[3]; + const auto* residual = dst->src[0] == nodes[2] ? dst->src[1] : dst->src[0]; + const bool modulation = residual == nodes[1]; + const bool alias = (overlaps(dst, x) && dst->data != x->data) || overlaps(dst, w) || overlaps(dst, gate) || + (!modulation && overlaps(dst, residual) && dst->data != residual->data); + auto* output = static_cast(alias ? ggml_backend_buffer_get_base(execution->workspace) : dst->data); + float eps; + std::memcpy(&eps, nodes[0]->op_params, sizeof(eps)); + auto launch = x->ne[0] < 1024 ? (modulation ? kernel<256, true> : kernel<256, false>) + : (modulation ? kernel<1024, true> : kernel<1024, false>); + const auto stream = static_cast(context.stream); + const int threads = x->ne[0] < 1024 ? 256 : 1024; + launch<<>>( + static_cast(x->data), static_cast(w->data), static_cast(gate->data), + modulation ? nullptr : static_cast(residual->data), output, int(x->ne[0]), eps); + auto error = cuda_last_error(); + 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; + } + } + + int rms_affine_match(const ggml_cgraph* graph, int first) { + return sd_rms_affine::match(nullptr, graph, first, nullptr); + } + + ggml_status rms_affine_pack(CUDAStream stream, ggml_tensor* const* nodes, void* output) { + const auto* x = nodes[0]->src[0]; + const auto* residual = nodes[3]->src[0] == nodes[2] ? nodes[3]->src[1] : nodes[3]->src[0]; + const bool modulation = residual == nodes[1]; + auto launch = x->ne[0] < 1024 ? (modulation ? sd_rms_affine::kernel<256, true, CUDABFloat16> : sd_rms_affine::kernel<256, false, CUDABFloat16>) + : (modulation ? sd_rms_affine::kernel<1024, true, CUDABFloat16> : sd_rms_affine::kernel<1024, false, CUDABFloat16>); + float eps; + std::memcpy(&eps, nodes[0]->op_params, sizeof(eps)); + const int threads = x->ne[0] < 1024 ? 256 : 1024; + launch<<>>( + static_cast(x->data), static_cast(nodes[1]->src[1]->data), + static_cast(nodes[2]->src[1]->data), modulation ? nullptr : static_cast(residual->data), + static_cast(output), int(x->ne[0]), eps); + return cuda_last_error() == cuda_success ? GGML_STATUS_SUCCESS : GGML_STATUS_FAILED; + } + + void register_rms_affine(ggml_backend_dev_t device) { + ggml_ext_kernel_desc desc = {}; + desc.abi_version = GGML_EXT_ABI_VERSION; + desc.name = "sd.rms_affine_residual." SD_CUDA_NAME ".v1"; + desc.op = GGML_OP_RMS_NORM; + desc.priority = 300; + desc.match = sd_rms_affine::match; + desc.optimize = sd_rms_affine::optimize; + desc.prepare = sd_rms_affine::prepare; + desc.destroy = sd_rms_affine::destroy; + desc.execute = sd_rms_affine::execute; + GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE)); + } +} diff --git a/src/core/cuda/rms_affine.h b/src/core/cuda/rms_affine.h new file mode 100644 index 000000000..2bd1abcdc --- /dev/null +++ b/src/core/cuda/rms_affine.h @@ -0,0 +1,11 @@ +#ifndef __SD_CORE_CUDA_RMS_AFFINE_H__ +#define __SD_CORE_CUDA_RMS_AFFINE_H__ + +#include "runtime.h" + +namespace SD_CUDA_NAMESPACE { + int rms_affine_match(const ggml_cgraph* graph, int first); + ggml_status rms_affine_pack(CUDAStream stream, ggml_tensor* const* nodes, void* output); +} + +#endif