char-fork/llama.cpp

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

Jeff Bolzvulkan: Support Q2_0 (#25430)788e07dc9

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
4.8 KiB141 linesraw
1#version 450
2
3#extension GL_EXT_shader_explicit_arithmetic_types_int32 : require
4#extension GL_EXT_integer_dot_product : require
5
6#define MMQ
7#define NEEDS_IQ1S_GRID_GPU
8#define B_TYPE block_q8_1_x4
9
10#include "mul_mat_vec_base.glsl"
11
12layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in;
13
14#if defined(DATA_A_Q2_0) || defined(DATA_A_QUANT_K)
15#define K_PER_ITER 16
16#elif defined(DATA_A_QUANT_LEGACY) || defined(DATA_A_MXFP4)
17#define K_PER_ITER 8
18#elif defined(DATA_A_IQ1_S) || defined(DATA_A_IQ1_M)
19#define K_PER_ITER 32
20#else
21#error unimplemented
22#endif
23
24uint a_offset, b_offset, d_offset;
25
26int32_t cache_b_qs[K_PER_ITER / 4];
27vec2 cache_b_ds;
28
29#include "mul_mat_vecq_funcs.glsl"
30
31void iter(inout FLOAT_TYPE temp[NUM_COLS][NUM_ROWS], const uint first_row, const uint num_rows, const uint col, const uint b_qs_idx) {
32    [[unroll]] for (uint j = 0; j < NUM_COLS; ++j) {
33        // Preload data_b block
34        const uint b_block_idx = (j*p.batch_stride_b + col) / QUANT_K_Q8_1 + b_offset;
35        const uint b_block_idx_outer = b_block_idx / 4;
36        const uint b_block_idx_inner = b_block_idx % 4;
37        cache_b_ds = vec2(data_b[b_block_idx_outer].ds[b_block_idx_inner]);
38
39#if QUANT_R == 2
40        // Assumes K_PER_ITER == 8
41        cache_b_qs[0] = data_b[b_block_idx_outer].qs[b_block_idx_inner * 8 + b_qs_idx];
42        cache_b_qs[1] = data_b[b_block_idx_outer].qs[b_block_idx_inner * 8 + b_qs_idx + 4];
43#else
44#if K_PER_ITER == 8
45        cache_b_qs[0] = data_b[b_block_idx_outer].qs[b_block_idx_inner * 8 + b_qs_idx * 2];
46        cache_b_qs[1] = data_b[b_block_idx_outer].qs[b_block_idx_inner * 8 + b_qs_idx * 2 + 1];
47#elif K_PER_ITER == 16
48        cache_b_qs[0] = data_b[b_block_idx_outer].qs[b_block_idx_inner * 8 + b_qs_idx * 4    ];
49        cache_b_qs[1] = data_b[b_block_idx_outer].qs[b_block_idx_inner * 8 + b_qs_idx * 4 + 1];
50        cache_b_qs[2] = data_b[b_block_idx_outer].qs[b_block_idx_inner * 8 + b_qs_idx * 4 + 2];
51        cache_b_qs[3] = data_b[b_block_idx_outer].qs[b_block_idx_inner * 8 + b_qs_idx * 4 + 3];
52#elif K_PER_ITER == 32
53        cache_b_qs[0] = data_b[b_block_idx_outer].qs[b_block_idx_inner * 8    ];
54        cache_b_qs[1] = data_b[b_block_idx_outer].qs[b_block_idx_inner * 8 + 1];
55        cache_b_qs[2] = data_b[b_block_idx_outer].qs[b_block_idx_inner * 8 + 2];
56        cache_b_qs[3] = data_b[b_block_idx_outer].qs[b_block_idx_inner * 8 + 3];
57        cache_b_qs[4] = data_b[b_block_idx_outer].qs[b_block_idx_inner * 8 + 4];
58        cache_b_qs[5] = data_b[b_block_idx_outer].qs[b_block_idx_inner * 8 + 5];
59        cache_b_qs[6] = data_b[b_block_idx_outer].qs[b_block_idx_inner * 8 + 6];
60        cache_b_qs[7] = data_b[b_block_idx_outer].qs[b_block_idx_inner * 8 + 7];
61#else
62#error unimplemented
63#endif
64#endif
65
66        uint ibi = first_row*p.ncols;
67        [[unroll]] for (uint n = 0; n < num_rows; ++n) {
68            const uint a_block_idx = (ibi + col)/QUANT_K_Q8_1 + a_offset;
69            ibi += p.ncols;
70
71            temp[j][n] += mmvq_dot_product(a_block_idx, b_qs_idx);
72        }
73    }
74}
75
76void compute_outputs(const uint32_t first_row, const uint32_t num_rows) {
77    const uint tid = gl_LocalInvocationID.x;
78
79    get_offsets(a_offset, b_offset, d_offset);
80    a_offset *= QUANT_K / QUANT_K_Q8_1;
81    b_offset /= QUANT_K_Q8_1;
82
83    FLOAT_TYPE temp[NUM_COLS][NUM_ROWS];
84
85    [[unroll]] for (uint j = 0; j < NUM_COLS; ++j) {
86        [[unroll]] for (uint n = 0; n < num_rows; ++n) {
87            temp[j][n] = FLOAT_TYPE(0.0f);
88        }
89    }
90
91    const uint col_stride = K_PER_ITER * BLOCK_SIZE;
92    uint num_iters = p.ncols / col_stride;
93    if (num_iters * col_stride + K_PER_ITER * tid < p.ncols) {
94        num_iters++;
95    }
96
97    const uint b_qs_idx = tid % (32 / K_PER_ITER);
98    uint col = tid * K_PER_ITER;
99    while (num_iters >= 4) {
100        // Manually partially unroll the loop
101        [[unroll]] for (uint k = 0; k < 4; ++k) {
102            iter(temp, first_row, num_rows, col, b_qs_idx);
103            col += col_stride;
104        }
105
106        num_iters -= 4;
107    }
108
109    if (num_iters >= 2) {
110        // Manually partially unroll the loop
111        iter(temp, first_row, num_rows, col, b_qs_idx);
112        col += col_stride;
113        iter(temp, first_row, num_rows, col, b_qs_idx);
114        col += col_stride;
115        num_iters -= 2;
116    }
117
118    if (num_iters > 0) {
119        iter(temp, first_row, num_rows, col, b_qs_idx);
120    }
121
122    reduce_result(temp, d_offset, first_row, num_rows, tid);
123}
124
125void main() {
126    const uint first_row = NUM_ROWS * (gl_WorkGroupID.x + gl_NumWorkGroups.x * gl_WorkGroupID.z);
127
128#ifdef NEEDS_INIT_IQ_SHMEM
129    init_iq_shmem(gl_WorkGroupSize);
130#endif
131
132    // do NUM_ROWS at a time, unless there aren't enough remaining rows
133    if (first_row + NUM_ROWS <= p.stride_d) {
134        compute_outputs(first_row, NUM_ROWS);
135    } else {
136        if (first_row >= p.stride_d) {
137            return;
138        }
139        compute_outputs(first_row, p.stride_d - first_row);
140    }
141}