Change Rust view ops to take element rather than byte offsets.

KerfuffleV2 · Apr 10, 2023 · 825a33b · 825a33b
1 parent 08b3172
commit 825a33b
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 10 deletions.
diff --git a/ggml/src/lib.rs b/ggml/src/lib.rs
@@ -272,6 +272,7 @@ impl Context {
 
     /// Creates a 1D view over `a`.
     pub fn op_view_1d(&self, a: &Tensor, ne0: usize, offset: usize) -> Tensor {
+        let offset = offset * a.element_size();
         let tensor = unsafe {
             ggml_sys::ggml_view_1d(self.ptr.as_ptr(), a.ptr.as_ptr(), usize_to_i64(ne0), offset)
         };
@@ -287,6 +288,10 @@ impl Context {
         nb1: usize,
         offset: usize,
     ) -> Tensor {
+        let elsize = a.element_size();
+        let offset = offset * elsize;
+        let nb1 = nb1 * elsize;
+
         let tensor = unsafe {
             ggml_sys::ggml_view_2d(
                 self.ptr.as_ptr(),
@@ -312,6 +317,11 @@ impl Context {
         nb2: usize,
         offset: usize,
     ) -> Tensor {
+        let elsize = a.element_size();
+        let offset = offset * a.element_size();
+        let nb1 = nb1 * elsize;
+        let nb2 = nb2 * elsize;
+
         let tensor = unsafe {
             ggml_sys::ggml_view_3d(
                 self.ptr.as_ptr(),

diff --git a/llama-rs/src/lib.rs b/llama-rs/src/lib.rs
@@ -1130,9 +1130,6 @@ impl Model {
         let n_past = session.n_past;
         let n_threads = params.n_threads;
 
-        let memk_elsize = session.memory_k.element_size();
-        let memv_elsize = session.memory_v.element_size();
-
         let Hyperparameters {
             n_vocab,
             n_ctx,
@@ -1213,15 +1210,15 @@ impl Model {
                     let k = ctx0.op_view_1d(
                         &session.memory_k,
                         n * n_embd,
-                        (memk_elsize * n_embd) * (il * n_ctx + n_past),
+                        n_embd * (il * n_ctx + n_past),
                     );
 
                     let v = ctx0.op_view_2d(
                         &session.memory_v,
                         n,
                         n_embd,
-                        n_ctx * memv_elsize,
-                        (il * n_ctx) * memv_elsize * n_embd + n_past * memv_elsize,
+                        n_ctx,
+                        (il * n_ctx) * n_embd + n_past,
                     );
 
                     // important: storing RoPE-ed version of K in the KV cache!
@@ -1236,7 +1233,7 @@ impl Model {
                         &ctx0.op_view_1d(
                             &session.memory_k,
                             (n_past + n) * n_embd,
-                            il * n_ctx * memk_elsize * n_embd,
+                            il * n_ctx * n_embd,
                         ),
                         n_embd / n_head,
                         n_head,
@@ -1269,9 +1266,9 @@ impl Model {
                     n_past + n,
                     n_embd / n_head,
                     n_head,
-                    n_ctx * memv_elsize,
-                    n_ctx * memv_elsize * n_embd / n_head,
-                    il * n_ctx * memv_elsize * n_embd,
+                    n_ctx,
+                    n_ctx * n_embd / n_head,
+                    il * n_ctx * n_embd,
                 );
 
                 let k_q_v = ctx0.op_mul_mat(&v, &k_q_soft_max);