Commit 58bb3c1 for stable-diffusion.cpp
commit 58bb3c1b1cd4749e94e39f537e9adfce295dcaee
Author: leejet <leejet714@gmail.com>
Date: Thu Nov 5 03:01:31 2026 +0800
perf(vulkan): accelerate MOSS normalization and SiTU staging (#2273)
Co-authored-by: mirek190 <mirek190@gmail.com>
diff --git a/CMakeLists.txt b/CMakeLists.txt
index 048f273..8c5866d 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 6b42ea1..2dc1e0f 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 ccb88f8..135ce1e 160000
--- a/ggml
+++ b/ggml
@@ -1 +1 @@
-Subproject commit ccb88f810be05bccc9621e706354611ccc2462c7
+Subproject commit 135ce1e0e3f2b0693d3de93d4f6b5fa0a1947820
diff --git a/src/core/vulkan/CMakeLists.txt b/src/core/vulkan/CMakeLists.txt
index 42ee931..55acf0b 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 0000000..2eee7c7
--- /dev/null
+++ b/src/core/vulkan/rms_affine.cpp
@@ -0,0 +1,192 @@
+#include <algorithm>
+#include <climits>
+#include <cstring>
+#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_ext_vulkan_api_fn>(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<uintptr_t>(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<ggml_cgraph*>(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 63b976e..194603e 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<ggml_cgraph*>(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/rms_modulate.comp b/src/core/vulkan/shaders/rms_modulate.comp
new file mode 100644
index 0000000..53d68bf
--- /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 7d08f74..ef2f581 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;