bugfix: corrected theta calculation in RoPE CUDA kernel (#290)

Author: jiaau

kernels/rope/rope.cu

@@ -23,9 +23,9 @@ __global__ void rope_f32_kernel(float* x, float* out, int seq_len, int N){
   float x2 = x[idx * 2 + 1];
   int token_pos = idx / N; 
   int token_idx = idx % N;
-  float exp_v = 1.0f / powf(theta, token_idx / (N * 2));
-  float sin_v = sinf(token_pos / exp_v);
-  float cos_v = cosf(token_pos / exp_v);
+  float exp_v = 1.0f / powf(theta, 2 * token_idx / (N * 2.0f));
+  float sin_v = sinf(token_pos * exp_v);
+  float cos_v = cosf(token_pos * exp_v);
   float out1 = x1 * cos_v - x2 * sin_v;
   float out2 = x1 * sin_v + x2 * cos_v;
   out[idx * 2] = out1;
@@ -38,9 +38,9 @@ __global__ void rope_f32_v2_kernel(float* x, float* out, int seq_len, int N){
   int tid = threadIdx.x;
   float x1 = x[token_pos * N * 2 + tid * 2];
   float x2 = x[token_pos * N * 2 + tid * 2 + 1];
-  float exp_v = 1.0f / powf(theta, (int)(tid / 2) / (N * 2));
-  float sin_v = sinf(token_pos / exp_v);
-  float cos_v = cosf(token_pos / exp_v);
+  float exp_v = 1.0f / powf(theta, 2 * tid / (N * 2.0f));
+  float sin_v = sinf(token_pos * exp_v);
+  float cos_v = cosf(token_pos * exp_v);
   float out1 = x1 * cos_v - x2 * sin_v;
   float out2 = x1 * sin_v + x2 * cos_v;
   out[token_pos * N * 2 + tid * 2] = out1;
@@ -52,12 +52,12 @@ __global__ void rope_f32x4_pack_kernel(float* x, float* out, int seq_len, int N)
   float4 x_v = FLOAT4(x[idx * 4]);
   int token_pos = idx / N; 
   int token_idx = idx % N;
-  float exp_f_v = 1.0f / powf(theta, token_idx * 2 / (N * 4));
-  float exp_s_v = 1.0f / powf(theta, ((token_idx * 2) + 1) / (N * 4));
-  float sin_f_v = sinf(token_pos / exp_f_v);
-  float cos_f_v = cosf(token_pos / exp_f_v);
-  float sin_s_v = sinf(token_pos / exp_s_v);
-  float cos_s_v = cosf(token_pos / exp_s_v);
+  float exp_f_v = 1.0f / powf(theta, 2 * token_idx * 2 / (N * 4.0f));
+  float exp_s_v = 1.0f / powf(theta, 2 * (token_idx * 2 + 1) / (N * 4.0f));
+  float sin_f_v = sinf(token_pos * exp_f_v);
+  float cos_f_v = cosf(token_pos * exp_f_v);
+  float sin_s_v = sinf(token_pos * exp_s_v);
+  float cos_s_v = cosf(token_pos * exp_s_v);
   float4 out_v;
   out_v.x = x_v.x * cos_f_v - x_v.y * sin_f_v;
   out_v.y = x_v.x * sin_f_v + x_v.y * cos_f_v;