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.8 KiB127 linesraw
1#include "common_decls.tmpl"
2enable f16;
3
4@group(0) @binding(0)
5var<storage, read_write> weights: array<WEIGHT_TYPE>;
6@group(0) @binding(1)
7var<storage, read_write> input: array<INPUT_TYPE>;
8@group(0) @binding(2)
9var<storage, read_write> output: array<OUTPUT_TYPE>;
10
11struct Params {
12    offset_w: u32,
13    offset_i: u32,
14    offset_o: u32,
15
16    // element strides
17    sw0: u32, sw1: u32, sw2: u32, sw3: u32,
18    si0: u32, si1: u32, si2: u32, si3: u32,
19    so0: u32, so1: u32, so2: u32, so3: u32,
20
21    // kernel dimensions
22    KW: u32, KH: u32, IC: u32,
23    // input dimensions
24    IW: u32, IH: u32,
25    // output dimensions
26    OW: u32, OH: u32, OC_out: u32, N_out: u32,
27
28    // stride
29    s0: u32, s1: u32,
30    // padding
31    p0: u32, p1: u32,
32    // dilation
33    d0: u32, d1: u32,
34};
35
36@group(0) @binding(3)
37var<uniform> params: Params;
38
39fn ceil_div_u32(x: u32, y: u32) -> u32 {
40    return (x + y - 1) / y;
41}
42
43// returns the first valid kernel index k such that base + k * step >= 0
44fn first_valid_k(base: i32, step: u32) -> u32 {
45    if (base >= 0) {
46        return 0;
47    }
48
49    return ceil_div_u32(u32(-base), step);
50}
51
52// returns the first invalid kernel index k such that base + k * step >= limit so valid k are in [0, end_valid_k)
53fn end_valid_k(base: i32, step: u32, limit: u32, k_max: u32) -> u32 {
54    let remaining = i32(limit) - base;
55    if (remaining <= 0) {
56        return 0;
57    }
58
59    return min(k_max, ceil_div_u32(u32(remaining), step));
60}
61
62@compute @workgroup_size(WG_SIZE)
63fn main(
64    @builtin(global_invocation_id) gid: vec3<u32>,
65    @builtin(num_workgroups) num_wg: vec3<u32>
66) {
67
68    let threads_per_group = u32(WG_SIZE);
69    let i_out = gid.x + (num_wg.x * threads_per_group) * gid.y;
70    let n_out = params.OW * params.OH * params.OC_out * params.N_out;
71
72    var sum: f32 = 0.0;
73    if (i_out >= n_out) {
74        return;
75    }
76
77    // Kernel layout: [KW, KH, IC, ..]
78    // Input layout:  [IW, IH, .., ..]
79    // Output layout: [OW, OH, OC, N]
80
81    var i = i_out;
82    let n = i / (params.OC_out * params.OH * params.OW);
83    i = i % (params.OC_out * params.OH * params.OW);
84    let oc = i / (params.OH * params.OW);
85    i = i % (params.OH * params.OW);
86    let oh = i / params.OW;
87    let ow = i % params.OW;
88
89    let ow_base = i32(ow * params.s0) - i32(params.p0);
90    let oh_base = i32(oh * params.s1) - i32(params.p1);
91
92    // clip the valid kernel window once
93    let kw_begin = first_valid_k(ow_base, params.d0);
94    let kw_end = end_valid_k(ow_base, params.d0, params.IW, params.KW);
95    let kh_begin = first_valid_k(oh_base, params.d1);
96    let kh_end = end_valid_k(oh_base, params.d1, params.IH, params.KH);
97
98    // entire receptive field is out of bounds
99    if (kw_begin >= kw_end || kh_begin >= kh_end) {
100        let out_idx = params.offset_o + ow * params.so0 + oh * params.so1 + oc * params.so2 + n * params.so3;
101        output[out_idx] = OUTPUT_TYPE(0.0);
102        return;
103    }
104
105    let weight_oc_base = params.offset_w + oc * params.sw3;
106    let input_n_base = params.offset_i + n * params.si3;
107
108    for (var ic: u32 = 0; ic < params.IC; ic += 1) {
109        let w_base_ic = ic * params.sw2 + weight_oc_base;
110        let in_base = ic * params.si2 + input_n_base;
111
112        for (var kh: u32 = kh_begin; kh < kh_end; kh += 1) {
113            let ih = u32(oh_base + i32(kh * params.d1));
114            let w_row_base = w_base_ic + kh * params.sw1;
115            let in_row_base = in_base + ih * params.si1;
116            for (var kw: u32 = kw_begin; kw < kw_end; kw += 1) {
117                let iw = u32(ow_base + i32(kw * params.d0));
118                let w_idx = w_row_base + kw * params.sw0;
119                let in_idx = in_row_base + iw * params.si0;
120                sum += f32(weights[w_idx]) * f32(input[in_idx]);
121            }
122        }
123    }
124
125    let out_idx = params.offset_o + ow * params.so0 + oh * params.so1 + oc * params.so2 + n * params.so3;
126    output[out_idx] = OUTPUT_TYPE(sum);
127}