Skip to content

Commit 6841888

Browse files
ikawrakowIwan Kawrakow
authored and
Nexesenex
committed
IQ4_KSS improvements (LostRuins#642)
* iq4_kss: slightly better quantization * iq4_kss: CUDA MMQ * iq4_kss: repack/convert to q8_k_r8 (AVX2) * iq4_kss: repack/convert to q8_k_r8 (NEON) --------- Co-authored-by: Iwan Kawrakow <[email protected]>
1 parent 7a93932 commit 6841888

File tree

6 files changed

+195
-18
lines changed

6 files changed

+195
-18
lines changed

‎ggml/src/ggml-cuda/mmq.cu‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -104,6 +104,9 @@ void ggml_cuda_op_mul_mat_q(
104104
case GGML_TYPE_IQ3_KS:
105105
mul_mat_q_case<GGML_TYPE_IQ3_KS>(ctx, args, stream);
106106
break;
107+
case GGML_TYPE_IQ4_KSS:
108+
mul_mat_q_case<GGML_TYPE_IQ4_KSS>(ctx, args, stream);
109+
break;
107110
case GGML_TYPE_IQ4_KS:
108111
mul_mat_q_case<GGML_TYPE_IQ4_KS>(ctx, args, stream);
109112
break;
@@ -218,6 +221,7 @@ bool ggml_cuda_should_use_mmq(enum ggml_type type, int cc, int64_t ne11) {
218221
case GGML_TYPE_IQ4_NL:
219222
case GGML_TYPE_IQ2_KL:
220223
case GGML_TYPE_IQ3_KS:
224+
case GGML_TYPE_IQ4_KSS:
221225
case GGML_TYPE_IQ4_KS:
222226
case GGML_TYPE_IQ4_KS_R4:
223227
case GGML_TYPE_IQ5_KS:

‎ggml/src/ggml-cuda/mmq.cuh‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -88,6 +88,7 @@ static mmq_q8_1_ds_layout mmq_get_q8_1_ds_layout(const ggml_type type_x) {
8888
case GGML_TYPE_IQ2_KL:
8989
case GGML_TYPE_IQ3_KS:
9090
case GGML_TYPE_IQ3_K_R4:
91+
case GGML_TYPE_IQ4_KSS:
9192
case GGML_TYPE_IQ4_KS:
9293
case GGML_TYPE_IQ4_KS_R4:
9394
case GGML_TYPE_IQ4_K:
@@ -205,6 +206,7 @@ static constexpr __host__ __device__ tile_x_sizes mmq_get_dp4a_tile_x_sizes(ggml
205206
case GGML_TYPE_IQ4_NL : return MMQ_DP4A_TXS_Q8_0;
206207
case GGML_TYPE_IQ2_KL : return MMQ_DP4A_TXS_Q8_0;
207208
case GGML_TYPE_IQ3_KS : return MMQ_DP4A_TXS_Q8_0;
209+
case GGML_TYPE_IQ4_KSS : return MMQ_DP4A_TXS_Q8_0;
208210
case GGML_TYPE_IQ4_KS : return MMQ_DP4A_TXS_Q8_0;
209211
case GGML_TYPE_IQ4_KS_R4 : return MMQ_DP4A_TXS_Q8_0;
210212
case GGML_TYPE_IQ5_KS : return MMQ_DP4A_TXS_Q8_0;
@@ -264,6 +266,7 @@ static constexpr __host__ __device__ int mmq_get_mma_tile_x_k(ggml_type type) {
264266
case GGML_TYPE_IQ4_NL : return MMQ_MMA_TILE_X_K_Q8_0;
265267
case GGML_TYPE_IQ2_KL : return MMQ_MMA_TILE_X_K_Q8_0;
266268
case GGML_TYPE_IQ3_KS : return MMQ_MMA_TILE_X_K_Q8_0;
269+
case GGML_TYPE_IQ4_KSS : return MMQ_MMA_TILE_X_K_Q8_0;
267270
case GGML_TYPE_IQ4_KS : return MMQ_MMA_TILE_X_K_Q8_0;
268271
case GGML_TYPE_IQ4_KS_R4 : return MMQ_MMA_TILE_X_K_Q8_0;
269272
case GGML_TYPE_IQ5_KS : return MMQ_MMA_TILE_X_K_Q8_0;
@@ -4301,6 +4304,7 @@ extern DECL_MMQ_CASE(GGML_TYPE_IQ4_NL);
43014304
extern DECL_MMQ_CASE(GGML_TYPE_IQ4_XS);
43024305
extern DECL_MMQ_CASE(GGML_TYPE_IQ2_KL);
43034306
extern DECL_MMQ_CASE(GGML_TYPE_IQ3_KS);
4307+
extern DECL_MMQ_CASE(GGML_TYPE_IQ4_KSS);
43044308
extern DECL_MMQ_CASE(GGML_TYPE_IQ4_KS);
43054309
extern DECL_MMQ_CASE(GGML_TYPE_IQ4_KS_R4);
43064310
extern DECL_MMQ_CASE(GGML_TYPE_IQ5_KS_R4);

‎ggml/src/ggml-cuda/template-instances/mmq-instance-iq4_kss.cu‎

Lines changed: 68 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,68 @@
1+
#include "../mmq.cuh"
2+
3+
template <int mmq_y, int nwarps, bool need_check> static __device__ __forceinline__ void load_tiles_iq4_kss(
4+
const char * __restrict__ x, int * __restrict__ x_tile, const int & kbx0, const int & i_max, const int & stride) {
5+
6+
#ifdef INT8_MMA_AVAILABLE
7+
int * x_qs = (int *) x_tile;
8+
float * x_df = (float *) (x_qs + WARP_SIZE*2);
9+
#else
10+
constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_IQ4_XS, mmq_y);
11+
int * x_qs = (int *) x_tile;
12+
float * x_df = (float *) (x_qs + txs.qs);
13+
#endif // INT8_MMA_AVAILABLE
14+
15+
const int kqsx = threadIdx.x / 4;
16+
17+
uint32_t aux32[2];
18+
auto a8 = (const uint8_t *)aux32;
19+
20+
#pragma unroll
21+
for (int i0 = 0; i0 < mmq_y; i0 += 4*nwarps) {
22+
int i = i0 + 4*threadIdx.y + threadIdx.x%4;
23+
24+
if (need_check) {
25+
i = min(i, i_max);
26+
}
27+
28+
const float * dptr = (const float *)(x + i*stride);
29+
const block_iq4_kss * bxi = (const block_iq4_kss *)(dptr + 1) + kbx0;
30+
const uint32_t * q4 = bxi->qs + 4*kqsx;
31+
uint32_t s32 = (q4[0] & 0x00010001) | ((q4[1] & 0x00010001) << 2) | ((q4[2] & 0x00010001) << 4) | ((q4[3] & 0x00010001) << 6);
32+
uint8_t ls = (s32 | (s32 >> 15)) & 0xff;
33+
34+
auto values = iq4k_table + ((ls & 1) << 8);
35+
36+
#pragma unroll
37+
for (int j = 0; j < 4; ++j) {
38+
uint32_t val = q4[j] & 0xfffefffe;
39+
val = val ^ (val >> 1);
40+
aux32[0] = (val >> 0) & 0x0f0f0f0f;
41+
aux32[1] = (val >> 4) & 0x0f0f0f0f;
42+
#ifdef INT8_MMA_AVAILABLE
43+
x_qs[i*MMQ_MMA_TILE_X_K_Q8_0 + 8*kqsx + j + 0] = int_from_table_x(a8+0, values);
44+
x_qs[i*MMQ_MMA_TILE_X_K_Q8_0 + 8*kqsx + j + 4] = int_from_table_x(a8+4, values);
45+
#else
46+
x_qs[i*(2*WARP_SIZE + 1) + 8*kqsx + j + 0] = int_from_table_x(a8+0, values);
47+
x_qs[i*(2*WARP_SIZE + 1) + 8*kqsx + j + 4] = int_from_table_x(a8+4, values);
48+
#endif // INT8_MMA_AVAILABLE
49+
}
50+
#ifdef INT8_MMA_AVAILABLE
51+
x_df[i*MMQ_MMA_TILE_X_K_Q8_0 + kqsx] = dptr[0] * ((ls & 254) - 127);
52+
#else
53+
x_df[i*(WARP_SIZE/4) + i/4 + kqsx] = dptr[0] * ((ls & 254) - 127);
54+
#endif // INT8_MMA_AVAILABLE
55+
}
56+
57+
}
58+
59+
60+
template <int mmq_x, int mmq_y, int nwarps, bool need_check>
61+
struct mmq_type_traits<mmq_x, mmq_y, nwarps, need_check, GGML_TYPE_IQ4_KSS> {
62+
static constexpr load_tiles_mmq_t load_tiles = load_tiles_iq4_kss<mmq_y, nwarps, need_check>;
63+
static constexpr vec_dot_mmq_t vec_dot_mma = vec_dot_q8_0_q8_1_mma<mmq_x, mmq_y, nwarps, MMQ_Q8_1_DS_LAYOUT_D4>;
64+
static constexpr vec_dot_mmq_t vec_dot_dp4a = vec_dot_q8_0_q8_1_dp4a<mmq_x, mmq_y, nwarps>;
65+
};
66+
67+
DECL_MMQ_CASE(GGML_TYPE_IQ4_KSS);
68+

‎ggml/src/iqk/iqk_gemm_iqk_quants.cpp‎

Lines changed: 107 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2706,6 +2706,60 @@ void iqk_convert_iq3_ks_q8_k_r8(int n, const void * vx, size_t bx, void * vy, in
27062706
}
27072707
}
27082708

2709+
void iqk_convert_iq4_kss_q8_k_r8(int n, const void * vx, size_t bx, void * vy, int nrc_x) {
2710+
GGML_ASSERT(n%QK_K == 0);
2711+
GGML_ASSERT(nrc_x%8 == 0);
2712+
2713+
int nb = n/QK_K;
2714+
2715+
const block_iq4_kss * x8[8];
2716+
2717+
block_q8_k_r8 * y = (block_q8_k_r8 *)vy;
2718+
2719+
__m256i values[2];
2720+
{
2721+
auto v1 = _mm_loadu_si128((const __m128i *)iq4k_values+0);
2722+
auto v2 = _mm_loadu_si128((const __m128i *)iq4k_values+1);
2723+
values[0] = MM256_SET_M128I(v1, v1);
2724+
values[1] = MM256_SET_M128I(v2, v2);
2725+
}
2726+
2727+
float drow[8];
2728+
float dnew[8];
2729+
int16_t ls[16];
2730+
2731+
__m256i xv[8];
2732+
uint32_t block[8];
2733+
2734+
for (int ix = 0; ix < nrc_x; ix += 8) {
2735+
for (int k = 0; k < 8; ++k) {
2736+
const float * dptr = (const float *)((const char *)vx + (ix + k)*bx);
2737+
drow[k] = dptr[0];
2738+
x8[k] = (const block_iq4_kss *)(dptr + 1);
2739+
}
2740+
auto vd = _mm256_loadu_ps(drow);
2741+
for (int i = 0; i < nb; ++i) {
2742+
for (int k = 0; k < 8; ++k) {
2743+
for (int ib32 = 0; ib32 < 8; ++ib32) {
2744+
auto val = _mm_loadu_si128((const __m128i *)x8[k][i].qs+ib32);
2745+
auto val_q = _mm_and_si128(val, _mm_set1_epi32(0xfffefffe));
2746+
val_q = _mm_xor_si128(val_q, _mm_srli_epi16(val_q, 1));
2747+
xv[ib32] = _mm256_and_si256(MM256_SET_M128I(_mm_srli_epi16(val_q, 4), val_q), _mm256_set1_epi8(0xf));
2748+
auto q4 = x8[k][i].qs + 4*ib32;
2749+
uint32_t s32 = (q4[0] & 0x00010001) | ((q4[1] & 0x00010001) << 2) | ((q4[2] & 0x00010001) << 4) | ((q4[3] & 0x00010001) << 6);
2750+
uint8_t s8 = (s32 | (s32 >> 15)) & 0xff;
2751+
//auto val_s = _mm_madd_epi16(_mm_and_si128(val, _mm_set1_epi32(0x00010001)), _mm_set1_epi64x(0x0008000400020001));
2752+
ls[2*ib32+0] = ls[2*ib32+1] = ((s8 & 254) - 127);
2753+
xv[ib32] = _mm256_shuffle_epi8(values[s8 & 1], xv[ib32]);
2754+
}
2755+
dnew[k] = convert_to_q8_k_r8(k, 1.f/127, xv, ls, block, y[i].qs);
2756+
}
2757+
_mm_storeu_si128((__m128i *)y[i].d, _mm256_cvtps_ph(_mm256_mul_ps(vd, _mm256_loadu_ps(dnew)), _MM_ROUND_NEAREST));
2758+
}
2759+
y += nb;
2760+
}
2761+
}
2762+
27092763
void iqk_convert_iq4_ks_q8_k_r8(int n, const void * vx, size_t bx, void * vy, int nrc_x) {
27102764
GGML_ASSERT(n%QK_K == 0);
27112765
GGML_ASSERT(nrc_x%8 == 0);
@@ -3132,6 +3186,7 @@ bool iqk_convert_iqk_quants_q80_r8(int type, int n, const void * vx, size_t bx,
31323186
case GGML_TYPE_IQ2_KL : iqk_convert_iq2_kl_q8_k_r8(n, vx, bx, vy, nrc_x); break;
31333187
case GGML_TYPE_IQ3_KS : iqk_convert_iq3_ks_q8_k_r8(n, vx, bx, vy, nrc_x); break;
31343188
case GGML_TYPE_IQ3_K : iqk_convert_iq3_k_q8_k_r8 (n, vx, bx, vy, nrc_x); break;
3189+
case GGML_TYPE_IQ4_KSS: iqk_convert_iq4_kss_q8_k_r8(n, vx, bx, vy, nrc_x); break;
31353190
case GGML_TYPE_IQ4_KS : iqk_convert_iq4_ks_q8_k_r8(n, vx, bx, vy, nrc_x); break;
31363191
case GGML_TYPE_IQ4_K : iqk_convert_iq4_k_q8_k_r8 (n, vx, bx, vy, nrc_x); break;
31373192
case GGML_TYPE_IQ5_KS : iqk_convert_iq5_ks_q8_k_r8(n, vx, bx, vy, nrc_x); break;
@@ -4718,6 +4773,57 @@ void iqk_convert_iq2_kl_q8_k_r8(int n, const void * vx, size_t bx, void * vy, in
47184773
}
47194774
}
47204775

4776+
void iqk_convert_iq4_kss_q8_k_r8(int n, const void * vx, size_t bx, void * vy, int nrc_x) {
4777+
GGML_ASSERT(n%QK_K == 0);
4778+
GGML_ASSERT(nrc_x%8 == 0);
4779+
4780+
int nb = n/QK_K;
4781+
4782+
const block_iq4_kss * x8[8];
4783+
4784+
block_q8_k_r8 * y = (block_q8_k_r8 *)vy;
4785+
4786+
auto values = vld1q_s8_x2(iq4k_values);
4787+
4788+
float drow[8];
4789+
float dnew[8];
4790+
int8_t ls[16];
4791+
4792+
int8x16x2_t xv[8];
4793+
uint32_t block[8];
4794+
4795+
for (int ix = 0; ix < nrc_x; ix += 8) {
4796+
for (int k = 0; k < 8; ++k) {
4797+
const float * dptr = (const float *)((const char *)vx + (ix + k)*bx);
4798+
drow[k] = dptr[0];
4799+
x8[k] = (const block_iq4_kss *)(dptr + 1);
4800+
}
4801+
auto vd = vld1q_f32_x2(drow);
4802+
for (int i = 0; i < nb; ++i) {
4803+
for (int k = 0; k < 8; ++k) {
4804+
for (int ib32 = 0; ib32 < 8; ++ib32) {
4805+
auto q4 = x8[k][i].qs + 4*ib32;
4806+
uint32_t s32 = (q4[0] & 0x00010001) | ((q4[1] & 0x00010001) << 2) | ((q4[2] & 0x00010001) << 4) | ((q4[3] & 0x00010001) << 6);
4807+
uint8_t s8 = (s32 | (s32 >> 15)) & 0xff;
4808+
ls[2*ib32+0] = ls[2*ib32+1] = ((s8 & 254) - 127);
4809+
auto val16 = vandq_u16(vld1q_u16((const uint16_t *)q4), vdupq_n_u16(0xfffe));
4810+
auto val8 = vreinterpretq_u8_u16(veorq_u16(val16, vshrq_n_u16(val16, 1)));
4811+
auto& block_values = values.val[s8 & 1];
4812+
xv[ib32].val[0] = vqtbl1q_s8(block_values, vandq_u8(val8, vdupq_n_u8(0xf)));
4813+
xv[ib32].val[1] = vqtbl1q_s8(block_values, vshrq_n_u8(val8, 4));
4814+
}
4815+
dnew[k] = convert_to_q8_k_r8(1.f/127, xv, ls, block, (uint32_t *)y[i].qs + k);
4816+
}
4817+
auto d = vld1q_f32_x2(dnew);
4818+
d.val[0] = vmulq_f32(d.val[0], vd.val[0]);
4819+
d.val[1] = vmulq_f32(d.val[1], vd.val[1]);
4820+
vst1_f16((float16_t *)y[i].d + 0, vcvt_f16_f32(d.val[0]));
4821+
vst1_f16((float16_t *)y[i].d + 4, vcvt_f16_f32(d.val[1]));
4822+
}
4823+
y += nb;
4824+
}
4825+
}
4826+
47214827
void iqk_convert_iq4_ks_q8_k_r8(int n, const void * vx, size_t bx, void * vy, int nrc_x) {
47224828
GGML_ASSERT(n%QK_K == 0);
47234829
GGML_ASSERT(nrc_x%8 == 0);
@@ -5163,6 +5269,7 @@ bool iqk_convert_iqk_quants_q80_r8(int type, int n, const void * vx, size_t bx,
51635269
case GGML_TYPE_IQ2_KL : iqk_convert_iq2_kl_q8_k_r8(n, vx, bx, vy, nrc_x); break;
51645270
case GGML_TYPE_IQ3_KS : iqk_convert_iq3_ks_q8_k_r8(n, vx, bx, vy, nrc_x); break;
51655271
case GGML_TYPE_IQ3_K : iqk_convert_iq3_k_q8_k_r8 (n, vx, bx, vy, nrc_x); break;
5272+
case GGML_TYPE_IQ4_KSS: iqk_convert_iq4_kss_q8_k_r8(n, vx, bx, vy, nrc_x); break;
51665273
case GGML_TYPE_IQ4_KS : iqk_convert_iq4_ks_q8_k_r8(n, vx, bx, vy, nrc_x); break;
51675274
case GGML_TYPE_IQ4_K : iqk_convert_iq4_k_q8_k_r8 (n, vx, bx, vy, nrc_x); break;
51685275
case GGML_TYPE_IQ5_KS : iqk_convert_iq5_ks_q8_k_r8(n, vx, bx, vy, nrc_x); break;

‎ggml/src/iqk/iqk_mul_mat.cpp‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -255,6 +255,7 @@ struct MulMat {
255255
case GGML_TYPE_IQ3_KS : return nrc_y >= 32 ? GGML_TYPE_Q8_K_R8 : type;
256256
case GGML_TYPE_IQ3_K : return nrc_y >= 32 ? GGML_TYPE_Q8_K_R8 : type;
257257
case GGML_TYPE_IQ4_KS : return nrc_y >= 32 ? GGML_TYPE_Q8_K_R8 : type;
258+
case GGML_TYPE_IQ4_KSS: return nrc_y >= 32 ? GGML_TYPE_Q8_K_R8 : type;
258259
case GGML_TYPE_IQ4_K : return nrc_y >= 32 ? GGML_TYPE_Q8_K_R8 : type;
259260
case GGML_TYPE_IQ5_KS : return nrc_y >= 32 ? GGML_TYPE_Q8_K_R8 : type;
260261
case GGML_TYPE_IQ5_K : return nrc_y >= 32 ? GGML_TYPE_Q8_K_R8 : type;
@@ -301,6 +302,7 @@ struct MulMat {
301302
case GGML_TYPE_IQ2_KS : return nrc_y >= 32 ? GGML_TYPE_Q8_K_R8 : type;
302303
case GGML_TYPE_IQ2_KL : return nrc_y >= 32 ? GGML_TYPE_Q8_K_R8 : type;
303304
case GGML_TYPE_IQ3_KS : return nrc_y >= 32 ? GGML_TYPE_Q8_K_R8 : type;
305+
case GGML_TYPE_IQ4_KSS: return nrc_y >= 32 ? GGML_TYPE_Q8_K_R8 : type;
304306
case GGML_TYPE_IQ4_KS : return nrc_y >= 32 ? GGML_TYPE_Q8_K_R8 : type;
305307
case GGML_TYPE_IQ5_KS : return nrc_y >= 32 ? GGML_TYPE_Q8_K_R8 : type;
306308
case GGML_TYPE_IQ2_K : return nrc_y >= 32 ? GGML_TYPE_Q8_K_R8 : type;

‎ggml/src/iqk/iqk_quantize.cpp‎

Lines changed: 10 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -4521,25 +4521,17 @@ uint16_t prune_iq4ks(uint16_t v, const int8_t * values, const float * x, const f
45214521
q4[j] = q;
45224522
auto pc = popcount(q);
45234523
float diff0 = dl*iq4k_values[q] - x[j];
4524-
if (q > 0) {
4525-
uint8_t qm = q - 1u;
4526-
int pcm = popcount(qm);
4527-
if (pcm == pc-1 || pcm == pc+1) {
4528-
float diff1 = dl*values[qm] - x[j];
4524+
int qmin = std::max(int(q)-2, 0);
4525+
int qmax = std::min(int(q)+2, 15);
4526+
for (int iq = qmin; iq <= qmax; ++iq) {
4527+
uint8_t qq = iq;
4528+
if (qq == q) continue;
4529+
int pci = popcount(qq);
4530+
if (std::abs(pci - pc)%2) {
4531+
float diff1 = dl*values[qq] - x[j];
45294532
float score = w[j]*(diff1*diff1 - diff0*diff0);
45304533
if (score < best_score) {
4531-
best_score = score; jbest = j; bestq = qm;
4532-
}
4533-
}
4534-
}
4535-
if (q < 15) {
4536-
uint8_t qp = q + 1u;
4537-
int pcp = popcount(qp);
4538-
if (pcp == pc-1 || pcp == pc+1) {
4539-
float diff1 = dl*values[qp] - x[j];
4540-
float score = w[j]*(diff1*diff1 - diff0*diff0);
4541-
if (score < best_score) {
4542-
best_score = score; jbest = j; bestq = qp;
4534+
best_score = score; jbest = j; bestq = qq;
45434535
}
45444536
}
45454537
}
@@ -4760,7 +4752,7 @@ static void quantize_row_iq4_kss_impl(int n_per_row, const float * x, char * cy,
47604752
}
47614753
}
47624754
}
4763-
if (sumq2 > 0) *dptr = sumqx/sumq2;
4755+
if (sumq2 > 0) *dptr = sumqx/sumq2 * 1.01f;
47644756
}
47654757

47664758
void prune_iq4ks_to_iq4kss(int n_per_row, const uint16_t * table, const char * cx, const float * x, char *cy,

0 commit comments

Comments
 (0)