perf(vulkan): accelerate MOSS normalization and SiTU staging (#2273)

Co-authored-by: mirek190 <mirek190@gmail.com>

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}")

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) = {

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}")

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));
+}

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;

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);
+}

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