char-fork/llama.cpp

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

Reese Levineggml webgpu: ops support for qwen3.5 (SET, TRI_SOLVE, SSM_CONV, GATED_DELTA_NET) + GET_ROWS optimization (#20687)c1258830b

char/aliceai-t5-moe
2.8 KiB109 linesraw
1#ifdef TYPE_I32
2#define TYPE i32
3#else
4#define TYPE f32
5#endif
6
7#ifndef INPLACE
8@group(0) @binding(0)
9var<storage, read_write> src0: array<TYPE>;
10#define SRC1_BINDING 1
11#else
12#define SRC1_BINDING 0
13#endif
14
15#define DST_BINDING SRC1_BINDING + 1
16#define PARAMS_BINDING SRC1_BINDING + 2
17
18@group(0) @binding(SRC1_BINDING)
19var<storage, read_write> src1: array<TYPE>;
20
21@group(0) @binding(DST_BINDING)
22var<storage, read_write> dst: array<TYPE>;
23
24struct Params {
25    ne: u32,
26    offset_src0: u32,
27    offset_src1: u32,
28    offset_view: u32,
29
30    stride_src10: u32,
31    stride_src11: u32,
32    stride_src12: u32,
33    stride_src13: u32,
34
35    stride_dst10: u32,
36    stride_dst11: u32,
37    stride_dst12: u32,
38    stride_dst13: u32,
39
40    src1_ne0: u32,
41    src1_ne1: u32,
42    src1_ne2: u32,
43    src1_ne3: u32,
44};
45
46@group(0) @binding(PARAMS_BINDING)
47var<uniform> params: Params;
48
49fn decode_src1_coords(idx: u32) -> vec4<u32> {
50    var i = idx;
51    let plane = params.src1_ne2 * params.src1_ne1 * params.src1_ne0;
52    let i3 = i / plane;
53    i = i % plane;
54    let row = params.src1_ne1 * params.src1_ne0;
55    let i2 = i / row;
56    i = i % row;
57    let i1 = i / params.src1_ne0;
58    let i0 = i % params.src1_ne0;
59    return vec4<u32>(i0, i1, i2, i3);
60}
61
62fn decode_view_coords(rel: u32) -> vec4<u32> {
63    let i3 = rel / params.stride_dst13;
64    let rem3 = rel % params.stride_dst13;
65    let i2 = rem3 / params.stride_dst12;
66    let rem2 = rem3 % params.stride_dst12;
67    let i1 = rem2 / params.stride_dst11;
68    let i0 = rem2 % params.stride_dst11;
69    return vec4<u32>(i0, i1, i2, i3);
70}
71
72fn view_rel_from_coords(coords: vec4<u32>) -> u32 {
73    return coords.x * params.stride_dst10 + coords.y * params.stride_dst11 +
74           coords.z * params.stride_dst12 + coords.w * params.stride_dst13;
75}
76
77fn src1_idx_from_coords(coords: vec4<u32>) -> u32 {
78    return coords.x * params.stride_src10 + coords.y * params.stride_src11 +
79           coords.z * params.stride_src12 + coords.w * params.stride_src13;
80}
81
82fn in_set_view(rel: u32, coords: vec4<u32>) -> bool {
83    return view_rel_from_coords(coords) == rel;
84}
85
86@compute @workgroup_size(WG_SIZE)
87fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
88    if (gid.x >= params.ne) {
89        return;
90    }
91
92#ifdef INPLACE
93    let coords = decode_src1_coords(gid.x);
94
95    let src1_idx = params.offset_src1 + src1_idx_from_coords(coords);
96    let dst_idx = params.offset_view + view_rel_from_coords(coords);
97
98    dst[dst_idx] = src1[src1_idx];
99#else
100    let rel = select(params.ne, gid.x - params.offset_view, gid.x >= params.offset_view);
101    let coords = decode_view_coords(rel);
102
103    if (rel < params.stride_dst13 * params.src1_ne3 && in_set_view(rel, coords)) {
104        dst[gid.x] = src1[params.offset_src1 + src1_idx_from_coords(coords)];
105    } else {
106        dst[gid.x] = src0[params.offset_src0 + gid.x];
107    }
108#endif
109}