long term effort: Serve safetensors models on Metal via MLX (Apple Silicon)

Cargo.toml

@@ -37,10 +37,12 @@ members = [
     "crates/model-hf",
     "crates/model-resolver",
     "crates/skippy-protocol",
+    "crates/skippy-engine",
     "crates/skippy-coordinator",
     "crates/skippy-topology",
     "crates/skippy-cache",
     "crates/skippy-metrics",
+    "crates/skippy-engine-mlx",
     "crates/openai-frontend",
     "crates/skippy-ffi",
     "crates/skippy-runtime",
@@ -81,6 +83,9 @@ strum = { version = "0.28", features = ["derive"] }
 
 [patch.crates-io]
 hf-hub = { git = "https://github.com/Mesh-LLM/hf-hub", branch = "mesh-llm" }
+safemlx = { git = "https://github.com/michaelneale/safemlx", rev = "c6b4741" }
+safemlx-lm = { git = "https://github.com/michaelneale/safemlx", rev = "c6b4741" }
+safemlx-lm-utils = { git = "https://github.com/michaelneale/safemlx", rev = "c6b4741" }
 
 [workspace.lints.clippy]
 cognitive_complexity = "warn"

Justfile

@@ -239,6 +239,39 @@ skippy-quantize-standalone-build backend="cpu":
 skippy-quantize-standalone-release-build backend="cpu":
     LLAMA_STAGE_BACKEND="{{ backend }}" LLAMA_STAGE_LINK_MODE=static just with-lld cargo build --release --locked -p skippy-quantize
 
+# Build and test the standalone SafeTensors stage-range research spike.
+mlx-safetensors-stage-plan-test:
+    just with-lld cargo test --manifest-path spikes/mlx-safetensors-stages/Cargo.toml
+
+# Inspect a remote checkpoint without downloading tensor payloads.
+mlx-safetensors-stage-plan *ARGS:
+    just with-lld cargo run --manifest-path spikes/mlx-safetensors-stages/Cargo.toml -- {{ ARGS }}
+
+# Compare whole-model MLX against two partial SafeTensors stages on Metal.
+mlx-safetensors-split-proof *ARGS:
+    DEVELOPER_DIR=/Applications/Xcode.app/Contents/Developer just with-lld cargo run --release --manifest-path spikes/mlx-solo/Cargo.toml --bin mlx-split-proof -- {{ ARGS }}
+
+# Build the MLX stage binary and its required sibling Metal library.
+mlx-stage-build:
+    DEVELOPER_DIR=/Applications/Xcode.app/Contents/Developer just with-lld cargo build --release -p skippy-engine-mlx --features mlx --bin mlx-stage
+    test -s target/release/safemlx-resources/mlx.metallib
+    cp target/release/safemlx-resources/mlx.metallib target/release/mlx.metallib
+    cmp -s target/release/safemlx-resources/mlx.metallib target/release/mlx.metallib
+
+# Build the shipped mesh-llm binary with whole-model and distributed MLX serving.
+[macos]
+mlx-build:
+    DEVELOPER_DIR=/Applications/Xcode.app/Contents/Developer MESH_LLM_CARGO_FEATURES=mlx scripts/build-mac.sh
+
+# Release-mode MLX build; writes mesh-llm and its sibling mlx.metallib.
+[macos]
+mlx-release-build:
+    DEVELOPER_DIR=/Applications/Xcode.app/Contents/Developer MESH_LLM_BUILD_PROFILE=release MESH_LLM_CARGO_FEATURES=mlx scripts/build-mac.sh
+
+# Run `mlx-stage serve ...` or `mlx-stage prove ...` after `just mlx-stage-build`.
+mlx-stage *ARGS:
+    target/release/mlx-stage {{ ARGS }}
+
 # Generate a reproducible benchmark corpus for skippy bench tooling.
 bench-corpus tier="smoke" *ARGS="":
     scripts/generate-bench-corpus.py "{{ tier }}" {{ ARGS }}

crates/mesh-llm-host-runtime/Cargo.toml

@@ -17,6 +17,18 @@ web-ui = ["mesh-llm-ui/embed-assets"]
 gpu-bench-cuda = ["mesh-llm-system/gpu-bench-cuda"]
 gpu-bench-hip = ["mesh-llm-system/gpu-bench-hip"]
 gpu-bench-intel = ["mesh-llm-system/gpu-bench-intel"]
+# Serve HF safetensors models on Apple Silicon via the MLX (Metal) engine.
+# macOS-only in practice: the dep is declared under a macOS target table, and all
+# engine code is additionally cfg(target_os = "macos") gated. Off by default so
+# the heavy MLX native build never runs in normal builds/CI.
+#
+# Implies `dynamic-native-runtime`: MLX statically links its own vendored GGUF
+# parser (antirez gguflib), which collides with the patched llama.cpp GGUF
+# symbols (e.g. `gguf_get_key`) when llama.cpp is linked statically. Loading the
+# llama.cpp runtime dynamically (as release builds already do) keeps the two
+# native GGUF parsers in separate link units and avoids the duplicate-symbol
+# link failure.
+mlx = ["dep:skippy-engine-mlx", "dynamic-native-runtime"]
 dynamic-native-runtime = [
     "mesh-llm-system/dynamic-native-runtime",
     "skippy-runtime/dynamic-native-runtime",
@@ -114,6 +126,15 @@ hf_hub = { package = "hf-hub", version = "1.0.0-rc.1", default-features = false,
 tabwriter = "1"
 tempfile = "3"
 
+# MLX (Metal) serving engine — Apple Silicon only. Declared under a macOS target
+# table so non-macOS builds never see it, and optional so it is pulled in only
+# with `--features mlx`. Enabling the feature turns on the crate's own `mlx`
+# feature (which runs the heavy MLX native build).
+[target.'cfg(target_os = "macos")'.dependencies]
+skippy-engine-mlx = { path = "../skippy-engine-mlx", version = "0.72.1", features = [
+    "mlx",
+], optional = true }
+
 [dev-dependencies]
 serial_test = "3"
 mesh-client = { package = "mesh-llm-client", path = "../mesh-client", version = "0.72.1"  }

crates/mesh-llm-host-runtime/src/api/tests.rs

@@ -663,6 +663,7 @@ fn make_test_state_peer(seed: u8, role: mesh::NodeRole) -> mesh::PeerInfo {
         artifact_transfer_supported: false,
         stage_protocol_generation_supported: false,
         stage_status_list_supported: false,
+        mlx_stage_supported: false,
         owner_summary: crate::crypto::OwnershipSummary::default(),
         first_joined_mesh_ts: None,
         advertised_model_throughput: vec![],
@@ -1886,6 +1887,7 @@ fn make_test_peer(
         artifact_transfer_supported: false,
         stage_protocol_generation_supported: false,
         stage_status_list_supported: false,
+        mlx_stage_supported: false,
         owner_summary: crate::crypto::OwnershipSummary::default(),
         advertised_model_throughput: vec![],
 
@@ -2400,6 +2402,7 @@ async fn seed_runtime_data_api_state(state: &MeshApi) {
             n_batch: Some(2048),
             n_ubatch: Some(512),
             flash_attn_type: skippy_protocol::FlashAttentionType::Enabled,
+            weight_quantization: crate::inference::skippy::StageWeightQuantization::Auto,
             error: None,
             shutdown_generation: 7,
             coordinator_term: 11,

crates/mesh-llm-host-runtime/src/inference/mlx.rs

@@ -0,0 +1,136 @@
+//! MLX (Metal) serving integration — Apple Silicon only.
+//!
+//! Bridges the `skippy-engine-mlx` crate (which serves HF safetensors models on
+//! Metal via `safemlx`) into the host runtime's local-model launch path. Gated
+//! behind both the `mlx` cargo feature and `target_os = "macos"`, so it is
+//! entirely absent from every other build.
+//!
+//! Unlike the skippy/GGUF path (which owns an embedded HTTP server in
+//! `skippy-server`), MLX serves over the plain `openai-frontend` router; this
+//! module stands up that router on a local port and manages its lifecycle with a
+//! graceful-shutdown handle mirroring `SkippyHttpHandle`.
+
+use std::net::SocketAddr;
+use std::path::PathBuf;
+use std::sync::Arc;
+
+use anyhow::{Context, Result};
+use skippy_engine_mlx::{
+    MlxBackend, MlxDistributedEngine, MlxDistributedEngineConfig, MlxEngine, MlxEngineConfig,
+    automatic_weight_quantization,
+};
+use skippy_protocol::binary::WireActivationDType;
+
+const DEFAULT_MAX_GENERATION_TOKENS: usize = 512;
+
+/// A loaded MLX model plus the OpenAI backend that serves it.
+pub(crate) struct MlxModelHandle {
+    backend: Arc<MlxBackend>,
+}
+
+impl MlxModelHandle {
+    /// Loads a safetensors model directory on the MLX (Metal) engine. Blocking:
+    /// call from `spawn_blocking`.
+    pub(crate) fn load(model_dir: PathBuf, model_id: String, context_length: u32) -> Result<Self> {
+        let weight_quantization = automatic_weight_quantization(&model_dir)?;
+        let config = MlxEngineConfig {
+            model_dir,
+            model_id,
+            default_max_tokens: (context_length.max(1) as usize).min(DEFAULT_MAX_GENERATION_TOKENS),
+            max_tokens_cap: context_length.max(1) as usize,
+            weight_quantization,
+            allow_native_quantization_fallback: true,
+        };
+        let engine = MlxEngine::spawn(config)?;
+        Ok(Self {
+            backend: Arc::new(MlxBackend::new(engine)),
+        })
+    }
+
+    /// Loads tokenizer/chat-template sidecars and drives a mesh-managed MLX
+    /// stage chain rooted at `stage_addr`. No model weights are loaded here.
+    pub(crate) fn load_distributed(
+        model_dir: PathBuf,
+        model_id: String,
+        context_length: u32,
+        stage_addr: SocketAddr,
+        wire_dtype: WireActivationDType,
+    ) -> Result<Self> {
+        let engine = MlxDistributedEngine::spawn(MlxDistributedEngineConfig {
+            model_dir,
+            model_id,
+            stage_addr,
+            wire_dtype,
+            default_max_tokens: (context_length.max(1) as usize).min(DEFAULT_MAX_GENERATION_TOKENS),
+            max_tokens_cap: context_length.max(1) as usize,
+            context_tokens: context_length.max(1) as usize,
+        })?;
+        Ok(Self {
+            backend: Arc::new(MlxBackend::new_distributed(engine)),
+        })
+    }
+
+    /// Starts an `openai-frontend` HTTP server for this model on `port`.
+    pub(crate) async fn start_http(
+        &self,
+        port: u16,
+        death_tx: tokio::sync::oneshot::Sender<()>,
+    ) -> Result<MlxHttpHandle> {
+        let addr: SocketAddr = ([127, 0, 0, 1], port).into();
+        let app = openai_frontend::router::router_for(self.backend.clone());
+        let (shutdown_tx, shutdown_rx) = tokio::sync::oneshot::channel::<()>();
+        let listener = tokio::net::TcpListener::bind(addr)
+            .await
+            .with_context(|| format!("bind MLX OpenAI frontend at {addr}"))?;
+
+        let server = tokio::spawn(async move {
+            let serve = axum::serve(listener, app).with_graceful_shutdown(async move {
+                let _ = shutdown_rx.await;
+            });
+            if let Err(error) = serve.await {
+                tracing::error!(%error, "MLX openai frontend server error");
+            }
+            let _ = death_tx.send(());
+        });
+
+        Ok(MlxHttpHandle {
+            port,
+            shutdown_tx: Some(shutdown_tx),
+            server: Some(server),
+        })
+    }
+}
+
+/// Lifecycle handle for the MLX model's HTTP server.
+pub(crate) struct MlxHttpHandle {
+    port: u16,
+    shutdown_tx: Option<tokio::sync::oneshot::Sender<()>>,
+    server: Option<tokio::task::JoinHandle<()>>,
+}
+
+impl MlxHttpHandle {
+    pub(crate) fn port(&self) -> u16 {
+        self.port
+    }
+
+    pub(crate) async fn shutdown(mut self) -> Result<()> {
+        if let Some(tx) = self.shutdown_tx.take() {
+            let _ = tx.send(());
+        }
+        if let Some(server) = self.server.take() {
+            server.await.context("join MLX openai frontend task")?;
+        }
+        Ok(())
+    }
+}
+
+impl Drop for MlxHttpHandle {
+    fn drop(&mut self) {
+        if let Some(tx) = self.shutdown_tx.take() {
+            let _ = tx.send(());
+        }
+        if let Some(server) = self.server.take() {
+            server.abort();
+        }
+    }
+}

crates/mesh-llm-host-runtime/src/inference/mod.rs

@@ -1,5 +1,7 @@
 pub(crate) mod consult;
 pub(crate) mod election;
+#[cfg(all(feature = "mlx", target_os = "macos"))]
+pub(crate) mod mlx;
 pub(crate) mod pipeline;
 pub(crate) mod skippy;
 pub(crate) mod virtual_llm;

crates/mesh-llm-host-runtime/src/inference/skippy/deployment.rs

@@ -59,6 +59,7 @@ pub(crate) fn remote_stage_load_request(
         n_gpu_layers: -1,
         mmap: context.mmap,
         mlock: context.mlock,
+        weight_quantization: super::StageWeightQuantization::Auto,
         cache_type_k: context.kv_cache.cache_type_k().to_string(),
         cache_type_v: context.kv_cache.cache_type_v().to_string(),
         flash_attn_type: context.flash_attn_type,

crates/mesh-llm-host-runtime/src/inference/skippy/materialization.rs

@@ -1491,6 +1491,7 @@ fn cache_key(input: &str) -> String {
 #[cfg(test)]
 mod tests {
     use super::*;
+    use crate::inference::skippy::StageWeightQuantization;
     use std::ffi::OsString;
 
     use serial_test::serial;
@@ -1604,6 +1605,7 @@ mod tests {
             n_gpu_layers: -1,
             mmap: None,
             mlock: false,
+            weight_quantization: StageWeightQuantization::Auto,
             cache_type_k: "f16".to_string(),
             cache_type_v: "f16".to_string(),
             flash_attn_type: FlashAttentionType::Auto,
@@ -2116,6 +2118,7 @@ mod tests {
             n_gpu_layers: 0,
             mmap: None,
             mlock: false,
+            weight_quantization: StageWeightQuantization::Auto,
             cache_type_k: "f16".to_string(),
             cache_type_v: "f16".to_string(),
             flash_attn_type: skippy_protocol::FlashAttentionType::Auto,

crates/mesh-llm-host-runtime/src/inference/skippy/mod.rs

@@ -68,8 +68,8 @@ pub(crate) use stage::{
     StageInventoryRequest, StageLayerInventory, StageLoadRequest, StagePackagePrefetcher,
     StagePeerDescriptor, StagePreparationState, StagePreparationStatus,
     StagePrepareAcceptedResponse, StagePrepareRequest, StageReadyResponse, StageRuntimeState,
-    StageStatusAck, StageStatusFilter, StageStatusSnapshot, StageStopRequest, StageWireDType,
-    spawn_stage_control_loop, stage_load_timeout,
+    StageStatusAck, StageStatusFilter, StageStatusSnapshot, StageStopRequest,
+    StageWeightQuantization, StageWireDType, spawn_stage_control_loop, stage_load_timeout,
 };
 #[cfg(test)]
 pub(crate) use topology::{StageTopologyParticipant, plan_package_identity_topology};

crates/mesh-llm-host-runtime/src/inference/skippy/stage/inventory.rs

@@ -116,7 +116,7 @@ pub(super) async fn run_stage_prepare_task(
     {
         return;
     }
-    let result = prepare_stage_source(&load).await;
+    let result = prepare_stage_source(&load, Arc::clone(&cancelled)).await;
     if cancelled.load(Ordering::Acquire) {
         return;
     }
@@ -187,7 +187,23 @@ struct PrepareSourceResult {
     bytes_total: Option<u64>,
 }
 
-async fn prepare_stage_source(load: &StageLoadRequest) -> Result<PrepareSourceResult> {
+async fn prepare_stage_source(
+    load: &StageLoadRequest,
+    cancelled: Arc<AtomicBool>,
+) -> Result<PrepareSourceResult> {
+    #[cfg(not(all(feature = "mlx", target_os = "macos")))]
+    let _ = &cancelled;
+    if load.backend == "mlx" {
+        #[cfg(all(feature = "mlx", target_os = "macos"))]
+        {
+            let artifact = super::mlx::prepare_stage(load, cancelled).await?;
+            return Ok(PrepareSourceResult {
+                bytes_total: Some(artifact.report.artifact_file_bytes),
+            });
+        }
+        #[cfg(not(all(feature = "mlx", target_os = "macos")))]
+        anyhow::bail!("unsupported stage backend 'mlx' on this build");
+    }
     if load.load_mode == LoadMode::LayerPackage || is_layer_package_ref(&load.package_ref) {
         let load = load.clone();
         let package = tokio::task::spawn_blocking(move || resolve_stage_load_package(&load))