char-fork/llama.cpp

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

Masashi Yoshimuraggml-webgpu: Enables running gpt-oss-20b (#22906)927dada6c

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
1.7 KiB64 linesraw
1struct Params {
2    offset_src0: u32,
3    offset_src1: u32,
4    offset_ids: u32,
5    offset_dst: u32,
6
7    nb01: u32,
8    nb02: u32,
9    nb11: u32,
10    nb20: u32,
11    nb21: u32,
12
13    ne0: u32,
14    ne1: u32,
15    ne2: u32,
16};
17
18@group(0) @binding(0) var<storage, read_write> src0: array<f32>; // [n_embd, n_experts_used, n_token]
19@group(0) @binding(1) var<storage, read_write> src1: array<f32>; // [n_embd, n_experts]
20@group(0) @binding(2) var<storage, read_write> ids:  array<i32>; // [n_experts_used, n_token]
21
22#ifdef INPLACE
23
24@group(0) @binding(3)
25var<uniform> params: Params;
26
27#else
28
29@group(0) @binding(3)
30var<storage, read_write> dst: array<f32>;
31
32@group(0) @binding(4)
33var<uniform> params: Params;
34
35#endif
36
37@compute @workgroup_size(WG_SIZE)
38fn main(@builtin(workgroup_id) wg_id: vec3<u32>,
39        @builtin(num_workgroups) num_wg: vec3<u32>,
40        @builtin(local_invocation_id) local_id: vec3<u32>) {
41
42    let wg_linear = wg_id.x + wg_id.y * num_wg.x;
43
44    if (wg_linear < params.ne1 * params.ne2) {
45        let thread_id = local_id.x;
46        let i2 = wg_linear / params.ne1;
47        let i1 = wg_linear % params.ne1;
48
49        let i11 = u32(ids[params.offset_ids + i1 * params.nb20 + i2 * params.nb21]);
50
51        let src0_row = params.offset_src0 + i1 * params.nb01 + i2 * params.nb02;
52        let src1_row = params.offset_src1 + i11 * params.nb11;
53        let dst_row = params.offset_dst + i1 * params.ne0 + i2 * (params.ne0 * params.ne1);
54
55        for (var i = thread_id;i < params.ne0; i += WG_SIZE) {
56#ifdef INPLACE
57            src0[src0_row + i] = src0[src0_row + i] + src1[src1_row + i];
58#else
59            dst[dst_row + i] = src0[src0_row + i] + src1[src1_row + i];
60#endif
61        }
62    }
63
64}