Adding IQ3_KS quants (#566)

* iq3_ks: basics

* iq3_ks: CUDA dequantize

* iq3_ks: CUDA mmvq

* iq3_ks: mmq

* iq3_ks: faster mmq

* iq3_ks: Zen4

* iq3_ks: AVX2 convert to q8_k_r8

This gives usPP-512 = 360 t/s.

* iq3_ks: AVX2 GEMM/GEMV

* iq3_ks: NEON GEMM/GEMV

* iq3_ks: NEON convert to q8_k_r8

This gives us PP-512 = 164 t/s.

* iq3_ks: Metal dequantize

* iq3_ks: Metal gemv - pathetic performance

---------

Co-authored-by: Iwan Kawrakow <[email protected]>

Author: ikawrakow

examples/quantize/quantize.cpp

@@ -76,6 +76,7 @@ static const std::vector<struct quant_option> QUANT_OPTIONS = {
     { "IQ2_K_R4", LLAMA_FTYPE_MOSTLY_IQ2_K_R4, "IQ2_K repacked",},
     { "IQ2_KS",   LLAMA_FTYPE_MOSTLY_IQ2_KS,   " 2.1875 bpw non-linear quantization",},
     { "IQ2_KT",   LLAMA_FTYPE_MOSTLY_IQ2_KT,   " 2.125 bpw trellis quantization",   },
+    { "IQ3_KS",   LLAMA_FTYPE_MOSTLY_IQ3_KS,   " 3.19 bpw non-linear quantization", },
     { "IQ3_K",    LLAMA_FTYPE_MOSTLY_IQ3_K,    " 3.44 bpw non-linear quantization", },
     { "IQ3_K_R4", LLAMA_FTYPE_MOSTLY_IQ3_K_R4, "IQ3_K repacked", },
     { "IQ3_KL",   LLAMA_FTYPE_MOSTLY_IQ3_KL,   " 4 bpw non-linear quantization mix",},

ggml/include/ggml.h

@@ -429,6 +429,7 @@ extern "C" {
         GGML_TYPE_IQ2_KT  = 153,
         GGML_TYPE_IQ3_KT  = 154,
         GGML_TYPE_IQ4_KT  = 155,
+        GGML_TYPE_IQ3_KS  = 156,
 
         GGML_TYPE_Q4_0_R8   = 202,
         GGML_TYPE_Q5_0_R4   = 206,
@@ -521,6 +522,7 @@ extern "C" {
         GGML_FTYPE_MOSTLY_IQ2_KT  = 142, // except 1d tensors
         GGML_FTYPE_MOSTLY_IQ3_KT  = 143, // except 1d tensors
         GGML_FTYPE_MOSTLY_IQ4_KT  = 144, // except 1d tensors
+        GGML_FTYPE_MOSTLY_IQ3_KS  = 145, // except 1d tensors
                                          //
         GGML_FTYPE_MOSTLY_Q4_0_R8   = 202, // except 1d tensors
         GGML_FTYPE_MOSTLY_Q8_0_R8   = 207, // except 1d tensors

ggml/src/ggml-common.h

@@ -649,6 +649,14 @@ typedef struct {
 } block_iq3_k;
 static_assert(sizeof(block_iq3_k) == sizeof(ggml_half) + 2*sizeof(uint16_t) + QK_K/32 + QK_K/4 + QK_K/8, "wrong iq3_k block size/padding");
 
+typedef struct {
+    uint16_t extra;
+    uint8_t scales[QK_K/64];
+    uint8_t qs[QK_K/4];
+    uint8_t qh[QK_K/8];
+} block_iq3_ks;
+static_assert(sizeof(block_iq3_ks) == sizeof(uint16_t) + QK_K/64 + QK_K/4 + QK_K/8, "wrong iq3_ks block size/padding");
+
 typedef struct {
     ggml_half d[4];
     uint8_t extra[8];

ggml/src/ggml-cuda.cu

@@ -3465,6 +3465,7 @@ GGML_CALL static bool ggml_backend_cuda_supports_op(ggml_backend_t backend, cons
                     case GGML_TYPE_IQ3_XXS:
                     case GGML_TYPE_IQ4_NL:
                     case GGML_TYPE_IQ4_XS:
+                    case GGML_TYPE_IQ3_KS:
                     case GGML_TYPE_IQ4_KS:
                     case GGML_TYPE_IQ4_KSS:
                     case GGML_TYPE_IQ5_KS:

ggml/src/ggml-cuda/common.cuh

@@ -599,6 +599,13 @@ struct ggml_cuda_type_traits<GGML_TYPE_IQ3_K> {
     static constexpr int qi = QI4_XS;
 };
 
+template<>
+struct ggml_cuda_type_traits<GGML_TYPE_IQ3_KS> {
+    static constexpr int qk = QK_K;
+    static constexpr int qr = QR4_XS;
+    static constexpr int qi = QI4_XS;
+};
+
 template<>
 struct ggml_cuda_type_traits<GGML_TYPE_IQ4_K> {
     static constexpr int qk = QK_K;

ggml/src/ggml-cuda/convert.cu

@@ -1333,6 +1333,51 @@ static __global__ void dequantize_block_iq3_k(const void * __restrict__ vx, dst_
     }
 }
 
+template<typename dst_t>
+static __global__ void dequantize_block_iq3_ks(const void * __restrict__ vx, dst_t * __restrict__ yy, int64_t n_per_row, int64_t row_size) {
+
+    int64_t ii  = blockIdx.x;
+    int64_t row = (QK_K * ii) / n_per_row;
+    const char * cx = (const char *)vx + row * row_size;
+    float scale = *(const ggml_half *)cx;
+    const block_iq3_ks * x = (const block_iq3_ks *)(cx + sizeof(ggml_half));
+    const int64_t i   = ii - (row*n_per_row)/QK_K;
+
+    const int64_t tid = threadIdx.x;
+    const int64_t is = tid/16;
+    const int64_t il = tid%16;
+    dst_t * y = yy + ii*QK_K + 128*is + 2*il;
+    const uint8_t  * qs = x[i].qs + 32*is + 2*il;
+    const uint8_t  * qh = x[i].qh + 2*il;
+    uint16_t extra = x[i].extra >> 4*is;
+    const float d0 = scale * (int(((x[i].scales[0] >> 4*is) & 0xf) | ((extra << 4) & 0x10)) - 16);
+    const float d1 = scale * (int(((x[i].scales[1] >> 4*is) & 0xf) | ((extra << 3) & 0x10)) - 16);
+    const float d2 = scale * (int(((x[i].scales[2] >> 4*is) & 0xf) | ((extra << 2) & 0x10)) - 16);
+    const float d3 = scale * (int(((x[i].scales[3] >> 4*is) & 0xf) | ((extra << 1) & 0x10)) - 16);
+    extra >>= 8;
+    const int8_t * values0 = iq3nl_values + ((extra & 1) << 3);
+    const int8_t * values1 = iq3nl_values + ((extra & 2) << 2);
+    const int8_t * values2 = iq3nl_values + ((extra & 4) << 1);
+    const int8_t * values3 = iq3nl_values + ((extra & 8) << 0);
+    if constexpr (std::is_same_v<dst_t, nv_bfloat16>) {
+        for (int j = 0; j < 2; ++j) {
+            uint8_t h = qh[j] >> 4*is;
+            y[j+ 0] = __float2bfloat16(d0 * values0[((qs[j] >> 0) & 3) | ((h << 2) & 4)]);
+            y[j+32] = __float2bfloat16(d1 * values1[((qs[j] >> 2) & 3) | ((h << 1) & 4)]);
+            y[j+64] = __float2bfloat16(d2 * values2[((qs[j] >> 4) & 3) | ((h >> 0) & 4)]);
+            y[j+96] = __float2bfloat16(d3 * values3[((qs[j] >> 6) & 3) | ((h >> 1) & 4)]);
+        }
+    } else {
+        for (int j = 0; j < 2; ++j) {
+            uint8_t h = qh[j] >> 4*is;
+            y[j+ 0] = d0 * values0[((qs[j] >> 0) & 3) | ((h << 2) & 4)];
+            y[j+32] = d1 * values1[((qs[j] >> 2) & 3) | ((h << 1) & 4)];
+            y[j+64] = d2 * values2[((qs[j] >> 4) & 3) | ((h >> 0) & 4)];
+            y[j+96] = d3 * values3[((qs[j] >> 6) & 3) | ((h >> 1) & 4)];
+        }
+    }
+}
+
 template <int qk, int qr, dequantize_kernel_t dequantize_kernel, typename dst_t>
 static void dequantize_block_cuda(const void * __restrict__ vx, dst_t * __restrict__ y, const int64_t nrows, const int64_t n_per_row, cudaStream_t stream) {
     const int64_t k = nrows * n_per_row;
@@ -1573,6 +1618,14 @@ static void dequantize_row_iq3_k_cuda(const void * vx, dst_t * y, const int64_t
     dequantize_block_iq3_k<<<nb, 32, 0, stream>>>(vx, y);
 }
 
+template<typename dst_t>
+static void dequantize_row_iq3_ks_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 int64_t row_size = ggml_row_size(GGML_TYPE_IQ3_KS, n_per_row);
+    const int nb = (k + QK_K - 1) / QK_K;
+    dequantize_block_iq3_ks<<<nb, 32, 0, stream>>>(vx, y, n_per_row, row_size);
+}
+
 template<typename dst_t>
 static void dequantize_row_iq3_k_r4_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;
@@ -1719,6 +1772,8 @@ to_bf16_cuda_t ggml_get_to_bf16_cuda(ggml_type type) {
             return dequantize_row_iq2_k_cuda<nv_bfloat16>;
         case GGML_TYPE_IQ3_K:
             return dequantize_row_iq3_k_cuda<nv_bfloat16>;
+        case GGML_TYPE_IQ3_KS:
+            return dequantize_row_iq3_ks_cuda<nv_bfloat16>;
         case GGML_TYPE_IQ4_KSS:
             return dequantize_row_iq4_kss_cuda<nv_bfloat16>;
         case GGML_TYPE_IQ4_KS:
@@ -1821,6 +1876,8 @@ to_fp16_cuda_t ggml_get_to_fp16_cuda(ggml_type type) {
             return dequantize_row_iq2_k_cuda;
         case GGML_TYPE_IQ3_K:
             return dequantize_row_iq3_k_cuda;
+        case GGML_TYPE_IQ3_KS:
+            return dequantize_row_iq3_ks_cuda;
         case GGML_TYPE_IQ4_K:
             return dequantize_row_iq4_k_cuda;
         case GGML_TYPE_IQ5_K:
@@ -1916,6 +1973,8 @@ to_fp32_cuda_t ggml_get_to_fp32_cuda(ggml_type type) {
             return dequantize_row_iq2_k_cuda;
         case GGML_TYPE_IQ3_K:
             return dequantize_row_iq3_k_cuda;
+        case GGML_TYPE_IQ3_KS:
+            return dequantize_row_iq3_ks_cuda;
         case GGML_TYPE_IQ4_K:
             return dequantize_row_iq4_k_cuda;
         case GGML_TYPE_IQ5_K:

ggml/src/ggml-cuda/iqk_mmvq.cu

@@ -1104,6 +1104,73 @@ __device__ __forceinline__ void vec_dot_iq3_k_q8_1(
 
 }
 
+__device__ __forceinline__ void vec_dot_iq3_ks_q8_1(
+    const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & kbx, const int & iiqs, float * result) {
+
+    float d = __half2float(*(const half *)vbq);
+    const block_iq3_ks * bq3 = (const block_iq3_ks *)((const char *)vbq + sizeof(half)) + kbx;
+
+    int iqs = iiqs/4;
+    const int ib128 = iqs/4;  // 0 or 1. 0 works on quants 0...127, 1 on quants 128...255
+                              // Each thread processes 8 quants in each of the 4 32-blocks
+    const int il8   = iqs%4;  // 0...3. 0 works on quants 0...7, 1 on quants 8...15, 2 on 16...23, 3 on 24...31
+    const int shift = 4*(il8/2);
+
+    const uint16_t * ql = (const uint16_t *)bq3->qs + 16*ib128 + 4*il8;
+    const uint16_t * qh = (const uint16_t *)bq3->qh + 4*il8;
+
+    int32_t aux32;
+    const uint8_t * aux8 = (const uint8_t *)&aux32;
+
+    uint16_t extra = bq3->extra >> 4*ib128;
+    uint16_t extra_v = extra >> 8;
+
+    const uint16_t * values1 = iq3k_table + ((extra_v << 6) & 0x40);
+    const uint16_t * values2 = iq3k_table + ((extra_v << 5) & 0x40);
+    const uint16_t * values3 = iq3k_table + ((extra_v << 4) & 0x40);
+    const uint16_t * values4 = iq3k_table + ((extra_v << 3) & 0x40);
+
+    const int * q8;
+    int sumi[4] = {0, 0, 0, 0};
+    int v;
+    for (int i = 0; i < 2; ++i) {
+        uint32_t vl = ql[2*i+0] | (ql[2*i+1] << 16);
+        uint32_t vh = ((qh[2*i+0] | (qh[2*i+1] << 16)) >> 4*ib128) << 2;
+
+        q8 = (const int *)bq8_1[4*ib128+0].qs + 2*il8;
+        aux32 = (vl & 0x03030303) | (vh & 0x04040404);
+        v = int_from_table_2(aux8, values1);
+        sumi[0] = ggml_cuda_dp4a(v, q8[i], sumi[0]);
+        vl >>= 2; vh >>= 1;
+
+        q8 += sizeof(block_q8_1)/4;
+        aux32 = (vl & 0x03030303) | (vh & 0x04040404);
+        v = int_from_table_2(aux8, values2);
+        sumi[1] = ggml_cuda_dp4a(v, q8[i], sumi[1]);
+        vl >>= 2; vh >>= 1;
+
+        q8 += sizeof(block_q8_1)/4;
+        aux32 = (vl & 0x03030303) | (vh & 0x04040404);
+        v = int_from_table_2(aux8, values3);
+        sumi[2] = ggml_cuda_dp4a(v, q8[i], sumi[2]);
+        vl >>= 2; vh >>= 1;
+
+        q8 += sizeof(block_q8_1)/4;
+        aux32 = (vl & 0x03030303) | (vh & 0x04040404);
+        v = int_from_table_2(aux8, values4);
+        sumi[3] = ggml_cuda_dp4a(v, q8[i], sumi[3]);
+
+    }
+    const uint16_t * sl16 = (const uint16_t *)bq3->scales;
+    aux32 = __vsub4(((sl16[0] | (sl16[1] << 16)) >> 4*ib128) & 0x0f0f0f0f, 0x10101010);
+    const int8_t * a8 = (const int8_t *)&aux32;
+    *result += d * (__low2float(bq8_1[4*ib128+0].ds) * (a8[0] + ((extra << 4) & 0x10)) * sumi[0] +
+                    __low2float(bq8_1[4*ib128+1].ds) * (a8[1] + ((extra << 3) & 0x10)) * sumi[1] +
+                    __low2float(bq8_1[4*ib128+2].ds) * (a8[2] + ((extra << 2) & 0x10)) * sumi[2] +
+                    __low2float(bq8_1[4*ib128+3].ds) * (a8[3] + ((extra << 1) & 0x10)) * sumi[3]);
+
+}
+
 __device__ __forceinline__ void vec_dot_iq1_bn_q8_1(
     const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & kbx, const int & iqs, float * result) {
 
@@ -1302,6 +1369,14 @@ void mul_mat_vec_iq4_ks_q8_1_cuda(
     iqk_mul_mat_vec_q_cuda<GGML_TYPE_IQ4_KS, VDR_IQ4_KS_Q8_1_MMVQ, vec_dot_iq4_ks_q8_1>(vx, vy, dst, ids_data, ncols_x, nrows_x, nrows_y, ncols_y, nrows_dst, ne2, nb02, nb12, nb2, ids_nb0, stream);
 }
 
+void mul_mat_vec_iq3_ks_q8_1_cuda(
+    const void * vx, const void * vy, float * dst, const char * ids_data,
+    const int ncols_x, const int nrows_x, const int nrows_y, const int ncols_y, const int nrows_dst,
+    const int ne2, const uint64_t nb02, const uint64_t nb12, const uint64_t nb2, int64_t ids_nb0, cudaStream_t stream) {
+
+    iqk_mul_mat_vec_q_cuda<GGML_TYPE_IQ3_KS, VDR_IQ3_K_Q8_1_MMVQ, vec_dot_iq3_ks_q8_1>(vx, vy, dst, ids_data, ncols_x, nrows_x, nrows_y, ncols_y, nrows_dst, ne2, nb02, nb12, nb2, ids_nb0, stream);
+}
+
 void mul_mat_vec_iq4_kt_q8_1_cuda(
     const void * vx, const void * vy, float * dst, const char * ids_data,
     const int ncols_x, const int nrows_x, const int nrows_y, const int ncols_y, const int nrows_dst,

ggml/src/ggml-cuda/iqk_mmvq.cuh

@@ -16,6 +16,11 @@ void mul_mat_vec_iq3_k_q8_1_cuda(
     const int ncols_x, const int nrows_x, const int nrows_y, const int ncols_y, const int nrows_dst,
     const int ne2, const uint64_t nb02, const uint64_t nb12, const uint64_t nb2, const int64_t ids_nb0, cudaStream_t stream);
 
+void mul_mat_vec_iq3_ks_q8_1_cuda(
+    const void * vx, const void * vy, float * dst, const char * ids_data,
+    const int ncols_x, const int nrows_x, const int nrows_y, const int ncols_y, const int nrows_dst,
+    const int ne2, const uint64_t nb02, const uint64_t nb12, const uint64_t nb2, const int64_t ids_nb0, cudaStream_t stream);
+
 void mul_mat_vec_iq4_k_q8_1_cuda(
     const void * vx, const void * vy, float * dst, const char * ids_data,
     const int ncols_x, const int nrows_x, const int nrows_y, const int ncols_y, const int nrows_dst,

ggml/src/ggml-cuda/mmq.cu

@@ -94,6 +94,9 @@ void ggml_cuda_op_mul_mat_q(
         case GGML_TYPE_IQ4_NL:
             mul_mat_q_case<GGML_TYPE_IQ4_NL>(ctx, args, stream);
             break;
+        case GGML_TYPE_IQ3_KS:
+            mul_mat_q_case<GGML_TYPE_IQ3_KS>(ctx, args, stream);
+            break;
         case GGML_TYPE_IQ4_KS:
             mul_mat_q_case<GGML_TYPE_IQ4_KS>(ctx, args, stream);
             break;
@@ -196,6 +199,7 @@ bool ggml_cuda_should_use_mmq(enum ggml_type type, int cc, int64_t ne11) {
         case GGML_TYPE_IQ1_S_R4:
         case GGML_TYPE_IQ4_XS:
         case GGML_TYPE_IQ4_NL:
+        case GGML_TYPE_IQ3_KS:
         case GGML_TYPE_IQ4_KS:
         case GGML_TYPE_IQ4_KS_R4:
         case GGML_TYPE_IQ5_KS: