fix metal implementation

Author: cebtenzzre

ggml-metal.metal

@@ -605,7 +605,7 @@ static float rope_ntkv2(
         const float theta_base,
         const float theta_linear,
         const float theta_ntk,
-        device const float corr_factors[4],
+        const float corr_factors[4],
         const int64_t i0,
         const float ntk_factor,
         const float ext_factor) {
@@ -620,6 +620,29 @@ static float rope_ntkv2(
     return theta;
 }
 
+// Interpolation constants found experimentally for LLaMA (might not be totally optimal though)
+// Do not change unless there is a good reason for doing so!
+constant float BETA_0  = 1.75f;
+constant float BETA_1  = 1.25f;
+constant float GAMMA_0 = 16.0f;
+constant float GAMMA_1 = 2.0f;
+
+constant float max_pos_emb = 2048;
+
+// Apparently solving `n_rot = 2pi * x * base^((2 * max_pos_emb) / n_dims)` for x, we get
+// `corr_fac(n_rot) = n_dims * log(max_pos_emb / (n_rot * 2pi)) / (2 * log(base))`
+static float rope_ntkv2_corr_factor(const int n_dims, const float n_rot, const float base) {
+    return n_dims * log(max_pos_emb / (n_rot * 2 * M_PI_F)) / (2 * log(base));
+}
+
+static void rope_ntkv2_corr_factors(int n_dims, const float freq_base, float factors[4]) {
+    // start and end correction factors
+    factors[0] = max(0.0f,         floor(rope_ntkv2_corr_factor(n_dims, BETA_0,  freq_base)));
+    factors[1] = min(n_dims - 1.0f, ceil(rope_ntkv2_corr_factor(n_dims, BETA_1,  freq_base)));
+    factors[2] = max(0.0f,         floor(rope_ntkv2_corr_factor(n_dims, GAMMA_0, freq_base)));
+    factors[3] = min(n_dims - 1.0f, ceil(rope_ntkv2_corr_factor(n_dims, GAMMA_1, freq_base)));
+}
+
 kernel void kernel_rope(
         device const  void * src0,
         device       float * dst,
@@ -651,10 +674,10 @@ kernel void kernel_rope(
     const int64_t i2 = tpig[1];
     const int64_t i1 = tpig[0];
 
-    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 float theta_scale = pow(freq_base, -2.0f/n_dims);
+    const float theta_ntk_scale = pow(freq_base * pow(freq_scale, (n_dims / (n_dims - 2.0f))), -2.0f/n_dims);
+    float corr_factors[4];
+    rope_ntkv2_corr_factors(n_dims, freq_base, corr_factors);
 
     float theta_base = (mode & 1) == 0 ? n_past + i2 : i2;
     float theta_ntk = theta_base;
@@ -666,8 +689,8 @@ kernel void kernel_rope(
             const float theta_linear = freq_scale * theta_base;
             const float theta = rope_ntkv2(theta_base, theta_linear, theta_ntk, corr_factors,
                     i0, ntk_factor, ext_factor);
-            const float cos_theta = cosf(theta);
-            const float sin_theta = sinf(theta);
+            const float cos_theta = cos(theta);
+            const float sin_theta = sin(theta);
 
             theta_base *= theta_scale;
             theta_ntk *= theta_ntk_scale;