Integrate DS4 execution parity in Rust

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) },