char-fork/llama.cpp

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

Jeff Bolzvulkan: support solve_tri with larger N/K values (#17781)c6c5e8597

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
2.6 KiB81 linesraw
1#version 450
2
3#include "types.glsl"
4#include "generic_binary_head.glsl"
5
6layout (constant_id = 1) const uint N = 64;
7layout (constant_id = 2) const uint K = 32;
8layout (constant_id = 3) const uint BATCH_N = 32;
9
10layout(local_size_x_id = 4, local_size_y = 1, local_size_z = 1) in;
11
12uint a_base, b_base, x_base;
13
14FLOAT_TYPE get_a(uint r, uint c) {
15    return FLOAT_TYPE(data_a[a_base + r * p.nb01 + c * p.nb00]);
16}
17
18FLOAT_TYPE get_b(uint r, uint c) {
19    return FLOAT_TYPE(data_b[b_base + r * p.nb11 + c * p.nb10]);
20}
21
22void store_x(uint r, uint c, FLOAT_TYPE v) {
23    data_d[x_base + r * p.nb21 + c * p.nb20] = D_TYPE(v);
24}
25
26shared FLOAT_TYPE shA[BATCH_N * N];
27shared FLOAT_TYPE shB[BATCH_N * K];
28
29void main() {
30    const uint batch = gl_WorkGroupID.z * 262144 + gl_WorkGroupID.y * 512 + gl_WorkGroupID.x;
31    const uint tid = gl_LocalInvocationID.x;
32
33    if (batch >= p.ne02 * p.ne03) {
34        return;
35    }
36
37    const uint i3 = batch / p.ne22;
38    const uint i2 = batch % p.ne22;
39    a_base = get_aoffset() + i2 * p.nb02 + i3 * p.nb03;
40    b_base = get_boffset() + i2 * p.nb12 + i3 * p.nb13;
41    x_base = get_doffset() + i2 * p.nb22 + i3 * p.nb23;
42
43    FLOAT_TYPE X[N];
44
45    // Loop over batches of rows
46    [[unroll]] for (uint row_base = 0; row_base < N; row_base += BATCH_N) {
47        const uint cur_N = min(BATCH_N, N - row_base);
48
49        // Load the A matrix batch into shA
50        [[unroll]] for (uint i = 0; i < cur_N * N; i += gl_WorkGroupSize.x) {
51            uint idx = i + tid;
52            if (((cur_N * N) % gl_WorkGroupSize.x == 0) || idx < cur_N * N) {
53                shA[idx] = get_a(row_base + idx / N, idx % N);
54            }
55        }
56        // Load the B matrix batch into shB
57        [[unroll]] for (uint i = 0; i < cur_N * K; i += gl_WorkGroupSize.x) {
58            uint idx = i + tid;
59            if (((cur_N * K) % gl_WorkGroupSize.x == 0) || idx < cur_N * K) {
60                shB[idx] = get_b(row_base + idx / K, idx % K);
61            }
62        }
63        barrier();
64
65        // Each thread solves one column
66        if (tid < K) {
67            [[unroll]] for (uint row_offset = 0; row_offset < cur_N; ++row_offset) {
68                uint r = row_base + row_offset;
69                FLOAT_TYPE b = shB[row_offset * K + tid];
70                // Compute x[r,c] = (b[r,c] - sum(a[r,c]*x[c])) / a[r,r]
71                [[unroll]] for (int c = 0; c < r; ++c) {
72                    b -= shA[row_offset * N + c] * X[c];
73                }
74                FLOAT_TYPE x = b / shA[row_offset * N + r];
75                X[r] = x;
76                store_x(r, tid, x);
77            }
78        }
79        barrier();
80    }
81}