perf: fuse Z-Image GPU operations and optimize BF16 GEMM (#2265)

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()

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")

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));
+    }
+}

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

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
+}

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

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);

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));
+    }
+}

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