iq4_k_r4: Zen4 and hopefully AVX2

On Zen4 we get PP-512(LLaMA-3.1-8B) = 232.6 t/s, up from 182.2 t/s
for iq4_k. Applying the extra shift costs a ~6 performance penalty.
This commit is contained in:
Iwan Kawrakow
2024-12-12 11:39:29 +02:00
parent fb79167eec
commit ba9a9a1655
3 changed files with 17 additions and 26 deletions

View File

@@ -543,7 +543,7 @@ static_assert(sizeof(block_iq4_k) == sizeof(ggml_half) + sizeof(uint16_t) + QK_K
typedef struct {
ggml_half d[4];
uint16_t extra[4];
uint8_t extra[8];
uint8_t scales_h[QK_K/16];
uint8_t scales_l[QK_K/8];
uint8_t qs[QK_K*2];

View File

@@ -3792,6 +3792,8 @@ static void mul_mat_iq4_k_r4_q8_k(int n, const void * vx, size_t bx, const DataI
auto m4 = _mm256_set1_epi8(0xf);
auto m30 = _mm256_set1_epi8(0x30);
auto m32 = _mm256_set1_epi8(32);
auto ms = _mm256_set1_epi8(4);
auto shift_shuffle = _mm256_set_epi64x(0x0303030302020202, 0x0101010100000000, 0x0303030302020202, 0x0101010100000000);
static const uint8_t k_shuff[32] = {0, 1, 8, 9, 2, 3, 10, 11, 4, 5, 12, 13, 6, 7, 14, 15, 0, 1, 8, 9, 2, 3, 10, 11, 4, 5, 12, 13, 6, 7, 14, 15};
// r0b0, r0b0, r1b0, r1b0, r2b0, r2b0, r3b0, r3b0, r0b0, r0b0, r1b0, r1b0, r2b0, r2b0, r3b0, r3b0,
auto shuff = _mm256_loadu_si256((const __m256i *)k_shuff);
@@ -3808,13 +3810,12 @@ static void mul_mat_iq4_k_r4_q8_k(int n, const void * vx, size_t bx, const DataI
__m256 acc[nrc_y] = {};
__m256i qx[4];
int8_t stored_scales[64];
//uint64_t aux64;
//uint8_t * aux8 = (uint8_t *)&aux64;
for (int ix = 0; ix < nrc_x; ix += 4) {
const block_iq4_k_r4 * iq4 = (const block_iq4_k_r4 *)((const char *)vx + (ix+0)*bx);
for (int ibl = 0; ibl < nbl; ++ibl) { // Block of 256
auto dl = _mm_cvtph_ps(_mm_loadl_epi64((const __m128i *)iq4[ibl].d));
auto d4 = _mm256_set_m128(dl, dl);
auto extra = _mm256_set1_epi64x(*(const uint64_t *)iq4[ibl].extra);
#ifdef HAVE_FANCY_SIMD
for (int iy = 0; iy < nrc_y; ++iy) {
d4s[iy] = _mm256_mul_ps(d4, _mm256_set1_ps(q8.scale(iy, ibl)));
@@ -3839,7 +3840,6 @@ static void mul_mat_iq4_k_r4_q8_k(int n, const void * vx, size_t bx, const DataI
auto t2 = _mm256_shuffle_epi8(_mm256_cvtepi8_epi16(_mm256_extracti128_si256(i8scales1, 1)), shuff); // blocks 4, 5, 6, 7 for each row
auto t3 = _mm256_shuffle_epi8(_mm256_cvtepi8_epi16(_mm256_extracti128_si256(i8scales2, 0)), shuff); // blocks 8, 9, 10, 11 for each row
auto t4 = _mm256_shuffle_epi8(_mm256_cvtepi8_epi16(_mm256_extracti128_si256(i8scales2, 1)), shuff); // blocks 12, 13, 14, 15 for each row
//auto scales_s = _mm256_mullo_epi16(scales16, _mm256_mask_add_epi16(min, extra, min, eshift));
auto s1 = MM256_SET_M128I(_mm256_extracti128_si256(t3, 0), _mm256_extracti128_si256(t1, 0)); // blocks 0, 1, 8, 9
auto s2 = MM256_SET_M128I(_mm256_extracti128_si256(t3, 1), _mm256_extracti128_si256(t1, 1)); // blocks 2, 3, 10, 11
auto s3 = MM256_SET_M128I(_mm256_extracti128_si256(t4, 0), _mm256_extracti128_si256(t2, 0)); // blocks 4, 5, 12, 13
@@ -3864,22 +3864,12 @@ static void mul_mat_iq4_k_r4_q8_k(int n, const void * vx, size_t bx, const DataI
#endif
auto bits1 = _mm256_loadu_si256((const __m256i *)iq4[ibl].qs+2*ib+0);
auto bits2 = _mm256_loadu_si256((const __m256i *)iq4[ibl].qs+2*ib+1);
//for (int k = 0; k < 4; ++k) {
// auto extra = iq4[ibl].extra[k] >> 2*ib;
// aux8[k+0] = (extra & 1) << 2;
// aux8[k+4] = (extra & 2) << 1;
//}
//auto shift = _mm256_set1_epi64x(aux64);
//// TODO: extra
//qx[0] = _mm256_add_epi8(shift, _mm256_shuffle_epi8(values, _mm256_and_si256(bits1, m4)));
//qx[1] = _mm256_add_epi8(shift, _mm256_shuffle_epi8(values, _mm256_and_si256(bits2, m4)));
//qx[2] = _mm256_add_epi8(shift, _mm256_shuffle_epi8(values, _mm256_and_si256(_mm256_srli_epi16(bits1, 4), m4)));
//qx[3] = _mm256_add_epi8(shift, _mm256_shuffle_epi8(values, _mm256_and_si256(_mm256_srli_epi16(bits2, 4), m4)));
// TODO: extra
qx[0] = _mm256_shuffle_epi8(values, _mm256_and_si256(bits1, m4));
qx[1] = _mm256_shuffle_epi8(values, _mm256_and_si256(bits2, m4));
qx[2] = _mm256_shuffle_epi8(values, _mm256_and_si256(_mm256_srli_epi16(bits1, 4), m4));
qx[3] = _mm256_shuffle_epi8(values, _mm256_and_si256(_mm256_srli_epi16(bits2, 4), m4));
auto shift = _mm256_and_si256(ms, _mm256_slli_epi16(extra, 2)); extra = _mm256_srli_epi16(extra, 1);
shift = _mm256_shuffle_epi8(shift, shift_shuffle);
qx[0] = _mm256_add_epi8(shift, _mm256_shuffle_epi8(values, _mm256_and_si256(bits1, m4)));
qx[1] = _mm256_add_epi8(shift, _mm256_shuffle_epi8(values, _mm256_and_si256(bits2, m4)));
qx[2] = _mm256_add_epi8(shift, _mm256_shuffle_epi8(values, _mm256_and_si256(_mm256_srli_epi16(bits1, 4), m4)));
qx[3] = _mm256_add_epi8(shift, _mm256_shuffle_epi8(values, _mm256_and_si256(_mm256_srli_epi16(bits2, 4), m4)));
#ifndef HAVE_FANCY_SIMD
auto s1 = _mm256_sign_epi8(qx[0], qx[0]);
auto s2 = _mm256_sign_epi8(qx[1], qx[1]);
@@ -3896,7 +3886,6 @@ static void mul_mat_iq4_k_r4_q8_k(int n, const void * vx, size_t bx, const DataI
sumi = _mm256_dpbusd_epi32(sumi, qx[3], _mm256_shuffle_epi32(y, 0xff));
acc[iy] = _mm256_fmadd_ps(_mm256_mul_ps(scales, d4s[iy]), _mm256_cvtepi32_ps(sumi), acc[iy]);
#else
// This is wrong! (it may overflow the int16_t range
auto sumi1 = _mm256_maddubs_epi16(s1, _mm256_sign_epi8(_mm256_shuffle_epi32(y, 0x00), qx[0]));
auto sumi2 = _mm256_maddubs_epi16(s2, _mm256_sign_epi8(_mm256_shuffle_epi32(y, 0x55), qx[1]));
auto sumi3 = _mm256_maddubs_epi16(s3, _mm256_sign_epi8(_mm256_shuffle_epi32(y, 0xaa), qx[2]));

View File

@@ -4572,12 +4572,16 @@ static void repack_iq4_k(int nrows, int n_per_row, const block_iq4_k * x, block_
for (int row = 0; row < nrows; row += 4) {
for (int k = 0; k < 4; ++k) x4[k] = x + nblock*k;
for (int ibl = 0; ibl < nblock; ++ibl) {
std::memset(y[ibl].extra, 0, 8);
std::memset(y[ibl].scales_l, 0, QK_K/8);
std::memset(y[ibl].scales_h, 0, QK_K/16);
for (int k = 0; k < 4; ++k) {
y[ibl].d[k] = x4[k][ibl].d;
y[ibl].extra[k] = x4[k][ibl].extra;
auto extra = x4[k][ibl].extra;
for (int ib = 0; ib < QK_K/32; ++ib) {
if (extra & 1) y[ibl].extra[k+0] |= (1 << ib);
if (extra & 2) y[ibl].extra[k+4] |= (1 << ib);
extra >>= 2;
uint8_t sl1 = x4[k][ibl].scales_l[ib] & 0xf;
uint8_t sl2 = x4[k][ibl].scales_l[ib] >> 4;
uint8_t sh = x4[k][ibl].scales_h[ib/2] >> 4*(ib%2);
@@ -4627,15 +4631,13 @@ void dequantize_row_iq4_k_r4(const block_iq4_k_r4 * x, float * y, int64_t k) {
for (int ibl = 0; ibl < nblock; ++ibl) {
for (int k = 0; k < 4; ++k) {
const float d = GGML_FP16_TO_FP32(x[ibl].d[k]);
auto extra = x[ibl].extra[k];
for (int ib = 0; ib < QK_K/32; ++ib) {
int is = 8*ib + k;
float dl1 = d * ((((x[ibl].scales_l[is%32] >> 4*(is/32)) & 0xf) | (((x[ibl].scales_h[is%16] >> 2*(is/16)) & 3) << 4)) - 32);
is += 4;
float dl2 = d * ((((x[ibl].scales_l[is%32] >> 4*(is/32)) & 0xf) | (((x[ibl].scales_h[is%16] >> 2*(is/16)) & 3) << 4)) - 32);
auto values1 = iq4k_values + ((extra & 1) << 4);
auto values2 = iq4k_values + ((extra & 2) << 3);
extra >>= 2;
auto values1 = iq4k_values + (x[ibl].extra[k+0] & (1 << ib) ? 16 : 0);
auto values2 = iq4k_values + (x[ibl].extra[k+4] & (1 << ib) ? 16 : 0);
for (int i = 0; i < 4; ++i) {
y4[k][QK_K*ibl+32*ib+i+ 0] = dl1 * values1[x[ibl].qs[64*ib+4*k+i+ 0] & 0xf];
y4[k][QK_K*ibl+32*ib+i+ 8] = dl1 * values1[x[ibl].qs[64*ib+4*k+i+ 0] >> 4];