CPU Optimizations - Prefill, Tokenization, and Token Generation

ggml/src/ggml-cpu/ggml-cpu.c

@@ -3053,6 +3053,21 @@ static int ggml_cpu_try_fuse_ops(
             }
         }
     }
+    if (node->op == GGML_OP_GET_ROWS) {
+        // GET_ROWS + ADD fusion
+        const enum ggml_op fuse_ops[] = { GGML_OP_GET_ROWS, GGML_OP_ADD };
+        if (ggml_can_fuse(cgraph, node_n, fuse_ops, 2)) {
+            struct ggml_tensor * add_node = cgraph->nodes[node_n + 1];
+            struct ggml_tensor * pos_embd = (add_node->src[0] == node) ? add_node->src[1] : add_node->src[0];
+
+            if (node->src[1]->type == GGML_TYPE_I32 &&
+                pos_embd->type == GGML_TYPE_F32 &&
+                add_node->type == GGML_TYPE_F32) {
+
+                return ggml_compute_forward_get_rows_add_fused(params, node, add_node);
+            }
+        }
+    }
 
     return 0;
 }

ggml/src/ggml-cpu/ops.cpp

@@ -4877,17 +4877,178 @@ static void ggml_compute_forward_get_rows_q(
     const int ir0 = dr*ith;
     const int ir1 = MIN(ir0 + dr, nr);
 
-    for (int64_t i = ir0; i < ir1; ++i) {
-        const int64_t i12 = i/(ne11*ne10);
-        const int64_t i11 = (i - i12*ne11*ne10)/ne10;
-        const int64_t i10 = (i - i12*ne11*ne10 - i11*ne10);
-        const int64_t i01 = *(int32_t *) ((char *) src1->data + i10*nb10 + i11*nb11 + i12*nb12);
+    const bool is_1d = (ne12 == 1 && ne11 == 1);
+    if (is_1d) {
+        const char * src1_ptr = (const char *) src1->data + ir0 * nb10;
+        char * dst_ptr = (char *) dst->data + ir0 * nb1;
+
+        const int64_t src1_step = nb10;
+        const int64_t dst_step = nb1;
+
+        int64_t last_i01 = -1;
+        const float * last_dst_ptr = nullptr;
+
+        for (int64_t i = ir0; i < ir1; ++i) {
+            const int64_t i01 = *(const int32_t *) src1_ptr;
+            GGML_ASSERT(i01 >= 0 && i01 < ne01);
+
+            // Cache hit (with prev)
+            if (i01 == last_i01 && last_dst_ptr != nullptr) {
+                memcpy(dst_ptr, last_dst_ptr, nc * sizeof(float));
+            } else {
+                // Prefetch next row
+                if (i + 1 < ir1) {
+                    const int64_t next_i01 = *(const int32_t *)(src1_ptr + src1_step);
+                    const char * next_src0_row = (const char *) src0->data + next_i01 * nb01;
+                    __builtin_prefetch(next_src0_row, 0, 1);
+                    __builtin_prefetch(next_src0_row + 64, 0, 1);
+                    __builtin_prefetch(next_src0_row + 128, 0, 1);
+                }
+
+                dequantize_row_q(
+                    (const void *) ((const char *) src0->data + i01 * nb01),
+                    (float *) dst_ptr,
+                    nc
+                );
+
+                last_i01 = i01;
+                last_dst_ptr = (const float *) dst_ptr;
+            }
+
+            src1_ptr += src1_step;
+            dst_ptr += dst_step;
+        }
+    }
+    else {
+        int64_t i10 = ir0 % ne10;
+        int64_t i11 = (ir0 / ne10) % ne11;
+        int64_t i12 = ir0 / (ne10 * ne11);
+
+        const char * ptr_src1 = (const char *) src1->data + i10*nb10 + i11*nb11 + i12*nb12;
+        char * ptr_dst        = (char *) dst->data        + i10*nb1  + i11*nb2  + i12*nb3;
+        const char * base_src0 = (const char *) src0->data + i11*nb02 + i12*nb03;
+
+        for (int64_t i = ir0; i < ir1; ++i) {
+            const int64_t i01 = *(const int32_t *) ptr_src1;
+            GGML_ASSERT(i01 >= 0 && i01 < ne01);
+
+            dequantize_row_q(
+                (const void *) (base_src0 + i01*nb01),
+                (float *) ptr_dst,
+                nc
+            );
+
+            ptr_src1 += nb10;
+            ptr_dst  += nb1;
+            i10++;
+
+            if (__builtin_expect(i10 == ne10, 0)) {
+                i10 = 0;
+
+                ptr_src1 += nb11 - (ptrdiff_t)ne10 * (ptrdiff_t)nb10;
+                ptr_dst  += nb2  - (ptrdiff_t)ne10 * (ptrdiff_t)nb1;
+
+                i11++;
+                if (__builtin_expect(i11 == ne11, 0)) {
+                    i11 = 0;
+                    ptr_src1 += nb12 - (ptrdiff_t)ne11 * nb11;
+                    ptr_dst  += nb3  - (ptrdiff_t)ne11 * nb2;
+
+                    i12++;
+                    base_src0 = (const char *) src0->data + i12*nb03;
+                } else {
+                    base_src0 += nb02;
+                }
+            }
+        }
+    }
+}
+
+static void ggml_compute_forward_get_rows_add_fused_q(
+        const ggml_compute_params * params,
+        const ggml_tensor * dst_get_rows,
+        ggml_tensor * dst_add) {
+
+    const ggml_tensor * src0 = dst_get_rows->src[0];
+    const ggml_tensor * src1 = dst_get_rows->src[1];
+    const ggml_tensor * add_src = (dst_add->src[0] == dst_get_rows)
+                                   ? dst_add->src[1]
+                                   : dst_add->src[0];
+    auto dst = dst_add;
+    GGML_TENSOR_BINARY_OP_LOCALS
+
+    const int64_t nc = ne00;
+    const int64_t nr = ggml_nelements(src1);
+
+    const ggml_type type = src0->type;
+    ggml_to_float_t const dequantize_row_q = ggml_get_type_traits(type)->to_float;
+
+    const int ith = params->ith;
+    const int nth = params->nth;
+
+    const int dr = (nr + nth - 1)/nth;
+    const int ir0 = dr*ith;
+    const int ir1 = MIN(ir0 + dr, nr);
+
+    GGML_ASSERT(ne0 == nc);
+    GGML_ASSERT(ne1 == nc);
+    GGML_ASSERT(nb00 == ggml_type_size(type));
+    GGML_ASSERT(add_src->type == GGML_TYPE_F32);
+
+    const int64_t add_nelem = ggml_nelements(add_src);
+    const bool is_broadcast_single = (add_nelem == nc);
+    const bool is_positional = (add_nelem >= nr * nc);
 
+    const char * src1_ptr = (const char *) src1->data + ir0 * nb10;
+    char * dst_ptr = (char *) dst_add->data + ir0 * nb1;
+    const int64_t src1_step = nb10;
+    const int64_t dst_step = nb1;
+
+    int64_t last_i01 = -1;
+    int64_t last_add_i = -1;
+    const float * last_dst_ptr = nullptr;
+
+    for (int64_t i = ir0; i < ir1; ++i) {
+        const int64_t i01 = *(const int32_t *) src1_ptr;
         GGML_ASSERT(i01 >= 0 && i01 < ne01);
 
-        dequantize_row_q(
-                (const void *) ((char *) src0->data + i01*nb01 + i11*nb02 + i12*nb03),
-                     (float *) ((char *)  dst->data + i10*nb1  + i11*nb2  + i12*nb3), nc);
+        int64_t add_i{0};
+        if (!is_broadcast_single) {
+            if (is_positional) {
+                add_i = i;
+            } else {
+                add_i = i % (add_nelem / nc);
+            }
+        }
+
+        const float * add_row = (const float *)((const char *)add_src->data + add_i * nb11);
+
+        if (i01 == last_i01 && add_i == last_add_i && last_dst_ptr != nullptr) {
+            memcpy(dst_ptr, last_dst_ptr, nc * sizeof(float));
+        } else {
+            if (i + 1 < ir1) {
+                const int64_t next_i01 = *(const int32_t *)(src1_ptr + src1_step);
+                const char * next_src0_row = (const char *) src0->data + next_i01 * nb01;
+                __builtin_prefetch(next_src0_row, 0, 1);
+                __builtin_prefetch(next_src0_row + 64, 0, 1);
+                __builtin_prefetch(next_src0_row + 128, 0, 1);
+            }
+
+            dequantize_row_q(
+                (const void *) ((const char *) src0->data + i01 * nb01),
+                (float *) dst_ptr,
+                nc
+            );
+
+            ggml_vec_acc_f32(nc, (float *) dst_ptr, add_row);
+
+            last_i01 = i01;
+            last_add_i = add_i;
+            last_dst_ptr = (const float *) dst_ptr;
+        }
+
+        src1_ptr += src1_step;
+        dst_ptr += dst_step;
     }
 }
 
@@ -5088,6 +5249,47 @@ void ggml_compute_forward_get_rows(
     //}
 }
 
+int ggml_compute_forward_get_rows_add_fused(const struct ggml_compute_params * params, struct ggml_tensor * dst_get_rows, struct ggml_tensor * dst_add) {
+    const ggml_tensor * src0 = dst_get_rows->src[0];
+
+    switch (src0->type) {
+        case GGML_TYPE_Q1_0:
+        case GGML_TYPE_Q2_0:
+        case GGML_TYPE_Q4_0:
+        case GGML_TYPE_Q4_1:
+        case GGML_TYPE_Q5_0:
+        case GGML_TYPE_Q5_1:
+        case GGML_TYPE_Q8_0:
+        case GGML_TYPE_Q8_1:
+        case GGML_TYPE_MXFP4:
+        case GGML_TYPE_NVFP4:
+        case GGML_TYPE_Q2_K:
+        case GGML_TYPE_Q3_K:
+        case GGML_TYPE_Q4_K:
+        case GGML_TYPE_Q5_K:
+        case GGML_TYPE_Q6_K:
+        case GGML_TYPE_TQ1_0:
+        case GGML_TYPE_TQ2_0:
+        case GGML_TYPE_IQ2_XXS:
+        case GGML_TYPE_IQ2_XS:
+        case GGML_TYPE_IQ3_XXS:
+        case GGML_TYPE_IQ1_S:
+        case GGML_TYPE_IQ1_M:
+        case GGML_TYPE_IQ4_NL:
+        case GGML_TYPE_IQ4_XS:
+        case GGML_TYPE_IQ3_S:
+        case GGML_TYPE_IQ2_S:
+        {
+            ggml_compute_forward_get_rows_add_fused_q(params, dst_get_rows, dst_add);
+        } break;
+        default:
+        {
+            return 0;
+        }
+    }
+    return 1;
+}
+
 template<typename src_t, typename idx_t>
 static void ggml_compute_forward_set_rows_impl(
         const ggml_compute_params * params,

ggml/src/ggml-cpu/ops.h

@@ -54,6 +54,7 @@ void ggml_compute_forward_set(const struct ggml_compute_params * params, struct
 void ggml_compute_forward_cpy(const struct ggml_compute_params * params, struct ggml_tensor * dst);
 void ggml_compute_forward_cont(const struct ggml_compute_params * params, struct ggml_tensor * dst);
 void ggml_compute_forward_get_rows(const struct ggml_compute_params * params, struct ggml_tensor * dst);
+int ggml_compute_forward_get_rows_add_fused(const struct ggml_compute_params * params, struct ggml_tensor * dst_get_rows, struct ggml_tensor * dst_add);
 void ggml_compute_forward_get_rows_back(const struct ggml_compute_params * params, struct ggml_tensor * dst);
 void ggml_compute_forward_set_rows(const struct ggml_compute_params * params, struct ggml_tensor * dst);
 void ggml_compute_forward_diag(const struct ggml_compute_params * params, struct ggml_tensor * dst);

ggml/src/ggml-impl.h

@@ -11,6 +11,9 @@
 #include <stdbool.h>
 #include <stdint.h>
 #include <string.h>
+#if defined(__x86_64__) || defined(_M_X64) || defined(__i386__)
+    #include <immintrin.h>
+#endif
 
 #ifdef __ARM_FEATURE_SVE
 #include <arm_sve.h>
@@ -382,6 +385,12 @@ static inline uint32_t fp32_to_bits(float f) {
 }
 
 static inline float ggml_compute_fp16_to_fp32(ggml_fp16_t h) {
+#ifdef __F16C__
+    return _cvtsh_ss(h);
+#elif defined(__aarch64__) && defined(__ARM_FP) && (__ARM_FP & 2)
+    union { uint16_t u; __fp16 f; } u = { .u = h };
+    return (float)u.f;
+#else
     const uint32_t w = (uint32_t) h << 16;
     const uint32_t sign = w & UINT32_C(0x80000000);
     const uint32_t two_w = w + w;
@@ -402,6 +411,7 @@ static inline float ggml_compute_fp16_to_fp32(ggml_fp16_t h) {
     const uint32_t result = sign |
         (two_w < denormalized_cutoff ? fp32_to_bits(denormalized_value) : fp32_to_bits(normalized_value));
     return fp32_from_bits(result);
+#endif
 }
 
 static inline ggml_fp16_t ggml_compute_fp32_to_fp16(float f) {