From 1de954b579fb4ef9c961bd95a1a35043ab6b2b96 Mon Sep 17 00:00:00 2001 From: Georg Bauer Date: Sun, 26 Jul 2026 12:02:06 +0200 Subject: [PATCH] Support DeepSeek V4 Pro execution --- src/engine.rs | 4 +-- src/engine/metal.rs | 55 +++++++++++++++++++++------------- src/engine/metal/checkpoint.rs | 4 +-- src/server.rs | 2 +- 4 files changed, 40 insertions(+), 25 deletions(-) diff --git a/src/engine.rs b/src/engine.rs index d8f0cc2..e6a6a18 100644 --- a/src/engine.rs +++ b/src/engine.rs @@ -372,8 +372,8 @@ pub(crate) struct CompactionOutput { #[cfg(target_os = "macos")] impl Generator { pub(crate) fn open(settings: &EngineSettings, metrics: Arc) -> Result { - if settings.model != ModelChoice::DeepSeekV4Flash { - return Err("local generation currently supports DeepSeek V4 Flash only".into()); + if settings.model == ModelChoice::Glm52 { + return Err("GLM 5.2 generation is not initialized by the DeepSeek executor".into()); } if settings.speculative.dspark || settings.ssd.enabled || settings.steering.file.is_some() { return Err( diff --git a/src/engine/metal.rs b/src/engine/metal.rs index 2ac3d6e..941dc6d 100644 --- a/src/engine/metal.rs +++ b/src/engine/metal.rs @@ -150,10 +150,10 @@ struct IndexerWeights { } impl Layer { - fn bind(model: &Gguf, index: u32) -> Result { + fn bind(model: &Gguf, shape: super::Shape, index: u32) -> Result { let required = |suffix: &str| Weight::bind(model, &format!("blk.{index}.{suffix}")); let optional = |suffix: &str| Weight::optional(model, &format!("blk.{index}.{suffix}")); - let attn_compressor = (index >= 2) + let attn_compressor = (compression_ratio(shape, index) != 0) .then(|| { Ok::<_, String>(CompressorWeights { ape: required("attn_compressor_ape.weight")?, @@ -163,7 +163,7 @@ impl Layer { }) }) .transpose()?; - let indexer = (compression_ratio(index) == 4) + let indexer = (compression_ratio(shape, index) == 4) .then(|| { Ok::<_, String>(IndexerWeights { q: required("indexer.attn_q_b.weight")?, @@ -215,13 +215,13 @@ impl Layer { ("shared down", layer.shared_down), ] { if weight.kind != Q8_0 { - return Err(format!("Flash Metal path requires Q8_0 {name} weights")); + return Err(format!("DeepSeek Metal path requires Q8_0 {name} weights")); } } if let Some(indexer) = layer.indexer && (!matches!(indexer.q.kind, F16 | Q8_0) || indexer.proj.kind != F16) { - return Err("Flash Metal path requires F16/Q8 indexer weights".into()); + return Err("DeepSeek Metal path requires F16/Q8 indexer weights".into()); } Ok(layer) } @@ -240,15 +240,15 @@ struct Weights { impl Weights { fn bind(model: &Model) -> Result { if model.shape.family != ModelFamily::DeepSeek { - return Err("the Rust Metal executor currently supports DeepSeek Flash only".into()); + return Err("the DeepSeek Metal executor received a different model family".into()); } let main = &model.main; let output = Weight::bind(main, "output.weight")?; if output.kind != Q8_0 { - return Err("Flash Metal path requires a Q8_0 output weight".into()); + return Err("DeepSeek Metal path requires a Q8_0 output weight".into()); } let layers = (0..model.shape.layers) - .map(|index| Layer::bind(main, index)) + .map(|index| Layer::bind(main, model.shape, index)) .collect::>()?; Ok(Self { token_embedding: Weight::bind(main, "token_embd.weight")?, @@ -486,7 +486,7 @@ struct CompressionState { impl LayerState { fn allocate(model: &Model, index: u32, context: u32, raw_cap: u32) -> Result { let shape = model.shape; - let ratio = compression_ratio(index); + let ratio = compression_ratio(shape, index); let compression = if ratio == 0 { None } else { @@ -1284,8 +1284,8 @@ fn encode_batch_layer( let hc_dim = shape.hc * shape.embd; let mix_hc = 2 * shape.hc + shape.hc * shape.hc; let q_dim = shape.heads * shape.head_dim; - let compressed = layer >= 2; - let ratio = compression_ratio(layer); + let ratio = compression_ratio(shape, layer); + let compressed = ratio != 0; let freq_base = if compressed { shape.compress_rope_base } else { @@ -2515,7 +2515,7 @@ fn encode_layer( n_comp, shape.indexer_top_k as u32, shape.sliding_window as u32, - compression_ratio(layer), + compression_ratio(shape, layer), shape.heads as u32, shape.head_dim as u32, ) @@ -2965,13 +2965,19 @@ fn update_compressor_stage( Ok((pos + 1).is_multiple_of(state.ratio)) } -fn compression_ratio(layer: u32) -> u32 { - if layer < 2 { - 0 - } else if layer.is_multiple_of(2) { - 4 - } else { - 128 +fn compression_ratio(shape: super::Shape, layer: u32) -> u32 { + match shape.model { + crate::model::ModelChoice::DeepSeekV4Flash if layer < 2 => 0, + crate::model::ModelChoice::DeepSeekV4Pro if layer < 2 => 128, + crate::model::ModelChoice::DeepSeekV4Flash | crate::model::ModelChoice::DeepSeekV4Pro + if layer.is_multiple_of(2) => + { + 4 + } + crate::model::ModelChoice::DeepSeekV4Flash | crate::model::ModelChoice::DeepSeekV4Pro => { + 128 + } + crate::model::ModelChoice::Glm52 => 0, } } @@ -3250,7 +3256,16 @@ fn check(result: i32, operation: &str) -> Result<(), String> { #[cfg(test)] mod tests { - use super::{raw_batch_span, raw_decode_span}; + use super::{compression_ratio, raw_batch_span, raw_decode_span}; + use crate::engine::{FLASH, PRO}; + + #[test] + fn compression_schedule_tracks_the_deepseek_model_shape() { + assert_eq!(compression_ratio(FLASH, 0), 0); + assert_eq!(compression_ratio(FLASH, 2), 4); + assert_eq!(compression_ratio(PRO, 0), 128); + assert_eq!(compression_ratio(PRO, 1), 128); + } #[test] fn decode_raw_span_tracks_the_logical_sliding_window() { diff --git a/src/engine/metal/checkpoint.rs b/src/engine/metal/checkpoint.rs index 7d04a3d..f8168e7 100644 --- a/src/engine/metal/checkpoint.rs +++ b/src/engine/metal/checkpoint.rs @@ -225,7 +225,7 @@ impl Executor { let mut indexer_rows = Vec::with_capacity(shape.layers as usize); for layer in 0..shape.layers { let rows = read_u32(file)?; - let ratio = compression_ratio(layer); + let ratio = compression_ratio(shape, layer); if rows != token_count.checked_div(ratio).unwrap_or(0) { return Err("KV checkpoint compressed row count is invalid".into()); } @@ -233,7 +233,7 @@ impl Executor { } for layer in 0..shape.layers { let rows = read_u32(file)?; - let expected = if compression_ratio(layer) == 4 { + let expected = if compression_ratio(shape, layer) == 4 { token_count / 4 } else { 0 diff --git a/src/server.rs b/src/server.rs index b935efd..00af802 100644 --- a/src/server.rs +++ b/src/server.rs @@ -544,7 +544,7 @@ fn compatible_completion( fn installed_endpoint_models(models_path: &std::path::Path) -> Vec { model::installed_models(models_path) .into_iter() - .filter(|model| *model == ModelChoice::DeepSeekV4Flash) + .filter(|model| *model != ModelChoice::Glm52) .collect() }