Llama 3.1 RoPE

runq.c

@@ -650,11 +650,35 @@ float* forward(Transformer* transformer, int token, int pos) {
         matmul(s->k, &s->xq, w->wk + l, dim, kv_dim);
         matmul(s->v, &s->xq, w->wv + l, dim, kv_dim);
 
+    // Llama 3.1 RoPE Scaling
+        const float scale_factor = 8;
+        const float low_freq_factor = 1;
+        const float high_freq_factor = 4;
+        const float old_context_len = 8192;
+
+        const float low_freq_wavelen  = old_context_len / low_freq_factor;
+        const float high_freq_wavelen = old_context_len / high_freq_factor;
+    // END Llama 3.1 RoPE Scaling
+
         // RoPE relative positional encoding: complex-valued rotate q and k in each head
         for (int i = 0; i < dim; i+=2) {
             int head_dim = i % head_size;
 // L2E Addition
             float freq = 1.0f / powf(rope_tf, head_dim / (float)head_size);
+
+// Llama 3.1
+            // Check if M_PI is defined? Sorry haven't coded in C in ages lol
+            float wavelen = 2.0f * M_PI / freq;
+            if      (wavelen < high_freq_wavelen) { freq = freq; }
+            else if (wavelen > low_freq_wavelen)  { freq = freq / scale_factor; }
+            else {
+                float smooth = (old_context_len / wavelen - low_freq_factor) / \
+                    (high_freq_factor - low_freq_factor);
+
+                freq = (1.0 - smooth) * freq / scale_factor + smooth * freq;
+            }
+// END Llama 3.1
+
 // END L2E Addition
             float val = pos * freq;
             float fcr = cosf(val);