chore: RoPE doesn't work

huggingface · Oct 4, 2024 · 4d63692 · 4d63692
1 parent 88e7c07
commit 4d63692
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 5 deletions.
diff --git a/crates/ratchet-core/src/cpu/rope.rs b/crates/ratchet-core/src/cpu/rope.rs
@@ -187,7 +187,6 @@ fn rope(src: Vec<f32>, shape: &Shape, dim: usize, base: f32, offset: usize) -> V
     println!("R2: {:?}", r2);
 
     if dim < shape[3] {
-        //outs.push_back(slice(x, {0, 0, dims}, x.shape(), s));
         outs.push(slice(&src, &[0, 0, dim], &[num_heads, seq_len, head_dim]));
     }
 

diff --git a/crates/ratchet-core/src/ops/rope.rs b/crates/ratchet-core/src/ops/rope.rs
@@ -377,10 +377,10 @@ def mlx_rope(input, dim, offset):
     fn debug_rope_cpu() {
         let prob = RoPEProblem {
             BS: 1,
-            NH: 2,
-            SL: 16,
-            HD: 16,
-            dim: 8,
+            NH: 1,
+            SL: 2,
+            HD: 128,
+            dim: 96,
             offset: 0,
         };
         println!("{prob:?}");