Add GLM 5.2 Metal execution

src/engine/metal.rs

@@ -1,6 +1,8 @@
 mod checkpoint;
+mod glm;
 mod gpu;
 
+use glm::GlmExecutor;
 use gpu::*;
 
 use super::gguf::{F16, Gguf, Q8_0, Tensor as GgufTensor};
@@ -563,7 +565,7 @@ impl Session {
 // and `_context` must drop before `model` unmaps memory wrapped without copying
 // by native/metal/ds4_metal.m:10329. This intentionally differs from
 // ../ds4/ds4.c:56287-56288; do not reorder these fields to match it.
-pub(super) struct Executor {
+pub(super) struct DeepSeekExecutor {
     weights: Weights,
     session: Session,
     logits: Vec<f32>,
@@ -576,7 +578,7 @@ pub(super) struct Executor {
     model: Model,
 }
 
-impl Executor {
+impl DeepSeekExecutor {
     pub(super) fn open(
         model: Model,
         context: u32,
@@ -584,7 +586,7 @@ impl Executor {
         prefill_chunk: u32,
     ) -> Result<Self, String> {
         let weights = Weights::bind(&model)?;
-        let context_handle = Context::open(&model, quality)?;
+        let context_handle = Context::open(&model, quality, false, 0)?;
         let session = Session::new(
             &model,
             context,
@@ -849,6 +851,136 @@ impl Executor {
     }
 }
 
+/// Model-family dispatch over the two Rust-owned Metal graphs.
+pub(super) enum Executor {
+    DeepSeek(Box<DeepSeekExecutor>),
+    Glm(Box<GlmExecutor>),
+}
+
+impl Executor {
+    #[allow(dead_code)]
+    pub(super) fn open(
+        model: Model,
+        context: u32,
+        quality: bool,
+        prefill_chunk: u32,
+    ) -> Result<Self, String> {
+        Self::open_configured(
+            model,
+            context,
+            quality,
+            prefill_chunk,
+            crate::settings::EngineSsdSettings {
+                enabled: false,
+                cold: false,
+                cache_experts: 0,
+                cache_bytes: 0,
+                full_layers: 0,
+                full_layers_set: false,
+                preload_experts: 0,
+            },
+        )
+    }
+
+    pub(super) fn open_configured(
+        model: Model,
+        context: u32,
+        quality: bool,
+        prefill_chunk: u32,
+        ssd: crate::settings::EngineSsdSettings,
+    ) -> Result<Self, String> {
+        match model.shape.family {
+            ModelFamily::DeepSeek => DeepSeekExecutor::open(model, context, quality, prefill_chunk)
+                .map(Box::new)
+                .map(Self::DeepSeek),
+            ModelFamily::Glm => GlmExecutor::open(model, context, quality, ssd)
+                .map(Box::new)
+                .map(Self::Glm),
+        }
+    }
+
+    pub(super) fn eval(&mut self, token: i32) -> Result<(), String> {
+        match self {
+            Self::DeepSeek(executor) => executor.eval(token),
+            Self::Glm(executor) => executor.eval(token),
+        }
+    }
+
+    pub(super) fn prefill(
+        &mut self,
+        tokens: &[i32],
+        progress: impl FnMut(u32) -> bool,
+    ) -> Result<usize, String> {
+        match self {
+            Self::DeepSeek(executor) => executor.prefill(tokens, progress),
+            Self::Glm(executor) => executor.prefill(tokens, progress),
+        }
+    }
+
+    pub(super) fn logits(&self) -> &[f32] {
+        match self {
+            Self::DeepSeek(executor) => executor.logits(),
+            Self::Glm(executor) => executor.logits(),
+        }
+    }
+
+    pub(super) fn model(&self) -> &Model {
+        match self {
+            Self::DeepSeek(executor) => executor.model(),
+            Self::Glm(executor) => executor.model(),
+        }
+    }
+
+    pub(super) fn context(&self) -> u32 {
+        match self {
+            Self::DeepSeek(executor) => executor.context(),
+            Self::Glm(executor) => executor.context(),
+        }
+    }
+
+    pub(super) fn position(&self) -> u32 {
+        match self {
+            Self::DeepSeek(executor) => executor.position(),
+            Self::Glm(executor) => executor.position(),
+        }
+    }
+
+    pub(super) fn reset(&mut self) -> Result<(), String> {
+        match self {
+            Self::DeepSeek(executor) => executor.reset(),
+            Self::Glm(executor) => executor.reset(),
+        }
+    }
+
+    pub(super) fn align_prompt(&mut self, tokens: &[i32]) -> Result<usize, String> {
+        match self {
+            Self::DeepSeek(executor) => executor.align_prompt(tokens),
+            Self::Glm(executor) => executor.align_prompt(tokens),
+        }
+    }
+
+    pub(super) fn tokens(&self) -> &[i32] {
+        match self {
+            Self::DeepSeek(executor) => executor.tokens(),
+            Self::Glm(executor) => executor.tokens(),
+        }
+    }
+
+    pub(super) fn checkpoint_tag(&self) -> [u8; 32] {
+        match self {
+            Self::DeepSeek(executor) => executor.checkpoint_tag(),
+            Self::Glm(executor) => executor.checkpoint_tag(),
+        }
+    }
+
+    pub(super) fn note_checkpoint_tag(&mut self, tag: [u8; 32]) {
+        match self {
+            Self::DeepSeek(executor) => executor.note_checkpoint_tag(tag),
+            Self::Glm(executor) => executor.note_checkpoint_tag(tag),
+        }
+    }
+}
+
 #[allow(clippy::too_many_arguments)]
 fn refresh_ratio4_compressor_state(
     s: &BatchScratch,

src/engine/metal/checkpoint.rs

@@ -1,6 +1,6 @@
 use super::*;
 
-impl Executor {
+impl DeepSeekExecutor {
     pub(in crate::engine) fn save_checkpoint(
         &mut self,
         path: &Path,
@@ -305,12 +305,37 @@ impl Executor {
     }
 }
 
+impl Executor {
+    pub(in crate::engine) fn save_checkpoint(
+        &mut self,
+        path: &Path,
+        tag: [u8; 32],
+        progress: &mut impl FnMut(u64),
+    ) -> Result<(), String> {
+        match self {
+            Self::DeepSeek(executor) => executor.save_checkpoint(path, tag, progress),
+            Self::Glm(executor) => executor.save_checkpoint(path, tag, progress),
+        }
+    }
+
+    pub(in crate::engine) fn load_checkpoint(
+        &mut self,
+        path: &Path,
+        progress: &mut impl FnMut(u64),
+    ) -> Result<bool, String> {
+        match self {
+            Self::DeepSeek(executor) => executor.load_checkpoint(path, progress),
+            Self::Glm(executor) => executor.load_checkpoint(path, progress),
+        }
+    }
+}
+
 fn compressor_state_bytes(ratio: u32, head_dim: u64) -> u64 {
     let coefficient = if ratio == 4 { 2 } else { 1 };
     coefficient * head_dim * coefficient * u64::from(ratio) * 4
 }
 
-fn write_buffer(
+pub(super) fn write_buffer(
     file: &mut File,
     buffer: &Buffer,
     mut offset: u64,
@@ -330,7 +355,7 @@ fn write_buffer(
     Ok(())
 }
 
-fn read_buffer(
+pub(super) fn read_buffer(
     file: &mut File,
     buffer: &Buffer,
     mut offset: u64,
@@ -350,24 +375,24 @@ fn read_buffer(
     Ok(())
 }
 
-fn write_u32(file: &mut File, value: u32) -> Result<(), String> {
+pub(super) fn write_u32(file: &mut File, value: u32) -> Result<(), String> {
     file.write_all(&value.to_le_bytes())
         .map_err(|error| error.to_string())
 }
 
-fn write_u64(file: &mut File, value: u64) -> Result<(), String> {
+pub(super) fn write_u64(file: &mut File, value: u64) -> Result<(), String> {
     file.write_all(&value.to_le_bytes())
         .map_err(|error| error.to_string())
 }
 
-fn read_u32(file: &mut File) -> Result<u32, String> {
+pub(super) fn read_u32(file: &mut File) -> Result<u32, String> {
     let mut bytes = [0; 4];
     file.read_exact(&mut bytes)
         .map_err(|error| error.to_string())?;
     Ok(u32::from_le_bytes(bytes))
 }
 
-fn read_u64(file: &mut File) -> Result<u64, String> {
+pub(super) fn read_u64(file: &mut File) -> Result<u64, String> {
     let mut bytes = [0; 8];
     file.read_exact(&mut bytes)
         .map_err(|error| error.to_string())?;

src/engine/metal/gpu.rs

@@ -5,6 +5,19 @@ pub(super) struct GpuTensor {
     _private: [u8; 0],
 }
 
+#[repr(C)]
+pub(super) struct StreamExpertTable {
+    pub(super) model_map: *const c_void,
+    pub(super) model_size: u64,
+    pub(super) layer: u32,
+    pub(super) total_experts: u32,
+    pub(super) gate_offset: u64,
+    pub(super) up_offset: u64,
+    pub(super) down_offset: u64,
+    pub(super) gate_expert_bytes: u64,
+    pub(super) down_expert_bytes: u64,
+}
+
 unsafe extern "C" {
     pub(super) fn ds4_gpu_init() -> i32;
     pub(super) fn ds4_gpu_cleanup();
@@ -16,6 +29,22 @@ unsafe extern "C" {
         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);
+    pub(super) fn ds4_gpu_set_model_fd(fd: i32) -> i32;
+    pub(super) fn ds4_gpu_set_streaming_expert_cache_budget(experts: u32);
+    pub(super) fn ds4_gpu_set_streaming_expert_cache_expert_bytes(bytes: u64);
+    pub(super) fn ds4_gpu_recommended_working_set_size() -> u64;
+    pub(super) fn ds4_gpu_stream_expert_cache_budget_for_expert_size(
+        gate_expert_bytes: u64,
+        down_expert_bytes: u64,
+    ) -> u32;
+    pub(super) fn ds4_gpu_glm_stream_expert_cache_begin_selected_load_tensor(
+        table: *const StreamExpertTable,
+        selected: *const GpuTensor,
+        count: u32,
+    ) -> i32;
+    pub(super) fn ds4_gpu_flush_commands() -> i32;
     pub(super) fn ds4_gpu_tensor_alloc(bytes: u64) -> *mut GpuTensor;
     pub(super) fn ds4_gpu_tensor_view(
         base: *const GpuTensor,
@@ -122,6 +151,214 @@ unsafe extern "C" {
         x: *const GpuTensor,
         rows: u64,
     ) -> i32;
+    pub(super) fn ds4_gpu_matmul_quant_tensor(
+        out: *mut GpuTensor,
+        map: *const c_void,
+        size: u64,
+        weight: u64,
+        kind: u32,
+        input: u64,
+        output: u64,
+        x: *const GpuTensor,
+        rows: u64,
+    ) -> i32;
+    pub(super) fn ds4_gpu_matmul_quant_decode_mpp_model_view_tensor(
+        out: *mut GpuTensor,
+        map: *const c_void,
+        size: u64,
+        weight: u64,
+        kind: u32,
+        input: u64,
+        output: u64,
+        x: *const GpuTensor,
+        rows: u64,
+    ) -> i32;
+    pub(super) fn ds4_gpu_matmul_f32_tensor(
+        out: *mut GpuTensor,
+        map: *const c_void,
+        size: u64,
+        weight: u64,
+        input: u64,
+        output: u64,
+        x: *const GpuTensor,
+        rows: u64,
+    ) -> i32;
+    pub(super) fn ds4_gpu_embed_token_quant_tensor(
+        out: *mut GpuTensor,
+        map: *const c_void,
+        size: u64,
+        weight: u64,
+        kind: u32,
+        vocab: u32,
+        token: u32,
+        embd: u32,
+    ) -> i32;
+    pub(super) fn ds4_gpu_glm_rope_tail_tensor(
+        x: *mut GpuTensor,
+        tokens: u32,
+        heads: u32,
+        head_dim: u32,
+        rot: u32,
+        pos: u32,
+        original: u32,
+        freq_base: f32,
+        freq_scale: f32,
+        ext: f32,
+        attn_factor: f32,
+        beta_fast: f32,
+        beta_slow: f32,
+    ) -> i32;
+    pub(super) fn ds4_gpu_glm_qkv_norm_store_compact_kv_tensor(
+        q_out: *mut GpuTensor,
+        q: *const GpuTensor,
+        map: *const c_void,
+        size: u64,
+        q_weight: u64,
+        q_n: u32,
+        kv_cache: *mut GpuTensor,
+        rope_cache: *mut GpuTensor,
+        kv_raw: *const GpuTensor,
+        kv_weight: u64,
+        pos: u32,
+        tokens: u32,
+        cache_cap: u32,
+        kv_raw_dim: u32,
+        kv_lora: u32,
+        rot: u32,
+        cache_f16: bool,
+        eps: f32,
+    ) -> i32;
+    pub(super) fn ds4_gpu_glm_store_indexer_k_tensor(
+        cache: *mut GpuTensor,
+        raw: *const GpuTensor,
+        map: *const c_void,
+        size: u64,
+        weight: u64,
+        bias: u64,
+        pos: u32,
+        tokens: u32,
+        cache_cap: u32,
+        head_dim: u32,
+        rot: u32,
+        original: u32,
+        eps: f32,
+        freq_base: f32,
+        freq_scale: f32,
+        ext: f32,
+        attn_factor: f32,
+        beta_fast: f32,
+        beta_slow: f32,
+        cache_f16: bool,
+    ) -> i32;
+    pub(super) fn ds4_gpu_glm_fill_selected_range_tensor(
+        selected: *mut GpuTensor,
+        count: u32,
+    ) -> i32;
+    pub(super) fn ds4_gpu_glm_indexer_rope_tail_tensor(
+        x: *mut GpuTensor,
+        tokens: u32,
+        heads: u32,
+        head_dim: u32,
+        rot: u32,
+        pos: u32,
+        original: u32,
+        freq_base: f32,
+        freq_scale: f32,
+        ext: f32,
+        attn_factor: f32,
+        beta_fast: f32,
+        beta_slow: f32,
+    ) -> i32;
+    pub(super) fn ds4_gpu_glm_indexer_score_one_tensor(
+        scores: *mut GpuTensor,
+        q: *const GpuTensor,
+        weights: *const GpuTensor,
+        cache: *const GpuTensor,
+        rows: u32,
+        heads: u32,
+        head_dim: u32,
+        scale: f32,
+        cache_f16: bool,
+    ) -> i32;
+    pub(super) fn ds4_gpu_glm_qk_lowrank_typed_tensor(
+        out: *mut GpuTensor,
+        q: *const GpuTensor,
+        map: *const c_void,
+        size: u64,
+        weight: u64,
+        kind: u32,
+        heads: u32,
+        kv_lora: u32,
+        q_nope: u32,
+        q_dim: u32,
+    ) -> i32;
+    pub(super) fn ds4_gpu_glm_attention_indexed_decode_typed_tensor(
+        heads_out: *mut GpuTensor,
+        q: *const GpuTensor,
+        qk_low: *const GpuTensor,
+        kv_cache: *const GpuTensor,
+        rope_cache: *const GpuTensor,
+        map: *const c_void,
+        size: u64,
+        value_weight: u64,
+        value_kind: u32,
+        selected: *const GpuTensor,
+        selected_count: u32,
+        cache_cap: u32,
+        cache_f16: bool,
+        heads: u32,
+        kv_lora: u32,
+        q_nope: u32,
+        rot: u32,
+        value_dim: u32,
+        original: u32,
+        freq_base: f32,
+        freq_scale: f32,
+        ext: f32,
+        attn_factor: f32,
+        beta_fast: f32,
+        beta_slow: f32,
+    ) -> i32;
+    pub(super) fn ds4_gpu_glm_router_select_tensor(
+        selected: *mut GpuTensor,
+        weights: *mut GpuTensor,
+        probs: *mut GpuTensor,
+        map: *const c_void,
+        size: u64,
+        bias: u64,
+        logits: *const GpuTensor,
+        experts: u32,
+        used: u32,
+        scale: f32,
+    ) -> i32;
+    pub(super) fn ds4_gpu_glm_routed_moe_one_tensor(
+        out: *mut GpuTensor,
+        mid: *mut GpuTensor,
+        map: *const c_void,
+        size: u64,
+        gate: u64,
+        up: u64,
+        down: u64,
+        gate_kind: u32,
+        up_kind: u32,
+        down_kind: u32,
+        gate_expert_bytes: u64,
+        gate_row_bytes: u64,
+        up_expert_bytes: u64,
+        up_row_bytes: u64,
+        down_expert_bytes: u64,
+        down_row_bytes: u64,
+        input: u32,
+        hidden: u32,
+        output: u32,
+        selected: *const GpuTensor,
+        weights: *const GpuTensor,
+        total_experts: u32,
+        used: u32,
+        layer: u32,
+        x: *const GpuTensor,
+        force_resident: bool,
+    ) -> i32;
     pub(super) fn ds4_gpu_matmul_q8_0_pair_tensor(
         out_a: *mut GpuTensor,
         out_b: *mut GpuTensor,
@@ -615,6 +852,30 @@ unsafe extern "C" {
         clamp: f32,
         scale: f32,
     ) -> i32;
+    pub(super) fn ds4_gpu_add_tensor(
+        out: *mut GpuTensor,
+        a: *const GpuTensor,
+        b: *const GpuTensor,
+        count: u32,
+    ) -> i32;
+    pub(super) fn ds4_gpu_add3_tensor(
+        out: *mut GpuTensor,
+        a: *const GpuTensor,
+        b: *const GpuTensor,
+        c: *const GpuTensor,
+        count: u32,
+    ) -> i32;
+    pub(super) fn ds4_gpu_add_rms_norm_weight_tensor(
+        norm: *mut GpuTensor,
+        sum: *mut GpuTensor,
+        a: *const GpuTensor,
+        b: *const GpuTensor,
+        map: *const c_void,
+        size: u64,
+        weight: u64,
+        count: u32,
+        eps: f32,
+    ) -> i32;
     pub(super) fn ds4_gpu_hc_expand_split_half_tensor(
         out: *mut GpuTensor,
         block_half: *const GpuTensor,
@@ -772,11 +1033,31 @@ unsafe extern "C" {
     ) -> i32;
 }
 
-pub(super) struct Context;
+pub(super) struct Context {
+    _model_file: File,
+}
 
 impl Context {
-    pub(super) fn open(model: &Model, quality: bool) -> Result<Self, String> {
+    pub(super) fn open(
+        model: &Model,
+        quality: bool,
+        ssd_streaming: bool,
+        admission_bytes: u64,
+    ) -> Result<Self, String> {
         check(unsafe { ds4_gpu_init() }, "Metal initialization")?;
+        unsafe {
+            ds4_gpu_set_glm_model(model.shape.family == ModelFamily::Glm);
+            ds4_gpu_set_ssd_streaming(ssd_streaming);
+        }
+        let recommended = unsafe { ds4_gpu_recommended_working_set_size() };
+        if admission_bytes != 0 && recommended != 0 && admission_bytes > recommended {
+            unsafe { ds4_gpu_cleanup() };
+            return Err(format!(
+                "model load needs {:.1} GiB including context and scratch, but Metal recommends at most {:.1} GiB",
+                admission_bytes as f64 / 1_073_741_824.0,
+                recommended as f64 / 1_073_741_824.0,
+            ));
+        }
         let data_offset = model.main.data_offset();
         if let Err(error) = check(
             unsafe {
@@ -794,7 +1075,24 @@ impl Context {
             return Err(error);
         }
         unsafe { ds4_gpu_set_quality(quality) };
-        Ok(Self)
+        let model_file = File::open(model.main.path()).map_err(|error| {
+            unsafe { ds4_gpu_cleanup() };
+            error.to_string()
+        })?;
+        #[cfg(target_os = "macos")]
+        {
+            use std::os::fd::AsRawFd;
+            if let Err(error) = check(
+                unsafe { ds4_gpu_set_model_fd(model_file.as_raw_fd()) },
+                "model file registration",
+            ) {
+                unsafe { ds4_gpu_cleanup() };
+                return Err(error);
+            }
+        }
+        Ok(Self {
+            _model_file: model_file,
+        })
     }
 }