Q2_K_R4 (#136)

* q2_k_r4: Zen4

PP-512(LLaMA-3.1-8B) = 256 t/s

* q3_k_r4: AVX2

* q2_k_r4: AVX2

We get PP-512(LLaMA-3.1-8B) = 287 t/s.

Also cherry-picked the q3_k_r4 AVX2 adaptation that I somehow
forgot to push upstream.

* q2_k_r4: NEON

We get PP-512(LLaMA-3.1-8B) = 106.2 t/s.
TG-128 is 36.02 t/s, which is ~10% higher than q2_K_S.

* Make sure rows per thread are a multiple of 4

---------

Co-authored-by: Iwan Kawrakow <iwan.kawrakow@gmail.com>

1 files changed, 23 insertions, 1 deletions

diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c
index 0bb59d2b..6c574933 100644
--- a/ggml/src/ggml.c
+++ b/ggml/src/ggml.c
@@ -862,6 +862,19 @@ static const ggml_type_traits_t type_traits[GGML_TYPE_COUNT] = {
         .nrows                    = 1,
         .row_meta_size            = 0,
     },
+    [GGML_TYPE_Q2_K_R4] = {
+        .type_name                = "q2_k_r4",
+        .blck_size                = QK_K,
+        .type_size                = sizeof(block_q2_K),
+        .is_quantized             = true,
+        .to_float                 = (ggml_to_float_t) dequantize_row_q2_k_r4,
+        .from_float               = quantize_row_q2_k_r4,
+        .from_float_ref           = (ggml_from_float_t) quantize_row_q2_k_r4_ref,
+        .vec_dot                  = vec_dot_q2_k_r4_q8_k,
+        .vec_dot_type             = GGML_TYPE_Q8_K,
+        .nrows                    = 1,
+        .row_meta_size            = 0,
+    },
     [GGML_TYPE_Q3_K] = {
         .type_name                = "q3_K",
         .blck_size                = QK_K,
@@ -4070,7 +4083,8 @@ enum ggml_type ggml_ftype_to_ggml_type(enum ggml_ftype ftype) {
         case GGML_FTYPE_MOSTLY_Q6_0:          wtype = GGML_TYPE_Q6_0;  break;
         case GGML_FTYPE_MOSTLY_Q8_0:          wtype = GGML_TYPE_Q8_0;  break;
         case GGML_FTYPE_MOSTLY_Q2_K:          wtype = GGML_TYPE_Q2_K;  break;
-        case GGML_FTYPE_MOSTLY_Q3_K:          wtype = GGML_TYPE_Q3_K;  break;
+        case GGML_FTYPE_MOSTLY_Q2_K_R4:       wtype = GGML_TYPE_Q2_K_R4;  break;
+        case GGML_FTYPE_MOSTLY_Q3_K:          wtype = GGML_TYPE_Q3_K;     break;
         case GGML_FTYPE_MOSTLY_Q3_K_R4:       wtype = GGML_TYPE_Q3_K_R4;  break;
         case GGML_FTYPE_MOSTLY_Q4_K:          wtype = GGML_TYPE_Q4_K;     break;
         case GGML_FTYPE_MOSTLY_Q4_K_R4:       wtype = GGML_TYPE_Q4_K_R4;  break;
@@ -10604,6 +10618,7 @@ static void ggml_compute_forward_add(
         case GGML_TYPE_Q6_0:
         case GGML_TYPE_Q8_0:
         case GGML_TYPE_Q2_K:
+        case GGML_TYPE_Q2_K_R4:
         case GGML_TYPE_Q3_K:
         case GGML_TYPE_Q3_K_R4:
         case GGML_TYPE_Q4_K:
@@ -11057,6 +11072,7 @@ static void ggml_compute_forward_add1(
         case GGML_TYPE_Q8_0:
         case GGML_TYPE_Q8_1:
         case GGML_TYPE_Q2_K:
+        case GGML_TYPE_Q2_K_R4:
         case GGML_TYPE_Q3_K:
         case GGML_TYPE_Q3_K_R4:
         case GGML_TYPE_Q4_K:
@@ -11207,6 +11223,7 @@ static void ggml_compute_forward_acc(
         case GGML_TYPE_Q8_0:
         case GGML_TYPE_Q8_1:
         case GGML_TYPE_Q2_K:
+        case GGML_TYPE_Q2_K_R4:
         case GGML_TYPE_Q3_K:
         case GGML_TYPE_Q3_K_R4:
         case GGML_TYPE_Q4_K:
@@ -14403,6 +14420,7 @@ static void ggml_compute_forward_out_prod(
         case GGML_TYPE_Q6_0:
         case GGML_TYPE_Q8_0:
         case GGML_TYPE_Q2_K:
+        case GGML_TYPE_Q2_K_R4:
         case GGML_TYPE_Q3_K:
         case GGML_TYPE_Q3_K_R4:
         case GGML_TYPE_Q4_K:
@@ -14793,6 +14811,7 @@ static void ggml_compute_forward_set(
         case GGML_TYPE_Q8_0:
         case GGML_TYPE_Q8_1:
         case GGML_TYPE_Q2_K:
+        case GGML_TYPE_Q2_K_R4:
         case GGML_TYPE_Q3_K:
         case GGML_TYPE_Q3_K_R4:
         case GGML_TYPE_Q4_K:
@@ -15077,6 +15096,7 @@ static void ggml_compute_forward_get_rows(
         case GGML_TYPE_Q8_0:
         case GGML_TYPE_Q8_1:
         case GGML_TYPE_Q2_K:
+        case GGML_TYPE_Q2_K_R4:
         case GGML_TYPE_Q3_K:
         case GGML_TYPE_Q3_K_R4:
         case GGML_TYPE_Q4_K:
@@ -15688,6 +15708,7 @@ static void ggml_compute_forward_clamp(
         case GGML_TYPE_Q8_0:
         case GGML_TYPE_Q8_1:
         case GGML_TYPE_Q2_K:
+        case GGML_TYPE_Q2_K_R4:
         case GGML_TYPE_Q3_K:
         case GGML_TYPE_Q3_K_R4:
         case GGML_TYPE_Q4_K:
@@ -22527,6 +22548,7 @@ size_t ggml_quantize_chunk(
         case GGML_TYPE_Q6_0:    result = quantize_q6_0(src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break;
         case GGML_TYPE_Q8_0:    result = quantize_q8_0(src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break;
         case GGML_TYPE_Q2_K:    result = quantize_q2_K(src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break;
+        case GGML_TYPE_Q2_K_R4: result = quantize_q2_k_r4(src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break;
         case GGML_TYPE_Q3_K:    result = quantize_q3_K(src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break;
         case GGML_TYPE_Q3_K_R4: result = quantize_q3_k_r4(src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break;
         case GGML_TYPE_Q4_K:    result = quantize_q4_K(src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break;