char-fork/llama.cpp

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

Jeff Bolzvulkan: combine duplicated fastdiv functions, rename the one optimizing small divs (#27526)c9ca51c1f

char/aliceai-t5-moe
3.4 KiB75 linesraw
1#extension GL_EXT_shader_16bit_storage : require
2#extension GL_EXT_control_flow_attributes : require
3
4#include "utils.glsl"
5
6layout (push_constant) uniform parameter
7{
8    uint ne;
9    uint ne00; uint ne01; uint ne02; uint ne03; uint nb00; uint nb01; uint nb02; uint nb03;
10    uint ne10; uint ne11; uint ne12; uint ne13; uint nb10; uint nb11; uint nb12; uint nb13;
11    uint misalign_offsets;
12    float param1; float param2; float param3; float param4;
13
14    // The three L values are packed as bytes to keep this layout under the 128B
15    // push constant limit while still leaving room for four float parameters.
16    uint ne0_012mp; uint ne0_01mp;  uint ne0_0mp;  uint ne0_Ls;
17    uint ne1_012mp; uint ne1_01mp;  uint ne1_0mp;  uint ne1_Ls;
18} p;
19
20layout (binding = 0) readonly buffer A {A_TYPE data_a[];};
21#if defined(A_TYPE_PACKED16)
22layout (binding = 0) readonly buffer A_PACKED16 {A_TYPE_PACKED16 data_a_packed16[];};
23#endif
24#if defined(A_TYPE_PACKED32)
25layout (binding = 0) readonly buffer A_PACKED32 {A_TYPE_PACKED32 data_a_packed32[];};
26#endif
27
28layout (binding = 1) writeonly buffer D {D_TYPE data_d[];};
29
30uint get_idx() {
31    return gl_GlobalInvocationID.z * 262144 + gl_GlobalInvocationID.y * 512 + gl_GlobalInvocationID.x;
32}
33
34uint get_aoffset() { return p.misalign_offsets >> 16; }
35uint get_doffset() { return p.misalign_offsets & 0xFFFF; }
36
37uint src0_idx(uint idx) {
38    const uint i03 = fastdiv(idx, p.ne0_012mp, fastdiv_L(p.ne0_Ls, 0));
39    const uint i03_offset = i03 * p.ne02*p.ne01*p.ne00;
40    const uint i02 = fastdiv(idx - i03_offset, p.ne0_01mp, fastdiv_L(p.ne0_Ls, 1));
41    const uint i02_offset = i02*p.ne01*p.ne00;
42    const uint i01 = fastdiv(idx - i03_offset - i02_offset, p.ne0_0mp, fastdiv_L(p.ne0_Ls, 2));
43    const uint i00 = idx - i03_offset - i02_offset - i01*p.ne00;
44    return i03*p.nb03 + i02*p.nb02 + i01*p.nb01 + i00*p.nb00;
45}
46
47uint dst_idx(uint idx) {
48    const uint i13 = fastdiv(idx, p.ne1_012mp, fastdiv_L(p.ne1_Ls, 0));
49    const uint i13_offset = i13 * p.ne12*p.ne11*p.ne10;
50    const uint i12 = fastdiv(idx - i13_offset, p.ne1_01mp, fastdiv_L(p.ne1_Ls, 1));
51    const uint i12_offset = i12*p.ne11*p.ne10;
52    const uint i11 = fastdiv(idx - i13_offset - i12_offset, p.ne1_0mp, fastdiv_L(p.ne1_Ls, 2));
53    const uint i10 = idx - i13_offset - i12_offset - i11*p.ne10;
54    return i13*p.nb13 + i12*p.nb12 + i11*p.nb11 + i10*p.nb10;
55}
56
57uint src0_idx_quant(uint idx, uint qk) {
58    const uint i03 = fastdiv(idx, p.ne0_012mp, fastdiv_L(p.ne0_Ls, 0));
59    const uint i03_offset = i03 * p.ne02*p.ne01*p.ne00;
60    const uint i02 = fastdiv(idx - i03_offset, p.ne0_01mp, fastdiv_L(p.ne0_Ls, 1));
61    const uint i02_offset = i02*p.ne01*p.ne00;
62    const uint i01 = fastdiv(idx - i03_offset - i02_offset, p.ne0_0mp, fastdiv_L(p.ne0_Ls, 2));
63    const uint i00 = idx - i03_offset - i02_offset - i01*p.ne00;
64    return i03*p.nb03 + i02*p.nb02 + i01*p.nb01 + (i00/qk)*p.nb00;
65}
66
67uint dst_idx_quant(uint idx, uint qk) {
68    const uint i13 = fastdiv(idx, p.ne1_012mp, fastdiv_L(p.ne1_Ls, 0));
69    const uint i13_offset = i13 * p.ne12*p.ne11*p.ne10;
70    const uint i12 = fastdiv(idx - i13_offset, p.ne1_01mp, fastdiv_L(p.ne1_Ls, 1));
71    const uint i12_offset = i12*p.ne11*p.ne10;
72    const uint i11 = fastdiv(idx - i13_offset - i12_offset, p.ne1_0mp, fastdiv_L(p.ne1_Ls, 2));
73    const uint i10 = idx - i13_offset - i12_offset - i11*p.ne10;
74    return i13*p.nb13 + i12*p.nb12 + i11*p.nb11 + (i10/qk)*p.nb10;
75}