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 →
18.3 KiB488 linesraw
1#extension GL_EXT_shader_explicit_arithmetic_types_int32 : require
2#extension GL_EXT_shader_explicit_arithmetic_types_int16 : require
3#extension GL_EXT_shader_explicit_arithmetic_types_int8 : require
4
5#include "types.glsl"
6
7// Each iqs value maps to a 32-bit integer
8
9#if defined(DATA_A_Q2_0)
10void block_a_to_shmem(const uint buf_ib, const uint ib, const uint iqs) {
11    const uint block_idx = ib / 2;
12    const uint byte_idx = (ib & 1u) * 8u + iqs;
13    const uint bits = uint(data_a[block_idx].qs[byte_idx]);
14    buf_a[buf_ib].qs[iqs] = pack32(i8vec4(
15        int8_t(bits & 3u),
16        int8_t((bits >> 2u) & 3u),
17        int8_t((bits >> 4u) & 3u),
18        int8_t(bits >> 6u)));
19
20    if (iqs == 0) {
21        buf_a[buf_ib].dm = FLOAT_TYPE(data_a[block_idx].d);
22    }
23}
24
25void block_a_to_registers(const uint reg_ib, const uint buf_ib) {
26    cache_a[reg_ib].dm = buf_a[buf_ib].dm;
27
28    [[unroll]] for (uint iqs = 0; iqs < 8; ++iqs) {
29        cache_a[reg_ib].qs[iqs] = buf_a[buf_ib].qs[iqs];
30    }
31}
32
33ACC_TYPE mmq_dot_product(const uint ib_a) {
34    int32_t q_sum = 0;
35    [[unroll]] for (uint iqs = 0; iqs < 8; ++iqs) {
36        q_sum += dotPacked4x8EXT(cache_a[ib_a].qs[iqs], cache_b.qs[iqs]);
37    }
38
39    return ACC_TYPE(float(cache_a[ib_a].dm) * (float(q_sum) * float(cache_b.ds.x) - float(cache_b.ds.y)));
40}
41#endif
42
43#if defined(DATA_A_Q4_0) || defined(DATA_A_Q4_1)
44// 2-byte loads for Q4_0 blocks (18 bytes)
45// 4-byte loads for Q4_1 blocks (20 bytes)
46void block_a_to_shmem(const uint buf_ib, const uint ib, const uint iqs) {
47#ifdef DATA_A_Q4_0
48    buf_a[buf_ib].qs[iqs] = pack32(u16vec2(data_a_packed16[ib].qs[iqs * 2],
49                                           data_a_packed16[ib].qs[iqs * 2 + 1]));
50
51    if (iqs == 0) {
52        buf_a[buf_ib].dm = FLOAT_TYPE(data_a_packed16[ib].d);
53    }
54#else // DATA_A_Q4_1
55    buf_a[buf_ib].qs[iqs] = data_a_packed32[ib].qs[iqs];
56
57    if (iqs == 0) {
58        buf_a[buf_ib].dm = FLOAT_TYPEV2(data_a_packed32[ib].dm);
59    }
60#endif
61}
62
63void block_a_to_registers(const uint reg_ib, const uint buf_ib) {
64    cache_a[reg_ib].dm = buf_a[buf_ib].dm;
65
66    [[unroll]] for (uint iqs = 0; iqs < 4; iqs++) {
67        cache_a[reg_ib].qs[iqs] = buf_a[buf_ib].qs[iqs];
68    }
69}
70
71ACC_TYPE mmq_dot_product(const uint ib_a) {
72    int32_t q_sum = 0;
73    [[unroll]] for (uint iqs = 0; iqs < 4; iqs++) {
74        const uint32_t vui = cache_a[ib_a].qs[iqs];
75        const i32vec2 qs_a = i32vec2( vui       & 0x0F0F0F0F,
76                                     (vui >> 4) & 0x0F0F0F0F);
77
78        const int32_t qs_b0 = cache_b.qs[iqs];
79        const int32_t qs_b1 = cache_b.qs[iqs + 4];
80
81        q_sum += dotPacked4x8EXT(qs_a.x, qs_b0);
82        q_sum += dotPacked4x8EXT(qs_a.y, qs_b1);
83    }
84
85#ifdef DATA_A_Q4_0
86    return ACC_TYPE(float(cache_a[ib_a].dm) * (float(q_sum) * float(cache_b.ds.x) - 8.0 * float(cache_b.ds.y)));
87#else // DATA_A_Q4_1
88    return ACC_TYPE(float(q_sum) * float(cache_a[ib_a].dm.x) * float(cache_b.ds.x) + float(cache_a[ib_a].dm.y) * float(cache_b.ds.y));
89#endif
90}
91#endif
92
93#if defined(DATA_A_Q5_0) || defined(DATA_A_Q5_1)
94// 2-byte loads for Q5_0 blocks (22 bytes)
95// 4-byte loads for Q5_1 blocks (24 bytes)
96void block_a_to_shmem(const uint buf_ib, const uint ib, const uint iqs) {
97#ifdef DATA_A_Q5_0
98    buf_a[buf_ib].qs[iqs] = pack32(u16vec2(data_a_packed16[ib].qs[iqs * 2],
99                                           data_a_packed16[ib].qs[iqs * 2 + 1]));
100
101    if (iqs == 0) {
102        buf_a[buf_ib].dm = FLOAT_TYPE(data_a_packed16[ib].d);
103        buf_a[buf_ib].qh = pack32(u16vec2(data_a_packed16[ib].qh[0], data_a_packed16[ib].qh[1]));
104    }
105#else // DATA_A_Q5_1
106    buf_a[buf_ib].qs[iqs] = data_a_packed32[ib].qs[iqs];
107
108    if (iqs == 0) {
109        buf_a[buf_ib].dm = FLOAT_TYPEV2(data_a_packed32[ib].dm);
110        buf_a[buf_ib].qh = data_a_packed32[ib].qh;
111    }
112#endif
113}
114
115void block_a_to_registers(const uint reg_ib, const uint buf_ib) {
116    cache_a[reg_ib].dm = buf_a[buf_ib].dm;
117    cache_a[reg_ib].qh = buf_a[buf_ib].qh;
118
119    [[unroll]] for (uint iqs = 0; iqs < 4; iqs++) {
120        cache_a[reg_ib].qs[iqs] = buf_a[buf_ib].qs[iqs];
121    }
122}
123
124ACC_TYPE mmq_dot_product(const uint ib_a) {
125    int32_t q_sum = 0;
126    [[unroll]] for (uint iqs = 0; iqs < 4; iqs++) {
127        const uint32_t vui = cache_a[ib_a].qs[iqs];
128        const int32_t qh = int32_t(cache_a[ib_a].qh >> (4 * iqs));
129        const int32_t qs_a0 = int32_t(vui & 0x0F0F0F0F)
130                         | ((qh & 0xF) * 0x02040810) & 0x10101010; // (0,1,2,3) -> (4,12,20,28)
131        const int32_t qs_a1 = int32_t((vui >> 4) & 0x0F0F0F0F)
132                         | (((qh >> 16) & 0xF) * 0x02040810) & 0x10101010; // (16,17,18,19) -> (4,12,20,28)
133
134        const int32_t qs_b0 = cache_b.qs[iqs];
135        const int32_t qs_b1 = cache_b.qs[iqs + 4];
136
137        q_sum += dotPacked4x8EXT(qs_a0, qs_b0);
138        q_sum += dotPacked4x8EXT(qs_a1, qs_b1);
139    }
140
141#ifdef DATA_A_Q5_0
142    return ACC_TYPE(float(cache_a[ib_a].dm) * (float(q_sum) * float(cache_b.ds.x) - 16.0 * float(cache_b.ds.y)));
143#else // DATA_A_Q5_1
144    return ACC_TYPE(float(q_sum) * float(cache_a[ib_a].dm.x) * float(cache_b.ds.x) + float(cache_a[ib_a].dm.y) * float(cache_b.ds.y));
145#endif
146}
147#endif
148
149#if defined(DATA_A_Q8_0)
150// 2-byte loads for Q8_0 blocks (34 bytes)
151void block_a_to_shmem(const uint buf_ib, const uint ib, const uint iqs) {
152    buf_a[buf_ib].qs[iqs] = pack32(i16vec2(data_a_packed16[ib].qs[iqs * 2],
153                                           data_a_packed16[ib].qs[iqs * 2 + 1]));
154
155    if (iqs == 0) {
156        buf_a[buf_ib].dm = FLOAT_TYPE(data_a_packed16[ib].d);
157    }
158}
159
160void block_a_to_registers(const uint reg_ib, const uint buf_ib) {
161    cache_a[reg_ib].dm = buf_a[buf_ib].dm;
162
163    [[unroll]] for (uint iqs = 0; iqs < 8; iqs++) {
164        cache_a[reg_ib].qs[iqs] = buf_a[buf_ib].qs[iqs];
165    }
166}
167
168ACC_TYPE mmq_dot_product(const uint ib_a) {
169    int32_t q_sum = 0;
170    [[unroll]] for (uint iqs = 0; iqs < 8; iqs++) {
171        const int32_t qs_a = cache_a[ib_a].qs[iqs];
172        const int32_t qs_b = cache_b.qs[iqs];
173
174        q_sum += dotPacked4x8EXT(qs_a, qs_b);
175    }
176
177    return ACC_TYPE(float(q_sum) * float(cache_a[ib_a].dm) * float(cache_b.ds.x));
178}
179#endif
180
181#if defined(DATA_A_MXFP4)
182// 1-byte loads for mxfp4 blocks (17 bytes)
183void block_a_to_shmem(const uint buf_ib, const uint ib, const uint iqs) {
184    const uint32_t qs = pack32(u8vec4(data_a[ib].qs[iqs * 4    ],
185                                      data_a[ib].qs[iqs * 4 + 1],
186                                      data_a[ib].qs[iqs * 4 + 2],
187                                      data_a[ib].qs[iqs * 4 + 3]));
188
189    const u8vec4 i_a0 = unpack8( qs       & 0x0F0F0F0F);
190    const u8vec4 i_a1 = unpack8((qs >> 4) & 0x0F0F0F0F);
191
192    buf_a[buf_ib].qs[iqs    ] = pack32(i8vec4(kvalues_mxfp4[i_a0.x], kvalues_mxfp4[i_a0.y], kvalues_mxfp4[i_a0.z], kvalues_mxfp4[i_a0.w]));
193    buf_a[buf_ib].qs[iqs + 4] = pack32(i8vec4(kvalues_mxfp4[i_a1.x], kvalues_mxfp4[i_a1.y], kvalues_mxfp4[i_a1.z], kvalues_mxfp4[i_a1.w]));
194
195    if (iqs == 0) {
196        buf_a[buf_ib].d = FLOAT_TYPE(e8m0_to_fp32(data_a[ib].e) * 0.5);
197    }
198}
199
200void block_a_to_registers(const uint reg_ib, const uint buf_ib) {
201    cache_a[reg_ib].d = buf_a[buf_ib].d;
202
203    [[unroll]] for (uint iqs = 0; iqs < 8; iqs++) {
204        cache_a[reg_ib].qs[iqs] = buf_a[buf_ib].qs[iqs];
205    }
206}
207
208ACC_TYPE mmq_dot_product(const uint ib_a) {
209    int32_t q_sum = 0;
210    [[unroll]] for (uint iqs = 0; iqs < 8; iqs++) {
211        const int32_t qs_a = cache_a[ib_a].qs[iqs];
212
213        q_sum += dotPacked4x8EXT(qs_a, cache_b.qs[iqs]);
214    }
215
216    return ACC_TYPE(float(cache_a[ib_a].d) * float(cache_b.ds.x) * float(q_sum));
217}
218#endif
219
220// For k-quants, ib and iqs still assume 32-wide blocks, but k-quants are 256-wide
221// iqs still refers to a 32-bit integer, meaning 0..7 for 32-wide quants
222#if defined(DATA_A_Q2_K)
223// 4-byte loads for Q2_K blocks (84 bytes)
224void block_a_to_shmem(const uint buf_ib, const uint ib, const uint iqs) {
225    const uint ib_k = ib / 8;
226    const uint iqs_k = (ib % 8) * 8 + iqs * QUANT_R_MMQ;
227
228    const uint qs_idx = (iqs_k / 32) * 8 + (iqs_k % 8);
229    const uint qs_shift = ((iqs_k % 32) / 8) * 2;
230
231    // Repack 4x4 quants into one int
232    const uint32_t vals0 = (data_a_packed32[ib_k].qs[qs_idx    ] >> qs_shift) & 0x03030303;
233    const uint32_t vals1 = (data_a_packed32[ib_k].qs[qs_idx + 1] >> qs_shift) & 0x03030303;
234    const uint32_t vals2 = (data_a_packed32[ib_k].qs[qs_idx + 2] >> qs_shift) & 0x03030303;
235    const uint32_t vals3 = (data_a_packed32[ib_k].qs[qs_idx + 3] >> qs_shift) & 0x03030303;
236
237    buf_a[buf_ib].qs[iqs] = vals0 | (vals1 << 2) | (vals2 << 4) | (vals3 << 6);
238
239    if (iqs == 0) {
240        buf_a[buf_ib].dm = FLOAT_TYPEV2(data_a_packed32[ib_k].dm);
241        buf_a[buf_ib].scales = unpack8(uint32_t(data_a_packed16[ib_k].scales[iqs_k / 8])).xy; // vec4 used due to #12147
242    }
243}
244
245void block_a_to_registers(const uint reg_ib, const uint buf_ib) {
246    cache_a[reg_ib].dm = buf_a[buf_ib].dm;
247    cache_a[reg_ib].scales = buf_a[buf_ib].scales;
248
249    [[unroll]] for (uint iqs = 0; iqs < 2; iqs++) {
250        cache_a[reg_ib].qs[iqs] = buf_a[buf_ib].qs[iqs];
251    }
252}
253
254ACC_TYPE mmq_dot_product(const uint ib_a) {
255    int32_t sum_d = 0;
256    int32_t sum_m = 0;
257
258    [[unroll]] for (uint iqs = 0; iqs < 8; iqs++) {
259        const uint8_t scale = cache_a[ib_a].scales[iqs / 4];
260        const int32_t scale_m = int32_t(scale >> 4) * 0x01010101; // Duplicate 8-bit value across 32-bits.
261        const int32_t qs_a = int32_t((cache_a[ib_a].qs[iqs / 4] >> ((iqs % 4) * 2)) & 0x03030303);
262
263        sum_d += dotPacked4x8EXT(qs_a, cache_b.qs[iqs]) * (scale & 0xF);
264        sum_m += dotPacked4x8EXT(scale_m, cache_b.qs[iqs]);
265    }
266
267    return ACC_TYPE(float(cache_b.ds.x) * (float(cache_a[ib_a].dm.x) * float(sum_d) - float(cache_a[ib_a].dm.y) * float(sum_m)));
268}
269#endif
270
271#if defined(DATA_A_Q3_K)
272// 2-byte loads for Q3_K blocks (110 bytes)
273void block_a_to_shmem(const uint buf_ib, const uint ib, const uint iqs) {
274    const uint ib_k = ib / 8;
275    const uint hm_idx = iqs * QUANT_R_MMQ;
276    const uint iqs_k = (ib % 8) * 8 + hm_idx;
277
278    const uint qs_idx = (iqs_k / 32) * 8 + (iqs_k % 8);
279    const uint qs_shift = ((iqs_k % 32) / 8) * 2;
280    const uint hm_shift = iqs_k / 8;
281
282    // Repack 2x4 quants into one int
283    // Add the 3rd bit instead of subtracting it to allow packing the quants
284    // vec4 for unpack8 used due to #12147
285    const i8vec2 vals00 = unpack8(int32_t(int16_t((data_a_packed16[ib_k].qs[qs_idx * 2        ] >> qs_shift) & uint16_t(0x0303)))).xy |
286                          unpack8(int32_t(int16_t(((data_a_packed16[ib_k].hmask[hm_idx * 2    ] >> hm_shift) & uint16_t(0x0101))) << 2)).xy;
287    const i8vec2 vals01 = unpack8(int32_t(int16_t((data_a_packed16[ib_k].qs[qs_idx * 2 + 1    ] >> qs_shift) & uint16_t(0x0303)))).xy |
288                          unpack8(int32_t(int16_t(((data_a_packed16[ib_k].hmask[hm_idx * 2 + 1] >> hm_shift) & uint16_t(0x0101))) << 2)).xy;
289    const i8vec2 vals10 = unpack8(int32_t(int16_t((data_a_packed16[ib_k].qs[qs_idx * 2 + 2    ] >> qs_shift) & uint16_t(0x0303)))).xy |
290                          unpack8(int32_t(int16_t(((data_a_packed16[ib_k].hmask[hm_idx * 2 + 2] >> hm_shift) & uint16_t(0x0101))) << 2)).xy;
291    const i8vec2 vals11 = unpack8(int32_t(int16_t((data_a_packed16[ib_k].qs[qs_idx * 2 + 3    ] >> qs_shift) & uint16_t(0x0303)))).xy |
292                          unpack8(int32_t(int16_t(((data_a_packed16[ib_k].hmask[hm_idx * 2 + 3] >> hm_shift) & uint16_t(0x0101))) << 2)).xy;
293    buf_a[buf_ib].qs[iqs] = pack32(u8vec4(vals00.x, vals00.y, vals01.x, vals01.y)) |
294                           (pack32(u8vec4(vals10.x, vals10.y, vals11.x, vals11.y)) << 4);
295
296    if (iqs == 0) {
297        const uint is = iqs_k / 4;
298        const i8vec2 scales = i8vec2(unpack8(uint32_t(((data_a_packed16[ib_k].scales[(is % 8      ) / 2] >> (4 * (is / 8))) & 0x0F0F) |
299                                                     (((data_a_packed16[ib_k].scales[(8 + (is % 4)) / 2] >> (2 * (is / 4))) & 0x0303) << 4))).xy); // vec4 used due to #12147
300
301        buf_a[buf_ib].d_scales = FLOAT_TYPEV2(float(data_a_packed16[ib_k].d) * vec2(scales - 32));
302    }
303}
304
305void block_a_to_registers(const uint reg_ib, const uint buf_ib) {
306    cache_a[reg_ib].d_scales = buf_a[buf_ib].d_scales;
307
308    [[unroll]] for (uint iqs = 0; iqs < 4; iqs++) {
309        cache_a[reg_ib].qs[iqs] = buf_a[buf_ib].qs[iqs];
310    }
311}
312
313ACC_TYPE mmq_dot_product(const uint ib_a) {
314    float result = 0.0;
315    int32_t q_sum = 0;
316
317    [[unroll]] for (uint iqs = 0; iqs < 4; iqs++) {
318        // Subtract 4 from the quants to correct the 3rd bit offset
319        const int32_t qs_a = pack32(unpack8(int32_t((cache_a[ib_a].qs[iqs / 2] >> ((iqs % 2) * 4)) & 0x0F0F0F0F)) - int8_t(4));
320
321        q_sum += dotPacked4x8EXT(qs_a, cache_b.qs[iqs]);
322    }
323    result += float(cache_a[ib_a].d_scales[0]) * float(q_sum);
324    q_sum = 0;
325
326    [[unroll]] for (uint iqs = 4; iqs < 8; iqs++) {
327        const int32_t qs_a = pack32(unpack8(int32_t((cache_a[ib_a].qs[iqs / 2] >> ((iqs % 2) * 4)) & 0x0F0F0F0F)) - int8_t(4));
328
329        q_sum += dotPacked4x8EXT(qs_a, cache_b.qs[iqs]);
330    }
331    result += float(cache_a[ib_a].d_scales[1]) * float(q_sum);
332
333    return ACC_TYPE(float(cache_b.ds.x) * result);
334}
335#endif
336
337#if defined(DATA_A_Q4_K) || defined(DATA_A_Q5_K)
338// 4-byte loads for Q4_K blocks (144 bytes) and Q5_K blocks (176 bytes)
339void block_a_to_shmem(const uint buf_ib, const uint ib, const uint iqs) {
340    const uint ib_k = ib / 8;
341    const uint iqs_k = (ib % 8) * 8 + iqs * QUANT_R_MMQ;
342
343    const uint qs_idx = (iqs_k / 16) * 8 + (iqs_k % 8);
344    const uint qs_shift = ((iqs_k % 16) / 8) * 4;
345
346    // Repack 2x4 quants into one int
347#if defined(DATA_A_Q4_K)
348    const uint32_t vals0 = (data_a_packed32[ib_k].qs[qs_idx    ] >> qs_shift) & 0x0F0F0F0F;
349    const uint32_t vals1 = (data_a_packed32[ib_k].qs[qs_idx + 1] >> qs_shift) & 0x0F0F0F0F;
350
351    buf_a[buf_ib].qs[iqs] = vals0 | (vals1 << 4);
352#else // defined(DATA_A_Q5_K)
353    const uint qh_idx = iqs * QUANT_R_MMQ;
354    const uint qh_shift = iqs_k / 8;
355
356    buf_a[buf_ib].qs[iqs] = int32_t(((data_a_packed32[ib_k].qs[qs_idx] >> qs_shift) & 0x0F0F0F0F) |
357                                   (((data_a_packed32[ib_k].qh[qh_idx] >> qh_shift) & 0x01010101) << 4));
358#endif
359
360    if (iqs == 0) {
361        // Scale index
362        const uint is = iqs_k / 8;
363        u8vec2 scale_dm;
364        if (is < 4) {
365            scale_dm = u8vec2(data_a[ib_k].scales[is] & 0x3F, data_a[ib_k].scales[is + 4] & 0x3F);
366        } else {
367            scale_dm = u8vec2((data_a[ib_k].scales[is+4] & 0xF) | ((data_a[ib_k].scales[is-4] & 0xC0) >> 2),
368                              (data_a[ib_k].scales[is+4] >>  4) | ((data_a[ib_k].scales[is  ] & 0xC0) >> 2));
369        }
370
371        buf_a[buf_ib].dm = FLOAT_TYPEV2(vec2(data_a_packed32[ib_k].dm) * vec2(scale_dm));
372    }
373}
374
375void block_a_to_registers(const uint reg_ib, const uint buf_ib) {
376    cache_a[reg_ib].dm = buf_a[buf_ib].dm;
377
378    [[unroll]] for (uint iqs = 0; iqs < 8 / QUANT_R_MMQ; iqs++) {
379        cache_a[reg_ib].qs[iqs] = buf_a[buf_ib].qs[iqs];
380    }
381}
382
383ACC_TYPE mmq_dot_product(const uint ib_a) {
384    int32_t q_sum = 0;
385
386    [[unroll]] for (uint iqs = 0; iqs < 8; iqs++) {
387#if defined(DATA_A_Q4_K)
388        const int32_t qs_a = int32_t((cache_a[ib_a].qs[iqs / 2] >> ((iqs % 2) * 4)) & 0x0F0F0F0F);
389#else // defined(DATA_A_Q5_K)
390        const int32_t qs_a = cache_a[ib_a].qs[iqs];
391#endif
392
393        q_sum += dotPacked4x8EXT(qs_a, cache_b.qs[iqs]);
394    }
395
396    return ACC_TYPE(float(cache_b.ds.x) * float(cache_a[ib_a].dm.x) * float(q_sum) - float(cache_a[ib_a].dm.y) * float(cache_b.ds.y));
397}
398#endif
399
400#if defined(DATA_A_Q6_K)
401// 2-byte loads for Q6_K blocks (210 bytes)
402void block_a_to_shmem(const uint buf_ib, const uint ib, const uint iqs) {
403    const uint ib_k = ib / 8;
404    const uint iqs_k = (ib % 8) * 8 + iqs;
405
406    const uint ql_idx = (iqs_k / 32) * 16 + iqs_k % 16;
407    const uint ql_shift = ((iqs_k % 32) / 16) * 4;
408
409    const uint qh_idx = (iqs_k / 32) * 8 + iqs;
410    const uint qh_shift = ((iqs_k % 32) / 8) * 2;
411
412    const i8vec2 vals00 = (unpack8(int32_t((data_a_packed16[ib_k].ql[ql_idx * 2    ] >> ql_shift) & uint16_t(0x0F0F))).xy |
413                          unpack8(int32_t(((data_a_packed16[ib_k].qh[qh_idx * 2    ] >> qh_shift) & uint16_t(0x0303)) << 4)).xy) - int8_t(32);
414    const i8vec2 vals01 = (unpack8(int32_t((data_a_packed16[ib_k].ql[ql_idx * 2 + 1] >> ql_shift) & uint16_t(0x0F0F))).xy |
415                          unpack8(int32_t(((data_a_packed16[ib_k].qh[qh_idx * 2 + 1] >> qh_shift) & uint16_t(0x0303)) << 4)).xy) - int8_t(32);
416    buf_a[buf_ib].qs[iqs] = pack32(i8vec4(vals00.x, vals00.y, vals01.x, vals01.y));
417
418    if (iqs == 0) {
419        const uint is = iqs_k / 4;
420        const i8vec2 scales = unpack8(int32_t(data_a_packed16[ib_k].scales[is / 2])).xy;
421
422        buf_a[buf_ib].d_scales = FLOAT_TYPEV2(float(data_a_packed16[ib_k].d) * vec2(scales));
423    }
424}
425
426void block_a_to_registers(const uint reg_ib, const uint buf_ib) {
427    cache_a[reg_ib].d_scales = buf_a[buf_ib].d_scales;
428
429    [[unroll]] for (uint iqs = 0; iqs < 8; iqs++) {
430        cache_a[reg_ib].qs[iqs] = buf_a[buf_ib].qs[iqs];
431    }
432}
433
434ACC_TYPE mmq_dot_product(const uint ib_a) {
435    float result = 0.0;
436    int32_t q_sum = 0;
437
438    [[unroll]] for (uint iqs = 0; iqs < 4; iqs++) {
439        const int32_t qs_a = cache_a[ib_a].qs[iqs];
440
441        q_sum += dotPacked4x8EXT(qs_a, cache_b.qs[iqs]);
442    }
443    result += float(cache_a[ib_a].d_scales[0]) * float(q_sum);
444    q_sum = 0;
445
446    [[unroll]] for (uint iqs = 4; iqs < 8; iqs++) {
447        const int32_t qs_a = cache_a[ib_a].qs[iqs];
448
449        q_sum += dotPacked4x8EXT(qs_a, cache_b.qs[iqs]);
450    }
451    result += float(cache_a[ib_a].d_scales[1]) * float(q_sum);
452
453    return ACC_TYPE(float(cache_b.ds.x) * result);
454}
455#endif
456
457void block_b_to_shmem(const uint buf_ib, const uint ib, const uint iqs, const bool is_in_bounds) {
458    if (is_in_bounds) {
459        const uint ib_outer = ib / 4;
460        const uint ib_inner = ib % 4;
461
462        if (iqs == 0) {
463            buf_b[buf_ib].ds = FLOAT_TYPEV2(data_b[ib_outer].ds[ib_inner]);
464        }
465
466        const ivec4 values = data_b[ib_outer].qs[ib_inner * 2 + iqs];
467        buf_b[buf_ib].qs[iqs * 4    ] = values.x;
468        buf_b[buf_ib].qs[iqs * 4 + 1] = values.y;
469        buf_b[buf_ib].qs[iqs * 4 + 2] = values.z;
470        buf_b[buf_ib].qs[iqs * 4 + 3] = values.w;
471    } else {
472        if (iqs == 0) {
473            buf_b[buf_ib].ds = FLOAT_TYPEV2(0.0f);
474        }
475
476        buf_b[buf_ib].qs[iqs * 4    ] = 0;
477        buf_b[buf_ib].qs[iqs * 4 + 1] = 0;
478        buf_b[buf_ib].qs[iqs * 4 + 2] = 0;
479        buf_b[buf_ib].qs[iqs * 4 + 3] = 0;
480    }
481}
482
483void block_b_to_registers(const uint ib) {
484    cache_b.ds = buf_b[ib].ds;
485    [[unroll]] for (uint iqs = 0; iqs < BK / 4; iqs++) {
486        cache_b.qs[iqs] = buf_b[ib].qs[iqs];
487    }
488}