char-fork/llama.cpp

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

Anand Patilvulkan: add POOL_1D op (#25431)876a43211

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
1.4 KiB65 linesraw
1#version 450
2
3#include "types.glsl"
4
5#extension GL_EXT_shader_16bit_storage : require
6
7layout(push_constant) uniform parameter {
8    uint IL;
9    uint OL;
10    uint OC;
11    uint pelements;
12    uint op;
13    int k0;
14    int s0;
15    int p0;
16} p;
17
18#define BLOCK_SIZE 512
19#define FLT_MAX 3.402823466e+38F
20#define OP_POOL_MAX 0u
21#define OP_POOL_AVG 1u
22
23layout (local_size_x = BLOCK_SIZE, local_size_y = 1, local_size_z = 1) in;
24
25layout(binding = 0) readonly buffer X {A_TYPE data_a[];};
26layout(binding = 1) writeonly buffer D {D_TYPE data_d[];};
27
28void main() {
29    const uint idx = gl_GlobalInvocationID.x;
30    if (idx >= p.pelements) {
31        return;
32    }
33
34    const uint nc = idx / p.OL;
35    const uint cur_ol = idx % p.OL;
36
37    const int start = int(cur_ol) * p.s0 - p.p0;
38    const int bl = max(start, 0);
39    const int el = min(max(start + p.k0, 0), int(p.IL));
40
41    const int window_size = el - bl;
42    const float scale = window_size > 0 ? 1.0 / float(window_size) : 0.0;
43    float res;
44
45    if (p.op == OP_POOL_AVG) {
46        res = 0.0;
47    } else if (p.op == OP_POOL_MAX) {
48        res = -FLT_MAX;
49    } else {
50        return;
51    }
52
53    #pragma unroll
54    for (uint i = bl; i < el; i++) {
55        const float cur = D_TYPE(data_a[nc * p.IL + i]);
56
57        if (p.op == OP_POOL_AVG) {
58            res += cur * scale;
59        } else if (p.op == OP_POOL_MAX) {
60            res = max(res, cur);
61        }
62    }
63
64    data_d[nc * p.OL + cur_ol] = res;
65}