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)
if (NOT SD_USE_UPSTREAM_GGML)
enable_language(CUDA)
target_sources(${SD_LIB} PRIVATE src/core/cuda/rope.cu src/core/cuda/int8.cu src/core/cuda/convrot.cu src/core/cuda/affine.cu)
set_source_files_properties(src/core/cuda/rope.cu src/core/cuda/int8.cu src/core/cuda/convrot.cu src/core/cuda/affine.cu PROPERTIES COMPILE_OPTIONS "--use_fast_math")
set_target_properties(${SD_LIB} PROPERTIES CUDA_STANDARD 17 CUDA_STANDARD_REQUIRED ON)
if (GGML_CUDA_CUTLASS_PATH)
target_sources(${SD_LIB} PRIVATE src/core/cuda/cute_attention.cu src/core/cuda/i8_projection.cu src/core/cuda/i8_gated_residual.cu src/core/cuda/i8_gemm.cu)
set_source_files_properties(src/core/cuda/cute_attention.cu src/core/cuda/i8_projection.cu src/core/cuda/i8_gated_residual.cu src/core/cuda/i8_gemm.cu PROPERTIES
@@ -358,17 +357,28 @@ if (SD_CUDA)
if (GGML_CUDA_FORCE_CUBLAS)
set_property(SOURCE src/core/cuda/i8_projection.cu src/core/cuda/i8_gated_residual.cu src/core/cuda/i8_gemm.cu APPEND PROPERTY COMPILE_DEFINITIONS SD_CUDA_FORCE_CUBLAS)
endif()
set_property(SOURCE src/core/ggml_extensions.cpp APPEND PROPERTY COMPILE_DEFINITIONS SD_CUDA_CUTE_ATTENTION)
set_property(SOURCE src/core/cuda/int8.cu APPEND PROPERTY COMPILE_DEFINITIONS SD_CUDA_I8_CUTLASS)
endif()
set_target_properties(${SD_LIB} PROPERTIES CUDA_STANDARD 17 CUDA_STANDARD_REQUIRED ON)
target_link_libraries(${SD_LIB} PRIVATE CUDA::cudart CUDA::cublas)
set_property(SOURCE src/core/ggml_extensions.cpp APPEND PROPERTY COMPILE_DEFINITIONS SD_CUDA_EXTENSIONS)
endif()
endif()
set(SD_BACKEND_REGISTRATIONS)
if (NOT SD_USE_UPSTREAM_GGML AND (SD_CUDA OR SD_HIPBLAS OR SD_MUSA))
add_subdirectory(src/core/cuda)
foreach(VENDOR CUDA HIP MUSA)
if(SD_${VENDOR} OR (VENDOR STREQUAL "HIP" AND SD_HIPBLAS))
string(TOLOWER "${VENDOR}" NAMESPACE)
list(APPEND SD_BACKEND_REGISTRATIONS "sd_${NAMESPACE}::register_extensions")
endif()
endforeach()
endif()
if (SD_VULKAN AND NOT SD_USE_UPSTREAM_GGML)
add_subdirectory(src/core/vulkan)
set_property(SOURCE src/core/ggml_extensions.cpp APPEND PROPERTY COMPILE_DEFINITIONS SD_VULKAN_EXTENSIONS)
list(APPEND SD_BACKEND_REGISTRATIONS sd_ggml_register_vulkan_rope)
endif()
if(SD_BACKEND_REGISTRATIONS)
string(REPLACE ";" ",\n " SD_BACKEND_REGISTRATIONS "${SD_BACKEND_REGISTRATIONS}")
configure_file(cmake/backend_extensions.h.in generated/sd_backend_extensions.h @ONLY)
set_property(SOURCE src/core/ggml_extensions.cpp APPEND PROPERTY INCLUDE_DIRECTORIES "${CMAKE_CURRENT_BINARY_DIR}/generated")
set_property(SOURCE src/core/ggml_extensions.cpp APPEND PROPERTY COMPILE_DEFINITIONS SD_BACKEND_EXTENSIONS)
endif()
if (SD_HIPBLAS)
find_package(hip REQUIRED CONFIG HINTS "$ENV{ROCM_PATH}" /opt/rocm)
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_device_v1` | Get the physical CUDA ordinal for capability checks before backend creation, including virtual ggml devices. |
| `ggml_ext_cuda_blas_v1` | Borrow the current stream's cuBLAS handle during execution, with host scalar pointers. |
| `ggml_ext_gpu_device_v1` | Query CUDA/HIP/MUSA runtime kind, physical ordinal, shared-memory limit, warp size, and INT8 BLAS capability. |
| `ggml_ext_gpu_context_v1` | Borrow the matching runtime's device and execution stream. |
| `ggml_ext_gpu_blas_v1` | Borrow the matching cuBLAS, hipBLAS, or muBLAS handle on that stream. |
The borrowed cuBLAS handle remains owned by the backend. Extensions must not retain
or destroy it, or change its stream, pointer mode, math mode, or workspace. This
service is exposed by the CUDA backend, not HIP/MUSA.
CUDA-specific BLAS service is retained for existing callers; the runtime-aware
services support CUDA, HIP, and MUSA with the same ownership rules. Extensions
call vendor runtime and BLAS APIs directly. These services do not execute GEMM,
quantization, or model fusion on the application's behalf.
Workspace is a backend buffer. Its base is a CUDA device pointer on CUDA; callers
must use the representation appropriate to their backend. A Vulkan command
recording service is not implemented yet.
must use the representation appropriate to their backend. Vulkan extensions use
the command recording services described below.
## sd.cpp integration
@@ -127,8 +133,8 @@ recording service is not implemented yet.
`sd.rope_apply.v1`, an override for existing `GGML_OP_ROPE_APPLY` nodes, and a
`RMSNorm -> Mul -> RoPE` fusion. The fusion accepts a one-dimensional weight
vector. Both interleaved and split-half RoPE, partial rotation, and strided input
rows are supported. Model construction uses the custom operator on CPU and CUDA
and retains existing backend implementations elsewhere.
rows are supported. Model construction uses the custom operator on CPU, CUDA,
HIP, MUSA, and Vulkan, and retains existing backend implementations elsewhere.
`src/core/cuda/rope.cu` supplies standalone CUDA RoPE, the same builtin override,
and `RMSNorm -> Mul -> RoPE [-> Scale -> FP16 Copy]` fusion. CuTe attention shares
@@ -136,8 +142,16 @@ its RMSNorm/RoPE kernel. Partial rotation and arbitrary float-aligned input and
position-encoding strides are supported by the standalone kernel; RMSNorm fusion
requires contiguous position encoding and weights. Copy fusion uses workspace
when the destination overlaps an input. These registrations do not require
CUTLASS. The old CUDA RoPE dispatch is removed; HIP and MUSA retain their existing
implementations.
CUTLASS. The old CUDA/HIP/MUSA RoPE dispatch is removed.
Shared CUDA-style kernels live in `src/core/cuda/` and are compiled separately
for each enabled runtime. `SD_CUDA_NAMESPACE` expands to `sd_cuda`, `sd_hip`, or
`sd_musa`, so implementations can coexist. `runtime.h` handles the small runtime
differences, including logical 32-lane shuffles on HIP. The application defines
and launches its kernels directly; it includes only public ggml headers.
CMake generates the registration table from enabled backend options. The shared
entry point uses one `SD_BACKEND_EXTENSIONS` macro; backend registrars verify the
device kind, and disabled backends add no symbol references.
`SD_USE_UPSTREAM_GGML` excludes these registrations. The stable-diffusion public C
API is unchanged. INT8 ConvRot retains its existing packed activation layout and
@@ -216,10 +230,12 @@ Custom I8 nodes also retain cuBLAS when the reduction could overflow I32, to
preserve their previous accumulation behavior. Fusion disable and force-cuBLAS
settings are honored. No CUTLASS headers are needed for the cuBLAS path.
The native CUDA CUTLASS GEMM implementation and its SwiGLU/row-concat matchers
have been removed from ggml. Native cuBLAS GEMM and standalone ConvRot remain
compatibility implementations for callers without these registrations. HIP
retains its native INT8/activation fusion paths.
The native CUDA/HIP INT8 GEMM, standalone ConvRot, and associated fusion paths
have been removed from ggml. The shared application implementation calls cuBLAS
or hipBLAS directly. HIP retains the existing architecture eligibility policy;
MUSA does not advertise INT8 BLAS support and uses scheduler fallback for GEMM.
Standalone ConvRot kernels are compiled for all three runtimes. CUTLASS
implementations remain CUDA-specific.
`src/core/cuda/i8_gated_residual.cu` implements
`INT8 GEMM -> [Scale] -> Mul(gate) -> Add(residual) -> [Concat]` as a registered
@@ -255,8 +271,8 @@ overlapping outputs. Long rows or external intermediate consumers retain
ordinary execution. These fusions do not require CUTLASS and honor
`GGML_CUDA_DISABLE_FUSION`.
The corresponding scale, normalization, and unary-activation fusion rules in
ggml now only serve HIP. CUDA GEMM/SwiGLU/ConvRot chains are selected by the
The corresponding scale, normalization, and unary-activation fusion rules have
been removed from ggml for CUDA and HIP. GEMM/SwiGLU/ConvRot chains are selected by the
application INT8 extension before the standalone activation fusion.
`src/core/cuda/affine.cu` registers FP32 `Mul(gate) -> Add(residual)`,
@@ -268,8 +284,8 @@ Add inputs, and fusion disabling retain ordinary operators. Internal in-place
Add outputs are accepted when they overwrite the complete fused Mul allocation. Exact in-place residual
or input addresses are safe for the pointwise kernel; partial overlap, broadcast
gate overlap, and FP16 copy overlap use scheduler workspace. These registrations
require no CUTLASS and support reusable CUDA Graph replay. Their former native
fusion implementations now only serve HIP/MUSA.
require no CUTLASS, compile for CUDA/HIP/MUSA, and support reusable CUDA Graph
replay. Their former native fusion implementations have been removed from ggml.
CUDA kernels can opt into device replay with `ggml_ext_register_kernel_ex` and
`GGML_EXT_KERNEL_CAPTURE_SAFE`. The original registration function retains its
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(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 int32_t uses = index != GGML_HASHSET_FULL && ggml_bitset_get(graph->visited_hash_set.used, index)
@@ -187,7 +190,8 @@ namespace sd {
ggml_backend_buffer_set_usage(buffer, GGML_BACKEND_BUFFER_USAGE_WEIGHTS);
}
copy->buffer = buffer;
copy->data = reinterpret_cast<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;
}
}
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 <cstdlib>
#include <cstring>
#include "ggml-ext.h"
#include "runtime.h"
namespace sd_affine_cuda {
static bool disabled() {
static const bool value = [] { const char* v = std::getenv("GGML_CUDA_DISABLE_FUSION"); return v && std::atoi(v) != 0; }();
return value;
}
static bool aligned(const ggml_tensor* tensor, size_t alignment) {
size_t offset = 0;
while (!tensor->data && tensor->view_src) {
offset += tensor->view_offs;
tensor = tensor->view_src;
namespace SD_CUDA_NAMESPACE {
namespace sd_affine_cuda {
static bool disabled() {
static const bool value = [] { const char* v = std::getenv("GGML_CUDA_DISABLE_FUSION"); return v && std::atoi(v) != 0; }();
return value;
}
return (reinterpret_cast<uintptr_t>(tensor->data) + offset) % alignment == 0;
}
static float param(const ggml_tensor* tensor, int index) {
float value;
std::memcpy(&value, tensor->op_params + index, sizeof(value));
return value;
}
static bool contiguous_f32(const ggml_tensor* tensor) {
return tensor && tensor->type == GGML_TYPE_F32 && ggml_is_contiguous(tensor) && aligned(tensor, 16);
}
static bool bounded(const ggml_tensor* tensor) {
return ggml_nelements(tensor) > 0 && ggml_nelements(tensor) <= int64_t(INT_MAX) * 1024;
}
static int match(ggml_backend_dev_t, const ggml_cgraph* graph, int first, void*) {
auto* g = const_cast<ggml_cgraph*>(graph);
if (disabled() || first + 1 >= ggml_graph_n_nodes(g))
static bool aligned(const ggml_tensor* tensor, size_t alignment) {
size_t offset = 0;
while (!tensor->data && tensor->view_src) {
offset += tensor->view_offs;
tensor = tensor->view_src;
}
return (reinterpret_cast<uintptr_t>(tensor->data) + offset) % alignment == 0;
}
static float param(const ggml_tensor* tensor, int index) {
float value;
std::memcpy(&value, tensor->op_params + index, sizeof(value));
return value;
}
static bool contiguous_f32(const ggml_tensor* tensor) {
return tensor && tensor->type == GGML_TYPE_F32 && ggml_is_contiguous(tensor) && aligned(tensor, 16);
}
static bool bounded(const ggml_tensor* tensor) {
return ggml_nelements(tensor) > 0 && ggml_nelements(tensor) <= int64_t(INT_MAX) * 1024;
}
static int match(ggml_backend_dev_t, const ggml_cgraph* graph, int first, void*) {
auto* g = const_cast<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;
const auto* node = ggml_graph_node(g, first);
const auto* next = ggml_graph_node(g, first + 1);
if (node->op == GGML_OP_MUL) {
if (next->op != GGML_OP_ADD || (next->src[0] != node && next->src[1] != node))
return 0;
const auto* x = node->src[0];
const auto* gate = node->src[1];
const auto* residual = next->src[0] == node ? next->src[1] : next->src[0];
return residual != node && contiguous_f32(x) && contiguous_f32(gate) && contiguous_f32(residual) && contiguous_f32(next) &&
node->type == GGML_TYPE_F32 && ggml_are_same_shape(node, next) && ggml_are_same_shape(x, next) &&
ggml_are_same_shape(residual, next) && gate->ne[0] == x->ne[0] && ggml_nelements(gate) == x->ne[0] &&
x->ne[0] % 4 == 0 && bounded(next)
? 2
: 0;
}
if (node->op == GGML_OP_CONCAT) {
const bool permuted = next->op == GGML_OP_PERMUTE;
const int count = permuted ? 3 : 2;
if (first + count > ggml_graph_n_nodes(g) || node->type != GGML_TYPE_F32 || node->ne[0] % 4 ||
node->op_params[0] != (permuted ? 2 : 1) || !ggml_is_contiguous(node) ||
!contiguous_f32(node->src[0]) || !contiguous_f32(node->src[1]))
return 0;
if (permuted && (next->src[0] != node || next->ne[0] != node->ne[0] || next->ne[1] != node->ne[2] ||
next->ne[2] != node->ne[1] || next->ne[3] != node->ne[3] || next->nb[0] != node->nb[0] ||
next->nb[1] != node->nb[2] || next->nb[2] != node->nb[1] || next->nb[3] != node->nb[3]))
return 0;
const auto* dst = ggml_graph_node(g, first + count - 1);
return dst->op == GGML_OP_CPY && dst->src[0] == (permuted ? next : node) && dst->type == GGML_TYPE_F16 &&
ggml_is_contiguous(dst) && ggml_are_same_shape(dst, dst->src[0]) && aligned(dst, 4) && bounded(dst)
? count
: 0;
static __global__ void gated_residual(const float4* x, const float4* gate, const float4* residual, float4* dst, int64_t n4, int64_t count4) {
const int64_t i = int64_t(blockIdx.x) * blockDim.x + threadIdx.x;
if (i >= count4)
return;
const float4 a = x[i], g = gate[i % n4], r = residual[i];
// The graph performs two FP32 operations; contracting them into FMA changes rounding.
dst[i] = make_float4(__fadd_rn(__fmul_rn(a.x, g.x), r.x), __fadd_rn(__fmul_rn(a.y, g.y), r.y),
__fadd_rn(__fmul_rn(a.z, g.z), r.z), __fadd_rn(__fmul_rn(a.w, g.w), r.w));
}
if (node->op == GGML_OP_CONT && first + 3 < ggml_graph_n_nodes(g)) {
const auto* src = node->src[0];
const auto* permute = ggml_graph_node(g, first + 2);
const auto* dst = ggml_graph_node(g, first + 3);
return ggml_are_same_shape(src, node) && src->type == GGML_TYPE_F32 && src->nb[0] == 4 && src->ne[0] % 4 == 0 &&
aligned(src, 16) && (src->nb[1] | src->nb[2] | src->nb[3]) % 16 == 0 &&
next->op == GGML_OP_SCALE && next->src[0] == node && next->type == GGML_TYPE_F32 && param(next, 1) == 0.f &&
permute->op == GGML_OP_PERMUTE && permute->src[0] == next && permute->op_params[0] == 0 &&
permute->op_params[1] == 2 && permute->op_params[2] == 1 && permute->op_params[3] == 3 &&
dst->op == GGML_OP_CPY && dst->src[0] == permute && dst->type == GGML_TYPE_F16 &&
ggml_are_same_shape(dst, permute) && ggml_is_contiguous(dst) && aligned(dst, 4) && bounded(dst)
? 4
: 0;
template <bool Permuted>
static __global__ void concat_cast(const float4* prefix, const float4* current, half2* dst, int64_t columns4, int64_t prefix_tokens, int64_t tokens, int64_t heads, int64_t count4) {
const int64_t i = int64_t(blockIdx.x) * blockDim.x + threadIdx.x;
if (i >= count4)
return;
const int64_t col = i % columns4, row = i / columns4;
const int64_t token = row % tokens, head = row / tokens % heads, batch = row / (tokens * heads);
const bool from_prefix = token < prefix_tokens;
const int64_t src_tokens = from_prefix ? prefix_tokens : tokens - prefix_tokens;
const int64_t src_token = from_prefix ? token : token - prefix_tokens;
const int64_t src_row = Permuted ? (batch * src_tokens + src_token) * heads + head
: (batch * heads + head) * src_tokens + src_token;
const float4 value = (from_prefix ? prefix : current)[src_row * columns4 + col];
dst[2 * i] = __floats2half2_rn(value.x, value.y);
dst[2 * i + 1] = __floats2half2_rn(value.z, value.w);
}
return 0;
}
static __global__ void gated_residual(const float4* x, const float4* gate, const float4* residual, float4* dst, int64_t n4, int64_t count4) {
const int64_t i = int64_t(blockIdx.x) * blockDim.x + threadIdx.x;
if (i >= count4)
return;
const float4 a = x[i], g = gate[i % n4], r = residual[i];
// The graph performs two FP32 operations; contracting them into FMA changes rounding.
dst[i] = make_float4(__fadd_rn(__fmul_rn(a.x, g.x), r.x), __fadd_rn(__fmul_rn(a.y, g.y), r.y),
__fadd_rn(__fmul_rn(a.z, g.z), r.z), __fadd_rn(__fmul_rn(a.w, g.w), r.w));
}
template <bool Permuted>
static __global__ void concat_cast(const float4* prefix, const float4* current, half2* dst, int64_t columns4, int64_t prefix_tokens, int64_t tokens, int64_t heads, int64_t count4) {
const int64_t i = int64_t(blockIdx.x) * blockDim.x + threadIdx.x;
if (i >= count4)
return;
const int64_t col = i % columns4, row = i / columns4;
const int64_t token = row % tokens, head = row / tokens % heads, batch = row / (tokens * heads);
const bool from_prefix = token < prefix_tokens;
const int64_t src_tokens = from_prefix ? prefix_tokens : tokens - prefix_tokens;
const int64_t src_token = from_prefix ? token : token - prefix_tokens;
const int64_t src_row = Permuted ? (batch * src_tokens + src_token) * heads + head
: (batch * heads + head) * src_tokens + src_token;
const float4 value = (from_prefix ? prefix : current)[src_row * columns4 + col];
dst[2 * i] = __floats2half2_rn(value.x, value.y);
dst[2 * i + 1] = __floats2half2_rn(value.z, value.w);
}
static __global__ void scale_permute(const float* src, half2* dst, int64_t dim, int64_t tokens, int64_t heads, int64_t count, int64_t nb1, int64_t nb2, int64_t nb3, float scale) {
const int64_t index = (int64_t(blockIdx.x) * blockDim.x + threadIdx.x) * 4;
if (index >= count)
return;
const int64_t col = index % dim, row = index / dim;
const int64_t token = row % tokens, head = row / tokens % heads, batch = row / (tokens * heads);
const float4 value = *reinterpret_cast<const float4*>(src + col + head * nb1 + token * nb2 + batch * nb3);
dst[index / 2] = __floats2half2_rn(__fmul_rn(value.x, scale), __fmul_rn(value.y, scale));
dst[index / 2 + 1] = __floats2half2_rn(__fmul_rn(value.z, scale), __fmul_rn(value.w, scale));
}
static bool overlaps(const ggml_tensor* a, const ggml_tensor* b) {
const auto x = reinterpret_cast<uintptr_t>(a->data), y = reinterpret_cast<uintptr_t>(b->data);
return x <= y ? y - x < ggml_nbytes(a) : x - y < ggml_nbytes(b);
}
static ggml_status prepare(ggml_backend_t, ggml_tensor* const* nodes, int count, void**, size_t* workspace, void*) {
*workspace = ggml_nbytes(nodes[count - 1]);
return GGML_STATUS_SUCCESS;
}
static void destroy(void*, void*) {}
static ggml_status execute(const ggml_ext_execution_context* execution, ggml_tensor* const* nodes, int count, void*, void*) {
auto reg = ggml_backend_dev_backend_reg(ggml_backend_get_device(execution->backend));
auto query = reinterpret_cast<ggml_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))
return 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));
static __global__ void scale_permute(const float* src, half2* dst, int64_t dim, int64_t tokens, int64_t heads, int64_t count, int64_t nb1, int64_t nb2, int64_t nb3, float scale) {
const int64_t index = (int64_t(blockIdx.x) * blockDim.x + threadIdx.x) * 4;
if (index >= count)
return;
const int64_t col = index % dim, row = index / dim;
const int64_t token = row % tokens, head = row / tokens % heads, batch = row / (tokens * heads);
const float4 value = *reinterpret_cast<const float4*>(src + col + head * nb1 + token * nb2 + batch * nb3);
dst[index / 2] = __floats2half2_rn(__fmul_rn(value.x, scale), __fmul_rn(value.y, scale));
dst[index / 2 + 1] = __floats2half2_rn(__fmul_rn(value.z, scale), __fmul_rn(value.w, scale));
}
static bool overlaps(const ggml_tensor* a, const ggml_tensor* b) {
const auto x = reinterpret_cast<uintptr_t>(a->data), y = reinterpret_cast<uintptr_t>(b->data);
return x <= y ? y - x < ggml_nbytes(a) : x - y < ggml_nbytes(b);
}
static ggml_status prepare(ggml_backend_t, ggml_tensor* const* nodes, int count, void**, size_t* workspace, void*) {
*workspace = ggml_nbytes(nodes[count - 1]);
return GGML_STATUS_SUCCESS;
}
static void destroy(void*, void*) {}
static ggml_status execute(const ggml_ext_execution_context* execution, ggml_tensor* const* nodes, int count, void*, void*) {
ggml_ext_gpu_context context;
if (!cuda_context(execution->backend, &context))
return GGML_STATUS_FAILED;
auto stream = static_cast<CUDAStream>(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 = cuda_last_error();
if (error == cuda_success && alias)
error = cuda_copy_async(dst->data, output, ggml_nbytes(dst), stream);
return error == cuda_success ? GGML_STATUS_SUCCESS : GGML_STATUS_FAILED;
}
auto error = cudaGetLastError();
if (error == cudaSuccess && alias)
error = cudaMemcpyAsync(dst->data, output, ggml_nbytes(dst), cudaMemcpyDeviceToDevice, stream);
return error == cudaSuccess ? GGML_STATUS_SUCCESS : GGML_STATUS_FAILED;
}
}
void sd_ggml_register_cuda_affine(ggml_backend_dev_t device) {
ggml_ext_kernel_desc desc = {};
desc.abi_version = GGML_EXT_ABI_VERSION;
desc.priority = 200;
desc.match = sd_affine_cuda::match;
desc.prepare = sd_affine_cuda::prepare;
desc.destroy = sd_affine_cuda::destroy;
desc.execute = sd_affine_cuda::execute;
desc.name = "sd.gated_residual.cuda.v1";
desc.op = GGML_OP_MUL;
GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
desc.name = "sd.concat_cast.cuda.v1";
desc.op = GGML_OP_CONCAT;
GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
desc.name = "sd.cont_scale_permute_cast.cuda.v1";
desc.op = GGML_OP_CONT;
GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
void register_affine(ggml_backend_dev_t device) {
ggml_ext_kernel_desc desc = {};
desc.abi_version = GGML_EXT_ABI_VERSION;
desc.priority = 200;
desc.match = sd_affine_cuda::match;
desc.prepare = sd_affine_cuda::prepare;
desc.destroy = sd_affine_cuda::destroy;
desc.execute = sd_affine_cuda::execute;
desc.name = "sd.gated_residual." SD_CUDA_NAME ".v1";
desc.op = GGML_OP_MUL;
GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
desc.name = "sd.concat_cast." SD_CUDA_NAME ".v1";
desc.op = GGML_OP_CONCAT;
GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
desc.name = "sd.cont_scale_permute_cast." SD_CUDA_NAME ".v1";
desc.op = GGML_OP_CONT;
GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
}
}
src/core/cuda/convrot.cu
+271 -290
View File
@@ -2,314 +2,295 @@
#include <climits>
#include <cstdlib>
#include <cstring>
#include <mutex>
#include <string>
#include <unordered_map>
#include "convrot_kernel.h"
#include "core/int8_ops.h"
#include "runtime.h"
namespace sd_convrot {
struct DeviceInfo {
int ordinal = -1;
size_t shared = 0;
};
static DeviceInfo device_info(ggml_backend_dev_t device) {
static std::mutex mutex;
static std::unordered_map<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;
namespace SD_CUDA_NAMESPACE {
namespace sd_convrot {
static bool fusion_disabled() {
static const bool disabled = [] { const char* v = std::getenv("GGML_CUDA_DISABLE_FUSION"); return v && std::atoi(v) != 0; }();
return disabled;
}
return true;
}
static bool supported(const Operation& op, const DeviceInfo& device) {
const size_t extra = op.kind == NORM ? 384 : op.kind == RMS_MOD ? 256
: 128;
return device.ordinal >= 0 && size_t(op.k) * sizeof(float) + extra <= device.shared;
}
static int match(ggml_backend_dev_t device, const ggml_cgraph* graph, int first, void*) {
if (fusion_disabled())
return 0;
auto* g = const_cast<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;
return describe(nodes, count, op) && supported(op, device_info(device)) ? count : 0;
}
static void optimize(ggml_backend_dev_t device, ggml_cgraph* graph, void*) {
if (fusion_disabled())
return;
for (int i = 2; i < ggml_graph_n_nodes(graph); ++i) {
auto* quant = ggml_graph_node(graph, i);
static float param_f32(const ggml_tensor* tensor, int index) {
float value;
std::memcpy(&value, tensor->op_params + index, sizeof(value));
return value;
}
static bool f32(const ggml_tensor* tensor) {
return tensor && tensor->type == GGML_TYPE_F32 && ggml_is_contiguous(tensor);
}
static bool vector(const ggml_tensor* tensor, int64_t k) {
return f32(tensor) && tensor->ne[0] == k && ggml_nelements(tensor) == k;
}
enum Kind { SCALE,
NORM,
RMS_MOD,
SWIGLU,
SILU,
SIGMOID };
struct Operation {
Kind kind;
const ggml_tensor* input = nullptr;
const ggml_tensor* up = nullptr;
const ggml_tensor* modulation = nullptr;
const ggml_tensor* weight = nullptr;
const ggml_tensor* shift = nullptr;
int64_t k, rows, stride, up_stride, gate_offset = 0, up_offset = 0;
float eps = 0.f, scale = 1.f, bias = 0.f;
};
static bool describe(ggml_tensor* const* nodes, int count, Operation& op) {
int group;
const auto* first = nodes[0];
const auto* quant = nodes[count - 1];
if (!sd_ggml_convrot_describe(quant, &group) || group != 256)
continue;
ggml_tensor* chain[6] = {};
int count = 0;
auto* value = quant->src[0];
if (value->op == GGML_OP_ADD && value->src[0]->op == GGML_OP_ADD) {
auto* add = value->src[0];
auto* weight = add->src[0];
if (weight->op == GGML_OP_MUL) {
chain[0] = weight->src[0];
chain[1] = weight;
chain[2] = add->src[1];
chain[3] = add;
chain[4] = value;
chain[5] = quant;
count = 6;
return false;
const auto* value = quant->src[0];
op.k = value->ne[0];
op.rows = ggml_nrows(value);
op.input = first->src[0];
if (!f32(op.input) || !f32(first) || op.k > INT_MAX || op.rows > INT_MAX)
return false;
op.stride = op.input->ne[0];
op.up_stride = op.k;
if (count == 2 && first->op == GGML_OP_SCALE && value == first) {
op.kind = SCALE;
op.scale = param_f32(first, 0);
op.bias = param_f32(first, 1);
} else if (count == 3 && first->op == GGML_OP_NORM && nodes[1]->op == GGML_OP_MUL && nodes[1]->src[0] == first && value == nodes[1]) {
op.kind = NORM;
op.modulation = nodes[1]->src[1];
op.eps = param_f32(first, 0);
if (op.k < 1024 || !vector(op.modulation, op.k))
return false;
} else if (count == 6 && first->op == GGML_OP_RMS_NORM && nodes[1]->op == GGML_OP_MUL && nodes[1]->src[0] == first &&
nodes[2]->op == GGML_OP_MUL && nodes[2]->src[0] == nodes[1] &&
nodes[3]->op == GGML_OP_ADD && nodes[3]->src[0] == nodes[1] && nodes[3]->src[1] == nodes[2] &&
nodes[4]->op == GGML_OP_ADD && nodes[4]->src[0] == nodes[3] && value == nodes[4]) {
op.kind = RMS_MOD;
op.weight = nodes[1]->src[1];
op.modulation = nodes[2]->src[1];
op.shift = nodes[4]->src[1];
op.eps = param_f32(first, 0);
if (op.k < 1024 || !vector(op.weight, op.k) || !vector(op.modulation, op.k) || !vector(op.shift, op.k))
return false;
} else if (count == 2 && first->op == GGML_OP_GLU && ggml_get_glu_op(first) == GGML_GLU_OP_SWIGLU && value == first) {
op.kind = SWIGLU;
op.up = first->src[1] ? first->src[1] : op.input;
if (!f32(op.up))
return false;
op.up_stride = op.up->ne[0];
if (!first->src[1]) {
if (first->op_params[1])
op.gate_offset = op.k;
else
op.up_offset = op.k;
}
} else if (value->op == GGML_OP_MUL) {
for (int src = 0; src < 2; ++src)
if (value->src[src]->op == GGML_OP_UNARY) {
chain[0] = value->src[src];
chain[1] = value;
chain[2] = quant;
count = 3;
break;
}
} else if (count == 3 && first->op == GGML_OP_UNARY && nodes[1]->op == GGML_OP_MUL && value == nodes[1] &&
(nodes[1]->src[0] == first || nodes[1]->src[1] == first)) {
auto unary = ggml_get_unary_op(first);
if (unary != GGML_UNARY_OP_SILU && unary != GGML_UNARY_OP_SIGMOID)
return false;
op.kind = unary == GGML_UNARY_OP_SILU ? SILU : SIGMOID;
op.up = nodes[1]->src[0] == first ? nodes[1]->src[1] : nodes[1]->src[0];
if (!f32(op.up) || !ggml_are_same_shape(op.input, op.up))
return false;
} else
return false;
for (int i = 0; i < count - 1; ++i) {
if (!f32(nodes[i]) || !ggml_are_same_shape(nodes[i], value))
return false;
for (const auto* input : {op.input, op.up, op.modulation, op.weight, op.shift})
if (input == nodes[i])
return false;
}
return true;
}
static bool supported(const Operation& op, const ggml_ext_gpu_device_info& device) {
const size_t extra = op.kind == NORM ? 384 : op.kind == RMS_MOD ? 256
: 128;
return device.device >= 0 && size_t(op.k) * sizeof(float) + extra <= device.shared_memory;
}
static int match(ggml_backend_dev_t device, const ggml_cgraph* graph, int first, void*) {
if (fusion_disabled())
return 0;
auto* g = const_cast<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;
if (!count || !describe(chain, count, op) || !supported(op, device_info(device)))
continue;
int begin = -1;
for (int j = std::max(0, i - 64); j < i; ++j)
if (ggml_graph_node(graph, j) == chain[0]) {
begin = j;
break;
return describe(nodes, count, op) && supported(op, cuda_device_info(device)) ? count : 0;
}
static void optimize(ggml_backend_dev_t device, ggml_cgraph* graph, void*) {
if (fusion_disabled())
return;
for (int i = 2; i < ggml_graph_n_nodes(graph); ++i) {
auto* quant = ggml_graph_node(graph, i);
int group;
if (!sd_ggml_convrot_describe(quant, &group) || group != 256)
continue;
ggml_tensor* chain[6] = {};
int count = 0;
auto* value = quant->src[0];
if (value->op == GGML_OP_ADD && value->src[0]->op == GGML_OP_ADD) {
auto* add = value->src[0];
auto* weight = add->src[0];
if (weight->op == GGML_OP_MUL) {
chain[0] = weight->src[0];
chain[1] = weight;
chain[2] = add->src[1];
chain[3] = add;
chain[4] = value;
chain[5] = quant;
count = 6;
}
} else if (value->op == GGML_OP_MUL) {
for (int src = 0; src < 2; ++src)
if (value->src[src]->op == GGML_OP_UNARY) {
chain[0] = value->src[src];
chain[1] = value;
chain[2] = quant;
count = 3;
break;
}
}
if (begin < 0)
continue;
// Move the owning producer with its in-place consumers so alias checks retain their storage relationship.
for (int part = 1; part < count; ++part) {
int next = -1;
for (int j = begin + part; j <= i; ++j)
if (ggml_graph_node(graph, j) == chain[part]) {
next = j;
Operation op;
if (!count || !describe(chain, count, op) || !supported(op, cuda_device_info(device)))
continue;
int begin = -1;
for (int j = std::max(0, i - 64); j < i; ++j)
if (ggml_graph_node(graph, j) == chain[0]) {
begin = j;
break;
}
if (next < 0 || (next > begin + part && !ggml_ext_graph_delay(graph, begin, part, next)))
break;
begin = next - part;
if (begin < 0)
continue;
// Move the owning producer with its in-place consumers so alias checks retain their storage relationship.
for (int part = 1; part < count; ++part) {
int next = -1;
for (int j = begin + part; j <= i; ++j)
if (ggml_graph_node(graph, j) == chain[part]) {
next = j;
break;
}
if (next < 0 || (next > begin + part && !ggml_ext_graph_delay(graph, begin, part, next)))
break;
begin = next - part;
}
}
}
}
struct Plan {
DeviceInfo device;
Operation op;
bool configured = false;
};
static ggml_status prepare(ggml_backend_t backend, ggml_tensor* const* nodes, int count, void** state, size_t* workspace, void*) {
auto* plan = new Plan{};
*state = plan;
plan->device = device_info(ggml_backend_get_device(backend));
if (!describe(nodes, count, plan->op) || !supported(plan->op, plan->device))
return GGML_STATUS_FAILED;
*workspace = ggml_nbytes(nodes[count - 1]);
return GGML_STATUS_SUCCESS;
}
static void destroy(void* state, void*) {
delete static_cast<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:
struct Plan {
ggml_ext_gpu_device_info device;
Operation op;
bool configured = false;
};
static ggml_status prepare(ggml_backend_t backend, ggml_tensor* const* nodes, int count, void** state, size_t* workspace, void*) {
auto* plan = new Plan{};
*state = plan;
plan->device = cuda_device_info(ggml_backend_get_device(backend));
if (!describe(nodes, count, plan->op) || !supported(plan->op, plan->device))
return GGML_STATUS_FAILED;
*workspace = ggml_nbytes(nodes[count - 1]);
return GGML_STATUS_SUCCESS;
}
static void destroy(void* state, void*) {
delete static_cast<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) {
for (auto op : {GGML_OP_SCALE, GGML_OP_NORM, GGML_OP_RMS_NORM, GGML_OP_GLU, GGML_OP_UNARY}) {
ggml_ext_kernel_desc desc = {};
desc.abi_version = GGML_EXT_ABI_VERSION;
const std::string name = std::string("sd.convrot_fusion.") + ggml_op_name(op) + ".cuda.v1";
desc.name = name.c_str();
desc.op = op;
desc.priority = 225;
desc.optimize = op == GGML_OP_SCALE ? sd_convrot::optimize : nullptr;
desc.match = sd_convrot::match;
desc.prepare = sd_convrot::prepare;
desc.execute = sd_convrot::execute;
desc.destroy = sd_convrot::destroy;
GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
void register_convrot(ggml_backend_dev_t device) {
for (auto op : {GGML_OP_SCALE, GGML_OP_NORM, GGML_OP_RMS_NORM, GGML_OP_GLU, GGML_OP_UNARY}) {
ggml_ext_kernel_desc desc = {};
desc.abi_version = GGML_EXT_ABI_VERSION;
const std::string name = std::string("sd.convrot_fusion.") + ggml_op_name(op) + "." SD_CUDA_NAME ".v1";
desc.name = name.c_str();
desc.op = op;
desc.priority = 225;
desc.optimize = op == GGML_OP_SCALE ? sd_convrot::optimize : nullptr;
desc.match = sd_convrot::match;
desc.prepare = sd_convrot::prepare;
desc.execute = sd_convrot::execute;
desc.destroy = sd_convrot::destroy;
GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
}
}
}
src/core/cuda/convrot_kernel.h
+224 -222
View File
@@ -1,249 +1,251 @@
#ifndef __SD_CORE_CUDA_CONVROT_KERNEL_H__
#define __SD_CORE_CUDA_CONVROT_KERNEL_H__
#include <cuda_runtime.h>
#include <cstdint>
namespace sd_convrot {
static __device__ __forceinline__ float warp_reduce_sum(float x) {
#include "runtime.h"
namespace SD_CUDA_NAMESPACE {
namespace sd_convrot {
static __device__ __forceinline__ float warp_reduce_sum(float x) {
#pragma unroll
for (int offset = 16; offset > 0; offset >>= 1)
x += __shfl_xor_sync(0xffffffff, x, offset, 32);
return x;
}
static __device__ __forceinline__ float2 warp_reduce_sum(float2 x) {
#pragma unroll
for (int offset = 16; offset > 0; offset >>= 1) {
x.x += __shfl_xor_sync(0xffffffff, x.x, offset, 32);
x.y += __shfl_xor_sync(0xffffffff, x.y, offset, 32);
for (int offset = 16; offset > 0; offset >>= 1)
x += cuda_shuffle_xor(x, offset, 32);
return x;
}
return x;
}
template <int Width>
static __device__ __forceinline__ float warp_reduce_max(float x) {
static __device__ __forceinline__ float2 warp_reduce_sum(float2 x) {
#pragma unroll
for (int offset = Width / 2; offset > 0; offset >>= 1)
x = fmaxf(x, __shfl_xor_sync(0xffffffff, x, offset, Width));
return x;
}
static __device__ float2 block_sum(float2 x, float2* shared) {
x = warp_reduce_sum(x);
if (threadIdx.x % 32 == 0)
shared[threadIdx.x / 32] = x;
__syncthreads();
return warp_reduce_sum(shared[threadIdx.x % 32]);
}
static __device__ __forceinline__ float convrot_h4_cuda(float a, float b, float c, float d, int index) {
a = index == 3 ? -a : a;
b = index == 2 ? -b : b;
c = index == 1 ? -c : c;
d = index == 0 ? -d : d;
return __fadd_rn(__fadd_rn(__fadd_rn(a, b), c), d);
}
template <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;
}
for (int offset = 16; offset > 0; offset >>= 1) {
x.x += cuda_shuffle_xor(x.x, offset, 32);
x.y += cuda_shuffle_xor(x.y, offset, 32);
}
return x;
}
template <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();
const float sum = warp_reduce_sum(sums[tid % 32]);
rms_scale = rsqrtf(sum / k + eps);
return warp_reduce_sum(shared[threadIdx.x % 32]);
}
static __device__ __forceinline__ float convrot_h4_cuda(float a, float b, float c, float d, int index) {
a = index == 3 ? -a : a;
b = index == 2 ? -b : b;
c = index == 1 ? -c : c;
d = index == 0 ? -d : d;
return __fadd_rn(__fadd_rn(__fadd_rn(a, b), c), d);
}
if constexpr (swiglu || warp_rotation) {
const int lane = tid % 32;
for (int64_t group = tid / 32; group < k / group_size; group += blockDim.x / 32) {
float v[8];
#pragma unroll
for (int j = 0; j < 8; ++j) {
const int64_t i = group * group_size + j * 32 + lane;
if constexpr (!swiglu) {
float value = src_row[i];
if constexpr (rms_modulation) {
value = __fmul_rn(__fmul_rn(value, rms_scale), rms_weight[i]);
value = __fadd_rn(__fadd_rn(value, __fmul_rn(value, modulation[i])), rms_shift[i]);
}
if constexpr (scaled) {
value = __fmaf_rn(value, input_scale, input_bias);
}
v[j] = value * (1.0f / 16.0f);
} else {
float gate_value, up_value;
if constexpr (dequantize) {
gate_value = (float)((const int32_t*)src_row)[i] * activation_scales[row];
up_value = (float)((const int32_t*)up)[row * up_stride + i] * activation_scales[row];
gate_value *= weight_scales[i];
up_value *= weight_scales[k + i];
} else {
gate_value = src_row[i];
up_value = up[row * up_stride + i];
}
const float activated = sigmoid_gate ? 1.0f / (1.0f + expf(-gate_value)) : (gate_value / (1.0f + expf(-gate_value)));
v[j] = (activated * up_value) * (1.0f / 16.0f);
template <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;
}
}
// One logical warp owns H256, keeping every butterfly in registers.
#pragma unroll
for (int stride = 1; stride <= 4; stride *= 4) {
const int base = lane & ~(3 * stride);
__syncthreads();
const float sum = warp_reduce_sum(sums[tid % 32]);
rms_scale = rsqrtf(sum / k + eps);
}
if constexpr (swiglu || warp_rotation) {
const int lane = tid % 32;
for (int64_t group = tid / 32; group < k / group_size; group += blockDim.x / 32) {
float v[8];
#pragma unroll
for (int j = 0; j < 8; ++j) {
const float a = __shfl_sync(0xffffffff, v[j], base, 32);
const float b = __shfl_sync(0xffffffff, v[j], base + stride, 32);
const float c = __shfl_sync(0xffffffff, v[j], base + 2 * stride, 32);
const float d = __shfl_sync(0xffffffff, v[j], base + 3 * stride, 32);
v[j] = convrot_h4_cuda(a, b, c, d, (lane / stride) % 4);
const int64_t i = group * group_size + j * 32 + lane;
if constexpr (!swiglu) {
float value = src_row[i];
if constexpr (rms_modulation) {
value = __fmul_rn(__fmul_rn(value, rms_scale), rms_weight[i]);
value = __fadd_rn(__fadd_rn(value, __fmul_rn(value, modulation[i])), rms_shift[i]);
}
if constexpr (scaled) {
value = __fmaf_rn(value, input_scale, input_bias);
}
v[j] = value * (1.0f / 16.0f);
} else {
float gate_value, up_value;
if constexpr (dequantize) {
gate_value = (float)((const int32_t*)src_row)[i] * activation_scales[row];
up_value = (float)((const int32_t*)up)[row * up_stride + i] * activation_scales[row];
gate_value *= weight_scales[i];
up_value *= weight_scales[k + i];
} else {
gate_value = src_row[i];
up_value = up[row * up_stride + i];
}
const float activated = sigmoid_gate ? 1.0f / (1.0f + expf(-gate_value)) : (gate_value / (1.0f + expf(-gate_value)));
v[j] = (activated * up_value) * (1.0f / 16.0f);
}
}
// One logical warp owns H256, keeping every butterfly in registers.
#pragma unroll
for (int stride = 1; stride <= 4; stride *= 4) {
const int base = lane & ~(3 * stride);
#pragma unroll
for (int j = 0; j < 8; ++j) {
const float a = cuda_shuffle(v[j], base, 32);
const float b = cuda_shuffle(v[j], base + stride, 32);
const float c = cuda_shuffle(v[j], base + 2 * stride, 32);
const float d = cuda_shuffle(v[j], base + 3 * stride, 32);
v[j] = convrot_h4_cuda(a, b, c, d, (lane / stride) % 4);
}
}
#pragma unroll
for (int j = 0; j < 8; j += 2) {
const float a = cuda_shuffle(v[j], lane % 16, 32);
const float b = cuda_shuffle(v[j], lane % 16 + 16, 32);
const float c = cuda_shuffle(v[j + 1], lane % 16, 32);
const float d = cuda_shuffle(v[j + 1], lane % 16 + 16, 32);
v[j] = convrot_h4_cuda(a, b, c, d, lane / 16);
v[j + 1] = convrot_h4_cuda(a, b, c, d, 2 + lane / 16);
}
#pragma unroll
for (int j = 0; j < 2; ++j) {
const float a = v[j], b = v[j + 2], c = v[j + 4], d = v[j + 6];
#pragma unroll
for (int t = 0; t < 4; ++t) {
values[group * group_size + (j + 2 * t) * 32 + lane] = convrot_h4_cuda(a, b, c, d, t);
}
}
}
#pragma unroll
for (int j = 0; j < 8; j += 2) {
const float a = __shfl_sync(0xffffffff, v[j], lane % 16, 32);
const float b = __shfl_sync(0xffffffff, v[j], lane % 16 + 16, 32);
const float c = __shfl_sync(0xffffffff, v[j + 1], lane % 16, 32);
const float d = __shfl_sync(0xffffffff, v[j + 1], lane % 16 + 16, 32);
v[j] = convrot_h4_cuda(a, b, c, d, lane / 16);
v[j + 1] = convrot_h4_cuda(a, b, c, d, 2 + lane / 16);
}
#pragma unroll
for (int j = 0; j < 2; ++j) {
const float a = v[j], b = v[j + 2], c = v[j + 4], d = v[j + 6];
#pragma unroll
for (int t = 0; t < 4; ++t) {
values[group * group_size + (j + 2 * t) * 32 + lane] = convrot_h4_cuda(a, b, c, d, t);
}
}
}
__syncthreads();
} else {
constexpr int group_threads = group_size / 4;
constexpr int groups_per_wave = 1024 / group_threads;
if constexpr (normalize) {
__shared__ float2 sums[32];
float2 mean_var = make_float2(0.0f, 0.0f);
for (int64_t i = tid; i < k; i += 1024) {
const float value = src_row[i];
mean_var.x += value;
mean_var.y += value * value;
}
mean_var = block_sum(mean_var, sums);
const float mean = mean_var.x / k;
const float variance = mean_var.y / k - mean * mean;
const float inv_std = rsqrtf(variance + eps);
for (int64_t i = tid; i < k; i += 1024) {
const float normalized = __fmul_rn(__fsub_rn(src_row[i], mean), inv_std);
values[i] = __fmul_rn(__fmul_rn(normalized, modulation[i]), 1.0f / 16.0f);
}
__syncthreads();
} else {
for (int64_t i = tid; i < k; i += blockDim.x) {
const float value = scaled ? __fmaf_rn(src_row[i], input_scale, input_bias) : src_row[i];
values[i] = value * (1.0f / 16.0f);
constexpr int group_threads = group_size / 4;
constexpr int groups_per_wave = 1024 / group_threads;
if constexpr (normalize) {
__shared__ float2 sums[32];
float2 mean_var = make_float2(0.0f, 0.0f);
for (int64_t i = tid; i < k; i += 1024) {
const float value = src_row[i];
mean_var.x += value;
mean_var.y += value * value;
}
mean_var = block_sum(mean_var, sums);
const float mean = mean_var.x / k;
const float variance = mean_var.y / k - mean * mean;
const float inv_std = rsqrtf(variance + eps);
for (int64_t i = tid; i < k; i += 1024) {
const float normalized = __fmul_rn(__fsub_rn(src_row[i], mean), inv_std);
values[i] = __fmul_rn(__fmul_rn(normalized, modulation[i]), 1.0f / 16.0f);
}
} else {
for (int64_t i = tid; i < k; i += blockDim.x) {
const float value = scaled ? __fmaf_rn(src_row[i], input_scale, input_bias) : src_row[i];
values[i] = value * (1.0f / 16.0f);
}
}
__syncthreads();
const int group_lane = tid % group_threads;
const int wave_group = tid / group_threads;
const int64_t groups = k / group_size;
for (int64_t group_base = 0; group_base < groups; group_base += groups_per_wave) {
const int64_t group = group_base + wave_group;
#pragma unroll
for (int stride = 1; stride < group_size; stride *= 4) {
if (group < groups) {
float* group_values = values + group * group_size;
const int base = (group_lane / stride) * 4 * stride + group_lane % stride;
const int i0 = base;
const int i1 = i0 + stride;
const int i2 = i1 + stride;
const int i3 = i2 + stride;
const float a = group_values[i0];
const float b = group_values[i1];
const float c = group_values[i2];
const float d = group_values[i3];
group_values[i0] = a + b + c - d;
group_values[i1] = a + b - c + d;
group_values[i2] = a - b + c + d;
group_values[i3] = -a + b + c + d;
}
__syncthreads();
}
}
}
float local_max = 0.0f;
for (int64_t i = tid; i < k; i += blockDim.x) {
local_max = fmaxf(local_max, fabsf(values[i]));
}
local_max = warp_reduce_max<32>(local_max);
const int lane = tid & 31;
const int warp = tid >> 5;
if (lane == 0) {
maxima[warp] = local_max;
}
__syncthreads();
if (warp == 0) {
local_max = lane < blockDim.x / 32 ? maxima[lane] : 0.0f;
local_max = warp_reduce_max<32>(local_max);
if (lane == 0) {
maxima[0] = local_max / 127.0f;
if (scales != nullptr) {
scales[row] = maxima[0];
} else {
*((float*)(dst_row + k)) = maxima[0];
}
}
}
__syncthreads();
const int group_lane = tid % group_threads;
const int wave_group = tid / group_threads;
const int64_t groups = k / group_size;
for (int64_t group_base = 0; group_base < groups; group_base += groups_per_wave) {
const int64_t group = group_base + wave_group;
#pragma unroll
for (int stride = 1; stride < group_size; stride *= 4) {
if (group < groups) {
float* group_values = values + group * group_size;
const int base = (group_lane / stride) * 4 * stride + group_lane % stride;
const int i0 = base;
const int i1 = i0 + stride;
const int i2 = i1 + stride;
const int i3 = i2 + stride;
const float a = group_values[i0];
const float b = group_values[i1];
const float c = group_values[i2];
const float d = group_values[i3];
group_values[i0] = a + b + c - d;
group_values[i1] = a + b - c + d;
group_values[i2] = a - b + c + d;
group_values[i3] = -a + b + c + d;
}
__syncthreads();
}
const float row_scale = maxima[0];
const float inv_scale = row_scale == 0.0f ? 0.0f : 1.0f / row_scale;
for (int64_t i = tid; i < k; i += blockDim.x) {
int value = row_scale == 0.0f ? 0 : __float2int_rn(values[i] * inv_scale);
value = max(-127, min(127, value));
dst_row[i] = (int8_t)value;
}
}
float local_max = 0.0f;
for (int64_t i = tid; i < k; i += blockDim.x) {
local_max = fmaxf(local_max, fabsf(values[i]));
}
local_max = warp_reduce_max<32>(local_max);
const int lane = tid & 31;
const int warp = tid >> 5;
if (lane == 0) {
maxima[warp] = local_max;
}
__syncthreads();
if (warp == 0) {
local_max = lane < blockDim.x / 32 ? maxima[lane] : 0.0f;
local_max = warp_reduce_max<32>(local_max);
if (lane == 0) {
maxima[0] = local_max / 127.0f;
if (scales != nullptr) {
scales[row] = maxima[0];
} else {
*((float*)(dst_row + k)) = maxima[0];
}
}
}
__syncthreads();
const float row_scale = maxima[0];
const float inv_scale = row_scale == 0.0f ? 0.0f : 1.0f / row_scale;
for (int64_t i = tid; i < k; i += blockDim.x) {
int value = row_scale == 0.0f ? 0 : __float2int_rn(values[i] * inv_scale);
value = max(-127, min(127, value));
dst_row[i] = (int8_t)value;
}
}
}
#endif
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;
} else if (graph_node(graph, first)->op == GGML_OP_RMS_NORM) {
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))
return 0;
offset = 3;
@@ -608,7 +608,7 @@ static void attention_optimize(ggml_backend_dev_t device, ggml_cgraph* graph, vo
continue;
for (int j = i - 3; j >= 0 && j >= i - 32; --j) {
if (attn->src[0] == graph_node(graph, j + 2) &&
sd_cuda_rope_fusion_supports(graph_node(graph, j), graph_node(graph, j + 1), graph_node(graph, j + 2))) {
sd_cuda::rope_fusion_supports(graph_node(graph, j), graph_node(graph, j + 1), graph_node(graph, j + 2))) {
if (j >= 2 && sd_cuda_i8_rope_supports(device, graph_node(graph, j - 2), graph_node(graph, j - 1), graph_node(graph, j), graph_node(graph, j + 1), graph_node(graph, j + 2)) &&
ggml_ext_graph_delay(graph, j - 2, 5, i))
break;
@@ -687,7 +687,7 @@ static ggml_status attention_execute(const ggml_ext_execution_context* execution
if (sd_cuda_i8_rope_launch(stream, nodes, work[0], scale) != cudaSuccess)
return GGML_STATUS_FAILED;
} else if (i == 0 && plan.index == 3) {
if (sd_cuda_rms_rope_launch(stream, nodes, work[0], true, scale) != cudaSuccess)
if (sd_cuda::rms_rope_launch(stream, nodes, work[0], true, scale) != cudaSuccess)
return GGML_STATUS_FAILED;
} else if (i != 0 && src->type == GGML_TYPE_F16 && ggml_is_contiguous(src) && uintptr_t(src->data) % 16 == 0) {
inputs[i] = static_cast<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) {
if (!projection_supports(device, mat, view) || norm->src[0] != view ||
!sd_cuda_rope_fusion_supports(norm, mul, rope) || rope->src[1]->ne[2] != 64 || !aligned(mul->src[1]) || !aligned(rope->src[1]))
!sd_cuda::rope_fusion_supports(norm, mul, rope) || rope->src[1]->ne[2] != 64 || !aligned(mul->src[1]) || !aligned(rope->src[1]))
return false;
int32_t interleaved;
std::memcpy(&interleaved, ggml_ext_op_get(rope) ? ggml_ext_op_attributes(rope) : rope->op_params, sizeof(interleaved));
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/ggml_extensions.h"
#include <cuda_fp16.h>
#include <climits>
#include <cstdlib>
#include <cstring>
#include <type_traits>
static __device__ __forceinline__ float sd_warp_sum(float v) {
for (int mask = 16; mask > 0; mask >>= 1)
v = __fadd_rn(v, __shfl_xor_sync(0xffffffff, v, mask));
return v;
}
template <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;
namespace SD_CUDA_NAMESPACE {
static __device__ __forceinline__ float sd_warp_sum(float v) {
for (int mask = 16; mask > 0; mask >>= 1)
v = __fadd_rn(v, cuda_shuffle_xor(v, mask));
return v;
}
const int64_t head = row % heads;
const int64_t token = (row / heads) % tokens;
const int64_t batch = row / (heads * tokens);
const float* x_row = x + head * nb1 + token * nb2 + batch * nb3;
float v[D / 32], sums[8] = {};
template <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;
const int64_t token = (row / heads) % tokens;
const int64_t batch = row / (heads * tokens);
const float* x_row = x + head * nb1 + token * nb2 + batch * nb3;
float v[D / 32], sums[8] = {};
#pragma unroll
for (int j = 0; j < D / 32; ++j) {
v[j] = x_row[j * 32 + lane];
sums[j] = sd_warp_sum(__fmul_rn(v[j], v[j]));
}
// Preserve the 256-thread RMSNorm kernel's reduction tree, including inactive warps.
const float sum = __fadd_rn(
__fadd_rn(__fadd_rn(sums[0], sums[4]), __fadd_rn(sums[2], sums[6])),
__fadd_rn(__fadd_rn(sums[1], sums[5]), __fadd_rn(sums[3], sums[7])));
const float scale = rsqrtf(sum / float(D) + eps);
T* out = dst + ((batch * heads + head) * tokens + token) * D;
for (int j = 0; j < D / 32; ++j) {
v[j] = x_row[j * 32 + lane];
sums[j] = sd_warp_sum(__fmul_rn(v[j], v[j]));
}
// Preserve the 256-thread RMSNorm kernel's reduction tree, including inactive warps.
const float sum = __fadd_rn(
__fadd_rn(__fadd_rn(sums[0], sums[4]), __fadd_rn(sums[2], sums[6])),
__fadd_rn(__fadd_rn(sums[1], sums[5]), __fadd_rn(sums[3], sums[7])));
const float scale = rsqrtf(sum / float(D) + eps);
T* out = dst + ((batch * heads + head) * tokens + token) * D;
#pragma unroll
for (int j = 0; j < D / 32; ++j) {
const int col = j * 32 + lane;
const float value = __fmul_rn(__fmul_rn(scale, v[j]), weight[col]);
float rotated = value;
if constexpr (interleaved) {
const float other = __shfl_xor_sync(0xffffffff, value, 1, 32);
if (col < 2 * rot_half) {
const float* m = pe + 4 * (token * rot_half + col / 2);
const float4 matrix = make_float4(m[0], m[1], m[2], m[3]);
rotated = lane % 2 == 0
? __fadd_rn(__fmul_rn(value, matrix.x), __fmul_rn(other, matrix.y))
: __fadd_rn(__fmul_rn(other, matrix.z), __fmul_rn(value, matrix.w));
for (int j = 0; j < D / 32; ++j) {
const int col = j * 32 + lane;
const float value = __fmul_rn(__fmul_rn(scale, v[j]), weight[col]);
float rotated = value;
if constexpr (interleaved) {
const float other = cuda_shuffle_xor(value, 1, 32);
if (col < 2 * rot_half) {
const float* m = pe + 4 * (token * rot_half + col / 2);
const float4 matrix = make_float4(m[0], m[1], m[2], m[3]);
rotated = lane % 2 == 0
? __fadd_rn(__fmul_rn(value, matrix.x), __fmul_rn(other, matrix.y))
: __fadd_rn(__fmul_rn(other, matrix.z), __fmul_rn(value, matrix.w));
}
} else {
if (col < 2 * rot_half) {
const int partner = col < rot_half ? col + rot_half : col - rot_half;
const float other = __fmul_rn(__fmul_rn(scale, x_row[partner]), weight[partner]);
const float* m = pe + 4 * (token * rot_half + col % rot_half);
const float4 matrix = make_float4(m[0], m[1], m[2], m[3]);
rotated = col < rot_half
? __fsub_rn(__fmul_rn(value, matrix.x), __fmul_rn(other, matrix.z))
: __fadd_rn(__fmul_rn(other, matrix.z), __fmul_rn(value, matrix.x));
}
}
} else {
if (col < 2 * rot_half) {
const int partner = col < rot_half ? col + rot_half : col - rot_half;
const float other = __fmul_rn(__fmul_rn(scale, x_row[partner]), weight[partner]);
const float* m = pe + 4 * (token * rot_half + col % rot_half);
const float4 matrix = make_float4(m[0], m[1], m[2], m[3]);
rotated = col < rot_half
? __fsub_rn(__fmul_rn(value, matrix.x), __fmul_rn(other, matrix.z))
: __fadd_rn(__fmul_rn(other, matrix.z), __fmul_rn(value, matrix.x));
if constexpr (std::is_same<T, half>::value) {
out[col] = __float2half_rn(__fmul_rn(rotated, output_scale));
} else {
out[col] = rotated;
}
}
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 {
out[col] = rotated;
out[i0] = __fsub_rn(__fmul_rn(a, c), __fmul_rn(b, s));
out[i1] = __fadd_rn(__fmul_rn(a, s), __fmul_rn(b, c));
}
}
}
struct RopeApplyStrides {
int64_t x[4];
int64_t pe[4];
};
template <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;
static bool rope_interleaved(const ggml_tensor* rope) {
int32_t value;
std::memcpy(&value, ggml_ext_op_get(rope) ? ggml_ext_op_attributes(rope) : rope->op_params, sizeof(value));
return value != 0;
}
const int64_t j = index % half_dim;
const int64_t row = index / half_dim;
const int64_t token = row % tokens;
const int64_t head = (row / tokens) % heads;
const int64_t batch = row / (tokens * heads);
const int64_t i0 = interleaved || j >= rot_half ? 2 * j : j;
const int64_t i1 = interleaved || j >= rot_half ? 2 * j + 1 : j + rot_half;
const float* x_row = x + head * strides.x[1] + token * strides.x[2] + batch * strides.x[3];
const float a = x_row[i0 * strides.x[0]];
const float b = x_row[i1 * strides.x[0]];
float* out = dst + row * half_dim * 2;
if (j >= rot_half) {
out[i0] = a;
out[i1] = b;
return;
static float param_f32(const ggml_tensor* node, int index) {
float value;
std::memcpy(&value, node->op_params + index, sizeof(value));
return value;
}
const float* matrix = pe + j * strides.pe[2] + token * strides.pe[3];
const float c = matrix[0];
const float s = matrix[strides.pe[1]];
// Match the separate mul/add rounding of the unfused graph.
if (interleaved) {
out[i0] = __fadd_rn(__fmul_rn(a, c), __fmul_rn(b, matrix[strides.pe[0]]));
out[i1] = __fadd_rn(__fmul_rn(a, s), __fmul_rn(b, matrix[strides.pe[0] + strides.pe[1]]));
} else {
out[i0] = __fsub_rn(__fmul_rn(a, c), __fmul_rn(b, s));
out[i1] = __fadd_rn(__fmul_rn(a, s), __fmul_rn(b, c));
static bool rope_is_node(const ggml_tensor* node) {
return node && (node->op == GGML_OP_ROPE_APPLY || ggml_ext_op_get(node) == sd_ggml_rope_op());
}
}
static bool rope_interleaved(const ggml_tensor* rope) {
int32_t value;
std::memcpy(&value, ggml_ext_op_get(rope) ? ggml_ext_op_attributes(rope) : rope->op_params, sizeof(value));
return value != 0;
}
static float param_f32(const ggml_tensor* node, int index) {
float value;
std::memcpy(&value, node->op_params + index, sizeof(value));
return value;
}
bool sd_cuda_rope_is_node(const ggml_tensor* node) {
return node && (node->op == GGML_OP_ROPE_APPLY || ggml_ext_op_get(node) == sd_ggml_rope_op());
}
static bool rope_supports(ggml_backend_dev_t, const ggml_tensor* rope, void*) {
const auto* x = rope->src[0];
const auto* pe = rope->src[1];
if (!x || !pe || x->type != GGML_TYPE_F32 || pe->type != GGML_TYPE_F32 || rope->type != GGML_TYPE_F32 ||
x->ne[0] % 2 || pe->ne[0] != 2 || pe->ne[1] != 2 || pe->ne[2] <= 0 || pe->ne[2] > x->ne[0] / 2 ||
pe->ne[3] != x->ne[2] || !ggml_is_contiguous(rope) || rope->ne[0] != x->ne[0] ||
rope->ne[1] != x->ne[2] || rope->ne[2] != x->ne[1] * x->ne[3] || rope->ne[3] != 1 ||
(ggml_nelements(x) / 2 + 255) / 256 > INT_MAX) {
return false;
}
for (int i = 0; i < 4; ++i) {
if (x->nb[i] % sizeof(float) || pe->nb[i] % sizeof(float))
static bool rope_supports(ggml_backend_dev_t, const ggml_tensor* rope, void*) {
const auto* x = rope->src[0];
const auto* pe = rope->src[1];
if (!x || !pe || x->type != GGML_TYPE_F32 || pe->type != GGML_TYPE_F32 || rope->type != GGML_TYPE_F32 ||
x->ne[0] % 2 || pe->ne[0] != 2 || pe->ne[1] != 2 || pe->ne[2] <= 0 || pe->ne[2] > x->ne[0] / 2 ||
pe->ne[3] != x->ne[2] || !ggml_is_contiguous(rope) || rope->ne[0] != x->ne[0] ||
rope->ne[1] != x->ne[2] || rope->ne[2] != x->ne[1] * x->ne[3] || rope->ne[3] != 1 ||
(ggml_nelements(x) / 2 + 255) / 256 > INT_MAX) {
return false;
}
for (int i = 0; i < 4; ++i) {
if (x->nb[i] % sizeof(float) || pe->nb[i] % sizeof(float))
return false;
}
return true;
}
return true;
}
bool sd_cuda_rope_fusion_supports(const ggml_tensor* norm, const ggml_tensor* mul, const ggml_tensor* rope) {
if (norm->op != GGML_OP_RMS_NORM || mul->op != GGML_OP_MUL || !sd_cuda_rope_is_node(rope) ||
mul->src[0] != norm || rope->src[0] != mul || !rope_supports(nullptr, rope, nullptr))
return false;
const auto* x = norm->src[0];
const auto* weight = mul->src[1];
return x && weight && x->ne[0] >= 32 && x->ne[0] <= 256 && x->ne[0] % 32 == 0 &&
x->type == GGML_TYPE_F32 && weight->type == GGML_TYPE_F32 && norm->type == GGML_TYPE_F32 &&
ggml_are_same_shape(norm, mul) && ggml_are_same_shape(x, norm) &&
weight->ne[0] == x->ne[0] && ggml_nelements(weight) == x->ne[0] &&
x->nb[0] == sizeof(float) && x->nb[1] % sizeof(float) == 0 && x->nb[2] % sizeof(float) == 0 &&
x->nb[3] % sizeof(float) == 0 && ggml_is_contiguous(weight) && ggml_is_contiguous(rope->src[1]) &&
(ggml_nrows(x) + 3) / 4 <= INT_MAX;
}
bool rope_fusion_supports(const ggml_tensor* norm, const ggml_tensor* mul, const ggml_tensor* rope) {
if (norm->op != GGML_OP_RMS_NORM || mul->op != GGML_OP_MUL || !rope_is_node(rope) ||
mul->src[0] != norm || rope->src[0] != mul || !rope_supports(nullptr, rope, nullptr))
return false;
const auto* x = norm->src[0];
const auto* weight = mul->src[1];
return x && weight && x->ne[0] >= 32 && x->ne[0] <= 256 && x->ne[0] % 32 == 0 &&
x->type == GGML_TYPE_F32 && weight->type == GGML_TYPE_F32 && norm->type == GGML_TYPE_F32 &&
ggml_are_same_shape(norm, mul) && ggml_are_same_shape(x, norm) &&
weight->ne[0] == x->ne[0] && ggml_nelements(weight) == x->ne[0] &&
x->nb[0] == sizeof(float) && x->nb[1] % sizeof(float) == 0 && x->nb[2] % sizeof(float) == 0 &&
x->nb[3] % sizeof(float) == 0 && ggml_is_contiguous(weight) && ggml_is_contiguous(rope->src[1]) &&
(ggml_nrows(x) + 3) / 4 <= INT_MAX;
}
template <typename T, bool Interleaved>
static cudaError_t rms_rope_launch(cudaStream_t stream, ggml_tensor* const* nodes, T* output, float scale) {
const auto* x = nodes[0]->src[0];
const auto* pe = nodes[2]->src[1];
const int64_t rows = ggml_nrows(x);
template <typename T, bool Interleaved>
static CUDAError rms_rope_launch(CUDAStream stream, ggml_tensor* const* nodes, T* output, float scale) {
const auto* x = nodes[0]->src[0];
const auto* pe = nodes[2]->src[1];
const int64_t rows = ggml_nrows(x);
#define SD_ROPE(D) \
case D: \
rms_norm_mul_rope_apply_fused<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, \
x->nb[1] / 4, x->nb[2] / 4, x->nb[3] / 4, int(pe->ne[2]), param_f32(nodes[0], 0), scale); \
break
switch (x->ne[0]) {
SD_ROPE(32);
SD_ROPE(64);
SD_ROPE(96);
SD_ROPE(128);
SD_ROPE(160);
SD_ROPE(192);
SD_ROPE(224);
SD_ROPE(256);
default:
return cudaErrorInvalidValue;
}
switch (x->ne[0]) {
SD_ROPE(32);
SD_ROPE(64);
SD_ROPE(96);
SD_ROPE(128);
SD_ROPE(160);
SD_ROPE(192);
SD_ROPE(224);
SD_ROPE(256);
default:
return cuda_invalid_value;
}
#undef SD_ROPE
return cudaGetLastError();
}
cudaError_t sd_cuda_rms_rope_launch(cudaStream_t stream, ggml_tensor* const* nodes, void* output, bool fp16, float scale) {
if (fp16) {
return rope_interleaved(nodes[2]) ? rms_rope_launch<half, true>(stream, nodes, static_cast<half*>(output), scale)
: rms_rope_launch<half, false>(stream, nodes, static_cast<half*>(output), scale);
return cuda_last_error();
}
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*) {
static const bool disabled = [] {const char* value=std::getenv("GGML_CUDA_DISABLE_FUSION"); return value && std::atoi(value)!=0; }();
if (disabled)
return 0;
auto* g = const_cast<ggml_cgraph*>(graph);
if (first + 2 >= ggml_graph_n_nodes(g) || !sd_cuda_rope_fusion_supports(
ggml_graph_node(g, first), ggml_graph_node(g, first + 1), ggml_graph_node(g, first + 2)))
return 0;
if (first + 4 < ggml_graph_n_nodes(g)) {
const auto* rope = ggml_graph_node(g, first + 2);
const auto* scale = ggml_graph_node(g, first + 3);
const auto* copy = ggml_graph_node(g, first + 4);
if (scale->op == GGML_OP_SCALE && copy->op == GGML_OP_CPY && scale->src[0] == rope && copy->src[0] == scale &&
scale->type == GGML_TYPE_F32 && copy->type == GGML_TYPE_F16 && param_f32(scale, 1) == 0.f &&
ggml_are_same_shape(rope, scale) && ggml_are_same_shape(rope, copy) && ggml_is_contiguous(copy))
return 5;
CUDAError rms_rope_launch(CUDAStream stream, ggml_tensor* const* nodes, void* output, bool fp16, float scale) {
if (fp16) {
return rope_interleaved(nodes[2]) ? rms_rope_launch<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);
}
return 3;
}
static void rope_optimize(ggml_backend_dev_t, ggml_cgraph* graph, void*) {
static const bool disabled = [] {
const char* value = std::getenv("GGML_CUDA_DISABLE_FUSION");
return value && std::atoi(value) != 0;
}();
if (disabled)
return;
const int total = ggml_graph_n_nodes(graph);
for (int i = 2; i + 1 < total; ++i) {
const auto* mul = ggml_graph_node(graph, i);
const auto* rope = ggml_graph_node(graph, i + 1);
if (mul->op != GGML_OP_MUL || !sd_cuda_rope_is_node(rope) || rope->src[0] != mul ||
!mul->src[0] || mul->src[0]->op != GGML_OP_RMS_NORM)
continue;
const auto* norm = mul->src[0];
for (int j = i - 1; j >= 0 && j >= i - 64; --j) {
if (ggml_graph_node(graph, j) != norm)
static int rope_match(ggml_backend_dev_t, const ggml_cgraph* graph, int first, void*) {
static const bool disabled = [] {const char* value=std::getenv("GGML_CUDA_DISABLE_FUSION"); return value && std::atoi(value)!=0; }();
if (disabled)
return 0;
auto* g = const_cast<ggml_cgraph*>(graph);
if (first + 2 >= ggml_graph_n_nodes(g) || !rope_fusion_supports(
ggml_graph_node(g, first), ggml_graph_node(g, first + 1), ggml_graph_node(g, first + 2)))
return 0;
if (first + 4 < ggml_graph_n_nodes(g)) {
const auto* rope = ggml_graph_node(g, first + 2);
const auto* scale = ggml_graph_node(g, first + 3);
const auto* copy = ggml_graph_node(g, first + 4);
if (scale->op == GGML_OP_SCALE && copy->op == GGML_OP_CPY && scale->src[0] == rope && copy->src[0] == scale &&
scale->type == GGML_TYPE_F32 && copy->type == GGML_TYPE_F16 && param_f32(scale, 1) == 0.f &&
ggml_are_same_shape(rope, scale) && ggml_are_same_shape(rope, copy) && ggml_is_contiguous(copy))
return 5;
}
return 3;
}
static void rope_optimize(ggml_backend_dev_t, ggml_cgraph* graph, void*) {
static const bool disabled = [] {
const char* value = std::getenv("GGML_CUDA_DISABLE_FUSION");
return value && std::atoi(value) != 0;
}();
if (disabled)
return;
const int total = ggml_graph_n_nodes(graph);
for (int i = 2; i + 1 < total; ++i) {
const auto* mul = ggml_graph_node(graph, i);
const auto* rope = ggml_graph_node(graph, i + 1);
if (mul->op != GGML_OP_MUL || !rope_is_node(rope) || rope->src[0] != mul ||
!mul->src[0] || mul->src[0]->op != GGML_OP_RMS_NORM)
continue;
if (j + 1 < i) {
if (j >= 2) {
const auto* mat = ggml_graph_node(graph, j - 2);
const auto* view = ggml_graph_node(graph, j - 1);
if (mat->op == GGML_OP_MUL_MAT && mat->src[0]->type == GGML_TYPE_I8 &&
view->op == GGML_OP_RESHAPE && view->src[0] == mat && norm->src[0] == view &&
ggml_ext_graph_delay(graph, j - 2, 3, i))
break;
const auto* norm = mul->src[0];
for (int j = i - 1; j >= 0 && j >= i - 64; --j) {
if (ggml_graph_node(graph, j) != norm)
continue;
if (j + 1 < i) {
if (j >= 2) {
const auto* mat = ggml_graph_node(graph, j - 2);
const auto* view = ggml_graph_node(graph, j - 1);
if (mat->op == GGML_OP_MUL_MAT && mat->src[0]->type == GGML_TYPE_I8 &&
view->op == GGML_OP_RESHAPE && view->src[0] == mat && norm->src[0] == view &&
ggml_ext_graph_delay(graph, j - 2, 3, i))
break;
}
ggml_ext_graph_delay(graph, j, 1, i);
}
ggml_ext_graph_delay(graph, j, 1, i);
break;
}
break;
}
}
}
static ggml_status rope_prepare(ggml_backend_t, ggml_tensor* const* nodes, int count, void**, size_t* workspace, void*) {
*workspace = ggml_nbytes(nodes[count - 1]);
return GGML_STATUS_SUCCESS;
}
static ggml_status rope_prepare(ggml_backend_t, ggml_tensor* const* nodes, int count, void**, size_t* workspace, void*) {
*workspace = ggml_nbytes(nodes[count - 1]);
return GGML_STATUS_SUCCESS;
}
static void rope_destroy(void*, void*) {}
static void rope_destroy(void*, void*) {}
static bool overlaps(const ggml_tensor* a, const ggml_tensor* b) {
const auto x = reinterpret_cast<uintptr_t>(a->data), y = reinterpret_cast<uintptr_t>(b->data);
return x <= y ? y - x < ggml_nbytes(a) : x - y < ggml_nbytes(b);
}
static bool overlaps(const ggml_tensor* a, const ggml_tensor* b) {
const auto x = reinterpret_cast<uintptr_t>(a->data), y = reinterpret_cast<uintptr_t>(b->data);
return x <= y ? y - x < ggml_nbytes(a) : x - y < ggml_nbytes(b);
}
static ggml_status rope_execute(const ggml_ext_execution_context* execution, ggml_tensor* const* nodes, int count, void*, void*) {
auto reg = ggml_backend_dev_backend_reg(ggml_backend_get_device(execution->backend));
auto query = reinterpret_cast<ggml_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))
return GGML_STATUS_FAILED;
auto stream = static_cast<cudaStream_t>(context.stream);
const auto* rope = nodes[count == 1 ? 0 : 2];
const auto* x = count == 1 ? rope->src[0] : nodes[0]->src[0];
auto* dst = nodes[count - 1];
const bool alias = overlaps(dst, x) || overlaps(dst, rope->src[1]) || (count > 1 && overlaps(dst, nodes[1]->src[1]));
void* output = alias ? ggml_backend_buffer_get_base(execution->workspace) : dst->data;
cudaError_t error;
if (count > 1) {
error = sd_cuda_rms_rope_launch(stream, nodes, output, count == 5, count == 5 ? param_f32(nodes[3], 0) : 1.f);
} else {
RopeApplyStrides strides;
const auto* pe = rope->src[1];
for (int i = 0; i < 4; ++i) {
strides.x[i] = x->nb[i] / 4;
strides.pe[i] = pe->nb[i] / 4;
static ggml_status rope_execute(const ggml_ext_execution_context* execution, ggml_tensor* const* nodes, int count, void*, void*) {
ggml_ext_gpu_context context;
if (!cuda_context(execution->backend, &context))
return GGML_STATUS_FAILED;
auto stream = static_cast<CUDAStream>(context.stream);
const auto* rope = nodes[count == 1 ? 0 : 2];
const auto* x = count == 1 ? rope->src[0] : nodes[0]->src[0];
auto* dst = nodes[count - 1];
const bool alias = overlaps(dst, x) || overlaps(dst, rope->src[1]) || (count > 1 && overlaps(dst, nodes[1]->src[1]));
void* output = alias ? ggml_backend_buffer_get_base(execution->workspace) : dst->data;
CUDAError error;
if (count > 1) {
error = rms_rope_launch(stream, nodes, output, count == 5, count == 5 ? param_f32(nodes[3], 0) : 1.f);
} else {
RopeApplyStrides strides;
const auto* pe = rope->src[1];
for (int i = 0; i < 4; ++i) {
strides.x[i] = x->nb[i] / 4;
strides.pe[i] = pe->nb[i] / 4;
}
const int64_t pairs = ggml_nelements(x) / 2;
auto kernel = rope_interleaved(rope) ? rope_apply_f32<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;
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 = cudaGetLastError();
if (error == cuda_success && alias)
error = cuda_copy_async(dst->data, output, ggml_nbytes(dst), stream);
return error == cuda_success ? GGML_STATUS_SUCCESS : GGML_STATUS_FAILED;
}
void register_rope(ggml_backend_dev_t device) {
ggml_ext_kernel_desc desc = {};
desc.abi_version = GGML_EXT_ABI_VERSION;
desc.name = "sd.rope_apply." SD_CUDA_NAME ".v1";
desc.op = GGML_OP_CUSTOM;
desc.custom_op = sd_ggml_rope_op();
desc.priority = 100;
desc.supports = rope_supports;
desc.prepare = rope_prepare;
desc.destroy = rope_destroy;
desc.execute = rope_execute;
GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
desc.name = "sd.rope_apply." SD_CUDA_NAME ".override.v1";
desc.op = GGML_OP_ROPE_APPLY;
desc.custom_op = nullptr;
GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
desc.name = "sd.rms_norm_mul_rope." SD_CUDA_NAME ".v1";
desc.op = GGML_OP_RMS_NORM;
desc.priority = 250;
desc.supports = nullptr;
desc.match = rope_match;
desc.optimize = rope_optimize;
GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
}
if (error == cudaSuccess && alias)
error = cudaMemcpyAsync(dst->data, output, ggml_nbytes(dst), cudaMemcpyDeviceToDevice, stream);
return error == cudaSuccess ? GGML_STATUS_SUCCESS : GGML_STATUS_FAILED;
}
void sd_ggml_register_cuda_rope(ggml_backend_dev_t device) {
ggml_ext_kernel_desc desc = {};
desc.abi_version = GGML_EXT_ABI_VERSION;
desc.name = "sd.rope_apply.cuda.v1";
desc.op = GGML_OP_CUSTOM;
desc.custom_op = sd_ggml_rope_op();
desc.priority = 100;
desc.supports = rope_supports;
desc.prepare = rope_prepare;
desc.destroy = rope_destroy;
desc.execute = rope_execute;
GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
desc.name = "sd.rope_apply.cuda.override.v1";
desc.op = GGML_OP_ROPE_APPLY;
desc.custom_op = nullptr;
GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
desc.name = "sd.rms_norm_mul_rope.cuda.v1";
desc.op = GGML_OP_RMS_NORM;
desc.priority = 250;
desc.supports = nullptr;
desc.match = rope_match;
desc.optimize = rope_optimize;
GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
}
src/core/cuda/rope.h
+5 -6
View File
@@ -1,12 +1,11 @@
#ifndef __SD_CORE_CUDA_ROPE_H__
#define __SD_CORE_CUDA_ROPE_H__
#include <cuda_runtime.h>
#include "ggml-ext.h"
#include "runtime.h"
bool sd_cuda_rope_is_node(const ggml_tensor* node);
bool sd_cuda_rope_fusion_supports(const ggml_tensor* norm, const ggml_tensor* mul, const ggml_tensor* rope);
cudaError_t sd_cuda_rms_rope_launch(cudaStream_t stream, ggml_tensor* const* nodes, void* output, bool fp16, float scale);
void sd_ggml_register_cuda_rope(ggml_backend_dev_t device);
namespace SD_CUDA_NAMESPACE {
bool rope_fusion_supports(const ggml_tensor* norm, const ggml_tensor* mul, const ggml_tensor* rope);
CUDAError rms_rope_launch(CUDAStream stream, ggml_tensor* const* nodes, void* output, bool fp16, float scale);
}
#endif
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 <unordered_set>
#ifdef SD_VULKAN_EXTENSIONS
void sd_ggml_register_vulkan_rope(ggml_backend_dev_t device);
#endif
#ifdef SD_CUDA_EXTENSIONS
void sd_ggml_register_cuda_affine(ggml_backend_dev_t device);
void sd_ggml_register_cuda_i8(ggml_backend_dev_t device);
void sd_ggml_register_cuda_convrot(ggml_backend_dev_t device);
void sd_ggml_register_cuda_rope(ggml_backend_dev_t device);
#endif
#ifdef SD_CUDA_CUTE_ATTENTION
void sd_ggml_register_cuda_i8_gated_residual(ggml_backend_dev_t device);
void sd_ggml_register_cuda_i8_projection(ggml_backend_dev_t device);
void sd_ggml_register_cuda_attention(ggml_backend_dev_t device);
#ifdef SD_BACKEND_EXTENSIONS
#include "sd_backend_extensions.h"
#endif
static bool rope_infer(ggml_tensor* const* inputs, int count, const void*, ggml_ext_tensor_desc* out) {
@@ -171,37 +160,19 @@ void sd_ggml_register_extensions(ggml_backend_t backend) {
if (!backend) {
return;
}
auto device = ggml_backend_get_device(backend);
auto reg = ggml_backend_dev_backend_reg(device);
const bool cpu = ggml_backend_reg_get_proc_address(reg, "ggml_ext_cpu_parallel_v1") != nullptr;
const bool cuda = ggml_backend_reg_get_proc_address(reg, "ggml_ext_cuda_context_v1") != nullptr;
const bool vulkan = ggml_backend_reg_get_proc_address(reg, "ggml_ext_vulkan_v1") != nullptr;
if (!cpu && !cuda && !vulkan) {
return;
}
auto device = ggml_backend_get_device(backend);
auto reg = ggml_backend_dev_backend_reg(device);
const bool cpu = ggml_backend_reg_get_proc_address(reg, "ggml_ext_cpu_parallel_v1") != nullptr;
static std::mutex mutex;
static std::unordered_set<ggml_backend_dev_t> registered;
std::lock_guard<std::mutex> lock(mutex);
if (!registered.insert(device).second) {
return;
}
if (vulkan) {
#ifdef SD_VULKAN_EXTENSIONS
sd_ggml_register_vulkan_rope(device);
#endif
return;
}
if (cuda) {
#ifdef SD_CUDA_EXTENSIONS
sd_ggml_register_cuda_affine(device);
sd_ggml_register_cuda_i8(device);
sd_ggml_register_cuda_convrot(device);
sd_ggml_register_cuda_rope(device);
#endif
#ifdef SD_CUDA_CUTE_ATTENTION
sd_ggml_register_cuda_i8_projection(device);
sd_ggml_register_cuda_i8_gated_residual(device);
sd_ggml_register_cuda_attention(device);
if (!cpu) {
#ifdef SD_BACKEND_EXTENSIONS
for (auto register_backend : sd_backend_extensions)
register_backend(device);
#endif
return;
}
@@ -234,11 +205,8 @@ ggml_tensor* sd_ggml_rope_apply(ggml_context* ctx, ggml_backend_t backend, ggml_
if (backend) {
auto reg = ggml_backend_dev_backend_reg(ggml_backend_get_device(backend));
if (ggml_backend_reg_get_proc_address(reg, "ggml_ext_cpu_parallel_v1")
#ifdef SD_CUDA_EXTENSIONS
|| ggml_backend_reg_get_proc_address(reg, "ggml_ext_cuda_context_v1")
#endif
#ifdef SD_VULKAN_EXTENSIONS
|| ggml_backend_reg_get_proc_address(reg, "ggml_ext_vulkan_v1")
#ifdef SD_BACKEND_EXTENSIONS
|| ggml_backend_reg_get_proc_address(reg, "ggml_ext_gpu_context_v1") || ggml_backend_reg_get_proc_address(reg, "ggml_ext_vulkan_v1")
#endif
) {
ggml_tensor* inputs[] = {x, pe};
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) {
if (!sd_vulkan_rope::api(device))
return;
ggml_ext_kernel_desc desc = {};
desc.abi_version = GGML_EXT_ABI_VERSION;
desc.name = "sd.rope_apply.vulkan.v1";