Integrate DS4 execution parity in Rust

src/engine/gguf.rs

@@ -1,4 +1,4 @@
-use memmap2::{Mmap, MmapOptions};
+use memmap2::{Advice, Mmap, MmapOptions};
 use sha2::{Digest, Sha256};
 use std::collections::HashMap;
 use std::fs::File;
@@ -215,6 +215,23 @@ impl Gguf {
         self.max_tensor_bytes
     }
 
+    pub(super) fn warm(&self) -> Result<(), String> {
+        let start = self.data_offset as usize;
+        if start >= self.map.len() {
+            return Ok(());
+        }
+        self.map
+            .advise_range(Advice::WillNeed, start, self.map.len() - start)
+            .map_err(|error| format!("Cannot warm {}: {error}", self.path.display()))?;
+        let mut checksum = 0_u64;
+        for offset in (start..self.map.len()).step_by(16 * 1024) {
+            checksum = checksum.wrapping_add(u64::from(self.map[offset]));
+        }
+        checksum = checksum.wrapping_add(u64::from(self.map[self.map.len() - 1]));
+        std::hint::black_box(checksum);
+        Ok(())
+    }
+
     pub(super) fn tensor(&self, name: &str) -> Result<&Tensor, String> {
         self.tensors
             .get(name)
@@ -544,6 +561,7 @@ mod tests {
         assert_eq!(model.bytes("general.architecture").unwrap(), b"deepseek4");
         assert_eq!(model.tensor("weight").unwrap().dims, [1]);
         assert_eq!(model.tensor_data("weight").unwrap(), 1_f32.to_le_bytes());
+        model.warm().unwrap();
         fs::remove_file(path).unwrap();
     }
 

src/engine/metal/checkpoint.rs

@@ -301,6 +301,15 @@ impl DeepSeekExecutor {
         self.tokens = tokens;
         self.logits = logits;
         self.checkpoint_tag = checkpoint_tag;
+        if let Some(mtp) = &mut self.legacy_mtp {
+            mtp.draft_token = None;
+            mtp.raw_rows = 0;
+        }
+        if let Some(dspark) = &mut self.dspark {
+            dspark.capture_mask = 0;
+            dspark.cache_start = 0;
+            dspark.cache_len = 0;
+        }
         Ok(())
     }
 }

src/engine/metal/glm.rs

@@ -231,7 +231,7 @@ impl GlmExecutor {
         }
         let weights = GlmWeights::bind(&model)?;
         let admission = admission_bytes(&model, &weights, context, ssd)?;
-        let context_handle = Context::open(&model, quality, ssd.enabled, admission)?;
+        let context_handle = Context::open(&model, quality, ssd.enabled, admission, None)?;
         configure_streaming(&model, &weights, ssd)?;
         let scratch = GlmScratch::allocate(&model, context)?;
         let caches = (0..weights.layers.len())

src/engine/metal/gpu.rs

@@ -28,6 +28,14 @@ unsafe extern "C" {
         map_size: u64,
         max_tensor_bytes: u64,
     ) -> i32;
+    pub(super) fn ds4_gpu_set_model_map_spans(
+        model_map: *const c_void,
+        model_size: u64,
+        offsets: *const u64,
+        sizes: *const u64,
+        count: u32,
+        max_tensor_bytes: u64,
+    ) -> i32;
     pub(super) fn ds4_gpu_set_quality(quality: bool);
     pub(super) fn ds4_gpu_set_glm_model(enabled: bool);
     pub(super) fn ds4_gpu_set_ssd_streaming(enabled: bool);
@@ -39,6 +47,19 @@ unsafe extern "C" {
         gate_expert_bytes: u64,
         down_expert_bytes: u64,
     ) -> u32;
+    pub(super) fn ds4_gpu_stream_expert_cache_seed_experts(
+        table: *const StreamExpertTable,
+        expert_ids: *const i32,
+        expert_priorities: *const u32,
+        experts: u32,
+    ) -> i32;
+    pub(super) fn ds4_gpu_stream_expert_cache_begin_selected_load(
+        table: *const StreamExpertTable,
+        selected_ids: *const i32,
+        selected: u32,
+    ) -> i32;
+    pub(super) fn ds4_gpu_stream_expert_cache_note_service_thread();
+    pub(super) fn ds4_gpu_stream_expert_cache_reset_route_hotness();
     pub(super) fn ds4_gpu_glm_stream_expert_cache_begin_selected_load_tensor(
         table: *const StreamExpertTable,
         selected: *const GpuTensor,
@@ -79,8 +100,23 @@ unsafe extern "C" {
         src_offset: u64,
         count: u64,
     ) -> i32;
+    pub(super) fn ds4_gpu_pack_slot_rows_f32_tensor(
+        out: *mut GpuTensor,
+        slots: *const GpuTensor,
+        rows: u32,
+        width: u32,
+        slot_count: u32,
+        slot_stride: u32,
+    ) -> i32;
     pub(super) fn ds4_gpu_begin_commands() -> i32;
     pub(super) fn ds4_gpu_end_commands() -> i32;
+    pub(super) fn ds4_gpu_signal_selected_readback_ready(event: *mut u64) -> i32;
+    pub(super) fn ds4_gpu_wait_selected_readback_ready(
+        event: u64,
+        label: *const std::ffi::c_char,
+    ) -> i32;
+    pub(super) fn ds4_gpu_routed_moe_set_selected_override(selected: *const i32, count: u32)
+    -> i32;
 
     pub(super) fn ds4_gpu_embed_tokens_hc_tensor(
         out: *mut GpuTensor,
@@ -110,6 +146,13 @@ unsafe extern "C" {
         n: u32,
         eps: f32,
     ) -> i32;
+    pub(super) fn ds4_gpu_rms_norm_plain_rows_tensor(
+        out: *mut GpuTensor,
+        x: *const GpuTensor,
+        n: u32,
+        rows: u32,
+        eps: f32,
+    ) -> i32;
     pub(super) fn ds4_gpu_rms_norm_weight_tensor(
         out: *mut GpuTensor,
         x: *const GpuTensor,
@@ -119,6 +162,36 @@ unsafe extern "C" {
         n: u32,
         eps: f32,
     ) -> i32;
+    pub(super) fn ds4_gpu_rms_norm_weight_rows_tensor(
+        out: *mut GpuTensor,
+        x: *const GpuTensor,
+        map: *const c_void,
+        size: u64,
+        weight: u64,
+        n: u32,
+        rows: u32,
+        eps: f32,
+    ) -> i32;
+    pub(super) fn ds4_gpu_repeat_hc_tensor(
+        out: *mut GpuTensor,
+        x: *const GpuTensor,
+        embd: u32,
+        hc: u32,
+    ) -> i32;
+    pub(super) fn ds4_gpu_attention_noncausal_raw_batch_heads_tensor(
+        out: *mut GpuTensor,
+        map: *const c_void,
+        size: u64,
+        sinks: u64,
+        q: *const GpuTensor,
+        raw_cache: *const GpuTensor,
+        rows: u32,
+        visible_rows: u32,
+        cache_cap: u32,
+        raw_start: u32,
+        heads: u32,
+        head_dim: u32,
+    ) -> i32;
     pub(super) fn ds4_gpu_hc_rms_scale_project_f16_tensor(
         out: *mut GpuTensor,
         scale: *mut GpuTensor,
@@ -865,6 +938,14 @@ unsafe extern "C" {
         c: *const GpuTensor,
         count: u32,
     ) -> i32;
+    pub(super) fn ds4_gpu_directional_steering_project_tensor(
+        x: *mut GpuTensor,
+        directions: *const GpuTensor,
+        layer: u32,
+        width: u32,
+        rows: u32,
+        scale: f32,
+    ) -> i32;
     pub(super) fn ds4_gpu_add_rms_norm_weight_tensor(
         norm: *mut GpuTensor,
         sum: *mut GpuTensor,
@@ -1043,6 +1124,7 @@ impl Context {
         quality: bool,
         ssd_streaming: bool,
         admission_bytes: u64,
+        model_spans: Option<&[(u64, u64)]>,
     ) -> Result<Self, String> {
         check(unsafe { ds4_gpu_init() }, "Metal initialization")?;
         unsafe {
@@ -1059,7 +1141,19 @@ impl Context {
             ));
         }
         let data_offset = model.main.data_offset();
-        if let Err(error) = check(
+        let mapped = if let Some(spans) = model_spans {
+            let (offsets, sizes): (Vec<_>, Vec<_>) = spans.iter().copied().unzip();
+            unsafe {
+                ds4_gpu_set_model_map_spans(
+                    model.main.map_ptr().cast(),
+                    model.main.len(),
+                    offsets.as_ptr(),
+                    sizes.as_ptr(),
+                    spans.len() as u32,
+                    model.main.max_tensor_bytes(),
+                )
+            }
+        } else {
             unsafe {
                 ds4_gpu_set_model_map_range(
                     model.main.map_ptr().cast(),
@@ -1068,12 +1162,27 @@ impl Context {
                     model.main.len() - data_offset,
                     model.main.max_tensor_bytes(),
                 )
-            },
-            "model mapping",
-        ) {
+            }
+        };
+        if let Err(error) = check(mapped, "model mapping") {
             unsafe { ds4_gpu_cleanup() };
             return Err(error);
         }
+        if let Some(support) = &model.support {
+            let mapped = unsafe {
+                ds4_gpu_set_model_map_range(
+                    support.map_ptr().cast(),
+                    support.len(),
+                    support.data_offset(),
+                    support.len() - support.data_offset(),
+                    support.max_tensor_bytes(),
+                )
+            };
+            if let Err(error) = check(mapped, "support-model mapping") {
+                unsafe { ds4_gpu_cleanup() };
+                return Err(error);
+            }
+        }
         unsafe { ds4_gpu_set_quality(quality) };
         let model_file = File::open(model.main.path()).map_err(|error| {
             unsafe { ds4_gpu_cleanup() };
@@ -1209,6 +1318,20 @@ impl Buffer {
         )
     }
 
+    pub(super) fn write_f32(&self, values: &[f32]) -> Result<(), String> {
+        check(
+            unsafe {
+                ds4_gpu_tensor_write(
+                    self.raw(),
+                    0,
+                    values.as_ptr().cast(),
+                    std::mem::size_of_val(values) as u64,
+                )
+            },
+            "uploading floats",
+        )
+    }
+
     pub(super) fn fill(&self, value: f32, count: u64) -> Result<(), String> {
         call(
             unsafe { ds4_gpu_tensor_fill_f32(self.raw(), value, count) },

src/engine/validation.rs

@@ -7,7 +7,7 @@ pub(crate) fn validate_model_artifact(
 ) -> Result<(), String> {
     if support {
         let model = Gguf::open(path)?;
-        validate_dspark(&model, &FLASH)
+        validate_support(&model, &FLASH).map(|_| ())
     } else {
         let model = Model::open_main(path, expected)?;
         let summary = model.summary();
@@ -29,6 +29,192 @@ pub(crate) fn validate_model_artifact(
     }
 }
 
+#[derive(Clone, Copy, Debug, Eq, PartialEq)]
+pub(super) enum SupportKind {
+    LegacyMtp,
+    DSpark,
+}
+
+#[derive(Clone, Debug, Eq, PartialEq)]
+pub(super) struct DsparkConfig {
+    pub(super) block_size: u32,
+    pub(super) markov_rank: u32,
+    pub(super) noise_token: u32,
+    pub(super) target_layers: Vec<u32>,
+    pub(super) stages: u32,
+}
+
+pub(super) fn dspark_config(model: &Gguf) -> Result<DsparkConfig, String> {
+    let block_size = first_u32(
+        model,
+        &[
+            "deepseek4.dspark.block_size",
+            "deepseek4.dspark_block_size",
+            "dspark.block_size",
+        ],
+    )?;
+    let markov_rank = first_u32(
+        model,
+        &[
+            "deepseek4.dspark.markov_rank",
+            "deepseek4.dspark_markov_rank",
+            "dspark.markov_rank",
+        ],
+    )?;
+    let noise_token = first_u32(
+        model,
+        &[
+            "deepseek4.dspark.noise_token_id",
+            "deepseek4.dspark_noise_token_id",
+            "dspark.noise_token_id",
+        ],
+    )?;
+    let target_layers = first_u32s(
+        model,
+        &[
+            "deepseek4.dspark.target_layer_ids",
+            "deepseek4.dspark_target_layer_ids",
+            "dspark.target_layer_ids",
+        ],
+    )?;
+    let stages = model
+        .tensors
+        .keys()
+        .filter_map(|name| {
+            name.strip_prefix("mtp.")?
+                .split('.')
+                .next()?
+                .parse::<u32>()
+                .ok()
+        })
+        .max()
+        .map_or(0, |stage| stage + 1);
+    Ok(DsparkConfig {
+        block_size,
+        markov_rank,
+        noise_token,
+        target_layers: target_layers.to_vec(),
+        stages,
+    })
+}
+
+pub(super) fn validate_support(model: &Gguf, shape: &Shape) -> Result<SupportKind, String> {
+    if model.tensors.contains_key("mtp.0.e_proj.weight")
+        && model.tensors.contains_key("mtp.0.h_proj.weight")
+        && model.tensors.contains_key("mtp.0.hc_head_base.weight")
+    {
+        validate_legacy_mtp(model, shape)?;
+        Ok(SupportKind::LegacyMtp)
+    } else if model.metadata.contains_key("deepseek4.dspark.block_size")
+        || model.metadata.contains_key("deepseek4.dspark_block_size")
+        || model.metadata.contains_key("dspark.block_size")
+    {
+        validate_dspark(model, shape)?;
+        Ok(SupportKind::DSpark)
+    } else {
+        Err("support GGUF is neither legacy MTP nor DSpark".into())
+    }
+}
+
+fn validate_legacy_mtp(model: &Gguf, shape: &Shape) -> Result<(), String> {
+    if shape.model != ModelChoice::DeepSeekV4Flash {
+        return Err("legacy MTP support is available only for DeepSeek V4 Flash".into());
+    }
+    let prefix = "mtp.0";
+    let hc_dim = shape.embd * shape.hc;
+    let hc_mix = 2 * shape.hc + shape.hc * shape.hc;
+    let q_dim = shape.heads * shape.head_dim;
+    let output_low = shape.out_groups * shape.lora_o;
+    for (suffix, types, dims) in [
+        ("hc_head_base.weight", &[F32][..], vec![shape.hc]),
+        ("hc_head_fn.weight", PLAIN, vec![hc_dim, shape.hc]),
+        ("hc_head_scale.weight", &[F32][..], vec![1]),
+        ("e_proj.weight", &[Q8_0][..], vec![shape.embd, shape.embd]),
+        ("h_proj.weight", &[Q8_0][..], vec![shape.embd, shape.embd]),
+        ("enorm.weight", &[F32][..], vec![shape.embd]),
+        ("hnorm.weight", &[F32][..], vec![shape.embd]),
+        ("norm.weight", &[F32][..], vec![shape.embd]),
+        ("hc_attn_fn.weight", PLAIN, vec![hc_dim, hc_mix]),
+        ("hc_attn_scale.weight", &[F32][..], vec![3]),
+        ("hc_attn_base.weight", &[F32][..], vec![hc_mix]),
+        ("attn_norm.weight", &[F32][..], vec![shape.embd]),
+        (
+            "attn_q_a.weight",
+            &[Q8_0][..],
+            vec![shape.embd, shape.lora_q],
+        ),
+        ("attn_q_a_norm.weight", &[F32][..], vec![shape.lora_q]),
+        ("attn_q_b.weight", &[Q8_0][..], vec![shape.lora_q, q_dim]),
+        (
+            "attn_kv.weight",
+            &[Q8_0][..],
+            vec![shape.embd, shape.head_dim],
+        ),
+        ("attn_kv_a_norm.weight", &[F32][..], vec![shape.head_dim]),
+        ("attn_sinks.weight", &[F32][..], vec![shape.heads]),
+        (
+            "attn_output_a.weight",
+            &[Q8_0][..],
+            vec![
+                shape.head_dim * (shape.heads / shape.out_groups),
+                output_low,
+            ],
+        ),
+        (
+            "attn_output_b.weight",
+            &[Q8_0][..],
+            vec![output_low, shape.embd],
+        ),
+        ("hc_ffn_fn.weight", PLAIN, vec![hc_dim, hc_mix]),
+        ("hc_ffn_scale.weight", &[F32][..], vec![3]),
+        ("hc_ffn_base.weight", &[F32][..], vec![hc_mix]),
+        ("ffn_norm.weight", &[F32][..], vec![shape.embd]),
+        (
+            "ffn_gate_inp.weight",
+            PLAIN,
+            vec![shape.embd, shape.experts],
+        ),
+        ("exp_probs_b.bias", &[F32][..], vec![shape.experts]),
+        (
+            "ffn_gate_exps.weight",
+            ROUTED,
+            vec![shape.embd, shape.ff_expert, shape.experts],
+        ),
+        (
+            "ffn_up_exps.weight",
+            ROUTED,
+            vec![shape.embd, shape.ff_expert, shape.experts],
+        ),
+        (
+            "ffn_down_exps.weight",
+            ROUTED,
+            vec![shape.ff_expert, shape.embd, shape.experts],
+        ),
+        (
+            "ffn_gate_shexp.weight",
+            &[Q8_0][..],
+            vec![shape.embd, shape.ff_expert],
+        ),
+        (
+            "ffn_up_shexp.weight",
+            &[Q8_0][..],
+            vec![shape.embd, shape.ff_expert],
+        ),
+        (
+            "ffn_down_shexp.weight",
+            &[Q8_0][..],
+            vec![shape.ff_expert, shape.embd],
+        ),
+    ] {
+        expect(model, &format!("{prefix}.{suffix}"), types, &dims)?;
+    }
+    same_type(
+        model,
+        "mtp.0.ffn_gate_exps.weight",
+        "mtp.0.ffn_up_exps.weight",
+    )
+}
+
 pub(super) fn validate_main(model: &Gguf, expected: ModelChoice) -> Result<Shape, String> {
     let family = if model.bytes("general.architecture").ok() == Some(b"glm-dsa") {
         ModelFamily::Glm
@@ -575,38 +761,13 @@ pub(super) fn validate_dspark(model: &Gguf, shape: &Shape) -> Result<(), String>
     if shape.model != ModelChoice::DeepSeekV4Flash {
         return Err("DSpark support is available only for DeepSeek V4 Flash".into());
     }
-    let block_size = first_u32(
-        model,
-        &[
-            "deepseek4.dspark.block_size",
-            "deepseek4.dspark_block_size",
-            "dspark.block_size",
-        ],
-    )?;
-    let markov_rank = first_u32(
-        model,
-        &[
-            "deepseek4.dspark.markov_rank",
-            "deepseek4.dspark_markov_rank",
-            "dspark.markov_rank",
-        ],
-    )?;
-    let noise_token = first_u32(
-        model,
-        &[
-            "deepseek4.dspark.noise_token_id",
-            "deepseek4.dspark_noise_token_id",
-            "dspark.noise_token_id",
-        ],
-    )?;
-    let targets = first_u32s(
-        model,
-        &[
-            "deepseek4.dspark.target_layer_ids",
-            "deepseek4.dspark_target_layer_ids",
-            "dspark.target_layer_ids",
-        ],
-    )?;
+    let DsparkConfig {
+        block_size,
+        markov_rank,
+        noise_token,
+        target_layers: targets,
+        stages,
+    } = dspark_config(model)?;
     if !(1..=16).contains(&block_size) || markov_rank == 0 || noise_token >= shape.vocab as u32 {
         return Err("invalid DSpark block, Markov, or noise-token metadata".into());
     }
@@ -617,18 +778,6 @@ pub(super) fn validate_dspark(model: &Gguf, shape: &Shape) -> Result<(), String>
     {
         return Err("invalid DSpark target-layer metadata".into());
     }
-    let stages = model
-        .tensors
-        .keys()
-        .filter_map(|name| {
-            name.strip_prefix("mtp.")?
-                .split('.')
-                .next()?
-                .parse::<u32>()
-                .ok()
-        })
-        .max()
-        .map_or(0, |stage| stage + 1);
     if !(1..=8).contains(&stages) {
         return Err(format!("invalid DSpark stage count: {stages}"));
     }
@@ -1079,4 +1228,15 @@ mod tests {
             validate_model_artifact(path, ModelChoice::DeepSeekV4Flash, true).unwrap();
         }
     }
+
+    #[test]
+    fn installed_legacy_mtp_fixture_passes_the_target_layout() {
+        let path = Path::new("../ds4/gguf/DeepSeek-V4-Flash-MTP-Q4K-Q8_0-F32.gguf");
+        if path.exists() {
+            assert_eq!(
+                validate_support(&Gguf::open(path).unwrap(), &FLASH).unwrap(),
+                SupportKind::LegacyMtp
+            );
+        }
+    }
 }