initial Metal implementation

ggerganov · Jul 21, 2023 · fe2413c · fe2413c
1 parent be50066
commit fe2413c
Show file tree

Hide file tree

Showing 2 changed files with 50 additions and 11 deletions.
diff --git a/ggml-metal.m b/ggml-metal.m
@@ -887,10 +887,11 @@ void ggml_metal_graph_compute(
 
                             const int n_past = ((int32_t *)(src1->data))[0];
 
-                            float freq_base;
-                            float freq_scale;
-                            memcpy(&freq_base,  (int32_t *) src1->data + 4, sizeof(float));
-                            memcpy(&freq_scale, (int32_t *) src1->data + 5, sizeof(float));
+                            float freq_base, freq_scale, ntk_factor, extrapolation_factor;
+                            memcpy(&freq_base,            (int32_t *) src1->data + 4, sizeof(float));
+                            memcpy(&freq_scale,           (int32_t *) src1->data + 5, sizeof(float));
+                            memcpy(&ntk_factor,           (int32_t *) src1->data + 6, sizeof(float));
+                            memcpy(&extrapolation_factor, (int32_t *) src1->data + 7, sizeof(float));
 
                             [encoder setComputePipelineState:ctx->pipeline_rope];
                             [encoder setBuffer:id_src0 offset:offs_src0 atIndex:0];
@@ -916,6 +917,8 @@ void ggml_metal_graph_compute(
                             [encoder setBytes:&mode    length:sizeof(     int) atIndex:20];
                             [encoder setBytes:&freq_base  length:sizeof(float) atIndex:21];
                             [encoder setBytes:&freq_scale length:sizeof(float) atIndex:22];
+                            [encoder setBytes:&ntk_factor length:sizeof(float) atIndex:23];
+                            [encoder setBytes:&extrapolation_factor length:sizeof(float) atIndex:24];
 
                             [encoder dispatchThreadgroups:MTLSizeMake(ne01, ne02, ne03) threadsPerThreadgroup:MTLSizeMake(1, 1, 1)];
                         } break;

diff --git a/ggml-metal.metal b/ggml-metal.metal
@@ -580,6 +580,32 @@ kernel void kernel_alibi_f32(
     }
 }
 
+static float rope_ntkv2_ramp(const float low, const float high, const int i0) {
+    const float y = (i0 / 2 - low) / min(0.001f, high - low);
+    return 1.0f - min(1.0f, max(0.0f, y));
+}
+
+// NTKv2 algorithm based on LlamaPartNTKScaledRotaryEmbedding.py from https://github.com/jquesnelle/scaled-rope
+// MIT licensed. Copyright (c) 2023 Jeffrey Quesnelle and Bowen Peng.
+static float rope_ntkv2(
+        const float theta_base,
+        const float theta_linear,
+        const float theta_ntk,
+        device const float corr_factors[4],
+        const int64_t i0,
+        const float ntk_factor,
+        const float extrapolation_factor) {
+    float ramp_mix;
+    float theta;
+
+    ramp_mix = rope_ntkv2_ramp(corr_factors[0], corr_factors[1], i0) * ntk_factor;
+    theta = theta_linear * (1 - ramp_mix) + theta_ntk  * ramp_mix;
+
+    ramp_mix = rope_ntkv2_ramp(corr_factors[2], corr_factors[3], i0) * extrapolation_factor;
+    theta = theta        * (1 - ramp_mix) + theta_base * ramp_mix;
+    return theta;
+}
+
 kernel void kernel_rope(
         device const  void * src0,
         device       float * dst,
@@ -604,24 +630,33 @@ kernel void kernel_rope(
         constant       int & mode,
         constant     float & freq_base,
         constant     float & freq_scale,
+        constant     float & ntk_factor,
+        constant     float & extrapolation_factor,
         uint3 tpig[[thread_position_in_grid]]) {
     const int64_t i3 = tpig[2];
     const int64_t i2 = tpig[1];
     const int64_t i1 = tpig[0];
 
-    const bool is_neox = mode & 2;
-    const float theta_scale = pow(freq_base, -2.0f/n_dims);
+    const float theta_scale = powf(freq_base, -2.0f/n_dims);
+    const float theta_ntk_scale = powf(freq_base * powf(freq_scale, (n_dims / (n_dims - 2.0f))), -2.0f/n_dims);
+    device float corr_factors[4];
+    ggml_rope_ntkv2_corr_factors(n_dims, freq_base, corr_factors);
 
-    const int64_t p = ((mode & 1) == 0 ? n_past + i2 : i2);
+    float theta_base = (mode & 1) == 0 ? n_past + i2 : i2;
+    float theta_ntk = theta_base;
 
-    float theta = freq_scale * (float)p;
+    const bool is_neox = mode & 2;
 
     if (!is_neox) {
         for (int64_t i0 = 0; i0 < ne0; i0 += 2) {
-            const float cos_theta = cos(theta);
-            const float sin_theta = sin(theta);
+            const float theta_linear = freq_scale * theta_base;
+            const float theta = rope_ntkv2(theta_base, theta_linear, theta_ntk, corr_factors,
+                    i0, ntk_factor, extrapolation_factor);
+            const float cos_theta = cosf(theta);
+            const float sin_theta = sinf(theta);
 
-            theta *= theta_scale;
+            theta_base *= theta_scale;
+            theta_ntk *= theta_ntk_scale;
 
             device const float * const src = (device float *)((device char *) src0 + i3*nb03 + i2*nb02 + i1*nb01 + i0*nb00);
             device       float * dst_data  = (device float *)((device char *)  dst + i3*nb3  + i2*nb2  + i1*nb1  + i0*nb0);
@@ -633,6 +668,7 @@ kernel void kernel_rope(
             dst_data[1] = x0*sin_theta + x1*cos_theta;
         }
     } else {
+        theta_base *= freq_scale;
         // TODO: implement
     }
 }