long term effort: Serve safetensors models on Metal via MLX (Apple Silicon) - #1009
michaelneale wants to merge 37 commits into
Conversation
Design proposal evaluating Apple MLX (via safemlx/safemlx-lm) as an additive second engine behind a StageEngine trait. Covers the staged execution boundary, safemlx-lm layer-split seam, JIT-quant-from-safetensors artifact strategy, selective download, platform/dependency footprint, and a solo-first phased plan gated on partial-load and boundary-fence spikes.
Standalone spikes/mlx-solo crate (own workspace, NOT in the mesh-llm workspace) that loads an HF safetensors model via safemlx-lm and generates tokens in Rust. Results (Qwen3-0.6B, Apple Silicon CPU-only, no Metal compiler on this box): - source precision: 18.1 tok/s decode, coherent — no GGUF, no ahead-of-time quant - JIT 4/8-bit affine: correct but ~0.4 tok/s (MLX quant matmul is Metal-optimized; no fast CPU kernel) -> JIT quant must be gated behind Metal/CUDA Also validates MLX C++ builds under cmake 4.4 CPU-only (the #1 feared build risk). Findings folded into the plan (status, Phase 2, risks). Full writeup and the required safemlx fork edits are in spikes/mlx-solo/FINDINGS.md.
Redo of the solo spike the goose way (metal backend, Device::Gpu) after installing the Metal toolchain (Xcode 26.6 + MetalToolchain component). Metal results (Qwen3-0.6B, Apple Silicon): - source precision bf16: 321 tok/s decode, coherent, ZERO fork patches - pre-quantized 4-bit mlx-community repo: 603 tok/s - JIT 4-bit quant on load: 604 tok/s Headline: JIT-quantize-on-load (604) == pre-quantized artifact (603), so quantizing on load is free at inference time. Supersedes the earlier CPU-only run (18/0.4 tok/s), which was an unrepresentative CPU-kernel artifact. The goose baseline (source precision) needs no fork changes. Two small safemlx-lm fixes are only needed to go beyond it (JIT quant of a tied-embedding checkpoint; loading published quant repos that omit config mode) and are upstream-PR candidates for jbg/safemlx, not mesh-llm drift. Details + repro in spikes/mlx-solo/FINDINGS.md.
…le yet - ROCm is a real but unmerged upstream MLX experiment (PR #2300), not 'no signal'; Vulkan is wishlist-only; hardware coverage has two gates (upstream mlx -> safemlx). - safemlx supply chain: published crates.io 0.4.1 collides version strings with a different, older codebase than fork HEAD 0.4.1 (851 vs 2221 lines qwen3.rs). Fork-free build ran but produced gibberish for Qwen3 + crashed on pre-quant repo; working dense-model + JIT-quant code is fork-HEAD-only. Must pin a git rev. - Adds spikes/mlx-solo-published (crates.io-only) demonstrating the breakage.
Adds crates/skippy-engine-mlx: a working MLX serving engine that serves HF safetensors models over mesh-llm's REAL OpenAI frontend (openai-frontend router_for), on Apple Silicon. Verified against Qwen3-0.6B: - GET /v1/models lists the model - POST /v1/chat/completions returns a real generation with usage - streaming returns proper SSE (role + content deltas + final finish_reason) Design: a dedicated OS worker thread owns the non-Send MLX objects (model, streams, arrays) and communicates via Send channels; this also serializes GPU access. MlxBackend implements OpenAiBackend; incremental detokenization via decode-prefix-diff. Reuses safemlx-lm (pinned fork checkout) and ports goose's generation-loop patterns rather than depending on goose-local-inference (which force-compiles a second llama.cpp and is pinned to a safemlx version proven broken for Qwen3). Standalone by design: own [workspace], not a main-workspace member, so it does not perturb the main build or CI (verified: cargo metadata on the main workspace still resolves and excludes this crate). All MLX code is gated behind both the mlx feature and target_os=macos. WIRING.md documents the promotion path into the shipped binary: git-pin safemlx, add an Mlx variant to LocalRuntimeBackendHandle, route ModelFormat::Safetensors to the MLX engine at launch, and auto-enable on macOS. Model discovery/listing already handles MLX safetensors; the missing piece was the serving engine. Refs branch micn/mlx-redux.
Pins safemlx/safemlx-lm to a public commit of jbg/safemlx (4e53c5e) instead of a path-dep on a local fork, and switches to goose-style plain source-precision serving (drops JIT quantization). Why the git-rev pin (not crates.io): safemlx's published crates (0.1.5 and 0.4.1) emit repeated-token gibberish for dense models (Qwen3 AND Llama) with this exact crate code, while the pinned upstream commit serves them correctly. Verified to be a library-version issue, not prompting: the Qwen3 chat template is confirmed applied and greedy sampling (temp=0 -> argmax) is used; swapping only the safemlx version flips output coherent<->gibberish. No private patches (plain LoadedModel::load avoids the quant-path loader quirks entirely). Will swap to a version pin once safemlx cuts a working dense-model release. Verified on Apple Silicon (Metal) over the real openai-frontend router: - Qwen3-0.6B: coherent non-stream + streaming (19 SSE chunks) - SmolLM2 (Llama arch): coherent WIRING.md updated with the git-pin rationale, the published-is-broken finding, and a Linux+NVIDIA (CUDA) future note. Crate stays a standalone workspace so the heavy MLX native build never runs in unrelated builds/CI.
|
This pull request is currently a draft. Reviews will not take place until the PR is marked as ready for review. |
|
Important Review skippedDraft detected. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Plus Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
Stacks the serve-integration on top of the standalone MLX engine crate. On a Mac, routes safetensors models to the MLX (Metal) engine over the real openai-frontend; non-macOS / no-feature builds unaffected. Done & verified locally (Apple Silicon): - skippy-engine-mlx is a workspace member (out of default-members); CI crate lists updated; xtask ci-crate-lists passes. - `mlx` feature on mesh-llm-host-runtime + mesh-llm as a macOS-target-gated optional dep; implies dynamic-native-runtime. - inference/mlx.rs (MlxModelHandle + MlxHttpHandle over openai-frontend router_for); LocalRuntimeBackendHandle::Mlx variant + arms; start_runtime_mlx_model + is_safetensors_model_path route safetensors before the GGUF path. - cargo build -p mesh-llm --features mlx -> exit 0. no-feature/default clean. fmt + clippy clean with and without the feature. Link gate resolved: static MLX + patched llama.cpp collide on gguf_get_key (MLX vendors antirez gguflib; llama.cpp exports the same). Fixed by making the `mlx` feature imply dynamic-native-runtime so llama.cpp loads as a dylib. BLOCKED: xtask release-targets forbids the published mesh-llm-host-runtime from depending on non-publishable skippy-engine-mlx (git-pinned safemlx -> publish=false; crates.io safemlx is broken for dense models). See crates/skippy-engine-mlx/SERVE_INTEGRATION_STATUS.md for options and reasoning. Not for merge as-is; captures the working integration + the publish-invariant decision for a maintainer.
i386
left a comment
i386
left a comment
There was a problem hiding this comment.
One partial-download safety issue in the new exact-range reader.
| pub fn into_bytes(mut self) -> Result<Vec<u8>> { | ||
| let mut bytes = Vec::with_capacity(usize::try_from(self.expected_bytes)?); | ||
| self.response | ||
| .read_to_end(&mut bytes) | ||
| .context("read HTTP range response")?; | ||
| ensure!( | ||
| bytes.len() as u64 == self.expected_bytes, | ||
| "HTTP range returned {} bytes, expected {}", | ||
| bytes.len(), | ||
| self.expected_bytes | ||
| ); | ||
| Ok(bytes) | ||
| } | ||
|
|
||
| pub fn copy_to(mut self, writer: &mut impl Write) -> Result<u64> { | ||
| let written = | ||
| std::io::copy(&mut self.response, writer).context("stream HTTP tensor range")?; |
There was a problem hiding this comment.
These paths validate the byte count only after fully consuming the response. A server can return 206 with a matching Content-Range and no Content-Length, then stream far more than requested; read_to_end/io::copy will allocate or write all of it before the equality check, defeating the full-shard-download guard and potentially exhausting memory or disk. Please cap both readers at expected_bytes + 1 and reject as soon as an extra byte is observed.
i386
left a comment
i386
left a comment
There was a problem hiding this comment.
Review of current head 75fb8b24.
One blocking safety issue remains in the new exact-range download path: ExactRangeResponse::into_bytes() and copy_to() consume the response without a limit and only compare the byte count afterward. A 206 response with the requested Content-Range but no Content-Length can stream arbitrarily more data, causing unbounded memory growth or disk writes before rejection. Cap both readers at expected_bytes + 1 and reject immediately when the extra byte is observed. Detailed inline comment: #1009 (comment).
The staged-protocol changes I inspected are additive: older peers can ignore the new fields, and missing weight-quantization values default to Auto. The whole-model MLX and explicit split evidence is useful, but this remains a very large 93-file integration (engine abstraction, partial HTTP/materialization, host wiring, docs, and three spikes) and is conflicted with current main; splitting/removing the spike work would make the production path substantially easier to review and maintain.
michaelneale
changed the title
2 participants
What this draft now proves
This branch has a connected Apple-Silicon path from Hugging Face SafeTensors
to both ordinary single-node serving and mesh-managed explicit split serving.
Off-the-shelf whole-model MLX serving
The normal shipped command now serves both canonical dense checkpoints and
unchanged MLX-LM repositories through mesh-llm's shared OpenAI frontend:
The second command was exercised at HF revision
73e3e38d981303bc594367cd910ea6eb48349da8. It loaded the cached model inabout 1.2 seconds, appeared in
/v1/models, answered23 * 17with391,completed SSE with
[DONE], and completed a Goose OpenAI-provider request.The same path now reports the repository as 0.6B rather than mistaking its
4bitsuffix for a 4B parameter label.Published MLX-LM affine checkpoints commonly omit
quantization.mode. Thepinned safemlx correction treats that omission as the standard
affinedefault. The one-commit fork is proposed upstream as
jbg/safemlx#2.
For eligible unquantized dense checkpoints, automatic mode applies
affine-4/group-64 while loading. Inkling, Nemotron-H, GGUF files, and sources
already declaring quantization/compression preserve their native
representation. Qwen3-0.6B's exact benign tied-
lm_head.weightrejection usesa narrow native-load retry; all unrelated strict-load failures remain
fail-closed.
Mesh-managed range-only split serving
With explicit
--split, the coordinator:checkpoint tensor payloads;
backend-mlxstage capability;embedding/readout boundary costs;
activation wire.
A real two-host SmolLM2 run used independent identities and ordinary stage
control:
0..2929..30The remote did not store the complete roughly 269 MB source checkpoint.
Non-streaming chat traversed both stages and returned coherent text; SSE
completed with
[DONE].This is deliberately different from whole-model prequantized loading:
prequantized
mlx-communityrepositories work unchanged on one node, whiledistributed
--splitcurrently starts from canonical dense SafeTensors andderives stage-local quantized artifacts. It does not yet slice a published
prequantized repository across nodes.
Lifecycle and compatibility hardening
for split tensor ranges.
downstream.
incremental decoding with final suffix reconciliation.
MLX build.
GGUF/Skippy loads still require it.
Frontier evidence retained
109.84 GiB across 942.99 GiB of source shards, avoiding 833.15 GiB.
affine expert banks, strict-loaded, executed, and validated through the
engine-neutral stage contract and binary wire.
cancellation, and portable
mlx.metallibevidence are documented incrates/skippy-engine-mlx/STAGED_EXECUTION.md.Validation for the latest checkpoint
just mlx-build: pass; emits and signsmesh-llmplus siblingmlx.metallib.Goose, and clean shutdown: pass; independently reproduced by the MLX expert.
skippy-engine-mlx --features mlx --lib: 42 passed.mesh-llm-host-runtime --lib: 1,602 passed, 8 ignored.safemlx-lm --lib: 119 passed, 66 ignored.CI-crate-list consistency, and publish dry-run checks: pass.
are now stated below rather than hidden.
Deliberately incomplete
--split; automatic splitselection is not implemented.
model_type=llama, one greedylane. Inkling partial stages and complete Nemotron hybrid topology remain.
metadata, not yet a user-facing hardware/quality policy.
tools, sampling controls, logprobs, parallel lanes, or multimodal requests.
patches do not propagate, and registry
safemlx-lm 0.4.1lacks APIs used bythis engine. Standalone crates.io
skippy-engine-mlx --features mlxsupporttherefore remains unusable until compatible safemlx releases land.