DarthAffe/stable-diffusion.cpp
1
0 Fork 0
mirror of https://github.com/leejet/stable-diffusion.cpp.git synced 2026-11-12 11:56:36 +00:00
Code Issues Packages Projects Releases Wiki Activity

refactor: share CUDA-style extension kernels across backends

Browse Source
This commit is contained in:
leejet committed 2026-11-01 00:21:35 +08:00
1 parent 7ab323ef25
commit d11410eee7
19 files changed
+1961 -1719
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

CMakeLists.txt
+18 -8
View File
@@ -345,8 +345,7 @@ if (SD_CUDA)
set_property(SOURCE src/core/ggml_extend_backend.cpp APPEND PROPERTY COMPILE_DEFINITIONS SD_USE_CUDA) set_property(SOURCE src/core/ggml_extend_backend.cpp APPEND PROPERTY COMPILE_DEFINITIONS SD_USE_CUDA)
if (NOT SD_USE_UPSTREAM_GGML) if (NOT SD_USE_UPSTREAM_GGML)
enable_language(CUDA) 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_target_properties(${SD_LIB} PROPERTIES CUDA_STANDARD 17 CUDA_STANDARD_REQUIRED ON)
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")
if (GGML_CUDA_CUTLASS_PATH) 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) 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 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) 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 APPEND PROPERTY COMPILE_DEFINITIONS SD_CUDA_FORCE_CUBLAS)
endif() 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() 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()
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) if (SD_VULKAN AND NOT SD_USE_UPSTREAM_GGML)
add_subdirectory(src/core/vulkan) 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() endif()
if (SD_HIPBLAS) if (SD_HIPBLAS)
find_package(hip REQUIRED CONFIG HINTS "$ENV{ROCM_PATH}" /opt/rocm) find_package(hip REQUIRED CONFIG HINTS "$ENV{ROCM_PATH}" /opt/rocm)
cmake/backend_extensions.h.in
+12
View File
@@ -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
docs/ggml_extensions.md
+31 -15
View File
@@ -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_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_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_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 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 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 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 must use the representation appropriate to their backend. Vulkan extensions use
recording service is not implemented yet. the command recording services described below.
## sd.cpp integration ## 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 `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 `RMSNorm -> Mul -> RoPE` fusion. The fusion accepts a one-dimensional weight
vector. Both interleaved and split-half RoPE, partial rotation, and strided input 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 rows are supported. Model construction uses the custom operator on CPU, CUDA,
and retains existing backend implementations elsewhere. HIP, MUSA, and Vulkan, and retains existing backend implementations elsewhere.
`src/core/cuda/rope.cu` supplies standalone CUDA RoPE, the same builtin override, `src/core/cuda/rope.cu` supplies standalone CUDA RoPE, the same builtin override,
and `RMSNorm -> Mul -> RoPE [-> Scale -> FP16 Copy]` fusion. CuTe attention shares 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 position-encoding strides are supported by the standalone kernel; RMSNorm fusion
requires contiguous position encoding and weights. Copy fusion uses workspace requires contiguous position encoding and weights. Copy fusion uses workspace
when the destination overlaps an input. These registrations do not require 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 CUTLASS. The old CUDA/HIP/MUSA RoPE dispatch is removed.
implementations.
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 `SD_USE_UPSTREAM_GGML` excludes these registrations. The stable-diffusion public C
API is unchanged. INT8 ConvRot retains its existing packed activation layout and 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 preserve their previous accumulation behavior. Fusion disable and force-cuBLAS
settings are honored. No CUTLASS headers are needed for the cuBLAS path. settings are honored. No CUTLASS headers are needed for the cuBLAS path.
The native CUDA CUTLASS GEMM implementation and its SwiGLU/row-concat matchers The native CUDA/HIP INT8 GEMM, standalone ConvRot, and associated fusion paths
have been removed from ggml. Native cuBLAS GEMM and standalone ConvRot remain have been removed from ggml. The shared application implementation calls cuBLAS
compatibility implementations for callers without these registrations. HIP or hipBLAS directly. HIP retains the existing architecture eligibility policy;
retains its native INT8/activation fusion paths. 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 `src/core/cuda/i8_gated_residual.cu` implements
`INT8 GEMM -> [Scale] -> Mul(gate) -> Add(residual) -> [Concat]` as a registered `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 ordinary execution. These fusions do not require CUTLASS and honor
`GGML_CUDA_DISABLE_FUSION`. `GGML_CUDA_DISABLE_FUSION`.
The corresponding scale, normalization, and unary-activation fusion rules in The corresponding scale, normalization, and unary-activation fusion rules have
ggml now only serve HIP. CUDA GEMM/SwiGLU/ConvRot chains are selected by the been removed from ggml for CUDA and HIP. GEMM/SwiGLU/ConvRot chains are selected by the
application INT8 extension before the standalone activation fusion. application INT8 extension before the standalone activation fusion.
`src/core/cuda/affine.cu` registers FP32 `Mul(gate) -> Add(residual)`, `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 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 or input addresses are safe for the pointwise kernel; partial overlap, broadcast
gate overlap, and FP16 copy overlap use scheduler workspace. These registrations gate overlap, and FP16 copy overlap use scheduler workspace. These registrations
require no CUTLASS and support reusable CUDA Graph replay. Their former native require no CUTLASS, compile for CUDA/HIP/MUSA, and support reusable CUDA Graph
fusion implementations now only serve HIP/MUSA. 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 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 `GGML_EXT_KERNEL_CAPTURE_SAFE`. The original registration function retains its
ggml
+1 -1
Submodule ggml updated: 64edc3b784...a294003351.
src/core/compute_workspace.cpp
+5 -1
View File
@@ -147,6 +147,9 @@ namespace sd {
append(usage); append(usage);
append(type); append(type);
} }
} else {
const auto alignment_offset = reinterpret_cast<uintptr_t>(tensor->data) % ggml_backend_get_alignment(external);
append(alignment_offset);
} }
const size_t index = ggml_hash_find(&graph->visited_hash_set, tensor); 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) 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); ggml_backend_buffer_set_usage(buffer, GGML_BACKEND_BUFFER_USAGE_WEIGHTS);
} }
copy->buffer = buffer; copy->buffer = buffer;
copy->data = reinterpret_cast<void*>(static_cast<uintptr_t>(1)); // Capability checks inspect alignment even before weights are resident.
copy->data = source->data ? source->data : reinterpret_cast<void*>(ggml_backend_get_alignment(external));
copy->extra = nullptr; copy->extra = nullptr;
} }
} }
src/core/cuda/CMakeLists.txt
+90
View File
@@ -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()
src/core/cuda/affine.cu
+169 -169
View File
@@ -1,183 +1,183 @@
#include <cuda_fp16.h>
#include <cuda_runtime.h>
#include <climits> #include <climits>
#include <cstdlib> #include <cstdlib>
#include <cstring> #include <cstring>
#include "ggml-ext.h" #include "ggml-ext.h"
#include "runtime.h"
namespace sd_affine_cuda { namespace SD_CUDA_NAMESPACE {
static bool disabled() { namespace sd_affine_cuda {
static const bool value = [] { const char* v = std::getenv("GGML_CUDA_DISABLE_FUSION"); return v && std::atoi(v) != 0; }(); static bool disabled() {
return value; 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;
} }
return (reinterpret_cast<uintptr_t>(tensor->data) + offset) % alignment == 0; static bool aligned(const ggml_tensor* tensor, size_t alignment) {
} size_t offset = 0;
static float param(const ggml_tensor* tensor, int index) { while (!tensor->data && tensor->view_src) {
float value; offset += tensor->view_offs;
std::memcpy(&value, tensor->op_params + index, sizeof(value)); tensor = tensor->view_src;
return value; }
} return (reinterpret_cast<uintptr_t>(tensor->data) + offset) % alignment == 0;
static bool contiguous_f32(const ggml_tensor* tensor) { }
return tensor && tensor->type == GGML_TYPE_F32 && ggml_is_contiguous(tensor) && aligned(tensor, 16); static float param(const ggml_tensor* tensor, int index) {
} float value;
static bool bounded(const ggml_tensor* tensor) { std::memcpy(&value, tensor->op_params + index, sizeof(value));
return ggml_nelements(tensor) > 0 && ggml_nelements(tensor) <= int64_t(INT_MAX) * 1024; return value;
} }
static int match(ggml_backend_dev_t, const ggml_cgraph* graph, int first, void*) { static bool contiguous_f32(const ggml_tensor* tensor) {
auto* g = const_cast<ggml_cgraph*>(graph); return tensor && tensor->type == GGML_TYPE_F32 && ggml_is_contiguous(tensor) && aligned(tensor, 16);
if (disabled() || first + 1 >= ggml_graph_n_nodes(g)) }
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<ggml_cgraph*>(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; 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) { static __global__ void gated_residual(const float4* x, const float4* gate, const float4* residual, float4* dst, int64_t n4, int64_t count4) {
const bool permuted = next->op == GGML_OP_PERMUTE; const int64_t i = int64_t(blockIdx.x) * blockDim.x + threadIdx.x;
const int count = permuted ? 3 : 2; if (i >= count4)
if (first + count > ggml_graph_n_nodes(g) || node->type != GGML_TYPE_F32 || node->ne[0] % 4 || return;
node->op_params[0] != (permuted ? 2 : 1) || !ggml_is_contiguous(node) || const float4 a = x[i], g = gate[i % n4], r = residual[i];
!contiguous_f32(node->src[0]) || !contiguous_f32(node->src[1])) // The graph performs two FP32 operations; contracting them into FMA changes rounding.
return 0; dst[i] = make_float4(__fadd_rn(__fmul_rn(a.x, g.x), r.x), __fadd_rn(__fmul_rn(a.y, g.y), r.y),
if (permuted && (next->src[0] != node || next->ne[0] != node->ne[0] || next->ne[1] != node->ne[2] || __fadd_rn(__fmul_rn(a.z, g.z), r.z), __fadd_rn(__fmul_rn(a.w, g.w), r.w));
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)) { template <bool Permuted>
const auto* src = node->src[0]; 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 auto* permute = ggml_graph_node(g, first + 2); const int64_t i = int64_t(blockIdx.x) * blockDim.x + threadIdx.x;
const auto* dst = ggml_graph_node(g, first + 3); if (i >= count4)
return ggml_are_same_shape(src, node) && src->type == GGML_TYPE_F32 && src->nb[0] == 4 && src->ne[0] % 4 == 0 && return;
aligned(src, 16) && (src->nb[1] | src->nb[2] | src->nb[3]) % 16 == 0 && const int64_t col = i % columns4, row = i / columns4;
next->op == GGML_OP_SCALE && next->src[0] == node && next->type == GGML_TYPE_F32 && param(next, 1) == 0.f && const int64_t token = row % tokens, head = row / tokens % heads, batch = row / (tokens * heads);
permute->op == GGML_OP_PERMUTE && permute->src[0] == next && permute->op_params[0] == 0 && const bool from_prefix = token < prefix_tokens;
permute->op_params[1] == 2 && permute->op_params[2] == 1 && permute->op_params[3] == 3 && const int64_t src_tokens = from_prefix ? prefix_tokens : tokens - prefix_tokens;
dst->op == GGML_OP_CPY && dst->src[0] == permute && dst->type == GGML_TYPE_F16 && const int64_t src_token = from_prefix ? token : token - prefix_tokens;
ggml_are_same_shape(dst, permute) && ggml_is_contiguous(dst) && aligned(dst, 4) && bounded(dst) const int64_t src_row = Permuted ? (batch * src_tokens + src_token) * heads + head
? 4 : (batch * heads + head) * src_tokens + src_token;
: 0; 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 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;
static __global__ void gated_residual(const float4* x, const float4* gate, const float4* residual, float4* dst, int64_t n4, int64_t count4) { if (index >= count)
const int64_t i = int64_t(blockIdx.x) * blockDim.x + threadIdx.x; return;
if (i >= count4) const int64_t col = index % dim, row = index / dim;
return; const int64_t token = row % tokens, head = row / tokens % heads, batch = row / (tokens * heads);
const float4 a = x[i], g = gate[i % n4], r = residual[i]; const float4 value = *reinterpret_cast<const float4*>(src + col + head * nb1 + token * nb2 + batch * nb3);
// The graph performs two FP32 operations; contracting them into FMA changes rounding. dst[index / 2] = __floats2half2_rn(__fmul_rn(value.x, scale), __fmul_rn(value.y, scale));
dst[i] = make_float4(__fadd_rn(__fmul_rn(a.x, g.x), r.x), __fadd_rn(__fmul_rn(a.y, g.y), r.y), dst[index / 2 + 1] = __floats2half2_rn(__fmul_rn(value.z, scale), __fmul_rn(value.w, scale));
__fadd_rn(__fmul_rn(a.z, g.z), r.z), __fadd_rn(__fmul_rn(a.w, g.w), r.w)); }
} static bool overlaps(const ggml_tensor* a, const ggml_tensor* b) {
template <bool Permuted> const auto x = reinterpret_cast<uintptr_t>(a->data), y = reinterpret_cast<uintptr_t>(b->data);
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) { return x <= y ? y - x < ggml_nbytes(a) : x - y < ggml_nbytes(b);
const int64_t i = int64_t(blockIdx.x) * blockDim.x + threadIdx.x; }
if (i >= count4) static ggml_status prepare(ggml_backend_t, ggml_tensor* const* nodes, int count, void**, size_t* workspace, void*) {
return; *workspace = ggml_nbytes(nodes[count - 1]);
const int64_t col = i % columns4, row = i / columns4; return GGML_STATUS_SUCCESS;
const int64_t token = row % tokens, head = row / tokens % heads, batch = row / (tokens * heads); }
const bool from_prefix = token < prefix_tokens; static void destroy(void*, void*) {}
const int64_t src_tokens = from_prefix ? prefix_tokens : tokens - prefix_tokens; static ggml_status execute(const ggml_ext_execution_context* execution, ggml_tensor* const* nodes, int count, void*, void*) {
const int64_t src_token = from_prefix ? token : token - prefix_tokens; ggml_ext_gpu_context context;
const int64_t src_row = Permuted ? (batch * src_tokens + src_token) * heads + head if (!cuda_context(execution->backend, &context))
: (batch * heads + head) * src_tokens + src_token; return GGML_STATUS_FAILED;
const float4 value = (from_prefix ? prefix : current)[src_row * columns4 + col]; auto stream = static_cast<CUDAStream>(context.stream);
dst[2 * i] = __floats2half2_rn(value.x, value.y); const auto* node = nodes[0];
dst[2 * i + 1] = __floats2half2_rn(value.z, value.w); const auto* a = node->src[0];
} const auto* b = node->op == GGML_OP_CONT ? nullptr : node->src[1];
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) { auto* dst = nodes[count - 1];
const int64_t index = (int64_t(blockIdx.x) * blockDim.x + threadIdx.x) * 4; const auto* residual = node->op == GGML_OP_MUL ? (dst->src[0] == node ? dst->src[1] : dst->src[0]) : nullptr;
if (index >= count) const bool gate = residual != nullptr;
return; const bool alias = (overlaps(dst, a) && !(gate && dst->data == a->data)) || (b && overlaps(dst, b)) ||
const int64_t col = index % dim, row = index / dim; (residual && overlaps(dst, residual) && dst->data != residual->data);
const int64_t token = row % tokens, head = row / tokens % heads, batch = row / (tokens * heads); void* output = alias ? ggml_backend_buffer_get_base(execution->workspace) : dst->data;
const float4 value = *reinterpret_cast<const float4*>(src + col + head * nb1 + token * nb2 + batch * nb3); const int64_t count4 = ggml_nelements(dst) / 4;
dst[index / 2] = __floats2half2_rn(__fmul_rn(value.x, scale), __fmul_rn(value.y, scale)); const unsigned blocks = unsigned((count4 + 255) / 256);
dst[index / 2 + 1] = __floats2half2_rn(__fmul_rn(value.z, scale), __fmul_rn(value.w, scale)); if (gate) {
} gated_residual<<<blocks, 256, 0, stream>>>(static_cast<const float4*>(a->data), static_cast<const float4*>(b->data),
static bool overlaps(const ggml_tensor* a, const ggml_tensor* b) { static_cast<const float4*>(residual->data), static_cast<float4*>(output), a->ne[0] / 4, count4);
const auto x = reinterpret_cast<uintptr_t>(a->data), y = reinterpret_cast<uintptr_t>(b->data); } else if (node->op == GGML_OP_CONCAT) {
return x <= y ? y - x < ggml_nbytes(a) : x - y < ggml_nbytes(b); const bool permuted = count == 3;
} auto kernel = permuted ? concat_cast<true> : concat_cast<false>;
static ggml_status prepare(ggml_backend_t, ggml_tensor* const* nodes, int count, void**, size_t* workspace, void*) { kernel<<<blocks, 256, 0, stream>>>(static_cast<const float4*>(a->data), static_cast<const float4*>(b->data),
*workspace = ggml_nbytes(nodes[count - 1]); static_cast<half2*>(output), node->ne[0] / 4, a->ne[permuted ? 2 : 1], node->ne[permuted ? 2 : 1],
return GGML_STATUS_SUCCESS; node->ne[permuted ? 1 : 2], count4);
} } else {
static void destroy(void*, void*) {} scale_permute<<<blocks, 256, 0, stream>>>(static_cast<const float*>(a->data), static_cast<half2*>(output),
static ggml_status execute(const ggml_ext_execution_context* execution, ggml_tensor* const* nodes, int count, void*, void*) { 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 reg = ggml_backend_dev_backend_reg(ggml_backend_get_device(execution->backend)); }
auto query = reinterpret_cast<ggml_ext_cuda_context_fn>(ggml_backend_reg_get_proc_address(reg, "ggml_ext_cuda_context_v1")); auto error = cuda_last_error();
ggml_ext_cuda_context context; if (error == cuda_success && alias)
if (!query || !query(execution->backend, &context)) error = cuda_copy_async(dst->data, output, ggml_nbytes(dst), stream);
return GGML_STATUS_FAILED; return error == cuda_success ? GGML_STATUS_SUCCESS : GGML_STATUS_FAILED;
auto stream = static_cast<cudaStream_t>(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<<<blocks, 256, 0, stream>>>(static_cast<const float4*>(a->data), static_cast<const float4*>(b->data),
static_cast<const float4*>(residual->data), static_cast<float4*>(output), a->ne[0] / 4, count4);
} else if (node->op == GGML_OP_CONCAT) {
const bool permuted = count == 3;
auto kernel = permuted ? concat_cast<true> : concat_cast<false>;
kernel<<<blocks, 256, 0, stream>>>(static_cast<const float4*>(a->data), static_cast<const float4*>(b->data),
static_cast<half2*>(output), node->ne[0] / 4, a->ne[permuted ? 2 : 1], node->ne[permuted ? 2 : 1],
node->ne[permuted ? 1 : 2], count4);
} else {
scale_permute<<<blocks, 256, 0, stream>>>(static_cast<const float*>(a->data), static_cast<half2*>(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 = 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) { void register_affine(ggml_backend_dev_t device) {
ggml_ext_kernel_desc desc = {}; ggml_ext_kernel_desc desc = {};
desc.abi_version = GGML_EXT_ABI_VERSION; desc.abi_version = GGML_EXT_ABI_VERSION;
desc.priority = 200; desc.priority = 200;
desc.match = sd_affine_cuda::match; desc.match = sd_affine_cuda::match;
desc.prepare = sd_affine_cuda::prepare; desc.prepare = sd_affine_cuda::prepare;
desc.destroy = sd_affine_cuda::destroy; desc.destroy = sd_affine_cuda::destroy;
desc.execute = sd_affine_cuda::execute; desc.execute = sd_affine_cuda::execute;
desc.name = "sd.gated_residual.cuda.v1"; desc.name = "sd.gated_residual." SD_CUDA_NAME ".v1";
desc.op = GGML_OP_MUL; desc.op = GGML_OP_MUL;
GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE)); 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.name = "sd.concat_cast." SD_CUDA_NAME ".v1";
desc.op = GGML_OP_CONCAT; desc.op = GGML_OP_CONCAT;
GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE)); 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.name = "sd.cont_scale_permute_cast." SD_CUDA_NAME ".v1";
desc.op = GGML_OP_CONT; desc.op = GGML_OP_CONT;
GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE)); GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
}
} }
src/core/cuda/convrot.cu
+271 -290
View File
@@ -2,314 +2,295 @@
#include <climits> #include <climits>
#include <cstdlib> #include <cstdlib>
#include <cstring> #include <cstring>
#include <mutex>
#include <string> #include <string>
#include <unordered_map>
#include "convrot_kernel.h" #include "convrot_kernel.h"
#include "core/int8_ops.h" #include "core/int8_ops.h"
#include "runtime.h"
namespace sd_convrot { namespace SD_CUDA_NAMESPACE {
struct DeviceInfo { namespace sd_convrot {
int ordinal = -1;
size_t shared = 0; static bool fusion_disabled() {
}; static const bool disabled = [] { const char* v = std::getenv("GGML_CUDA_DISABLE_FUSION"); return v && std::atoi(v) != 0; }();
static DeviceInfo device_info(ggml_backend_dev_t device) { return disabled;
static std::mutex mutex;
static std::unordered_map<ggml_backend_dev_t, DeviceInfo> cache;
std::lock_guard<std::mutex> lock(mutex);
auto found = cache.find(device);
if (found != cache.end())
return found->second;
DeviceInfo info;
auto query = reinterpret_cast<ggml_ext_cuda_device_fn>(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;
} }
return true; static float param_f32(const ggml_tensor* tensor, int index) {
} float value;
static bool supported(const Operation& op, const DeviceInfo& device) { std::memcpy(&value, tensor->op_params + index, sizeof(value));
const size_t extra = op.kind == NORM ? 384 : op.kind == RMS_MOD ? 256 return value;
: 128; }
return device.ordinal >= 0 && size_t(op.k) * sizeof(float) + extra <= device.shared; static bool f32(const ggml_tensor* tensor) {
} return tensor && tensor->type == GGML_TYPE_F32 && ggml_is_contiguous(tensor);
static int match(ggml_backend_dev_t device, const ggml_cgraph* graph, int first, void*) { }
if (fusion_disabled()) static bool vector(const ggml_tensor* tensor, int64_t k) {
return 0; return f32(tensor) && tensor->ne[0] == k && ggml_nelements(tensor) == k;
auto* g = const_cast<ggml_cgraph*>(graph); }
const int available = ggml_graph_n_nodes(g) - first; enum Kind { SCALE,
const auto* node = ggml_graph_node(g, first); NORM,
const int count = node->op == GGML_OP_RMS_NORM ? 6 : node->op == GGML_OP_NORM || node->op == GGML_OP_UNARY ? 3 RMS_MOD,
: 2; SWIGLU,
if (available < count) SILU,
return 0; SIGMOID };
ggml_tensor* nodes[6]; struct Operation {
for (int j = 0; j < count; ++j) Kind kind;
nodes[j] = ggml_graph_node(g, first + j); const ggml_tensor* input = nullptr;
Operation op; const ggml_tensor* up = nullptr;
return describe(nodes, count, op) && supported(op, device_info(device)) ? count : 0; const ggml_tensor* modulation = nullptr;
} const ggml_tensor* weight = nullptr;
static void optimize(ggml_backend_dev_t device, ggml_cgraph* graph, void*) { const ggml_tensor* shift = nullptr;
if (fusion_disabled()) int64_t k, rows, stride, up_stride, gate_offset = 0, up_offset = 0;
return; float eps = 0.f, scale = 1.f, bias = 0.f;
for (int i = 2; i < ggml_graph_n_nodes(graph); ++i) { };
auto* quant = ggml_graph_node(graph, i); static bool describe(ggml_tensor* const* nodes, int count, Operation& op) {
int group; int group;
const auto* first = nodes[0];
const auto* quant = nodes[count - 1];
if (!sd_ggml_convrot_describe(quant, &group) || group != 256) if (!sd_ggml_convrot_describe(quant, &group) || group != 256)
continue; return false;
ggml_tensor* chain[6] = {}; const auto* value = quant->src[0];
int count = 0; op.k = value->ne[0];
auto* value = quant->src[0]; op.rows = ggml_nrows(value);
if (value->op == GGML_OP_ADD && value->src[0]->op == GGML_OP_ADD) { op.input = first->src[0];
auto* add = value->src[0]; if (!f32(op.input) || !f32(first) || op.k > INT_MAX || op.rows > INT_MAX)
auto* weight = add->src[0]; return false;
if (weight->op == GGML_OP_MUL) { op.stride = op.input->ne[0];
chain[0] = weight->src[0]; op.up_stride = op.k;
chain[1] = weight; if (count == 2 && first->op == GGML_OP_SCALE && value == first) {
chain[2] = add->src[1]; op.kind = SCALE;
chain[3] = add; op.scale = param_f32(first, 0);
chain[4] = value; op.bias = param_f32(first, 1);
chain[5] = quant; } else if (count == 3 && first->op == GGML_OP_NORM && nodes[1]->op == GGML_OP_MUL && nodes[1]->src[0] == first && value == nodes[1]) {
count = 6; 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) { } else if (count == 3 && first->op == GGML_OP_UNARY && nodes[1]->op == GGML_OP_MUL && value == nodes[1] &&
for (int src = 0; src < 2; ++src) (nodes[1]->src[0] == first || nodes[1]->src[1] == first)) {
if (value->src[src]->op == GGML_OP_UNARY) { auto unary = ggml_get_unary_op(first);
chain[0] = value->src[src]; if (unary != GGML_UNARY_OP_SILU && unary != GGML_UNARY_OP_SIGMOID)
chain[1] = value; return false;
chain[2] = quant; op.kind = unary == GGML_UNARY_OP_SILU ? SILU : SIGMOID;
count = 3; op.up = nodes[1]->src[0] == first ? nodes[1]->src[1] : nodes[1]->src[0];
break; 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<ggml_cgraph*>(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; Operation op;
if (!count || !describe(chain, count, op) || !supported(op, device_info(device))) return describe(nodes, count, op) && supported(op, cuda_device_info(device)) ? count : 0;
continue; }
int begin = -1; static void optimize(ggml_backend_dev_t device, ggml_cgraph* graph, void*) {
for (int j = std::max(0, i - 64); j < i; ++j) if (fusion_disabled())
if (ggml_graph_node(graph, j) == chain[0]) { return;
begin = j; for (int i = 2; i < ggml_graph_n_nodes(graph); ++i) {
break; 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) Operation op;
continue; if (!count || !describe(chain, count, op) || !supported(op, cuda_device_info(device)))
// Move the owning producer with its in-place consumers so alias checks retain their storage relationship. continue;
for (int part = 1; part < count; ++part) { int begin = -1;
int next = -1; for (int j = std::max(0, i - 64); j < i; ++j)
for (int j = begin + part; j <= i; ++j) if (ggml_graph_node(graph, j) == chain[0]) {
if (ggml_graph_node(graph, j) == chain[part]) { begin = j;
next = j;
break; break;
} }
if (next < 0 || (next > begin + part && !ggml_ext_graph_delay(graph, begin, part, next))) if (begin < 0)
break; continue;
begin = next - part; // 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 {
struct Plan { ggml_ext_gpu_device_info device;
DeviceInfo device; Operation op;
Operation op; bool configured = false;
bool configured = false; };
}; static ggml_status prepare(ggml_backend_t backend, ggml_tensor* const* nodes, int count, void** state, size_t* workspace, void*) {
static ggml_status prepare(ggml_backend_t backend, ggml_tensor* const* nodes, int count, void** state, size_t* workspace, void*) { auto* plan = new Plan{};
auto* plan = new Plan{}; *state = plan;
*state = plan; plan->device = cuda_device_info(ggml_backend_get_device(backend));
plan->device = device_info(ggml_backend_get_device(backend)); if (!describe(nodes, count, plan->op) || !supported(plan->op, plan->device))
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<Plan*>(state);
}
static bool overlaps(const ggml_tensor* a, const ggml_tensor* b) {
if (!b)
return false;
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);
}
template <bool Swiglu = false, bool Normalize = false, bool Scaled = false, bool Sigmoid = false, bool RMS = false>
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<Swiglu, false, Normalize, Scaled, warp, float, Sigmoid, RMS>,
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<const float*>(t->data) : nullptr; };
quantize<Swiglu, false, Normalize, Scaled, warp, float, Sigmoid, RMS><<<unsigned(op.rows), threads, size_t(op.k) * sizeof(float), stream>>>(
data(op.input) + op.gate_offset, output, reinterpret_cast<float*>(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<Plan*>(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_ext_cuda_context_fn>(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<cudaStream_t>(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<int8_t*>(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<false, false, true>(plan, op, output, stream);
break;
case NORM:
error = launch<false, true>(plan, op, output, stream);
break;
case RMS_MOD:
error = launch<false, false, false, false, true>(plan, op, output, stream);
break;
case SWIGLU:
case SILU:
error = launch<true>(plan, op, output, stream);
break;
case SIGMOID:
error = launch<true, false, false, true>(plan, op, output, stream);
break;
default:
return GGML_STATUS_FAILED; return GGML_STATUS_FAILED;
*workspace = ggml_nbytes(nodes[count - 1]);
return GGML_STATUS_SUCCESS;
}
static void destroy(void* state, void*) {
delete static_cast<Plan*>(state);
}
static bool overlaps(const ggml_tensor* a, const ggml_tensor* b) {
if (!b)
return false;
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);
}
template <bool Swiglu = false, bool Normalize = false, bool Scaled = false, bool Sigmoid = false, bool RMS = false>
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<Swiglu, false, Normalize, Scaled, warp, float, Sigmoid, RMS>,
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<const float*>(t->data) : nullptr; };
quantize<Swiglu, false, Normalize, Scaled, warp, float, Sigmoid, RMS><<<unsigned(op.rows), threads, size_t(op.k) * sizeof(float), stream>>>(
data(op.input) + op.gate_offset, output, reinterpret_cast<float*>(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<Plan*>(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<CUDAStream>(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<int8_t*>(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<false, false, true>(plan, op, output, stream);
break;
case NORM:
error = launch<false, true>(plan, op, output, stream);
break;
case RMS_MOD:
error = launch<false, false, false, false, true>(plan, op, output, stream);
break;
case SWIGLU:
case SILU:
error = launch<true>(plan, op, output, stream);
break;
case SIGMOID:
error = launch<true, false, false, true>(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) { 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}) { for (auto op : {GGML_OP_SCALE, GGML_OP_NORM, GGML_OP_RMS_NORM, GGML_OP_GLU, GGML_OP_UNARY}) {
ggml_ext_kernel_desc desc = {}; ggml_ext_kernel_desc desc = {};
desc.abi_version = GGML_EXT_ABI_VERSION; desc.abi_version = GGML_EXT_ABI_VERSION;
const std::string name = std::string("sd.convrot_fusion.") + ggml_op_name(op) + ".cuda.v1"; const std::string name = std::string("sd.convrot_fusion.") + ggml_op_name(op) + "." SD_CUDA_NAME ".v1";
desc.name = name.c_str(); desc.name = name.c_str();
desc.op = op; desc.op = op;
desc.priority = 225; desc.priority = 225;
desc.optimize = op == GGML_OP_SCALE ? sd_convrot::optimize : nullptr; desc.optimize = op == GGML_OP_SCALE ? sd_convrot::optimize : nullptr;
desc.match = sd_convrot::match; desc.match = sd_convrot::match;
desc.prepare = sd_convrot::prepare; desc.prepare = sd_convrot::prepare;
desc.execute = sd_convrot::execute; desc.execute = sd_convrot::execute;
desc.destroy = sd_convrot::destroy; desc.destroy = sd_convrot::destroy;
GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE)); GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
}
} }
} }
src/core/cuda/convrot_kernel.h
+224 -222
View File
@@ -1,249 +1,251 @@
#ifndef __SD_CORE_CUDA_CONVROT_KERNEL_H__ #ifndef __SD_CORE_CUDA_CONVROT_KERNEL_H__
#define __SD_CORE_CUDA_CONVROT_KERNEL_H__ #define __SD_CORE_CUDA_CONVROT_KERNEL_H__
#include <cuda_runtime.h>
#include <cstdint> #include <cstdint>
namespace sd_convrot { #include "runtime.h"
static __device__ __forceinline__ float warp_reduce_sum(float x) { namespace SD_CUDA_NAMESPACE {
namespace sd_convrot {
static __device__ __forceinline__ float warp_reduce_sum(float x) {
#pragma unroll #pragma unroll
for (int offset = 16; offset > 0; offset >>= 1) for (int offset = 16; offset > 0; offset >>= 1)
x += __shfl_xor_sync(0xffffffff, x, offset, 32); x += cuda_shuffle_xor(x, offset, 32);
return x; 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);
} }
return x; static __device__ __forceinline__ float2 warp_reduce_sum(float2 x) {
}
template <int Width>
static __device__ __forceinline__ float warp_reduce_max(float x) {
#pragma unroll #pragma unroll
for (int offset = Width / 2; offset > 0; offset >>= 1) for (int offset = 16; offset > 0; offset >>= 1) {
x = fmaxf(x, __shfl_xor_sync(0xffffffff, x, offset, Width)); x.x += cuda_shuffle_xor(x.x, offset, 32);
return x; x.y += cuda_shuffle_xor(x.y, offset, 32);
}
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 <bool swiglu, bool dequantize = false, bool normalize = false, bool scaled = false, bool warp_rotation = false, typename T = float, bool sigmoid_gate = false, bool rms_modulation = false>
__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;
}
} }
return x;
}
template <int Width>
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(); __syncthreads();
const float sum = warp_reduce_sum(sums[tid % 32]); return warp_reduce_sum(shared[threadIdx.x % 32]);
rms_scale = rsqrtf(sum / k + eps); }
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) { template <bool swiglu, bool dequantize = false, bool normalize = false, bool scaled = false, bool warp_rotation = false, typename T = float, bool sigmoid_gate = false, bool rms_modulation = false>
const int lane = tid % 32; __global__ void quantize(
for (int64_t group = tid / 32; group < k / group_size; group += blockDim.x / 32) { const T* src,
float v[8]; int8_t* dst,
#pragma unroll float* scales,
for (int j = 0; j < 8; ++j) { int64_t k,
const int64_t i = group * group_size + j * 32 + lane; int64_t rows,
if constexpr (!swiglu) { int64_t dst_stride,
float value = src_row[i]; const float* up,
if constexpr (rms_modulation) { int64_t src_stride,
value = __fmul_rn(__fmul_rn(value, rms_scale), rms_weight[i]); int64_t up_stride,
value = __fadd_rn(__fadd_rn(value, __fmul_rn(value, modulation[i])), rms_shift[i]); const float* activation_scales = nullptr,
} const float* weight_scales = nullptr,
if constexpr (scaled) { const float* modulation = nullptr,
value = __fmaf_rn(value, input_scale, input_bias); float eps = 0.0f,
} float input_scale = 1.0f,
v[j] = value * (1.0f / 16.0f); float input_bias = 0.0f,
} else { const float* rms_weight = nullptr,
float gate_value, up_value; const float* rms_shift = nullptr) {
if constexpr (dequantize) { constexpr int group_size = 256;
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]; const int64_t row = blockIdx.x;
gate_value *= weight_scales[i]; const int tid = threadIdx.x;
up_value *= weight_scales[k + i]; if (row >= rows) {
} else { return;
gate_value = src_row[i]; }
up_value = up[row * up_stride + i];
} extern __shared__ float values[];
const float activated = sigmoid_gate ? 1.0f / (1.0f + expf(-gate_value)) : (gate_value / (1.0f + expf(-gate_value))); __shared__ float maxima[32];
v[j] = (activated * up_value) * (1.0f / 16.0f); 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. __syncthreads();
#pragma unroll const float sum = warp_reduce_sum(sums[tid % 32]);
for (int stride = 1; stride <= 4; stride *= 4) { rms_scale = rsqrtf(sum / k + eps);
const int base = lane & ~(3 * stride); }
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 #pragma unroll
for (int j = 0; j < 8; ++j) { for (int j = 0; j < 8; ++j) {
const float a = __shfl_sync(0xffffffff, v[j], base, 32); const int64_t i = group * group_size + j * 32 + lane;
const float b = __shfl_sync(0xffffffff, v[j], base + stride, 32); if constexpr (!swiglu) {
const float c = __shfl_sync(0xffffffff, v[j], base + 2 * stride, 32); float value = src_row[i];
const float d = __shfl_sync(0xffffffff, v[j], base + 3 * stride, 32); if constexpr (rms_modulation) {
v[j] = convrot_h4_cuda(a, b, c, d, (lane / stride) % 4); 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 __syncthreads();
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);
}
} else { } else {
for (int64_t i = tid; i < k; i += blockDim.x) { constexpr int group_threads = group_size / 4;
const float value = scaled ? __fmaf_rn(src_row[i], input_scale, input_bias) : src_row[i]; constexpr int groups_per_wave = 1024 / group_threads;
values[i] = value * (1.0f / 16.0f); 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(); __syncthreads();
const int group_lane = tid % group_threads; const float row_scale = maxima[0];
const int wave_group = tid / group_threads; const float inv_scale = row_scale == 0.0f ? 0.0f : 1.0f / row_scale;
const int64_t groups = k / group_size; for (int64_t i = tid; i < k; i += blockDim.x) {
for (int64_t group_base = 0; group_base < groups; group_base += groups_per_wave) { int value = row_scale == 0.0f ? 0 : __float2int_rn(values[i] * inv_scale);
const int64_t group = group_base + wave_group; value = max(-127, min(127, value));
#pragma unroll dst_row[i] = (int8_t)value;
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 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 #endif
src/core/cuda/cute_attention.cu
+3 -3
View File
@@ -588,7 +588,7 @@ static int attention_match(ggml_backend_dev_t device, const ggml_cgraph* graph,
offset = 5; offset = 5;
} else if (graph_node(graph, first)->op == GGML_OP_RMS_NORM) { } else if (graph_node(graph, first)->op == GGML_OP_RMS_NORM) {
if (first + 3 >= ggml_graph_n_nodes(const_cast<ggml_cgraph*>(graph)) || if (first + 3 >= ggml_graph_n_nodes(const_cast<ggml_cgraph*>(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)) graph_node(graph, first + 3)->src[0] != graph_node(graph, first + 2))
return 0; return 0;
offset = 3; offset = 3;
@@ -608,7 +608,7 @@ static void attention_optimize(ggml_backend_dev_t device, ggml_cgraph* graph, vo
continue; continue;
for (int j = i - 3; j >= 0 && j >= i - 32; --j) { for (int j = i - 3; j >= 0 && j >= i - 32; --j) {
if (attn->src[0] == graph_node(graph, j + 2) && 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)) && 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)) ggml_ext_graph_delay(graph, j - 2, 5, i))
break; 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) if (sd_cuda_i8_rope_launch(stream, nodes, work[0], scale) != cudaSuccess)
return GGML_STATUS_FAILED; return GGML_STATUS_FAILED;
} else if (i == 0 && plan.index == 3) { } 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; return GGML_STATUS_FAILED;
} else if (i != 0 && src->type == GGML_TYPE_F16 && ggml_is_contiguous(src) && uintptr_t(src->data) % 16 == 0) { } else if (i != 0 && src->type == GGML_TYPE_F16 && ggml_is_contiguous(src) && uintptr_t(src->data) % 16 == 0) {
inputs[i] = static_cast<const half*>(src->data); inputs[i] = static_cast<const half*>(src->data);
src/core/cuda/i8_projection.cu
+1 -1
View File
@@ -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) { 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 || 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; return false;
int32_t interleaved; int32_t interleaved;
std::memcpy(&interleaved, ggml_ext_op_get(rope) ? ggml_ext_op_attributes(rope) : rope->op_params, sizeof(interleaved)); std::memcpy(&interleaved, ggml_ext_op_get(rope) ? ggml_ext_op_attributes(rope) : rope->op_params, sizeof(interleaved));
src/core/cuda/int8.cu
+677 -671
View File
File diff suppressed because it is too large. Load diff
src/core/cuda/registration.cu
+28
View File
@@ -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
}
}
src/core/cuda/registration.h
+13
View File
@@ -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
src/core/cuda/rope.cu
+289 -289
View File
@@ -1,188 +1,188 @@
#include "core/cuda/rope.h" #include "core/cuda/rope.h"
#include "core/ggml_extensions.h" #include "core/ggml_extensions.h"
#include <cuda_fp16.h>
#include <climits> #include <climits>
#include <cstdlib> #include <cstdlib>
#include <cstring> #include <cstring>
#include <type_traits> #include <type_traits>
static __device__ __forceinline__ float sd_warp_sum(float v) { namespace SD_CUDA_NAMESPACE {
for (int mask = 16; mask > 0; mask >>= 1) static __device__ __forceinline__ float sd_warp_sum(float v) {
v = __fadd_rn(v, __shfl_xor_sync(0xffffffff, v, mask)); for (int mask = 16; mask > 0; mask >>= 1)
return v; v = __fadd_rn(v, cuda_shuffle_xor(v, mask));
} return v;
template <int D, bool interleaved, typename T>
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; template <int D, bool interleaved, typename T>
const int64_t token = (row / heads) % tokens; static __global__ void rms_norm_mul_rope_apply_fused(
const int64_t batch = row / (heads * tokens); const float* x,
const float* x_row = x + head * nb1 + token * nb2 + batch * nb3; const float* weight,
float v[D / 32], sums[8] = {}; 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 #pragma unroll
for (int j = 0; j < D / 32; ++j) { for (int j = 0; j < D / 32; ++j) {
v[j] = x_row[j * 32 + lane]; v[j] = x_row[j * 32 + lane];
sums[j] = sd_warp_sum(__fmul_rn(v[j], v[j])); sums[j] = sd_warp_sum(__fmul_rn(v[j], v[j]));
} }
// Preserve the 256-thread RMSNorm kernel's reduction tree, including inactive warps. // Preserve the 256-thread RMSNorm kernel's reduction tree, including inactive warps.
const float sum = __fadd_rn( const float sum = __fadd_rn(
__fadd_rn(__fadd_rn(sums[0], sums[4]), __fadd_rn(sums[2], sums[6])), __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]))); __fadd_rn(__fadd_rn(sums[1], sums[5]), __fadd_rn(sums[3], sums[7])));
const float scale = rsqrtf(sum / float(D) + eps); const float scale = rsqrtf(sum / float(D) + eps);
T* out = dst + ((batch * heads + head) * tokens + token) * D; T* out = dst + ((batch * heads + head) * tokens + token) * D;
#pragma unroll #pragma unroll
for (int j = 0; j < D / 32; ++j) { for (int j = 0; j < D / 32; ++j) {
const int col = j * 32 + lane; const int col = j * 32 + lane;
const float value = __fmul_rn(__fmul_rn(scale, v[j]), weight[col]); const float value = __fmul_rn(__fmul_rn(scale, v[j]), weight[col]);
float rotated = value; float rotated = value;
if constexpr (interleaved) { if constexpr (interleaved) {
const float other = __shfl_xor_sync(0xffffffff, value, 1, 32); const float other = cuda_shuffle_xor(value, 1, 32);
if (col < 2 * rot_half) { if (col < 2 * rot_half) {
const float* m = pe + 4 * (token * rot_half + col / 2); const float* m = pe + 4 * (token * rot_half + col / 2);
const float4 matrix = make_float4(m[0], m[1], m[2], m[3]); const float4 matrix = make_float4(m[0], m[1], m[2], m[3]);
rotated = lane % 2 == 0 rotated = lane % 2 == 0
? __fadd_rn(__fmul_rn(value, matrix.x), __fmul_rn(other, matrix.y)) ? __fadd_rn(__fmul_rn(value, matrix.x), __fmul_rn(other, matrix.y))
: __fadd_rn(__fmul_rn(other, matrix.z), __fmul_rn(value, matrix.w)); : __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 constexpr (std::is_same<T, half>::value) {
if (col < 2 * rot_half) { out[col] = __float2half_rn(__fmul_rn(rotated, output_scale));
const int partner = col < rot_half ? col + rot_half : col - rot_half; } else {
const float other = __fmul_rn(__fmul_rn(scale, x_row[partner]), weight[partner]); out[col] = rotated;
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<T, half>::value) { }
out[col] = __float2half_rn(__fmul_rn(rotated, output_scale));
struct RopeApplyStrides {
int64_t x[4];
int64_t pe[4];
};
template <bool interleaved>
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 { } 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 { static bool rope_interleaved(const ggml_tensor* rope) {
int64_t x[4]; int32_t value;
int64_t pe[4]; std::memcpy(&value, ggml_ext_op_get(rope) ? ggml_ext_op_attributes(rope) : rope->op_params, sizeof(value));
}; return value != 0;
template <bool interleaved>
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; static float param_f32(const ggml_tensor* node, int index) {
const int64_t token = row % tokens; float value;
const int64_t head = (row / tokens) % heads; std::memcpy(&value, node->op_params + index, sizeof(value));
const int64_t batch = row / (tokens * heads); return value;
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]; static bool rope_is_node(const ggml_tensor* node) {
const float s = matrix[strides.pe[1]]; return node && (node->op == GGML_OP_ROPE_APPLY || ggml_ext_op_get(node) == sd_ggml_rope_op());
// 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_interleaved(const ggml_tensor* rope) { static bool rope_supports(ggml_backend_dev_t, const ggml_tensor* rope, void*) {
int32_t value; const auto* x = rope->src[0];
std::memcpy(&value, ggml_ext_op_get(rope) ? ggml_ext_op_attributes(rope) : rope->op_params, sizeof(value)); const auto* pe = rope->src[1];
return value != 0; 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] ||
static float param_f32(const ggml_tensor* node, int index) { rope->ne[1] != x->ne[2] || rope->ne[2] != x->ne[1] * x->ne[3] || rope->ne[3] != 1 ||
float value; (ggml_nelements(x) / 2 + 255) / 256 > INT_MAX) {
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))
return false; 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) { 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 || !sd_cuda_rope_is_node(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)) mul->src[0] != norm || rope->src[0] != mul || !rope_supports(nullptr, rope, nullptr))
return false; return false;
const auto* x = norm->src[0]; const auto* x = norm->src[0];
const auto* weight = mul->src[1]; const auto* weight = mul->src[1];
return x && weight && x->ne[0] >= 32 && x->ne[0] <= 256 && x->ne[0] % 32 == 0 && 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 && 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) && ggml_are_same_shape(norm, mul) && ggml_are_same_shape(x, norm) &&
weight->ne[0] == x->ne[0] && ggml_nelements(weight) == x->ne[0] && 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[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]) && x->nb[3] % sizeof(float) == 0 && ggml_is_contiguous(weight) && ggml_is_contiguous(rope->src[1]) &&
(ggml_nrows(x) + 3) / 4 <= INT_MAX; (ggml_nrows(x) + 3) / 4 <= INT_MAX;
} }
template <typename T, bool Interleaved> template <typename T, bool Interleaved>
static cudaError_t rms_rope_launch(cudaStream_t stream, ggml_tensor* const* nodes, T* output, float scale) { static CUDAError rms_rope_launch(CUDAStream stream, ggml_tensor* const* nodes, T* output, float scale) {
const auto* x = nodes[0]->src[0]; const auto* x = nodes[0]->src[0];
const auto* pe = nodes[2]->src[1]; const auto* pe = nodes[2]->src[1];
const int64_t rows = ggml_nrows(x); const int64_t rows = ggml_nrows(x);
#define SD_ROPE(D) \ #define SD_ROPE(D) \
case D: \ case D: \
rms_norm_mul_rope_apply_fused<D, Interleaved, T><<<static_cast<unsigned int>((rows + 3) / 4), dim3(32, 4), 0, stream>>>( \ rms_norm_mul_rope_apply_fused<D, Interleaved, T><<<static_cast<unsigned int>((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<const float*>(pe->data), output, x->ne[1], x->ne[2], rows, \ static_cast<const float*>(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); \ x->nb[1] / 4, x->nb[2] / 4, x->nb[3] / 4, int(pe->ne[2]), param_f32(nodes[0], 0), scale); \
break break
switch (x->ne[0]) { switch (x->ne[0]) {
SD_ROPE(32); SD_ROPE(32);
SD_ROPE(64); SD_ROPE(64);
SD_ROPE(96); SD_ROPE(96);
SD_ROPE(128); SD_ROPE(128);
SD_ROPE(160); SD_ROPE(160);
SD_ROPE(192); SD_ROPE(192);
SD_ROPE(224); SD_ROPE(224);
SD_ROPE(256); SD_ROPE(256);
default: default:
return cudaErrorInvalidValue; return cuda_invalid_value;
} }
#undef SD_ROPE #undef SD_ROPE
return cudaGetLastError(); return cuda_last_error();
}
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<half, true>(stream, nodes, static_cast<half*>(output), scale)
: rms_rope_launch<half, false>(stream, nodes, static_cast<half*>(output), scale);
} }
return rope_interleaved(nodes[2]) ? rms_rope_launch<float, true>(stream, nodes, static_cast<float*>(output), scale)
: rms_rope_launch<float, false>(stream, nodes, static_cast<float*>(output), scale);
}
static int rope_match(ggml_backend_dev_t, const ggml_cgraph* graph, int first, void*) { CUDAError rms_rope_launch(CUDAStream stream, ggml_tensor* const* nodes, void* output, bool fp16, float scale) {
static const bool disabled = [] {const char* value=std::getenv("GGML_CUDA_DISABLE_FUSION"); return value && std::atoi(value)!=0; }(); if (fp16) {
if (disabled) return rope_interleaved(nodes[2]) ? rms_rope_launch<half, true>(stream, nodes, static_cast<half*>(output), scale)
return 0; : rms_rope_launch<half, false>(stream, nodes, static_cast<half*>(output), scale);
auto* g = const_cast<ggml_cgraph*>(graph); }
if (first + 2 >= ggml_graph_n_nodes(g) || !sd_cuda_rope_fusion_supports( return rope_interleaved(nodes[2]) ? rms_rope_launch<float, true>(stream, nodes, static_cast<float*>(output), scale)
ggml_graph_node(g, first), ggml_graph_node(g, first + 1), ggml_graph_node(g, first + 2))) : rms_rope_launch<float, false>(stream, nodes, static_cast<float*>(output), scale);
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 int rope_match(ggml_backend_dev_t, const ggml_cgraph* graph, int first, void*) {
static const bool disabled = [] { static const bool disabled = [] {const char* value=std::getenv("GGML_CUDA_DISABLE_FUSION"); return value && std::atoi(value)!=0; }();
const char* value = std::getenv("GGML_CUDA_DISABLE_FUSION"); if (disabled)
return value && std::atoi(value) != 0; return 0;
}(); auto* g = const_cast<ggml_cgraph*>(graph);
if (disabled) if (first + 2 >= ggml_graph_n_nodes(g) || !rope_fusion_supports(
return; ggml_graph_node(g, first), ggml_graph_node(g, first + 1), ggml_graph_node(g, first + 2)))
const int total = ggml_graph_n_nodes(graph); return 0;
for (int i = 2; i + 1 < total; ++i) { if (first + 4 < ggml_graph_n_nodes(g)) {
const auto* mul = ggml_graph_node(graph, i); const auto* rope = ggml_graph_node(g, first + 2);
const auto* rope = ggml_graph_node(graph, i + 1); const auto* scale = ggml_graph_node(g, first + 3);
if (mul->op != GGML_OP_MUL || !sd_cuda_rope_is_node(rope) || rope->src[0] != mul || const auto* copy = ggml_graph_node(g, first + 4);
!mul->src[0] || mul->src[0]->op != GGML_OP_RMS_NORM) if (scale->op == GGML_OP_SCALE && copy->op == GGML_OP_CPY && scale->src[0] == rope && copy->src[0] == scale &&
continue; scale->type == GGML_TYPE_F32 && copy->type == GGML_TYPE_F16 && param_f32(scale, 1) == 0.f &&
const auto* norm = mul->src[0]; ggml_are_same_shape(rope, scale) && ggml_are_same_shape(rope, copy) && ggml_is_contiguous(copy))
for (int j = i - 1; j >= 0 && j >= i - 64; --j) { return 5;
if (ggml_graph_node(graph, j) != norm) }
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; continue;
if (j + 1 < i) { const auto* norm = mul->src[0];
if (j >= 2) { for (int j = i - 1; j >= 0 && j >= i - 64; --j) {
const auto* mat = ggml_graph_node(graph, j - 2); if (ggml_graph_node(graph, j) != norm)
const auto* view = ggml_graph_node(graph, j - 1); continue;
if (mat->op == GGML_OP_MUL_MAT && mat->src[0]->type == GGML_TYPE_I8 && if (j + 1 < i) {
view->op == GGML_OP_RESHAPE && view->src[0] == mat && norm->src[0] == view && if (j >= 2) {
ggml_ext_graph_delay(graph, j - 2, 3, i)) const auto* mat = ggml_graph_node(graph, j - 2);
break; 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*) { 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]); *workspace = ggml_nbytes(nodes[count - 1]);
return GGML_STATUS_SUCCESS; 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) { 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); 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); 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*) { 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)); ggml_ext_gpu_context context;
auto query = reinterpret_cast<ggml_ext_cuda_context_fn>(ggml_backend_reg_get_proc_address(reg, "ggml_ext_cuda_context_v1")); if (!cuda_context(execution->backend, &context))
ggml_ext_cuda_context context; return GGML_STATUS_FAILED;
if (!query || !query(execution->backend, &context)) auto stream = static_cast<CUDAStream>(context.stream);
return GGML_STATUS_FAILED; const auto* rope = nodes[count == 1 ? 0 : 2];
auto stream = static_cast<cudaStream_t>(context.stream); const auto* x = count == 1 ? rope->src[0] : nodes[0]->src[0];
const auto* rope = nodes[count == 1 ? 0 : 2]; auto* dst = nodes[count - 1];
const auto* x = count == 1 ? rope->src[0] : nodes[0]->src[0]; const bool alias = overlaps(dst, x) || overlaps(dst, rope->src[1]) || (count > 1 && overlaps(dst, nodes[1]->src[1]));
auto* dst = nodes[count - 1]; void* output = alias ? ggml_backend_buffer_get_base(execution->workspace) : dst->data;
const bool alias = overlaps(dst, x) || overlaps(dst, rope->src[1]) || (count > 1 && overlaps(dst, nodes[1]->src[1])); CUDAError error;
void* output = alias ? ggml_backend_buffer_get_base(execution->workspace) : dst->data; if (count > 1) {
cudaError_t error; error = rms_rope_launch(stream, nodes, output, count == 5, count == 5 ? param_f32(nodes[3], 0) : 1.f);
if (count > 1) { } else {
error = sd_cuda_rms_rope_launch(stream, nodes, output, count == 5, count == 5 ? param_f32(nodes[3], 0) : 1.f); RopeApplyStrides strides;
} else { const auto* pe = rope->src[1];
RopeApplyStrides strides; for (int i = 0; i < 4; ++i) {
const auto* pe = rope->src[1]; strides.x[i] = x->nb[i] / 4;
for (int i = 0; i < 4; ++i) { strides.pe[i] = pe->nb[i] / 4;
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<true> : rope_apply_f32<false>;
kernel<<<static_cast<unsigned int>((pairs + 255) / 256), 256, 0, stream>>>(static_cast<const float*>(x->data), static_cast<const float*>(pe->data),
static_cast<float*>(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; if (error == cuda_success && alias)
auto kernel = rope_interleaved(rope) ? rope_apply_f32<true> : rope_apply_f32<false>; error = cuda_copy_async(dst->data, output, ggml_nbytes(dst), stream);
kernel<<<static_cast<unsigned int>((pairs + 255) / 256), 256, 0, stream>>>(static_cast<const float*>(x->data), static_cast<const float*>(pe->data), return error == cuda_success ? GGML_STATUS_SUCCESS : GGML_STATUS_FAILED;
static_cast<float*>(output), x->ne[0] / 2, pe->ne[2], x->ne[1], x->ne[2], pairs, strides); }
error = cudaGetLastError();
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));
} }
src/core/cuda/rope.h
+5 -6
View File
@@ -1,12 +1,11 @@
#ifndef __SD_CORE_CUDA_ROPE_H__ #ifndef __SD_CORE_CUDA_ROPE_H__
#define __SD_CORE_CUDA_ROPE_H__ #define __SD_CORE_CUDA_ROPE_H__
#include <cuda_runtime.h> #include "runtime.h"
#include "ggml-ext.h"
bool sd_cuda_rope_is_node(const ggml_tensor* node); namespace SD_CUDA_NAMESPACE {
bool sd_cuda_rope_fusion_supports(const ggml_tensor* norm, const ggml_tensor* mul, const ggml_tensor* rope); bool 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); CUDAError rms_rope_launch(CUDAStream stream, ggml_tensor* const* nodes, void* output, bool fp16, float scale);
void sd_ggml_register_cuda_rope(ggml_backend_dev_t device);
}
#endif #endif
src/core/cuda/runtime.h
+111
View File
@@ -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 <hip/hip_fp16.h>
#include <hip/hip_runtime.h>
#define SD_CUDA_NAMESPACE sd_hip
#define SD_CUDA_NAME "hip"
#elif defined(SD_CUDA_VENDOR_MUSA)
#include <musa_fp16.h>
#include <musa_runtime.h>
#define SD_CUDA_NAMESPACE sd_musa
#define SD_CUDA_NAME "musa"
#else
#include <cuda_fp16.h>
#include <cuda_runtime.h>
#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 <typename Kernel>
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 <typename Kernel>
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 <typename Kernel>
static CUDAError cuda_set_shared_memory(Kernel kernel, int bytes) {
return cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, bytes);
}
#endif
template <typename T>
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 <typename T>
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_ext_gpu_context_fn>(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_ext_gpu_device_fn>(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
src/core/ggml_extensions.cpp
+11 -43
View File
@@ -9,19 +9,8 @@
#include <mutex> #include <mutex>
#include <unordered_set> #include <unordered_set>
#ifdef SD_VULKAN_EXTENSIONS #ifdef SD_BACKEND_EXTENSIONS
void sd_ggml_register_vulkan_rope(ggml_backend_dev_t device); #include "sd_backend_extensions.h"
#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);
#endif #endif
static bool rope_infer(ggml_tensor* const* inputs, int count, const void*, ggml_ext_tensor_desc* out) { 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) { if (!backend) {
return; return;
} }
auto device = ggml_backend_get_device(backend); auto device = ggml_backend_get_device(backend);
auto reg = ggml_backend_dev_backend_reg(device); 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 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;
}
static std::mutex mutex; static std::mutex mutex;
static std::unordered_set<ggml_backend_dev_t> registered; static std::unordered_set<ggml_backend_dev_t> registered;
std::lock_guard<std::mutex> lock(mutex); std::lock_guard<std::mutex> lock(mutex);
if (!registered.insert(device).second) { if (!registered.insert(device).second) {
return; return;
} }
if (vulkan) { if (!cpu) {
#ifdef SD_VULKAN_EXTENSIONS #ifdef SD_BACKEND_EXTENSIONS
sd_ggml_register_vulkan_rope(device); for (auto register_backend : sd_backend_extensions)
#endif register_backend(device);
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);
#endif #endif
return; return;
} }
@@ -234,11 +205,8 @@ ggml_tensor* sd_ggml_rope_apply(ggml_context* ctx, ggml_backend_t backend, ggml_
if (backend) { if (backend) {
auto reg = ggml_backend_dev_backend_reg(ggml_backend_get_device(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") if (ggml_backend_reg_get_proc_address(reg, "ggml_ext_cpu_parallel_v1")
#ifdef SD_CUDA_EXTENSIONS #ifdef SD_BACKEND_EXTENSIONS
|| ggml_backend_reg_get_proc_address(reg, "ggml_ext_cuda_context_v1") || ggml_backend_reg_get_proc_address(reg, "ggml_ext_gpu_context_v1") || ggml_backend_reg_get_proc_address(reg, "ggml_ext_vulkan_v1")
#endif
#ifdef SD_VULKAN_EXTENSIONS
|| ggml_backend_reg_get_proc_address(reg, "ggml_ext_vulkan_v1")
#endif #endif
) { ) {
ggml_tensor* inputs[] = {x, pe}; ggml_tensor* inputs[] = {x, pe};
src/core/vulkan/rope.cpp
+2
View File
@@ -200,6 +200,8 @@ namespace sd_vulkan_rope {
} }
void sd_ggml_register_vulkan_rope(ggml_backend_dev_t device) { void sd_ggml_register_vulkan_rope(ggml_backend_dev_t device) {
if (!sd_vulkan_rope::api(device))
return;
ggml_ext_kernel_desc desc = {}; ggml_ext_kernel_desc desc = {};
desc.abi_version = GGML_EXT_ABI_VERSION; desc.abi_version = GGML_EXT_ABI_VERSION;
desc.name = "sd.rope_apply.vulkan.v1"; desc.name = "sd.rope_apply.vulkan.v1";