From 65c9cbfc45d112806fdb8d08b34f6e89e7e188f1 Mon Sep 17 00:00:00 2001 From: Georg Bauer Date: Sun, 26 Jul 2026 12:16:30 +0200 Subject: [PATCH] Add GLM 5.2 Metal execution --- src/engine.rs | 14 +- src/engine/metal.rs | 138 +++- src/engine/metal/checkpoint.rs | 39 +- src/engine/metal/glm.rs | 1341 ++++++++++++++++++++++++++++++++ src/engine/metal/gpu.rs | 304 +++++++- src/server.rs | 3 - 6 files changed, 1819 insertions(+), 20 deletions(-) create mode 100644 src/engine/metal/glm.rs diff --git a/src/engine.rs b/src/engine.rs index e6a6a18..8e06d95 100644 --- a/src/engine.rs +++ b/src/engine.rs @@ -372,21 +372,27 @@ pub(crate) struct CompactionOutput { #[cfg(target_os = "macos")] impl Generator { pub(crate) fn open(settings: &EngineSettings, metrics: Arc) -> Result { - if settings.model == ModelChoice::Glm52 { - return Err("GLM 5.2 generation is not initialized by the DeepSeek executor".into()); + if settings.speculative.glm_mtp { + return Err( + "GLM MTP requires the shared speculative verifier, which is not enabled".into(), + ); } - if settings.speculative.dspark || settings.ssd.enabled || settings.steering.file.is_some() { + if settings.speculative.dspark + || (settings.ssd.enabled && settings.model != ModelChoice::Glm52) + || settings.steering.file.is_some() + { return Err( "DSpark, SSD streaming, and steering are not yet available in the Rust executor" .into(), ); } let model = Model::open(settings)?; - let executor = metal::Executor::open( + let executor = metal::Executor::open_configured( model, settings.context_tokens.max(1) as u32, settings.execution.quality, settings.execution.prefill_chunk, + settings.ssd, )?; Ok(Self { executor, diff --git a/src/engine/metal.rs b/src/engine/metal.rs index 941dc6d..64b1fe2 100644 --- a/src/engine/metal.rs +++ b/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, @@ -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 { 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), + Glm(Box), +} + +impl Executor { + #[allow(dead_code)] + pub(super) fn open( + model: Model, + context: u32, + quality: bool, + prefill_chunk: u32, + ) -> Result { + 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 { + 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 { + 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 { + 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, diff --git a/src/engine/metal/checkpoint.rs b/src/engine/metal/checkpoint.rs index f8168e7..84c4afb 100644 --- a/src/engine/metal/checkpoint.rs +++ b/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 { + 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 { +pub(super) fn read_u32(file: &mut File) -> Result { 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 { +pub(super) fn read_u64(file: &mut File) -> Result { let mut bytes = [0; 8]; file.read_exact(&mut bytes) .map_err(|error| error.to_string())?; diff --git a/src/engine/metal/glm.rs b/src/engine/metal/glm.rs new file mode 100644 index 0000000..921bf32 --- /dev/null +++ b/src/engine/metal/glm.rs @@ -0,0 +1,1341 @@ +use super::checkpoint::{read_buffer, read_u32, read_u64, write_buffer, write_u32, write_u64}; +use super::*; +use crate::settings::EngineSsdSettings; + +const CHECKPOINT_MAGIC: &[u8; 8] = b"DS4GLM01"; +const CHECKPOINT_VERSION: u32 = 1; +const CACHE_F16: bool = true; +const DECODE_FLUSH_LAYERS: usize = 4; + +#[derive(Clone, Copy)] +struct SparseWeights { + router: Weight, + bias: Weight, + gate: Weight, + up: Weight, + down: Weight, + shared_gate: Weight, + shared_up: Weight, + shared_down: Weight, +} + +#[derive(Clone, Copy)] +struct DenseWeights { + gate: Weight, + up: Weight, + down: Weight, +} + +struct GlmLayer { + attn_norm: Weight, + q_a: Weight, + q_a_norm: Weight, + q_b: Weight, + kv_a: Weight, + kv_norm: Weight, + k_b: Weight, + v_b: Weight, + output: Weight, + indexer_k: Weight, + indexer_q: Weight, + indexer_k_norm: Weight, + indexer_k_bias: Weight, + indexer_proj: Weight, + ffn_norm: Weight, + dense: Option, + sparse: Option, +} + +struct GlmWeights { + embedding: Weight, + output_norm: Weight, + output: Weight, + layers: Vec, +} + +impl GlmWeights { + fn bind(model: &Model) -> Result { + let main = &model.main; + let normal_layers = model.shape.layers - model.shape.nextn; + let layers = (0..normal_layers) + .map(|index| { + let required = |suffix: &str| Weight::bind(main, &format!("blk.{index}.{suffix}")); + let dense = (index < model.shape.leading_dense) + .then(|| { + Ok::<_, String>(DenseWeights { + gate: required("ffn_gate.weight")?, + up: required("ffn_up.weight")?, + down: required("ffn_down.weight")?, + }) + }) + .transpose()?; + let sparse = (index >= model.shape.leading_dense) + .then(|| { + Ok::<_, String>(SparseWeights { + router: required("ffn_gate_inp.weight")?, + bias: required("exp_probs_b.bias")?, + gate: required("ffn_gate_exps.weight")?, + up: required("ffn_up_exps.weight")?, + down: required("ffn_down_exps.weight")?, + shared_gate: required("ffn_gate_shexp.weight")?, + shared_up: required("ffn_up_shexp.weight")?, + shared_down: required("ffn_down_shexp.weight")?, + }) + }) + .transpose()?; + Ok(GlmLayer { + attn_norm: required("attn_norm.weight")?, + q_a: required("attn_q_a.weight")?, + q_a_norm: required("attn_q_a_norm.weight")?, + q_b: required("attn_q_b.weight")?, + kv_a: required("attn_kv_a_mqa.weight")?, + kv_norm: required("attn_kv_a_norm.weight")?, + k_b: required("attn_k_b.weight")?, + v_b: required("attn_v_b.weight")?, + output: required("attn_output.weight")?, + indexer_k: required("indexer.attn_k.weight")?, + indexer_q: required("indexer.attn_q_b.weight")?, + indexer_k_norm: required("indexer.k_norm.weight")?, + indexer_k_bias: required("indexer.k_norm.bias")?, + indexer_proj: required("indexer.proj.weight")?, + ffn_norm: required("ffn_norm.weight")?, + dense, + sparse, + }) + }) + .collect::>()?; + Ok(Self { + embedding: Weight::bind(main, "token_embd.weight")?, + output_norm: Weight::bind(main, "output_norm.weight")?, + output: Weight::bind(main, "output.weight")?, + layers, + }) + } +} + +struct LayerCache { + kv: Buffer, + rope: Buffer, + indexer: Option, +} + +impl LayerCache { + fn allocate(shape: super::super::Shape, layer: usize, context: u32) -> Result { + Ok(Self { + kv: Buffer::bytes(u64::from(context) * shape.kv_lora * 2)?, + rope: Buffer::bytes(u64::from(context) * shape.rot * 2)?, + indexer: full_indexer_layer(shape, layer) + .then(|| Buffer::bytes(u64::from(context) * shape.indexer_head_dim * 2)) + .transpose()?, + }) + } +} + +struct GlmScratch { + current: Buffer, + next: Buffer, + attn_norm: Buffer, + q_rank: Buffer, + q_rank_norm: Buffer, + q: Buffer, + kv_raw: Buffer, + indexer_k: Buffer, + indexer_q: Buffer, + indexer_weights: Buffer, + indexer_scores: Buffer, + indexer_selected: Buffer, + qk_low: Buffer, + heads: Buffer, + attn_out: Buffer, + after_attn: Buffer, + ffn_norm: Buffer, + ffn_gate: Buffer, + ffn_up: Buffer, + ffn_mid: Buffer, + ffn_out: Buffer, + ffn_sum: Buffer, + router_logits: Buffer, + router_probs: Buffer, + router_selected: Buffer, + router_weights: Buffer, + output_norm: Buffer, + logits: Buffer, +} + +impl GlmScratch { + fn allocate(model: &Model, context: u32) -> Result { + let shape = model.shape; + let q = shape.heads * shape.key_mla; + let heads = shape.heads * shape.value_mla; + Ok(Self { + current: Buffer::floats(shape.embd)?, + next: Buffer::floats(shape.embd)?, + attn_norm: Buffer::floats(shape.embd)?, + q_rank: Buffer::floats(shape.lora_q)?, + q_rank_norm: Buffer::floats(shape.lora_q)?, + q: Buffer::floats(q)?, + kv_raw: Buffer::floats(shape.head_dim)?, + indexer_k: Buffer::floats(shape.indexer_head_dim)?, + indexer_q: Buffer::floats(shape.indexer_heads * shape.indexer_head_dim)?, + indexer_weights: Buffer::floats(shape.indexer_heads)?, + indexer_scores: Buffer::floats(u64::from(context))?, + indexer_selected: Buffer::bytes(shape.indexer_top_k * 4)?, + qk_low: Buffer::floats(shape.heads * shape.kv_lora)?, + heads: Buffer::floats(heads)?, + attn_out: Buffer::floats(shape.embd)?, + after_attn: Buffer::floats(shape.embd)?, + ffn_norm: Buffer::floats(shape.embd)?, + ffn_gate: Buffer::floats(shape.ff_dense.max(shape.experts_used * shape.ff_expert))?, + ffn_up: Buffer::floats(shape.ff_dense.max(shape.experts_used * shape.ff_expert))?, + ffn_mid: Buffer::floats(shape.ff_dense.max(shape.experts_used * shape.ff_expert))?, + ffn_out: Buffer::floats(shape.embd)?, + ffn_sum: Buffer::floats(shape.embd)?, + router_logits: Buffer::floats(shape.experts)?, + router_probs: Buffer::floats(shape.experts)?, + router_selected: Buffer::bytes(shape.experts_used * 4)?, + router_weights: Buffer::floats(shape.experts_used)?, + output_norm: Buffer::floats(shape.embd)?, + logits: Buffer::floats(shape.vocab)?, + }) + } +} + +pub(in crate::engine) struct GlmExecutor { + weights: GlmWeights, + scratch: GlmScratch, + caches: Vec, + logits: Vec, + tokens: Vec, + context: u32, + quality: bool, + ssd: EngineSsdSettings, + checkpoint_tag: [u8; 32], + model_modified: (u64, u32), + model_identity: [u8; 32], + _context: Context, + model: Model, +} + +impl GlmExecutor { + pub(super) fn open( + model: Model, + context: u32, + quality: bool, + ssd: EngineSsdSettings, + ) -> Result { + if context == 0 || u64::from(context) > model.shape.original_context { + return Err(format!( + "GLM context must be between 1 and {} tokens", + model.shape.original_context + )); + } + let weights = GlmWeights::bind(&model)?; + let admission = admission_bytes(&model, &weights, context, ssd)?; + let context_handle = Context::open(&model, quality, ssd.enabled, admission)?; + configure_streaming(&model, &weights, ssd)?; + let scratch = GlmScratch::allocate(&model, context)?; + let caches = (0..weights.layers.len()) + .map(|layer| LayerCache::allocate(model.shape, layer, context)) + .collect::>()?; + let model_modified = fs::metadata(model.main.path()) + .and_then(|metadata| metadata.modified()) + .ok() + .and_then(|modified| modified.duration_since(UNIX_EPOCH).ok()) + .map(|duration| (duration.as_secs(), duration.subsec_nanos())) + .unwrap_or_default(); + let model_identity = model.checkpoint_identity(); + Ok(Self { + weights, + scratch, + caches, + logits: vec![0.0; model.shape.vocab as usize], + tokens: Vec::new(), + context, + quality, + ssd, + checkpoint_tag: [0; 32], + model_modified, + model_identity, + _context: context_handle, + model, + }) + } + + pub(super) fn eval(&mut self, token: i32) -> Result<(), String> { + let shape = self.model.shape; + if token < 0 || token as u64 >= shape.vocab { + return Err(format!("token {token} is outside the vocabulary")); + } + let pos = u32::try_from(self.tokens.len()).map_err(|_| "GLM position overflow")?; + if pos >= self.context { + return Err(format!( + "the GLM Metal executor supports {} tokens", + self.context + )); + } + let map = self.model.main.map_ptr().cast(); + let size = self.model.main.len(); + let commands = Commands::begin()?; + call( + unsafe { + ds4_gpu_embed_token_quant_tensor( + self.scratch.current.raw(), + map, + size, + self.weights.embedding.offset, + self.weights.embedding.kind, + shape.vocab as u32, + token as u32, + shape.embd as u32, + ) + }, + "GLM token embedding", + )?; + let mut selected_count = 0; + for (layer_index, ((layer, cache), ordinal)) in self + .weights + .layers + .iter() + .zip(&self.caches) + .zip(0_u32..) + .enumerate() + { + self.encode_layer(layer, cache, layer_index, ordinal, pos, &mut selected_count)?; + std::mem::swap(&mut self.scratch.current, &mut self.scratch.next); + if (layer_index + 1).is_multiple_of(DECODE_FLUSH_LAYERS) + && layer_index + 1 < self.weights.layers.len() + { + call( + unsafe { ds4_gpu_flush_commands() }, + "flushing the GLM decode graph", + )?; + } + } + norm( + &self.scratch.output_norm, + &self.scratch.current, + self.weights.output_norm, + shape.embd as u32, + shape.rms_epsilon, + map, + size, + )?; + project( + &self.scratch.logits, + self.weights.output, + shape.embd, + shape.vocab, + &self.scratch.output_norm, + map, + size, + self.ssd.enabled, + )?; + commands.finish()?; + self.scratch.logits.read_f32(&mut self.logits)?; + self.tokens.push(token); + Ok(()) + } + + fn encode_layer( + &self, + layer: &GlmLayer, + cache: &LayerCache, + _layer_index: usize, + ordinal: u32, + pos: u32, + selected_count: &mut u32, + ) -> Result<(), String> { + let shape = self.model.shape; + let map = self.model.main.map_ptr().cast(); + let size = self.model.main.len(); + let q_dim = shape.heads * shape.key_mla; + let q_nope = shape.key_mla - shape.rot; + norm( + &self.scratch.attn_norm, + &self.scratch.current, + layer.attn_norm, + shape.embd as u32, + shape.rms_epsilon, + map, + size, + )?; + project( + &self.scratch.q_rank, + layer.q_a, + shape.embd, + shape.lora_q, + &self.scratch.attn_norm, + map, + size, + self.ssd.enabled, + )?; + project( + &self.scratch.kv_raw, + layer.kv_a, + shape.embd, + shape.head_dim, + &self.scratch.attn_norm, + map, + size, + self.ssd.enabled, + )?; + call( + unsafe { + ds4_gpu_glm_qkv_norm_store_compact_kv_tensor( + self.scratch.q_rank_norm.raw(), + self.scratch.q_rank.raw(), + map, + size, + layer.q_a_norm.offset, + shape.lora_q as u32, + cache.kv.raw(), + cache.rope.raw(), + self.scratch.kv_raw.raw(), + layer.kv_norm.offset, + pos, + 1, + self.context, + shape.head_dim as u32, + shape.kv_lora as u32, + shape.rot as u32, + CACHE_F16, + shape.rms_epsilon, + ) + }, + "storing GLM compact KV", + )?; + project( + &self.scratch.q, + layer.q_b, + shape.lora_q, + q_dim, + &self.scratch.q_rank_norm, + map, + size, + self.ssd.enabled, + )?; + call( + unsafe { + ds4_gpu_glm_rope_tail_tensor( + self.scratch.q.raw(), + 1, + shape.heads as u32, + shape.key_mla as u32, + shape.rot as u32, + pos, + 0, + shape.rope_base, + 1.0, + 0.0, + 1.0, + 0.0, + 0.0, + ) + }, + "applying GLM query RoPE", + )?; + + if let Some(indexer_cache) = &cache.indexer { + project( + &self.scratch.indexer_k, + layer.indexer_k, + shape.embd, + shape.indexer_head_dim, + &self.scratch.current, + map, + size, + self.ssd.enabled, + )?; + call( + unsafe { + ds4_gpu_glm_store_indexer_k_tensor( + indexer_cache.raw(), + self.scratch.indexer_k.raw(), + map, + size, + layer.indexer_k_norm.offset, + layer.indexer_k_bias.offset, + pos, + 1, + self.context, + shape.indexer_head_dim as u32, + shape.rot as u32, + 0, + 1.0e-6, + shape.rope_base, + 1.0, + 0.0, + 1.0, + 0.0, + 0.0, + CACHE_F16, + ) + }, + "storing the GLM indexer key", + )?; + let visible = pos + 1; + *selected_count = visible.min(shape.indexer_top_k as u32); + if visible <= shape.indexer_top_k as u32 { + call( + unsafe { + ds4_gpu_glm_fill_selected_range_tensor( + self.scratch.indexer_selected.raw(), + *selected_count, + ) + }, + "selecting the visible GLM context", + )?; + } else { + project( + &self.scratch.indexer_q, + layer.indexer_q, + shape.lora_q, + shape.indexer_heads * shape.indexer_head_dim, + &self.scratch.q_rank_norm, + map, + size, + self.ssd.enabled, + )?; + call( + unsafe { + ds4_gpu_glm_indexer_rope_tail_tensor( + self.scratch.indexer_q.raw(), + 1, + shape.indexer_heads as u32, + shape.indexer_head_dim as u32, + shape.rot as u32, + pos, + 0, + shape.rope_base, + 1.0, + 0.0, + 1.0, + 0.0, + 0.0, + ) + }, + "applying GLM indexer RoPE", + )?; + f32_project( + &self.scratch.indexer_weights, + layer.indexer_proj, + shape.embd, + shape.indexer_heads, + &self.scratch.current, + map, + size, + )?; + let scale = 1.0 / ((shape.indexer_heads * shape.indexer_head_dim) as f32).sqrt(); + call( + unsafe { + ds4_gpu_glm_indexer_score_one_tensor( + self.scratch.indexer_scores.raw(), + self.scratch.indexer_q.raw(), + self.scratch.indexer_weights.raw(), + indexer_cache.raw(), + visible, + shape.indexer_heads as u32, + shape.indexer_head_dim as u32, + scale, + CACHE_F16, + ) + }, + "scoring the GLM indexer", + )?; + call( + unsafe { + ds4_gpu_indexer_topk_tensor( + self.scratch.indexer_selected.raw(), + self.scratch.indexer_scores.raw(), + visible, + 1, + *selected_count, + ) + }, + "selecting GLM indexed attention rows", + )?; + } + } + if *selected_count == 0 { + return Err("GLM indexer did not select an attention context".into()); + } + call( + unsafe { + ds4_gpu_glm_qk_lowrank_typed_tensor( + self.scratch.qk_low.raw(), + self.scratch.q.raw(), + map, + size, + layer.k_b.offset, + layer.k_b.kind, + shape.heads as u32, + shape.kv_lora as u32, + q_nope as u32, + shape.key_mla as u32, + ) + }, + "projecting the GLM low-rank query", + )?; + call( + unsafe { + ds4_gpu_glm_attention_indexed_decode_typed_tensor( + self.scratch.heads.raw(), + self.scratch.q.raw(), + self.scratch.qk_low.raw(), + cache.kv.raw(), + cache.rope.raw(), + map, + size, + layer.v_b.offset, + layer.v_b.kind, + self.scratch.indexer_selected.raw(), + *selected_count, + self.context, + CACHE_F16, + shape.heads as u32, + shape.kv_lora as u32, + q_nope as u32, + shape.rot as u32, + shape.value_mla as u32, + 0, + shape.rope_base, + 1.0, + 0.0, + 1.0, + 0.0, + 0.0, + ) + }, + "running GLM indexed attention", + )?; + project( + &self.scratch.attn_out, + layer.output, + shape.heads * shape.value_mla, + shape.embd, + &self.scratch.heads, + map, + size, + self.ssd.enabled, + )?; + call( + unsafe { + ds4_gpu_add_rms_norm_weight_tensor( + self.scratch.ffn_norm.raw(), + self.scratch.after_attn.raw(), + self.scratch.current.raw(), + self.scratch.attn_out.raw(), + map, + size, + layer.ffn_norm.offset, + shape.embd as u32, + shape.rms_epsilon, + ) + }, + "normalizing the GLM FFN input", + )?; + + if let Some(dense) = layer.dense { + project( + &self.scratch.ffn_gate, + dense.gate, + shape.embd, + shape.ff_dense, + &self.scratch.ffn_norm, + map, + size, + self.ssd.enabled, + )?; + project( + &self.scratch.ffn_up, + dense.up, + shape.embd, + shape.ff_dense, + &self.scratch.ffn_norm, + map, + size, + self.ssd.enabled, + )?; + call( + unsafe { + ds4_gpu_swiglu_tensor( + self.scratch.ffn_mid.raw(), + self.scratch.ffn_gate.raw(), + self.scratch.ffn_up.raw(), + shape.ff_dense as u32, + 0.0, + 1.0, + ) + }, + "activating the dense GLM FFN", + )?; + project( + &self.scratch.ffn_out, + dense.down, + shape.ff_dense, + shape.embd, + &self.scratch.ffn_mid, + map, + size, + self.ssd.enabled, + )?; + call( + unsafe { + ds4_gpu_add_tensor( + self.scratch.next.raw(), + self.scratch.after_attn.raw(), + self.scratch.ffn_out.raw(), + shape.embd as u32, + ) + }, + "adding the dense GLM residual", + )?; + } else if let Some(sparse) = layer.sparse { + f32_project( + &self.scratch.router_logits, + sparse.router, + shape.embd, + shape.experts, + &self.scratch.ffn_norm, + map, + size, + )?; + call( + unsafe { + ds4_gpu_glm_router_select_tensor( + self.scratch.router_selected.raw(), + self.scratch.router_weights.raw(), + self.scratch.router_probs.raw(), + map, + size, + sparse.bias.offset, + self.scratch.router_logits.raw(), + shape.experts as u32, + shape.experts_used as u32, + shape.expert_weight_scale, + ) + }, + "routing the GLM experts", + )?; + let force_resident = self.ssd.enabled + && ordinal.saturating_sub(shape.leading_dense) < self.ssd.full_layers; + if self.ssd.enabled && !force_resident { + let table = expert_table(map, size, ordinal, shape, sparse); + call( + unsafe { + ds4_gpu_glm_stream_expert_cache_begin_selected_load_tensor( + &table, + self.scratch.router_selected.raw(), + shape.experts_used as u32, + ) + }, + "loading selected GLM experts", + )?; + } + let (gate_expert, gate_row) = expert_layout(sparse.gate, shape.experts); + let (up_expert, up_row) = expert_layout(sparse.up, shape.experts); + let (down_expert, down_row) = expert_layout(sparse.down, shape.experts); + call( + unsafe { + ds4_gpu_glm_routed_moe_one_tensor( + self.scratch.ffn_out.raw(), + self.scratch.ffn_mid.raw(), + map, + size, + sparse.gate.offset, + sparse.up.offset, + sparse.down.offset, + sparse.gate.kind, + sparse.up.kind, + sparse.down.kind, + gate_expert, + gate_row, + up_expert, + up_row, + down_expert, + down_row, + shape.embd as u32, + shape.ff_expert as u32, + shape.embd as u32, + self.scratch.router_selected.raw(), + self.scratch.router_weights.raw(), + shape.experts as u32, + shape.experts_used as u32, + ordinal, + self.scratch.ffn_norm.raw(), + force_resident, + ) + }, + "running the routed GLM experts", + )?; + project( + &self.scratch.ffn_gate, + sparse.shared_gate, + shape.embd, + shape.ff_expert, + &self.scratch.ffn_norm, + map, + size, + self.ssd.enabled, + )?; + project( + &self.scratch.ffn_up, + sparse.shared_up, + shape.embd, + shape.ff_expert, + &self.scratch.ffn_norm, + map, + size, + self.ssd.enabled, + )?; + call( + unsafe { + ds4_gpu_swiglu_tensor( + self.scratch.ffn_mid.raw(), + self.scratch.ffn_gate.raw(), + self.scratch.ffn_up.raw(), + shape.ff_expert as u32, + 0.0, + 1.0, + ) + }, + "activating the shared GLM expert", + )?; + project( + &self.scratch.ffn_sum, + sparse.shared_down, + shape.ff_expert, + shape.embd, + &self.scratch.ffn_mid, + map, + size, + self.ssd.enabled, + )?; + call( + unsafe { + ds4_gpu_add3_tensor( + self.scratch.next.raw(), + self.scratch.after_attn.raw(), + self.scratch.ffn_out.raw(), + self.scratch.ffn_sum.raw(), + shape.embd as u32, + ) + }, + "adding the sparse GLM residual", + )?; + } + Ok(()) + } + + pub(super) fn prefill( + &mut self, + tokens: &[i32], + mut progress: impl FnMut(u32) -> bool, + ) -> Result { + let mut completed = 0; + for &token in tokens { + if !progress(self.position()) { + break; + } + self.eval(token)?; + completed += 1; + } + Ok(completed) + } + + pub(super) fn logits(&self) -> &[f32] { + &self.logits + } + pub(super) fn model(&self) -> &Model { + &self.model + } + pub(super) fn context(&self) -> u32 { + self.context + } + pub(super) fn position(&self) -> u32 { + self.tokens.len() as u32 + } + pub(super) fn tokens(&self) -> &[i32] { + &self.tokens + } + pub(super) fn checkpoint_tag(&self) -> [u8; 32] { + self.checkpoint_tag + } + pub(super) fn note_checkpoint_tag(&mut self, tag: [u8; 32]) { + self.checkpoint_tag = tag; + } + + pub(super) fn reset(&mut self) -> Result<(), String> { + self.scratch = GlmScratch::allocate(&self.model, self.context)?; + self.caches = (0..self.weights.layers.len()) + .map(|layer| LayerCache::allocate(self.model.shape, layer, self.context)) + .collect::>()?; + self.tokens.clear(); + self.checkpoint_tag = [0; 32]; + Ok(()) + } + + pub(super) fn align_prompt(&mut self, tokens: &[i32]) -> Result { + if !tokens.starts_with(&self.tokens) { + self.reset()?; + } + Ok(self.tokens.len()) + } + + pub(super) fn save_checkpoint( + &mut self, + path: &Path, + tag: [u8; 32], + progress: &mut impl FnMut(u64), + ) -> Result<(), String> { + if let Some(parent) = path.parent() { + fs::create_dir_all(parent).map_err(|e| e.to_string())?; + } + let temporary = path.with_extension("tmp"); + let mut file = File::create(&temporary).map_err(|e| e.to_string())?; + file.write_all(CHECKPOINT_MAGIC) + .map_err(|e| e.to_string())?; + let shape = self.model.shape; + for value in [ + CHECKPOINT_VERSION, + self.context, + self.position(), + self.quality as u32, + self.weights.layers.len() as u32, + shape.kv_lora as u32, + shape.rot as u32, + shape.indexer_head_dim as u32, + shape.vocab as u32, + ] { + write_u32(&mut file, value)?; + } + write_u64(&mut file, self.model.main.len())?; + write_u64(&mut file, self.model_modified.0)?; + write_u32(&mut file, self.model_modified.1)?; + file.write_all(&self.model_identity) + .map_err(|e| e.to_string())?; + file.write_all(&tag).map_err(|e| e.to_string())?; + for &token in &self.tokens { + write_u32(&mut file, token as u32)?; + } + for &logit in &self.logits { + write_u32(&mut file, logit.to_bits())?; + } + let mut chunk = vec![0; CHECKPOINT_IO_CHUNK]; + let rows = u64::from(self.position()); + for cache in &self.caches { + write_buffer( + &mut file, + &cache.kv, + 0, + rows * shape.kv_lora * 2, + &mut chunk, + progress, + )?; + write_buffer( + &mut file, + &cache.rope, + 0, + rows * shape.rot * 2, + &mut chunk, + progress, + )?; + if let Some(indexer) = &cache.indexer { + write_buffer( + &mut file, + indexer, + 0, + rows * shape.indexer_head_dim * 2, + &mut chunk, + progress, + )?; + } + } + file.sync_all().map_err(|e| e.to_string())?; + fs::rename(&temporary, path).map_err(|e| e.to_string())?; + self.checkpoint_tag = tag; + Ok(()) + } + + pub(super) fn load_checkpoint( + &mut self, + path: &Path, + progress: &mut impl FnMut(u64), + ) -> Result { + let mut file = match File::open(path) { + Ok(file) => file, + Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(false), + Err(error) => return Err(error.to_string()), + }; + let mut magic = [0; 8]; + file.read_exact(&mut magic).map_err(|e| e.to_string())?; + if &magic != CHECKPOINT_MAGIC + || read_u32(&mut file)? != CHECKPOINT_VERSION + || read_u32(&mut file)? != self.context + { + return Err("GLM checkpoint does not match the current executor".into()); + } + let token_count = read_u32(&mut file)?; + let shape = self.model.shape; + if token_count > self.context + || read_u32(&mut file)? != self.quality as u32 + || read_u32(&mut file)? != self.weights.layers.len() as u32 + || read_u32(&mut file)? != shape.kv_lora as u32 + || read_u32(&mut file)? != shape.rot as u32 + || read_u32(&mut file)? != shape.indexer_head_dim as u32 + || read_u32(&mut file)? != shape.vocab as u32 + || read_u64(&mut file)? != self.model.main.len() + || read_u64(&mut file)? != self.model_modified.0 + || read_u32(&mut file)? != self.model_modified.1 + { + return Err("GLM checkpoint was written for a different model or configuration".into()); + } + let mut identity = [0; 32]; + file.read_exact(&mut identity).map_err(|e| e.to_string())?; + if identity != self.model_identity { + return Err("GLM checkpoint model identity changed".into()); + } + let mut tag = [0; 32]; + file.read_exact(&mut tag).map_err(|e| e.to_string())?; + let mut tokens = Vec::with_capacity(token_count as usize); + for _ in 0..token_count { + let token = read_u32(&mut file)?; + if u64::from(token) >= self.model.shape.vocab { + return Err("GLM checkpoint token is invalid".into()); + } + tokens.push(token as i32); + } + let mut logits = Vec::with_capacity(shape.vocab as usize); + for _ in 0..shape.vocab { + logits.push(f32::from_bits(read_u32(&mut file)?)); + } + self.reset()?; + let mut chunk = vec![0; CHECKPOINT_IO_CHUNK]; + let rows = u64::from(token_count); + for cache in &self.caches { + read_buffer( + &mut file, + &cache.kv, + 0, + rows * shape.kv_lora * 2, + &mut chunk, + progress, + )?; + read_buffer( + &mut file, + &cache.rope, + 0, + rows * shape.rot * 2, + &mut chunk, + progress, + )?; + if let Some(indexer) = &cache.indexer { + read_buffer( + &mut file, + indexer, + 0, + rows * shape.indexer_head_dim * 2, + &mut chunk, + progress, + )?; + } + } + let mut trailing = [0]; + if file.read(&mut trailing).map_err(|e| e.to_string())? != 0 { + return Err("GLM checkpoint has trailing data".into()); + } + self.tokens = tokens; + self.logits = logits; + self.checkpoint_tag = tag; + Ok(true) + } +} + +fn full_indexer_layer(shape: super::super::Shape, layer: usize) -> bool { + layer < (shape.layers - shape.nextn) as usize + && (layer < shape.leading_dense as usize || (layer >= 6 && (layer - 6).is_multiple_of(4))) +} + +fn expert_layout(weight: Weight, experts: u64) -> (u64, u64) { + let expert = weight.bytes / experts; + (expert, expert / weight.dims[1]) +} + +fn expert_table( + map: *const c_void, + size: u64, + layer: u32, + shape: super::super::Shape, + weights: SparseWeights, +) -> StreamExpertTable { + StreamExpertTable { + model_map: map, + model_size: size, + layer, + total_experts: shape.experts as u32, + gate_offset: weights.gate.offset, + up_offset: weights.up.offset, + down_offset: weights.down.offset, + gate_expert_bytes: weights.gate.bytes / shape.experts, + down_expert_bytes: weights.down.bytes / shape.experts, + } +} + +fn configure_streaming( + model: &Model, + weights: &GlmWeights, + ssd: EngineSsdSettings, +) -> Result<(), String> { + if !ssd.enabled { + return Ok(()); + } + let sparse = weights + .layers + .iter() + .find_map(|layer| layer.sparse) + .ok_or("GLM model has no routed expert layers")?; + let gate = sparse.gate.bytes / model.shape.experts; + let up = sparse.up.bytes / model.shape.experts; + let down = sparse.down.bytes / model.shape.experts; + let per_expert = gate + .checked_add(up) + .and_then(|bytes| bytes.checked_add(down)) + .ok_or("GLM expert size overflow")?; + let budget = if ssd.cache_experts != 0 { + ssd.cache_experts + } else if ssd.cache_bytes != 0 { + u32::try_from((ssd.cache_bytes / per_expert).max(1)).unwrap_or(u32::MAX) + } else { + unsafe { ds4_gpu_stream_expert_cache_budget_for_expert_size(gate, down) } + }; + if budget == 0 { + return Err("GLM SSD streaming has no memory for an expert cache".into()); + } + unsafe { + ds4_gpu_set_streaming_expert_cache_expert_bytes(per_expert); + ds4_gpu_set_streaming_expert_cache_budget(budget); + } + Ok(()) +} + +fn admission_bytes( + model: &Model, + weights: &GlmWeights, + context: u32, + ssd: EngineSsdSettings, +) -> Result { + let shape = model.shape; + let normal_layers = u64::from(shape.layers - shape.nextn); + let indexer_layers = (0..normal_layers as usize) + .filter(|layer| full_indexer_layer(shape, *layer)) + .count() as u64; + let per_layer = u64::from(context) + .checked_mul((shape.kv_lora + shape.rot) * 2) + .ok_or("GLM compact-cache size overflow")?; + let kv = normal_layers + .checked_mul(per_layer) + .and_then(|bytes| { + bytes.checked_add(indexer_layers * u64::from(context) * shape.indexer_head_dim * 2) + }) + .ok_or("GLM compact-cache size overflow")?; + let resident = if ssd.enabled { + model + .main + .tensors + .iter() + .filter(|(name, _)| { + !name.ends_with("ffn_gate_exps.weight") + && !name.ends_with("ffn_up_exps.weight") + && !name.ends_with("ffn_down_exps.weight") + }) + .try_fold(0_u64, |total, (_, tensor)| total.checked_add(tensor.bytes)) + .ok_or("GLM resident tensor size overflow")? + } else { + model.main.len() - model.main.data_offset() + }; + let sparse_expert = weights + .layers + .iter() + .find_map(|layer| layer.sparse) + .map(|weights| { + weights.gate.bytes / shape.experts + + weights.up.bytes / shape.experts + + weights.down.bytes / shape.experts + }) + .unwrap_or(0); + let cache = if ssd.enabled { + if ssd.cache_bytes != 0 { + ssd.cache_bytes + } else if ssd.cache_experts != 0 { + sparse_expert.saturating_mul(u64::from(ssd.cache_experts)) + } else { + // Reference auto policy reserves roughly one seventh of the available + // budget for full layers and uses the remainder for selected experts. + 16 * 1024 * 1024 * 1024_u64 + } + } else { + 0 + }; + let full_layers = sparse_expert + .saturating_mul(shape.experts) + .saturating_mul(u64::from(ssd.full_layers)); + resident + .checked_add(kv) + .and_then(|bytes| bytes.checked_add(cache)) + .and_then(|bytes| bytes.checked_add(full_layers)) + .and_then(|bytes| bytes.checked_add(512 * 1024 * 1024)) + .ok_or_else(|| "GLM runtime memory size overflow".into()) +} + +fn norm( + out: &Buffer, + input: &Buffer, + weight: Weight, + width: u32, + epsilon: f32, + map: *const c_void, + size: u64, +) -> Result<(), String> { + call( + unsafe { + ds4_gpu_rms_norm_weight_tensor( + out.raw(), + input.raw(), + map, + size, + weight.offset, + width, + epsilon, + ) + }, + "normalizing GLM activations", + ) +} + +#[allow(clippy::too_many_arguments)] +fn project( + out: &Buffer, + weight: Weight, + input: u64, + output: u64, + x: &Buffer, + map: *const c_void, + size: u64, + streaming: bool, +) -> Result<(), String> { + let result = unsafe { + if streaming { + ds4_gpu_matmul_quant_tensor( + out.raw(), + map, + size, + weight.offset, + weight.kind, + input, + output, + x.raw(), + 1, + ) + } else { + ds4_gpu_matmul_quant_decode_mpp_model_view_tensor( + out.raw(), + map, + size, + weight.offset, + weight.kind, + input, + output, + x.raw(), + 1, + ) + } + }; + call(result, "projecting GLM activations") +} + +#[allow(clippy::too_many_arguments)] +fn f32_project( + out: &Buffer, + weight: Weight, + input: u64, + output: u64, + x: &Buffer, + map: *const c_void, + size: u64, +) -> Result<(), String> { + call( + unsafe { + ds4_gpu_matmul_f32_tensor( + out.raw(), + map, + size, + weight.offset, + input, + output, + x.raw(), + 1, + ) + }, + "projecting GLM F32 activations", + ) +} + +#[cfg(test)] +mod tests { + use super::{GlmExecutor, full_indexer_layer}; + use crate::engine::{GLM, Model}; + use crate::model::ModelChoice; + use crate::settings::EngineSsdSettings; + use std::path::Path; + + #[test] + fn dsa_indexer_schedule_matches_the_reference() { + assert!(full_indexer_layer(GLM, 0)); + assert!(full_indexer_layer(GLM, 2)); + assert!(!full_indexer_layer(GLM, 3)); + assert!(full_indexer_layer(GLM, 6)); + assert!(full_indexer_layer(GLM, 10)); + assert!(!full_indexer_layer(GLM, 11)); + assert!(!full_indexer_layer(GLM, 78)); + } + + #[test] + #[ignore = "requires the 197 GiB GLM 5.2 checkpoint and Apple Metal"] + fn resident_and_streamed_glm_match_the_short_code_fixture() { + let path = std::env::var("DS4_GLM_MODEL").unwrap_or_else(|_| { + "../ds4/models/GLM-5.2-UD-IQ2_XXS_RoutedIQ2XXS_blk78Q2K.gguf".into() + }); + let prompt = b"Complete the C statement with the next exact token only:\nreturn snprintf(buf, sizeof(buf), \"%d\", value"; + for streamed in [false, true] { + let model = Model::open_main(Path::new(&path), ModelChoice::Glm52).unwrap(); + let tokens = model.tokenize(std::str::from_utf8(prompt).unwrap()); + let mut executor = GlmExecutor::open( + model, + 4096, + true, + EngineSsdSettings { + enabled: streamed, + cold: false, + cache_experts: 0, + cache_bytes: 0, + full_layers: 0, + full_layers_set: false, + preload_experts: 0, + }, + ) + .unwrap(); + assert_eq!(executor.prefill(&tokens, |_| true).unwrap(), tokens.len()); + let token = executor + .logits() + .iter() + .enumerate() + .max_by(|a, b| a.1.total_cmp(b.1)) + .unwrap() + .0 as i32; + assert_eq!( + executor.model().token_bytes(token).as_deref(), + Some(b")".as_slice()) + ); + } + } +} diff --git a/src/engine/metal/gpu.rs b/src/engine/metal/gpu.rs index 3c2b0e8..a0a0b6a 100644 --- a/src/engine/metal/gpu.rs +++ b/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 { + pub(super) fn open( + model: &Model, + quality: bool, + ssd_streaming: bool, + admission_bytes: u64, + ) -> Result { 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, + }) } } diff --git a/src/server.rs b/src/server.rs index 00af802..9c39996 100644 --- a/src/server.rs +++ b/src/server.rs @@ -543,9 +543,6 @@ fn compatible_completion( fn installed_endpoint_models(models_path: &std::path::Path) -> Vec { model::installed_models(models_path) - .into_iter() - .filter(|model| *model != ModelChoice::Glm52) - .collect() } fn models_json(state: &State) -> Value {