Merge pull request #1606 from FL33TW00D/feature/larger-batches

fix: larger batches
huggingface · Jan 29, 2024 · 6d83d42 · 6d83d42
2 parents fd7c856 + b6afb46
commit 6d83d42
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 7 deletions.
diff --git a/candle-metal-kernels/src/lib.rs b/candle-metal-kernels/src/lib.rs
@@ -1364,13 +1364,12 @@ pub fn call_gemm(
         // TODO byte_stride_d
         let byte_stride_d = 0;
 
-        let mut buffer: Vec<u64> = Vec::with_capacity(b * 4);
-        for i in 0..b {
-            buffer.push((i * byte_stride_a) as u64);
-            buffer.push((i * byte_stride_b) as u64);
-            buffer.push((i * byte_stride_c) as u64);
-            buffer.push((i * byte_stride_d) as u64);
-        }
+        let buffer: Vec<u64> = vec![
+            byte_stride_a as _,
+            byte_stride_b as _,
+            byte_stride_c as _,
+            byte_stride_d as _,
+        ];
         encoder.set_bytes(
             10,
             (buffer.len() * core::mem::size_of::<u64>()) as NSUInteger,

diff --git a/candle-metal-kernels/src/libMetalFlashAttention.metallib b/candle-metal-kernels/src/libMetalFlashAttention.metallib