char-fork/llama.cpp

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

Nikhil Jain[ggml-webgpu] Implement 2D workgroups for scale, binary, and unary ops (#24044)3b3da01dc

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
1.5 KiB65 linesraw
1#ifdef INPLACE
2@group(0) @binding(1)
3var<uniform> params: Params;
4
5fn store_scale(val: f32, offset: u32) {
6    src[offset] = val;
7}
8#else
9@group(0) @binding(1)
10var<storage, read_write> dst: array<f32>;
11
12@group(0) @binding(2)
13var<uniform> params: Params;
14
15fn store_scale(val: f32, offset: u32) {
16    dst[offset] = val;
17}
18#endif
19
20struct Params {
21    offset_src: u32,
22    offset_dst: u32,
23
24    // Strides (in elements)
25    stride_src1: u32,
26    stride_src2: u32,
27    stride_src3: u32,
28
29    stride_dst1: u32,
30    stride_dst2: u32,
31    stride_dst3: u32,
32
33    ne: u32,
34    ne0: u32,
35    ne1: u32,
36    ne2: u32,
37
38    scale: f32,
39    bias: f32
40};
41
42@group(0) @binding(0)
43var<storage, read_write> src: array<f32>;
44
45@compute @workgroup_size(WG_SIZE)
46fn main(
47    @builtin(global_invocation_id) gid: vec3<u32>,
48    @builtin(num_workgroups)       num_wg:  vec3<u32>) {
49    let threads_per_group = u32(WG_SIZE);
50    var i = gid.x + (num_wg.x * threads_per_group) * gid.y;
51    if (i >= params.ne) {
52        return;
53    }
54    let i3 = i / (params.ne2 * params.ne1 * params.ne0);
55    i = i % (params.ne2 * params.ne1 * params.ne0);
56    let i2 = i / (params.ne1 * params.ne0);
57    i = i % (params.ne1 * params.ne0);
58    let i1 = i / params.ne0;
59    let i0 = i % params.ne0;
60
61    let i_src = params.offset_src + i3 * params.stride_src3 + i2 * params.stride_src2 + i1 * params.stride_src1 + i0;
62    let i_dst = params.offset_dst + i3 * params.stride_dst3 + i2 * params.stride_dst2 + i1 * params.stride_dst1 + i0;
63
64    store_scale(src[i_src] * params.scale + params.bias, i_dst);
65}