char-fork/llama.cpp

git clone https://git.t4t.associates/char-fork/llama.cpp

Pascalvulkan: fuse snake activation (mul, sin, sqr, mul, add) (#22855)47c0eda9d

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
1.4 KiB49 linesraw
1#version 450
2
3#include "types.glsl"
4
5// Fused snake activation: y = x + sin(b * x)^2 * c
6//   data_a [ne0, ne1]   per element activation x    (A_TYPE)
7//   data_b [1,   ne1]   per channel multiplier      (float)
8//   data_c [1,   ne1]   per channel inverse scale   (float, precomputed as 1 / freq)
9//   data_d [ne0, ne1]   output y                    (D_TYPE)
10layout (binding = 0) readonly  buffer A {A_TYPE data_a[];};
11layout (binding = 1) readonly  buffer B {float  data_b[];};
12layout (binding = 2) readonly  buffer C {float  data_c[];};
13layout (binding = 3) writeonly buffer D {D_TYPE data_d[];};
14
15layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in;
16
17layout (push_constant) uniform parameter {
18    uint32_t ne0;
19    uint32_t ne1;
20} p;
21
22// Load A_TYPE to float
23float load_val(uint32_t idx) {
24#if defined(DATA_A_BF16)
25    return bf16_to_fp32(uint32_t(data_a[idx]));
26#else
27    return float(data_a[idx]);
28#endif
29}
30
31// Store float as D_TYPE
32void store_val(uint32_t idx, float v) {
33#if defined(DATA_D_BF16)
34    data_d[idx] = D_TYPE(fp32_to_bf16(v));
35#else
36    data_d[idx] = D_TYPE(v);
37#endif
38}
39
40void main() {
41    const uint32_t i0 = gl_GlobalInvocationID.x;
42    const uint32_t i1 = gl_GlobalInvocationID.y;
43    if (i0 >= p.ne0 || i1 >= p.ne1) return;
44
45    const uint32_t idx = i0 + i1 * p.ne0;
46    const float xi = load_val(idx);
47    const float s  = sin(data_b[i1] * xi);
48    store_val(idx, xi + s * s * data_c[i1]);
49}