From 93e39d7018b34aefb313ab161532bc3fb52ca028 Mon Sep 17 00:00:00 2001 From: leejet Date: Thu, 5 Nov 2026 02:47:44 +0800 Subject: [PATCH] perf(vulkan): accelerate MOSS normalization and SiTU staging Co-authored-by: mirek190 --- CMakeLists.txt | 2 +- cmake/backend_extensions.h.in | 1 + ggml | 2 +- src/core/layer_split_partition.cpp | 2 +- src/core/vulkan/CMakeLists.txt | 3 +- src/core/vulkan/rms_affine.cpp | 192 ++++++++++++++++++ src/core/vulkan/rope.cpp | 32 +-- src/core/vulkan/shaders/modulate.comp | 19 +- src/core/vulkan/shaders/modulate_convrot.comp | 4 +- src/core/vulkan/shaders/rms_modulate.comp | 136 +++++++++++++ src/core/vulkan/shaders/rms_rope_apply.comp | 7 +- 11 files changed, 375 insertions(+), 25 deletions(-) create mode 100644 src/core/vulkan/rms_affine.cpp create mode 100644 src/core/vulkan/shaders/rms_modulate.comp diff --git a/CMakeLists.txt b/CMakeLists.txt index 048f273de..8c5866d49 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -378,7 +378,7 @@ if (NOT SD_USE_UPSTREAM_GGML AND (SD_CUDA OR SD_HIPBLAS OR SD_MUSA)) endif() if (SD_VULKAN AND NOT SD_USE_UPSTREAM_GGML) add_subdirectory(src/core/vulkan) - list(APPEND SD_BACKEND_REGISTRATIONS sd_ggml_register_vulkan_rope sd_ggml_register_vulkan_i8 sd_ggml_register_vulkan_modulate) + list(APPEND SD_BACKEND_REGISTRATIONS sd_ggml_register_vulkan_rope sd_ggml_register_vulkan_i8 sd_ggml_register_vulkan_modulate sd_ggml_register_vulkan_rms_affine) endif() if(SD_BACKEND_REGISTRATIONS) string(REPLACE ";" ",\n " SD_BACKEND_REGISTRATIONS "${SD_BACKEND_REGISTRATIONS}") diff --git a/cmake/backend_extensions.h.in b/cmake/backend_extensions.h.in index 6b42ea117..2dc1e0f6c 100644 --- a/cmake/backend_extensions.h.in +++ b/cmake/backend_extensions.h.in @@ -5,6 +5,7 @@ void sd_ggml_register_vulkan_rope(ggml_backend_dev_t device); void sd_ggml_register_vulkan_i8(ggml_backend_dev_t device); +void sd_ggml_register_vulkan_rms_affine(ggml_backend_dev_t device); void sd_ggml_register_vulkan_modulate(ggml_backend_dev_t device); static void (*const sd_backend_extensions[])(ggml_backend_dev_t) = { diff --git a/ggml b/ggml index ccb88f810..135ce1e0e 160000 --- a/ggml +++ b/ggml @@ -1 +1 @@ -Subproject commit ccb88f810be05bccc9621e706354611ccc2462c7 +Subproject commit 135ce1e0e3f2b0693d3de93d4f6b5fa0a1947820 diff --git a/src/core/layer_split_partition.cpp b/src/core/layer_split_partition.cpp index 6ef24357a..ff8aaeb60 100644 --- a/src/core/layer_split_partition.cpp +++ b/src/core/layer_split_partition.cpp @@ -206,7 +206,7 @@ namespace sd { for (size_t seg_idx = 0; seg_idx < plan.segments.size(); seg_idx++) { ggml_backend_t backend = backend_by_segment[seg_idx]; const auto& segment = plan.segments[seg_idx]; - auto assign_node = [&](int node_index) { + auto assign_node = [&](int node_index) { if (node_index < 0 || node_index >= n_nodes) { return; } diff --git a/src/core/vulkan/CMakeLists.txt b/src/core/vulkan/CMakeLists.txt index 42ee931a2..55acf0bb9 100644 --- a/src/core/vulkan/CMakeLists.txt +++ b/src/core/vulkan/CMakeLists.txt @@ -49,6 +49,7 @@ endfunction() # and is disabled for KHR cooperative matrix shaders. sd_vulkan_shader(rope_apply rope_apply ON) sd_vulkan_shader(modulate modulate ON -O) +sd_vulkan_shader(rms_modulate rms_modulate ON -O) sd_vulkan_shader(rms_rope_apply rms_rope_apply ON) sd_vulkan_shader(rms_rope_cast rms_rope_apply ON -DCAST_F16) sd_vulkan_shader(quantize_i8_convrot quantize_i8_convrot ${SD_VK_DOT} -O) @@ -66,5 +67,5 @@ sd_vulkan_shader(mul_mat_i8_swiglu_split_cm2 mul_mat_i8_tensorwise_cm2 ${SD_VK_C sd_vulkan_shader(mul_mat_i8_residual_cm2 mul_mat_i8_tensorwise_cm2 ${SD_VK_CM2} -O -DRESIDUAL=1) 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" "${CMAKE_CURRENT_SOURCE_DIR}/int8.cpp" "${CMAKE_CURRENT_SOURCE_DIR}/modulate.cpp") +target_sources(${SD_LIB} PRIVATE "${CMAKE_CURRENT_SOURCE_DIR}/rope.cpp" "${CMAKE_CURRENT_SOURCE_DIR}/int8.cpp" "${CMAKE_CURRENT_SOURCE_DIR}/modulate.cpp" "${CMAKE_CURRENT_SOURCE_DIR}/rms_affine.cpp") target_include_directories(${SD_LIB} PRIVATE "${SD_VULKAN_GENERATED}") diff --git a/src/core/vulkan/rms_affine.cpp b/src/core/vulkan/rms_affine.cpp new file mode 100644 index 000000000..2eee7c7e3 --- /dev/null +++ b/src/core/vulkan/rms_affine.cpp @@ -0,0 +1,192 @@ +#include +#include +#include +#include "core/ggml_extensions.h" +#include "ggml-ext-vulkan.h" +#include "rms_modulate.h" + +namespace sd_vulkan_rms_affine { + struct Parameters { + uint32_t ne[4], stride[4], offsets[5]; + float eps; + uint32_t ne10; + }; + static_assert(sizeof(Parameters) == 60, "RMS affine push constant layout"); + + 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 bool bounded(const ggml_tensor* t, const ggml_ext_vulkan_device_info& info) { + if (!t || t->type != GGML_TYPE_F32 || t->nb[0] != 4 || ggml_nbytes(t) > info.max_buffer_range - (info.buffer_alignment - 1)) + return false; + size_t offset = 0; + const auto* owner = t; + while (!owner->data && owner->view_src) { + offset += owner->view_offs; + owner = owner->view_src; + } + if ((reinterpret_cast(owner->data) + offset) % 4) + return false; + for (int d = 0; d < 4; ++d) + if (t->ne[d] > UINT32_MAX || t->nb[d] % 4 || t->nb[d] / 4 > UINT32_MAX) + return false; + return true; + } + static const ggml_tensor* other(const ggml_tensor* node, const ggml_tensor* input) { + return node->src[0] == input ? node->src[1] : node->src[1] == input ? node->src[0] + : nullptr; + } + struct Operation { + const ggml_tensor* x; + const ggml_tensor* weight; + const ggml_tensor* scale; + const ggml_tensor* shift; + const ggml_tensor* output; + }; + static Operation describe(ggml_tensor* const* nodes, int count) { + Operation op = {nodes[0]->src[0], other(nodes[1], nodes[0]), nullptr, nullptr, nodes[count - 1]}; + if (count == 3) { + op.scale = nodes[2]->src[1]; + op.shift = nodes[2]->src[2]; + } else if (count == 5) { + op.scale = other(nodes[2], nodes[1]); + op.shift = other(nodes[4], nodes[3]); + } + return op; + } + static int match(ggml_backend_dev_t device, const ggml_cgraph* graph, int first, void*) { + const auto* service = api(device); + ggml_ext_vulkan_device_info info; + auto* g = const_cast(graph); + if (!service || !service->device_info(device, &info) || info.fusion_disabled || first + 1 >= ggml_graph_n_nodes(g)) + return 0; + ggml_tensor* nodes[5] = {}; + const int available = std::min(5, ggml_graph_n_nodes(g) - first); + for (int i = 0; i < available; ++i) + nodes[i] = ggml_graph_node(g, first + i); + if (nodes[0]->op != GGML_OP_RMS_NORM || nodes[1]->op != GGML_OP_MUL) + return 0; + const auto* x = nodes[0]->src[0]; + if (!bounded(x, info) || x->ne[0] != 3072 || available < 3 || !ggml_is_contiguous(x) || + x->ne[1] <= 1 || x->ne[1] > info.max_group_count[0] || x->ne[2] != 1 || x->ne[3] != 1 || + info.max_group_size[0] < 512 || info.max_group_invocations < 512) + return 0; + int count; + if (sd_ggml_modulate_describe(nodes[2]) && nodes[2]->src[0] == nodes[1]) { + count = 3; + } else if (available == 5 && nodes[2]->op == GGML_OP_MUL && nodes[3]->op == GGML_OP_ADD && + other(nodes[3], nodes[1]) == nodes[2] && nodes[4]->op == GGML_OP_ADD) { + count = 5; + } else { + return 0; + } + const auto op = describe(nodes, count); + if (!op.weight || !op.scale || (count == 5 && !op.shift)) + return 0; + for (const auto* parameter : {op.weight, op.scale, op.shift}) { + if (!parameter) + continue; + if (!bounded(parameter, info) || !ggml_is_contiguous(parameter) || parameter->ne[0] != x->ne[0] || ggml_nelements(parameter) != x->ne[0]) + return 0; + for (int i = 0; i < count; ++i) + if (parameter == nodes[i]) + return 0; + } + for (int i = 0; i < count; ++i) + if (!bounded(nodes[i], info) || !ggml_are_same_shape(nodes[i], x) || !ggml_is_contiguous(nodes[i])) + return 0; + return count; + } + static void optimize(ggml_backend_dev_t device, ggml_cgraph* graph, void*) { + const auto* service = api(device); + ggml_ext_vulkan_device_info info; + if (!service || !service->device_info(device, &info) || info.fusion_disabled) + return; + for (int i = 2; i < ggml_graph_n_nodes(graph); ++i) { + const auto* node = ggml_graph_node(graph, i); + if (!sd_ggml_modulate_describe(node) && node->op != GGML_OP_MUL) + continue; + const auto* weighted = node->src[0]; + if (!weighted || weighted->op != GGML_OP_MUL || !weighted->src[0] || + weighted->src[0]->op != GGML_OP_RMS_NORM || weighted->ne[0] != 3072) + continue; + for (int j = i - 2; j >= std::max(0, i - 32); --j) { + if (ggml_graph_node(graph, j) == weighted->src[0] && ggml_graph_node(graph, j + 1) == weighted) { + if (j + 2 < i) + ggml_ext_graph_delay(graph, j, 2, i); + break; + } + } + } + } + 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 ggml_status execute(const ggml_ext_execution_context* e, ggml_tensor* const* nodes, int count, void*, void*) { + auto device = ggml_backend_get_device(e->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 op = describe(nodes, count); + ggml_ext_vulkan_binding bindings[] = {{op.x}, {op.weight}, {op.output}, {op.scale}, {op.shift ? op.shift : op.scale}}; + ggml_ext_vulkan_buffer_info buffers[5]; + const int n = 5; + for (int i = 0; i < n; ++i) + if (!service->buffer_info(e->backend, bindings + i, buffers + i)) + return GGML_STATUS_FAILED; + bool alias = false; + for (int i = 0; i < n; ++i) + if (i != 2) + alias |= buffers[i].storage == buffers[2].storage && buffers[i].offset < buffers[2].offset + buffers[2].size && + buffers[2].offset < buffers[i].offset + buffers[i].size; + const auto destination = bindings[2]; + if (alias) { + bindings[2] = {nullptr, e->workspace, 0, ggml_nbytes(op.output)}; + if (!service->buffer_info(e->backend, bindings + 2, buffers + 2)) + return GGML_STATUS_FAILED; + } + Parameters p = {}; + p.ne10 = uint32_t(op.weight->ne[0]); + for (int i = 0; i < 4; ++i) { + p.ne[i] = uint32_t(op.x->ne[i]); + p.stride[i] = uint32_t(op.x->nb[i] / 4); + } + for (int i = 0; i < n; ++i) + p.offsets[i] = buffers[i].shader_offset / 4; + std::memcpy(&p.eps, nodes[0]->op_params, sizeof(p.eps)); + const uint32_t has_shift = op.shift != nullptr; + ggml_ext_vulkan_pipeline pipeline = {}; + pipeline.name = "sd.rms_modulate.vulkan"; + pipeline.spirv = sd_vulkan_shaders::rms_modulate; + pipeline.spirv_size = sizeof(sd_vulkan_shaders::rms_modulate); + pipeline.binding_count = n; + pipeline.push_constant_size = sizeof(p); + pipeline.constants = &has_shift; + pipeline.constant_count = 1; + uint32_t groups[] = {p.ne[1], std::min(p.ne[2], info.max_group_count[1]), std::min(p.ne[3], info.max_group_count[2])}; + auto status = service->dispatch(e->backend, &pipeline, bindings, &p, groups); + if (status == GGML_STATUS_SUCCESS && alias) + status = service->copy(e->backend, bindings + 2, &destination); + return status; + } +} +void sd_ggml_register_vulkan_rms_affine(ggml_backend_dev_t device) { + if (!sd_vulkan_rms_affine::api(device)) + return; + ggml_ext_kernel_desc desc = {}; + desc.abi_version = GGML_EXT_ABI_VERSION; + desc.name = "sd.rms_affine.vulkan.v1"; + desc.op = GGML_OP_RMS_NORM; + desc.priority = 150; + desc.match = sd_vulkan_rms_affine::match; + desc.optimize = sd_vulkan_rms_affine::optimize; + desc.prepare = sd_vulkan_rms_affine::prepare; + desc.execute = sd_vulkan_rms_affine::execute; + desc.destroy = sd_vulkan_rms_affine::destroy; + GGML_ASSERT(ggml_ext_register_kernel(device, &desc)); +} diff --git a/src/core/vulkan/rope.cpp b/src/core/vulkan/rope.cpp index 63b976e79..194603e13 100644 --- a/src/core/vulkan/rope.cpp +++ b/src/core/vulkan/rope.cpp @@ -72,7 +72,7 @@ namespace sd_vulkan_rope { 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 && + return x && weight && weight != norm && 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]; @@ -81,32 +81,35 @@ namespace sd_vulkan_rope { 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); + return info.float16 && (!scale || (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 ? scale : rope) && copy->type == GGML_TYPE_F16 && + ggml_is_contiguous(copy) && aligned(copy, 4) && ggml_are_same_shape(rope, 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; + if (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))) + return 5; + return first + 3 < ggml_graph_n_nodes(g) && cast_supports(device, ggml_graph_node(g, first + 2), nullptr, ggml_graph_node(g, first + 3)) ? 4 : 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; + const auto* source = output->op == GGML_OP_CPY ? output->src[0] : nullptr; + const auto* scale = source && source->op == GGML_OP_SCALE ? source : nullptr; + const auto* rope = source ? (scale ? scale->src[0] : source) : 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 + const int count = output == rope ? 3 : cast_supports(device, rope, scale, output) ? (scale ? 5 : 4) : 0; if (!count) continue; - const ggml_tensor* nodes[] = {norm, mul, rope, scale, output}; + const ggml_tensor* nodes[] = {norm, mul, rope, scale ? scale : output, 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]) @@ -170,6 +173,7 @@ namespace sd_vulkan_rope { pc.dst_offset = buffers[2].shader_offset / uint32_t(ggml_type_size(output->type)); std::memcpy(&pc.interleaved, ggml_ext_op_attributes(rope), 4); uint32_t groups[3] = {1, 1, 1}; + const uint32_t apply_output_scale = count == 5; ggml_ext_vulkan_pipeline pipeline = {}; pipeline.name = "sd.rope_apply.vulkan"; pipeline.spirv = sd_vulkan_shaders::rope_apply; @@ -181,9 +185,13 @@ namespace sd_vulkan_rope { 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.name = count > 3 ? "sd.rms_rope_cast.vulkan" : "sd.rms_rope_apply.vulkan"; + pipeline.spirv = count > 3 ? sd_vulkan_shaders::rms_rope_cast : sd_vulkan_shaders::rms_rope_apply; + pipeline.spirv_size = count > 3 ? sizeof(sd_vulkan_shaders::rms_rope_cast) : sizeof(sd_vulkan_shaders::rms_rope_apply); + if (count > 3) { + pipeline.constants = &apply_output_scale; + pipeline.constant_count = 1; + } pipeline.push_constant_size = sizeof(FusedParameters); pipeline.full_subgroups = true; pipeline.required_subgroup_size = info.subgroup_size_control ? 32 : 0; diff --git a/src/core/vulkan/shaders/modulate.comp b/src/core/vulkan/shaders/modulate.comp index df07626d2..81f910d07 100644 --- a/src/core/vulkan/shaders/modulate.comp +++ b/src/core/vulkan/shaders/modulate.comp @@ -1,16 +1,25 @@ #version 450 layout(local_size_x = 256) in; -layout(binding = 0) readonly buffer X { float x[]; }; -layout(binding = 1) readonly buffer S { float scale[]; }; -layout(binding = 2) readonly buffer H { float shift[]; }; -layout(binding = 3) writeonly buffer D { float dst[]; }; +layout(binding = 0) readonly buffer X { + float x[]; +}; +layout(binding = 1) readonly buffer S { + float scale[]; +}; +layout(binding = 2) readonly buffer H { + float shift[]; +}; +layout(binding = 3) writeonly buffer D { + float dst[]; +}; layout(push_constant) uniform Parameters { uint ne[4], scale_ne[4], shift_ne[4]; uint x_stride[4], scale_stride[4], shift_stride[4]; uint offsets[4]; uint count, has_shift; -} p; +} +p; void main() { uint group = gl_WorkGroupID.y * gl_NumWorkGroups.x + gl_WorkGroupID.x; diff --git a/src/core/vulkan/shaders/modulate_convrot.comp b/src/core/vulkan/shaders/modulate_convrot.comp index 3cac110c9..4e85b1d0b 100644 --- a/src/core/vulkan/shaders/modulate_convrot.comp +++ b/src/core/vulkan/shaders/modulate_convrot.comp @@ -2,8 +2,8 @@ #extension GL_EXT_control_flow_attributes : require -layout(constant_id = 0) const uint K = 4096; -layout(constant_id = 1) const bool AFFINE = false; +layout(constant_id = 0) const uint K = 4096; +layout(constant_id = 1) const bool AFFINE = false; layout(constant_id = 2) const bool HAS_SHIFT = true; layout(local_size_x = 512, local_size_y = 1, local_size_z = 1) in; layout(binding = 0) readonly buffer X { diff --git a/src/core/vulkan/shaders/rms_modulate.comp b/src/core/vulkan/shaders/rms_modulate.comp new file mode 100644 index 000000000..53d68bf2b --- /dev/null +++ b/src/core/vulkan/shaders/rms_modulate.comp @@ -0,0 +1,136 @@ +#version 450 +#extension GL_EXT_control_flow_attributes : enable +#define BLOCK_SIZE 512 +#define FLOAT_TYPE float +#define D_TYPE float +layout(local_size_x = BLOCK_SIZE) in; +layout(binding = 0) readonly buffer X { + float data_a[]; +}; +layout(binding = 1) readonly buffer W { + float data_b[]; +}; +layout(binding = 2) writeonly buffer D { + float data_d[]; +}; +layout(binding = 3) readonly buffer S { + float data_c[]; +}; +layout(binding = 4) readonly buffer H { + float data_e[]; +}; +layout(constant_id = 0) const bool has_shift = true; +layout(constant_id = 1) const bool do_multiply = true; +layout(push_constant) uniform Parameters { + uint ne[4], stride[4], offsets[5]; + float eps; + uint ne10; +} +p; +shared float sumsh[BLOCK_SIZE]; +shared float modulation_input[BLOCK_SIZE]; +uint get_aoffset() { + return p.offsets[0]; +} +uint get_boffset() { + return p.offsets[1]; +} +uint get_doffset() { + return p.offsets[2]; +} +uint src1_idx(uint a, uint b, uint c, uint d) { + return 0; +} +uint fastmod(uint a, uint b) { + return a % b; +} +float modulate(float value, uint col) { + // Keep precise modulation from changing the native normalization's FMA contraction. + modulation_input[gl_LocalInvocationIndex] = value; + memoryBarrierShared(); + float weighted = modulation_input[gl_LocalInvocationIndex]; + precise float product = weighted * data_c[p.offsets[3] + col]; + precise float result = weighted + product; + if (has_shift) + result += data_e[p.offsets[4] + col]; + return result; +} +void rms_norm(uint num_iters) { + const uint ncols = p.ne[0]; + const uint nrows = gl_NumWorkGroups.x; + const uint nchannels = p.ne[2]; + const uint nsamples = p.ne[3]; + + const uint row = gl_WorkGroupID.x; + const uint tid = gl_LocalInvocationID.x; + + const uint stride_row = p.stride[1]; + const uint stride_channel = p.stride[2]; + const uint stride_sample = p.stride[3]; + + // grid.y/z are clamped to the device workgroup limit, iterate over the excess channels/samples + for (uint samp = gl_WorkGroupID.z; samp < nsamples; samp += gl_NumWorkGroups.z) { + for (uint channel = gl_WorkGroupID.y; channel < nchannels; channel += gl_NumWorkGroups.y) { + barrier(); + + uint a_offset = samp * stride_sample + channel * stride_channel + row * stride_row + get_aoffset(); + uint b_offset = src1_idx(0, row, channel, samp) + get_boffset(); + uint d_offset = ((samp * nchannels + channel) * nrows + row) * ncols + get_doffset(); + FLOAT_TYPE sum = FLOAT_TYPE(0.0f); // partial sum for thread in warp + + [[unroll]] for (uint col = tid, idx = 0; idx < num_iters; col += BLOCK_SIZE, ++idx) { + FLOAT_TYPE xi = FLOAT_TYPE(0); + if (col < ncols) { + xi = FLOAT_TYPE(data_a[a_offset + col]); + } + sum += xi * xi; + } + + sumsh[tid] = sum; + // sum up partial sums and write back result + barrier(); + [[unroll]] for (int s = BLOCK_SIZE / 2; s > 0; s >>= 1) { + if (tid < s) { + sum += sumsh[tid + s]; + sumsh[tid] = sum; + } + barrier(); + } + sum = sumsh[0]; + + const FLOAT_TYPE mean = sum / FLOAT_TYPE(ncols); + const FLOAT_TYPE scale = inversesqrt(mean + FLOAT_TYPE(p.eps)); + + if (do_multiply) { + if (ncols > p.ne10) { + [[unroll]] for (uint col = tid, idx = 0; idx < num_iters; col += BLOCK_SIZE, ++idx) { + if (col >= ncols) { + continue; + } + FLOAT_TYPE value = scale * FLOAT_TYPE(data_a[a_offset + col]) * FLOAT_TYPE(data_b[b_offset + fastmod(col, p.ne10)]); + data_d[d_offset + col] = modulate(value, col); + } + } else { + [[unroll]] for (uint col = tid, idx = 0; idx < num_iters; col += BLOCK_SIZE, ++idx) { + if (col >= ncols) { + continue; + } + FLOAT_TYPE value = scale * FLOAT_TYPE(data_a[a_offset + col]) * FLOAT_TYPE(data_b[b_offset + col]); + data_d[d_offset + col] = modulate(value, col); + } + } + } else { + [[unroll]] for (uint col = tid, idx = 0; idx < num_iters; col += BLOCK_SIZE, ++idx) { + if (col >= ncols) { + continue; + } + data_d[d_offset + col] = D_TYPE(scale * FLOAT_TYPE(data_a[a_offset + col])); + } + } + } + } +} + +void main() { + rms_norm(8); +} diff --git a/src/core/vulkan/shaders/rms_rope_apply.comp b/src/core/vulkan/shaders/rms_rope_apply.comp index 7d08f7414..ef2f581c7 100644 --- a/src/core/vulkan/shaders/rms_rope_apply.comp +++ b/src/core/vulkan/shaders/rms_rope_apply.comp @@ -12,6 +12,7 @@ layout(binding = 1) readonly buffer PE { float data_pe[]; }; #ifdef CAST_F16 +layout(constant_id = 0) const bool apply_output_scale = false; #extension GL_EXT_shader_16bit_storage : require #extension GL_EXT_shader_explicit_arithmetic_types_float16 : require layout(binding = 2) writeonly buffer D { @@ -96,8 +97,10 @@ void main() { } } #ifdef CAST_F16 - value *= pc.output_scale; - value += 0.0; + if (apply_output_scale) { + 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;