diff --git a/src/multimodal/vlm_runtime.rs b/src/multimodal/vlm_runtime.rs index 796638fae..c0ad9aa93 100644 --- a/src/multimodal/vlm_runtime.rs +++ b/src/multimodal/vlm_runtime.rs @@ -1256,13 +1256,25 @@ where total_image_tokens: stats.total_image_tokens, }); - // LFM2-VL runs every image in one preparation pass; skip the - // opportunistic vision cache for this first integration. - let _ = active_caches; - let _ = image_cache_keys; + // LFM2-VL runs every image in one preparation pass. The cache key + // covers the concatenated pixel bytes for all images in the + // request; multi-turn conversations that revisit the same image + // set skip the vision tower + connector on subsequent turns. + let lfm2_cache_key = if active_caches.is_some() { + first_explicit_key(image_cache_keys) + .or_else(|| Some(CacheKey::from_hash(image_hash_from_pixels(&pixel_values)))) + } else { + None + }; let input_ids_arr = prompt_ids_array(prompt_tokens); - let embeddings = lfm2vl.get_input_embeddings(&input_ids_arr, &pixel_values, &grids); + let embeddings = lfm2vl.get_input_embeddings_with_cache( + &input_ids_arr, + &pixel_values, + &grids, + lfm2_cache_key.as_ref(), + active_caches.map(|c| &c.single), + ); Ok(Some(PreparedVlmEmbeddings { embeddings, diff --git a/src/vision/lfm2_vl.rs b/src/vision/lfm2_vl.rs index 291c9243a..dd453099a 100644 --- a/src/vision/lfm2_vl.rs +++ b/src/vision/lfm2_vl.rs @@ -40,6 +40,7 @@ use mlxcel_core::{MlxArray, UniquePtr}; use crate::models::lfm2::Lfm2Model; use crate::vision::connectors::lfm2_vl::Lfm2VlConnector; use crate::vision::encoders::lfm2_vl::Lfm2VlVisionTower; +use crate::vision::feature_cache::{CacheKey, SingleArrayFeatures, VisionFeatureCache}; use crate::vision::merge::{self, InputEmbeddings}; use crate::vision::processors::lfm2_vl::Lfm2VlProcessor; @@ -72,31 +73,82 @@ impl Lfm2VlModel { input_ids: &MlxArray, pixel_values: &MlxArray, grids: &[(i32, i32)], + ) -> InputEmbeddings { + self.get_input_embeddings_with_cache(input_ids, pixel_values, grids, None, None) + } + + /// Get input embeddings with optional vision feature caching. + /// + /// LFM2-VL runs every image through the tower + connector independently + /// (variable patch count per image), and the resulting per-image feature + /// rows are concatenated before the LLaVA-style merge. The cache stores + /// the concatenated post-connector tensor keyed by the request-supplied + /// image bytes (SHA-256 + soft-token budget), so multi-turn conversations + /// that revisit the same image skip the vision tower + connector on + /// subsequent turns. + /// + /// Note: the connector output shape depends on the image's native patch + /// count, which is folded into the key by the soft-token budget; two + /// different resolutions of the same image cannot collide. + pub fn get_input_embeddings_with_cache( + &self, + input_ids: &MlxArray, + pixel_values: &MlxArray, + grids: &[(i32, i32)], + cache_key: Option<&CacheKey>, + vision_cache: Option<&std::sync::Mutex>>, ) -> InputEmbeddings { let inputs_embeds = self.text_model.input_embeddings(input_ids); let embed_dtype = mlxcel_core::array_dtype(&inputs_embeds); - let pv = mlxcel_core::astype(pixel_values, embed_dtype); - - let mut features: Vec> = Vec::with_capacity(grids.len()); - let mut offset = 0i32; - for &(h, w) in grids { - let n = h * w; - // patches_i: [1, n, patch_dim] - let patches = - mlxcel_core::slice(&pv, &[0, offset, 0], &[1, offset + n, self.patch_dim]); - offset += n; - let vision_out = self.vision_tower.forward(&patches, (h, w)); - features.push(self.connector.forward(&vision_out, (h, w))); - } - - let image_features = match features.len() { - 0 => mlxcel_core::astype(&inputs_embeds, embed_dtype), // unused; no images - 1 => features.into_iter().next().unwrap(), - _ => { - let mut iter = features.into_iter(); - let first = iter.next().unwrap(); - iter.fold(first, |acc, next| mlxcel_core::concatenate(&acc, &next, 0)) + + // Cache lookup: reuse the concatenated post-connector features when + // the same image bytes have already been encoded. + let cached_features = match (cache_key, vision_cache) { + (Some(key), Some(cache)) => cache + .lock() + .ok() + .and_then(|mut guard| guard.get(key)) + .map(|c| c.features), + _ => None, + }; + + let image_features: UniquePtr = if let Some(cached) = cached_features { + cached + } else { + let pv = mlxcel_core::astype(pixel_values, embed_dtype); + let mut features: Vec> = Vec::with_capacity(grids.len()); + let mut offset = 0i32; + for &(h, w) in grids { + let n = h * w; + let patches = + mlxcel_core::slice(&pv, &[0, offset, 0], &[1, offset + n, self.patch_dim]); + offset += n; + let vision_out = self.vision_tower.forward(&patches, (h, w)); + features.push(self.connector.forward(&vision_out, (h, w))); } + + let combined = match features.len() { + 0 => mlxcel_core::astype(&inputs_embeds, embed_dtype), + 1 => features.into_iter().next().unwrap(), + _ => { + let mut iter = features.into_iter(); + let first = iter.next().unwrap(); + iter.fold(first, |acc, next| mlxcel_core::concatenate(&acc, &next, 0)) + } + }; + + // Populate the cache on miss so subsequent turns can skip the + // vision tower + connector. Materialize before copying so we + // cache a concrete array rather than a deferred graph node. + if let (Some(key), Some(cache)) = (cache_key, vision_cache) { + mlxcel_core::eval(&combined); + let snapshot = SingleArrayFeatures::new(mlxcel_core::copy(combined.as_ref().unwrap())); + if let Ok(mut guard) = cache.lock() { + guard.put(key.clone(), &snapshot); + } + } + + combined }; merge::merge_llava(