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.2 KiB120 linesraw
1@group(0) @binding(0)
2var<storage, read_write> src0: array<f32>;
3
4@group(0) @binding(1)
5var<storage, read_write> src1: array<f32>;
6
7@group(0) @binding(2)
8var<storage, read_write> dst: array<f32>;
9
10struct Params {
11    offset_src0: u32,
12    offset_src1: u32,
13    offset_dst: u32,
14
15    stride_src00: u32,
16    stride_src01: u32,
17    stride_src02: u32,
18    stride_src03: u32,
19
20    stride_src10: u32,
21    stride_src11: u32,
22    stride_src12: u32,
23    stride_src13: u32,
24
25    stride_dst0: u32,
26    stride_dst1: u32,
27    stride_dst2: u32,
28    stride_dst3: u32,
29
30    k: u32,
31    ne2: u32,
32};
33
34@group(0) @binding(3)
35var<uniform> params: Params;
36
37var<workgroup> shA: array<f32, BATCH_N * N>;
38var<workgroup> shB: array<f32, BATCH_N * K_TILE>;
39
40fn src0_idx(row: u32, col: u32, i2: u32, i3: u32) -> u32 {
41    return params.offset_src0 +
42           col * params.stride_src00 +
43           row * params.stride_src01 +
44           i2 * params.stride_src02 +
45           i3 * params.stride_src03;
46}
47
48fn src1_idx(row: u32, col: u32, i2: u32, i3: u32) -> u32 {
49    return params.offset_src1 +
50           col * params.stride_src10 +
51           row * params.stride_src11 +
52           i2 * params.stride_src12 +
53           i3 * params.stride_src13;
54}
55
56fn dst_idx(row: u32, col: u32, i2: u32, i3: u32) -> u32 {
57    return params.offset_dst +
58           col * params.stride_dst0 +
59           row * params.stride_dst1 +
60           i2 * params.stride_dst2 +
61           i3 * params.stride_dst3;
62}
63
64@compute @workgroup_size(WG_SIZE)
65fn main(
66    @builtin(workgroup_id) workgroup_id: vec3<u32>,
67    @builtin(local_invocation_id) local_id: vec3<u32>
68) {
69    let batch = workgroup_id.y;
70    let col = workgroup_id.x * WG_SIZE + local_id.x;
71    let i3 = batch / params.ne2;
72    let i2 = batch % params.ne2;
73    let active_lane = local_id.x < K_TILE;
74    let active_col = active_lane && col < params.k;
75
76    var X: array<f32, N>;
77
78    for (var row_base = 0u; row_base < N; row_base += BATCH_N) {
79        let cur_n = min(BATCH_N, N - row_base);
80
81        for (var i = local_id.x; i < cur_n * N; i += WG_SIZE) {
82            let tile_row = i / N;
83            let tile_col = i % N;
84            shA[i] = src0[src0_idx(row_base + tile_row, tile_col, i2, i3)];
85        }
86
87        for (var i = local_id.x; i < cur_n * K_TILE; i += WG_SIZE) {
88            let tile_row = i / K_TILE;
89            let tile_col = i % K_TILE;
90            let global_col = workgroup_id.x * WG_SIZE + tile_col;
91            let sh_idx = tile_row * K_TILE + tile_col;
92
93            if (global_col < params.k) {
94                shB[sh_idx] = src1[src1_idx(row_base + tile_row, global_col, i2, i3)];
95            } else {
96                shB[sh_idx] = 0.0;
97            }
98        }
99
100        workgroupBarrier();
101
102        if (active_col) {
103            for (var row_offset = 0u; row_offset < cur_n; row_offset++) {
104                let r = row_base + row_offset;
105                var b = shB[row_offset * K_TILE + local_id.x];
106                let a_row = row_offset * N;
107
108                for (var t = 0u; t < r; t++) {
109                    b -= shA[a_row + t] * X[t];
110                }
111
112                let x = b / shA[a_row + r];
113                X[r] = x;
114                dst[dst_idx(r, col, i2, i3)] = x;
115            }
116        }
117
118        workgroupBarrier();
119    }
120}