Optimize Krea Vulkan - fusion and misc. changes

ggml/src/ggml-vulkan/vulkan-shaders/mul_add.comp

@@ -0,0 +1,37 @@
+#version 450
+
+#extension GL_EXT_control_flow_attributes : require
+
+layout (push_constant) uniform parameter {
+    uint ne0_vec4;
+    uint nrows;
+} p;
+
+layout (binding = 0) readonly buffer Residual { vec4 data_residual[]; };
+layout (binding = 1) readonly buffer Value    { vec4 data_value[]; };
+layout (binding = 2) readonly buffer Gate     { vec4 data_gate[]; };
+layout (binding = 3) writeonly buffer D       { vec4 data_d[]; };
+
+const uint num_threads = 256;
+const uint num_iter = 2;
+
+layout(local_size_x = num_threads, local_size_y = 1, local_size_z = 1) in;
+
+void main() {
+    const uint row = gl_WorkGroupID.y;
+    uint channel = gl_WorkGroupID.x * num_threads * num_iter + gl_LocalInvocationID.x;
+
+    if (row >= p.nrows) {
+        return;
+    }
+
+    [[unroll]] for (uint i = 0; i < num_iter; ++i) {
+        if (channel < p.ne0_vec4) {
+            const uint idx = row * p.ne0_vec4 + channel;
+            precise vec4 product = data_value[idx] * data_gate[channel];
+            precise vec4 sum = data_residual[idx] + product;
+            data_d[idx] = sum;
+        }
+        channel += num_threads;
+    }
+}

ggml/src/ggml-vulkan/vulkan-shaders/quantize_q8_1.comp

@@ -24,10 +24,16 @@ layout(constant_id = 0) const uint GROUP_SIZE = 32;
 layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in;
 
 layout (binding = 0) readonly buffer A {vec4 data_a[];};
+#if defined(FUSED_SWIGLU) || defined(FUSED_SIGMOID_MUL)
+layout (binding = 1) readonly buffer B {vec4 data_b[];};
+#define D_BINDING 2
+#else
+#define D_BINDING 1
+#endif
 #ifndef QBLOCK_X4
-layout (binding = 1) writeonly buffer D {block_q8_1_packed32 data_b[];};
+layout (binding = D_BINDING) writeonly buffer D {block_q8_1_packed32 data_d[];};
 #else
-layout (binding = 1) writeonly buffer D {block_q8_1_x4 data_b[];};
+layout (binding = D_BINDING) writeonly buffer D {block_q8_1_x4 data_d[];};
 #endif
 
 #ifndef USE_SUBGROUPS
@@ -57,7 +63,20 @@ void quantize(const uint wgid) {
 
     const uint a_idx = ib * 8 + iqs;
 
-    vec4 vals = a_idx < p.ne / 4 ? data_a[a_idx] : vec4(0.0f);
+    vec4 vals = vec4(0.0f);
+    if (a_idx < p.ne / 4) {
+#ifdef FUSED_SWIGLU
+        precise vec4 gate = data_a[a_idx];
+        precise vec4 activated = gate / (vec4(1.0f) + exp(-gate));
+        vals = activated * data_b[a_idx];
+#elif defined(FUSED_SIGMOID_MUL)
+        precise vec4 gate = data_b[a_idx];
+        precise vec4 activated = vec4(1.0f) / (vec4(1.0f) + exp(-gate));
+        vals = data_a[a_idx] * activated;
+#else
+        vals = data_a[a_idx];
+#endif
+    }
     const vec4 abs_vals = abs(vals);
 
     // Find absolute max for each block
@@ -82,9 +101,9 @@ void quantize(const uint wgid) {
     vals = round(vals * d_inv);
 
 #ifndef QBLOCK_X4
-    data_b[ib].qs[iqs] = pack32(i8vec4(round(vals)));
+    data_d[ib].qs[iqs] = pack32(i8vec4(round(vals)));
 #else
-    data_b[ibx4_outer].qs[ibx4_inner * 8 + iqs] = pack32(i8vec4(round(vals)));
+    data_d[ibx4_outer].qs[ibx4_inner * 8 + iqs] = pack32(i8vec4(round(vals)));
 #endif
 
 #ifndef USE_SUBGROUPS
@@ -111,9 +130,9 @@ void quantize(const uint wgid) {
 #endif
 
 #ifndef QBLOCK_X4
-        data_b[ib].ds = f16vec2(vec2(d, sum * d));
+        data_d[ib].ds = f16vec2(vec2(d, sum * d));
 #else
-        data_b[ibx4_outer].ds[ibx4_inner] = f16vec2(vec2(d, sum * d));
+        data_d[ibx4_outer].ds[ibx4_inner] = f16vec2(vec2(d, sum * d));
 #endif
     }
 }

ggml/src/ggml-vulkan/vulkan-shaders/rms_norm.comp

@@ -33,6 +33,9 @@ layout (binding = 6) readonly buffer R_I {uvec2 rope_data_i[];}; // indices for
 #define BLOCK_SIZE 512
 
 layout (constant_id = 1) const bool do_multiply = false;
+#if RMS_NORM_ROPE_FUSION
+layout (constant_id = 2) const bool mrope_pack = false;
+#endif
 
 layout(local_size_x = BLOCK_SIZE, local_size_y = 1, local_size_z = 1) in;
 
@@ -91,29 +94,49 @@ void rms_norm(uint num_iters) {
                 if (col >= ncols) {
                     continue;
                 }
-                data_d[d_offset + col] = D_TYPE(scale * FLOAT_TYPE(data_a[a_offset + col]) * FLOAT_TYPE(data_b[b_offset + fastmod(col, p.ne10)]));
+                uint dst_col = col;
+#if RMS_NORM_ROPE_FUSION
+                if (mrope_pack) {
+                    dst_col = col / 2 + (col % 2) * (ncols / 2);
+                }
+#endif
+                data_d[d_offset + dst_col] = D_TYPE(scale * FLOAT_TYPE(data_a[a_offset + col]) * FLOAT_TYPE(data_b[b_offset + fastmod(col, p.ne10)]));
             }
         } 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]) * FLOAT_TYPE(data_b[b_offset + col]));
+                uint dst_col = col;
+#if RMS_NORM_ROPE_FUSION
+                if (mrope_pack) {
+                    dst_col = col / 2 + (col % 2) * (ncols / 2);
+                }
+#endif
+                data_d[d_offset + dst_col] = D_TYPE(scale * FLOAT_TYPE(data_a[a_offset + col]) * FLOAT_TYPE(data_b[b_offset + 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]));
+            uint dst_col = col;
+#if RMS_NORM_ROPE_FUSION
+            if (mrope_pack) {
+                dst_col = col / 2 + (col % 2) * (ncols / 2);
+            }
+#endif
+            data_d[d_offset + dst_col] = D_TYPE(scale * FLOAT_TYPE(data_a[a_offset + col]));
         }
     }
 #if RMS_NORM_ROPE_FUSION
     barrier();
     rope_params rp = p.rope;
     for (uint t = 2*tid; t < ncols; t += 2*BLOCK_SIZE) {
-        if (rp.rope_mode == GGML_ROPE_TYPE_NEOX) {
+        if (mrope_pack) {
+            rope_multi(t, row, channel, samp, rp);
+        } else if (rp.rope_mode == GGML_ROPE_TYPE_NEOX) {
             rope_neox(t, row, channel, samp, rp);
         } else if (rp.rope_mode == GGML_ROPE_TYPE_NORMAL) {
             rope_norm(t, row, channel, samp, rp);

ggml/src/ggml-vulkan/vulkan-shaders/rms_norm_channel_last.comp

@@ -0,0 +1,50 @@
+#version 450
+
+#extension GL_EXT_control_flow_attributes : enable
+
+layout(constant_id = 0) const uint BLOCK_SIZE = 512;
+
+layout(push_constant) uniform parameter {
+    uint width;
+    uint height;
+    uint depth;
+    uint channels;
+    float eps;
+} p;
+
+layout(binding = 0) readonly buffer A    { float data_a[];    };
+layout(binding = 1) readonly buffer NORM { float data_norm[]; };
+layout(binding = 2) writeonly buffer D   { float data_d[];    };
+
+layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in;
+
+shared float sumsh[BLOCK_SIZE];
+
+void main() {
+    const uint tid = gl_LocalInvocationID.x;
+    const uint plane = p.width * p.height * p.depth;
+    const uint spatial = (gl_WorkGroupID.z * p.height + gl_WorkGroupID.y) * p.width + gl_WorkGroupID.x;
+
+    precise float sum = 0.0f;
+    [[unroll]] for (uint channel = tid; channel < p.channels; channel += BLOCK_SIZE) {
+        const float value = data_a[spatial + channel * plane];
+        precise float square = value * value;
+        sum += square;
+    }
+
+    sumsh[tid] = sum;
+    barrier();
+    [[unroll]] for (uint s = BLOCK_SIZE / 2; s > 0; s >>= 1) {
+        if (tid < s) {
+            sum += sumsh[tid + s];
+            sumsh[tid] = sum;
+        }
+        barrier();
+    }
+
+    const float scale = inversesqrt(sumsh[0] / float(p.channels) + p.eps);
+    [[unroll]] for (uint channel = tid; channel < p.channels; channel += BLOCK_SIZE) {
+        const uint offset = spatial + channel * plane;
+        data_d[offset] = scale * data_a[offset] * data_norm[channel];
+    }
+}

ggml/src/ggml-vulkan/vulkan-shaders/rms_norm_modulate.comp

@@ -0,0 +1,93 @@
+#version 450
+
+#extension GL_EXT_control_flow_attributes : enable
+
+#define BLOCK_SIZE 512
+
+layout(push_constant) uniform parameter {
+    uint ne00;
+    uint ne01;
+    uint ne02;
+    uint ne03;
+    uint nb01;
+    uint nb02;
+    uint nb03;
+    float eps;
+} p;
+
+layout(binding = 0) readonly buffer A     { float data_a[];     };
+layout(binding = 1) readonly buffer NORM  { float data_norm[];  };
+layout(binding = 2) readonly buffer SCALE { float data_scale[]; };
+layout(binding = 3) readonly buffer SHIFT { float data_shift[]; };
+layout(binding = 4) writeonly buffer D    { float data_d[];     };
+
+layout(local_size_x = BLOCK_SIZE, local_size_y = 1, local_size_z = 1) in;
+
+shared float sumsh[BLOCK_SIZE];
+
+void rms_norm_modulate(uint num_iters) {
+    const uint row     = gl_WorkGroupID.x;
+    const uint channel = gl_WorkGroupID.y;
+    const uint samp    = gl_WorkGroupID.z;
+    const uint tid     = gl_LocalInvocationID.x;
+
+    const uint a_offset = samp*p.nb03 + channel*p.nb02 + row*p.nb01;
+    const uint d_offset = ((samp*p.ne02 + channel)*p.ne01 + row)*p.ne00;
+
+    float sum = 0.0f;
+    [[unroll]] for (uint col = tid, idx = 0; idx < num_iters; col += BLOCK_SIZE, ++idx) {
+        float xi = 0.0f;
+        if (col < p.ne00) {
+            xi = data_a[a_offset + col];
+        }
+        sum += xi * xi;
+    }
+
+    sumsh[tid] = sum;
+    barrier();
+    [[unroll]] for (int s = BLOCK_SIZE / 2; s > 0; s >>= 1) {
+        if (tid < s) {
+            sum += sumsh[tid + s];
+            sumsh[tid] = sum;
+        }
+        barrier();
+    }
+
+    const float inv_rms = inversesqrt(sumsh[0] / float(p.ne00) + p.eps);
+
+    [[unroll]] for (uint col = tid, idx = 0; idx < num_iters; col += BLOCK_SIZE, ++idx) {
+        if (col >= p.ne00) {
+            continue;
+        }
+
+        precise float normalized = inv_rms * data_a[a_offset + col] * data_norm[col];
+        precise float factor     = data_scale[col] * 1.0f + 1.0f;
+        precise float modulated  = normalized * factor;
+        data_d[d_offset + col]   = modulated + data_shift[col];
+    }
+}
+
+void main() {
+    uint num_blocks = (p.ne00 + BLOCK_SIZE - 1) / BLOCK_SIZE;
+    if (num_blocks > 32) {
+        rms_norm_modulate(num_blocks);
+    } else if (num_blocks > 16) {
+        rms_norm_modulate(32);
+    } else if (num_blocks > 12) {
+        rms_norm_modulate(16);
+    } else if (num_blocks > 10) {
+        rms_norm_modulate(12);
+    } else if (num_blocks > 8) {
+        rms_norm_modulate(10);
+    } else if (num_blocks > 4) {
+        rms_norm_modulate(8);
+    } else if (num_blocks == 4) {
+        rms_norm_modulate(4);
+    } else if (num_blocks == 3) {
+        rms_norm_modulate(3);
+    } else if (num_blocks == 2) {
+        rms_norm_modulate(2);
+    } else if (num_blocks == 1) {
+        rms_norm_modulate(1);
+    }
+}

ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp

@@ -805,6 +805,9 @@ void process_shaders() {
     string_to_spv("norm_f32", "norm.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "float"}}));
     string_to_spv("group_norm_f32", "group_norm.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "float"}}));
     string_to_spv("rms_norm_f32", "rms_norm.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}}));
+    string_to_spv("rms_norm_channel_last_f32", "rms_norm_channel_last.comp", {});
+    string_to_spv("rms_norm_modulate_f32", "rms_norm_modulate.comp", {});
+    string_to_spv("mul_add_f32", "mul_add.comp", {});
     string_to_spv("rms_norm_partials_f32", "rms_norm_partials.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}}));
     string_to_spv("rms_norm_mul_rope_f32_f32", "rms_norm.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}, {"ROPE_D_TYPE", "float"}, {"RMS_NORM_ROPE_FUSION", "1"}}));
     string_to_spv("rms_norm_mul_rope_f32_f16", "rms_norm.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}, {"ROPE_D_TYPE", "float16_t"}, {"RMS_NORM_ROPE_FUSION", "1"}}));
@@ -882,6 +885,10 @@ void process_shaders() {
 
     string_to_spv("quantize_q8_1_x4", "quantize_q8_1.comp", {{"QBLOCK_X4", "1"}});
     string_to_spv("quantize_q8_1_x4_subgroup", "quantize_q8_1.comp", {{"QBLOCK_X4", "1"}, {"USE_SUBGROUPS", "1"}});
+    string_to_spv("quantize_q8_1_x4_swiglu", "quantize_q8_1.comp", {{"QBLOCK_X4", "1"}, {"FUSED_SWIGLU", "1"}});
+    string_to_spv("quantize_q8_1_x4_swiglu_subgroup", "quantize_q8_1.comp", {{"QBLOCK_X4", "1"}, {"FUSED_SWIGLU", "1"}, {"USE_SUBGROUPS", "1"}});
+    string_to_spv("quantize_q8_1_x4_sigmoid_mul", "quantize_q8_1.comp", {{"QBLOCK_X4", "1"}, {"FUSED_SIGMOID_MUL", "1"}});
+    string_to_spv("quantize_q8_1_x4_sigmoid_mul_subgroup", "quantize_q8_1.comp", {{"QBLOCK_X4", "1"}, {"FUSED_SIGMOID_MUL", "1"}, {"USE_SUBGROUPS", "1"}});
 
     string_to_spv("mul_f32", "mul.comp", {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}, {"FLOAT_TYPE", "float"}});