char-fork/llama.cpp

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

Reese Levineggml-webgpu: Fix dequantization helpers to not pass in pointers (#21872)20d3bc2cc

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
7.1 KiB195 linesraw
1enable f16;
2
3#define DECLARE_BYTE_LOADERS_SRC0
4#include "common_decls.tmpl"
5
6#include "mul_mat_decls.tmpl"
7
8#ifdef VEC
9fn store_val(acc: array<array<f16, TILE_M>, TILE_N>, tn: u32, tm: u32) -> vec4<f32> {
10    return vec4<f32>(f32(acc[tn][tm]), f32(acc[tn][tm + 1]), f32(acc[tn][tm + 2]), f32(acc[tn][tm + 3]));
11}
12#endif
13
14#ifdef SCALAR
15fn store_val(acc: array<array<f16, TILE_M>, TILE_N>, tn: u32, tm: u32) -> f32 {
16    return f32(acc[tn][tm]);
17}
18#endif
19
20struct MulMatIdParams {
21    offset_src0: u32,
22    offset_src1: u32,
23    offset_dst: u32,
24
25    k: u32,
26    m: u32,
27    n_expert: u32,
28    n_expert_used: u32,
29    n_tokens: u32,
30    b_ne1: u32,
31
32    stride_01: u32,
33    stride_11: u32,
34    stride_02: u32,
35    stride_12: u32,
36};
37
38@group(0) @binding(0) var<storage, read_write> src0: array<SRC0_TYPE>; // [cols, rows, n_expert]
39@group(0) @binding(1) var<storage, read_write> src1: array<SRC1_TYPE>; // [cols, b_ne1, n_tokens]
40@group(0) @binding(2) var<storage, read_write> dst: array<DST_TYPE>;   // [rows, n_expert_used, n_tokens]
41@group(0) @binding(3) var<storage, read_write> global_gathered_expert_used: array<u32>; // [n_expert][n_tokens]
42@group(0) @binding(4) var<storage, read_write> global_gathered_tokens: array<u32>; // [n_expert][n_tokens]
43@group(0) @binding(5) var<storage, read_write> gathered_count_ids: array<u32>; // [n_expert]
44
45@group(0) @binding(6) var<uniform> params: MulMatIdParams;
46
47fn get_local_n(thread_id: u32) -> u32 {
48    return thread_id / WORKGROUP_SIZE_M;
49}
50fn get_local_m(thread_id: u32) -> u32 {
51    return thread_id % WORKGROUP_SIZE_M;
52}
53
54const TOTAL_WORKGROUP_SIZE = WORKGROUP_SIZE_M * WORKGROUP_SIZE_N;
55const TILE_SRC0_SHMEM = TILE_K * WORKGROUP_SIZE_M * TILE_M;
56const TILE_SRC1_SHMEM = TILE_K * WORKGROUP_SIZE_N * TILE_N;
57
58var<workgroup> shmem: array<f16, TILE_SRC0_SHMEM + TILE_SRC1_SHMEM>;
59var<workgroup> gathered_expert_used: array<u32, TILE_N * WORKGROUP_SIZE_N>;
60var<workgroup> gathered_tokens: array<u32, TILE_N * WORKGROUP_SIZE_N>;
61
62#ifdef INIT_SRC1_SHMEM_FLOAT
63fn init_shmem_id_src1(thread_id: u32, offset_src1: u32, rest_token_n: u32, k_outer: u32) {
64    for (var elem_idx = thread_id * VEC_SIZE; elem_idx < TILE_SRC1_SHMEM; elem_idx += TOTAL_WORKGROUP_SIZE * VEC_SIZE) {
65        let tile_n = elem_idx / TILE_K;
66        let tile_k = elem_idx % TILE_K;
67        if (tile_n < rest_token_n) {
68            let global_src10 = k_outer + tile_k;
69            let expert_used_idx = gathered_expert_used[tile_n] % params.b_ne1;
70            let token_idx = gathered_tokens[tile_n];
71            let src1_idx = offset_src1 + token_idx * params.stride_12 + expert_used_idx * params.stride_11 + global_src10;
72            let src1_val = select(
73                SRC1_TYPE(0.0),
74                src1[src1_idx/VEC_SIZE],
75                global_src10 < params.k);
76            store_shmem(SHMEM_TYPE(src1_val), TILE_SRC0_SHMEM + elem_idx);
77        } else {
78            store_shmem(SHMEM_TYPE(0.0), TILE_SRC0_SHMEM + elem_idx);
79        }
80    }
81}
82#endif // INIT_SRC1_SHMEM_FLOAT
83
84@compute @workgroup_size(TOTAL_WORKGROUP_SIZE)
85fn main(@builtin(workgroup_id) wg_id: vec3<u32>,
86        @builtin(local_invocation_id) local_id: vec3<u32>,
87        @builtin(num_workgroups) num_wg: vec3<u32>) {
88
89    let thread_id = local_id.x;
90    let local_m = get_local_m(thread_id);
91    let local_n = get_local_n(thread_id);
92
93    var expert_idx:u32 = 0xFFFFFFFFu;
94    var wg_in_batch:u32 = 0;
95    var wg_sum:u32 = 0;
96    let wg_m_count = (params.m + WORKGROUP_SIZE_M * TILE_M - 1u) / (WORKGROUP_SIZE_M * TILE_M);
97    let wg_linear = wg_id.y * num_wg.x + wg_id.x;
98
99    for (var i = 0u;i < params.n_expert;i += 1) {
100        let wg_n_count = (gathered_count_ids[i] + WORKGROUP_SIZE_N * TILE_N - 1u) / (WORKGROUP_SIZE_N * TILE_N);
101        let wg_per_matrix = wg_m_count * wg_n_count;
102        if (wg_sum <= wg_linear && wg_linear < wg_sum + wg_per_matrix) {
103            expert_idx = i;
104            wg_in_batch = wg_linear - wg_sum;
105            break;
106        }
107        wg_sum += wg_per_matrix;
108    }
109
110    let is_valid = expert_idx != 0xFFFFFFFFu;
111
112    var wg_m: u32 = 0;
113    var wg_n: u32 = 0;
114    var offset_wg_m: u32 = 0;
115    var offset_wg_n: u32 = 0;
116    var rest_token_n: u32 = 0;
117    var src0_batch_offset: u32 = 0;
118
119    wg_m = wg_in_batch % wg_m_count;
120    wg_n = wg_in_batch / wg_m_count;
121
122    offset_wg_m = wg_m * WORKGROUP_SIZE_M * TILE_M;
123    offset_wg_n = wg_n * WORKGROUP_SIZE_N * TILE_N;
124
125    if (is_valid) {
126        rest_token_n = gathered_count_ids[expert_idx] - offset_wg_n;
127        let global_gathered_base = expert_idx * params.n_tokens + offset_wg_n;
128        for (var i = thread_id; i < TILE_N * WORKGROUP_SIZE_N && offset_wg_n + i < gathered_count_ids[expert_idx]; i += TOTAL_WORKGROUP_SIZE) {
129            gathered_expert_used[i] = global_gathered_expert_used[global_gathered_base + i];
130            gathered_tokens[i] = global_gathered_tokens[global_gathered_base + i];
131        }
132        src0_batch_offset = params.offset_src0 + expert_idx * params.stride_02;
133    }
134
135    workgroupBarrier();
136
137    let output_row_base = offset_wg_m + local_m * TILE_M;
138    let output_col_base = offset_wg_n + local_n * TILE_N;
139
140    let dst2_stride = params.m * params.n_expert_used;
141    let dst1_stride = params.m;
142
143    var acc: array<array<f16, TILE_M>, TILE_N>;
144
145    for (var k_outer = 0u; k_outer < params.k; k_outer += TILE_K) {
146
147        if (is_valid) {
148            init_shmem_src0(thread_id, src0_batch_offset, offset_wg_m, k_outer);
149            init_shmem_id_src1(thread_id, params.offset_src1, rest_token_n, k_outer);
150        }
151
152        workgroupBarrier();
153
154        if (is_valid) {
155            let k_end = min(TILE_K, params.k - k_outer);
156
157            for (var k_inner = 0u; k_inner < k_end; k_inner++) {
158                var src0_tile: array<f16, TILE_M>;
159                for (var tm = 0u; tm < TILE_M; tm++) {
160                    let src0_m = local_m * TILE_M + tm;
161                    let src0_idx = k_inner + src0_m * TILE_K;
162                    src0_tile[tm] = shmem[src0_idx];
163                }
164                for (var tn = 0u; tn < TILE_N; tn++) {
165                    let src1_n = local_n * TILE_N + tn;
166                    let src1_idx = src1_n * TILE_K + k_inner;
167                    let src1_val = shmem[TILE_SRC0_SHMEM + src1_idx];
168                    for (var tm = 0u; tm < TILE_M; tm++) {
169                        acc[tn][tm] += src0_tile[tm] * src1_val;
170                    }
171                }
172            }
173        }
174
175        workgroupBarrier();
176    }
177
178    if (is_valid) {
179        for (var tn = 0u; tn < TILE_N; tn++) {
180            let n_idx = output_col_base + tn;
181            if (n_idx < gathered_count_ids[expert_idx]) {
182                let dst1_idx = gathered_expert_used[n_idx - offset_wg_n];
183                let dst2_idx = gathered_tokens[n_idx - offset_wg_n];
184                let dst12_offset = params.offset_dst + dst2_idx * dst2_stride + dst1_idx * dst1_stride;
185                for (var tm = 0u; tm < TILE_M; tm += VEC_SIZE) {
186                    let global_row = output_row_base + tm;
187                    if (global_row < params.m) {
188                        let dst_idx = dst12_offset + global_row;
189                        dst[dst_idx/VEC_SIZE] = store_val(acc, tn, tm);
190                    }
191                }
192            }
193        }
194    }
195}