Fix

EricLBuehler · Nov 18, 2024 · bc3449a · bc3449a
1 parent 6dd7147
commit bc3449a
Show file tree

Hide file tree

Showing 6 changed files with 46 additions and 37 deletions.
diff --git a/mistralrs-core/src/pipeline/cache_manager.rs b/mistralrs-core/src/pipeline/cache_manager.rs
@@ -568,7 +568,7 @@ impl Cache {
     }
 }
 
-pub struct DefaultCacheManager;
+pub struct FullCacheManager;
 
 enum SeqCache {
     Normal,
@@ -655,7 +655,7 @@ fn clone_out_cache(
     }
 }
 
-impl<T: CacheManagerMixin + MetadataMixin + ?Sized> CacheManager<T> for DefaultCacheManager {
+impl<T: CacheManagerMixin + MetadataMixin + ?Sized> CacheManager<T> for FullCacheManager {
     fn clone_in_cache(
         &self,
         pipeline: &T,

diff --git a/mistralrs-core/src/pipeline/ggml.rs b/mistralrs-core/src/pipeline/ggml.rs
@@ -1,4 +1,4 @@
-use super::cache_manager::DefaultCacheManager;
+use super::cache_manager::FullCacheManager;
 use super::{
     get_model_paths, get_xlora_paths, text_models_inputs_processor::ModelInputs, AdapterKind,
     CacheManager, GeneralMetadata, Loader, ModelKind, ModelPaths, QuantizationKind, TokenSource,
@@ -456,10 +456,10 @@ impl IsqPipelineMixin for GGMLPipeline {
 
 impl CacheManagerMixin for GGMLPipeline {
     fn clone_in_cache(&self, seqs: &mut [&mut Sequence], modify_draft_cache: bool) {
-        DefaultCacheManager.clone_in_cache(self, seqs, modify_draft_cache)
+        FullCacheManager.clone_in_cache(self, seqs, modify_draft_cache)
     }
     fn clone_out_cache(&self, seqs: &mut [&mut Sequence], modify_draft_cache: bool) {
-        DefaultCacheManager.clone_out_cache(self, seqs, modify_draft_cache)
+        FullCacheManager.clone_out_cache(self, seqs, modify_draft_cache)
     }
     fn set_none_cache(
         &self,
@@ -468,7 +468,7 @@ impl CacheManagerMixin for GGMLPipeline {
         modify_draft_cache: bool,
         load_preallocated_cache: bool,
     ) {
-        DefaultCacheManager.set_none_cache(self, seqs, modify_draft_cache, load_preallocated_cache);
+        FullCacheManager.set_none_cache(self, seqs, modify_draft_cache, load_preallocated_cache);
         if reset_non_granular {
             self.reset_non_granular_state()
         }

diff --git a/mistralrs-core/src/pipeline/gguf.rs b/mistralrs-core/src/pipeline/gguf.rs
@@ -1,4 +1,4 @@
-use super::cache_manager::DefaultCacheManager;
+use super::cache_manager::{FullCacheManager, NormalCacheManager};
 use super::{
     get_model_paths, get_xlora_paths, text_models_inputs_processor::ModelInputs, AdapterKind,
     CacheManager, GeneralMetadata, Loader, ModelKind, ModelPaths, PrettyName, QuantizationKind,
@@ -487,12 +487,12 @@ impl Loader for GGUFLoader {
         let tok_trie: Arc<TokTrie> = build_tok_trie(tokenizer.clone()).into();
         let num_hidden_layers = match model {
             Model::Llama(ref model) => model.cache.normal().0.len(),
-            Model::Phi2(ref model) => model.cache.full().lock().len(),
+            Model::Phi2(ref model) => model.cache.normal().0.len(),
             Model::XLoraLlama(ref model) => model.cache.full().lock().len(),
-            Model::Phi3(ref model) => model.cache.full().lock().len(),
+            Model::Phi3(ref model) => model.cache.normal().0.len(),
             Model::XLoraPhi3(ref model) => model.cache.full().lock().len(),
-            Model::Starcoder2(ref model) => model.cache.full().lock().len(),
-            Model::Qwen2(ref model) => model.cache.full().lock().len(),
+            Model::Starcoder2(ref model) => model.cache.normal().0.len(),
+            Model::Qwen2(ref model) => model.cache.normal().0.len(),
         };
 
         if chat_template.bos_token.is_none() && bos.is_some() {
@@ -570,10 +570,18 @@ impl IsqPipelineMixin for GGUFPipeline {
 
 impl CacheManagerMixin for GGUFPipeline {
     fn clone_in_cache(&self, seqs: &mut [&mut Sequence], modify_draft_cache: bool) {
-        DefaultCacheManager.clone_in_cache(self, seqs, modify_draft_cache)
+        if matches!(self.cache(), EitherCache::Full(_)) {
+            FullCacheManager.clone_in_cache(self, seqs, modify_draft_cache)
+        } else {
+            NormalCacheManager.clone_in_cache(self, seqs, modify_draft_cache)
+        }
     }
     fn clone_out_cache(&self, seqs: &mut [&mut Sequence], modify_draft_cache: bool) {
-        DefaultCacheManager.clone_out_cache(self, seqs, modify_draft_cache)
+        if matches!(self.cache(), EitherCache::Full(_)) {
+            FullCacheManager.clone_out_cache(self, seqs, modify_draft_cache)
+        } else {
+            NormalCacheManager.clone_out_cache(self, seqs, modify_draft_cache)
+        }
     }
     fn set_none_cache(
         &self,
@@ -582,7 +590,16 @@ impl CacheManagerMixin for GGUFPipeline {
         modify_draft_cache: bool,
         load_preallocated_cache: bool,
     ) {
-        DefaultCacheManager.set_none_cache(self, seqs, modify_draft_cache, load_preallocated_cache);
+        if matches!(self.cache(), EitherCache::Full(_)) {
+            FullCacheManager.set_none_cache(self, seqs, modify_draft_cache, false);
+        } else {
+            NormalCacheManager.set_none_cache(
+                self,
+                seqs,
+                modify_draft_cache,
+                load_preallocated_cache,
+            );
+        }
         if reset_non_granular {
             self.reset_non_granular_state()
         }

diff --git a/mistralrs-core/src/pipeline/normal.rs b/mistralrs-core/src/pipeline/normal.rs
@@ -1,4 +1,4 @@
-use super::cache_manager::{DefaultCacheManager, NormalCacheManager};
+use super::cache_manager::{FullCacheManager, NormalCacheManager};
 use super::{
     get_model_paths, get_xlora_paths, text_models_inputs_processor::ModelInputs, AdapterKind,
     CacheManager, GeneralMetadata, Loader, ModelKind, ModelPaths, NormalModel, NormalModelLoader,
@@ -519,14 +519,14 @@ impl IsqPipelineMixin for NormalPipeline {
 impl CacheManagerMixin for NormalPipeline {
     fn clone_in_cache(&self, seqs: &mut [&mut Sequence], modify_draft_cache: bool) {
         if matches!(self.model.cache(), EitherCache::Full(_)) {
-            DefaultCacheManager.clone_in_cache(self, seqs, modify_draft_cache)
+            FullCacheManager.clone_in_cache(self, seqs, modify_draft_cache)
         } else {
             NormalCacheManager.clone_in_cache(self, seqs, modify_draft_cache)
         }
     }
     fn clone_out_cache(&self, seqs: &mut [&mut Sequence], modify_draft_cache: bool) {
         if matches!(self.model.cache(), EitherCache::Full(_)) {
-            DefaultCacheManager.clone_out_cache(self, seqs, modify_draft_cache)
+            FullCacheManager.clone_out_cache(self, seqs, modify_draft_cache)
         } else {
             NormalCacheManager.clone_out_cache(self, seqs, modify_draft_cache)
         }
@@ -539,7 +539,7 @@ impl CacheManagerMixin for NormalPipeline {
         load_preallocated_cache: bool,
     ) {
         if matches!(self.model.cache(), EitherCache::Full(_)) {
-            DefaultCacheManager.set_none_cache(self, seqs, modify_draft_cache, false);
+            FullCacheManager.set_none_cache(self, seqs, modify_draft_cache, false);
         } else {
             NormalCacheManager.set_none_cache(
                 self,

diff --git a/mistralrs-core/src/pipeline/speculative.rs b/mistralrs-core/src/pipeline/speculative.rs
@@ -27,7 +27,7 @@ use crate::{
 };
 
 use super::{
-    cache_manager::DefaultCacheManager, chat_template::ChatTemplate, sampling::SpeculativeSample,
+    cache_manager::FullCacheManager, chat_template::ChatTemplate, sampling::SpeculativeSample,
     AdapterActivationMixin, AnyMoePipelineMixin, CacheBackendMetadata, CacheInstruction,
     CacheManager, CacheManagerMixin, EitherCache, ForwardInputsResult, GeneralMetadata,
     IsqPipelineMixin, MetadataMixin, ModelCategory, ModelPaths, PreProcessingMixin,
@@ -246,20 +246,12 @@ impl IsqPipelineMixin for SpeculativePipeline {
 // TODO: correct handling of cloning in and out for normal cache
 impl CacheManagerMixin for SpeculativePipeline {
     fn clone_in_cache(&self, seqs: &mut [&mut Sequence], modify_draft_cache: bool) {
-        DefaultCacheManager.clone_in_cache(
-            &*get_mut_arcmutex!(self.draft),
-            seqs,
-            modify_draft_cache,
-        );
-        DefaultCacheManager.clone_in_cache(&*get_mut_arcmutex!(self.target), seqs, false);
+        FullCacheManager.clone_in_cache(&*get_mut_arcmutex!(self.draft), seqs, modify_draft_cache);
+        FullCacheManager.clone_in_cache(&*get_mut_arcmutex!(self.target), seqs, false);
     }
     fn clone_out_cache(&self, seqs: &mut [&mut Sequence], modify_draft_cache: bool) {
-        DefaultCacheManager.clone_out_cache(
-            &*get_mut_arcmutex!(self.draft),
-            seqs,
-            modify_draft_cache,
-        );
-        DefaultCacheManager.clone_out_cache(&*get_mut_arcmutex!(self.target), seqs, false);
+        FullCacheManager.clone_out_cache(&*get_mut_arcmutex!(self.draft), seqs, modify_draft_cache);
+        FullCacheManager.clone_out_cache(&*get_mut_arcmutex!(self.target), seqs, false);
     }
     fn set_none_cache(
         &self,
@@ -268,13 +260,13 @@ impl CacheManagerMixin for SpeculativePipeline {
         modify_draft_cache: bool,
         load_preallocated_cache: bool,
     ) {
-        DefaultCacheManager.set_none_cache(
+        FullCacheManager.set_none_cache(
             &*get_mut_arcmutex!(self.draft),
             seqs,
             modify_draft_cache,
             load_preallocated_cache,
         );
-        DefaultCacheManager.set_none_cache(
+        FullCacheManager.set_none_cache(
             &*get_mut_arcmutex!(self.target),
             seqs,
             false,

diff --git a/mistralrs-core/src/pipeline/vision.rs b/mistralrs-core/src/pipeline/vision.rs
@@ -1,4 +1,4 @@
-use super::cache_manager::DefaultCacheManager;
+use super::cache_manager::FullCacheManager;
 use super::isq::UqffFullSer;
 use super::{
     get_model_paths, get_xlora_paths, AdapterActivationMixin, AnyMoePipelineMixin, CacheManager,
@@ -436,10 +436,10 @@ impl IsqPipelineMixin for VisionPipeline {
 
 impl CacheManagerMixin for VisionPipeline {
     fn clone_in_cache(&self, seqs: &mut [&mut Sequence], modify_draft_cache: bool) {
-        DefaultCacheManager.clone_in_cache(self, seqs, modify_draft_cache)
+        FullCacheManager.clone_in_cache(self, seqs, modify_draft_cache)
     }
     fn clone_out_cache(&self, seqs: &mut [&mut Sequence], modify_draft_cache: bool) {
-        DefaultCacheManager.clone_out_cache(self, seqs, modify_draft_cache)
+        FullCacheManager.clone_out_cache(self, seqs, modify_draft_cache)
     }
     fn set_none_cache(
         &self,
@@ -448,7 +448,7 @@ impl CacheManagerMixin for VisionPipeline {
         modify_draft_cache: bool,
         load_preallocated_cache: bool,
     ) {
-        DefaultCacheManager.set_none_cache(self, seqs, modify_draft_cache, load_preallocated_cache);
+        FullCacheManager.set_none_cache(self, seqs, modify_draft_cache, load_preallocated_cache);
         if reset_non_granular {
             self.reset_non_granular_state()
         }