Fix

EricLBuehler · Dec 23, 2024 · 90e86a1 · 90e86a1
1 parent afb8ce4
commit 90e86a1
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 7 deletions.
diff --git a/mistralrs-paged-attn/src/metal/backend/paged_attention.rs b/mistralrs-paged-attn/src/metal/backend/paged_attention.rs
@@ -1,6 +1,6 @@
 use candle_core::{
-    backend::BackendStorage, CpuStorage, DType, Device, Layout, MetalStorage, Result, Shape,
-    Storage, Tensor,
+    backend::BackendStorage, CpuStorage, DType, Layout, MetalStorage, Result, Shape, Storage,
+    Tensor,
 };
 
 use crate::metal::kernels::{self, PagedAttentionDType};
@@ -412,9 +412,7 @@ pub fn reshape_and_cache(
     let key_stride = k_l.stride()[0] as i32;
     let value_stride = v_l.stride()[0] as i32;
 
-    let Device::Metal(dev) = key.device() else {
-        panic!("Expected the key to be on a Metal device.")
-    };
+    let dev = key.device().as_metal_device()?;
 
     let command_buffer = dev.command_buffer()?;
     command_buffer.set_label("reshape-and-cache");

diff --git a/mistralrs-paged-attn/src/metal/kernels/pagedattention.metal b/mistralrs-paged-attn/src/metal/kernels/pagedattention.metal
@@ -265,8 +265,7 @@ inline float qk_dot(device T* q[N], device T* k[N]) {
   // Compute the parallel products then sum for Q*K^T (treat vector lanes separately).
   float qk = 0;
 #pragma unroll
-  for (int ii = 1; ii < N; ++ii) {
-#pragma unroll
+  for (int ii = 0; ii < N; ++ii) {
     for (int vi = 0; vi < VEC_SIZE; ++vi) {
       qk = fma(float(q[ii][vi]), float(k[ii][vi]), qk);
     }