Commit d308544 for stable-diffusion.cpp
commit d30854463dbc75fb3671209b79016bb0d8ac321e
Author: leejet <leejet714@gmail.com>
Date: Tue Nov 3 23:59:21 2026 +0800
perf: fuse Z-Image GPU operations and optimize BF16 GEMM (#2265)
diff --git a/CMakeLists.txt b/CMakeLists.txt
index 15bd917..048f273 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 f366899..0a95d6c 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 0000000..59f5073
--- /dev/null
+++ b/src/core/cuda/bf16.cu
@@ -0,0 +1,277 @@
+#include "bf16.h"
+#if defined(SD_CUDA_VENDOR_HIP)
+#include <hipblas/hipblas.h>
+#elif defined(SD_CUDA_VENDOR_MUSA)
+#include <mublas.h>
+#else
+#include <cublas_v2.h>
+#endif
+#include <climits>
+#include <cstdlib>
+#include <cstring>
+#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<ggml_cgraph*>(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<uintptr_t>(a->data), y = reinterpret_cast<uintptr_t>(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<float*>(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<hipblasHandle_t>(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<mublasHandle_t>(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<cublasHandle_t>(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<const bool*>(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_ext_gpu_blas_fn>(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<CUDAStream>(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<char*>(ggml_backend_buffer_get_base(execution->workspace));
+ const Workspace work(nodes, count);
+ auto* packed = reinterpret_cast<CUDABFloat16*>(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<uintptr_t>(gate) | reinterpret_cast<uintptr_t>(up)) % 16 == 0) {
+ pack_swiglu4<<<unsigned((size / 4 + 255) / 256), 256, 0, stream>>>(static_cast<const float4*>(gate),
+ static_cast<const float4*>(up),
+ reinterpret_cast<CUDABFloat162*>(packed), size / 4, scale);
+ } else {
+ pack_swiglu<<<unsigned((size + 255) / 256), 256, 0, stream>>>(static_cast<const float*>(gate),
+ static_cast<const float*>(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 0000000..b5da370
--- /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 <hip/hip_bf16.h>
+#elif defined(SD_CUDA_VENDOR_MUSA)
+#include <musa_bf16.h>
+#else
+#include <cuda_bf16.h>
+#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 0000000..41d7570
--- /dev/null
+++ b/src/core/cuda/bf16_gemm.cu
@@ -0,0 +1,25 @@
+#include <cutlass/bfloat16.h>
+#include <cutlass/gemm/device/gemm.h>
+#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<uintptr_t>(weight) | reinterpret_cast<uintptr_t>(input) | reinterpret_cast<uintptr_t>(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<float, 4, float, float>, cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>, 3>;
+ Gemm gemm;
+ const Gemm::Arguments args({int(rows), int(n), int(k)},
+ {static_cast<const cutlass::bfloat16_t*>(input), k},
+ {static_cast<const cutlass::bfloat16_t*>(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 0000000..6cb3650
--- /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 <cuda_runtime.h>
+#include <cstdint>
+
+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 291edd6..d757280 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 0000000..d9ac351
--- /dev/null
+++ b/src/core/cuda/rms_affine.cu
@@ -0,0 +1,168 @@
+#include <climits>
+#include <cstdlib>
+#include <cstring>
+#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<ggml_cgraph*>(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 <int Block, bool Modulate, typename T = float>
+ 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<uintptr_t>(a->data), y = reinterpret_cast<uintptr_t>(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<float*>(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<CUDAStream>(context.stream);
+ const int threads = x->ne[0] < 1024 ? 256 : 1024;
+ launch<<<unsigned(ggml_nrows(x)), threads, 0, stream>>>(
+ static_cast<const float*>(x->data), static_cast<const float*>(w->data), static_cast<const float*>(gate->data),
+ modulation ? nullptr : static_cast<const float*>(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<<<unsigned(ggml_nrows(x)), threads, 0, stream>>>(
+ static_cast<const float*>(x->data), static_cast<const float*>(nodes[1]->src[1]->data),
+ static_cast<const float*>(nodes[2]->src[1]->data), modulation ? nullptr : static_cast<const float*>(residual->data),
+ static_cast<CUDABFloat16*>(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 0000000..2bd1abc
--- /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