diff --git a/candle-metal-kernels/src/lib.rs b/candle-metal-kernels/src/lib.rs
index 8cb3c16a..4c0f9223 100644
--- a/candle-metal-kernels/src/lib.rs
+++ b/candle-metal-kernels/src/lib.rs
@@ -1369,11 +1369,12 @@ pub fn call_gemm(
         }
 
         let matrix_offsets = device.new_buffer_with_data(
-            buffer.as_ptr() as *const NSUInteger as *const c_void,
+            buffer.as_ptr() as *const c_void,
             (buffer.len() * core::mem::size_of::<u64>()) as NSUInteger,
-            MTLResourceOptions::StorageModePrivate,
+            MTLResourceOptions::StorageModeManaged,
         );
         encoder.set_buffer(10, Some(&matrix_offsets), 0);
+        encoder.use_resource(&matrix_offsets, metal::MTLResourceUsage::Read);
     }
 
     let grid_size = MTLSize {
diff --git a/candle-metal-kernels/src/libMetalFlashAttention.metallib b/candle-metal-kernels/src/libMetalFlashAttention.metallib
index f5116ca6..5ed9d033 100644
Binary files a/candle-metal-kernels/src/libMetalFlashAttention.metallib and b/candle-metal-kernels/src/libMetalFlashAttention.metallib differ