Skip to content

Commit 045b7a1

Browse files
remyoudomphengjoelvaneenwyk
remyoudompheng
authored and
joelvaneenwyk
committed
vulkan: initial support for IQ1_S and IQ1_M quantizations (llama/11528)
* vulkan: initial support for IQ1_S and IQ1_M quantizations * vulkan: define MMV kernels for IQ1 quantizations * devops: increase timeout of Vulkan tests again * vulkan: simplify ifdef for init_iq_shmem
1 parent 7bc12ec commit 045b7a1

16 files changed

+710
-29
lines changed

‎ggml/src/ggml-vulkan/ggml-vulkan.cpp‎

Lines changed: 74 additions & 20 deletions
Large diffs are not rendered by default.

‎ggml/src/ggml-vulkan/vulkan-shaders/copy_from_quant.comp‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@ layout(local_size_x = 1, local_size_y = 1, local_size_z = 1) in;
1212
#endif
1313

1414
void main() {
15-
#if 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)
15+
#ifdef NEEDS_INIT_IQ_SHMEM
1616
init_iq_shmem(gl_WorkGroupSize);
1717
if (gl_LocalInvocationIndex.x != 0) {
1818
return;

‎ggml/src/ggml-vulkan/vulkan-shaders/copy_to_quant.comp‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -217,7 +217,7 @@ void quantize(uint dst_idx, uint src_idx)
217217
#endif
218218

219219
void main() {
220-
#if 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)
220+
#ifdef NEEDS_INIT_IQ_SHMEM
221221
init_iq_shmem(gl_WorkGroupSize);
222222
if (gl_LocalInvocationIndex.x != 0) {
223223
return;

‎ggml/src/ggml-vulkan/vulkan-shaders/dequant_funcs.comp‎

Lines changed: 87 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -88,6 +88,83 @@ vec4 dequantize4(uint ib, uint iqs, uint a_offset) {
8888
}
8989
#endif
9090

91+
#if defined(DATA_A_IQ1_S)
92+
vec2 dequantize(uint ib, uint iqs, uint a_offset) {
93+
const uint ib32 = iqs / 32;
94+
const uint ib8 = iqs / 8;
95+
const int i8 = int(iqs % 8);
96+
const uint qh = data_a[a_offset + ib].qh[ib32];
97+
const uint qs = data_a[a_offset + ib].qs[ib8];
98+
const float dl = float(2 * bitfieldExtract(qh, 12, 3) + 1);
99+
const float delta = ((qh & 0x8000) != 0) ? -IQ1S_DELTA : IQ1S_DELTA;
100+
const uint idxhi = bitfieldExtract(qh, 3 * int(ib8 & 3), 3);
101+
const int16_t grid = int16_t(iq1s_grid[qs | (idxhi << 8)]);
102+
// Signed bitfield extract.
103+
const ivec2 gvec = ivec2(
104+
bitfieldExtract(grid, 2 * (i8), 2),
105+
bitfieldExtract(grid, 2 * (i8 + 1), 2)
106+
);
107+
return dl * (vec2(gvec) + delta);
108+
}
109+
vec4 dequantize4(uint ib, uint iqs, uint a_offset) {
110+
const uint ib32 = iqs / 32;
111+
const uint ib8 = iqs / 8;
112+
const int i8 = int(iqs % 8);
113+
const uint qh = data_a[a_offset + ib].qh[ib32];
114+
const uint qs = data_a[a_offset + ib].qs[ib8];
115+
const float dl = 2 * bitfieldExtract(qh, 12, 3) + 1;
116+
const float delta = ((qh & 0x8000) != 0) ? -IQ1S_DELTA : IQ1S_DELTA;
117+
const int16_t grid = int16_t(iq1s_grid[qs | (bitfieldExtract(qh, 3 * int(ib8 & 3), 3) << 8)]);
118+
// Signed bitfield extract.
119+
const ivec4 gvec = ivec4(
120+
bitfieldExtract(grid, 2 * (i8), 2),
121+
bitfieldExtract(grid, 2 * (i8 + 1), 2),
122+
bitfieldExtract(grid, 2 * (i8 + 2), 2),
123+
bitfieldExtract(grid, 2 * (i8 + 3), 2)
124+
);
125+
return dl * (vec4(gvec) + delta);
126+
}
127+
#endif
128+
129+
#if defined(DATA_A_IQ1_M)
130+
vec2 dequantize(uint ib, uint iqs, uint a_offset) {
131+
const uint ib8 = iqs / 8;
132+
const uint ib16 = iqs / 16;
133+
const int i8 = int(iqs % 8);
134+
const uint sc = data_a[a_offset + ib].scales[iqs / 64];
135+
const uint qs = data_a[a_offset + ib].qs[ib8];
136+
const uint qh = data_a[a_offset + ib].qh[ib16] >> (4 * (ib8 & 1));
137+
const float dl = 2 * bitfieldExtract(sc, 3 * int(ib16 & 3), 3) + 1;
138+
const float delta = ((qh & 8) != 0) ? -IQ1M_DELTA : IQ1M_DELTA;
139+
const int16_t grid = int16_t(iq1s_grid[qs | ((qh & 7) << 8)]);
140+
// Signed bitfield extract.
141+
const ivec2 gvec = ivec2(
142+
bitfieldExtract(grid, 2 * (i8), 2),
143+
bitfieldExtract(grid, 2 * (i8 + 1), 2)
144+
);
145+
return dl * (vec2(gvec) + delta);
146+
}
147+
vec4 dequantize4(uint ib, uint iqs, uint a_offset) {
148+
const uint ib8 = iqs / 8;
149+
const uint ib16 = iqs / 16;
150+
const int i8 = int(iqs % 8);
151+
const uint sc = data_a[a_offset + ib].scales[iqs / 64];
152+
const uint qs = data_a[a_offset + ib].qs[ib8];
153+
const uint qh = data_a[a_offset + ib].qh[ib16] >> (4 * (ib8 & 1));
154+
const float dl = 2 * bitfieldExtract(sc, 3 * int(ib16 & 3), 3) + 1;
155+
const float delta = ((qh & 8) != 0) ? -IQ1M_DELTA : IQ1M_DELTA;
156+
const int16_t grid = int16_t(iq1s_grid[qs | ((qh & 7) << 8)]);
157+
// Signed bitfield extract.
158+
const ivec4 gvec = ivec4(
159+
bitfieldExtract(grid, 2 * (i8), 2),
160+
bitfieldExtract(grid, 2 * (i8 + 1), 2),
161+
bitfieldExtract(grid, 2 * (i8 + 2), 2),
162+
bitfieldExtract(grid, 2 * (i8 + 3), 2)
163+
);
164+
return dl * (vec4(gvec) + delta);
165+
}
166+
#endif
167+
91168
#if defined(DATA_A_IQ2_XXS)
92169
vec2 dequantize(uint ib, uint iqs, uint a_offset) {
93170
const uint ib32 = iqs / 32;
@@ -357,7 +434,16 @@ vec2 get_dm(uint ib, uint a_offset) {
357434
}
358435
#endif
359436

360-
#if defined(DATA_A_Q4_0) || defined(DATA_A_Q5_0) || defined(DATA_A_Q8_0) || 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)
437+
#if defined(DATA_A_IQ1_M)
438+
vec2 get_dm(uint ib, uint a_offset) {
439+
const uint16_t[4] scales = data_a[a_offset + ib].scales;
440+
const u16vec4 s = u16vec4(scales[0], scales[1], scales[2], scales[3]) >> 12;
441+
const float d = float(unpackHalf2x16(s.x | (s.y << 4) | (s.z << 8) | (s.w << 12)).x);
442+
return vec2(d, 0);
443+
}
444+
#endif
445+
446+
#if 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)
361447
vec2 get_dm(uint ib, uint a_offset) {
362448
return vec2(float(data_a[a_offset + ib].d), 0);
363449
}

‎ggml/src/ggml-vulkan/vulkan-shaders/dequant_funcs_cm2.comp‎

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -301,6 +301,56 @@ float16_t dequantFuncQ6_K(const in decodeBufQ6_K bl, const in uint blockCoords[2
301301
return ret;
302302
}
303303

304+
#if defined(DATA_A_IQ1_S)
305+
layout(buffer_reference, std430, buffer_reference_align = 2) buffer decodeBufIQ1_S {
306+
block_iq1_s block;
307+
};
308+
309+
float16_t dequantFuncIQ1_S(const in decodeBufIQ1_S bl, const in uint blockCoords[2], const in uint coordInBlock[2])
310+
{
311+
const float16_t d = bl.block.d;
312+
const uint idx = coordInBlock[1];
313+
314+
const uint ib32 = idx / 32;
315+
const uint ib8 = idx / 8;
316+
317+
const uint qh = bl.block.qh[ib32];
318+
const uint qs = bl.block.qs[ib8];
319+
const float dl = d * float(2 * bitfieldExtract(qh, 12, 3) + 1);
320+
const float delta = ((qh & 0x8000) != 0) ? -IQ1S_DELTA : IQ1S_DELTA;
321+
const uint grid = iq1s_grid[qs | (bitfieldExtract(qh, 3 * int(ib8 & 3), 3) << 8)];
322+
323+
float16_t ret = float16_t(dl) * (float16_t(bitfieldExtract(int(grid), 2 * int(idx % 8), 2)) + float16_t(delta));
324+
return ret;
325+
}
326+
#endif
327+
328+
#if defined(DATA_A_IQ1_M)
329+
layout(buffer_reference, std430, buffer_reference_align = 2) buffer decodeBufIQ1_M {
330+
block_iq1_m block;
331+
};
332+
333+
float16_t dequantFuncIQ1_M(const in decodeBufIQ1_M bl, const in uint blockCoords[2], const in uint coordInBlock[2])
334+
{
335+
const u16vec4 scales = u16vec4(bl.block.scales[0], bl.block.scales[1], bl.block.scales[2], bl.block.scales[3]) >> 12;
336+
const float16_t d = uint16BitsToHalf(scales.x | (scales.y << 4) | (scales.z << 8) | (scales.w << 12));
337+
const uint idx = coordInBlock[1];
338+
339+
const uint ib8 = idx / 8;
340+
const uint ib16 = idx / 16;
341+
const int i8 = int(idx % 8);
342+
const uint sc = bl.block.scales[ib8 / 8];
343+
const uint qs = bl.block.qs[ib8];
344+
const uint qh = bl.block.qh[ib16] >> (4 * (ib8 & 1));
345+
const float dl = 2 * bitfieldExtract(sc, 3 * int(ib16 & 3), 3) + 1;
346+
const float delta = ((qh & 8) != 0) ? -IQ1S_DELTA : IQ1S_DELTA;
347+
const uint grid = iq1s_grid[qs | ((qh & 7) << 8)];
348+
349+
float16_t ret = d * float16_t(dl) * (float16_t(bitfieldExtract(int(grid), 2 * i8, 2)) + float16_t(delta));
350+
return ret;
351+
}
352+
#endif
353+
304354
#if defined(DATA_A_IQ2_XXS)
305355
layout(buffer_reference, std430, buffer_reference_align = 2) buffer decodeBufIQ2_XXS {
306356
block_iq2_xxs block;
@@ -512,6 +562,10 @@ float16_t dequantFuncIQ4_NL(const in decodeBufIQ4_NL bl, const in uint blockCoor
512562
#define dequantFuncA dequantFuncQ5_K
513563
#elif defined(DATA_A_Q6_K)
514564
#define dequantFuncA dequantFuncQ6_K
565+
#elif defined(DATA_A_IQ1_S)
566+
#define dequantFuncA dequantFuncIQ1_S
567+
#elif defined(DATA_A_IQ1_M)
568+
#define dequantFuncA dequantFuncIQ1_M
515569
#elif defined(DATA_A_IQ2_XXS)
516570
#define dequantFuncA dequantFuncIQ2_XXS
517571
#elif defined(DATA_A_IQ2_XS)

‎ggml/src/ggml-vulkan/vulkan-shaders/dequant_iq1_m.comp‎

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,42 @@
1+
#version 450
2+
3+
#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require
4+
5+
#include "dequant_head.comp"
6+
7+
layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in;
8+
9+
layout (binding = 0) readonly buffer A {block_iq1_m data_a[];};
10+
layout (binding = 1) writeonly buffer D {D_TYPE data_b[];};
11+
12+
void main() {
13+
// Each thread handles 1 subblock (32 values with 2 scales)
14+
const uint ib = gl_WorkGroupID.x * 32 + gl_LocalInvocationID.x / 8;
15+
16+
init_iq_shmem(gl_WorkGroupSize);
17+
18+
if (ib >= p.nel / 256) {
19+
return;
20+
}
21+
22+
const uint ib32 = gl_LocalInvocationID.x % 8;
23+
const uint ib64 = ib32 / 2;
24+
const uint b_idx = 256 * ib + 32 * ib32;
25+
26+
const uint16_t[4] scales = data_a[ib].scales;
27+
const u16vec4 s = u16vec4(scales[0], scales[1], scales[2], scales[3]) >> 12;
28+
const float d = float(unpackHalf2x16(s.x | (s.y << 4) | (s.z << 8) | (s.w << 12)).x);
29+
30+
const uint sc = data_a[ib].scales[ib64];
31+
[[unroll]] for (int l = 0; l < 4; ++l) {
32+
const uint ib16 = 2 * ib32 + l / 2;
33+
const float dl = d * (2 * bitfieldExtract(sc, 3 * int(ib16 & 3), 3) + 1);
34+
const uint qh = data_a[ib].qh[ib16] >> (4 * (l & 1));
35+
const uint qs = data_a[ib].qs[4 * ib32 + l];
36+
const float delta = ((qh & 8) != 0) ? -IQ1M_DELTA : IQ1M_DELTA;
37+
const int16_t grid = int16_t(iq1s_grid[qs | ((qh & 7) << 8)]);
38+
[[unroll]] for (int j = 0; j < 8; ++j) {
39+
data_b[b_idx + 8 * l + j] = D_TYPE(dl * (bitfieldExtract(grid, 2*j, 2) + delta));
40+
}
41+
}
42+
}

‎ggml/src/ggml-vulkan/vulkan-shaders/dequant_iq1_s.comp‎

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,35 @@
1+
#version 450
2+
3+
#include "dequant_head.comp"
4+
5+
layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in;
6+
7+
layout (binding = 0) readonly buffer A {block_iq1_s data_a[];};
8+
layout (binding = 1) writeonly buffer D {D_TYPE data_b[];};
9+
10+
void main() {
11+
// Each thread handles 1 subblock (32 values with 2 scales)
12+
const uint ib = gl_WorkGroupID.x * 32 + gl_LocalInvocationID.x / 8;
13+
14+
init_iq_shmem(gl_WorkGroupSize);
15+
16+
if (ib >= p.nel / 256) {
17+
return;
18+
}
19+
20+
const uint ib32 = gl_LocalInvocationID.x % 8;
21+
const uint b_idx = 256 * ib + 32 * ib32;
22+
23+
uint qh = data_a[ib].qh[ib32];
24+
const float d = float(data_a[ib].d);
25+
const float dl = d * float(2 * bitfieldExtract(qh, 12, 3) + 1);
26+
const float delta = ((qh & 0x8000) != 0) ? -IQ1S_DELTA : IQ1S_DELTA;
27+
[[unroll]] for (uint l = 0; l < 4; ++l) {
28+
const uint qs = data_a[ib].qs[4 * ib32 + l];
29+
const uint hi = bitfieldExtract(qh, 3 * int(l), 3);
30+
const int16_t grid = int16_t(iq1s_grid[qs | (hi << 8)]);
31+
[[unroll]] for (int j = 0; j < 8; ++j) {
32+
data_b[b_idx + 8 * l + j] = D_TYPE(dl * (bitfieldExtract(grid, 2*j, 2) + delta));
33+
}
34+
}
35+
}

‎ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm2.comp‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -104,7 +104,7 @@ ACC_TYPE Max(const in uint32_t row, const in uint32_t col, const in ACC_TYPE ele
104104
#endif
105105

106106
void main() {
107-
#if 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)
107+
#ifdef NEEDS_INIT_IQ_SHMEM
108108
init_iq_shmem(gl_WorkGroupSize);
109109
#endif
110110

‎ggml/src/ggml-vulkan/vulkan-shaders/get_rows_quant.comp‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@ void main() {
1212
const uint i11 = (gl_GlobalInvocationID.z)/p.ne12;
1313
const uint i12 = (gl_GlobalInvocationID.z)%p.ne12;
1414

15-
#if 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)
15+
#ifdef NEEDS_INIT_IQ_SHMEM
1616
init_iq_shmem(gl_WorkGroupSize);
1717
#endif
1818

‎ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec.comp‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -133,7 +133,7 @@ void compute_outputs(const uint32_t first_row, const uint32_t num_rows) {
133133
void main() {
134134
const uint first_row = NUM_ROWS * (gl_WorkGroupID.x + gl_NumWorkGroups.x * gl_WorkGroupID.z);
135135

136-
#if 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)
136+
#ifdef NEEDS_INIT_IQ_SHMEM
137137
init_iq_shmem(gl_WorkGroupSize);
138138
#endif
139139

0 commit comments

Comments
 (0)