Better handling for kvcache set_len

mistralrs-core/src/pipeline/cache_manager.rs

@@ -101,8 +101,9 @@ impl SingleCache {
         self.all_data = None;
     }
 
-    pub fn set_len(&mut self, len: usize) {
+    pub fn set_len(&mut self, len: usize) -> candle_core::Result<()> {
         self.current_seq_len = len;
+        Ok(())
     }
 
     pub fn append(&mut self, src: &Tensor) -> Result<()> {
@@ -204,12 +205,17 @@ impl RotatingCache {
         self.all_data = None;
     }
 
-    pub fn set_len(&mut self, len: usize) {
-        if len < self.max_seq_len {
-            self.current_seq_len = len % self.max_seq_len;
-        } else {
-            // No change.
+    pub fn set_len(&mut self, len: usize) -> candle_core::Result<()> {
+        // If trying to roll it back past the boundary of max_seq_len, fail early.
+        if self.current_seq_len - len > self.max_seq_len {
+            candle_core::bail!(
+                "Rotating KV cache (usually for sliding window) tried to reset to len {len} while current is {} and max retained is {}",
+                self.current_seq_len,
+                self.max_seq_len
+            );
         }
+        self.current_seq_len = len % self.max_seq_len;
+        Ok(())
     }
 
     pub fn append(&mut self, src: &Tensor) -> Result<Tensor> {
@@ -354,15 +360,18 @@ impl KvCache {
         }
     }
 
-    pub fn set_len(&mut self, len: usize) {
+    /// Returns Ok if the length reassignment was successful, otherwise returns Err.
+    pub fn set_len(&mut self, len: usize) -> candle_core::Result<()> {
         match self {
             Self::Normal { k, v } => {
-                k.set_len(len);
-                v.set_len(len);
+                k.set_len(len)?;
+                v.set_len(len)?;
+                Ok(())
             }
             Self::Rotating { k, v } => {
-                k.set_len(len);
-                v.set_len(len);
+                k.set_len(len)?;
+                v.set_len(len)?;
+                Ok(())
             }
         }
     }

mistralrs-core/src/pipeline/normal.rs

@@ -661,7 +661,7 @@ impl Loader for NormalLoader {
                     }
                     EitherCache::Normal(normal) => {
                         for layer in &mut *normal.lock().unwrap().0 {
-                            layer.set_len(0);
+                            layer.reset();
                         }
                     }
                 }

mistralrs-core/src/pipeline/speculative.rs

@@ -559,7 +559,9 @@ impl Pipeline for SpeculativePipeline {
                     }
                     EitherCache::Normal(normal) => {
                         for cache in &mut *normal.lock().unwrap().0 {
-                            cache.set_len(cache.current_seq_len() - n_not_accepted);
+                            cache
+                                .set_len(cache.current_seq_len() - n_not_accepted)
+                                .map_err(|_| candle_core::Error::msg("KV cache set_len failed."))?;
                         }
                     }
                 }
@@ -585,7 +587,9 @@ impl Pipeline for SpeculativePipeline {
                     }
                     EitherCache::Normal(normal) => {
                         for cache in &mut *normal.lock().unwrap().0 {
-                            cache.set_len(cache.current_seq_len() - n_not_accepted);
+                            cache
+                                .set_len(cache.current_seq_len() - n_not_accepted)
+                                .map_err(|_| candle_core::Error::msg("KV cache set_len failed."))?;
                         }
                     }
                 }

mistralrs-core/src/pipeline/vision.rs

@@ -532,7 +532,7 @@ impl Loader for VisionLoader {
                     }
                     EitherCache::Normal(normal) => {
                         for layer in &mut *normal.lock().unwrap().0 {
-                            layer.set_len(0);
+                            layer.reset();
                         }
                     }
                 }

mistralrs-core/src/prefix_cacher.rs

@@ -239,7 +239,10 @@ impl PrefixCacheManagerV2 {
             let mut cache = longest_match.clone();
             Self::cache_to(&mut cache.cache, Either::Right(&cache.devices))?;
             for layer in cache.cache.iter_mut().flatten() {
-                layer.set_len(match_len);
+                match layer.set_len(match_len) {
+                    Ok(_) => (),
+                    Err(_) => return Ok(None),
+                }
             }
             Ok(Some(MatchingCache {
                 normal: cache.cache,