char-fork/llama.cpp
git clone https://git.t4t.associates/char-fork/llama.cpp
1f30ac0ce
char/aliceai-t5-moe
1#version 450 2 3#extension GL_EXT_shader_16bit_storage : require 4#extension GL_EXT_nonuniform_qualifier : enable 5#extension GL_EXT_control_flow_attributes : require 6#if ADD_RMS 7#extension GL_KHR_shader_subgroup_arithmetic : enable 8#extension GL_KHR_shader_subgroup_basic : enable 9#endif 10 11#include "types.glsl" 12#include "utils.glsl" 13 14layout (push_constant) uniform parameter2 15{ 16 // shape for dst 17 uint ne20; uint ne21; uint ne22; uint ne23; 18 19 // strides for srcs+dst 20 uint nb[12][4]; 21 22 uint rms_partials; 23} p; 24 25// No readonly/writeonly decorations. Workaround for MoltenVK Bug, see https://github.com/ggml-org/llama.cpp/issues/15498 26layout (binding = 0) buffer A0 {A_TYPE data_a[];} a0; 27layout (binding = 1) buffer A1 {A_TYPE data_a[];} a1; 28layout (binding = 2) buffer A2 {A_TYPE data_a[];} a2; 29layout (binding = 3) buffer A3 {A_TYPE data_a[];} a3; 30layout (binding = 4) buffer A4 {A_TYPE data_a[];} a4; 31layout (binding = 5) buffer A5 {A_TYPE data_a[];} a5; 32layout (binding = 6) buffer A6 {A_TYPE data_a[];} a6; 33layout (binding = 7) buffer A7 {A_TYPE data_a[];} a7; 34layout (binding = 8) buffer A8 {A_TYPE data_a[];} a8; 35layout (binding = 9) buffer A9 {A_TYPE data_a[];} a9; 36layout (binding = 10) buffer A10 {A_TYPE data_a[];} a10; 37layout (binding = 11) buffer A11 {A_TYPE data_a[];} a11; 38layout (binding = 0) buffer D0 {D_TYPE data_d[];} d0; 39layout (binding = 1) buffer D1 {D_TYPE data_d[];} d1; 40layout (binding = 2) buffer D2 {D_TYPE data_d[];} d2; 41layout (binding = 3) buffer D3 {D_TYPE data_d[];} d3; 42layout (binding = 4) buffer D4 {D_TYPE data_d[];} d4; 43layout (binding = 5) buffer D5 {D_TYPE data_d[];} d5; 44layout (binding = 6) buffer D6 {D_TYPE data_d[];} d6; 45layout (binding = 7) buffer D7 {D_TYPE data_d[];} d7; 46layout (binding = 8) buffer D8 {D_TYPE data_d[];} d8; 47layout (binding = 9) buffer D9 {D_TYPE data_d[];} d9; 48layout (binding = 10) buffer D10 {D_TYPE data_d[];} d10; 49layout (binding = 11) buffer D11 {D_TYPE data_d[];} d11; 50layout (binding = 0, std430) buffer PartialBuf0 {float partial_sums[];} partials0; 51layout (binding = 1, std430) buffer PartialBuf1 {float partial_sums[];} partials1; 52layout (binding = 2, std430) buffer PartialBuf2 {float partial_sums[];} partials2; 53layout (binding = 3, std430) buffer PartialBuf3 {float partial_sums[];} partials3; 54layout (binding = 4, std430) buffer PartialBuf4 {float partial_sums[];} partials4; 55layout (binding = 5, std430) buffer PartialBuf5 {float partial_sums[];} partials5; 56layout (binding = 6, std430) buffer PartialBuf6 {float partial_sums[];} partials6; 57layout (binding = 7, std430) buffer PartialBuf7 {float partial_sums[];} partials7; 58layout (binding = 8, std430) buffer PartialBuf8 {float partial_sums[];} partials8; 59layout (binding = 9, std430) buffer PartialBuf9 {float partial_sums[];} partials9; 60layout (binding = 10, std430) buffer PartialBuf10 {float partial_sums[];} partials10; 61layout (binding = 11, std430) buffer PartialBuf11 {float partial_sums[];} partials11; 62 63layout(constant_id = 0) const uint num_srcs = 2; 64 65FLOAT_TYPE load_a(uint b, uint i) { 66 switch (b) { 67 case 0: return FLOAT_TYPE(a0.data_a[i]); 68 case 1: return FLOAT_TYPE(a1.data_a[i]); 69 case 2: return FLOAT_TYPE(a2.data_a[i]); 70 case 3: return FLOAT_TYPE(a3.data_a[i]); 71 case 4: return FLOAT_TYPE(a4.data_a[i]); 72 case 5: return FLOAT_TYPE(a5.data_a[i]); 73 case 6: return FLOAT_TYPE(a6.data_a[i]); 74 case 7: return FLOAT_TYPE(a7.data_a[i]); 75 case 8: return FLOAT_TYPE(a8.data_a[i]); 76 case 9: return FLOAT_TYPE(a9.data_a[i]); 77 case 10: return FLOAT_TYPE(a10.data_a[i]); 78 case 11: return FLOAT_TYPE(a11.data_a[i]); 79 default: return FLOAT_TYPE(0); 80 } 81} 82 83void store_d(uint b, uint i, FLOAT_TYPE v) { 84 switch (b) { 85 case 0: d0.data_d[i] = D_TYPE(v); break; 86 case 1: d1.data_d[i] = D_TYPE(v); break; 87 case 2: d2.data_d[i] = D_TYPE(v); break; 88 case 3: d3.data_d[i] = D_TYPE(v); break; 89 case 4: d4.data_d[i] = D_TYPE(v); break; 90 case 5: d5.data_d[i] = D_TYPE(v); break; 91 case 6: d6.data_d[i] = D_TYPE(v); break; 92 case 7: d7.data_d[i] = D_TYPE(v); break; 93 case 8: d8.data_d[i] = D_TYPE(v); break; 94 case 9: d9.data_d[i] = D_TYPE(v); break; 95 case 10: d10.data_d[i] = D_TYPE(v); break; 96 case 11: d11.data_d[i] = D_TYPE(v); break; 97 default: break; 98 } 99} 100 101void store_partial(uint b, uint i, float v) { 102 switch (b) { 103 case 0: partials0.partial_sums[i] = v; break; 104 case 1: partials1.partial_sums[i] = v; break; 105 case 2: partials2.partial_sums[i] = v; break; 106 case 3: partials3.partial_sums[i] = v; break; 107 case 4: partials4.partial_sums[i] = v; break; 108 case 5: partials5.partial_sums[i] = v; break; 109 case 6: partials6.partial_sums[i] = v; break; 110 case 7: partials7.partial_sums[i] = v; break; 111 case 8: partials8.partial_sums[i] = v; break; 112 case 9: partials9.partial_sums[i] = v; break; 113 case 10: partials10.partial_sums[i] = v; break; 114 case 11: partials11.partial_sums[i] = v; break; 115 default: break; 116 } 117} 118 119uint src_idx(uint s, uint i00, uint i01, uint i02, uint i03) { 120 return i03*p.nb[s][3] + i02*p.nb[s][2] + i01*p.nb[s][1] + i00*p.nb[s][0]; 121} 122 123uint dst_idx(uint i00, uint i01, uint i02, uint i03) { 124 uint nb20 = p.nb[num_srcs][0]; 125 uint nb21 = p.nb[num_srcs][1]; 126 uint nb22 = p.nb[num_srcs][2]; 127 uint nb23 = p.nb[num_srcs][3]; 128 return i03*nb23 + i02*nb22 + i01*nb21 + i00*nb20; 129} 130 131uint get_idx() { 132 return gl_GlobalInvocationID.z * 262144 + gl_GlobalInvocationID.y * 512 + gl_GlobalInvocationID.x; 133} 134 135const uint num_threads = 256; 136 137layout(local_size_x = num_threads, local_size_y = 1, local_size_z = 1) in; 138 139#if ADD_RMS 140// XXX TODO this could be sized based on number of subgroups, but that't not considered a constant 141shared FLOAT_TYPE sumsh[num_threads]; 142#endif 143 144void main() { 145 uint idx = get_idx(); 146 uint orig_idx = idx; 147 148 uint ne = p.ne20 * p.ne21 * p.ne22 * p.ne23; 149 150 // num_threads * num_iter must equal 512, to match the wg_denoms and get_idx calculation 151 const uint num_iter = 2; 152 153 FLOAT_TYPE sum_sq = 0; 154 155 [[unroll]] for (uint i = 0; i < num_iter; ++i) { 156 if (idx >= ne) { 157 continue; 158 } 159 uint i00, i01, i02, i03; 160 get_indices(idx, i00, i01, i02, i03, p.ne20, p.ne21, p.ne22, p.ne23); 161 162 FLOAT_TYPE sum = FLOAT_TYPE(0); 163 [[unroll]] for (uint s = 0; s < num_srcs; ++s) { 164 sum += load_a(s, src_idx(s, i00, i01, i02, i03)); 165 } 166 sum_sq += sum*sum; 167 store_d(num_srcs, dst_idx(i00, i01, i02, i03), sum); 168 169 idx += num_threads; 170 } 171 172#if ADD_RMS 173 if (p.rms_partials != 0) { 174 // reduce the sum within each subgroup, then across subgroups 175 const uint NumSubgroups = num_threads / gl_SubgroupSize; 176 sum_sq = subgroupAdd(sum_sq); 177 if (gl_SubgroupInvocationID == 0) { 178 sumsh[gl_SubgroupID] = sum_sq; 179 } 180 barrier(); 181 [[unroll]] for (uint s = NumSubgroups / 2; s > 0; s >>= 1) { 182 if (gl_SubgroupID < s && gl_SubgroupInvocationID == 0) { 183 sum_sq += sumsh[gl_SubgroupID + s]; 184 sumsh[gl_SubgroupID] = sum_sq; 185 } 186 barrier(); 187 } 188 189 if (gl_SubgroupID == 0 && gl_SubgroupInvocationID == 0) { 190 store_partial(num_srcs + 1, orig_idx / (num_iter * num_threads), sum_sq); 191 } 192 } 193#endif 194}