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 →
1.9 KiB77 linesraw
1enable f16;
2
3@group(0) @binding(0)
4var<storage, read_write> input: array<INPUT_TYPE>;
5@group(0) @binding(1)
6var<storage, read_write> output: array<OUTPUT_TYPE>;
7
8struct Params {
9    offset_i: u32,
10    offset_o: u32,
11
12    // element strides
13    si0: u32, si1: u32, si2: u32, si3: u32,
14    so0: u32, so1: u32, so2: u32, so3: u32,
15
16    KW: u32, KH: u32, IC: u32,
17    IW: u32, IH: u32, N: u32,
18    OW: u32, OH: u32,
19
20    // stride
21    s0: u32, s1: u32,
22    // padding
23    p0: u32, p1: u32,
24    // dilation
25    d0: u32, d1: u32,
26}
27
28@group(0) @binding(2)
29var<uniform> params: Params;
30
31@compute @workgroup_size(WG_SIZE)
32fn main(
33    @builtin(global_invocation_id) gid: vec3<u32>,
34    @builtin(num_workgroups) num_wg: vec3<u32>
35) {
36
37    let threads_per_group = u32(WG_SIZE);
38    let i_out = gid.x + (num_wg.x * threads_per_group) * gid.y;
39    let K = params.KW * params.KH * params.IC;
40    let M = params.OW * params.OH;
41    let total = K * M * params.N;
42
43    if (i_out >= total) {
44        return;
45    }
46
47    // decode (k, m, n)
48    var i = i_out;
49    let n = i / (K * M);
50    i = i % (K * M);
51    let m = i / K;
52    let k = i % K;
53
54    // decode (oh, ow)
55    let oh = m / params.OW;
56    let ow = m % params.OW;
57
58    // decode (kw, kh, ic)
59    let kw = k % params.KW;
60    let tmp = k / params.KW;
61    let kh = tmp % params.KH;
62    let ic = tmp / params.KH;
63
64    let iw_i32 = i32(ow * params.s0 + kw * params.d0) - i32(params.p0);
65    let ih_i32 = i32(oh * params.s1 + kh * params.d1) - i32(params.p1);
66
67    let output_idx = params.offset_o + k * params.so0 + ow * params.so1 + oh * params.so2 + n * params.so3;
68
69    if (iw_i32 >= 0 && iw_i32 < i32(params.IW) && ih_i32 >= 0 && ih_i32 < i32(params.IH)) {
70        let iw = u32(iw_i32);
71        let ih = u32(ih_i32);
72        let in_idx = params.offset_i + iw * params.si0 + ih * params.si1 + ic * params.si2 + n * params.si3;
73        output[output_idx] = OUTPUT_TYPE(input[in_idx]);
74    } else {
75        output[output_idx] = OUTPUT_TYPE(0.0);
76    }
77}