Load GLM vision weights on demand

This commit is contained in:
Georg Bauer
2026-09-01 20:08:36 +02:00
parent 46d6a976a5
commit ee515ee824
10 changed files with 293 additions and 26 deletions

View File

@@ -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<vision::VisionEmbedding, String> {
pub(super) fn encode_visions(
&self,
encoded: &[Vec<u8>],
) -> Result<Vec<vision::VisionEmbedding>, 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,

View File

@@ -650,6 +650,7 @@ pub(in crate::engine) struct GlmExecutor {
model_identity: [u8; 32],
streaming_spans: Option<Vec<(u64, u64)>>,
vision: Option<VisionEncoder>,
keep_vision_loaded: bool,
vision_overlays: Vec<VisionOverlay>,
_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<VisionEmbedding, String> {
pub(super) fn encode_visions(
&self,
encoded: &[Vec<u8>],
) -> Result<Vec<VisionEmbedding>, 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::<Result<Vec<_>, _>>();
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<u64, String> {
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,

View File

@@ -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<Self, String> {
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(),