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 →
3.3 KiB152 linesraw
1#if defined(SRC_F16) || defined(DST_F16)
2enable f16;
3#endif
4
5#ifdef SRC_F16
6#define SRC_TYPE f16
7#else
8#define SRC_TYPE f32
9#endif
10
11#ifdef DST_F16
12#define DST_TYPE f16
13#else
14#define DST_TYPE f32
15#endif
16
17struct Params {
18    offset_src: u32, // in elements
19    offset_dst: u32, // in elements
20
21    // Strides (in elements)
22    stride_src1: u32,
23    stride_src2: u32,
24    stride_src3: u32,
25
26    stride_dst1: u32,
27    stride_dst2: u32,
28    stride_dst3: u32,
29
30    // Shape of src/dst
31    ne0: u32,
32    ne1: u32,
33    ne2: u32,
34
35    eps: f32
36};
37
38@group(0) @binding(0)
39var<storage, read_write> src: array<SRC_TYPE>;
40
41#ifdef INPLACE
42@group(0) @binding(1)
43var<uniform> params: Params;
44#else
45@group(0) @binding(1)
46var<storage, read_write> dst: array<DST_TYPE>;
47
48@group(0) @binding(2)
49var<uniform> params: Params;
50#endif
51
52var<workgroup> scratch: array<f32, WG_SIZE * 2u>;
53
54@compute @workgroup_size(WG_SIZE)
55fn main(@builtin(workgroup_id) wid: vec3<u32>,
56        @builtin(local_invocation_id) lid: vec3<u32>) {
57
58    // one thread per row
59    var i = wid.x;
60    let i3 = i / (params.ne2 * params.ne1);
61    i = i % (params.ne2 * params.ne1);
62    let i2 = i / params.ne1;
63    let i1 = i % params.ne1;
64    let i_src_row = params.offset_src + i3 * params.stride_src3 + i2 * params.stride_src2 + i1 * params.stride_src1;
65    let i_dst_row = params.offset_dst + i3 * params.stride_dst3 + i2 * params.stride_dst2 + i1 * params.stride_dst1;
66
67    let elems = (params.ne0 + WG_SIZE - 1) / WG_SIZE;
68
69    var sum = 0.0f;
70    var col = lid.x;
71    for (var j: u32 = 0; j < elems; j++) {
72        if (col >= params.ne0) {
73            break;
74        }
75        let v = f32(src[i_src_row + col]);
76#ifdef NORM
77        sum += v;
78#else
79        sum += v * v;
80#endif
81        col += WG_SIZE;
82    }
83
84    scratch[lid.x] = sum;
85    workgroupBarrier();
86
87    var offset: u32 = WG_SIZE / 2u;
88    while (offset > 0) {
89        if (lid.x < offset) {
90            scratch[lid.x] += scratch[lid.x + offset];
91        }
92        offset /= 2u;
93        workgroupBarrier();
94    }
95    sum = scratch[0];
96
97#ifdef NORM
98    let mean = sum / f32(params.ne0);
99    var sq_sum = 0.0f;
100    col = lid.x;
101    for (var j: u32 = 0; j < elems; j++) {
102        if (col >= params.ne0) {
103            break;
104        }
105        let v = f32(src[i_src_row + col]);
106        let d = v - mean;
107        sq_sum += d * d;
108        col += WG_SIZE;
109    }
110
111    workgroupBarrier();
112    scratch[lid.x] = sq_sum;
113    workgroupBarrier();
114    offset = WG_SIZE / 2u;
115    while (offset > 0) {
116        if (lid.x < offset) {
117            scratch[lid.x] += scratch[lid.x + offset];
118        }
119        offset /= 2u;
120        workgroupBarrier();
121    }
122
123    let variance = scratch[0] / f32(params.ne0);
124    let scale = 1.0 / sqrt(variance + params.eps);
125#elif defined(RMS_NORM)
126    let scale = 1.0/sqrt(sum/f32(params.ne0) + params.eps);
127#elif defined(L2_NORM)
128    let scale = 1.0/max(sqrt(sum), params.eps);
129#endif
130
131#ifdef NORM
132    let mean_val = mean;
133#else
134    let mean_val = 0.0f;
135#endif
136
137    col = lid.x;
138    for (var j: u32 = 0; j < elems; j++) {
139        if (col >= params.ne0) {
140            break;
141        }
142        let i_src = i_src_row + col;
143        let i_dst = i_dst_row + col;
144        let v = src[i_src];
145#ifdef INPLACE
146        src[i_dst] = scale * (v - mean_val);
147#else
148        dst[i_dst] = scale * (v - mean_val);
149#endif
150        col += WG_SIZE;
151    }
152}