diff --git a/native/metal/ds4_gpu.h b/native/metal/ds4_gpu.h index a88fc09..caf7b83 100644 --- a/native/metal/ds4_gpu.h +++ b/native/metal/ds4_gpu.h @@ -118,6 +118,10 @@ int ds4_gpu_build_derived_artifacts(const void *model_map, uint64_t model_size, int ds4_gpu_model_range_replaced(const void *model_map, uint64_t offset, uint64_t bytes); int ds4_gpu_set_model_map_range(const void *model_map, uint64_t model_size, uint64_t map_offset, uint64_t map_size, uint64_t max_tensor_bytes); +int ds4_gpu_set_transient_model_map_range(const void *model_map, uint64_t model_size, uint64_t map_offset, uint64_t map_size, uint64_t max_tensor_bytes); +/* Caller must finish every command that references this mapping first. */ +int ds4_gpu_release_transient_model_map(const void *model_map, uint64_t model_size); +int ds4_gpu_model_map_active(const void *model_map, uint64_t model_size); /* Add a secondary GGUF mapping without replacing the primary model mapping. */ int ds4_gpu_set_aux_model_map_range(const void *model_map, uint64_t model_size, diff --git a/native/metal/ds4_metal.m b/native/metal/ds4_metal.m index adae805..82d3746 100644 --- a/native/metal/ds4_metal.m +++ b/native/metal/ds4_metal.m @@ -1329,7 +1329,7 @@ static uint64_t ds4_gpu_effective_model_max_tensor_bytes(uint64_t map_size, uint } static id ds4_gpu_get_pipeline(const char *function_name); -static int ds4_gpu_warm_model_views(void); +static int ds4_gpu_warm_model_views(uint32_t first_view); static double ds4_gpu_gib(uint64_t bytes); static double ds4_gpu_now_ms(void) { @@ -1934,7 +1934,8 @@ static int ds4_gpu_add_model_view_range( static int ds4_gpu_finish_model_views( double t0, uint64_t mapped_model_size, - uint64_t display_offset) { + uint64_t display_offset, + uint32_t first_new_view) { const double t_mapped = ds4_gpu_now_ms(); const int request_residency = !g_ssd_streaming_mode && @@ -1968,7 +1969,7 @@ static int ds4_gpu_finish_model_views( warmed = 1; } else { ds4_gpu_progress_begin("warming Metal model views"); - warmed = ds4_gpu_warm_model_views(); + warmed = ds4_gpu_warm_model_views(first_new_view); if (warmed) ds4_gpu_progress_done(); else ds4_gpu_progress_failed(); } @@ -1994,6 +1995,7 @@ static int ds4_gpu_map_model_views( uint64_t map_size, uint64_t max_tensor_bytes) { const double t0 = ds4_gpu_now_ms(); + const uint32_t first_new_view = g_model_view_count; uint64_t mapped_model_size = 0; if (!ds4_gpu_add_model_view_range(model_map, model_size, @@ -2004,7 +2006,8 @@ static int ds4_gpu_map_model_views( &mapped_model_size)) { return 0; } - return ds4_gpu_finish_model_views(t0, mapped_model_size, map_offset); + return ds4_gpu_finish_model_views(t0, mapped_model_size, map_offset, + first_new_view); } static id ds4_gpu_new_transient_buffer(NSUInteger bytes, const char *label) { @@ -2535,8 +2538,8 @@ static void ds4_gpu_detect_metal4_features(void) { #endif } -static int ds4_gpu_warm_model_views(void) { - if (g_model_view_count == 0) return 1; +static int ds4_gpu_warm_model_views(uint32_t first_view) { + if (first_view >= g_model_view_count) return 1; id pipeline = ds4_gpu_get_pipeline("kernel_touch_u8_stride"); if (!pipeline) return 0; @@ -2562,7 +2565,7 @@ static int ds4_gpu_warm_model_views(void) { } uint64_t total_touches = 0; - for (uint32_t i = 0; i < g_model_view_count; i++) { + for (uint32_t i = first_view; i < g_model_view_count; i++) { total_touches += (g_model_views[i].bytes + stride - 1) / stride; } if (total_touches == 0 || total_touches > (uint64_t)NSUIntegerMax) return 0; @@ -2585,7 +2588,7 @@ static int ds4_gpu_warm_model_views(void) { id enc = ds4_gpu_compute_encoder(cb); [enc setComputePipelineState:pipeline]; uint64_t dst_offset = 0; - for (uint32_t i = 0; i < g_model_view_count; i++) { + for (uint32_t i = first_view; i < g_model_view_count; i++) { const uint64_t bytes = g_model_views[i].bytes; const uint64_t n = (bytes + stride - 1) / stride; [enc setBuffer:g_model_views[i].buffer offset:0 atIndex:0]; @@ -11336,6 +11339,109 @@ int ds4_gpu_set_model_map_range(const void *model_map, uint64_t model_size, uint } } +int ds4_gpu_model_map_active(const void *model_map, uint64_t model_size) { + if (!model_map || model_size == 0) return 0; + for (uint32_t i = 0; i < g_model_view_count; i++) { + if (g_model_views[i].model_map == model_map && + g_model_views[i].model_size == model_size) return 1; + } + return 0; +} + +int ds4_gpu_set_transient_model_map_range( + const void *model_map, + uint64_t model_size, + uint64_t map_offset, + uint64_t map_size, + uint64_t max_tensor_bytes) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!model_map || model_size == 0 || map_offset > model_size || map_size == 0 || + map_size > model_size - map_offset) return 0; + + const uint64_t end = map_offset + map_size; + for (uint32_t i = 0; i < g_model_view_count; i++) { + const uint64_t view_start = g_model_views[i].model_offset; + const uint64_t view_end = view_start + g_model_views[i].bytes; + if (g_model_views[i].model_map == model_map && + g_model_views[i].model_size == model_size && + map_offset >= view_start && end <= view_end) return 1; + } + + @autoreleasepool { + const uint32_t first_view = g_model_view_count; + uint64_t mapped_size = 0; + max_tensor_bytes = ds4_gpu_effective_model_max_tensor_bytes(map_size, + max_tensor_bytes); + if (ds4_gpu_add_model_view_range(model_map, model_size, map_offset, map_size, + max_tensor_bytes, false, &mapped_size)) { + return 1; + } + while (g_model_view_count > first_view) { + g_model_view_count--; + if (g_model_wrap_count > 0) g_model_wrap_count--; + if (g_model_wrap_bytes >= g_model_views[g_model_view_count].bytes) { + g_model_wrap_bytes -= g_model_views[g_model_view_count].bytes; + } else { + g_model_wrap_bytes = 0; + } + g_model_views[g_model_view_count].buffer = nil; + g_model_views[g_model_view_count].model_map = NULL; + g_model_views[g_model_view_count].model_size = 0; + g_model_views[g_model_view_count].model_offset = 0; + g_model_views[g_model_view_count].bytes = 0; + } + return 0; + } +} + +int ds4_gpu_release_transient_model_map(const void *model_map, uint64_t model_size) { + if (!model_map || model_size == 0) return 0; + if (!g_initialized || !ds4_gpu_model_map_active(model_map, model_size)) return 1; + + for (uint32_t i = 0; i < g_model_residency_count; i++) { + if (g_model_views[i].model_map == model_map && + g_model_views[i].model_size == model_size) return 0; + } + + @autoreleasepool { + uint32_t kept = 0; + for (uint32_t i = 0; i < g_model_view_count; i++) { + if (g_model_views[i].model_map == model_map && + g_model_views[i].model_size == model_size) { + if (g_model_wrap_count > 0) g_model_wrap_count--; + if (g_model_wrap_bytes >= g_model_views[i].bytes) { + g_model_wrap_bytes -= g_model_views[i].bytes; + } else { + g_model_wrap_bytes = 0; + } + g_model_views[i].buffer = nil; + g_model_views[i].model_map = NULL; + g_model_views[i].model_size = 0; + g_model_views[i].model_offset = 0; + g_model_views[i].bytes = 0; + continue; + } + if (kept != i) { + g_model_views[kept] = g_model_views[i]; + g_model_views[i].buffer = nil; + g_model_views[i].model_map = NULL; + g_model_views[i].model_size = 0; + g_model_views[i].model_offset = 0; + g_model_views[i].bytes = 0; + } + kept++; + } + g_model_view_count = kept; + g_model_wrap_max_bytes = 0; + for (uint32_t i = 0; i < g_model_view_count; i++) { + if (g_model_views[i].bytes > g_model_wrap_max_bytes) { + g_model_wrap_max_bytes = g_model_views[i].bytes; + } + } + return 1; + } +} + static int ds4_gpu_model_views_cover_spans( const void *model_map, uint64_t model_size, @@ -11394,7 +11500,7 @@ int ds4_gpu_set_model_map_spans( return 0; } } - if (!ds4_gpu_finish_model_views(t0, mapped_total, first_offset)) { + if (!ds4_gpu_finish_model_views(t0, mapped_total, first_offset, 0)) { ds4_gpu_model_residency_clear(); ds4_gpu_model_views_clear(); return 0; diff --git a/src/app.rs b/src/app.rs index 6e09930..3f15cdd 100644 --- a/src/app.rs +++ b/src/app.rs @@ -449,6 +449,7 @@ pub(crate) enum Message { PreferenceWarmWeightsChanged(bool), PreferenceGlmMtpChanged(bool), PreferenceGlmMtpTimingChanged(bool), + PreferenceKeepVisionLoadedChanged(bool), PreferenceDsparkConfidenceChanged(String), PreferenceDsparkStrictChanged(bool), PreferenceDsparkExactSamplingChanged(bool), diff --git a/src/app/preferences.rs b/src/app/preferences.rs index 87790e1..53f4549 100644 --- a/src/app/preferences.rs +++ b/src/app/preferences.rs @@ -40,6 +40,7 @@ pub(super) struct PreferenceDraft { pub(super) warm_weights: bool, pub(super) glm_mtp: bool, pub(super) glm_mtp_timing: bool, + pub(super) keep_vision_loaded: bool, pub(super) dspark_confidence_threshold: String, pub(super) dspark_strict: bool, pub(super) dspark_exact_sampling: bool, @@ -104,6 +105,7 @@ impl PreferenceDraft { warm_weights: execution.warm_weights, glm_mtp: speculative.glm_mtp, glm_mtp_timing: speculative.glm_mtp_timing, + keep_vision_loaded: speculative.keep_vision_loaded, dspark_confidence_threshold: optional_string(speculative.dspark_confidence_threshold), dspark_strict: speculative.dspark_strict, dspark_exact_sampling: speculative.dspark_exact_sampling, @@ -213,6 +215,7 @@ impl PreferenceDraft { Ok(SpeculativePreferences { glm_mtp: self.glm_mtp, glm_mtp_timing: self.glm_mtp_timing, + keep_vision_loaded: self.keep_vision_loaded, dspark_enabled: self.dspark_enabled, dspark_confidence_threshold: parse_optional_f32( "DSpark confidence", @@ -256,6 +259,7 @@ impl PreferenceDraft { self.dspark_enabled = speculative.dspark_enabled; self.glm_mtp = speculative.glm_mtp; self.glm_mtp_timing = speculative.glm_mtp_timing; + self.keep_vision_loaded = speculative.keep_vision_loaded; self.dspark_confidence_threshold = optional_string(speculative.dspark_confidence_threshold); self.dspark_strict = speculative.dspark_strict; self.dspark_exact_sampling = speculative.dspark_exact_sampling; @@ -870,6 +874,11 @@ impl App { } self.preference_error = None; } + Message::PreferenceKeepVisionLoadedChanged(value) => { + self.preference_draft.keep_vision_loaded = + self.preference_draft.acceleration_model == ModelChoice::Glm53Flash && value; + self.preference_error = None; + } Message::PreferenceDsparkConfidenceChanged(value) => { self.preference_draft.dspark_confidence_threshold = value; if self.preference_draft.acceleration_model.supports_dspark() @@ -1013,10 +1022,15 @@ mod tests { draft.select_acceleration(ModelChoice::Glm52).unwrap(); assert!(!draft.ssd_streaming); + draft.select_acceleration(ModelChoice::Glm53Flash).unwrap(); + draft.keep_vision_loaded = true; draft .select_acceleration(ModelChoice::DeepSeekV4Flash0731) .unwrap(); assert!(draft.ssd_streaming); + assert!(!draft.keep_vision_loaded); + draft.select_acceleration(ModelChoice::Glm53Flash).unwrap(); + assert!(draft.keep_vision_loaded); } #[test] diff --git a/src/app/view/preferences.rs b/src/app/view/preferences.rs index 8095502..dc603ce 100644 --- a/src/app/view/preferences.rs +++ b/src/app/view/preferences.rs @@ -24,6 +24,9 @@ impl App { .acceleration_model .supports_glm_mtp() .then_some(Message::PreferenceGlmMtpTimingChanged); + let keep_vision_loaded_toggle: Option Message> = + (self.preference_draft.acceleration_model == ModelChoice::Glm53Flash) + .then_some(Message::PreferenceKeepVisionLoadedChanged); let dspark_strict_toggle: Option Message> = self .preference_draft .acceleration_model @@ -615,6 +618,12 @@ impl App { .on_toggle_maybe(glm_mtp_timing_toggle), "Records per-stage timings of the speculative path to the log, to show where the acceleration actually goes. A diagnostic aid that costs a little throughput.", ), + hint( + toggle(self.preference_draft.keep_vision_loaded) + .label("Keep GLM 5.3 vision weights loaded") + .on_toggle_maybe(keep_vision_loaded_toggle), + "Keeps the vision encoder mapped between image turns for lower image latency. Off releases it after encoding all images in a turn, leaving more memory for long text contexts.", + ), dspark, preference_input_row( "DSpark confidence threshold", diff --git a/src/engine.rs b/src/engine.rs index 3bde425..e11e406 100644 --- a/src/engine.rs +++ b/src/engine.rs @@ -312,7 +312,7 @@ impl Model { if let Some(path) = &settings.artifacts.vision { validate_vision_artifact(path)?; let vision = Gguf::open(path)?; - if settings.execution.warm_weights { + if settings.execution.warm_weights && settings.speculative.keep_vision_loaded { vision.warm()?; } model.vision = Some(vision); @@ -1166,20 +1166,18 @@ impl Generator { reasoning: ReasoningMode, ) -> Result<(Vec, VisionOverlays), String> { let mut rendered = messages.to_vec(); - let mut embeddings = Vec::new(); + let mut images = Vec::new(); let mut total_images = 0_usize; let mut total_bytes = 0_usize; - for message in &mut rendered { + for message in &rendered { if !message.content.contains(VISION_DATA_START) { continue; } if !message.user && !message.tool { return Err("vision input is allowed only in user or tool messages".into()); } - let mut content = String::with_capacity(message.content.len()); let mut rest = message.content.as_str(); while let Some(start) = rest.find(VISION_DATA_START) { - content.push_str(&rest[..start]); let encoded = &rest[start + VISION_DATA_START.len()..]; let end = encoded .find(VISION_DATA_END) @@ -1202,7 +1200,27 @@ impl Generator { if total_bytes > 64 * 1024 * 1024 { return Err("image inputs exceed the 64 MiB request limit".into()); } - let embedding = self.executor.encode_vision(&bytes)?; + images.push(bytes); + rest = &encoded[end + VISION_DATA_END.len()..]; + } + } + let mut encoded_images = self.executor.encode_visions(&images)?.into_iter(); + let mut embeddings = Vec::with_capacity(images.len()); + for message in &mut rendered { + if !message.content.contains(VISION_DATA_START) { + continue; + } + let mut content = String::with_capacity(message.content.len()); + let mut rest = message.content.as_str(); + while let Some(start) = rest.find(VISION_DATA_START) { + content.push_str(&rest[..start]); + let encoded = &rest[start + VISION_DATA_START.len()..]; + let end = encoded + .find(VISION_DATA_END) + .ok_or("unterminated image input")?; + let embedding = encoded_images + .next() + .ok_or("vision encoder returned too few embeddings")?; content.push_str(VISION_TOKEN_START); content.push_str(&embedding.tokens.to_string()); content.push_str(VISION_TOKEN_END); @@ -1212,6 +1230,9 @@ impl Generator { content.push_str(rest); message.content = content; } + if encoded_images.next().is_some() { + return Err("vision encoder returned too many embeddings".into()); + } let tokens = self .executor .model() diff --git a/src/engine/metal.rs b/src/engine/metal.rs index d85ddef..40556bb 100644 --- a/src/engine/metal.rs +++ b/src/engine/metal.rs @@ -3056,7 +3056,7 @@ impl DeepSeekExecutor { } else { resident_deepseek_admission_bytes(&model, context, prefill_chunk)? }; - let context_handle = Context::open(&model, quality, ssd.enabled, admission, spans)?; + let context_handle = Context::open(&model, quality, ssd.enabled, admission, spans, false)?; let mxfp4_decode_fast_lookup = mxfp4_decode_fast_lookup_allowed(&model, &weights, quality, ssd.enabled); let steering = Steering::load(&model, steering)?; @@ -4387,6 +4387,7 @@ impl Executor { EngineSpeculativeSettings { glm_mtp: false, glm_mtp_timing: false, + keep_vision_loaded: false, dspark: false, dspark_confidence_threshold: 0.9, dspark_confidence_threshold_set: false, @@ -4490,9 +4491,12 @@ impl Executor { } } - pub(super) fn encode_vision(&self, encoded: &[u8]) -> Result { + pub(super) fn encode_visions( + &self, + encoded: &[Vec], + ) -> Result, String> { match self { - Self::Glm(executor) => executor.encode_vision(encoded), + Self::Glm(executor) => executor.encode_visions(encoded), Self::DeepSeek(_) => Err("vision input requires GLM 5.3 Flash".into()), } } @@ -8526,6 +8530,7 @@ mod tests { EngineSpeculativeSettings { glm_mtp: false, glm_mtp_timing: false, + keep_vision_loaded: false, dspark, dspark_confidence_threshold: confidence.unwrap_or(0.8), dspark_confidence_threshold_set: confidence.is_some(), @@ -8714,6 +8719,7 @@ mod tests { EngineSpeculativeSettings { glm_mtp: false, glm_mtp_timing: false, + keep_vision_loaded: false, dspark: false, dspark_confidence_threshold: 0.9, dspark_confidence_threshold_set: false, @@ -8799,6 +8805,7 @@ mod tests { EngineSpeculativeSettings { glm_mtp: false, glm_mtp_timing: false, + keep_vision_loaded: false, dspark: true, dspark_confidence_threshold: 0.9, dspark_confidence_threshold_set: false, @@ -8945,6 +8952,7 @@ mod tests { EngineSpeculativeSettings { glm_mtp: false, glm_mtp_timing: false, + keep_vision_loaded: false, dspark: true, dspark_confidence_threshold: 0.6, dspark_confidence_threshold_set: false, @@ -9057,6 +9065,7 @@ mod tests { let speculative = EngineSpeculativeSettings { glm_mtp: false, glm_mtp_timing: false, + keep_vision_loaded: false, dspark: false, dspark_confidence_threshold: 0.9, dspark_confidence_threshold_set: false, @@ -9149,6 +9158,7 @@ mod tests { EngineSpeculativeSettings { glm_mtp: false, glm_mtp_timing: false, + keep_vision_loaded: false, dspark: false, dspark_confidence_threshold: 0.9, dspark_confidence_threshold_set: false, @@ -9242,6 +9252,7 @@ mod tests { EngineSpeculativeSettings { glm_mtp: false, glm_mtp_timing: false, + keep_vision_loaded: false, dspark: false, dspark_confidence_threshold: 0.9, dspark_confidence_threshold_set: false, @@ -9325,6 +9336,7 @@ mod tests { EngineSpeculativeSettings { glm_mtp: false, glm_mtp_timing: false, + keep_vision_loaded: false, dspark: false, dspark_confidence_threshold: 0.9, dspark_confidence_threshold_set: false, @@ -9399,6 +9411,7 @@ mod tests { EngineSpeculativeSettings { glm_mtp: false, glm_mtp_timing: false, + keep_vision_loaded: false, dspark: false, dspark_confidence_threshold: 0.9, dspark_confidence_threshold_set: false, diff --git a/src/engine/metal/glm.rs b/src/engine/metal/glm.rs index 598afed..5bdd6ee 100644 --- a/src/engine/metal/glm.rs +++ b/src/engine/metal/glm.rs @@ -650,6 +650,7 @@ pub(in crate::engine) struct GlmExecutor { model_identity: [u8; 32], streaming_spans: Option>, vision: Option, + keep_vision_loaded: bool, vision_overlays: Vec, _context: Context, model: Model, @@ -680,6 +681,7 @@ impl GlmExecutor { EngineSpeculativeSettings { glm_mtp: false, glm_mtp_timing: false, + keep_vision_loaded: false, dspark: false, dspark_confidence_threshold: 0.9, dspark_confidence_threshold_set: false, @@ -713,7 +715,13 @@ impl GlmExecutor { let weights = GlmWeights::bind(&model)?; let streaming = glm_streaming_plan(&model, &weights, ssd)?; let effective_ssd = streaming.map_or(ssd, |plan| plan.settings); - let admission = admission_bytes(&model, &weights, context, streaming.as_ref())?; + let admission = admission_bytes( + &model, + &weights, + context, + streaming.as_ref(), + speculative.keep_vision_loaded, + )?; let model_spans = streaming .as_ref() .map(|plan| glm_streaming_model_spans(&model, &weights, plan)) @@ -727,6 +735,7 @@ impl GlmExecutor { effective_ssd.enabled, admission, context_spans, + speculative.keep_vision_loaded, )?; let model_spans = model_spans.map(|(spans, _)| spans); configure_streaming(&model, &weights, streaming.as_ref())?; @@ -826,13 +835,17 @@ impl GlmExecutor { model_identity, streaming_spans: model_spans, vision, + keep_vision_loaded: speculative.keep_vision_loaded, vision_overlays: Vec::new(), _context: context_handle, model, }) } - pub(super) fn encode_vision(&self, encoded: &[u8]) -> Result { + pub(super) fn encode_visions( + &self, + encoded: &[Vec], + ) -> Result, String> { let encoder = self .vision .as_ref() @@ -842,7 +855,38 @@ impl GlmExecutor { .vision .as_ref() .ok_or("GLM 5.3 vision GGUF is not loaded")?; - encoder.encode(vision, encoded) + if !self.keep_vision_loaded { + call( + unsafe { + ds4_gpu_set_transient_model_map_range( + vision.map_ptr().cast(), + vision.len(), + vision.data_offset(), + vision.len() - vision.data_offset(), + vision.max_tensor_bytes(), + ) + }, + "mapping GLM 5.3 vision weights", + )?; + } + let result = encoded + .iter() + .map(|image| encoder.encode(vision, image)) + .collect::, _>>(); + let released = if self.keep_vision_loaded { + Ok(()) + } else { + call( + unsafe { + ds4_gpu_release_transient_model_map(vision.map_ptr().cast(), vision.len()) + }, + "releasing GLM 5.3 vision weights", + ) + }; + match (result, released) { + (Ok(embeddings), Ok(())) => Ok(embeddings), + (Err(error), _) | (Ok(_), Err(error)) => Err(error), + } } pub(super) fn set_vision_overlays( @@ -5659,6 +5703,7 @@ fn admission_bytes( weights: &GlmWeights, context: u32, streaming: Option<&GlmStreamingPlan>, + keep_vision_loaded: bool, ) -> Result { let shape = model.shape; let normal_layers = u64::from(shape.layers - shape.nextn); @@ -5712,12 +5757,14 @@ fn admission_bytes( } else { model.main.len() - model.main.data_offset() }; - let resident = resident.saturating_add( + let resident = resident.saturating_add(if keep_vision_loaded { model .vision .as_ref() - .map_or(0, |vision| vision.len() - vision.data_offset()), - ); + .map_or(0, |vision| vision.len() - vision.data_offset()) + } else { + 0 + }); let cache = streaming.map_or(0, |plan| plan.planned_expert_bytes); let scratch = glm_batch_scratch_bytes(shape, context.min(INDEXED_PREFILL_CHUNK), context)? .max(512 * 1024 * 1024); @@ -6303,6 +6350,7 @@ mod tests { EngineSpeculativeSettings { glm_mtp: mtp, glm_mtp_timing: std::env::var_os("DS4SERVER_GLM53_MTP_TIMING").is_some(), + keep_vision_loaded: false, dspark: false, dspark_confidence_threshold: 0.9, dspark_confidence_threshold_set: false, @@ -6409,6 +6457,7 @@ mod tests { let mut model = Model::open_main(&model_path, ModelChoice::Glm53Flash).unwrap(); let eos = model.eos_token(); model.vision = Some(crate::engine::gguf::Gguf::open(&vision_path).unwrap()); + let keep_vision_loaded = std::env::var_os("DS4SERVER_GLM53_KEEP_VISION_LOADED").is_some(); let mut executor = GlmExecutor::open_profile( model, 4_096, @@ -6425,6 +6474,7 @@ mod tests { EngineSpeculativeSettings { glm_mtp: false, glm_mtp_timing: false, + keep_vision_loaded, dspark: false, dspark_confidence_threshold: 0.9, dspark_confidence_threshold_set: false, @@ -6439,12 +6489,32 @@ mod tests { None, ) .unwrap(); + let vision = executor.model().vision.as_ref().unwrap(); + assert_eq!( + unsafe { + super::super::gpu::ds4_gpu_model_map_active(vision.map_ptr().cast(), vision.len()) + != 0 + }, + keep_vision_loaded, + ); let encoded = std::fs::read(image_path).unwrap(); let mut runs = Vec::new(); let mut embedding = None; for _ in 0..3 { let started = Instant::now(); - embedding = Some(executor.encode_vision(&encoded).unwrap()); + embedding = executor + .encode_visions(std::slice::from_ref(&encoded)) + .unwrap() + .pop(); + assert_eq!( + unsafe { + super::super::gpu::ds4_gpu_model_map_active( + vision.map_ptr().cast(), + vision.len(), + ) != 0 + }, + keep_vision_loaded, + ); runs.push(started.elapsed().as_secs_f64()); } runs.sort_by(f64::total_cmp); @@ -6766,6 +6836,7 @@ mod tests { EngineSpeculativeSettings { glm_mtp: enabled, glm_mtp_timing: false, + keep_vision_loaded: false, dspark: false, dspark_confidence_threshold: 0.9, dspark_confidence_threshold_set: false, diff --git a/src/engine/metal/gpu.rs b/src/engine/metal/gpu.rs index 6892558..bafc937 100644 --- a/src/engine/metal/gpu.rs +++ b/src/engine/metal/gpu.rs @@ -81,6 +81,19 @@ unsafe extern "C" { map_size: u64, max_tensor_bytes: u64, ) -> i32; + pub(super) fn ds4_gpu_set_transient_model_map_range( + model_map: *const c_void, + model_size: u64, + map_offset: u64, + map_size: u64, + max_tensor_bytes: u64, + ) -> i32; + pub(super) fn ds4_gpu_release_transient_model_map( + model_map: *const c_void, + model_size: u64, + ) -> i32; + #[cfg(test)] + pub(super) fn ds4_gpu_model_map_active(model_map: *const c_void, model_size: u64) -> i32; pub(super) fn ds4_gpu_set_model_map_spans( model_map: *const c_void, model_size: u64, @@ -1801,6 +1814,7 @@ impl Context { ssd_streaming: bool, admission_bytes: u64, model_spans: Option<(&[(u64, u64)], u64)>, + map_vision: bool, ) -> Result { check(unsafe { ds4_gpu_init() }, "Metal initialization")?; unsafe { @@ -1862,7 +1876,7 @@ impl Context { return Err(error); } } - if let Some(vision) = &model.vision { + if map_vision && let Some(vision) = &model.vision { let mapped = unsafe { ds4_gpu_set_model_map_range( vision.map_ptr().cast(), diff --git a/src/settings.rs b/src/settings.rs index 6a9696c..c5f7a20 100644 --- a/src/settings.rs +++ b/src/settings.rs @@ -23,6 +23,7 @@ const DEFAULT_KV_CONTINUED_INTERVAL_TOKENS: u32 = 10_000; pub(crate) struct SpeculativePreferences { pub(crate) glm_mtp: bool, pub(crate) glm_mtp_timing: bool, + pub(crate) keep_vision_loaded: bool, pub(crate) dspark_enabled: bool, pub(crate) dspark_confidence_threshold: Option, pub(crate) dspark_strict: bool, @@ -37,6 +38,9 @@ impl SpeculativePreferences { if !model.supports_glm_mtp() && (self.glm_mtp || self.glm_mtp_timing) { return Err("GLM MTP is available only for GLM models.".into()); } + if self.keep_vision_loaded && model != ModelChoice::Glm53Flash { + return Err("Persistent vision weights are available only for GLM 5.3 Flash.".into()); + } if self.dspark_enabled && !model.supports_dspark() { return Err("DSpark is not available for the selected model.".into()); } @@ -57,6 +61,7 @@ impl SpeculativePreferences { EngineSpeculativeSettings { glm_mtp: self.glm_mtp, glm_mtp_timing: self.glm_mtp_timing, + keep_vision_loaded: self.keep_vision_loaded, dspark: self.dspark_enabled, dspark_confidence_threshold: self.dspark_confidence_threshold.unwrap_or(0.8), dspark_confidence_threshold_set: self.dspark_confidence_threshold.is_some(), @@ -70,6 +75,7 @@ impl SpeculativePreferences { pub(crate) struct EngineSpeculativeSettings { pub(crate) glm_mtp: bool, pub(crate) glm_mtp_timing: bool, + pub(crate) keep_vision_loaded: bool, pub(crate) dspark: bool, pub(crate) dspark_confidence_threshold: f32, pub(crate) dspark_confidence_threshold_set: bool, @@ -772,6 +778,14 @@ mod tests { assert!(glm.validate(ModelChoice::Glm52).is_ok()); assert!(glm.validate(ModelChoice::DeepSeekV4Pro).is_err()); + let persistent_vision = SpeculativePreferences { + keep_vision_loaded: true, + ..SpeculativePreferences::default() + }; + assert!(persistent_vision.validate(ModelChoice::Glm53Flash).is_ok()); + assert!(persistent_vision.validate(ModelChoice::Glm52).is_err()); + assert!(persistent_vision.engine_settings().keep_vision_loaded); + assert!( SpeculativePreferences { dspark_exact_sampling: true,