char-fork/llama.cpp

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

Matt Corallovulkan: Switch MUL_MAT_VEC to 4 K per iteration for F16/32 (#22887)c6e408837

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
8.9 KiB264 linesraw
1#version 450
2
3#extension GL_EXT_shader_explicit_arithmetic_types_int32 : require
4
5#include "mul_mat_vec_base.glsl"
6#include "dequant_funcs.glsl"
7
8layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in;
9
10#if !defined(DATA_A_F32) && !defined(DATA_A_F16) && !defined(DATA_A_BF16)
11#define K_PER_ITER 8
12#else
13#define K_PER_ITER 4
14#endif
15
16
17uint a_offset, b_offset, d_offset, y_offset;
18
19vec4 load_b(const uint j, const uint iybs, const uint iqs, const bool lastiter, out bool OOB_y, out bool OOB_z, out bool OOB_w) {
20    // Check if the latter elements are OOB, and don't fetch B or accumulate it.
21    OOB_y = lastiter && (iybs + iqs + y_offset >= p.ncols);
22    OOB_z = lastiter && (iybs + iqs + y_offset*2 >= p.ncols);
23    OOB_w = lastiter && (iybs + iqs + y_offset*3 >= p.ncols);
24
25    if (!OOB_w) {
26        return vec4(FLOAT_TYPE(data_b[j*p.batch_stride_b + b_offset + iybs + iqs]),
27                 FLOAT_TYPE(data_b[j*p.batch_stride_b + b_offset + iybs + iqs + y_offset]),
28                 FLOAT_TYPE(data_b[j*p.batch_stride_b + b_offset + iybs + iqs + y_offset*2]),
29                 FLOAT_TYPE(data_b[j*p.batch_stride_b + b_offset + iybs + iqs + y_offset*3]));
30    } else if (!OOB_z) {
31        return vec4(FLOAT_TYPE(data_b[j*p.batch_stride_b + b_offset + iybs + iqs]),
32                 FLOAT_TYPE(data_b[j*p.batch_stride_b + b_offset + iybs + iqs + y_offset]),
33                 FLOAT_TYPE(data_b[j*p.batch_stride_b + b_offset + iybs + iqs + y_offset*2]),
34                 0);
35    } else if (!OOB_y) {
36        return vec4(FLOAT_TYPE(data_b[j*p.batch_stride_b + b_offset + iybs + iqs]),
37                 FLOAT_TYPE(data_b[j*p.batch_stride_b + b_offset + iybs + iqs + y_offset]),
38                 0, 0);
39    } else {
40        return vec4(FLOAT_TYPE(data_b[j*p.batch_stride_b + b_offset + iybs + iqs]),
41                 0, 0, 0);
42    }
43}
44
45void iter(inout FLOAT_TYPE temp[NUM_COLS][NUM_ROWS], const uint first_row, const uint num_rows, const uint tid, const uint i, bool lastiter)
46{
47    [[unroll]] for (uint j = 0; j < NUM_COLS; ++j) {
48        const uint col = i*BLOCK_SIZE + K_PER_ITER*tid;
49        const uint iqs = (col%QUANT_K)/QUANT_R; // quant index
50        const uint iybs = col - col%QUANT_K; // y block start index
51
52#if K_PER_ITER == 8
53#if QUANT_R == 2
54        // Note that we end up fetching bogus elements here, but its fine as they'll be
55        // within an accessible block.
56        const vec4 bv02 = vec4(data_b_v4[(j*p.batch_stride_b + b_offset + iybs + iqs) / 4]);
57        const vec4 bv13 = vec4(data_b_v4[(j*p.batch_stride_b + b_offset + iybs + iqs + y_offset) / 4]);
58        const vec4 bv0 = vec4(bv02.x, bv13.x, bv02.y, bv13.y);
59        const vec4 bv1 = vec4(bv02.z, bv13.z, bv02.w, bv13.w);
60#else
61        const vec4 bv0 = vec4(data_b_v4[(j*p.batch_stride_b + b_offset + iybs + iqs) / 4]);
62        const vec4 bv1 = vec4(data_b_v4[(j*p.batch_stride_b + b_offset + iybs + iqs) / 4 + 1]);
63#endif
64#else
65        bool OOB_y;
66        bool OOB_z;
67        bool OOB_w;
68
69        const vec4 b = load_b(j, iybs, iqs, lastiter, OOB_y, OOB_z, OOB_w);
70#endif
71        uint ibi = first_row*p.ncols;
72        [[unroll]] for (uint n = 0; n < num_rows; ++n) {
73            const uint ib = (ibi + col)/QUANT_K; // block index
74            ibi += p.ncols;
75
76#if K_PER_ITER == 8
77            vec4 v = dequantize4(ib, iqs, a_offset);
78            vec4 v2 = dequantize4(ib, iqs+(4/QUANT_R), a_offset);
79
80            const vec2 dm = get_dm(ib, a_offset);
81            if (dm.y != 0) { // quant has min component
82                v = v * dm.x + dm.y;
83                v2 = v2 * dm.x + dm.y;
84            }
85
86            // matrix multiplication
87            FLOAT_TYPE rowtmp = dot(bv0, v);
88            rowtmp += dot(bv1, v2);
89
90            if (dm.y == 0)
91                rowtmp *= dm.x;
92
93            temp[j][n] += rowtmp;
94#else
95            if (!OOB_w) {
96                const vec4 v = dequantize4(ib, iqs, a_offset);
97                temp[j][n] += dot(v, b);
98            } else if (!OOB_z) {
99                const vec2 v0 = dequantize(ib, iqs, a_offset);
100                const FLOAT_TYPE v1 = dequantize1(ib + 2/QUANT_R, iqs, a_offset);
101                const vec3 v = vec3(v0.x, v0.y, v1);
102                const vec3 b0 = vec3(b.x, b.y, b.z);
103                temp[j][n] += dot(v, b0);
104            } else if (!OOB_y) {
105                const vec2 v0 = dequantize(ib, iqs, a_offset);
106                const vec2 b0 = vec2(b.x, b.y);
107                temp[j][n] += dot(v0, b0);
108            } else {
109                const FLOAT_TYPE v = dequantize1(ib, iqs, a_offset);
110                temp[j][n] = fma(v, b.x, temp[j][n]);
111            }
112#endif
113        }
114    }
115}
116
117#if defined(DATA_A_F32) || defined(DATA_A_F16) || defined(DATA_A_BF16)
118void iter_aligned_nonquant(inout FLOAT_TYPE temp[NUM_COLS][NUM_ROWS], const uint first_row, const uint num_rows, const uint tid, const uint i)
119{
120    [[unroll]] for (uint j = 0; j < NUM_COLS; ++j) {
121        const uint col = i*BLOCK_SIZE + K_PER_ITER*tid;
122        const uint iqs = 0; // quant index
123        const uint iybs = col; // y block start index
124
125        const vec4 b = data_b_v4[(j*p.batch_stride_b + b_offset + iybs + iqs) / 4];
126
127        uint ibi = first_row*p.ncols;
128        [[unroll]] for (uint n = 0; n < num_rows; ++n) {
129            const uint ib = (ibi + col)/QUANT_K; // block index
130            ibi += p.ncols;
131
132            const vec4 v = dequantize4_2aligned(ib, iqs, a_offset);
133
134            // matrix multiplication
135            temp[j][n] += dot(v, b);
136        }
137    }
138}
139#endif
140
141void compute_outputs(const uint32_t first_row, const uint32_t num_rows) {
142    const uint tid = gl_LocalInvocationID.x;
143
144    get_offsets(a_offset, b_offset, d_offset);
145    const bool is_aligned_nonquant =
146        p.batch_stride_b % 4 == 0 && b_offset % 4 == 0 &&
147        p.ncols % 4 == 0 && BLOCK_SIZE % 4 == 0 &&
148        K_PER_ITER == 4;
149
150    y_offset = QUANT_R == 1 ? 1 : QUANT_K/2;
151
152    FLOAT_TYPE temp[NUM_COLS][NUM_ROWS];
153
154    [[unroll]] for (uint j = 0; j < NUM_COLS; ++j) {
155        [[unroll]] for (uint i = 0; i < NUM_ROWS; ++i) {
156            temp[j][i] = FLOAT_TYPE(0);
157        }
158    }
159
160    uint num_iters = p.ncols / (K_PER_ITER * BLOCK_SIZE);
161    if (num_iters * K_PER_ITER * BLOCK_SIZE + K_PER_ITER*tid < p.ncols) {
162        num_iters++;
163    }
164    int unroll_count = 4;
165    uint unrolled_iters = num_iters & ~(unroll_count - 1);
166
167    uint i = 0;
168
169#if K_PER_ITER == 4
170    // If the K dimension is odd, we need lastiter==true on the last iteration
171    // so OOB is computed correctly. Skip some unrolling to make that happen.
172    if ((p.ncols & 3) != 0 &&
173        unrolled_iters == num_iters &&
174        unrolled_iters > 0) {
175        unrolled_iters -= unroll_count;
176    }
177    if (is_aligned_nonquant) {
178        while (i < unrolled_iters) {
179            // Manually partially unroll the loop
180            [[unroll]] for (uint k = 0; k < unroll_count; ++k) {
181                iter_aligned_nonquant(temp, first_row, num_rows, tid, i*K_PER_ITER);
182                i++;
183            }
184        }
185    } else {
186#endif
187    while (i < unrolled_iters) {
188        // Manually partially unroll the loop
189        [[unroll]] for (uint k = 0; k < unroll_count; ++k) {
190            iter(temp, first_row, num_rows, tid, i*K_PER_ITER, false);
191            i++;
192        }
193    }
194#if K_PER_ITER == 4
195    }
196#endif
197
198    unroll_count = 2;
199    unrolled_iters = num_iters & ~(unroll_count - 1);
200
201#if K_PER_ITER == 4
202    if ((p.ncols & 3) != 0 &&
203        unrolled_iters == num_iters &&
204        unrolled_iters > 0) {
205        unrolled_iters -= unroll_count;
206    }
207
208    if (is_aligned_nonquant) {
209        while (i < unrolled_iters && is_aligned_nonquant) {
210            // Manually partially unroll the loop
211            [[unroll]] for (uint k = 0; k < unroll_count; ++k) {
212                iter_aligned_nonquant(temp, first_row, num_rows, tid, i*K_PER_ITER);
213                i++;
214            }
215        }
216    } else {
217#endif
218    while (i < unrolled_iters) {
219        // Manually partially unroll the loop
220        [[unroll]] for (uint k = 0; k < unroll_count; ++k) {
221            iter(temp, first_row, num_rows, tid, i*K_PER_ITER, false);
222            i++;
223        }
224    }
225#if K_PER_ITER == 4
226    }
227#endif
228
229#if K_PER_ITER == 4
230    if (is_aligned_nonquant) {
231        while (i < num_iters) {
232            iter_aligned_nonquant(temp, first_row, num_rows, tid, i*K_PER_ITER);
233            i++;
234        }
235    } else {
236#endif
237    while (i < num_iters) {
238        iter(temp, first_row, num_rows, tid, i*K_PER_ITER, true);
239        i++;
240    }
241#if K_PER_ITER == 4
242    }
243#endif
244
245    reduce_result(temp, d_offset, first_row, num_rows, tid);
246}
247
248void main() {
249    const uint first_row = NUM_ROWS * (gl_WorkGroupID.x + gl_NumWorkGroups.x * gl_WorkGroupID.z);
250
251#ifdef NEEDS_INIT_IQ_SHMEM
252    init_iq_shmem(gl_WorkGroupSize);
253#endif
254
255    // do NUM_ROWS at a time, unless there aren't enough remaining rows
256    if (first_row + NUM_ROWS <= p.stride_d) {
257        compute_outputs(first_row, NUM_ROWS);
258    } else {
259        if (first_row >= p.stride_d) {
260            return;
261        }
262        compute_outputs(first_row, min(NUM_ROWS, p.stride_d - first_row));
263    }
264}