char-fork/llama.cpp
git clone https://git.t4t.associates/char-fork/llama.cpp
c1258830b
char/aliceai-t5-moe
1#ifdef TYPE_I32 2#define TYPE i32 3#else 4#define TYPE f32 5#endif 6 7#ifndef INPLACE 8@group(0) @binding(0) 9var<storage, read_write> src0: array<TYPE>; 10#define SRC1_BINDING 1 11#else 12#define SRC1_BINDING 0 13#endif 14 15#define DST_BINDING SRC1_BINDING + 1 16#define PARAMS_BINDING SRC1_BINDING + 2 17 18@group(0) @binding(SRC1_BINDING) 19var<storage, read_write> src1: array<TYPE>; 20 21@group(0) @binding(DST_BINDING) 22var<storage, read_write> dst: array<TYPE>; 23 24struct Params { 25 ne: u32, 26 offset_src0: u32, 27 offset_src1: u32, 28 offset_view: u32, 29 30 stride_src10: u32, 31 stride_src11: u32, 32 stride_src12: u32, 33 stride_src13: u32, 34 35 stride_dst10: u32, 36 stride_dst11: u32, 37 stride_dst12: u32, 38 stride_dst13: u32, 39 40 src1_ne0: u32, 41 src1_ne1: u32, 42 src1_ne2: u32, 43 src1_ne3: u32, 44}; 45 46@group(0) @binding(PARAMS_BINDING) 47var<uniform> params: Params; 48 49fn decode_src1_coords(idx: u32) -> vec4<u32> { 50 var i = idx; 51 let plane = params.src1_ne2 * params.src1_ne1 * params.src1_ne0; 52 let i3 = i / plane; 53 i = i % plane; 54 let row = params.src1_ne1 * params.src1_ne0; 55 let i2 = i / row; 56 i = i % row; 57 let i1 = i / params.src1_ne0; 58 let i0 = i % params.src1_ne0; 59 return vec4<u32>(i0, i1, i2, i3); 60} 61 62fn decode_view_coords(rel: u32) -> vec4<u32> { 63 let i3 = rel / params.stride_dst13; 64 let rem3 = rel % params.stride_dst13; 65 let i2 = rem3 / params.stride_dst12; 66 let rem2 = rem3 % params.stride_dst12; 67 let i1 = rem2 / params.stride_dst11; 68 let i0 = rem2 % params.stride_dst11; 69 return vec4<u32>(i0, i1, i2, i3); 70} 71 72fn view_rel_from_coords(coords: vec4<u32>) -> u32 { 73 return coords.x * params.stride_dst10 + coords.y * params.stride_dst11 + 74 coords.z * params.stride_dst12 + coords.w * params.stride_dst13; 75} 76 77fn src1_idx_from_coords(coords: vec4<u32>) -> u32 { 78 return coords.x * params.stride_src10 + coords.y * params.stride_src11 + 79 coords.z * params.stride_src12 + coords.w * params.stride_src13; 80} 81 82fn in_set_view(rel: u32, coords: vec4<u32>) -> bool { 83 return view_rel_from_coords(coords) == rel; 84} 85 86@compute @workgroup_size(WG_SIZE) 87fn main(@builtin(global_invocation_id) gid: vec3<u32>) { 88 if (gid.x >= params.ne) { 89 return; 90 } 91 92#ifdef INPLACE 93 let coords = decode_src1_coords(gid.x); 94 95 let src1_idx = params.offset_src1 + src1_idx_from_coords(coords); 96 let dst_idx = params.offset_view + view_rel_from_coords(coords); 97 98 dst[dst_idx] = src1[src1_idx]; 99#else 100 let rel = select(params.ne, gid.x - params.offset_view, gid.x >= params.offset_view); 101 let coords = decode_view_coords(rel); 102 103 if (rel < params.stride_dst13 * params.src1_ne3 && in_set_view(rel, coords)) { 104 dst[gid.x] = src1[params.offset_src1 + src1_idx_from_coords(coords)]; 105 } else { 106 dst[gid.x] = src0[params.offset_src0 + gid.x]; 107 } 108#endif 109}