iq2k improvement

This commit is contained in:
Iwan Kawrakow
2024-11-25 18:19:16 +02:00
parent 85d1011f52
commit 2f749cfdba
3 changed files with 222 additions and 34 deletions

View File

@@ -821,15 +821,18 @@ static __global__ void dequantize_block_iq6_k(const void * __restrict__ vx, dst_
}
template<typename dst_t>
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<typename dst_t>
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<<<nb, 32, 0, stream>>>(vx, y);
const int64_t row_size = ggml_row_size(GGML_TYPE_IQ2_K, n_per_row);
dequantize_block_iq2_k<<<nb, 32, 0, stream>>>(vx, y, n_per_row, row_size);
}
template<typename dst_t>

View File

@@ -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",

View File

@@ -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<uint64_t, 8> 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<std::mutex> 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<std::mutex> 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<std::mutex> 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<std::pair<float, int>>& 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<std::pair<float,int>, 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<int, 4> 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<std::pair<float,int>> 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];