char-fork/llama.cpp

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

Michael Lamotheggml: add f16 out_prod support for CPU and out_prod op for Vulkan (#23997)f955e394b

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
1.8 KiB59 linesraw
1#version 450
2
3#extension GL_EXT_shader_16bit_storage : require
4
5layout (push_constant) uniform parameter
6{
7    uint ne;
8    uint ne00; uint ne01; uint ne02; uint ne03; uint nb00; uint nb01; uint nb02; uint nb03;
9    uint ne10; uint ne11; uint ne12; uint ne13; uint nb10; uint nb11; uint nb12; uint nb13;
10    uint ne20; uint ne21; uint ne22; uint ne23; uint nb20; uint nb21; uint nb22; uint nb23;
11    uint misalign_offsets;
12    float param1; float param2; int param3;
13} p;
14
15layout (binding = 0) readonly buffer A {float data_a[];};
16layout (binding = 1) readonly buffer B {float data_b[];};
17layout (binding = 2) writeonly buffer D {float data_d[];};
18
19uint get_idx() {
20    return gl_GlobalInvocationID.z * 262144 + gl_GlobalInvocationID.y * 512 + gl_GlobalInvocationID.x;
21}
22
23uint get_aoffset() { return p.misalign_offsets >> 16; }
24uint get_boffset() { return (p.misalign_offsets >> 8) & 0xFF; }
25uint get_doffset() { return p.misalign_offsets & 0xFF; }
26
27layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in;
28
29void main() {
30    uint idx = get_idx();
31    if (idx >= p.ne) {
32        return;
33    }
34
35    uint tmp = idx;
36    uint i0 = tmp % p.ne20; tmp /= p.ne20;
37    uint i1 = tmp % p.ne21; tmp /= p.ne21;
38    uint i2 = tmp % p.ne22; tmp /= p.ne22;
39    uint i3 = tmp;
40
41    uint a_i0 = i0 % p.ne00;
42    uint a_i2 = i2 / (p.ne22 / p.ne02);
43    uint a_i3 = i3 / (p.ne23 / p.ne03);
44
45    uint b_i0 = i1 % p.ne10;
46    uint b_i2 = i2;
47    uint b_i3 = i3;
48
49    float sum = 0.0f;
50    uint K = p.ne01;
51    for (uint k = 0; k < K; k++) {
52        uint aoff = get_aoffset() + a_i3*p.nb03 + a_i2*p.nb02 + k*p.nb01 + a_i0*p.nb00;
53        uint boff = get_boffset() + b_i3*p.nb13 + b_i2*p.nb12 + k*p.nb11 + b_i0*p.nb10;
54        sum += data_a[aoff] * data_b[boff];
55    }
56
57    uint doff = get_doffset() + i3*p.nb23 + i2*p.nb22 + i1*p.nb21 + i0*p.nb20;
58    data_d[doff] = sum;
59}