Add device mapper

EricLBuehler · Dec 31, 2024 · 8a0177a · 8a0177a
1 parent efbd6f4
commit 8a0177a
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/mistralrs-core/src/pipeline/mod.rs b/mistralrs-core/src/pipeline/mod.rs
@@ -19,6 +19,7 @@ mod vision;
 
 pub use super::diffusion_models::DiffusionGenerationParams;
 use crate::amoe::{AnyMoeConfig, AnyMoeExpertType, AnyMoeTrainingInputs, AnyMoeTrainingResult};
+use crate::device_map::DeviceMapper;
 use crate::paged_attention::{CacheConfig, CacheEngine, ModelConfigLike};
 use crate::prefix_cacher_v2::PrefixCacheManagerV2;
 pub use amoe::{AnyMoeLoader, AnyMoePipeline};
@@ -148,6 +149,7 @@ pub trait MetadataMixin {
     fn name(&self) -> String;
     fn reset_non_granular_state(&self);
     fn get_metadata(&self) -> Arc<GeneralMetadata>;
+    fn device_mapper(&self) -> Option<&dyn DeviceMapper>;
 }
 
 /// Implemented by the base model of an AnyMoe.
@@ -326,6 +328,7 @@ pub trait Pipeline:
                     self.get_input_processor_config(),
                     None,
                     self.get_metadata().prompt_batchsize,
+                    self.device_mapper(),
                 );
 
                 let mut logits = vec![None; input_seqs.len()];
@@ -538,6 +541,7 @@ pub trait Pipeline:
                     self.get_input_processor_config(),
                     Some(metadata),
                     self.get_metadata().prompt_batchsize,
+                    self.device_mapper(),
                 );
 
                 let mut logits = vec![None; input_seqs.len()];