Support DeepSeek V4 Pro execution

src/engine/metal.rs

@@ -150,10 +150,10 @@ struct IndexerWeights {
 }
 
 impl Layer {
-    fn bind(model: &Gguf, index: u32) -> Result<Self, String> {
+    fn bind(model: &Gguf, shape: super::Shape, index: u32) -> Result<Self, String> {
         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<Self, String> {
         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::<Result<_, _>>()?;
         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<Self, String> {
         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() {

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