diff --git a/ggml/src/ggml-cuda/convert.cu b/ggml/src/ggml-cuda/convert.cu index fd43a537..26e7e532 100644 --- a/ggml/src/ggml-cuda/convert.cu +++ b/ggml/src/ggml-cuda/convert.cu @@ -821,15 +821,18 @@ static __global__ void dequantize_block_iq6_k(const void * __restrict__ vx, dst_ } template -static __global__ void dequantize_block_iq2_k(const void * __restrict__ vx, dst_t * __restrict__ yy) { +static __global__ void dequantize_block_iq2_k(const void * __restrict__ vx, dst_t * __restrict__ yy, int64_t n_per_row, int64_t row_size) { - const int i = blockIdx.x; - const block_iq2_k * x = (const block_iq2_k *) vx; + int64_t ii = blockIdx.x; + int64_t row = (QK_K * ii) / n_per_row; + const int8_t * row_values = (const int8_t *)vx + row * row_size; + const block_iq2_k * x = (const block_iq2_k *)(row_values + 8); + const int64_t i = ii - (row*n_per_row)/QK_K; const int tid = threadIdx.x; int ib128 = tid/16; // 0 or 1 int il = tid%16; // 0...15 - dst_t * y = yy + i*QK_K + 128*ib128 + 2*il; + dst_t * y = yy + ii*QK_K + 128*ib128 + 2*il; const float d = (float)x[i].d; const float dl1 = d * (((x[i].scales[4*ib128+0] >> 4*(il/8)) & 0xf) - 8); const float dl2 = d * (((x[i].scales[4*ib128+1] >> 4*(il/8)) & 0xf) - 8); @@ -838,10 +841,10 @@ static __global__ void dequantize_block_iq2_k(const void * __restrict__ vx, dst_ const uint8_t * qs = x[i].qs + 32*ib128 + 2*il; const int16_t extra = x[i].extra >> (8*ib128 + (il/8)); for (int j = 0; j < 2; ++j) { - y[j+ 0] = dl1 * iq2nl_values[((qs[j] >> 0) & 0x03) + ((extra << 2) & 4)]; - y[j+32] = dl2 * iq2nl_values[((qs[j] >> 2) & 0x03) + ((extra << 0) & 4)]; - y[j+64] = dl3 * iq2nl_values[((qs[j] >> 4) & 0x03) + ((extra >> 2) & 4)]; - y[j+96] = dl4 * iq2nl_values[((qs[j] >> 6) & 0x03) + ((extra >> 4) & 4)]; + y[j+ 0] = dl1 * row_values[((qs[j] >> 0) & 0x03) + ((extra << 2) & 4)]; + y[j+32] = dl2 * row_values[((qs[j] >> 2) & 0x03) + ((extra << 0) & 4)]; + y[j+64] = dl3 * row_values[((qs[j] >> 4) & 0x03) + ((extra >> 2) & 4)]; + y[j+96] = dl4 * row_values[((qs[j] >> 6) & 0x03) + ((extra >> 4) & 4)]; } } @@ -1113,7 +1116,8 @@ template static void dequantize_row_iq2_k_cuda(const void * vx, dst_t * y, const int64_t nrows, const int64_t n_per_row, cudaStream_t stream) { const int64_t k = nrows * n_per_row; const int nb = (k + QK_K - 1) / QK_K; - dequantize_block_iq2_k<<>>(vx, y); + const int64_t row_size = ggml_row_size(GGML_TYPE_IQ2_K, n_per_row); + dequantize_block_iq2_k<<>>(vx, y, n_per_row, row_size); } template diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c index 48dd2301..a636245d 100644 --- a/ggml/src/ggml.c +++ b/ggml/src/ggml.c @@ -1178,7 +1178,7 @@ static const ggml_type_traits_t type_traits[GGML_TYPE_COUNT] = { .vec_dot = vec_dot_iq2_k_q8_k, .vec_dot_type = GGML_TYPE_Q8_K, .nrows = 1, - .row_meta_size = 0, + .row_meta_size = 8, }, [GGML_TYPE_IQ2_KS] = { .type_name = "iq2_ks", diff --git a/ggml/src/iqk/iqk_quantize.cpp b/ggml/src/iqk/iqk_quantize.cpp index 24cd057e..3f62b85b 100644 --- a/ggml/src/iqk/iqk_quantize.cpp +++ b/ggml/src/iqk/iqk_quantize.cpp @@ -533,21 +533,199 @@ inline int best_index_iq2nl(const int8_t * values, float x) { return x - values[idx] < values[idx+1] - x ? idx : idx + 1; } -void quantize_row_iq2_k_impl(const float * x, void * vy, int n_per_row, const float * quant_weights) { +#ifdef IQ2K_STATS +struct IQ2KCollector { + std::array m_counts = {}; + std::array< int64_t, 8> m_values = {}; + int64_t m_nval = 0; + std::mutex m_mutex; +public: + IQ2KCollector() = default; + ~IQ2KCollector() { + printf("============== %s: bin counts:\n", __func__); + for (int j = 0; j < 8; ++j) printf("%d %g\n", j, 1.*m_counts[j]); + if (m_nval > 0) { + printf("============== %s: bin values from %g calls:\n", __func__, 1.*m_nval); + double norm = 1./m_nval; + for (int j = 0; j < 8; ++j) printf("%d %g\n", j, norm*m_values[j]); + } + } + void add(const int * counts, bool is_shifted) { + std::lock_guard lock(m_mutex); + int offset = is_shifted ? 4 : 0; + for (int j = 0; j < 4; ++j) m_counts[offset + j] += counts[j]; + } + void add_values(const int8_t * values) { + std::lock_guard lock(m_mutex); + for (int j = 0; j < 8; ++j) m_values[j] += values[j]; + ++m_nval; + } +}; + +IQ2KCollector& get_iq2k_collector() { + static std::mutex mutex; + std::lock_guard lock(mutex); + static IQ2KCollector collector; + return collector; +} +#endif + +void quantize_row_iq2_k_impl(const float * x, void * vy, int n_per_row, const float * quant_weights, std::vector>& all_pairs) { constexpr int kBlockSize = 16; - block_iq2_k * y = (block_iq2_k *)vy; + int8_t * row_values = (int8_t *)vy; + block_iq2_k * y = (block_iq2_k *)(row_values + 8); + + for (int j = 0; j < 8; ++j) row_values[j] = iq2nl_values[j]; float scales[QK_K/kBlockSize]; float weight[kBlockSize]; float sumx[kBlockSize+1], sumw[kBlockSize+1]; uint8_t L[QK_K]; - std::array, kBlockSize> pairs; + const int8_t * shifted_values = row_values + 4; - const int8_t * shifted_values = iq2nl_values + 4; + float sx[8], sw[8]; +#ifdef IQ2K_STATS + auto& collector = get_iq2k_collector(); +#endif + + if (int(all_pairs.size()) < n_per_row) all_pairs.resize(n_per_row); + + for (int ibl = 0; ibl < n_per_row/QK_K; ++ibl) { + const float * xbl = x + ibl*QK_K; + for (int ib = 0; ib < QK_K/kBlockSize; ++ib) { + const float * xb = xbl + kBlockSize*ib; + auto pairs = all_pairs.data() + ibl*QK_K + ib*kBlockSize; + for (int j = 0; j < kBlockSize; ++j) { + pairs[j] = {xb[j], j}; + } + std::sort(pairs, pairs + kBlockSize); + } + } + + for (int itry = 0; itry < 3; ++itry) { + std::memset(sx, 0, 8*sizeof(float)); + std::memset(sw, 0, 8*sizeof(float)); + for (int ibl = 0; ibl < n_per_row/QK_K; ++ibl) { + const float * xbl = x + ibl*QK_K; + float sumx2 = 0; + for (int j = 0; j < QK_K; ++j) sumx2 += xbl[j]*xbl[j]; + const float sigma2 = 1.5f*sumx2/QK_K; + + for (int ib = 0; ib < QK_K/kBlockSize; ++ib) { + const float * xb = xbl + kBlockSize*ib; + if (quant_weights) { + const float * qw = quant_weights + ibl*QK_K + ib*kBlockSize; + for (int j = 0; j < kBlockSize; ++j) weight[j] = qw[j] * sqrtf(sigma2 + xb[j]*xb[j]); + } else { + for (int j = 0; j < kBlockSize; ++j) weight[j] = 0.25f*sigma2 + xb[j]*xb[j]; + } + auto pairs = all_pairs.data() + ibl*QK_K + ib*kBlockSize; + sumx[0] = sumw[0] = 0; + for (int j = 0; j < kBlockSize; ++j) { + int jj = pairs[j].second; + sumw[j+1] = sumw[j] + weight[jj]; + sumx[j+1] = sumx[j] + weight[jj]*xb[jj]; + } + float best = 0, d = 0; + bool is_shifted = false; + int besti1 = -1, besti2 = -1, besti3 = -1; + bool reverse = false; + float sumqx, sumq2; + for (int i1 = 0; i1 < kBlockSize; ++i1) { + for (int i2 = i1; i2 < kBlockSize; ++i2) { + for (int i3 = i2; i3 < kBlockSize; ++i3) { + sumqx = (sumx[i1] - sumx[ 0])*row_values[0] + (sumx[i2] - sumx[i1])*row_values[1] + + (sumx[i3] - sumx[i2])*row_values[2] + (sumx[kBlockSize] - sumx[i3])*row_values[3]; + sumq2 = (sumw[i1] - sumw[ 0])*row_values[0]*row_values[0] + (sumw[i2] - sumw[i1])*row_values[1]*row_values[1] + + (sumw[i3] - sumw[i2])*row_values[2]*row_values[2] + (sumw[kBlockSize] - sumw[i3])*row_values[3]*row_values[3]; + if (sumq2 > 0 && sumqx*sumqx > best*sumq2) { + besti1 = i1; besti2 = i2; besti3 = i3; reverse = false; + d = sumqx/sumq2; best = d*sumqx; is_shifted = false; + } + sumqx = (sumx[i1] - sumx[ 0])*shifted_values[0] + (sumx[i2] - sumx[i1])*shifted_values[1] + + (sumx[i3] - sumx[i2])*shifted_values[2] + (sumx[kBlockSize] - sumx[i3])*shifted_values[3]; + sumq2 = (sumw[i1] - sumw[ 0])*shifted_values[0]*shifted_values[0] + (sumw[i2] - sumw[i1])*shifted_values[1]*shifted_values[1] + + (sumw[i3] - sumw[i2])*shifted_values[2]*shifted_values[2] + (sumw[kBlockSize] - sumw[i3])*shifted_values[3]*shifted_values[3]; + if (sumq2 > 0 && sumqx*sumqx > best*sumq2) { + besti1 = i1; besti2 = i2; besti3 = i3; reverse = false; + d = sumqx/sumq2; best = d*sumqx; is_shifted = true; + } + sumqx = (sumx[i1] - sumx[ 0])*row_values[3] + (sumx[i2] - sumx[i1])*row_values[2] + + (sumx[i3] - sumx[i2])*row_values[1] + (sumx[kBlockSize] - sumx[i3])*row_values[0]; + sumq2 = (sumw[i1] - sumw[ 0])*row_values[3]*row_values[3] + (sumw[i2] - sumw[i1])*row_values[2]*row_values[2] + + (sumw[i3] - sumw[i2])*row_values[1]*row_values[1] + (sumw[kBlockSize] - sumw[i3])*row_values[0]*row_values[0]; + if (sumq2 > 0 && sumqx*sumqx > best*sumq2) { + besti1 = i1; besti2 = i2; besti3 = i3; reverse = true; + d = sumqx/sumq2; best = d*sumqx; is_shifted = false; + } + sumqx = (sumx[i1] - sumx[ 0])*shifted_values[3] + (sumx[i2] - sumx[i1])*shifted_values[2] + + (sumx[i3] - sumx[i2])*shifted_values[1] + (sumx[kBlockSize] - sumx[i3])*shifted_values[0]; + sumq2 = (sumw[i1] - sumw[ 0])*shifted_values[3]*shifted_values[3] + (sumw[i2] - sumw[i1])*shifted_values[2]*shifted_values[2] + + (sumw[i3] - sumw[i2])*shifted_values[1]*shifted_values[1] + (sumw[kBlockSize] - sumw[i3])*shifted_values[0]*shifted_values[0]; + if (sumq2 > 0 && sumqx*sumqx > best*sumq2) { + besti1 = i1; besti2 = i2; besti3 = i3; reverse = true; + d = sumqx/sumq2; best = d*sumqx; is_shifted = true; + } + } + } + } + + //printf("Block %d: d = %g besti = %d, %d, %d, reverse = %d, is_shifted = %d\n", ib, d, besti1, besti2, besti3, reverse, is_shifted); + if (!d) continue; + + float id = 1/d; + + auto sxb = is_shifted ? sx + 4 : sx; + auto swb = is_shifted ? sw + 4 : sw; + int idx = reverse ? 3 : 0; + for (int j = 0; j < besti1; ++j) { + int jj = pairs[j].second; + sxb[idx] += weight[jj]*id*xb[jj]; + swb[idx] += weight[jj]; + } + idx = reverse ? 2 : 1; + for (int j = besti1; j < besti2; ++j) { + int jj = pairs[j].second; + sxb[idx] += weight[jj]*id*xb[jj]; + swb[idx] += weight[jj]; + } + idx = reverse ? 1 : 2; + for (int j = besti2; j < besti3; ++j) { + int jj = pairs[j].second; + sxb[idx] += weight[jj]*id*xb[jj]; + swb[idx] += weight[jj]; + } + idx = reverse ? 0 : 3; + for (int j = besti3; j < kBlockSize; ++j) { + int jj = pairs[j].second; + sxb[idx] += weight[jj]*id*xb[jj]; + swb[idx] += weight[jj]; + } + } + + } + + bool changed = false; + for (int j = 0; j < 8; ++j) { + float val = sw[j] > 0 ? sx[j]/sw[j] : iq2nl_values[j]; + //printf("Updated row value %d: %d -> %g (%g, %g)\n", j, row_values[j], val, sx[j], sw[j]); + int new_value = std::max(-48, std::min(48, nearest_int(val))); + if (row_values[j] != new_value) changed = true; + row_values[j] = new_value; + } + if (!changed) break; + + } + +#ifdef IQ2K_STATS + collector.add_values(row_values); +#endif + + std::array counts; for (int ibl = 0; ibl < n_per_row/QK_K; ++ibl) { memset(&y[ibl], 0, sizeof(block_iq2_k)); @@ -570,10 +748,7 @@ void quantize_row_iq2_k_impl(const float * x, void * vy, int n_per_row, const fl } else { for (int j = 0; j < kBlockSize; ++j) weight[j] = 0.25f*sigma2 + xb[j]*xb[j]; } - for (int j = 0; j < kBlockSize; ++j) { - pairs[j] = {xb[j], j}; - } - std::sort(pairs.begin(), pairs.end()); + auto pairs = all_pairs.data() + ibl*QK_K + ib*kBlockSize; sumx[0] = sumw[0] = 0; for (int j = 0; j < kBlockSize; ++j) { int jj = pairs[j].second; @@ -588,10 +763,10 @@ void quantize_row_iq2_k_impl(const float * x, void * vy, int n_per_row, const fl for (int i1 = 0; i1 < kBlockSize; ++i1) { for (int i2 = i1; i2 < kBlockSize; ++i2) { for (int i3 = i2; i3 < kBlockSize; ++i3) { - sumqx = (sumx[i1] - sumx[ 0])*iq2nl_values[0] + (sumx[i2] - sumx[i1])*iq2nl_values[1] - + (sumx[i3] - sumx[i2])*iq2nl_values[2] + (sumx[kBlockSize] - sumx[i3])*iq2nl_values[3]; - sumq2 = (sumw[i1] - sumw[ 0])*iq2nl_values[0]*iq2nl_values[0] + (sumw[i2] - sumw[i1])*iq2nl_values[1]*iq2nl_values[1] - + (sumw[i3] - sumw[i2])*iq2nl_values[2]*iq2nl_values[2] + (sumw[kBlockSize] - sumw[i3])*iq2nl_values[3]*iq2nl_values[3]; + sumqx = (sumx[i1] - sumx[ 0])*row_values[0] + (sumx[i2] - sumx[i1])*row_values[1] + + (sumx[i3] - sumx[i2])*row_values[2] + (sumx[kBlockSize] - sumx[i3])*row_values[3]; + sumq2 = (sumw[i1] - sumw[ 0])*row_values[0]*row_values[0] + (sumw[i2] - sumw[i1])*row_values[1]*row_values[1] + + (sumw[i3] - sumw[i2])*row_values[2]*row_values[2] + (sumw[kBlockSize] - sumw[i3])*row_values[3]*row_values[3]; if (sumq2 > 0 && sumqx*sumqx > best*sumq2) { besti1 = i1; besti2 = i2; besti3 = i3; reverse = false; d = sumqx/sumq2; best = d*sumqx; is_shifted = false; @@ -604,10 +779,10 @@ void quantize_row_iq2_k_impl(const float * x, void * vy, int n_per_row, const fl besti1 = i1; besti2 = i2; besti3 = i3; reverse = false; d = sumqx/sumq2; best = d*sumqx; is_shifted = true; } - sumqx = (sumx[i1] - sumx[ 0])*iq2nl_values[3] + (sumx[i2] - sumx[i1])*iq2nl_values[2] - + (sumx[i3] - sumx[i2])*iq2nl_values[1] + (sumx[kBlockSize] - sumx[i3])*iq2nl_values[0]; - sumq2 = (sumw[i1] - sumw[ 0])*iq2nl_values[3]*iq2nl_values[3] + (sumw[i2] - sumw[i1])*iq2nl_values[2]*iq2nl_values[2] - + (sumw[i3] - sumw[i2])*iq2nl_values[1]*iq2nl_values[1] + (sumw[kBlockSize] - sumw[i3])*iq2nl_values[0]*iq2nl_values[0]; + sumqx = (sumx[i1] - sumx[ 0])*row_values[3] + (sumx[i2] - sumx[i1])*row_values[2] + + (sumx[i3] - sumx[i2])*row_values[1] + (sumx[kBlockSize] - sumx[i3])*row_values[0]; + sumq2 = (sumw[i1] - sumw[ 0])*row_values[3]*row_values[3] + (sumw[i2] - sumw[i1])*row_values[2]*row_values[2] + + (sumw[i3] - sumw[i2])*row_values[1]*row_values[1] + (sumw[kBlockSize] - sumw[i3])*row_values[0]*row_values[0]; if (sumq2 > 0 && sumqx*sumqx > best*sumq2) { besti1 = i1; besti2 = i2; besti3 = i3; reverse = true; d = sumqx/sumq2; best = d*sumqx; is_shifted = false; @@ -663,7 +838,7 @@ void quantize_row_iq2_k_impl(const float * x, void * vy, int n_per_row, const fl } else { for (int j = 0; j < kBlockSize; ++j) weight[j] = 0.25f*sigma2 + xb[j]*xb[j]; } - auto block_values = extra & (1 << ib) ? shifted_values : iq2nl_values; + auto block_values = extra & (1 << ib) ? shifted_values : row_values; int ls = nearest_int(id*scales[ib]); ls = std::max(-8, std::min(7, ls)); for (int j = 0; j < kBlockSize; ++j) { @@ -685,12 +860,13 @@ void quantize_row_iq2_k_impl(const float * x, void * vy, int n_per_row, const fl float sumqx = 0, sumq2 = 0; for (int ib = 0; ib < QK_K/kBlockSize; ++ib) { + counts[0] = counts[1] = counts[2] = counts[3] = 0; int ls = nearest_int(best_id*scales[ib]); ls = std::max(-8, std::min(7, ls)); y[ibl].scales[ib/2] |= ((ls + 8) << 4*(ib%2)); float dl = d * ls; if (dl) { - const int8_t * block_values = y[ibl].extra & (1 << ib) ? shifted_values : iq2nl_values; + const int8_t * block_values = y[ibl].extra & (1 << ib) ? shifted_values : row_values; const float * xb = xbl + kBlockSize*ib; if (quant_weights) { const float * qw = quant_weights + ibl*QK_K + ib*kBlockSize; @@ -705,12 +881,16 @@ void quantize_row_iq2_k_impl(const float * x, void * vy, int n_per_row, const fl for (int j = 0; j < 16; ++j) { const float al = idl*xb[j]; int ibest = best_index_iq2nl(block_values, al); + ++counts[ibest]; qs[j] |= (ibest << 2*(ib32%4)); float w = weight[j]; float q = block_values[ibest]*ls; sumqx += w*q*xb[j]; sumq2 += w*q*q; } +#ifdef IQ2K_STATS + collector.add(counts.data(), y[ibl].extra & (1 << ib) ? true : false); +#endif } } y[ibl].d = GGML_FP32_TO_FP16(1.030f*(sumq2 > 0 ? sumqx/sumq2 : d)); @@ -732,20 +912,24 @@ void quantize_row_iq2_k(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, size_t quantize_iq2_k(const float * src, void * dst, int64_t nrows, int64_t n_per_row, const float * imatrix) { GGML_ASSERT(n_per_row%QK_K == 0); - int nblock = n_per_row/QK_K; + auto row_size = ggml_row_size(GGML_TYPE_IQ2_K, n_per_row); + std::vector> all_pairs(n_per_row); char * qrow = (char *)dst; for (int64_t row = 0; row < nrows; ++row) { - quantize_row_iq2_k_impl(src, (void *)qrow, n_per_row, imatrix); + quantize_row_iq2_k_impl(src, (void *)qrow, n_per_row, imatrix, all_pairs); src += n_per_row; - qrow += nblock*sizeof(block_iq2_k); + qrow += row_size; } - return nrows * nblock * sizeof(block_iq2_k); + return nrows * row_size; } void dequantize_row_iq2_k(const block_iq2_k * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) { assert(k % QK_K == 0); const int nb = k / QK_K; + const int8_t * row_values = (const int8_t *)x; + x = (const block_iq2_k *)(row_values + 8); + for (int i = 0; i < nb; i++) { const float d = GGML_FP16_TO_FP32(x[i].d); @@ -757,8 +941,8 @@ void dequantize_row_iq2_k(const block_iq2_k * GGML_RESTRICT x, float * GGML_RES for (int ib32 = 0; ib32 < QK_K/32; ++ib32) { float dl1 = d * ((x[i].scales[ib32] & 0xf) - 8); float dl2 = d * ((x[i].scales[ib32] >> 4) - 8); - const int8_t * values1 = extra & 1 ? iq2nl_values + 4 : iq2nl_values; - const int8_t * values2 = extra & 2 ? iq2nl_values + 4 : iq2nl_values; + const int8_t * values1 = extra & 1 ? row_values + 4 : row_values; + const int8_t * values2 = extra & 2 ? row_values + 4 : row_values; extra >>= 2; for (int j = 0; j < 16; ++j) { y[j+ 0] = dl1 * values1[(qs[j+ 0] >> shift) & 3];