char-fork/llama.cpp

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

Jeff Bolzvulkan: Change mul_mat_id to pad K rather than N (#27925)77f132cb1

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
9.5 KiB317 linesraw
1#version 450
2
3#extension GL_EXT_control_flow_attributes : enable
4#extension GL_EXT_shader_16bit_storage : require
5#extension GL_EXT_shader_explicit_arithmetic_types_int8 : require
6
7#extension GL_EXT_integer_dot_product : require
8
9#ifdef FLOAT16
10#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require
11#endif
12
13#if defined(MUL_MAT_ID_USE_SUBGROUPS)
14#extension GL_KHR_shader_subgroup_basic : enable
15#extension GL_KHR_shader_subgroup_ballot : enable
16#endif
17
18#ifdef MUL_MAT_ID
19#extension GL_EXT_shader_explicit_arithmetic_types_int16 : require
20#endif
21
22#include "types.glsl"
23
24layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in;
25
26layout (binding = 0) readonly buffer A {A_TYPE data_a[];};
27#if defined(A_TYPE_PACKED16)
28layout (binding = 0) readonly buffer A_PACKED16 {A_TYPE_PACKED16 data_a_packed16[];};
29#endif
30#if defined(A_TYPE_PACKED32)
31layout (binding = 0) readonly buffer A_PACKED32 {A_TYPE_PACKED32 data_a_packed32[];};
32#endif
33layout (binding = 1) readonly buffer B {block_q8_1_x4_packed128 data_b[];};
34layout (binding = 2) writeonly buffer D {D_TYPE data_d[];};
35
36#ifdef MUL_MAT_ID
37layout (binding = 3) readonly buffer IDS {int data_ids[];};
38layout (binding = 4) readonly buffer Counts {int data_expert_count[];};
39#endif
40
41layout (push_constant) uniform parameter
42{
43    uint M;
44    uint N;
45    uint K;
46    uint stride_a;
47    uint stride_b;
48    uint stride_d;
49
50    uint batch_stride_a;
51    uint batch_stride_b;
52    uint batch_stride_d;
53
54#ifdef MUL_MAT_ID
55    uint nei0;
56    uint nei1;
57    uint nbi1;
58    uint ne11;
59    uint n_experts;
60    uint hoist_row_ids;
61#else
62    uint base_work_group_z;
63    uint num_batches;
64    uint k_split;
65    uint ne02;
66    uint ne12;
67    uint broadcast2;
68    uint broadcast3;
69#endif
70} p;
71
72layout (constant_id = 0) const uint BLOCK_SIZE = 64;
73layout (constant_id = 1) const uint BM = 64;
74layout (constant_id = 2) const uint BN = 64;
75// layout (constant_id = 3) const uint BK = 32;
76layout (constant_id = 4) const uint WM = 32;
77layout (constant_id = 5) const uint WN = 32;
78layout (constant_id = 6) const uint WMITER = 2;
79layout (constant_id = 7) const uint TM = 4;
80layout (constant_id = 8) const uint TN = 2;
81layout (constant_id = 9) const uint TK = 1;  // Only needed for coopmat
82layout (constant_id = 10) const uint WARP = 32;
83
84#define BK 32
85
86#include "mul_mmq_shmem_types.glsl"
87
88#ifdef MUL_MAT_ID
89#define BK_STEP 1
90#else
91#ifndef BK_STEP
92#define BK_STEP 4
93#endif
94#endif
95
96// Shared memory cache
97shared block_a_cache buf_a[BM * BK_STEP];
98shared block_b_cache buf_b[BN * BK_STEP];
99// Register cache
100block_a_cache cache_a[WMITER * TM];
101block_b_cache cache_b;
102
103#define LOAD_VEC_A (4 * QUANT_R_MMQ)
104#define LOAD_VEC_B 16
105
106#define NUM_WARPS (BLOCK_SIZE / WARP)
107
108#include "mul_mm_id_funcs.glsl"
109#include "mul_mmq_funcs.glsl"
110
111void main() {
112    const uint ic = gl_WorkGroupID.y;
113
114#ifdef MUL_MAT_ID
115    const uint expert_idx = gl_WorkGroupID.z;
116    if (ic * BN >= data_expert_count[expert_idx]) {
117        return;
118    }
119#endif
120#ifdef NEEDS_INIT_IQ_SHMEM
121    init_iq_shmem(gl_WorkGroupSize);
122#endif
123
124#ifndef MUL_MAT_ID
125    const uint batch_idx = gl_WorkGroupID.z + p.base_work_group_z;
126
127    const uint i13 = batch_idx / p.ne12;
128    const uint i12 = batch_idx % p.ne12;
129
130    const uint i03 = i13 / p.broadcast3;
131    const uint i02 = i12 / p.broadcast2;
132
133    const uint batch_idx_a = i03 * p.ne02 + i02;
134#endif
135
136    const uint blocks_m = (p.M + BM - 1) / BM;
137    const uint ir = gl_WorkGroupID.x % blocks_m;
138    const uint ik = gl_WorkGroupID.x / blocks_m;
139
140    const uint WNITER = (WM * WN) / (WARP * TM * TN * WMITER);
141    const uint WSUBM = WM / WMITER;
142    const uint WSUBN = WN / WNITER;
143    const uint warp_i = gl_LocalInvocationID.x / WARP;
144
145    const uint tiw = gl_LocalInvocationID.x % WARP;
146
147    const uint tiwr = tiw % (WSUBM / TM);
148    const uint tiwc = tiw / (WSUBM / TM);
149
150    const uint warp_r = warp_i % (BM / WM);
151    const uint warp_c = warp_i / (BM / WM);
152
153    const uint loadr_a = gl_LocalInvocationID.x % (BK / LOAD_VEC_A);
154    const uint loadc_a = gl_LocalInvocationID.x / (BK / LOAD_VEC_A);
155    const uint loadr_b = gl_LocalInvocationID.x % (BK / LOAD_VEC_B);
156    const uint loadc_b = gl_LocalInvocationID.x / (BK / LOAD_VEC_B);
157
158    const uint loadstride_a = BLOCK_SIZE * LOAD_VEC_A / BK;
159    const uint loadstride_b = BLOCK_SIZE * LOAD_VEC_B / BK;
160
161#ifdef MUL_MAT_ID
162    if (p.hoist_row_ids != 0) {
163        load_row_ids_hoisted(expert_idx, ic);
164    } else {
165#ifdef MUL_MAT_ID_USE_SUBGROUPS
166        if (bitCount(p.nei0) == 1) {
167            load_row_ids(expert_idx, true, ic);
168        } else {
169            load_row_ids(expert_idx, false, ic);
170        }
171#else
172        _ne1 = 0;
173        for (uint ii1 = 0; ii1 < p.nei1 && _ne1 < (ic + 1) * BN; ii1++) {
174            for (uint ii0 = 0; ii0 < p.nei0 && _ne1 < (ic + 1) * BN; ii0++) {
175                if (data_ids[ii1*p.nbi1 + ii0] == expert_idx) {
176                    if (_ne1 >= ic * BN) {
177                        row_ids[_ne1 - ic * BN] = u16vec2(ii0, ii1);
178                    }
179                    _ne1++;
180                }
181            }
182        }
183
184        barrier();
185#endif
186    }
187
188    // Workgroup has no work
189    if (ic * BN >= _ne1) return;
190#endif
191
192#ifdef MUL_MAT_ID
193    const uint start_k = 0;
194    const uint end_k = p.K;
195#else
196    const uint start_k = ik * p.k_split;
197    const uint end_k = min(p.K, (ik + 1) * p.k_split);
198#endif
199
200    uint pos_a_ib =
201#ifdef MUL_MAT_ID
202        expert_idx * (p.batch_stride_a / BK) +
203#else
204        batch_idx_a * (p.batch_stride_a / BK) +
205#endif
206        (ir * BM * p.stride_a + start_k) / BK;
207#ifdef MUL_MAT_ID
208    uint pos_b_ib = 0;
209#else
210    uint pos_b_ib = (batch_idx * p.batch_stride_b + ic * BN * p.stride_b + start_k) / BK;
211#endif
212
213    ACC_TYPE sums[WMITER * TM * WNITER * TN];
214
215    [[unroll]] for (uint i = 0; i < WMITER*TM*WNITER*TN; i++) {
216        sums[i] = ACC_TYPE(0.0f);
217    }
218
219    for (uint block = start_k; block < end_k; block += BK * BK_STEP) {
220        [[unroll]] for (uint l = 0; loadc_a + l < BM; l += loadstride_a) {
221            const uint buf_ib = loadc_a + l;
222            const uint ib = pos_a_ib + buf_ib * p.stride_a / BK;
223            const uint iqs = loadr_a;
224
225            [[unroll]] for (uint k_step = 0; k_step < BK_STEP; k_step++) {
226                if (block + k_step * BK < end_k) {
227                    block_a_to_shmem(k_step * BM + buf_ib, ib + k_step, iqs);
228                }
229            }
230        }
231        [[unroll]] for (uint l = 0; loadc_b + l < BN; l += loadstride_b) {
232            const uint buf_ib = loadc_b + l;
233
234#ifdef MUL_MAT_ID
235            const u16vec2 row_idx = row_ids[buf_ib];
236            const uint ib = pos_b_ib + row_idx.y * p.batch_stride_b / BK + (row_idx.x % p.ne11) * p.stride_b / BK;
237#else
238            const uint ib = pos_b_ib + buf_ib * p.stride_b / BK;
239#endif
240            const uint iqs = loadr_b;
241
242            [[unroll]] for (uint k_step = 0; k_step < BK_STEP; k_step++) {
243                block_b_to_shmem(k_step * BN + buf_ib, ib + k_step, iqs, block + k_step * BK < end_k);
244            }
245        }
246
247        barrier();
248
249        pos_a_ib += BK_STEP;
250        pos_b_ib += BK_STEP;
251
252        for (uint k_step = 0; k_step < BK_STEP; k_step++) {
253            // Load from shared into cache
254            [[unroll]] for (uint wsir = 0; wsir < WMITER; wsir++) {
255                [[unroll]] for (uint cr = 0; cr < TM; cr++) {
256                    const uint reg_ib = wsir * TM + cr;
257                    const uint buf_ib = warp_r * WM + wsir * WSUBM + tiwr * TM + cr;
258
259                    block_a_to_registers(reg_ib, k_step * BM + buf_ib);
260                }
261            }
262
263            [[unroll]] for (uint wsic = 0; wsic < WNITER; wsic++) {
264                [[unroll]] for (uint cc = 0; cc < TN; cc++) {
265                    const uint ib = k_step * BN + warp_c * WN + wsic * WSUBN + tiwc * TN + cc;
266                    block_b_to_registers(ib);
267
268                    [[unroll]] for (uint wsir = 0; wsir < WMITER; wsir++) {
269                        [[unroll]] for (uint cr = 0; cr < TM; cr++) {
270                            const uint cache_a_idx = wsir * TM + cr;
271                            const uint sums_idx = (wsic * TN + cc) * (WMITER * TM) + wsir * TM + cr;
272
273                            sums[sums_idx] += mmq_dot_product(cache_a_idx);
274                        }
275                    }
276                }
277            }
278        }
279
280        barrier();
281    }
282
283    const uint dr = ir * BM + warp_r * WM;
284    const uint dc = ic * BN + warp_c * WN;
285
286#ifndef MUL_MAT_ID
287    const uint offsets = batch_idx * p.batch_stride_d + ik * p.batch_stride_d * p.num_batches;
288#endif
289
290    [[unroll]] for (uint wsic = 0; wsic < WNITER; wsic++) {
291        [[unroll]] for (uint wsir = 0; wsir < WMITER; wsir++) {
292
293            const uint dr_warp = dr + wsir * WSUBM + tiwr * TM;
294            const uint dc_warp = dc + wsic * WSUBN + tiwc * TN;
295            [[unroll]] for (uint cc = 0; cc < TN; cc++) {
296#ifdef MUL_MAT_ID
297                const uint row_i = dc_warp + cc;
298                if (row_i >= _ne1) break;
299
300                const u16vec2 row_idx = row_ids[row_i - ic * BN];
301#endif // MUL_MAT_ID
302                [[unroll]] for (uint cr = 0; cr < TM; cr++) {
303                    const uint sums_idx = (wsic * TN + cc) * WMITER * TM + wsir * TM + cr;
304#ifdef MUL_MAT_ID
305                    if (dr_warp + cr < p.M) {
306                        data_d[row_idx.y * p.batch_stride_d + row_idx.x * p.stride_d + dr_warp + cr] = D_TYPE(sums[sums_idx].x);
307                    }
308#else
309                    if (dr_warp + cr < p.M && dc_warp + cc < p.N) {
310                        data_d[offsets + (dc_warp + cc) * p.stride_d + dr_warp + cr] = D_TYPE(sums[sums_idx].x);
311                    }
312#endif // MUL_MAT_ID
313                }
314            }
315        }
316    }
317}