From 7ab323ef25efbb2d9cef01d00f795618786e5a24 Mon Sep 17 00:00:00 2001 From: leejet Date: Sat, 31 Oct 2026 23:39:19 +0800 Subject: [PATCH] refactor: move Vulkan RoPE kernels into extensions --- CMakeLists.txt | 4 + cmake/embed_spirv.cmake | 3 + docs/ggml_extensions.md | 61 +++++- ggml | 2 +- src/core/ggml_extensions.cpp | 23 +- src/core/vulkan/CMakeLists.txt | 27 +++ src/core/vulkan/rope.cpp | 225 ++++++++++++++++++++ src/core/vulkan/shaders/rms_rope_apply.comp | 97 +++++++++ src/core/vulkan/shaders/rope_apply.comp | 59 +++++ 9 files changed, 493 insertions(+), 8 deletions(-) create mode 100644 cmake/embed_spirv.cmake create mode 100644 src/core/vulkan/CMakeLists.txt create mode 100644 src/core/vulkan/rope.cpp create mode 100644 src/core/vulkan/shaders/rms_rope_apply.comp create mode 100644 src/core/vulkan/shaders/rope_apply.comp diff --git a/CMakeLists.txt b/CMakeLists.txt index 4b9bf559..123307d5 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -366,6 +366,10 @@ if (SD_CUDA) set_property(SOURCE src/core/ggml_extensions.cpp APPEND PROPERTY COMPILE_DEFINITIONS SD_CUDA_EXTENSIONS) endif() endif() +if (SD_VULKAN AND NOT SD_USE_UPSTREAM_GGML) + add_subdirectory(src/core/vulkan) + set_property(SOURCE src/core/ggml_extensions.cpp APPEND PROPERTY COMPILE_DEFINITIONS SD_VULKAN_EXTENSIONS) +endif() if (SD_HIPBLAS) find_package(hip REQUIRED CONFIG HINTS "$ENV{ROCM_PATH}" /opt/rocm) target_link_libraries(${SD_LIB} PRIVATE hip::host) diff --git a/cmake/embed_spirv.cmake b/cmake/embed_spirv.cmake new file mode 100644 index 00000000..a6642096 --- /dev/null +++ b/cmake/embed_spirv.cmake @@ -0,0 +1,3 @@ +file(READ "${INPUT}" SPIRV HEX) +string(REGEX REPLACE "(..)(..)(..)(..)" "0x\\4\\3\\2\\1," SPIRV_WORDS "${SPIRV}") +file(WRITE "${OUTPUT}" "#include \nnamespace sd_vulkan_shaders {\nstatic const uint32_t ${NAME}[] = {${SPIRV_WORDS}};\n}\n") diff --git a/docs/ggml_extensions.md b/docs/ggml_extensions.md index 5644858a..1ff829ea 100644 --- a/docs/ggml_extensions.md +++ b/docs/ggml_extensions.md @@ -313,7 +313,64 @@ CUDA VMM so native scratch addresses remain valid for the backend lifetime; legacy scratch pools retain the existing fragment path. HIP and MUSA do not yet provide this plan executor. sd.cpp's CUDA registrations opt in to reusable replay. -Remaining CUDA fusion rules and Vulkan INT8/ConvRot implementations still need -to be externalized. Vulkan execution services and kernels also remain pending. +## Vulkan execution services + +`ggml-ext-vulkan.h` exposes `ggml_ext_vulkan_v1` through the Vulkan backend +registry. The versioned function table reports device limits and subgroup +capabilities, resolves tensor/workspace byte ranges, dispatches application +SPIR-V, and copies nonoverlapping buffer ranges. Extensions do not include +private ggml Vulkan headers or own backend command buffers and descriptors. + +Shaders use the `main` entry point, consecutive storage-buffer bindings starting +at zero, and consecutive 32-bit specialization constants. Dispatch accepts actual +workgroup counts. The backend validates ranges and dispatch limits, caches the +complete shader/configuration per device, and releases pipelines with that +device. Push constants and specialization values are copied; the backend does +not retain application execution state. Pipeline creation occurs on execution, +never during host-only plan preparation. + +Bindings identify either a tensor or a backend buffer, with a byte offset and +optional byte count. `buffer_info` returns an opaque storage identity and byte +range for alias checks. Vulkan descriptor offsets are rounded down to the +device's required alignment; shaders must add the returned `shader_offset`. +Workspace bindings follow the same rules. Invalid or foreign-device ranges are +rejected before recording commands. Kernel bounds and shader capability +requirements remain the extension's responsibility. + +Dispatch and copy use the backend's existing compute command context, descriptor +pool, transfer waits, and memory barriers. They preserve ordering with native +fragments and invalidate cached converted operands after writes. Normal backend +synchronization completes queued work before host access or resource release. +Vulkan copy ranges must be equal in size, nonoverlapping, and aligned to four +bytes. There is no per-dispatch host wait or extension-owned queue. + +Scheduled Vulkan plans share submission counters and FLOP thresholds across +native fragments and extension callbacks. Native scratch is initialized once +per plan invocation, with its offsets retained across fragments. Fragment +boundaries do not force a +queue submission, and the final pending commands are submitted when the plan +returns. Extension dispatches and copies count toward the device's node limit. +Every callback still executes on every invocation; this batching neither +captures device replay nor retains plan state across scheduler rebuilds. +Callback errors propagate without retry. Evaluation callbacks, direct backend +calls, synchronous devices, serialized diagnostics, and timestamp profiling +retain fragment execution. + +`src/core/vulkan/rope.cpp` registers custom and builtin RoPE implementations, +plus `RMSNorm -> Mul -> RoPE` and `RMSNorm -> Mul -> RoPE -> Scale -> FP16 Copy` +fusion. These kernels retain the original strides, partial-rotation behavior, +128-channel subgroup reduction, and separate arithmetic rounding. Their SPIR-V +is generated from application shaders with the native RoPE compilation policy; +SPIR-V optimization is disabled because it changes the reduction results. +Safe graph-delay rules bring independent producers ahead of fusion candidates. +External intermediate consumers retain ordinary operators, and overlapping +outputs use plan workspace. `GGML_VK_DISABLE_FUSION` is honored. + +The corresponding Vulkan shader sources and native fusion rules have been +removed from ggml. Existing builtin RoPE graphs use the registered override; +unsupported cases retain scheduler fallback to another supported backend. + +Remaining model-specific CUDA paths and Vulkan INT8/ConvRot implementations still +need auditing or migration. This implementation does not yet externalize all model-specific kernels or add dynamic dtypes. diff --git a/ggml b/ggml index 8cc0d5a6..64edc3b7 160000 --- a/ggml +++ b/ggml @@ -1 +1 @@ -Subproject commit 8cc0d5a6ae9093754089aa795f53ac8296c867d4 +Subproject commit 64edc3b78462898ac6c7a0261945240ad2891706 diff --git a/src/core/ggml_extensions.cpp b/src/core/ggml_extensions.cpp index fc8004bf..8fe44279 100644 --- a/src/core/ggml_extensions.cpp +++ b/src/core/ggml_extensions.cpp @@ -9,6 +9,9 @@ #include #include +#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); @@ -168,11 +171,12 @@ 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; - if (!cpu && !cuda) { + 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; } static std::mutex mutex; @@ -181,6 +185,12 @@ void sd_ggml_register_extensions(ggml_backend_t backend) { 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); @@ -226,6 +236,9 @@ ggml_tensor* sd_ggml_rope_apply(ggml_context* ctx, ggml_backend_t backend, ggml_ 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") #endif ) { ggml_tensor* inputs[] = {x, pe}; diff --git a/src/core/vulkan/CMakeLists.txt b/src/core/vulkan/CMakeLists.txt new file mode 100644 index 00000000..2ae495e4 --- /dev/null +++ b/src/core/vulkan/CMakeLists.txt @@ -0,0 +1,27 @@ +find_package(Vulkan COMPONENTS glslc REQUIRED) +set(SD_VULKAN_GENERATED "${CMAKE_CURRENT_BINARY_DIR}/shaders") +file(MAKE_DIRECTORY "${SD_VULKAN_GENERATED}") +foreach(NAME rope_apply rms_rope_apply rms_rope_cast) + set(SHADER "${NAME}") + set(FLAGS) + if(NAME STREQUAL "rms_rope_cast") + set(SHADER rms_rope_apply) + set(FLAGS -DCAST_F16) + endif() + set(SOURCE "${CMAKE_CURRENT_SOURCE_DIR}/shaders/${SHADER}.comp") + set(SPV "${SD_VULKAN_GENERATED}/${NAME}.spv") + set(HEADER "${SD_VULKAN_GENERATED}/${NAME}.h") + # Keep native RoPE rounding: spirv-opt changes the subgroup reduction results. + add_custom_command(OUTPUT "${HEADER}" + COMMAND "${Vulkan_GLSLC_EXECUTABLE}" --target-env=vulkan1.2 ${FLAGS} "${SOURCE}" -o "${SPV}" + COMMAND "${CMAKE_COMMAND}" "-DINPUT=${SPV}" "-DOUTPUT=${HEADER}" "-DNAME=${NAME}" + -P "${PROJECT_SOURCE_DIR}/cmake/embed_spirv.cmake" + DEPENDS "${SOURCE}" "${PROJECT_SOURCE_DIR}/cmake/embed_spirv.cmake" + BYPRODUCTS "${SPV}" + VERBATIM) + list(APPEND SD_VULKAN_HEADERS "${HEADER}") +endforeach() +add_custom_target(sd-vulkan-shaders DEPENDS ${SD_VULKAN_HEADERS}) +add_dependencies(${SD_LIB} sd-vulkan-shaders) +target_sources(${SD_LIB} PRIVATE "${CMAKE_CURRENT_SOURCE_DIR}/rope.cpp") +target_include_directories(${SD_LIB} PRIVATE "${SD_VULKAN_GENERATED}") diff --git a/src/core/vulkan/rope.cpp b/src/core/vulkan/rope.cpp new file mode 100644 index 00000000..8a485139 --- /dev/null +++ b/src/core/vulkan/rope.cpp @@ -0,0 +1,225 @@ +#include "core/ggml_extensions.h" +#include "ggml-ext-vulkan.h" +#include "rms_rope_apply.h" +#include "rms_rope_cast.h" +#include "rope_apply.h" + +#include +#include +#include + +namespace sd_vulkan_rope { + struct Parameters { + uint32_t half_dim, rot_half, heads, tokens, pairs; + uint32_t x_stride[4], pe_stride[4]; + uint32_t x_offset, pe_offset, dst_offset, interleaved; + }; + struct FusedParameters { + Parameters rope; + float eps; + uint32_t w_offset; + float output_scale; + }; + static_assert(sizeof(Parameters) == 68 && sizeof(FusedParameters) == 80, "RoPE push constant layout mismatch"); + static const ggml_ext_vulkan_api* api(ggml_backend_dev_t device) { + auto query = reinterpret_cast(ggml_backend_reg_get_proc_address( + ggml_backend_dev_backend_reg(device), "ggml_ext_vulkan_v1")); + return query ? query() : nullptr; + } + static float param(const ggml_tensor* node, int i) { + float value; + std::memcpy(&value, node->op_params + i, sizeof(value)); + return value; + } + static bool is_rope(const ggml_tensor* node) { + return node && (node->op == GGML_OP_ROPE_APPLY || ggml_ext_op_get(node) == sd_ggml_rope_op()); + } + 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(tensor->data) + offset) % alignment == 0; + } + static bool bounded(const ggml_tensor* node, const ggml_ext_vulkan_device_info& info) { + if (!node || !aligned(node, 4) || ggml_nbytes(node) > info.max_buffer_range - (info.buffer_alignment - 1)) + return false; + for (int i = 0; i < 4; ++i) + if (node->nb[i] % 4 || node->nb[i] / 4 > UINT32_MAX || node->ne[i] > UINT32_MAX) + return false; + return true; + } + static bool supports(ggml_backend_dev_t device, const ggml_tensor* node, void*) { + const auto* service = api(device); + ggml_ext_vulkan_device_info info; + if (!service || !service->device_info(device, &info) || info.max_group_size[0] < 256 || info.max_group_invocations < 256) + return false; + const auto* x = node->src[0]; + const auto* pe = node->src[1]; + return x && pe && x->type == GGML_TYPE_F32 && pe->type == GGML_TYPE_F32 && node->type == GGML_TYPE_F32 && + x->ne[0] % 2 == 0 && 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] && node->ne[0] == x->ne[0] && node->ne[1] == x->ne[2] && + node->ne[2] == x->ne[1] * x->ne[3] && node->ne[3] == 1 && ggml_is_contiguous(node) && + bounded(x, info) && bounded(pe, info) && bounded(node, info) && ggml_nelements(node) / 2 <= UINT32_MAX && + (uint64_t(ggml_nelements(node) / 2) + 255) / 256 <= uint64_t(info.max_group_count[0]) * info.max_group_count[1]; + } + static bool fused_supports(ggml_backend_dev_t device, const ggml_tensor* norm, const ggml_tensor* mul, const ggml_tensor* rope) { + ggml_ext_vulkan_device_info info; + if (!api(device)->device_info(device, &info) || info.fusion_disabled || !info.subgroup_shuffle || !info.full_subgroups || + info.subgroup_size != 32 || !is_rope(rope) || !supports(device, rope, nullptr) || + norm->op != GGML_OP_RMS_NORM || mul->op != GGML_OP_MUL || mul->src[0] != norm || rope->src[0] != mul) + return false; + const auto* x = norm->src[0]; + const auto* weight = mul->src[1]; + return x && weight && x->type == GGML_TYPE_F32 && x->ne[0] == 128 && x->nb[0] == 4 && + norm->type == GGML_TYPE_F32 && mul->type == GGML_TYPE_F32 && ggml_are_same_shape(x, norm) && ggml_are_same_shape(norm, mul) && + weight->type == GGML_TYPE_F32 && weight->ne[0] == 128 && ggml_nelements(weight) == 128 && + ggml_is_contiguous(weight) && bounded(weight, info) && bounded(x, info) && (x->ne[1] + 3) / 4 <= info.max_group_count[0]; + } + static bool cast_supports(ggml_backend_dev_t device, const ggml_tensor* rope, const ggml_tensor* scale, const ggml_tensor* copy) { + ggml_ext_vulkan_device_info info; + api(device)->device_info(device, &info); + // Workspace writeback uses Vulkan's four-byte copy granularity. + return info.float16 && scale->op == GGML_OP_SCALE && scale->src[0] == rope && scale->type == GGML_TYPE_F32 && + param(scale, 1) == 0.f && copy->op == GGML_OP_CPY && copy->src[0] == scale && copy->type == GGML_TYPE_F16 && + ggml_is_contiguous(copy) && aligned(copy, 4) && ggml_are_same_shape(scale, copy); + } + static int match(ggml_backend_dev_t device, const ggml_cgraph* graph, int first, void*) { + auto* g = const_cast(graph); + if (first + 2 >= ggml_graph_n_nodes(g) || !fused_supports(device, ggml_graph_node(g, first), ggml_graph_node(g, first + 1), ggml_graph_node(g, first + 2))) + return 0; + return first + 4 < ggml_graph_n_nodes(g) && cast_supports(device, ggml_graph_node(g, first + 2), ggml_graph_node(g, first + 3), ggml_graph_node(g, first + 4)) ? 5 : 3; + } + static void optimize(ggml_backend_dev_t device, ggml_cgraph* graph, void*) { + for (int last = 2; last < ggml_graph_n_nodes(graph); ++last) { + const auto* output = ggml_graph_node(graph, last); + const auto* scale = output->op == GGML_OP_CPY ? output->src[0] : nullptr; + const auto* rope = scale && scale->op == GGML_OP_SCALE ? scale->src[0] : output; + if (!is_rope(rope)) + continue; + const auto* mul = rope->src[0]; + const auto* norm = mul->src[0]; + if (!norm || !fused_supports(device, norm, mul, rope)) + continue; + const int count = output == rope ? 3 : cast_supports(device, rope, scale, output) ? 5 + : 0; + if (!count) + continue; + const ggml_tensor* nodes[] = {norm, mul, rope, scale, output}; + int indices[5], next = count - 1; + for (int i = last; i >= std::max(0, last - 64) && next >= 0; --i) + if (ggml_graph_node(graph, i) == nodes[next]) + indices[next--] = i; + if (next >= 0) + continue; + int first = indices[0]; + for (int i = 1; i < count; ++i) { + if (!ggml_ext_graph_delay(graph, first, i, indices[i])) + break; + first = indices[i] - i; + } + } + } + static ggml_status prepare(ggml_backend_t, ggml_tensor* const* nodes, int count, void**, size_t* bytes, void*) { + *bytes = ggml_nbytes(nodes[count - 1]); + return GGML_STATUS_SUCCESS; + } + static void destroy(void*, void*) {} + static bool overlaps(const ggml_ext_vulkan_buffer_info& a, const ggml_ext_vulkan_buffer_info& b) { + return a.storage == b.storage && a.offset < b.offset + b.size && b.offset < a.offset + a.size; + } + static ggml_status execute(const ggml_ext_execution_context* execution, ggml_tensor* const* nodes, int count, void*, void*) { + auto device = ggml_backend_get_device(execution->backend); + const auto* service = api(device); + ggml_ext_vulkan_device_info info; + if (!service || !service->device_info(device, &info)) + return GGML_STATUS_FAILED; + const auto* rope = nodes[count == 1 ? 0 : 2]; + const auto* x = count == 1 ? rope->src[0] : nodes[0]->src[0]; + const auto* pe = rope->src[1]; + const auto* output = nodes[count - 1]; + ggml_ext_vulkan_binding bindings[] = {{x}, {pe}, {output}, {count == 1 ? nullptr : nodes[1]->src[1]}}; + ggml_ext_vulkan_buffer_info buffers[4]; + const int n_buffers = count == 1 ? 3 : 4; + for (int i = 0; i < n_buffers; ++i) + if (!service->buffer_info(execution->backend, bindings + i, buffers + i)) + return GGML_STATUS_FAILED; + bool alias = overlaps(buffers[2], buffers[0]) || overlaps(buffers[2], buffers[1]) || (count != 1 && overlaps(buffers[2], buffers[3])); + const auto destination = bindings[2]; + if (alias) { + bindings[2] = {nullptr, execution->workspace, 0, ggml_nbytes(output)}; + if (!service->buffer_info(execution->backend, bindings + 2, buffers + 2)) + return GGML_STATUS_FAILED; + } + FusedParameters params = {}; + auto& pc = params.rope; + pc.half_dim = uint32_t(x->ne[0] / 2); + pc.rot_half = uint32_t(pe->ne[2]); + pc.heads = uint32_t(x->ne[1]); + pc.tokens = uint32_t(x->ne[2]); + pc.pairs = uint32_t(ggml_nelements(x) / 2); + for (int i = 0; i < 4; ++i) { + pc.x_stride[i] = uint32_t(x->nb[i] / 4); + pc.pe_stride[i] = uint32_t(pe->nb[i] / 4); + } + pc.x_offset = buffers[0].shader_offset / 4; + pc.pe_offset = buffers[1].shader_offset / 4; + pc.dst_offset = buffers[2].shader_offset / uint32_t(ggml_type_size(output->type)); + std::memcpy(&pc.interleaved, ggml_ext_op_get(rope) ? ggml_ext_op_attributes(rope) : rope->op_params, 4); + uint32_t groups[3] = {1, 1, 1}; + ggml_ext_vulkan_pipeline pipeline = {}; + pipeline.name = "sd.rope_apply.vulkan"; + pipeline.spirv = sd_vulkan_shaders::rope_apply; + pipeline.spirv_size = sizeof(sd_vulkan_shaders::rope_apply); + pipeline.binding_count = n_buffers; + pipeline.push_constant_size = sizeof(Parameters); + if (count == 1) { + const uint32_t n_groups = uint32_t((uint64_t(pc.pairs) + 255) / 256); + groups[0] = std::min(n_groups, info.max_group_count[0]); + groups[1] = (n_groups + groups[0] - 1) / groups[0]; + } else { + pipeline.name = count == 5 ? "sd.rms_rope_cast.vulkan" : "sd.rms_rope_apply.vulkan"; + pipeline.spirv = count == 5 ? sd_vulkan_shaders::rms_rope_cast : sd_vulkan_shaders::rms_rope_apply; + pipeline.spirv_size = count == 5 ? sizeof(sd_vulkan_shaders::rms_rope_cast) : sizeof(sd_vulkan_shaders::rms_rope_apply); + pipeline.push_constant_size = sizeof(FusedParameters); + pipeline.full_subgroups = true; + pipeline.required_subgroup_size = info.subgroup_size_control ? 32 : 0; + params.eps = param(nodes[0], 0); + params.w_offset = buffers[3].shader_offset / 4; + params.output_scale = count == 5 ? param(nodes[3], 0) : 1.f; + groups[0] = (pc.heads + 3) / 4; + groups[1] = std::min(pc.tokens, info.max_group_count[1]); + groups[2] = std::min(uint32_t(x->ne[3]), info.max_group_count[2]); + } + auto status = service->dispatch(execution->backend, &pipeline, bindings, ¶ms, groups); + if (status == GGML_STATUS_SUCCESS && alias) + status = service->copy(execution->backend, bindings + 2, &destination); + return status; + } +} + +void sd_ggml_register_vulkan_rope(ggml_backend_dev_t device) { + ggml_ext_kernel_desc desc = {}; + desc.abi_version = GGML_EXT_ABI_VERSION; + desc.name = "sd.rope_apply.vulkan.v1"; + desc.op = GGML_OP_CUSTOM; + desc.custom_op = sd_ggml_rope_op(); + desc.priority = 100; + desc.supports = sd_vulkan_rope::supports; + desc.prepare = sd_vulkan_rope::prepare; + desc.execute = sd_vulkan_rope::execute; + desc.destroy = sd_vulkan_rope::destroy; + GGML_ASSERT(ggml_ext_register_kernel(device, &desc)); + desc.name = "sd.rope_apply.vulkan.override.v1"; + desc.op = GGML_OP_ROPE_APPLY; + desc.custom_op = nullptr; + GGML_ASSERT(ggml_ext_register_kernel(device, &desc)); + desc.name = "sd.rms_rope_apply.vulkan.v1"; + desc.op = GGML_OP_RMS_NORM; + desc.priority = 200; + desc.supports = nullptr; + desc.match = sd_vulkan_rope::match; + desc.optimize = sd_vulkan_rope::optimize; + GGML_ASSERT(ggml_ext_register_kernel(device, &desc)); +} diff --git a/src/core/vulkan/shaders/rms_rope_apply.comp b/src/core/vulkan/shaders/rms_rope_apply.comp new file mode 100644 index 00000000..5477c584 --- /dev/null +++ b/src/core/vulkan/shaders/rms_rope_apply.comp @@ -0,0 +1,97 @@ +#version 450 + +#extension GL_KHR_shader_subgroup_basic : require +#extension GL_KHR_shader_subgroup_shuffle : require +#extension GL_EXT_control_flow_attributes : require + +layout(local_size_x = 128, local_size_y = 1, local_size_z = 1) in; +layout(binding = 0) readonly buffer X { float data_x[]; }; +layout(binding = 1) readonly buffer PE { float data_pe[]; }; +#ifdef CAST_F16 +#extension GL_EXT_shader_16bit_storage : require +#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require +layout(binding = 2) writeonly buffer D { float16_t data_d[]; }; +#else +layout(binding = 2) writeonly buffer D { float data_d[]; }; +#endif +layout(binding = 3) readonly buffer W { float data_w[]; }; + +layout(push_constant) uniform parameter { + uint half_dim; + uint rot_half; + uint heads; + uint tokens; + uint pairs; + uint x_stride[4]; + uint pe_stride[4]; + uint x_offset; + uint pe_offset; + uint dst_offset; + uint interleaved; + float eps; + uint w_offset; + float output_scale; +} pc; + +float shuffle_channel(vec4 values, uint channel) { + const uint lane = channel & 31u; + const vec4 shuffled = vec4(subgroupShuffle(values.x, lane), subgroupShuffle(values.y, lane), + subgroupShuffle(values.z, lane), subgroupShuffle(values.w, lane)); + return shuffled[channel / 32]; +} + +void main() { + const uint head = gl_WorkGroupID.x * 4 + gl_SubgroupID; + if (head >= pc.heads) { + return; + } + const uint lane = gl_SubgroupInvocationID; + const uint batches = pc.pairs / (64 * pc.heads * pc.tokens); + for (uint batch = gl_WorkGroupID.z; batch < batches; batch += gl_NumWorkGroups.z) { + for (uint token = gl_WorkGroupID.y; token < pc.tokens; token += gl_NumWorkGroups.y) { + const uint base = pc.x_offset + head * pc.x_stride[1] + token * pc.x_stride[2] + batch * pc.x_stride[3]; + vec4 x = vec4(data_x[base + lane], data_x[base + lane + 32], + data_x[base + lane + 64], data_x[base + lane + 96]); + precise vec4 squares = x * x; + precise float sum = (squares.x + squares.z) + (squares.y + squares.w); + [[unroll]] for (uint stride = 16; stride > 0; stride >>= 1) { + sum += subgroupShuffle(sum, lane ^ stride); + } + sum = subgroupShuffle(sum, 0); + const float scale = inversesqrt(sum / 128.0 + pc.eps); + precise vec4 weighted = scale * x; + [[unroll]] for (uint i = 0; i < 4; ++i) { + weighted[i] *= data_w[pc.w_offset + lane + i * 32]; + } + [[unroll]] for (uint i = 0; i < 4; ++i) { + const uint col = lane + i * 32; + precise float value = weighted[i]; + if (pc.interleaved != 0) { + const float a = subgroupShuffle(value, lane & ~1u); + const float b = subgroupShuffle(value, lane | 1u); + if (col / 2 < pc.rot_half) { + const uint matrix = pc.pe_offset + (col / 2) * pc.pe_stride[2] + token * pc.pe_stride[3] + (col & 1u) * pc.pe_stride[1]; + value = a * data_pe[matrix] + b * data_pe[matrix + pc.pe_stride[0]]; + } + } else { + const uint pair = col % pc.rot_half; + const float a = shuffle_channel(weighted, pair); + const float b = shuffle_channel(weighted, pair + pc.rot_half); + if (col < 2 * pc.rot_half) { + const uint matrix = pc.pe_offset + pair * pc.pe_stride[2] + token * pc.pe_stride[3]; + const float c = data_pe[matrix]; + const float s = data_pe[matrix + pc.pe_stride[1]]; + value = col < pc.rot_half ? a * c - b * s : a * s + b * c; + } + } +#ifdef CAST_F16 + value *= pc.output_scale; + value += 0.0; + data_d[pc.dst_offset + ((batch * pc.heads + head) * pc.tokens + token) * 128 + col] = float16_t(value); +#else + data_d[pc.dst_offset + ((batch * pc.heads + head) * pc.tokens + token) * 128 + col] = value; +#endif + } + } + } +} diff --git a/src/core/vulkan/shaders/rope_apply.comp b/src/core/vulkan/shaders/rope_apply.comp new file mode 100644 index 00000000..33d9c67d --- /dev/null +++ b/src/core/vulkan/shaders/rope_apply.comp @@ -0,0 +1,59 @@ +#version 450 + +layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in; + +layout(binding = 0) readonly buffer X { float data_x[]; }; +layout(binding = 1) readonly buffer PE { float data_pe[]; }; +layout(binding = 2) writeonly buffer D { float data_d[]; }; + +layout(push_constant) uniform parameter { + uint half_dim; + uint rot_half; + uint heads; + uint tokens; + uint pairs; + uint x_stride[4]; + uint pe_stride[4]; + uint x_offset; + uint pe_offset; + uint dst_offset; + uint interleaved; +} pc; + +void main() { + const uint index = gl_GlobalInvocationID.x + gl_NumWorkGroups.x*gl_WorkGroupSize.x*gl_GlobalInvocationID.y; + if (index >= pc.pairs) { + return; + } + const uint j = index % pc.half_dim; + const uint row = index / pc.half_dim; + const uint token = row % pc.tokens; + const uint head = (row / pc.tokens) % pc.heads; + const uint batch = row / (pc.tokens*pc.heads); + const uint i0 = pc.interleaved != 0 || j >= pc.rot_half ? 2*j : j; + const uint i1 = pc.interleaved != 0 || j >= pc.rot_half ? 2*j + 1 : j + pc.rot_half; + const uint x_row = pc.x_offset + head*pc.x_stride[1] + token*pc.x_stride[2] + batch*pc.x_stride[3]; + const uint out_row = pc.dst_offset + row*pc.half_dim*2; + const float a = data_x[x_row + i0*pc.x_stride[0]]; + const float b = data_x[x_row + i1*pc.x_stride[0]]; + if (j >= pc.rot_half) { + data_d[out_row + i0] = a; + data_d[out_row + i1] = b; + return; + } + const uint matrix = pc.pe_offset + j*pc.pe_stride[2] + token*pc.pe_stride[3]; + const float c = data_pe[matrix]; + const float s = data_pe[matrix + pc.pe_stride[1]]; + // Preserve the separate multiply/add rounding of the unfused graph. + precise float y0; + precise float y1; + if (pc.interleaved != 0) { + y0 = a*c + b*data_pe[matrix + pc.pe_stride[0]]; + y1 = a*s + b*data_pe[matrix + pc.pe_stride[0] + pc.pe_stride[1]]; + } else { + y0 = a*c - b*s; + y1 = a*s + b*c; + } + data_d[out_row + i0] = y0; + data_d[out_row + i1] = y1; +}