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

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