char-fork/llama.cpp

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

Anjielonvulkan: add TQ1_0 support (mm, mat-vec, mat-vec-id, dequant, get_rows) (#27765)8fe90e1fb

char/aliceai-t5-moe
30.0 KiB756 linesraw
1#if !defined(DATA_A_F32) && !defined(DATA_A_F16)
2#extension GL_EXT_shader_explicit_arithmetic_types_int8 : require
3#endif
4
5#include "types.glsl"
6
7#if defined(DATA_A_F32)
8FLOAT_TYPE dequantize1(uint ib, uint iqs, uint a_offset) {
9    return data_a[a_offset + ib];
10}
11vec2 dequantize(uint ib, uint iqs, uint a_offset) {
12    return vec2(data_a[a_offset + ib], data_a[a_offset + ib + 1]);
13}
14vec4 dequantize4(uint ib, uint iqs, uint a_offset) {
15    return vec4(data_a[a_offset + ib    ], data_a[a_offset + ib + 1],
16                data_a[a_offset + ib + 2], data_a[a_offset + ib + 3]);
17}
18vec4 dequantize4_2aligned(uint ib, uint iqs, uint a_offset) {
19    return vec4(data_a[a_offset + ib    ], data_a[a_offset + ib + 1],
20                data_a[a_offset + ib + 2], data_a[a_offset + ib + 3]);
21}
22
23#endif
24
25#if defined(DATA_A_F16)
26FLOAT_TYPE dequantize1(uint ib, uint iqs, uint a_offset) {
27    return data_a[a_offset + ib];
28}
29vec2 dequantize(uint ib, uint iqs, uint a_offset) {
30    return vec2(data_a[a_offset + ib], data_a[a_offset + ib + 1]);
31}
32vec4 dequantize4(uint ib, uint iqs, uint a_offset) {
33    return vec4(data_a[a_offset + ib    ], data_a[a_offset + ib + 1],
34                data_a[a_offset + ib + 2], data_a[a_offset + ib + 3]);
35}
36vec4 dequantize4_2aligned(uint ib, uint iqs, uint a_offset) {
37    const vec2 a = data_a_packed32[(a_offset + ib)/2];
38    const vec2 b = data_a_packed32[(a_offset + ib)/2 + 1];
39    return vec4(a, b);
40}
41#endif
42
43#if defined(DATA_A_BF16)
44FLOAT_TYPE dequantize1(uint ib, uint iqs, uint a_offset) {
45    return bf16_to_fp32(data_a[a_offset + ib]);
46}
47vec2 dequantize(uint ib, uint iqs, uint a_offset) {
48    return vec2(bf16_to_fp32(data_a[a_offset + ib]), bf16_to_fp32(data_a[a_offset + ib + 1]));
49}
50vec4 dequantize4(uint ib, uint iqs, uint a_offset) {
51    return vec4(bf16_to_fp32(data_a[a_offset + ib    ]), bf16_to_fp32(data_a[a_offset + ib + 1]),
52                bf16_to_fp32(data_a[a_offset + ib + 2]), bf16_to_fp32(data_a[a_offset + ib + 3]));
53}
54vec4 dequantize4_2aligned(uint ib, uint iqs, uint a_offset) {
55    const uint a = data_a_packed32[(a_offset + ib)/2];
56    const uint b = data_a_packed32[(a_offset + ib)/2 + 1];
57    return vec4(uintBitsToFloat((a & 0x0000ffff) << 16),
58                uintBitsToFloat( a & 0xffff0000),
59                uintBitsToFloat((b & 0x0000ffff) << 16),
60                uintBitsToFloat( b & 0xffff0000));
61}
62#endif
63
64#if defined(DATA_A_Q4_0)
65vec2 dequantize(uint ib, uint iqs, uint a_offset) {
66    const uint vui = uint(data_a[a_offset + ib].qs[iqs]);
67    return (vec2(vui & 0xF, vui >> 4) - 8.0f);
68}
69vec4 dequantize4(uint ib, uint iqs, uint a_offset) {
70    const uint vui = uint(data_a_packed16[a_offset + ib].qs[iqs/2]);
71    return (vec4(vui & 0xF, (vui >> 4) & 0xF, (vui >> 8) & 0xF, vui >> 12) - 8.0f);
72}
73#endif
74
75#if defined(DATA_A_Q4_1)
76vec2 dequantize(uint ib, uint iqs, uint a_offset) {
77    const uint vui = uint(data_a[a_offset + ib].qs[iqs]);
78    return vec2(vui & 0xF, vui >> 4);
79}
80vec4 dequantize4(uint ib, uint iqs, uint a_offset) {
81    const uint vui = uint(data_a_packed16[a_offset + ib].qs[iqs/2]);
82    return vec4(vui & 0xF, (vui >> 4) & 0xF, (vui >> 8) & 0xF, vui >> 12);
83}
84#endif
85
86#if defined(DATA_A_Q5_0)
87vec2 dequantize(uint ib, uint iqs, uint a_offset) {
88    const uint uint_qh = uint(data_a[a_offset + ib].qh[1]) << 16 | data_a[a_offset + ib].qh[0];
89    const ivec2 qh = ivec2(((uint_qh >> iqs) << 4) & 0x10, (uint_qh >> (iqs + 12)) & 0x10);
90    const uint vui = uint(data_a[a_offset + ib].qs[iqs]);
91    return (vec2((vui & 0xF) | qh.x, (vui >> 4) | qh.y) - 16.0f);
92}
93vec4 dequantize4(uint ib, uint iqs, uint a_offset) {
94    const uint uint_qh = uint(data_a_packed16[a_offset + ib].qh[1]) << 16 | data_a_packed16[a_offset + ib].qh[0];
95    const ivec2 qh0 = ivec2(((uint_qh >> iqs) << 4) & 0x10, (uint_qh >> (iqs + 12)) & 0x10);
96    const ivec2 qh1 = ivec2(((uint_qh >> (iqs + 1)) << 4) & 0x10, (uint_qh >> (iqs + 13)) & 0x10);
97    const uint vui = uint(data_a_packed16[a_offset + ib].qs[iqs/2]);
98    return (vec4((vui & 0xF) | qh0.x, ((vui >> 4) & 0xF) | qh0.y, ((vui >> 8) & 0xF) | qh1.x, (vui >> 12) | qh1.y) - 16.0f);
99}
100#endif
101
102#if defined(DATA_A_Q5_1)
103vec2 dequantize(uint ib, uint iqs, uint a_offset) {
104    const uint uint_qh = data_a[a_offset + ib].qh;
105    const ivec2 qh = ivec2(((uint_qh >> iqs) << 4) & 0x10, (uint_qh >> (iqs + 12)) & 0x10);
106    const uint vui = uint(data_a[a_offset + ib].qs[iqs]);
107    return vec2((vui & 0xF) | qh.x, (vui >> 4) | qh.y);
108}
109vec4 dequantize4(uint ib, uint iqs, uint a_offset) {
110    const uint uint_qh = data_a_packed16[a_offset + ib].qh;
111    const ivec2 qh0 = ivec2(((uint_qh >> iqs) << 4) & 0x10, (uint_qh >> (iqs + 12)) & 0x10);
112    const ivec2 qh1 = ivec2(((uint_qh >> (iqs + 1)) << 4) & 0x10, (uint_qh >> (iqs + 13)) & 0x10);
113    const uint vui = uint(data_a_packed16[a_offset + ib].qs[iqs/2]);
114    return vec4((vui & 0xF) | qh0.x, ((vui >> 4) & 0xF) | qh0.y, ((vui >> 8) & 0xF) | qh1.x, (vui >> 12) | qh1.y);
115}
116#endif
117
118#if defined(DATA_A_Q8_0)
119vec2 dequantize(uint ib, uint iqs, uint a_offset) {
120    return vec2(int(data_a[a_offset + ib].qs[iqs]), int(data_a[a_offset + ib].qs[iqs + 1]));
121}
122vec4 dequantize4(uint ib, uint iqs, uint a_offset) {
123    const i8vec2 v0 = unpack8(int32_t(data_a_packed16[a_offset + ib].qs[iqs/2])).xy; // vec4 used due to #12147
124    const i8vec2 v1 = unpack8(int32_t(data_a_packed16[a_offset + ib].qs[iqs/2 + 1])).xy;
125    return vec4(v0.x, v0.y, v1.x, v1.y);
126}
127#endif
128
129#if defined(DATA_A_Q1_0)
130vec2 dequantize(uint ib, uint iqs, uint a_offset) {
131    const uint bits = uint(data_a[a_offset + ib].qs[iqs / 8u]) >> (iqs % 8u);
132    return vec2(
133        (bits & 1u) != 0u ? 1.0f : -1.0f,
134        (bits & 2u) != 0u ? 1.0f : -1.0f);
135}
136vec4 dequantize4(uint ib, uint iqs, uint a_offset) {
137    const uint bits = uint(data_a[a_offset + ib].qs[iqs / 8u]) >> (iqs % 8u);
138    return vec4(
139        (bits & 1u) != 0u ? 1.0f : -1.0f,
140        (bits & 2u) != 0u ? 1.0f : -1.0f,
141        (bits & 4u) != 0u ? 1.0f : -1.0f,
142        (bits & 8u) != 0u ? 1.0f : -1.0f);
143}
144#endif
145
146#if defined(DATA_A_Q2_0)
147vec2 dequantize(uint ib, uint iqs, uint a_offset) {
148    const uint bits = uint(data_a[a_offset + ib].qs[iqs / 4u]) >> (2u * (iqs % 4u));
149    return vec2(bits & 3u, (bits >> 2u) & 3u) - 1.0f;
150}
151vec4 dequantize4(uint ib, uint iqs, uint a_offset) {
152    const uint bits = uint(data_a[a_offset + ib].qs[iqs / 4u]);
153    return vec4(bits & 3u, (bits >> 2u) & 3u, (bits >> 4u) & 3u, bits >> 6u) - 1.0f;
154}
155#endif
156
157#if defined(DATA_A_IQ1_S)
158vec2 dequantize(uint ib, uint iqs, uint a_offset) {
159    const uint ib32 = iqs / 32;
160    const uint ib8 = iqs / 8;
161    const int i8 = int(iqs % 8);
162    const uint qh = data_a[a_offset + ib].qh[ib32];
163    const uint qs = data_a[a_offset + ib].qs[ib8];
164    const float dl = float(2 * bitfieldExtract(qh, 12, 3) + 1);
165    const float delta = ((qh & 0x8000) != 0) ? -IQ1S_DELTA : IQ1S_DELTA;
166    const uint idxhi = bitfieldExtract(qh, 3 * int(ib8 & 3), 3);
167    const int16_t grid = int16_t(iq1s_grid[qs | (idxhi << 8)]);
168    // Signed bitfield extract.
169    const ivec2 gvec = ivec2(
170      bitfieldExtract(grid, 2 * (i8), 2),
171      bitfieldExtract(grid, 2 * (i8 + 1), 2)
172    );
173    return dl * (vec2(gvec) + delta);
174}
175vec4 dequantize4(uint ib, uint iqs, uint a_offset) {
176    const uint ib32 = iqs / 32;
177    const uint ib8 = iqs / 8;
178    const int i8 = int(iqs % 8);
179    const uint qh = data_a[a_offset + ib].qh[ib32];
180    const uint qs = data_a[a_offset + ib].qs[ib8];
181    const float dl = 2 * bitfieldExtract(qh, 12, 3) + 1;
182    const float delta = ((qh & 0x8000) != 0) ? -IQ1S_DELTA : IQ1S_DELTA;
183    const int16_t grid = int16_t(iq1s_grid[qs | (bitfieldExtract(qh, 3 * int(ib8 & 3), 3) << 8)]);
184    // Signed bitfield extract.
185    const ivec4 gvec = ivec4(
186      bitfieldExtract(grid, 2 * (i8), 2),
187      bitfieldExtract(grid, 2 * (i8 + 1), 2),
188      bitfieldExtract(grid, 2 * (i8 + 2), 2),
189      bitfieldExtract(grid, 2 * (i8 + 3), 2)
190    );
191    return dl * (vec4(gvec) + delta);
192}
193#endif
194
195#if defined(DATA_A_IQ1_M)
196vec2 dequantize(uint ib, uint iqs, uint a_offset) {
197    const uint ib8 = iqs / 8;
198    const uint ib16 = iqs / 16;
199    const int i8 = int(iqs % 8);
200    const uint sc = data_a[a_offset + ib].scales[iqs / 64];
201    const uint qs = data_a[a_offset + ib].qs[ib8];
202    const uint qh = data_a[a_offset + ib].qh[ib16] >> (4 * (ib8 & 1));
203    const float dl = 2 * bitfieldExtract(sc, 3 * int(ib16 & 3), 3) + 1;
204    const float delta = ((qh & 8) != 0) ? -IQ1M_DELTA : IQ1M_DELTA;
205    const int16_t grid = int16_t(iq1s_grid[qs | ((qh & 7) << 8)]);
206    // Signed bitfield extract.
207    const ivec2 gvec = ivec2(
208      bitfieldExtract(grid, 2 * (i8), 2),
209      bitfieldExtract(grid, 2 * (i8 + 1), 2)
210    );
211    return dl * (vec2(gvec) + delta);
212}
213vec4 dequantize4(uint ib, uint iqs, uint a_offset) {
214    const uint ib8 = iqs / 8;
215    const uint ib16 = iqs / 16;
216    const int i8 = int(iqs % 8);
217    const uint sc = data_a[a_offset + ib].scales[iqs / 64];
218    const uint qs = data_a[a_offset + ib].qs[ib8];
219    const uint qh = data_a[a_offset + ib].qh[ib16] >> (4 * (ib8 & 1));
220    const float dl = 2 * bitfieldExtract(sc, 3 * int(ib16 & 3), 3) + 1;
221    const float delta = ((qh & 8) != 0) ? -IQ1M_DELTA : IQ1M_DELTA;
222    const int16_t grid = int16_t(iq1s_grid[qs | ((qh & 7) << 8)]);
223    // Signed bitfield extract.
224    const ivec4 gvec = ivec4(
225      bitfieldExtract(grid, 2 * (i8), 2),
226      bitfieldExtract(grid, 2 * (i8 + 1), 2),
227      bitfieldExtract(grid, 2 * (i8 + 2), 2),
228      bitfieldExtract(grid, 2 * (i8 + 3), 2)
229    );
230    return dl * (vec4(gvec) + delta);
231}
232#endif
233
234#if defined(DATA_A_IQ2_XXS)
235vec2 dequantize(uint ib, uint iqs, uint a_offset) {
236    const uint ib32 = iqs / 32;
237    const uint ib8 = (iqs / 8) % 4;
238    const uint qs = data_a[a_offset + ib].qs[8 * ib32 + ib8];
239    // Scales are stored as packed 7+7+7+7+4 bits (4 sign tuples and 1 int4 scale)
240    const uint signs = pack32(u16vec2(data_a_packed16[a_offset + ib].qs[4 * ib32 + 2],
241        data_a_packed16[a_offset + ib].qs[4 * ib32 + 3]));
242    const float db = 0.25 * (0.5 + (signs >> 28));
243    const uint sign7 = bitfieldExtract(signs, 7 * int(ib8), 7);
244    // Add parity bit
245    const uint sign8 = sign7 | (bitCount(sign7) << 7);
246    const uint sign = sign8 >> (iqs % 8);
247    const u8vec4 grid = unpack8(iq2xxs_grid[qs][(iqs % 8) / 4] >> (8 * (iqs % 4)));
248    bool sign0 = (sign & 1) != 0;
249    bool sign1 = (sign & 2) != 0;
250    return db * vec2(
251        grid.x * (sign0 ? -1.0 : 1.0),
252        grid.y * (sign1 ? -1.0 : 1.0)
253    );
254}
255vec4 dequantize4(uint ib, uint iqs, uint a_offset) {
256    const uint ib32 = iqs / 32;
257    const uint ib8 = (iqs / 8) % 4;
258    const uint qs = data_a[a_offset + ib].qs[8 * ib32 + ib8];
259    // Scales are stored as packed 7+7+7+7+4 bits (4 sign tuples and 1 int4 scale)
260    const uint signs = pack32(u16vec2(data_a_packed16[a_offset + ib].qs[4 * ib32 + 2],
261        data_a_packed16[a_offset + ib].qs[4 * ib32 + 3]));
262    const float db = 0.25 * (0.5 + (signs >> 28));
263    const uint sign7 = bitfieldExtract(signs, 7 * int(ib8), 7);
264    // Add parity bit
265    const uint sign8 = sign7 | (bitCount(sign7) << 7);
266    const uint sign = sign8 >> (iqs % 8);
267    const u8vec4 grid = unpack8(iq2xxs_grid[qs][(iqs % 8) / 4] >> (8 * (iqs % 4)));
268    bool sign0 = (sign & 1) != 0;
269    bool sign1 = (sign & 2) != 0;
270    bool sign2 = (sign & 4) != 0;
271    bool sign3 = (sign & 8) != 0;
272    return db * vec4(
273        grid.x * (sign0 ? -1.0 : 1.0),
274        grid.y * (sign1 ? -1.0 : 1.0),
275        grid.z * (sign2 ? -1.0 : 1.0),
276        grid.w * (sign3 ? -1.0 : 1.0)
277    );
278}
279#endif
280
281#if defined(DATA_A_IQ2_XS)
282vec2 dequantize(uint ib, uint iqs, uint a_offset) {
283    const uint scale = (data_a[a_offset + ib].scales[iqs / 32] >> (4 * ((iqs / 16) & 1))) & 0xf;
284    const uint qs = data_a[a_offset + ib].qs[iqs / 8];
285    const float db = 0.25 * (0.5 + scale);
286    const uint sign7 = qs >> 9;
287    // Add parity bit
288    const uint sign8 = sign7 | (bitCount(sign7) << 7);
289    const uint sign = sign8 >> (iqs % 8);
290    const u8vec4 grid = unpack8(iq2xs_grid[qs & 511][(iqs % 8) / 4] >> (8 * (iqs % 4)));
291    bool sign0 = (sign & 1) != 0;
292    bool sign1 = (sign & 2) != 0;
293    return db * vec2(
294        grid.x * (sign0 ? -1.0 : 1.0),
295        grid.y * (sign1 ? -1.0 : 1.0)
296    );
297}
298vec4 dequantize4(uint ib, uint iqs, uint a_offset) {
299    const uint scale = (data_a[a_offset + ib].scales[iqs / 32] >> (4 * ((iqs / 16) & 1))) & 0xf;
300    const uint qs = data_a[a_offset + ib].qs[iqs / 8];
301    const float db = 0.25 * (0.5 + scale);
302    const uint sign7 = qs >> 9;
303    // Add parity bit
304    const uint sign8 = sign7 | (bitCount(sign7) << 7);
305    const uint sign = sign8 >> (iqs % 8);
306    const u8vec4 grid = unpack8(iq2xs_grid[qs & 511][(iqs % 8) / 4] >> (8 * (iqs % 4)));
307    bool sign0 = (sign & 1) != 0;
308    bool sign1 = (sign & 2) != 0;
309    bool sign2 = (sign & 4) != 0;
310    bool sign3 = (sign & 8) != 0;
311    return db * vec4(
312        grid.x * (sign0 ? -1.0 : 1.0),
313        grid.y * (sign1 ? -1.0 : 1.0),
314        grid.z * (sign2 ? -1.0 : 1.0),
315        grid.w * (sign3 ? -1.0 : 1.0)
316    );
317}
318#endif
319
320#if defined(DATA_A_IQ2_S)
321vec2 dequantize(uint ib, uint iqs, uint a_offset) {
322    const uint ib32 = iqs / 32;
323    const uint ib8 = iqs / 8;
324
325    const uint scale = (data_a[a_offset + ib].scales[ib32] >> (4 * ((iqs / 16) & 1))) & 0xf;
326    const uint qs = data_a[a_offset + ib].qs[ib8];
327    const uint qh = data_a[a_offset + ib].qh[ib32];
328    const uint qhshift = 2 * (ib8 % 4);
329    const uint sign = data_a[a_offset + ib].qs[QUANT_K / 8 + ib8] >> (iqs % 8);
330
331    const float db = 0.25 * (0.5 + scale);
332    const u8vec4 grid = unpack8(iq2s_grid[qs | ((qh << (8 - qhshift)) & 0x300)][(iqs % 8) / 4]);
333    bool sign0 = (sign & 1) != 0;
334    bool sign1 = (sign & 2) != 0;
335    return db * vec2(
336        grid[iqs % 4] * (sign0 ? -1.0 : 1.0),
337        grid[(iqs % 4) + 1] * (sign1 ? -1.0 : 1.0)
338    );
339}
340vec4 dequantize4(uint ib, uint iqs, uint a_offset) {
341    const uint ib32 = iqs / 32;
342    const uint ib8 = iqs / 8;
343
344    const uint scale = (data_a[a_offset + ib].scales[ib32] >> (4 * ((iqs / 16) & 1))) & 0xf;
345    const uint qs = data_a[a_offset + ib].qs[ib8];
346    const uint qh = data_a[a_offset + ib].qh[ib32];
347    const uint qhshift = 2 * (ib8 % 4);
348    const uint sign = data_a[a_offset + ib].qs[QUANT_K / 8 + ib8] >> (iqs % 8);
349
350    const float db = 0.25 * (0.5 + scale);
351    const u8vec4 grid = unpack8(iq2s_grid[qs | ((qh << (8 - qhshift)) & 0x300)][(iqs % 8) / 4]);
352    bool sign0 = (sign & 1) != 0;
353    bool sign1 = (sign & 2) != 0;
354    bool sign2 = (sign & 4) != 0;
355    bool sign3 = (sign & 8) != 0;
356    return db * vec4(
357        grid.x * (sign0 ? -1.0 : 1.0),
358        grid.y * (sign1 ? -1.0 : 1.0),
359        grid.z * (sign2 ? -1.0 : 1.0),
360        grid.w * (sign3 ? -1.0 : 1.0)
361    );
362}
363#endif
364
365#if defined(DATA_A_IQ3_XXS)
366vec2 dequantize(uint ib, uint iqs, uint a_offset) {
367    const uint ib4 = iqs / 4;
368    const uint ib32 = iqs / 32;
369    const uint is = QUANT_K / 4 + 4 * ib32;
370    const uint qs = data_a[a_offset + ib].qs[ib4];
371    // Scales are stored as packed 7+7+7+7+4 bits (4 sign tuples and 1 int4 scale)
372    const uint signs = pack32(u16vec2(data_a_packed16[a_offset + ib].qs[is / 2],
373        data_a_packed16[a_offset + ib].qs[is / 2 + 1]));
374    const float db = 0.5 * (0.5 + (signs >> 28));
375    const uint sign7 = bitfieldExtract(signs, 7 * (int(ib4 / 2) % 4), 7);
376    // Add parity bit
377    const uint sign8 = sign7 | (bitCount(sign7) << 7);
378    const uint sign = sign8 >> (iqs % 8);
379    const u8vec4 grid = unpack8(iq3xxs_grid[qs] >> (8 * (iqs % 4)));
380    bool sign0 = (sign & 1) != 0;
381    bool sign1 = (sign & 2) != 0;
382    return db * vec2(
383        grid.x * (sign0 ? -1.0 : 1.0),
384        grid.y * (sign1 ? -1.0 : 1.0)
385    );
386}
387vec4 dequantize4(uint ib, uint iqs, uint a_offset) {
388    const uint ib4 = iqs / 4;
389    const uint ib32 = iqs / 32;
390    const uint is = QUANT_K / 4 + 4 * ib32;
391    const uint qs = data_a[a_offset + ib].qs[ib4];
392    const uint signs = pack32(u16vec2(data_a_packed16[a_offset + ib].qs[is / 2],
393        data_a_packed16[a_offset + ib].qs[is / 2 + 1]));
394    const float db = 0.5 * (0.5 + (signs >> 28));
395    const uint sign7 = bitfieldExtract(signs, 7 * (int(ib4 / 2) % 4), 7);
396    // Add parity bit
397    const uint sign8 = sign7 | (bitCount(sign7) << 7);
398    const uint sign = sign8 >> (iqs % 8);
399    const u8vec4 grid = unpack8(iq3xxs_grid[qs]);
400    bool sign0 = (sign & 1) != 0;
401    bool sign1 = (sign & 2) != 0;
402    bool sign2 = (sign & 4) != 0;
403    bool sign3 = (sign & 8) != 0;
404    return db * vec4(
405        grid.x * (sign0 ? -1.0 : 1.0),
406        grid.y * (sign1 ? -1.0 : 1.0),
407        grid.z * (sign2 ? -1.0 : 1.0),
408        grid.w * (sign3 ? -1.0 : 1.0)
409    );
410}
411#endif
412
413#if defined(DATA_A_IQ3_S)
414vec2 dequantize(uint ib, uint iqs, uint a_offset) {
415    const uint qs = data_a[a_offset + ib].qs[iqs / 4];
416    const uint qh = data_a[a_offset + ib].qh[iqs / 32];
417    const uint sign = data_a[a_offset + ib].signs[iqs / 8] >> (iqs % 8);
418    const uint scale = data_a[a_offset + ib].scales[iqs / 64];
419    bool sign0 = (sign & 1) != 0;
420    bool sign1 = (sign & 2) != 0;
421    const float db = 1 + 2 * ((scale >> (4 * ((iqs / 32) & 1))) & 0xf);
422    const uint32_t grid = iq3s_grid[qs | ((qh << (8 - ((iqs / 4) % 8))) & 256)] >> (8 * (iqs % 4));
423    return db * vec2(
424        int(grid & 0xFF) * (sign0 ? -1.0 : 1.0),
425        int((grid >> 8) & 0xFF) * (sign1 ? -1.0 : 1.0)
426    );
427}
428vec4 dequantize4(uint ib, uint iqs, uint a_offset) {
429    const uint ib4 = iqs / 4;
430    const uint ib32 = iqs / 32;
431    const uint qs = data_a[a_offset + ib].qs[ib4];
432    const uint qh = data_a[a_offset + ib].qh[ib32];
433    const uint sign = data_a[a_offset + ib].signs[iqs / 8] >> (iqs % 8);
434    const uint scale = data_a[a_offset + ib].scales[ib32 / 2];
435    bool sign0 = (sign & 1) != 0;
436    bool sign1 = (sign & 2) != 0;
437    bool sign2 = (sign & 4) != 0;
438    bool sign3 = (sign & 8) != 0;
439    const float db = 1 + 2 * ((scale >> (4 * (ib32 & 1))) & 0xf);
440    const uint32_t grid = iq3s_grid[qs | ((qh << (8 - ib4 % 8)) & 256)] >> (8 * (iqs % 4));
441    return db * vec4(
442        int(grid & 0xFF) * (sign0 ? -1.0 : 1.0),
443        int((grid >> 8) & 0xFF) * (sign1 ? -1.0 : 1.0),
444        int((grid >> 16) & 0xFF) * (sign2 ? -1.0 : 1.0),
445        int((grid >> 24) & 0xFF) * (sign3 ? -1.0 : 1.0)
446    );
447}
448#endif
449
450#if defined(DATA_A_IQ4_XS)
451vec2 dequantize(uint ib, uint iqs, uint a_offset) {
452    const uint ib32 = iqs / 32;
453    const uint iq = 16 * ib32 + (iqs % 16);
454
455    const uint sl = (data_a[a_offset + ib].scales_l[ib32/2] >> (4 * (ib32 & 1))) & 0xF;
456    const uint sh = (data_a[a_offset + ib].scales_h >> (2 * ib32)) & 3;
457    const uint qshift = (iqs & 16) >> 2;
458    u8vec2 qs = u8vec2(data_a[a_offset + ib].qs[iq], data_a[a_offset + ib].qs[iq + 1]);
459    qs = (qs >> qshift) & uint8_t(0xF);
460
461    const float dl = float(int(sl | (sh << 4)) - 32);
462    return dl * vec2(kvalues_iq4nl[qs.x], kvalues_iq4nl[qs.y]);
463}
464vec4 dequantize4(uint ib, uint iqs, uint a_offset) {
465    const uint ib32 = iqs / 32;
466    const uint iq = 16 * ib32 + (iqs % 16);
467
468    const uint sl = (data_a[a_offset + ib].scales_l[ib32/2] >> (4 * (ib32 & 1))) & 0xF;
469    const uint sh = (data_a[a_offset + ib].scales_h >> (2 * ib32)) & 3;
470    const uint qshift = (iqs & 16) >> 2;
471    const u8vec4 qs = unpack8((data_a_packed32[a_offset + ib].qs[iq/4] >> qshift) & 0x0F0F0F0F);
472
473    const float dl = float(int(sl | (sh << 4)) - 32);
474    return dl * vec4(
475        kvalues_iq4nl[qs.x], kvalues_iq4nl[qs.y],
476        kvalues_iq4nl[qs.z], kvalues_iq4nl[qs.w]);
477}
478#endif
479
480#if defined(DATA_A_IQ4_NL)
481vec2 dequantize(uint ib, uint iqs, uint a_offset) {
482    const uint vui = uint(data_a[a_offset + ib].qs[iqs]);
483    return vec2(kvalues_iq4nl[vui & 0xF], kvalues_iq4nl[vui >> 4]);
484}
485vec4 dequantize4(uint ib, uint iqs, uint a_offset) {
486    const uint vui = uint(data_a_packed16[a_offset + ib].qs[iqs/2]);
487    return vec4(kvalues_iq4nl[vui & 0xF], kvalues_iq4nl[(vui >> 4) & 0xF], kvalues_iq4nl[(vui >> 8) & 0xF], kvalues_iq4nl[vui >> 12]);
488}
489#endif
490
491#if defined(DATA_A_MXFP4)
492vec2 dequantize(uint ib, uint iqs, uint a_offset) {
493    const uint vui = uint(data_a[a_offset + ib].qs[iqs]);
494#ifdef USE_OCP_FP4
495    return vec2(unpackFloat2xfe2m1EXT(uint8_t(vui)));
496#else
497    return vec2(kvalues_mxfp4[vui & 0xF], kvalues_mxfp4[vui >> 4]) * 0.5;
498#endif
499}
500vec4 dequantize4(uint ib, uint iqs, uint a_offset) {
501#ifdef USE_OCP_FP4
502    const uint16_t vui = uint16_t(uint(data_a[a_offset + ib].qs[iqs]) |
503                                  uint(data_a[a_offset + ib].qs[iqs + 1]) << 8);
504    return vec4(unpackFloat4xfe2m1EXT(vui));
505#else
506    vec2 v0 = dequantize(ib, iqs, a_offset);
507    vec2 v1 = dequantize(ib, iqs + 1, a_offset);
508    return vec4(v0.x, v0.y, v1.x, v1.y);
509#endif
510}
511#endif
512
513#if defined(DATA_A_NVFP4)
514vec2 dequantize(uint ib, uint iqs, uint a_offset) {
515    const uint sub = iqs >> 4;
516    const float d = ue4m3_to_fp32(data_a[a_offset + ib].d[sub]);
517    const uint j = iqs & 7;
518    const uint shift = (iqs & 8) >> 1; // 0 or 4
519#ifdef USE_OCP_FP4
520    const uint vui = uint(data_a_packed16[a_offset + ib].qs[(sub * 8u + j) / 2u]);
521    return vec2(bitcastExtractfe2m1EXT(unpack8(vui).xy, shift)) * d;
522#else
523    const uint vui0 = uint(data_a[a_offset + ib].qs[sub * 8u + j]);
524    const uint vui1 = uint(data_a[a_offset + ib].qs[sub * 8u + j + 1]);
525    const uint qs0 = (vui0 >> shift) & 0xF;
526    const uint qs1 = (vui1 >> shift) & 0xF;
527    return vec2(float(kvalues_mxfp4[qs0]), float(kvalues_mxfp4[qs1])) * d * 0.5;
528#endif
529}
530vec4 dequantize4(uint ib, uint iqs, uint a_offset) {
531#ifdef USE_OCP_FP4
532    const uint sub = iqs >> 4;
533    const float d = ue4m3_to_fp32(data_a[a_offset + ib].d[sub]);
534    const uint j = iqs & 7;
535    const uint shift = (iqs & 8) >> 1; // 0 or 4
536    const uint vui = data_a_packed32[a_offset + ib].qs[(sub * 8u + j) / 4u];
537    return vec4(bitcastExtractfe2m1EXT(unpack8(vui), shift)) * d;
538#else
539    const vec2 v0 = dequantize(ib, iqs, a_offset);
540    const vec2 v1 = dequantize(ib, iqs + 2u, a_offset);
541    return vec4(v0.x, v0.y, v1.x, v1.y);
542#endif
543}
544#endif
545
546#if defined(DATA_A_F32) || defined(DATA_A_F16) || defined(DATA_A_BF16)
547vec2 get_dm(uint ib, uint a_offset) {
548    return vec2(0, 0);
549}
550#endif
551
552#if defined(DATA_A_IQ1_M)
553vec2 get_dm(uint ib, uint a_offset) {
554    const uint16_t[4] scales = data_a[a_offset + ib].scales;
555    const u16vec4 s = u16vec4(scales[0], scales[1], scales[2], scales[3]) >> 12;
556    const float d = float(unpackHalf2x16(s.x | (s.y << 4) | (s.z << 8) | (s.w << 12)).x);
557    return vec2(d, 0);
558}
559#endif
560
561#if defined(DATA_A_Q2_0) || defined(DATA_A_Q4_0) || defined(DATA_A_Q5_0) || defined(DATA_A_Q8_0) || defined(DATA_A_IQ1_S) || defined(DATA_A_IQ2_XXS) || defined(DATA_A_IQ2_XS) || defined(DATA_A_IQ2_S) || defined(DATA_A_IQ3_XXS) || defined(DATA_A_IQ3_S) || defined(DATA_A_IQ4_XS) || defined(DATA_A_IQ4_NL)
562vec2 get_dm(uint ib, uint a_offset) {
563    return vec2(float(data_a[a_offset + ib].d), 0);
564}
565#endif
566
567#if defined(DATA_A_Q1_0)
568vec2 get_dm(uint ib, uint a_offset) {
569    const float d = float(data_a[a_offset + ib].d);
570    return vec2(d, 0);
571}
572#endif
573
574#if defined(DATA_A_MXFP4)
575vec2 get_dm(uint ib, uint a_offset) {
576    return vec2(e8m0_to_fp32(data_a[a_offset + ib].e), 0);
577}
578#endif
579
580#if defined(DATA_A_NVFP4)
581vec2 get_dm(uint ib, uint a_offset) {
582    return vec2(1.0, 0.0);
583}
584#endif
585
586#if defined(DATA_A_Q4_1) || defined(DATA_A_Q5_1)
587vec2 get_dm(uint ib, uint a_offset) {
588    const vec2 dm = vec2(data_a_packed32[a_offset + ib].dm);
589    return dm;
590}
591#endif
592
593#if defined(DATA_A_Q2_K)
594vec2 dequantize(uint ib, uint iqs, uint a_offset) {
595    iqs /= 2;
596    const uint qsi = (iqs / 64) * 32 + (iqs % 16) * 2; // 0,2,4..30
597    const uint scalesi = iqs / 8;                      // 0..15
598    const uint qsshift = ((iqs % 64) / 16) * 2;        // 0,2,4,6
599
600    const uvec2 qs = uvec2(data_a[a_offset + ib].qs[qsi], data_a[a_offset + ib].qs[qsi + 1]);
601    const uint scales = data_a[a_offset + ib].scales[scalesi];
602    const vec2 dm = vec2(data_a[a_offset + ib].dm);
603
604    return dm.x * float(scales & 0xF) * vec2((qs >> qsshift) & 3) - dm.y * float(scales >> 4);
605}
606vec2 get_dm(uint ib, uint a_offset) {
607    return vec2(1, 0);
608}
609#endif
610
611#if defined(DATA_A_TQ1_0)
612float tq1_0_val(uint ib, uint e, uint a_offset) {
613    const uint bidx = tq1_0_byte_of(e);
614    const uint qbyte = uint(bidx < 48u ? data_a[a_offset + ib].qs[bidx]
615                                       : data_a[a_offset + ib].qh[bidx - 48u]);
616    return float(tq1_0_trit(qbyte, tq1_0_digit_of(e))) - 1.0;
617}
618vec2 dequantize(uint ib, uint iqs, uint a_offset) {
619    return vec2(tq1_0_val(ib, iqs, a_offset), tq1_0_val(ib, iqs + 1u, a_offset));
620}
621vec2 get_dm(uint ib, uint a_offset) {
622    return vec2(float(data_a[a_offset + ib].d), 0);
623}
624#endif
625
626#if defined(DATA_A_TQ2_0)
627vec2 dequantize(uint ib, uint iqs, uint a_offset) {
628    // elem e -> byte qs[(e/128)*32 + e%32], bits 2*((e%128)/32); w = q - 1 (d applied via get_dm)
629    const uint qsi   = (iqs / 128) * 32 + (iqs % 32);  // iqs even -> qsi, qsi+1 in same group/level
630    const uint shift = 2 * ((iqs % 128) / 32);
631
632    const uvec2 qs = uvec2(data_a[a_offset + ib].qs[qsi], data_a[a_offset + ib].qs[qsi + 1]);
633    return vec2((qs >> shift) & 3) - 1.0;
634}
635vec2 get_dm(uint ib, uint a_offset) {
636    return vec2(float(data_a[a_offset + ib].d), 0);
637}
638#endif
639
640#if defined(DATA_A_Q3_K)
641vec2 dequantize(uint ib, uint iqs, uint a_offset) {
642    iqs /= 2;
643    const uint n = iqs / 64;                     // 0,1
644    const uint qsi = n * 32 + (iqs % 16) * 2;    // 0,2,4..62
645    const uint hmi =          (iqs % 16) * 2;    // 0,2,4..30
646    const uint j = (iqs % 64) / 4;               // 0..3
647    const uint is = iqs / 8;                     // 0..15
648    const uint halfsplit = ((iqs % 64) / 16);    // 0,1,2,3
649    const uint qsshift = halfsplit * 2;          // 0,2,4,6
650    const uint m = 1 << (4 * n + halfsplit);     // 1,2,4,8,16,32,64,128
651
652    const int8_t us = int8_t(((data_a[a_offset + ib].scales[is % 8] >> (4 * int(is / 8))) & 0xF)
653                          | (((data_a[a_offset + ib].scales[8 + (is % 4)] >> (2 * int(is / 4))) & 3) << 4));
654    const float dl = float(data_a[a_offset + ib].d) * float(us - 32);
655
656    return vec2(dl * float(int8_t((data_a[a_offset + ib].qs[qsi    ] >> qsshift) & 3) - (((data_a[a_offset + ib].hmask[hmi    ] & m) != 0) ? 0 : 4)),
657                dl * float(int8_t((data_a[a_offset + ib].qs[qsi + 1] >> qsshift) & 3) - (((data_a[a_offset + ib].hmask[hmi + 1] & m) != 0) ? 0 : 4)));
658}
659vec2 get_dm(uint ib, uint a_offset) {
660    return vec2(1, 0);
661}
662#endif
663
664#if defined(DATA_A_Q4_K)
665vec2 dequantize(uint ib, uint iqs, uint a_offset) {
666    iqs /= 2;
667    const uint n = iqs / 32;                   // 0,1,2,3
668    const uint b = (iqs % 32) / 16;            // 0,1
669    const uint is = 2 * n + b;                 // 0..7
670    const uint qsi = n * 32 + (iqs % 16) * 2;  // 0,2,4..126
671
672    const vec2 loadd = vec2(data_a[a_offset + ib].dm);
673
674    const uint scidx0 = (is < 4) ? is : (is + 4);
675    const uint scidx1 = (is < 4) ? is : (is - 4);
676    const uint scidxmask1 = (is < 4) ? 0x30 : 0xC0;
677    const uint scidxshift1 = (is < 4) ? 0 : 2;
678    const uint mbidx0 = is + 4;
679    const uint mbidx1 = (is < 4) ? is + 4 : is;
680    const uint mbidxmask0 = (is < 4) ? 0xF : 0xF0;
681    const uint mbidxshift0 = (is < 4) ? 0 : 4;
682    const uint mbidxmask1 = (is < 4) ? 0x30 : 0xC0;
683    const uint mbidxshift1 = (is < 4) ? 0 : 2;
684
685    const uint8_t sc = uint8_t((data_a[a_offset + ib].scales[scidx0] & 0xF) | ((data_a[a_offset + ib].scales[scidx1] & scidxmask1) >> scidxshift1));
686    const uint8_t mbyte = uint8_t((data_a[a_offset + ib].scales[mbidx0] & mbidxmask0) >> mbidxshift0 | ((data_a[a_offset + ib].scales[mbidx1] & mbidxmask1) >> mbidxshift1));
687
688    const float d = loadd.x * sc;
689    const float m = -loadd.y * mbyte;
690
691    return vec2(fma(d, float((data_a[a_offset + ib].qs[qsi    ] >> (b * 4)) & 0xF), m),
692                fma(d, float((data_a[a_offset + ib].qs[qsi + 1] >> (b * 4)) & 0xF), m));
693}
694vec2 get_dm(uint ib, uint a_offset) {
695    return vec2(1, 0);
696}
697#endif
698
699#if defined(DATA_A_Q5_K)
700vec2 dequantize(uint ib, uint iqs, uint a_offset) {
701    iqs /= 2;
702    const uint n = iqs / 32;                   // 0,1,2,3
703    const uint b = (iqs % 32) / 16;            // 0,1
704    const uint is = 2 * n + b;                 // 0..7
705    const uint qsi = n * 32 + (iqs % 16) * 2;  // 0,2,4..126
706    const uint qhi = (iqs % 16) * 2;           // 0,2,4..30
707
708    const uint8_t hm = uint8_t(1 << (iqs / 16));
709
710    const vec2 loadd = vec2(data_a[a_offset + ib].dm);
711
712    const uint scidx0 = (is < 4) ? is : (is + 4);
713    const uint scidx1 = (is < 4) ? is : (is - 4);
714    const uint scidxmask1 = (is < 4) ? 0x30 : 0xC0;
715    const uint scidxshift1 = (is < 4) ? 0 : 2;
716    const uint mbidx0 = is + 4;
717    const uint mbidx1 = (is < 4) ? is + 4 : is;
718    const uint mbidxmask0 = (is < 4) ? 0xF : 0xF0;
719    const uint mbidxshift0 = (is < 4) ? 0 : 4;
720    const uint mbidxmask1 = (is < 4) ? 0x30 : 0xC0;
721    const uint mbidxshift1 = (is < 4) ? 0 : 2;
722
723    const uint8_t sc    = uint8_t((data_a[a_offset + ib].scales[scidx0] & 0xF)                         | ((data_a[a_offset + ib].scales[scidx1] & scidxmask1) >> scidxshift1));
724    const uint8_t mbyte = uint8_t(((data_a[a_offset + ib].scales[mbidx0] & mbidxmask0) >> mbidxshift0) | ((data_a[a_offset + ib].scales[mbidx1] & mbidxmask1) >> mbidxshift1));
725
726    const float d = loadd.x * sc;
727    const float m = -loadd.y * mbyte;
728
729    return vec2(fma(d, float((data_a[a_offset + ib].qs[qsi    ] >> (b * 4)) & 0xF) + float((data_a[a_offset + ib].qh[qhi    ] & hm) != 0 ? 16 : 0), m),
730                fma(d, float((data_a[a_offset + ib].qs[qsi + 1] >> (b * 4)) & 0xF) + float((data_a[a_offset + ib].qh[qhi + 1] & hm) != 0 ? 16 : 0), m));
731}
732vec2 get_dm(uint ib, uint a_offset) {
733    return vec2(1, 0);
734}
735#endif
736
737#if defined(DATA_A_Q6_K)
738vec2 dequantize(uint ib, uint iqs, uint a_offset) {
739    iqs /= 2;
740    const uint n = iqs / 64;                    // 0,1
741    const uint b = (iqs % 64) / 32;             // 0,1
742    const uint is_b = (iqs % 16) / 8;           // 0,1
743    const uint qhshift = ((iqs % 64) / 16) * 2; // 0,2,4,6
744    const uint is = 8 * n + qhshift + is_b;     // 0..15
745    const uint qsi = n * 64 + (iqs % 32) * 2;   // 0,2,4..126
746    const uint qhi = n * 32 + (iqs % 16) * 2;   // 0,2,4..62
747
748    const float dscale = float(data_a[a_offset + ib].d) * float(data_a[a_offset + ib].scales[is]);
749
750    return vec2(dscale * float(int8_t(((data_a[a_offset + ib].ql[qsi    ] >> (b * 4)) & 0xF) | (((data_a[a_offset + ib].qh[qhi    ] >> qhshift) & 3) << 4)) - 32),
751                dscale * float(int8_t(((data_a[a_offset + ib].ql[qsi + 1] >> (b * 4)) & 0xF) | (((data_a[a_offset + ib].qh[qhi + 1] >> qhshift) & 3) << 4)) - 32));
752}
753vec2 get_dm(uint ib, uint a_offset) {
754    return vec2(1, 0);
755}
756#endif