char-fork/llama.cpp

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

Masashi Yoshimuraggml-webgpu : refactor several wgsl files and simplify flash_attn wgsl. (#26134)f401bb139

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
2.4 KiB92 linesraw
1struct Params {
2    ne: u32,
3
4    offset_src0: u32,
5    offset_src1: u32,
6    offset_dst: u32,
7
8    stride_src0_0: u32,
9    stride_src0_1: u32,
10    stride_src0_2: u32,
11    stride_src0_3: u32,
12
13    stride_src1_0: u32,
14    stride_src1_1: u32,
15    stride_src1_2: u32,
16    stride_src1_3: u32,
17
18    ne0: u32,
19    ne1: u32,
20    ne2: u32,
21
22    dim: u32,
23    src0_nedim: u32
24};
25
26#ifdef TYPE_F32
27#define DataType f32
28#endif
29#ifdef TYPE_I32
30#define DataType i32
31#endif
32
33#ifdef SRC_OVERLAP
34@group(0) @binding(0)
35var<storage, read_write> merged_src: array<DataType>;
36
37@group(0) @binding(1)
38var<storage, read_write> dst: array<DataType>;
39
40@group(0) @binding(2)
41var<uniform> params: Params;
42#else
43@group(0) @binding(0)
44var<storage, read_write> src0: array<DataType>;
45
46@group(0) @binding(1)
47var<storage, read_write> src1 : array<DataType>;
48
49@group(0) @binding(2)
50var<storage, read_write> dst: array<DataType>;
51
52@group(0) @binding(3)
53var<uniform> params: Params;
54#endif
55@compute @workgroup_size(WG_SIZE)
56fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
57
58    if (gid.x < params.ne) {
59        var i = gid.x;
60        let i3 = i / (params.ne2 * params.ne1 * params.ne0);
61        i = i % (params.ne2 * params.ne1 * params.ne0);
62        let i2 = i / (params.ne1 * params.ne0);
63        i = i % (params.ne1 * params.ne0);
64        let i1 = i / params.ne0;
65        let i0 = i % params.ne0;
66
67        var ni = array<u32, 4>(i0, i1, i2, i3);
68
69        if (ni[params.dim] < params.src0_nedim) {
70            let src_i = ni[0] * params.stride_src0_0 +
71                             ni[1] * params.stride_src0_1 +
72                             ni[2] * params.stride_src0_2 +
73                             ni[3] * params.stride_src0_3;
74#ifdef SRC_OVERLAP
75            dst[params.offset_dst + gid.x] = merged_src[params.offset_src0 + src_i];
76#else
77            dst[params.offset_dst + gid.x] = src0[params.offset_src0 + src_i];
78#endif
79        } else {
80            ni[params.dim] -= params.src0_nedim;
81            let src_i = ni[0] * params.stride_src1_0 +
82                             ni[1] * params.stride_src1_1 +
83                             ni[2] * params.stride_src1_2 +
84                             ni[3] * params.stride_src1_3;
85#ifdef SRC_OVERLAP
86            dst[params.offset_dst + gid.x] = merged_src[params.offset_src1 + src_i];
87#else
88            dst[params.offset_dst + gid.x] = src1[params.offset_src1 + src_i];
89#endif
90        }
91    }
92}