From 0a0bf81ea21ac01723d2962efefc4c1e29fcc3b3 Mon Sep 17 00:00:00 2001 From: michaelneale Date: Thu, 16 Jul 2026 12:16:46 +1000 Subject: [PATCH 01/37] =?UTF-8?q?docs:=20MLX=20as=20a=20Skippy=20stage=20e?= =?UTF-8?q?ngine=20=E2=80=94=20deep=20dive=20and=20plan?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- docs/design/MLX_STAGE_ENGINE_PLAN.md | 691 +++++++++++++++++++++++++++ 1 file changed, 691 insertions(+) create mode 100644 docs/design/MLX_STAGE_ENGINE_PLAN.md diff --git a/docs/design/MLX_STAGE_ENGINE_PLAN.md b/docs/design/MLX_STAGE_ENGINE_PLAN.md new file mode 100644 index 0000000000..b9267f685a --- /dev/null +++ b/docs/design/MLX_STAGE_ENGINE_PLAN.md @@ -0,0 +1,691 @@ +# MLX as a Skippy Stage Engine — Deep Dive and Plan + +## Status: exploratory design proposal + +This document evaluates using Apple **MLX** (via the Rust `safemlx` / `safemlx-lm` +crates) as an alternative inference engine behind Skippy's staged execution +runtime, and proposes a phased plan. + +It combines a read of the current Skippy code (`skippy-ffi`, `skippy-runtime`, +`skippy-server`, `skippy-topology`), a read of the `safemlx` fork +(`../safemlx`), a read of goose's MLX backend (`../goose/crates/goose-local-inference`), +and a second-opinion review from an external model grounded against live +MLX/mlx-lm/safemlx documentation. + +--- + +## 1. Bottom line + +MLX is a **credible second engine** for Skippy, and `safemlx-lm` is a +surprisingly good fit because it implements each model in **pure Rust as an +explicit `embed → layers[..] → norm → lm_head` loop with a per-layer KV cache**. +That is exactly the seam Skippy needs for layer-range stage splitting, and it is +Rust-facing rather than buried in C++. + +It is **not** a drop-in replacement for the patched llama.cpp C ABI. The engine +boundary Skippy actually depends on is much larger than "generate tokens": it is +a **staged execution contract** (activation frames in/out, KV page +export/import, layer-range partial load, chunked prefill, single-token decode, +batched verify, trim/checkpoint, tokenizer/chat helpers). + +**Two distinct reasons to adopt MLX — keep them separate:** + +1. **Workflow / artifact win (biggest near-term value):** MLX loads HF + **safetensors directly** and can **JIT-quantize on load**, so we can serve + any supported model at a chosen bit-width *immediately* — no waiting for a + published GGUF and no pre-run quant pipeline. This is mostly independent of + the hard split work and pays off first in **solo serving** (§5.3). +2. **Apple-Silicon compute in a chain:** MLX runs straight to Metal and can add + Apple-Silicon nodes to a staged split — but this is the harder, later payoff, + gated on partial-load and boundary-fence behaviour. + +**Recommendation:** introduce a Rust `StageEngine` trait, keep the existing C ABI +as the `LlamaStageEngine` adapter, and add an Apple-Silicon-gated +`MlxStageEngine`. Do **not** extend the llama.cpp C ABI to host MLX, and do +**not** invent a separate MLX network protocol. **Lead with solo MLX + +JIT-quant serving** (immediate workflow win, minimal distributed work), then gate +the split work behind two go/no-go spikes: **partial layer loading** and +**per-token boundary fence latency**. + +--- + +## 2. What the Skippy "engine" boundary actually is + +Skippy's engine is not a token generator; it is a staged runtime. The contract +lives in `crates/skippy-ffi/src/lib.rs` (raw ABI) and is wrapped safely in +`crates/skippy-runtime/src/lib.rs`. The essential surface a new engine must +satisfy: + +**Stage-aware load** (`RuntimeConfig`, `skippy-runtime/src/lib.rs:840`): +- `stage_index`, `layer_start`, `layer_end` — this stage owns a contiguous + layer range only. +- `include_embeddings`, `include_output` — whether this stage owns the embedding + table and/or the final norm + lm_head. +- `filter_tensors_on_load` — the intent that a stage should load **only** its + tensors, not the whole model. +- backend device selection, KV cache dtype, ctx size, batch/ubatch, lanes. + +**Activation frame I/O** — the wire contract between stages +(`ActivationFrame` / `ActivationDesc`, `skippy-runtime/src/lib.rs:1046` and +`:1109`): + +``` +ActivationDesc { + version, dtype (F32|F16|BF16), layout (TokenMajor|Opaque), + producer_stage_index, layer_start, layer_end, + token_count, sequence_count, payload_bytes, flags +} +ActivationFrame { desc, payload: Vec } +``` + +- Stage 0 takes token IDs, runs its layers, emits an activation frame. +- Middle stages import a frame, run their layers, emit a new frame. +- Final stage runs last layers + readout, samples, returns the **predicted token + directly** to stage 0 (generation-3 protocol, see `skippy-server/README.md`). + +**Execution calls** (`skippy-runtime/src/lib.rs`): +- `prefill_chunk_frame*` (`:2998`), `decode_step_frame_sampled*` (`:3236`), + `verify_tokens_frame*` (`:3557`), `copy_output_activation_frame` (`:3652`). +- Sampled variants carry `SamplingConfig` (penalties, logit bias, grammar). + +**KV / state movement** (`skippy-runtime/src/lib.rs`): +- `export_kv_page` / `import_kv_page` (`:3866`, `:3956`) with + `RuntimeKvPageDesc` (`:1114`) — k/v dtype, row bytes, token range, layer range. +- `export_state` / `import_state` (`:3729`), full-state and recurrent-state + variants, `save_prefix` / `restore_prefix`, `trim_session`, checkpoint/restore. + +**Tokenizer / chat / introspection**: +- tokenize/detokenize, EOG check, chat-template apply (incl. JSON tools path), + chat-response parse, model-info tensor enumeration, GGUF slice writing. + +**ABI is versioned and feature-probed** (`skippy-ffi/src/lib.rs:1`): ABI +`0.1.30`, with a feature bitmask (`RUNTIME_SLICE`, `LAYER_PACKAGE`, +`ACTIVATION_FRAME`, `BATCH_VERIFY_FRAME`, `SESSION_CHECKPOINT`, +`NATIVE_MTP_N1`, …). This is the model for how MLX capabilities should be +advertised: **probed, not assumed**. + +**Key structural gap:** there is **no Rust `trait`** abstracting this today. +`skippy-server` binds the concrete `StageModel` / `StageSession` FFI structs +directly (`crates/skippy-server/src/frontend.rs:46`, `runtime_state.rs:26`). +Introducing that trait is the enabling refactor for any second engine. + +--- + +## 3. What MLX / safemlx actually gives us (evidence) + +### 3.1 The layer-split seam already exists in `safemlx-lm` + +Every model in `../safemlx/safemlx-lm/src/models/*.rs` is a Rust module with the +transformer decomposed into public fields and an explicit forward loop. From +`qwen3.rs`: + +```rust +pub struct Qwen3Model { + pub embed_tokens: MaybeQuantized, + pub layers: Vec, // per-layer blocks + pub norm: nn::RmsNorm, +} +// forward: +let mut h = self.embed_tokens.forward(inputs, stream)?; +for (layer, c) in self.layers.iter_mut().zip(cache.iter_mut()) { + h = layer.forward(/* x=h, mask, cache=c */, stream)?; +} +self.norm.forward(&h, stream) // then lm_head at the Model level +``` + +`pub embed_tokens` / `pub layers` / `pub norm` / `pub lm_head` are exposed across +`qwen3`, `llama`, `gpt_oss`, `gemma4`, `lfm2`, `nemotron_h`, `qwen3_5_moe`, etc. +The per-layer KV cache is a `Vec>` threaded through the loop +(`safemlx-lm/src/lib.rs`, `cache.rs`), and blocks accept an explicit mutable +cache slot. This means running **only** `layers[start..end]` and resuming from an +imported hidden state is mechanically straightforward — no C++ surgery. + +### 3.2 Activation observe / intervene hooks + +`safemlx-lm/src/inspection.rs` defines `ActivationObserver` with +`observe(name, &Array)` and `intervene(name, &Array) -> Option` at block +boundaries ([inspection API](https://docs.rs/safemlx-lm/latest/safemlx_lm/inspection/)). +This is useful plumbing/debugging, **but it is not a stage ABI** — it is +name-based and per-tensor. For production we want a first-class +`forward_range()` / `resume_from_hidden()` path, not a reliance on observer +names. + +### 3.3 KV cache primitives + +`safemlx-lm/src/cache.rs` defines `KeyValueCache` (offset, max_size, +`update_and_fetch`), with `ConcatKeyValueCache`, `SlidingKeyValueCache`, +quantized variants, and `truncate(len)`. Cache state is MLX arrays in unified +memory. This gives us the raw material for export/import/trim — but the state +layout is MLX/model-specific and **not** interchangeable with llama.cpp's ggml +page format. + +### 3.4 Loading, quant, and formats + +`safemlx-lm` loads Hugging Face-style dirs (`config.json` + `tokenizer.json` + +safetensors) and also **GGUF** (`Array::load_gguf_with_metadata`, +`models/mod.rs:1123`), with a **strict loader** (`weights.rs:155`) that errors on +missing/unused tensors unless explicitly allowed. MLX quant is affine +packed-weight-in-safetensors; there is load-time quantization from unquantized +F32/F16/BF16. + +### 3.5 Distributed primitives exist but are unwrapped + +`../safemlx/safemlx-sys/src/mlx-c/mlx/c/distributed.h` binds +`all_gather / all_sum / all_min / all_max / send / recv / sum_scatter` and +distributed groups. **The safe `safemlx` layer does not wrap them yet.** This +matters for tensor-parallel within a node, but Skippy's cross-machine model is +its own QUIC activation-frame protocol — we do **not** want to depend on MLX +distributed for the mesh boundary. + +### 3.6 goose already ships an MLX backend — what we can reuse + +`../goose/crates/goose-local-inference/src/mlx.rs` (1044 lines) uses `safemlx-lm` +(`LoadedModel::load`, `generate`, Gemma4 MTP draft/speculative) as a **single +node** backend behind goose's own `LocalInferenceBackend` trait, feature-gated +`mlx` on macOS. It proves safemlx-lm is production-usable for generation. + +**"Right to Metal" is accurate.** goose's MLX path pulls +`safemlx { features = ["accelerate", "metal", "safetensors"] }` and runs on +`Device::new(DeviceType::Gpu, 0)` (`mlx.rs:61,94`) — i.e. +`safemlx-sys → mlx-c → MLX → Metal`, with no llama.cpp/ggml in between. (goose +*also* keeps a separate `llama-cpp-2` Metal path; the two are independent +backends.) + +**What is actually reusable, and how:** + +| Asset | Reuse verdict | +| --- | --- | +| `safemlx` / `safemlx-lm` crates | **Reuse directly** — this is the real shared dependency. Both goose and Skippy just depend on the published crates. No goose code involved. | +| goose's `mlx.rs` generation flow (prompt build, sampling, MTP draft/verify, streaming, stop tokens, thinking-filter) | **Reuse as a reference template, port not lift.** It is the best worked example of driving safemlx-lm for chat + speculative, but it is coupled to goose types. | +| goose's `LocalInferenceBackend` trait (`backend.rs`, 50 lines) | **Reference only.** It is goose's shape (whole-model `load_model` + `generate`), not Skippy's staged contract. Skippy needs its own `StageEngine` (§6). | +| HF download + shard/registry (`hf_models.rs`, `goose-download-manager`) | **Reference / optional adapter.** `goose-download-manager` is cleanly separable (deps are just `reqwest`/`tokio`), but Skippy already has `model-hf` / `model-artifact`; prefer extending those. | + +**Coupling is the reason it is port-not-lift.** goose's `mlx.rs` and `backend.rs` +depend on `goose_provider_types` (`Message`, `MessageContent`, `ProviderError`, +`ProviderUsage`/`Usage`, `DraftStats`), `rmcp::model::Tool`/`Role`, and +`local_model_registry::ModelSettings`. Skippy speaks `ActivationFrame`, token +IDs, `SamplingConfig`, and its OpenAI frontend types instead. So the *algorithms* +(how to prefill, sample, run MTP draft/verify against safemlx-lm) transfer +directly; the *types and trait* do not. + +**License:** goose is **Apache-2.0**, so porting code with attribution is fine. + +**Bottom line:** the highest-leverage reuse is simply **sharing the `safemlx-lm` +crate** (and coordinating on/​contributing the stage-aware `forward_range` / +partial-load additions the fork needs — see Phase 3), plus using goose's `mlx.rs` +as the reference implementation for the single-stage generation path in Phase 2. +It is whole-model and single-stage, so it is not a template for the staged/ +KV-page work. + +--- + +## 4. Fit analysis and sharp edges + +The happy path fits: + +``` +tokens or imported hidden state + → optional embedding (stage 0 only) + → layers[start..end] with per-layer cache + → hidden state frame (middle/non-final) + or → final norm + lm_head → sample → token (final) +``` + +Sharp edges, in rough priority order: + +1. **Partial execution must mean partial loading.** Building a whole + `LoadedModel` and skipping layers can still materialize all weights. Skippy's + entire value proposition is fitting a big model across small machines, so a + stage must load only its layer range (plus embeddings/readout when it owns + them). This requires a **stage-aware loader** in `safemlx-lm` that + instantiates `layers` for `[start..end]` and only reads matching safetensors + shards. **This is the #1 go/no-go item.** + +2. **Every model family must define its exact residual-stream boundary.** + Embedding scale, final norm, tied vs untied output, RoPE position accounting + across a stage cut, attention mask construction, and any per-layer-type + sideband cannot be inferred generically. Each family is a separate + certification (mirrors `skippy-topology` family capability records, + `crates/skippy-topology/src/lib.rs:182`). + +3. **Hybrid / recurrent models need more than hidden states.** Mamba/RWKV/gated + DeltaNet-style layers (`nemotron_h`, `qwen3_5_moe` / `qwen3_next`, `lfm2` in + safemlx-lm) carry recurrent state, not page-addressable KV. Skippy already + has `export_recurrent_state`; MLX would need the analogous opaque sideband, + or those families are restricted to non-split. + +4. **The MLX model matrix ≠ the safemlx-lm matrix ≠ the Skippy-certified + matrix.** safemlx-lm implements models individually and is young. Start with + **dense Llama / Qwen**; do not promise arbitrary model coverage. + +5. **The boundary is more than execution.** Sampling, batched verify, trim, + checkpoint, tokenizer/chat, and state movement are all in the ABI today + (`crates/skippy-ffi/README.md`). The trait must cover them (some can be + engine-agnostic and moved above the engine). + +--- + +## 5. Hard problems (with concrete approaches) + +### 5.1 Lazy evaluation — the network boundary is an eval boundary + +MLX is lazy and uses unified memory. A stage boundary forces materialization. +The per-token non-final-stage sequence must be: + +1. Run local layer range (lazy). +2. Cast to negotiated wire dtype (`ActivationDType::F16` first). +3. Make contiguous + token-major (`ActivationLayout::TokenMajor`). +4. **Evaluate the outgoing array AND all updated cache arrays together.** +5. Get a host-readable slice, serialize into `ActivationFrame.payload`, send. + +In `safemlx`: `Array::evaluated()` materializes and `EvaluatedArray::as_slice()` +gives host access (host access / save also forces eval) +([safemlx lazy-eval source](https://docs.rs/safemlx/latest/src/safemlx/lib.rs.html), +[MLX lazy-evaluation guide](https://ml-explore.github.io/mlx/build/html/usage/lazy_evaluation.html)). + +Critical details: +- **Evaluate cache state even when it does not feed the outgoing hidden state**, + or lazy cache graphs grow unbounded across decode steps. +- Unified memory removes the explicit GPU→CPU copy but **not** GPU completion, + sync, layout conversion, or the QUIC copy. +- Final stage: evaluate the sampled token + cache; never materialize a + hidden-state frame. +- Evaluate params once at warmup. +- Consider separate **compiled** paths for fixed-shape decode vs bucketed + prefill (`safemlx` [`compile_with_state`](https://docs.rs/safemlx/latest/safemlx/transforms/compile/)), + watching recompilation from shape changes. + +**The benchmark that matters** is not MLX layer time; it is +`last layer → cast/contiguous → eval fence → host view → QUIC write`, per token, +at realistic hidden widths. **This is go/no-go item #2.** + +### 5.2 KV cache export/import/trim + +MLX cache arrays make movement possible, but mlx-lm/safemlx cache state is not +llama.cpp page-shaped (mlx-lm caches expose array state, metadata, and tail +trimming, and prompt caches serialize as safetensors — +[mlx-lm cache.py](https://raw.githubusercontent.com/ml-explore/mlx-lm/main/mlx_lm/models/cache.py)). +Do **not** reuse `RuntimeKvPageDesc` for MLX; keep that in the llama adapter. +Define an engine-general **cache codec** with a versioned descriptor: + +``` +engine + model digest, architecture revision +layer range, token range + absolute position +cache kind (concat | sliding/rotating | quantized | recurrent) +segments: { role, layer, dtype, shape, strides/layout, payload } +cache-specific metadata (rotating offset, quant scales/biases) +``` + +- Export: slice token range, order rotating caches temporally, contiguous, + evaluate, serialize. +- Import: validate engine/model/range/layout, rebuild arrays, restore offsets. +- Trim: offset change is cheap; reclaiming/compacting memory needs slicing. +- **Quantized KV** needs packed values + scales/biases, not just row sizes. +- **Do not promise KV interop between llama.cpp and MLX.** Import requires same + engine + model digest + quant + arch revision + cache policy. + +### 5.3 Artifact strategy: JIT safetensors vs pre-quantized layer packages + +This is arguably the **strongest first reason to adopt MLX**, and it is largely +independent of the hard split work. The benefit is really **two separate +things**, and they behave very differently for solo vs split serving: + +- **(A) Source freedom** — load HF **safetensors** directly instead of waiting + for someone to publish a GGUF (or running our own GGUF quant pipeline first). +- **(B) JIT quantization** — quantize at load time (`with_quantization(Q4/…)`) + instead of ahead of time. + +(A) applies equally to solo and splits (it's just "which file do we load"). +(B) is where solo and splits diverge sharply. + +**What safemlx-lm actually supports (confirmed):** +- `ModelLoadOptions::with_quantization(...)` quantizes eligible dense weights + **one tensor at a time** on load; checkpoints already carrying matching quant + metadata load directly without requantizing + (`../safemlx/safemlx-lm/src/models/mod.rs:137`). +- Sharded safetensors are understood via `model.safetensors.index.json`'s + `weight_map` (tensor name → shard file), so the loader knows which shard holds + `model.layers..*` (`../safemlx/safemlx-lm/src/weights.rs:463`). MLX quant is + affine packed-weights-in-safetensors, matching mlx-lm's converter + ([mlx-lm quantized loading/conversion](https://raw.githubusercontent.com/ml-explore/mlx-lm/main/mlx_lm/utils.py)). + +#### Solo serving → pure win, lands first + +Download BF16/FP16 safetensors → `with_quantization(Q4)` → serve. No wait for a +published GGUF, no pre-run of the quant pipeline. Any dense model safemlx-lm +supports is instantly servable at a chosen bit-width. **goose already does +exactly this** single-node (`../goose/crates/goose-local-inference/src/mlx.rs`). +Near-zero new distributed work; this is the cheapest, highest-value first step. + +#### Splits → works, but with four real conditions + +The premise of a split is that *no node holds the whole model*, which stresses +JIT quant: + +1. **Stage-aware partial load + quant.** Today safemlx-lm builds + `0..num_hidden_layers` and the strict loader expects all params, so JIT quant + is a *whole-model* op. A split node must instantiate only `layers[start..end]`, + read only the shards overlapping its range, and quantize only those tensors. + This is exactly go/no-go **Spike 1**, now with a quant step folded in. +2. **Selective shard download.** The `weight_map` makes fetching only the shards + for a layer range *possible*, but safetensors shards bundle several + consecutive layers, so nodes over-fetch at range boundaries — coarser than + GGUF layer-package parts, which Skippy slices exactly. +3. **Deterministic cross-stage quant.** Every stage must quantize *identically* + (same algo / group-size / bits / tie handling) or the split model drifts + numerically from the solo model. Affine quant is deterministic given its + params, so this is achievable — but the params must be pinned and folded into + family/topology certification. +4. **Cache the quantized slice.** Re-quantizing on every launch/replan across N + nodes is wasteful. Skippy's identity-bound materialized cache + (`crates/skippy-runtime/src/package/materialized_cache.rs`, keyed by + `model_id / topology_id / stage_id / layer_start / layer_end`) is the natural + home: first launch pays the JIT cost → materialize a per-stage quantized + artifact → reuse thereafter. + +#### The tension worth naming + +Skippy's existing chain (`skippy-quantize`, layer-package repos, BF16→GGUF) is +built around **pre-quantized, exactly-sliced GGUF parts** *specifically so split +nodes download their slice and never quantize at runtime*. +JIT-quant-from-safetensors trades that for flexible source + runtime quant + +coarser slicing. For solo there is no tension; for splits it is a genuine, but +acceptable, tradeoff. The two paths should **coexist**: + +- **JIT safetensors = fast coverage path** — try any supported HF model on the + mesh immediately, no publish step. +- **Pre-quantized layer packages = optimized path** — for models served + seriously (exact slices, no runtime quant, tailored partial download). + +#### One artifact identity, two physical encodings + +Use **one logical package identity, not one physical weight encoding**: + +``` +model identity + source revision + tokenizer/config/chat metadata + topology +variants: + llama-gguf: GGUF parts + quant (existing skippy-model-package path) + mlx-jit: HF safetensors + quant spec (quantize on load; cache the result) + mlx-packaged: pre-quantized MLX stage shards + index (optimized split path) +``` + +- Make **BF16/FP16 HF safetensors the canonical source**; derive all variants + reproducibly (this repo already has `skippy-quantize`, `model-hf`, + `model-package`, and BF16 GGUF conversion skills). +- **Never** transcode an already-quantized GGUF → MLX quant (dequant/requant + loses quality and still rebuilds arch metadata). +- Nodes download only their selected engine/stage variant, so catalog + duplication need not become per-node duplication. +- For true partial download in the packaged path, stage-specific MLX + safetensors shard/index generation is needed (parallel to today's GGUF slice + writing). + +### 5.4 Cross-machine execution + +Treat **MLX as the local compute engine and Skippy as the distributed runtime.** +mlx-lm's `pipeline()` / `sharded_load()` / `send`/`recv` + `all_gather` is useful +**reference**, but it uses static ranks and MLX collectives — it is not Skippy's +QUIC activation-frame protocol with independent stage lifecycle, capability +negotiation, and direct final-token return. Keep Skippy's transport; use MLX only +for compute +([MLX distributed docs](https://ml-explore.github.io/mlx/build/html/usage/distributed.html), +[mlx-lm utils](https://raw.githubusercontent.com/ml-explore/mlx-lm/main/mlx_lm/utils.py), +[pipeline mixin](https://raw.githubusercontent.com/ml-explore/mlx-lm/main/mlx_lm/models/pipeline.py)). + +For normal Ethernet/Wi-Fi: an 8192-wide F16 activation is ~16 KiB/token/boundary +(decode is latency-bound, not bandwidth-bound); prefill 512×8192×F16 is ~8 +MiB/boundary (bandwidth matters). Pipeline parallelism does **not** speed up +single-sequence decode — only concurrent sessions / speculative spans keep stages +busy. Skippy's topology wire sizing (`crates/skippy-topology/src/lib.rs:1415`) +already models F16 = `2 × hidden_width`. + +### 5.5 Platform and dependency footprint + +"To the metal like goose" and "lean dep" are both achievable, but the second has +a real catch: **MLX is lean at runtime and heavy at build time.** + +**Runtime footprint — genuinely lean.** goose's path is +`safemlx → safemlx-sys → mlx-c → MLX → Metal`, statically linked +(`safemlx-sys/build.rs` sets `BUILD_SHARED_LIBS=OFF`), running on +`Device::new(DeviceType::Gpu, 0)`. The vendored `mlx-c` is a ~1 MB C shim; there +is no runtime service or subprocess. Mesh can depend on the same crates for the +identical to-the-metal path — the metal-ness lives in `safemlx-sys`, nothing +goose-specific. + +**Build footprint — a second heavy native lane.** `safemlx-sys/build.rs` drives +**CMake**, and the bundled `CMakeLists.txt` uses `FetchContent` to **git-clone +the full MLX C++ core from `github.com/ml-explore/mlx.git` and compile it**. +Building therefore needs CMake ≥3.25, a C++20 compiler, network to fetch MLX, +and — for Metal — Apple's `metal` shader compiler (`xcrun -find metal`, producing +`mlx.metallib`). This sits alongside the existing llama.cpp patch-queue build and +becomes another native runtime artifact under the +`MESH_LLM_DYNAMIC_NATIVE_RUNTIME` packaging model. + +**Keeping it lean = isolation, not intrinsic lightness.** Put the engine in its +own crate (`skippy-engine-mlx`) gated by **both** a cargo `feature = "mlx"` +**and** `cfg(target)` (Apple Silicon, optionally Linux/CUDA). Then default, +Linux-ROCm, Vulkan, and Windows builds never pull MLX or run its CMake — exactly +how goose gates it. Lean by construction, for the platforms that don't use it. + +**Support matrix — broader than macOS, but not the full llama.cpp matrix** +(from `safemlx-sys/build.rs` + `safemlx-sys/README.md`): + +| Target | MLX support | +| --- | --- | +| macOS Apple Silicon | ✅ Metal + Accelerate | +| iOS / tvOS / visionOS | ✅ Metal | +| Linux x86_64 / aarch64 | ✅ CPU | +| Linux + NVIDIA | ✅ CUDA (the `cuda`/`nccl` features **panic** on non-Linux) | +| Linux + AMD (ROCm) | ❌ | +| Vulkan (any) | ❌ | +| Windows | ❌ | + +**Coverage is expanding, and safemlx tracks it fast.** `jbg/safemlx` is very +active (103 commits, latest 2026-07-15) and pins a recent MLX core (`v0.32.0`). +It wires in new backends quickly: the `Add CUDA support` commit landed a full +`build.rs` + CMake patch + Linux CI + `cuda.rs` module + smoke test in one go, +and there is `if(WIN32)` DLL-export scaffolding in the vendored `mlx-c`. So the +matrix above is a **snapshot, not a ceiling**. + +**But the gaps are gated by MLX upstream, not by safemlx.** safemlx can only +expose backends `ml-explore/mlx` itself provides. MLX's backend line is +**CPU + Metal + CUDA** — which is exactly why CUDA appeared here. There is **no +ROCm or Vulkan backend in MLX upstream**, and no in-repo signal that safemlx adds +them independently. So: Windows is the most plausible next addition +(scaffolding + working Linux/CUDA); ROCm/Vulkan depend on an upstream decision +with no current signal either way. + +**Strategic consequence (unchanged).** **Today**, the ROCm / Vulkan / Windows +gaps mean **MLX cannot be Skippy's sole engine** — which reinforces (not changes) +the plan: MLX is an **additive, feature+cfg-gated second engine**, +Apple-Silicon-first (with Linux/CUDA as a real second target), while llama.cpp +stays the cross-platform default. And the most durable reason to keep llama.cpp +is **not** platform coverage (which may well close as MLX upstream grows) but its +GGUF/imatrix quant maturity and the existing patch-queue investment — those would +argue against removing it even if MLX's backend matrix later caught up. + +--- + +## 6. Recommended architecture + +**Option (a): a second implementation behind a Rust `StageEngine` trait.** + +``` +Skippy protocol / skippy-server + └── StageEngine (new trait, engine-agnostic descriptors + byte buffers) + ├── LlamaStageEngine → existing skippy-ffi C ABI (unchanged) + └── MlxStageEngine → safemlx / safemlx-lm (Apple-Silicon-gated) +``` + +Trait covers: capability discovery + model inspection; stage-aware open/load; +session lifecycle; prefill / decode / batched verify; activation import/export; +trim / checkpoint / reset; opaque-or-segmented state export/import; final-stage +logits/sampling; tokenizer/chat (where not yet lifted above the engine). Backend +arrays and native handles stay private; the trait exchanges **Skippy-owned +descriptors and `Vec` payloads**. + +Rejected alternatives: +- **Extend the llama.cpp C ABI for MLX** — no. It embeds GGUF/ggml dtype and + llama loading concepts; MLX is already Rust-facing. This would degrade a good + native adapter into a lowest-common-denominator API. +- **Separate MLX server protocol** — no, initially. It duplicates lifecycle, + networking, and compatibility. If Metal/MLX crash isolation later becomes + necessary, add an **optional subprocess** implementation behind the *same* + `StageEngine` trait, reusing the existing Skippy stage protocol — not a new + public surface. + +Crate shape (fits the repo's semantic-ownership rules): +- `skippy-engine` (new): the `StageEngine` trait + shared descriptors + (activation frame, cache codec, capability probe). Engine-neutral. +- `skippy-runtime` becomes / provides `LlamaStageEngine` implementing the trait. +- `skippy-engine-mlx` (new, `cfg(all(target_os="macos", target_arch="aarch64"))`, + feature `mlx`): `MlxStageEngine` over `safemlx`/`safemlx-lm`. +- `skippy-server` depends on `dyn StageEngine`, not concrete `StageModel`. + +Protocol compatibility: MLX support is **additive** — a new engine capability +advertised via the existing feature-probe + gossip capability mechanism, with +llama.cpp remaining the default. No gossip/stream/ABI break. A mixed-engine chain +(llama stage ↔ MLX stage) must be a **separately certified** capability with +verified residual boundary, RoPE convention, activation dtype, and model +revision — default to **engine-homogeneous chains** first. + +--- + +## 7. Phased plan + +**Phase 0 — Spikes (go/no-go, no product wiring).** Standalone binaries in +`../safemlx` or a throwaway crate. See §8. Nothing merges to Skippy until Spike 1 +(partial load) and Spike 2 (boundary fence) pass. + +**Phase 1 — Introduce `StageEngine` trait (llama only).** Pure refactor: define +the trait in a new `skippy-engine` crate, implement it for the existing +`skippy-runtime` FFI, and switch `skippy-server` to `dyn StageEngine`. No +behavior change; ship this independently of MLX. Validate with existing +`skippy-correctness` and `mic-lab` runs. + +**Phase 2 — Solo MLX serving + JIT quant (the workflow win; lead here).** +`MlxStageEngine` as a single-stage/whole-model engine: open/load, session, +prefill, decode-sampled, tokenizer/chat, final-stage sampling — plus the +**source-freedom + JIT-quant** path (§5.3): download HF safetensors, quantize on +load at a chosen bit-width, serve. Port goose's `mlx.rs` generation flow (§3.6) +rather than lifting it. Wire behind `--serving-backend mlx` (parallel to the +existing skippy backend selector in `docs/SKIPPY.md`). Validate against +`skippy-correctness` vs llama.cpp logits for the same model. This delivers the +"serve any supported model instantly, no wait for quant" benefit with minimal +new distributed work, and de-risks the engine before any split work. + +**Phase 3 — Stage-aware partial load + activation frames.** Add `forward_range` +/ `resume_from_hidden` and the stage-aware loader to `safemlx-lm` (upstream to +the fork). Implement `prefill_chunk_frame` / `decode_step_frame` / +`copy_output_activation_frame` producing Skippy `ActivationFrame`s. Two-stage +single-machine parity first, then two Macs over the real network. + +**Phase 4 — KV/state codec + verify + trim/checkpoint.** Implement the +engine-general cache codec (§5.2), `verify_tokens_frame` for speculative decode, +trim/checkpoint/reset. Add speculative (safemlx-lm already has Gemma4 MTP draft +as a reference). + +**Phase 5 — Artifact/packaging + certification.** MLX variant in the model +package (§5.3), stage-shard partial download, per-family/quant certification into +`skippy-topology` capability records and `docs/skippy/` family docs. Mixed-engine +chain certification only if warranted. + +**Phase 6 — Promotion.** Only after correctness + performance parity on the +Apple-Silicon target does MLX become a default-selectable engine for +Apple-Silicon nodes. llama.cpp remains the cross-platform default. + +--- + +## 8. Spike gates (go/no-go before Phase 3) + +1. **Partial-loading proof (GO/NO-GO):** load only `layers[start..end]` (+ + embeddings/readout when owned) for a dense Qwen/Llama; confirm **peak RSS + contains only the selected range**. If a stage can't avoid loading the whole + model, the split story is dead for MLX. +2. **Boundary latency breakdown (GO/NO-GO):** measure layer compute, cast, + contiguous, **eval fence**, host readback, serialize, and receive-reconstruct + **independently**, at hidden widths 4096/8192/16384 and token counts + 1/32/512. Decode is single-sequence latency-bound; prove the fence doesn't + dominate. +3. **Two-stage dense parity:** Qwen/Llama, multiple split points, F32 + F16 + activations, chunked prefill, 128-token decode, compare logits to llama.cpp. +4. **Real network run:** two Macs over Wi-Fi and 1/10GbE (Thunderbolt if + relevant); report end-to-end tok/s + p50/p95 inter-token latency, not local + MLX throughput. +5. **KV round-trip:** export/import multiple token pages, resume decode, compare + logits; test trim + speculative rejection; include rotating + quantized cache. +6. **Concurrency:** multiple sessions, cancellation, repeated resets; verify + MLX stream/array ownership under the chosen Tokio / dedicated-thread model. +7. **Compilation stability:** separate fixed-shape decode vs bucketed prefill; + watch recompilation counts, observer overhead, long-run graph/memory growth. + +Spikes 1 and 2 are more decisive than any standalone token/s benchmark. + +--- + +## 9. Risks and unknowns + +- **Partial load may require nontrivial changes to `safemlx-lm`** (loader + model + constructors currently build `0..num_hidden_layers`). Upstreaming to the fork + is likely necessary. (Highest risk.) +- **Eval-fence latency** could erode the benefit of adding Apple-Silicon compute + to a chain, especially over Wi-Fi. +- **Model coverage churn:** safemlx-lm is young; each family is bespoke Rust and + a separate certification. +- **Recurrent/hybrid + MoE** splitting is materially harder than dense; scope + them out of early phases. +- **Two artifact pipelines** add storage + certification cost; mitigate with a + single canonical BF16 source and reproducible derivation. +- **safemlx maturity/maintenance** (external fork of mlx-rs) — pin carefully; + expect to contribute upstream. +- **Compat discipline:** MLX must stay additive (feature-probe + gossip + capability); homogeneous chains by default; mixed-engine only when certified. + +--- + +## 10. Immediate next steps + +1. Land **Phase 1** (the `StageEngine` trait refactor, llama-only) — valuable on + its own and the prerequisite for everything else. +2. Run **Spike 1 (partial load)** and **Spike 2 (boundary fence)** in + `../safemlx` against a dense Qwen/Llama. Treat both as go/no-go. +3. If both pass, proceed to Phase 2 reusing goose's `safemlx-lm` generation + patterns as the starting point. + +--- + +## Appendix — primary sources reviewed + +**This repo (Skippy):** +- `crates/skippy-ffi/src/lib.rs`, `crates/skippy-ffi/README.md` — staged C ABI +- `crates/skippy-runtime/src/lib.rs` — safe stage model/session, activation + frames, KV/state movement +- `crates/skippy-server/src/frontend*` — stage driver, generation-3 protocol +- `crates/skippy-topology/src/lib.rs` — split planning, wire sizing, family caps +- `crates/skippy-runtime/src/package/materialized_cache.rs` — identity-bound + stage artifact cache +- `docs/design/LLAMA_STAGE_INTEGRATION_PLAN.md`, `docs/SKIPPY.md` — why the ABI + is shaped this way; backend-selector parity + +**MLX Rust fork (`../safemlx`):** +- `safemlx-lm/src/models/{qwen3,llama,gpt_oss,gemma4,...}.rs` — per-layer forward +- `safemlx-lm/src/{cache,inspection,weights}.rs`, `models/mod.rs` — KV cache, + observer hooks, strict/sharded loading, JIT quant +- `safemlx-sys/src/mlx-c/mlx/c/distributed.h` — MLX collectives (unwrapped) + +**goose (`../goose`, Apache-2.0):** +- `crates/goose-local-inference/src/{mlx,backend,hf_models}.rs`, + `crates/goose-download-manager` — reference MLX backend + HF download + +**External (grounded via web search):** +- [MLX lazy evaluation](https://ml-explore.github.io/mlx/build/html/usage/lazy_evaluation.html) +- [MLX distributed](https://ml-explore.github.io/mlx/build/html/usage/distributed.html) +- [safemlx docs.rs](https://docs.rs/safemlx/latest/safemlx/) · + [safemlx-lm](https://docs.rs/safemlx-lm/latest/safemlx_lm/) · + [inspection](https://docs.rs/safemlx-lm/latest/safemlx_lm/inspection/) · + [compile](https://docs.rs/safemlx/latest/safemlx/transforms/compile/) +- mlx-lm reference: + [utils.py](https://raw.githubusercontent.com/ml-explore/mlx-lm/main/mlx_lm/utils.py) · + [pipeline.py](https://raw.githubusercontent.com/ml-explore/mlx-lm/main/mlx_lm/models/pipeline.py) · + [cache.py](https://raw.githubusercontent.com/ml-explore/mlx-lm/main/mlx_lm/models/cache.py) · + [deepseek_v3.py](https://raw.githubusercontent.com/ml-explore/mlx-lm/main/mlx_lm/models/deepseek_v3.py) From d2bf949eddddc81d0712e8290cc0d649ccd8d4da Mon Sep 17 00:00:00 2001 From: michaelneale Date: Thu, 16 Jul 2026 12:55:06 +1000 Subject: [PATCH 02/37] spike: prove goose-style solo MLX serving from raw safetensors MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- docs/design/MLX_STAGE_ENGINE_PLAN.md | 34 +- spikes/mlx-solo/Cargo.lock | 1556 ++++++++++++++++++++++++++ spikes/mlx-solo/Cargo.toml | 30 + spikes/mlx-solo/FINDINGS.md | 130 +++ spikes/mlx-solo/src/main.rs | 171 +++ 5 files changed, 1916 insertions(+), 5 deletions(-) create mode 100644 spikes/mlx-solo/Cargo.lock create mode 100644 spikes/mlx-solo/Cargo.toml create mode 100644 spikes/mlx-solo/FINDINGS.md create mode 100644 spikes/mlx-solo/src/main.rs diff --git a/docs/design/MLX_STAGE_ENGINE_PLAN.md b/docs/design/MLX_STAGE_ENGINE_PLAN.md index b9267f685a..af4f1f62a5 100644 --- a/docs/design/MLX_STAGE_ENGINE_PLAN.md +++ b/docs/design/MLX_STAGE_ENGINE_PLAN.md @@ -12,6 +12,11 @@ It combines a read of the current Skippy code (`skippy-ffi`, `skippy-runtime`, and a second-opinion review from an external model grounded against live MLX/mlx-lm/safemlx documentation. +**Update — a Phase-2 solo-serving spike has now run** (`spikes/mlx-solo/`, branch +`micn/mlx-redux`). It confirms the core workflow claim end to end and surfaced +two concrete findings (a CPU quant-perf cliff and two safemlx-lm papercuts). See +`spikes/mlx-solo/FINDINGS.md`; the results are folded into §5.3, Phase 2, and §9. + --- ## 1. Bottom line @@ -574,6 +579,14 @@ existing skippy backend selector in `docs/SKIPPY.md`). Validate against "serve any supported model instantly, no wait for quant" benefit with minimal new distributed work, and de-risks the engine before any split work. +> **Spike done (`spikes/mlx-solo/`).** The load→generate half is proven: Qwen3-0.6B +> from raw HF safetensors, in Rust, CPU-only, **18.1 tok/s** decode, coherent — no +> GGUF, no pre-quant. **But JIT quant on CPU is a trap:** 4-bit/8-bit are correct +> yet run at **~0.4 tok/s** (MLX quant matmul is Metal-optimized, no fast CPU +> kernel). So Phase 2 must be validated on **Metal**, and JIT quant should be +> gated behind a Metal (or CUDA) backend rather than offered as a CPU path. Two +> safemlx-lm papercuts were also found and fixed in the fork (see §9). + **Phase 3 — Stage-aware partial load + activation frames.** Add `forward_range` / `resume_from_hidden` and the stage-aware loader to `safemlx-lm` (upstream to the fork). Implement `prefill_chunk_frame` / `decode_step_frame` / @@ -625,19 +638,30 @@ Spikes 1 and 2 are more decisive than any standalone token/s benchmark. ## 9. Risks and unknowns +- **JIT quant has no fast CPU kernel — it needs Metal/CUDA (confirmed by spike).** + 4-bit/8-bit affine quant on Apple-Silicon **CPU** ran ~45× slower than source + precision (0.4 vs 18.1 tok/s on Qwen3-0.6B). The "serve any model instantly, + JIT-quantized" workflow is therefore only attractive on an accelerator backend; + do not expose CPU JIT quant as a serving path. Metal-backed re-measurement is + a prerequisite before leaning on the §5.3 workflow claim. - **Partial load may require nontrivial changes to `safemlx-lm`** (loader + model constructors currently build `0..num_hidden_layers`). Upstreaming to the fork - is likely necessary. (Highest risk.) + is likely necessary. (Highest risk for the split work.) - **Eval-fence latency** could erode the benefit of adding Apple-Silicon compute to a chain, especially over Wi-Fi. -- **Model coverage churn:** safemlx-lm is young; each family is bespoke Rust and - a separate certification. +- **Model coverage churn — confirmed by spike:** safemlx-lm is young; each family + is bespoke Rust and separately certified. The spike hit two papercuts on + Qwen3-0.6B alone: (1) the published crate hard-enables the `metal` feature, so + a Metal-less/CI build needs a **workspace-level** `default-features = false`; + (2) tied-embedding `lm_head.weight` fails the *quantized* strict loader + (dense load tolerates it). Both fixed in the fork; expect more per-family. - **Recurrent/hybrid + MoE** splitting is materially harder than dense; scope them out of early phases. - **Two artifact pipelines** add storage + certification cost; mitigate with a single canonical BF16 source and reproducible derivation. -- **safemlx maturity/maintenance** (external fork of mlx-rs) — pin carefully; - expect to contribute upstream. +- **safemlx maturity/maintenance** (external fork of mlx-rs) — the `MlxStageEngine` + will likely carry a small fork or need upstream PRs (backend feature exposure, + loader fixes, and eventually `forward_range`/partial-load). Pin carefully. - **Compat discipline:** MLX must stay additive (feature-probe + gossip capability); homogeneous chains by default; mixed-engine only when certified. diff --git a/spikes/mlx-solo/Cargo.lock b/spikes/mlx-solo/Cargo.lock new file mode 100644 index 0000000000..82a795d150 --- /dev/null +++ b/spikes/mlx-solo/Cargo.lock @@ -0,0 +1,1556 @@ +# This file is automatically @generated by Cargo. +# It is not intended for manual editing. +version = 4 + +[[package]] +name = "ahash" +version = "0.8.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5a15f179cd60c4584b8a8c596927aadc462e27f2ca70c04e0071964a73ba7a75" +dependencies = [ + "cfg-if", + "getrandom 0.3.4", + "once_cell", + "serde", + "version_check", + "zerocopy", +] + +[[package]] +name = "aho-corasick" +version = "1.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ddd31a130427c27518df266943a5308ed92d4b226cc639f5a8f1002816174301" +dependencies = [ + "memchr", +] + +[[package]] +name = "allocator-api2" +version = "0.2.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "683d7910e743518b0e34f1186f92494becacb047c7b6bf616c96772180fef923" + +[[package]] +name = "anstream" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "824a212faf96e9acacdbd09febd34438f8f711fb84e09a8916013cd7815ca28d" +dependencies = [ + "anstyle", + "anstyle-parse", + "anstyle-query", + "anstyle-wincon", + "colorchoice", + "is_terminal_polyfill", + "utf8parse", +] + +[[package]] +name = "anstyle" +version = "1.0.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "940b3a0ca603d1eade50a4846a2afffd5ef57a9feac2c0e2ec2e14f9ead76000" + +[[package]] +name = "anstyle-parse" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "52ce7f38b242319f7cabaa6813055467063ecdc9d355bbb4ce0c68908cd8130e" +dependencies = [ + "utf8parse", +] + +[[package]] +name = "anstyle-query" +version = "1.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "40c48f72fd53cd289104fc64099abca73db4166ad86ea0b4341abe65af83dadc" +dependencies = [ + "windows-sys", +] + +[[package]] +name = "anstyle-wincon" +version = "3.0.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "291e6a250ff86cd4a820112fb8898808a366d8f9f58ce16d1f538353ad55747d" +dependencies = [ + "anstyle", + "once_cell_polyfill", + "windows-sys", +] + +[[package]] +name = "anyhow" +version = "1.0.103" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2a4385e2e34eb35d6b3efe798b9eb88096925d87726c0798709bf56d9ed84af3" + +[[package]] +name = "autocfg" +version = "1.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f2032f911046de80f0a198e0901378627c33f59ea0ac00e363d481118bd70a53" + +[[package]] +name = "base64" +version = "0.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9e1b586273c5702936fe7b7d6896644d8be71e6314cfe09d3167c95f712589e8" + +[[package]] +name = "bit-set" +version = "0.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "08807e080ed7f9d5433fa9b275196cfc35414f66a0c79d864dc51a0d825231a3" +dependencies = [ + "bit-vec", +] + +[[package]] +name = "bit-vec" +version = "0.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5e764a1d40d510daf35e07be9eb06e75770908c27d411ee6c92109c9840eaaf7" + +[[package]] +name = "bitflags" +version = "2.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b588b76d00fde79687d7646a9b5bdf3cc0f655e0bbd080335a95d7e96f3587da" + +[[package]] +name = "bytemuck" +version = "1.25.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d6aedf8ae72766347502cf3cb4f41cf5e9cc37d28bee90f1fdaaae15f9cf9424" + +[[package]] +name = "castaway" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dec551ab6e7578819132c713a93c022a05d60159dc86e7a7050223577484c55a" +dependencies = [ + "rustversion", +] + +[[package]] +name = "cc" +version = "1.2.67" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e17dd265a7d0f31ef544e1b20e03add05d3b45b491b633b10d67145d2acc1a38" +dependencies = [ + "find-msvc-tools", + "shlex", +] + +[[package]] +name = "cfg-if" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" + +[[package]] +name = "clap" +version = "4.6.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dd059f9da4f5c36b3787f65d38ccaab1cc315f07b01f89abc8359ee6a8205011" +dependencies = [ + "clap_builder", + "clap_derive", +] + +[[package]] +name = "clap_builder" +version = "4.6.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f09628afdcc538b57f3c6341e9c8e9970f18e4a481690a64974d7023bd33548b" +dependencies = [ + "anstream", + "anstyle", + "clap_lex", + "strsim", +] + +[[package]] +name = "clap_derive" +version = "4.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f2ce8604710f6733aa641a2b3731eaa1e8b3d9973d5e3565da11800813f997a9" +dependencies = [ + "heck", + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "clap_lex" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8d4a3bb8b1e0c1050499d1815f5ab16d04f0959b233085fb31653fbfc9d98f9" + +[[package]] +name = "cmake" +version = "0.1.58" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c0f78a02292a74a88ac736019ab962ece0bc380e3f977bf72e376c5d78ff0678" +dependencies = [ + "cc", +] + +[[package]] +name = "colorchoice" +version = "1.0.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1d07550c9036bf2ae0c684c4297d503f838287c83c53686d05370d0e139ae570" + +[[package]] +name = "compact_str" +version = "0.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9dfdd1c2274d9aa354115b09dc9a901d6c5576818cdf70d14cae2bdb47df00ab" +dependencies = [ + "castaway", + "cfg-if", + "itoa", + "rustversion", + "ryu", + "serde", + "static_assertions", +] + +[[package]] +name = "crossbeam-deque" +version = "0.8.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5181e0de7b61eb03a81e347d6dd8797bae9da5146707b51077e2d71a54ec0ceb" +dependencies = [ + "crossbeam-epoch", + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-epoch" +version = "0.9.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2d6914041f254d6e9176c01941b21115dcfb7089e55135a35411081bd106ef3f" +dependencies = [ + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-utils" +version = "0.8.22" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "61803da095bee82a81bb1a452ecc25d3b2f1416d1897eb86430c6159ef717c17" + +[[package]] +name = "crunchy" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "460fbee9c2c2f33933d720630a6a0bac33ba7053db5344fac858d4b8952d77d5" + +[[package]] +name = "daachorse" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6f55d7153ba3b507595872a3874803f07a8a81d1e888abed8e5db7da0597d6e2" + +[[package]] +name = "darling" +version = "0.20.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc7f46116c46ff9ab3eb1597a45688b6715c6e628b5c133e288e709a29bcb4ee" +dependencies = [ + "darling_core 0.20.11", + "darling_macro 0.20.11", +] + +[[package]] +name = "darling" +version = "0.23.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "25ae13da2f202d56bd7f91c25fba009e7717a1e4a1cc98a76d844b65ae912e9d" +dependencies = [ + "darling_core 0.23.0", + "darling_macro 0.23.0", +] + +[[package]] +name = "darling_core" +version = "0.20.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0d00b9596d185e565c2207a0b01f8bd1a135483d02d9b7b0a54b11da8d53412e" +dependencies = [ + "fnv", + "ident_case", + "proc-macro2", + "quote", + "strsim", + "syn", +] + +[[package]] +name = "darling_core" +version = "0.23.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9865a50f7c335f53564bb694ef660825eb8610e0a53d3e11bf1b0d3df31e03b0" +dependencies = [ + "ident_case", + "proc-macro2", + "quote", + "strsim", + "syn", +] + +[[package]] +name = "darling_macro" +version = "0.20.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc34b93ccb385b40dc71c6fceac4b2ad23662c7eeb248cf10d529b7e055b6ead" +dependencies = [ + "darling_core 0.20.11", + "quote", + "syn", +] + +[[package]] +name = "darling_macro" +version = "0.23.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ac3984ec7bd6cfa798e62b4a642426a5be0e68f9401cfc2a01e3fa9ea2fcdb8d" +dependencies = [ + "darling_core 0.23.0", + "quote", + "syn", +] + +[[package]] +name = "dary_heap" +version = "0.3.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8b1e3a325bc115f096c8b77bbf027a7c2592230e70be2d985be950d3d5e60ebe" +dependencies = [ + "serde", +] + +[[package]] +name = "derive_builder" +version = "0.20.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "507dfb09ea8b7fa618fcf76e953f4f5e192547945816d5358edffe39f6f94947" +dependencies = [ + "derive_builder_macro", +] + +[[package]] +name = "derive_builder_core" +version = "0.20.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2d5bcf7b024d6835cfb3d473887cd966994907effbe9227e8c8219824d06c4e8" +dependencies = [ + "darling 0.20.11", + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "derive_builder_macro" +version = "0.20.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ab63b0e2bf4d5928aff72e83a7dace85d7bba5fe12dcc3c5a572d78caffd3f3c" +dependencies = [ + "derive_builder_core", + "syn", +] + +[[package]] +name = "displaydoc" +version = "0.2.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ac70aa55017e108007fbaf5aa0f54b021c98f92ff8af59d42eda9da96e3dd4f" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "dyn-clone" +version = "1.0.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d0881ea181b1df73ff77ffaaf9c7544ecc11e82fba9b5f27b262a3c73a332555" + +[[package]] +name = "either" +version = "1.16.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91622ff5e7162018101f2fea40d6ebf4a78bbe5a49736a2020649edf9693679e" + +[[package]] +name = "equivalent" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "877a4ace8713b0bcf2a4e7eec82529c029f1d0619886d18145fea96c3ffe5c0f" + +[[package]] +name = "errno" +version = "0.3.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" +dependencies = [ + "libc", + "windows-sys", +] + +[[package]] +name = "esaxx-rs" +version = "0.1.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d817e038c30374a4bcb22f94d0a8a0e216958d4c3dcde369b1439fec4bdda6e6" + +[[package]] +name = "fancy-regex" +version = "0.17.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "72cf461f865c862bb7dc573f643dd6a2b6842f7c30b07882b56bd148cc2761b8" +dependencies = [ + "bit-set", + "regex-automata", + "regex-syntax", +] + +[[package]] +name = "fastrand" +version = "2.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9f1f227452a390804cdb637b74a86990f2a7d7ba4b7d5693aac9b4dd6defd8d6" + +[[package]] +name = "find-msvc-tools" +version = "0.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5baebc0774151f905a1a2cc41989300b1e6fbb29aff0ceffa1064fdd3088d582" + +[[package]] +name = "fnv" +version = "1.0.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3f9eec918d3f24069decb9af1554cad7c880e2da24a9afd88aca000531ab82c1" + +[[package]] +name = "foldhash" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "77ce24cb58228fbb8aa041425bb1050850ac19177686ea6e0f41a70416f56fdb" + +[[package]] +name = "getrandom" +version = "0.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "899def5c37c4fd7b2664648c28120ecec138e4d395b459e5ca34f9cce2dd77fd" +dependencies = [ + "cfg-if", + "libc", + "r-efi 5.3.0", + "wasip2", +] + +[[package]] +name = "getrandom" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "300e883d756b2e4ec94e02791f39b04b522276138852cfc41d9fb7e904106099" +dependencies = [ + "cfg-if", + "libc", + "r-efi 6.0.0", +] + +[[package]] +name = "half" +version = "2.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" +dependencies = [ + "cfg-if", + "crunchy", + "zerocopy", +] + +[[package]] +name = "hashbrown" +version = "0.16.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "841d1cc9bed7f9236f321df977030373f4a4163ae1a7dbfe1a51a2c1a51d9100" +dependencies = [ + "allocator-api2", + "equivalent", + "foldhash", + "serde", + "serde_core", +] + +[[package]] +name = "hashbrown" +version = "0.17.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a" + +[[package]] +name = "heck" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" + +[[package]] +name = "icu_collections" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2984d1cd16c883d7935b9e07e44071dca8d917fd52ecc02c04d5fa0b5a3f191c" +dependencies = [ + "displaydoc", + "potential_utf", + "utf8_iter", + "yoke", + "zerofrom", + "zerovec", +] + +[[package]] +name = "icu_locale_core" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "92219b62b3e2b4d88ac5119f8904c10f8f61bf7e95b640d25ba3075e6cac2c29" +dependencies = [ + "displaydoc", + "litemap", + "tinystr", + "writeable", + "zerovec", +] + +[[package]] +name = "icu_normalizer" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c56e5ee99d6e3d33bd91c5d85458b6005a22140021cc324cea84dd0e72cff3b4" +dependencies = [ + "icu_collections", + "icu_provider", + "smallvec", + "zerovec", +] + +[[package]] +name = "icu_properties" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bee3b67d0ea5c2cca5003417989af8996f8604e34fb9ddf96208a033901e70de" +dependencies = [ + "icu_collections", + "icu_locale_core", + "icu_provider", + "zerotrie", + "zerovec", +] + +[[package]] +name = "icu_provider" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "139c4cf31c8b5f33d7e199446eff9c1e02decfc2f0eec2c8d71f65befa45b421" +dependencies = [ + "displaydoc", + "icu_locale_core", + "yoke", + "zerofrom", + "zerovec", +] + +[[package]] +name = "ident_case" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b9e0384b61958566e926dc50660321d12159025e767c18e043daf26b70104c39" + +[[package]] +name = "idna_adapter" +version = "1.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cb68373c0d6620ef8105e855e7745e18b0d00d3bdb07fb532e434244cdb9a714" +dependencies = [ + "icu_normalizer", + "icu_properties", +] + +[[package]] +name = "indexmap" +version = "2.14.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d466e9454f08e4a911e14806c24e16fba1b4c121d1ea474396f396069cf949d9" +dependencies = [ + "equivalent", + "hashbrown 0.17.1", +] + +[[package]] +name = "is_terminal_polyfill" +version = "1.70.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a6cb138bb79a146c1bd460005623e142ef0181e3d0219cb493e02f7d08a35695" + +[[package]] +name = "itertools" +version = "0.14.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2b192c782037fadd9cfa75548310488aabdbf3d2da73885b31bd0abd03351285" +dependencies = [ + "either", +] + +[[package]] +name = "itertools" +version = "0.15.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8b4baf93f58d4425749ca49a51c50ebab072c5df6994d08fed93541c331481dc" +dependencies = [ + "either", +] + +[[package]] +name = "itoa" +version = "1.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" + +[[package]] +name = "libc" +version = "0.2.186" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68ab91017fe16c622486840e4c83c9a37afeff978bd239b5293d61ece587de66" + +[[package]] +name = "linux-raw-sys" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32a66949e030da00e8c7d4434b251670a91556f4144941d37452769c25d58a53" + +[[package]] +name = "litemap" +version = "0.8.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "92daf443525c4cce67b150400bc2316076100ce0b3686209eb8cf3c31612e6f0" + +[[package]] +name = "lock_api" +version = "0.4.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "224399e74b87b5f3557511d98dff8b14089b3dadafcab6bb93eab67d3aace965" +dependencies = [ + "scopeguard", +] + +[[package]] +name = "log" +version = "0.4.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0ceec5bc11778974d1bcb055b18002eba7f4b3518b6a0081b3af5f21666da9ad" + +[[package]] +name = "macro_rules_attribute" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "65049d7923698040cd0b1ddcced9b0eb14dd22c5f86ae59c3740eab64a676520" +dependencies = [ + "macro_rules_attribute-proc_macro", + "paste", +] + +[[package]] +name = "macro_rules_attribute-proc_macro" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "670fdfda89751bc4a84ac13eaa63e205cf0fd22b4c9a5fbfa085b63c1f1d3a30" + +[[package]] +name = "memchr" +version = "2.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf8baf1c55e62ffcace7a9f06f4bd9cd3f0c4beb022d3b367256b91b87513d98" + +[[package]] +name = "memmap2" +version = "0.9.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d1219ed1b7f229ee7104d281dd01d6802fe28bb6e95d292942c4daacdeb798c0" +dependencies = [ + "libc", +] + +[[package]] +name = "memo-map" +version = "0.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "38d1115007560874e373613744c6fba374c17688327a71c1476d1a5954cc857b" + +[[package]] +name = "minijinja" +version = "2.21.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cb3d648e68cea56d9858d535ee28f9538404e2dd8cb08ed0bd05dca379477f39" +dependencies = [ + "memo-map", + "serde", +] + +[[package]] +name = "minijinja-contrib" +version = "2.21.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85342f6fac0be8ccd5bd00d9066be538f34f393f577b75d81b17c8398a6b43bb" +dependencies = [ + "minijinja", + "serde", +] + +[[package]] +name = "minimal-lexical" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68354c5c6bd36d73ff3feceb05efa59b6acb7626617f4962be322a825e61f79a" + +[[package]] +name = "mlx-solo-spike" +version = "0.0.0" +dependencies = [ + "anyhow", + "clap", + "safemlx", + "safemlx-lm", + "serde_json", +] + +[[package]] +name = "monostate" +version = "0.1.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3341a273f6c9d5bef1908f17b7267bbab0e95c9bf69a0d4dcf8e9e1b2c76ef67" +dependencies = [ + "monostate-impl", + "serde", + "serde_core", +] + +[[package]] +name = "monostate-impl" +version = "0.1.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e4db6d5580af57bf992f59068d4ea26fd518574ff48d7639b255a36f9de6e7e9" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "nom" +version = "7.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d273983c5a657a70a3e8f2a01329822f3b8c8172b73826411a55751e404a0a4a" +dependencies = [ + "memchr", + "minimal-lexical", +] + +[[package]] +name = "num-complex" +version = "0.4.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "73f88a1307638156682bada9d7604135552957b7818057dcef22705b4d509495" +dependencies = [ + "num-traits", +] + +[[package]] +name = "num-traits" +version = "0.2.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "071dfc062690e90b734c0b2273ce72ad0ffa95f0c74596bc250dcfd960262841" +dependencies = [ + "autocfg", +] + +[[package]] +name = "num_enum" +version = "0.7.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5d0bca838442ec211fa11de3a8b0e0e8f3a4522575b5c4c06ed722e005036f26" +dependencies = [ + "num_enum_derive", + "rustversion", +] + +[[package]] +name = "num_enum_derive" +version = "0.7.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "680998035259dcfcafe653688bf2aa6d3e2dc05e98be6ab46afb089dc84f1df8" +dependencies = [ + "proc-macro-crate", + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "once_cell" +version = "1.21.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" + +[[package]] +name = "once_cell_polyfill" +version = "1.70.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "384b8ab6d37215f3c5301a95a4accb5d64aa607f1fcb26a11b5303878451b4fe" + +[[package]] +name = "parking_lot" +version = "0.12.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "93857453250e3077bd71ff98b6a65ea6621a19bb0f559a85248955ac12c45a1a" +dependencies = [ + "lock_api", + "parking_lot_core", +] + +[[package]] +name = "parking_lot_core" +version = "0.9.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2621685985a2ebf1c516881c026032ac7deafcda1a2c9b7850dc81e3dfcb64c1" +dependencies = [ + "cfg-if", + "libc", + "redox_syscall", + "smallvec", + "windows-link", +] + +[[package]] +name = "paste" +version = "1.0.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "57c0d7b74b563b49d38dae00a0c37d4d6de9b432382b2892f0574ddcae73fd0a" + +[[package]] +name = "potential_utf" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0103b1cef7ec0cf76490e969665504990193874ea05c85ff9bab8b911d0a0564" +dependencies = [ + "zerovec", +] + +[[package]] +name = "ppv-lite86" +version = "0.2.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85eae3c4ed2f50dcfe72643da4befc30deadb458a9b590d720cde2f2b1e97da9" +dependencies = [ + "zerocopy", +] + +[[package]] +name = "proc-macro-crate" +version = "3.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e67ba7e9b2b56446f1d419b1d807906278ffa1a658a8a5d8a39dcb1f5a78614f" +dependencies = [ + "toml_edit", +] + +[[package]] +name = "proc-macro2" +version = "1.0.106" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8fd00f0bb2e90d81d1044c2b32617f68fcb9fa3bb7640c23e9c748e53fb30934" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "quote" +version = "1.0.46" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dfbc457d0c7a0759a614551b11a6409e5951f6c7537be1f1b7682b9ae9230368" +dependencies = [ + "proc-macro2", +] + +[[package]] +name = "r-efi" +version = "5.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "69cdb34c158ceb288df11e18b4bd39de994f6657d83847bdffdbd7f346754b0f" + +[[package]] +name = "r-efi" +version = "6.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8dcc9c7d52a811697d2151c701e0d08956f92b0e24136cf4cf27b57a6a0d9bf" + +[[package]] +name = "rand" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b9ef1d0d795eb7d84685bca4f72f3649f064e6641543d3a8c415898726a57b41" +dependencies = [ + "rand_chacha", + "rand_core", +] + +[[package]] +name = "rand_chacha" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d3022b5f1df60f26e1ffddd6c66e8aa15de382ae63b3a0c1bfc0e4d3e3f325cb" +dependencies = [ + "ppv-lite86", + "rand_core", +] + +[[package]] +name = "rand_core" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "76afc826de14238e6e8c374ddcc1fa19e374fd8dd986b0d2af0d02377261d83c" +dependencies = [ + "getrandom 0.3.4", +] + +[[package]] +name = "rayon" +version = "1.12.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fb39b166781f92d482534ef4b4b1b2568f42613b53e5b6c160e24cfbfa30926d" +dependencies = [ + "either", + "rayon-core", +] + +[[package]] +name = "rayon-cond" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2964d0cf57a3e7a06e8183d14a8b527195c706b7983549cd5462d5aa3747438f" +dependencies = [ + "either", + "itertools 0.14.0", + "rayon", +] + +[[package]] +name = "rayon-core" +version = "1.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "22e18b0f0062d30d4230b2e85ff77fdfe4326feb054b9783a3460d8435c8ab91" +dependencies = [ + "crossbeam-deque", + "crossbeam-utils", +] + +[[package]] +name = "redox_syscall" +version = "0.5.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed2bf2547551a7053d6fdfafda3f938979645c44812fbfcda098faae3f1a362d" +dependencies = [ + "bitflags", +] + +[[package]] +name = "regex" +version = "1.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f020237b6c8eed93db2e2cb53c00c60a8e1bc73da7d073199a1180401450218d" +dependencies = [ + "aho-corasick", + "memchr", + "regex-automata", + "regex-syntax", +] + +[[package]] +name = "regex-automata" +version = "0.4.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8fcfdb36bda0c880c5931cdc7a2bcdc8ba4556847b9d912bca70bc94708711ad" +dependencies = [ + "aho-corasick", + "memchr", + "regex-syntax", +] + +[[package]] +name = "regex-syntax" +version = "0.8.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d6f6ff9a378485b298a5286656da665ba74413d36db0979633275d2e708145d4" + +[[package]] +name = "rustix" +version = "1.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6fe4565b9518b83ef4f91bb47ce29620ca828bd32cb7e408f0062e9930ba190" +dependencies = [ + "bitflags", + "errno", + "libc", + "linux-raw-sys", + "windows-sys", +] + +[[package]] +name = "rustversion" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf54715a573b99ac80df0bc206da022bcd442c974952c7b9720069370852e21f" + +[[package]] +name = "ryu" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9774ba4a74de5f7b1c1451ed6cd5285a32eddb5cccb8cc655a4e50009e06477f" + +[[package]] +name = "safemlx" +version = "0.1.3" +dependencies = [ + "bytemuck", + "dyn-clone", + "half", + "itertools 0.15.0", + "libc", + "num-complex", + "num-traits", + "num_enum", + "parking_lot", + "paste", + "safemlx-internal-macros", + "safemlx-macros", + "safemlx-sys", + "safetensors", + "smallvec", + "strum", + "thiserror", +] + +[[package]] +name = "safemlx-internal-macros" +version = "0.1.1" +dependencies = [ + "darling 0.23.0", + "itertools 0.15.0", + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "safemlx-lm" +version = "0.4.1" +dependencies = [ + "anyhow", + "clap", + "idna_adapter", + "memmap2", + "minijinja", + "safemlx", + "safemlx-lm-utils", + "safetensors", + "serde", + "serde_json", + "thiserror", + "tokenizers", +] + +[[package]] +name = "safemlx-lm-utils" +version = "0.1.4" +dependencies = [ + "minijinja", + "minijinja-contrib", + "serde", + "serde_json", + "thiserror", + "tokenizers", +] + +[[package]] +name = "safemlx-macros" +version = "0.1.1" +dependencies = [ + "darling 0.23.0", + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "safemlx-sys" +version = "0.1.3" +dependencies = [ + "cc", + "cmake", +] + +[[package]] +name = "safetensors" +version = "0.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "79b079b829cb27a1c3c374341345ed2e8b2c0c839034522cee576c140bd7f846" +dependencies = [ + "hashbrown 0.16.1", + "libc", + "serde", + "serde_json", + "tempfile", +] + +[[package]] +name = "scopeguard" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" + +[[package]] +name = "serde" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9a8e94ea7f378bd32cbbd37198a4a91436180c5bb472411e48b5ec2e2124ae9e" +dependencies = [ + "serde_core", + "serde_derive", +] + +[[package]] +name = "serde_core" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "41d385c7d4ca58e59fc732af25c3983b67ac852c1a25000afe1175de458b67ad" +dependencies = [ + "serde_derive", +] + +[[package]] +name = "serde_derive" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d540f220d3187173da220f885ab66608367b6574e925011a9353e4badda91d79" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "serde_json" +version = "1.0.150" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e8014e44b4736ed0538adeecded0fce2a272f22dc9578a7eb6b2d9993c74cfb9" +dependencies = [ + "itoa", + "memchr", + "serde", + "serde_core", + "zmij", +] + +[[package]] +name = "shlex" +version = "2.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8fadd59c855ef2080decdef8ff161eb6661b86933c9d82e5ba29dc602a55aba" + +[[package]] +name = "smallvec" +version = "1.15.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8ed6a63f02c8539c91a8685a86f4099661ba3da017932f6ebbea6de3f0fa7c90" + +[[package]] +name = "spm_precompiled" +version = "0.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5851699c4033c63636f7ea4cf7b7c1f1bf06d0cc03cfb42e711de5a5c46cf326" +dependencies = [ + "base64", + "nom", + "serde", + "unicode-segmentation", +] + +[[package]] +name = "stable_deref_trait" +version = "1.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ce2be8dc25455e1f91df71bfa12ad37d7af1092ae736f3a6cd0e37bc7810596" + +[[package]] +name = "static_assertions" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a2eb9349b6444b326872e140eb1cf5e7c522154d69e7a0ffb0fb81c06b37543f" + +[[package]] +name = "strsim" +version = "0.11.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7da8b5736845d9f2fcb837ea5d9e2628564b3b043a70948a3f0b778838c5fb4f" + +[[package]] +name = "strum" +version = "0.28.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9628de9b8791db39ceda2b119bbe13134770b56c138ec1d3af810d045c04f9bd" +dependencies = [ + "strum_macros", +] + +[[package]] +name = "strum_macros" +version = "0.28.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ab85eea0270ee17587ed4156089e10b9e6880ee688791d45a905f5b1ca36f664" +dependencies = [ + "heck", + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "syn" +version = "2.0.119" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "872831b642d1a07999a962a351ed35b955ea2cfc8f3862091e2a240a84f17297" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "synstructure" +version = "0.13.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "728a70f3dbaf5bab7f0c4b1ac8d7ae5ea60a4b5549c8a5914361c99147a709d2" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "tempfile" +version = "3.27.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32497e9a4c7b38532efcdebeef879707aa9f794296a4f0244f6f69e9bc8574bd" +dependencies = [ + "fastrand", + "getrandom 0.4.3", + "once_cell", + "rustix", + "windows-sys", +] + +[[package]] +name = "thiserror" +version = "2.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4288b5bcbc7920c07a1149a35cf9590a2aa808e0bc1eafaade0b80947865fbc4" +dependencies = [ + "thiserror-impl", +] + +[[package]] +name = "thiserror-impl" +version = "2.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ebc4ee7f67670e9b64d05fa4253e753e016c6c95ff35b89b7941d6b856dec1d5" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "tinystr" +version = "0.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8323304221c2a851516f22236c5722a72eaa19749016521d6dff0824447d96d" +dependencies = [ + "displaydoc", + "zerovec", +] + +[[package]] +name = "tokenizers" +version = "0.23.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "44e5bea67576e04b6ff8564c5d9e09c2ef0cf476502245f2f120e497769d3112" +dependencies = [ + "ahash", + "compact_str", + "daachorse", + "dary_heap", + "derive_builder", + "esaxx-rs", + "fancy-regex", + "getrandom 0.3.4", + "itertools 0.14.0", + "log", + "macro_rules_attribute", + "monostate", + "paste", + "rand", + "rayon", + "rayon-cond", + "regex", + "regex-syntax", + "serde", + "serde_json", + "spm_precompiled", + "thiserror", + "unicode-normalization-alignments", + "unicode-segmentation", + "unicode_categories", +] + +[[package]] +name = "toml_datetime" +version = "1.1.1+spec-1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3165f65f62e28e0115a00b2ebdd37eb6f3b641855f9d636d3cd4103767159ad7" +dependencies = [ + "serde_core", +] + +[[package]] +name = "toml_edit" +version = "0.25.13+spec-1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6975367e4d2ef766d86af01ffad14b622fecc8d4357a998fbc4deb6e9bacaf9b" +dependencies = [ + "indexmap", + "toml_datetime", + "toml_parser", + "winnow", +] + +[[package]] +name = "toml_parser" +version = "1.1.2+spec-1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a2abe9b86193656635d2411dc43050282ca48aa31c2451210f4202550afb7526" +dependencies = [ + "winnow", +] + +[[package]] +name = "unicode-ident" +version = "1.0.24" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75" + +[[package]] +name = "unicode-normalization-alignments" +version = "0.1.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "43f613e4fa046e69818dd287fdc4bc78175ff20331479dab6e1b0f98d57062de" +dependencies = [ + "smallvec", +] + +[[package]] +name = "unicode-segmentation" +version = "1.13.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c6f5d3c3b1bf09027a88a6bc961fc00497d651009560b5463668dc81b0fa87a8" + +[[package]] +name = "unicode_categories" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "39ec24b3121d976906ece63c9daad25b85969647682eee313cb5779fdd69e14e" + +[[package]] +name = "utf8_iter" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6c140620e7ffbb22c2dee59cafe6084a59b5ffc27a8859a5f0d494b5d52b6be" + +[[package]] +name = "utf8parse" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "06abde3611657adf66d383f00b093d7faecc7fa57071cce2578660c9f1010821" + +[[package]] +name = "version_check" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" + +[[package]] +name = "wasip2" +version = "1.0.4+wasi-0.2.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b67efb37e106e55ce722a510d6b5f9c17f083e5fc79afc2badeb12cc313d9487" +dependencies = [ + "wit-bindgen", +] + +[[package]] +name = "windows-link" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" + +[[package]] +name = "windows-sys" +version = "0.61.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ae137229bcbd6cdf0f7b80a31df61766145077ddf49416a728b02cb3921ff3fc" +dependencies = [ + "windows-link", +] + +[[package]] +name = "winnow" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "23b97319f7b8343df12cc98938e5c3eb436064524c8d2b4e30a1d3a36eecdf81" +dependencies = [ + "memchr", +] + +[[package]] +name = "wit-bindgen" +version = "0.57.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ebf944e87a7c253233ad6766e082e3cd714b5d03812acc24c318f549614536e" + +[[package]] +name = "writeable" +version = "0.6.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ffae5123b2d3fc086436f8834ae3ab053a283cfac8fe0a0b8eaae044768a4c4" + +[[package]] +name = "yoke" +version = "0.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "709fe23a0424b6a435d82152b1bd3fdfb0833487d5fa90d05d42762a9891fef5" +dependencies = [ + "stable_deref_trait", + "yoke-derive", + "zerofrom", +] + +[[package]] +name = "yoke-derive" +version = "0.8.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "de844c262c8848816172cef550288e7dc6c7b7814b4ee56b3e1553f275f1858e" +dependencies = [ + "proc-macro2", + "quote", + "syn", + "synstructure", +] + +[[package]] +name = "zerocopy" +version = "0.8.54" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b7cbbc0a705a0fd05cc3676525980d2bf5a9bc4adac6d6475209a7887cf59d19" +dependencies = [ + "zerocopy-derive", +] + +[[package]] +name = "zerocopy-derive" +version = "0.8.54" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e2e817b7b52d0c7358d3246da9d69935ebb18116b2b102b4230dac079b4862f5" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "zerofrom" +version = "0.1.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0ec05a11813ea801ff6d75110ad09cd0824ddba17dfe17128ea0d5f68e6c5272" +dependencies = [ + "zerofrom-derive", +] + +[[package]] +name = "zerofrom-derive" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "11532158c46691caf0f2593ea8358fed6bbf68a0315e80aae9bd41fbade684a1" +dependencies = [ + "proc-macro2", + "quote", + "syn", + "synstructure", +] + +[[package]] +name = "zerotrie" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0f9152d31db0792fa83f70fb2f83148effb5c1f5b8c7686c3459e361d9bc20bf" +dependencies = [ + "displaydoc", + "yoke", + "zerofrom", +] + +[[package]] +name = "zerovec" +version = "0.11.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "90f911cbc359ab6af17377d242225f4d75119aec87ea711a880987b18cd7b239" +dependencies = [ + "yoke", + "zerofrom", + "zerovec-derive", +] + +[[package]] +name = "zerovec-derive" +version = "0.11.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "625dc425cab0dca6dc3c3319506e6593dcb08a9f387ea3b284dbd52a92c40555" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "zmij" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29666d0abbfad1e3dc4dcf6144730dd3a3ab225bbbdac83319345b1b44ccfc1b" diff --git a/spikes/mlx-solo/Cargo.toml b/spikes/mlx-solo/Cargo.toml new file mode 100644 index 0000000000..d1f9026e43 --- /dev/null +++ b/spikes/mlx-solo/Cargo.toml @@ -0,0 +1,30 @@ +# Standalone spike crate — deliberately its OWN workspace so it is NOT part of +# the mesh-llm cargo workspace (keeps MLX/CMake native deps off every other +# build). It path-depends on a sibling checkout of the safemlx fork: +# ../../../safemlx == /Users//Documents/code/safemlx +# It does not build in CI and is not meant to ship. +[package] +name = "mlx-solo-spike" +version = "0.0.0" +edition = "2021" +publish = false + +[[bin]] +name = "mlx-solo" +path = "src/main.rs" + +[workspace] + +[dependencies] +# CPU-only: accelerate backend, no metal (no Metal shader compiler on this box). +safemlx = { path = "../../../safemlx/safemlx", default-features = false, features = [ + "accelerate", + "safetensors", +] } +# Pure-Rust regex engine (fancy-regex) to avoid extra C/C++ tokenizer builds. +safemlx-lm = { path = "../../../safemlx/safemlx-lm", default-features = false, features = [ + "fancy-regex", +] } +anyhow = "1" +clap = { version = "4", features = ["derive"] } +serde_json = "1" diff --git a/spikes/mlx-solo/FINDINGS.md b/spikes/mlx-solo/FINDINGS.md new file mode 100644 index 0000000000..d82242e5ea --- /dev/null +++ b/spikes/mlx-solo/FINDINGS.md @@ -0,0 +1,130 @@ +# Spike: goose-style solo MLX serving — findings + +Validates the Phase-2 claim in `docs/design/MLX_STAGE_ENGINE_PLAN.md`: load an HF +**safetensors** model in Rust via `safemlx-lm`, optionally **JIT-quantize on +load**, and generate tokens — with no GGUF and no ahead-of-time quant step. + +Throwaway crate with its own `[workspace]`, deliberately **not** part of the +mesh-llm cargo workspace, so MLX/CMake native deps never touch other builds. + +## TL;DR + +- **Solo serving from raw safetensors works, in Rust, end to end.** ✅ + Qwen3-0.6B → coherent output at **18.1 tok/s** decode (CPU), no GGUF, no + pre-quant. +- **JIT quant is functional but has no fast CPU kernel.** 4-bit and 8-bit both + produce coherent output but collapse to **~0.4 tok/s** (~45× slower than + source precision). MLX quantized matmul is Metal-optimized; CPU is not a + viable quant path. **Metal is required to judge JIT-quant performance.** +- **safemlx-lm is young — hit (and fixed) a real loader bug** on the quant path + for tied-embedding checkpoints. +- **Feature wiring needs a fork/upstream change**: the published `safemlx-lm` + hard-enables the `metal` feature, so a Metal-less (CPU/CI) build is impossible + without a workspace-level `default-features = false`. + +## Environment + +| | | +| --- | --- | +| Machine | Apple Silicon (arm64), macOS 26.5 | +| Toolchain | rustc 1.97, cmake 4.4, Apple clang 21 | +| Metal shader compiler | **absent** (CommandLineTools only, no full Xcode) → CPU-only build | +| MLX backend | `accelerate` (CPU). MLX core `v0.32.0` via safemlx-sys FetchContent | +| safemlx fork | `jbg/safemlx` @ `0502a19` + 3 local edits (below) | +| Model | `Qwen/Qwen3-0.6B` (dense, safetensors, `tie_word_embeddings: true`) | + +## Results + +Command shape: +``` +mlx-solo --model -n 64 [--quantize 4|8] "" +``` + +| Mode | load | ttft | decode tok/s | output | +| --- | --- | --- | --- | --- | +| source precision | 0.30s | 0.174s | **18.1** | coherent ✅ | +| JIT 4-bit affine | 1.03s | 47.6s | **0.4** | coherent ✅ | +| JIT 8-bit affine | 1.13s | 48.4s | **0.4** | coherent ✅ | + +The quant slowdown is uniform across 4/8-bit and dominated by per-token decode +(ttft is essentially the first decode step), which is the signature of a missing +fast CPU quant-matmul kernel rather than a load-time cost. + +## Build notes + +- **MLX C++ builds cleanly under cmake 4.4, CPU-only** — this was the #1 feared + risk (cmake 4.x deprecates old-CMake compatibility). `libmlx.a` + `libmlxc.a` + produced, no policy errors. Cold build ~1m19s; warm rebuild after a Rust-only + change ~6.7s (MLX archives cached in the shared target dir). +- Only harmless linker warnings (`object file … has version 26.5.0, which is + newer than target minimum of 11.0.0`). + +## Issues found + +### 1. Published safemlx-lm hard-enables `metal` (build blocker off-Metal) + +`safemlx-lm`'s dependency on `safemlx` did not disable default features, and a +member-level `default-features = false` is **ignored** by Cargo (it warns and +requires the setting at the workspace root). Effect: any consumer inherits the +`metal` feature and the build `panic!`s when the Metal compiler is absent +(CI, Metal-less macOS, or a CPU-only lane). + +Fix in this spike (in the fork): +- root `Cargo.toml`: `safemlx = { …, default-features = false }` +- `safemlx-lm/Cargo.toml`: `safemlx = { workspace = true, default-features = false, features = ["accelerate", "safetensors"] }` + +**Plan implication:** the `MlxStageEngine` will need MLX backend selection as an +explicit cargo feature (`metal` / `accelerate` / `cuda`), which means either an +upstream PR to `safemlx-lm` or carrying a small fork. Real integration cost, not +a blocker. + +### 2. Tied-embedding `lm_head.weight` fails strict load on the quant path (qwen3) + +`Qwen3-0.6B` sets `tie_word_embeddings: true` but the checkpoint still ships a +redundant `lm_head.weight`. The **dense** load uses a lenient loader and +tolerates it; the **quantized** load uses `load_safetensors_dir_quantized_strict` +with a bare `StrictLoadConfig::default()` and rejects it as an unused tensor: + +``` +strict weight-load validation failed: 0 missing, 1 unused + unused: lm_head.weight +``` + +Fix in this spike (in the fork, `safemlx-lm/src/models/qwen3.rs` +`load_qwen3_model_quantized`): +```rust +let config = StrictLoadConfig::default().allow_unused_prefix("lm_head."); +``` +Harmless when untied (then `lm_head.weight` is a loaded param, not unused). + +**Plan implication:** confirms the "safemlx-lm is young; each model family is +bespoke and separately certified" risk. Expect per-family loader papercuts. + +## Reproduce + +1. Sibling checkout of the fork at `../safemlx` (repo-relative: + `/Users//Documents/code/safemlx`), base `0502a19`, with the two Cargo + edits and the qwen3 one-liner above. +2. Model dir with `config.json`, `tokenizer*.json`, `model.safetensors` + (e.g. `Qwen/Qwen3-0.6B`). +3. Build/run (own workspace, so it will trigger a one-time MLX C++ build): + ``` + cd spikes/mlx-solo + cargo build --release + ./target/release/mlx-solo --model -n 64 "..." + ./target/release/mlx-solo --model -n 64 --quantize 4 "..." + ``` + Using a shared `CARGO_TARGET_DIR` avoids rebuilding MLX across crates. + +## What this de-risks / what it does not + +De-risked: +- safemlx-lm as a Rust solo engine driving load + generate. +- The "no GGUF, no pre-quant, serve raw safetensors" workflow claim. +- MLX native build under this toolchain (cmake 4.4). + +Still open (needs Metal): +- Any performance judgement of JIT quant (CPU quant path is not representative). +- Metal throughput vs the llama.cpp backend. +- Everything staged/split (this spike is single-stage, whole-model) — the + partial-load and boundary-fence go/no-go spikes are unchanged. diff --git a/spikes/mlx-solo/src/main.rs b/spikes/mlx-solo/src/main.rs new file mode 100644 index 0000000000..c6bb9cf963 --- /dev/null +++ b/spikes/mlx-solo/src/main.rs @@ -0,0 +1,171 @@ +//! Spike: goose-style solo MLX serving via safemlx-lm. +//! +//! Proves the Phase-2 workflow claim from docs/design/MLX_STAGE_ENGINE_PLAN.md: +//! load an HF safetensors model in Rust, optionally JIT-quantize on load, and +//! generate tokens — with no GGUF and no ahead-of-time quant step. +//! +//! CPU-only (accelerate) because this box has no Metal shader compiler; the +//! generation *path* is what we are validating here, not Metal throughput. + +use std::io::Write; +use std::path::PathBuf; +use std::time::Instant; + +use anyhow::{bail, Context, Result}; +use clap::Parser; +use safemlx::transforms::async_eval; +use safemlx::{Device, DeviceType, Stream}; +use safemlx_lm::models::input::{InputPart, ModelInput}; +use safemlx_lm::models::{LoadedModel, ModelLoadOptions}; +use safemlx_lm::quantization::AffineQuantization; +use safemlx_lm::sampler::DefaultSampler; + +#[derive(Parser, Debug)] +#[command(about = "Solo MLX serving spike (load + optional JIT quant + generate)")] +struct Cli { + /// Model directory (HF safetensors) or GGUF file. + #[arg(short, long)] + model: PathBuf, + + /// Prompt text. + #[arg(default_value = "Explain what a mesh network is in two sentences.")] + prompt: String, + + /// Max tokens to generate. + #[arg(short = 'n', long, default_value_t = 128)] + max_tokens: usize, + + /// JIT-quantize eligible dense weights to this bit width on load (e.g. 4, 8). + /// Omit to load at source precision. + #[arg(short, long)] + quantize: Option, + + /// Group size for affine quantization. + #[arg(long, default_value_t = 64)] + quant_group_size: i32, + + /// Skip the chat template and feed the prompt raw. + #[arg(long)] + raw: bool, +} + +fn main() -> Result<()> { + let args = Cli::parse(); + + // CPU device: no Metal compiler on this box, so the GPU backend is absent. + let device = Device::new(DeviceType::Cpu, 0); + let stream = Stream::new_with_device(&device); + let weights_stream = Stream::new_with_device(&device); + + let load_options = match args.quantize { + Some(bits) => { + eprintln!("[load] JIT affine quantization: {bits}-bit, group_size={}", args.quant_group_size); + ModelLoadOptions::with_quantization(AffineQuantization::new(args.quant_group_size, bits)?) + } + None => { + eprintln!("[load] no quantization (source precision)"); + ModelLoadOptions::default() + } + }; + + eprintln!("[load] loading {} ...", args.model.display()); + let load_started = Instant::now(); + let mut model = LoadedModel::load_with_options(&args.model, load_options, &stream, &weights_stream) + .with_context(|| format!("failed to load model from {}", args.model.display()))?; + stream.synchronize()?; + let load_elapsed = load_started.elapsed(); + eprintln!( + "[load] ok: model_type={} in {:.2}s", + model.model_type(), + load_elapsed.as_secs_f64() + ); + + let (rendered, add_special) = if args.raw { + (args.prompt.clone(), true) + } else { + match model.apply_chat_template_json( + vec![vec![serde_json::json!({"role": "user", "content": args.prompt})]], + None, + true, + )? { + Some(rendered) => (rendered, false), + None => { + eprintln!("[prompt] no chat template; feeding raw"); + (args.prompt.clone(), true) + } + } + }; + + let tokens = model.encode_to_array(&rendered, add_special, &stream)?; + let prompt_len = tokens.shape()[1]; + if prompt_len == 0 { + bail!("prompt produced no tokens"); + } + eprintln!("[prompt] {prompt_len} tokens"); + + let eos = model.eos_token_ids().to_vec(); + let mut cache = model.new_cache(); + + let mut output_ids: Vec = Vec::with_capacity(args.max_tokens); + let gen_started = Instant::now(); + let mut ttft = None; + + { + let parts = [InputPart::text_token_ids(&tokens)]; + let input = ModelInput::new(&parts); + // Greedy sampling (temp=0.0, no prng key) keeps the spike deterministic. + let mut generator = + model.generate_input_with_cache_sampler(&mut cache, 0.0, input, None, &stream, DefaultSampler); + + let mut current = generator.next().transpose()?; + for index in 0..args.max_tokens { + let Some(token) = current.take() else { break }; + + // Kick off the next decode before reading this token back (mlx-lm's + // one-token async pipeline: overlaps compute with host readback). + let next = if index + 1 < args.max_tokens { + let next = generator.next(); + if let Some(Ok(next_token)) = next.as_ref() { + async_eval([next_token])?; + } + next + } else { + None + }; + + let token_id = token.item::(&stream); + if ttft.is_none() { + ttft = Some(gen_started.elapsed()); + } + output_ids.push(token_id); + if eos.contains(&token_id) { + break; + } + current = next.transpose()?; + } + } + + let gen_elapsed = gen_started.elapsed(); + let text = model.decode(&output_ids, true)?; + + let mut stdout = std::io::stdout().lock(); + writeln!(stdout, "\n===== OUTPUT =====\n{text}\n==================")?; + + let decode_tokens = output_ids.len().saturating_sub(1); + let decode_elapsed = gen_elapsed.saturating_sub(ttft.unwrap_or_default()); + let decode_rate = if decode_elapsed.is_zero() { + 0.0 + } else { + decode_tokens as f64 / decode_elapsed.as_secs_f64() + }; + eprintln!( + "[stats] load={:.2}s prompt_tokens={} generated={} ttft={:.3}s decode={:.1} tok/s", + load_elapsed.as_secs_f64(), + prompt_len, + output_ids.len(), + ttft.map(|d| d.as_secs_f64()).unwrap_or(0.0), + decode_rate, + ); + + Ok(()) +} From b63e10c0c65e0836acf067278febd29fd42d83b6 Mon Sep 17 00:00:00 2001 From: michaelneale Date: Thu, 16 Jul 2026 13:35:31 +1000 Subject: [PATCH 03/37] spike: real Metal numbers for goose-style solo MLX serving 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. --- docs/design/MLX_STAGE_ENGINE_PLAN.md | 44 ++++--- spikes/mlx-solo/Cargo.lock | 55 ++++---- spikes/mlx-solo/Cargo.toml | 21 +-- spikes/mlx-solo/FINDINGS.md | 185 +++++++++++++++------------ spikes/mlx-solo/src/main.rs | 8 +- 5 files changed, 174 insertions(+), 139 deletions(-) diff --git a/docs/design/MLX_STAGE_ENGINE_PLAN.md b/docs/design/MLX_STAGE_ENGINE_PLAN.md index af4f1f62a5..2ac7ca7b12 100644 --- a/docs/design/MLX_STAGE_ENGINE_PLAN.md +++ b/docs/design/MLX_STAGE_ENGINE_PLAN.md @@ -12,10 +12,15 @@ It combines a read of the current Skippy code (`skippy-ffi`, `skippy-runtime`, and a second-opinion review from an external model grounded against live MLX/mlx-lm/safemlx documentation. -**Update — a Phase-2 solo-serving spike has now run** (`spikes/mlx-solo/`, branch -`micn/mlx-redux`). It confirms the core workflow claim end to end and surfaced -two concrete findings (a CPU quant-perf cliff and two safemlx-lm papercuts). See -`spikes/mlx-solo/FINDINGS.md`; the results are folded into §5.3, Phase 2, and §9. +**Update — a Phase-2 solo-serving spike has now run on Metal** (`spikes/mlx-solo/`, +branch `micn/mlx-redux`). It confirms the core workflow claim end to end: Qwen3-0.6B +on Apple-Silicon Metal at **321 tok/s** (bf16) and **~604 tok/s** (4-bit), where +**JIT-quantize-on-load matches a pre-quantized artifact** (604 ≈ 603 tok/s) — so +quantizing on load is free at inference time. The goose baseline (source precision) +needs **zero fork patches**; two small `safemlx-lm` fixes are only needed to go +beyond it (JIT quant + loading arbitrary mlx-community repos) and are upstream-PR +candidates. See `spikes/mlx-solo/FINDINGS.md`; results are folded into §5.3, +Phase 2, and §9. --- @@ -579,13 +584,17 @@ existing skippy backend selector in `docs/SKIPPY.md`). Validate against "serve any supported model instantly, no wait for quant" benefit with minimal new distributed work, and de-risks the engine before any split work. -> **Spike done (`spikes/mlx-solo/`).** The load→generate half is proven: Qwen3-0.6B -> from raw HF safetensors, in Rust, CPU-only, **18.1 tok/s** decode, coherent — no -> GGUF, no pre-quant. **But JIT quant on CPU is a trap:** 4-bit/8-bit are correct -> yet run at **~0.4 tok/s** (MLX quant matmul is Metal-optimized, no fast CPU -> kernel). So Phase 2 must be validated on **Metal**, and JIT quant should be -> gated behind a Metal (or CUDA) backend rather than offered as a CPU path. Two -> safemlx-lm papercuts were also found and fixed in the fork (see §9). +> **Spike done on Metal (`spikes/mlx-solo/`).** The load→generate half is proven: +> Qwen3-0.6B from raw HF safetensors, in Rust, on Apple-Silicon Metal, matching +> goose's setup exactly (`["accelerate","metal","safetensors"]`, `Device::Gpu`). +> Measured decode: **321 tok/s** source precision (bf16), **~604 tok/s** at 4-bit — +> and crucially **JIT-quantize-on-load (604) ≈ a pre-quantized mlx-community repo +> (603)**, so quantizing on load is free at inference time. The source-precision +> path (goose's baseline) needs **zero fork patches**; two small `safemlx-lm` fixes +> are only needed to go beyond it (JIT quant of a tied-embedding checkpoint, and +> loading published quant repos that omit the `mode` field) — both upstream-PR +> candidates, not mesh-llm drift (see §9). CPU is not a serving path and was not +> benchmarked as one. **Phase 3 — Stage-aware partial load + activation frames.** Add `forward_range` / `resume_from_hidden` and the stage-aware loader to `safemlx-lm` (upstream to @@ -638,12 +647,13 @@ Spikes 1 and 2 are more decisive than any standalone token/s benchmark. ## 9. Risks and unknowns -- **JIT quant has no fast CPU kernel — it needs Metal/CUDA (confirmed by spike).** - 4-bit/8-bit affine quant on Apple-Silicon **CPU** ran ~45× slower than source - precision (0.4 vs 18.1 tok/s on Qwen3-0.6B). The "serve any model instantly, - JIT-quantized" workflow is therefore only attractive on an accelerator backend; - do not expose CPU JIT quant as a serving path. Metal-backed re-measurement is - a prerequisite before leaning on the §5.3 workflow claim. +- **JIT quant is free at inference time on Metal (confirmed by spike), but CPU is + not a serving path.** On Apple-Silicon Metal, JIT 4-bit (604 tok/s) matched a + pre-quantized mlx-community repo (603 tok/s), and source precision ran at 321 + tok/s — so the §5.3 "serve any model instantly, JIT-quantized" claim holds with + no runtime penalty. MLX quant matmul is Metal-optimized with no fast CPU kernel, + so JIT quant must be gated behind a Metal (or CUDA) backend; do not expose a CPU + quant serving path. (This supersedes an earlier CPU-only measurement.) - **Partial load may require nontrivial changes to `safemlx-lm`** (loader + model constructors currently build `0..num_hidden_layers`). Upstreaming to the fork is likely necessary. (Highest risk for the split work.) diff --git a/spikes/mlx-solo/Cargo.lock b/spikes/mlx-solo/Cargo.lock index 82a795d150..e9e35f0511 100644 --- a/spikes/mlx-solo/Cargo.lock +++ b/spikes/mlx-solo/Cargo.lock @@ -99,21 +99,6 @@ version = "0.13.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9e1b586273c5702936fe7b7d6896644d8be71e6314cfe09d3167c95f712589e8" -[[package]] -name = "bit-set" -version = "0.8.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "08807e080ed7f9d5433fa9b275196cfc35414f66a0c79d864dc51a0d825231a3" -dependencies = [ - "bit-vec", -] - -[[package]] -name = "bit-vec" -version = "0.8.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5e764a1d40d510daf35e07be9eb06e75770908c27d411ee6c92109c9840eaaf7" - [[package]] name = "bitflags" version = "2.13.1" @@ -411,16 +396,8 @@ name = "esaxx-rs" version = "0.1.10" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d817e038c30374a4bcb22f94d0a8a0e216958d4c3dcde369b1439fec4bdda6e6" - -[[package]] -name = "fancy-regex" -version = "0.17.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "72cf461f865c862bb7dc573f643dd6a2b6842f7c30b07882b56bd148cc2761b8" dependencies = [ - "bit-set", - "regex-automata", - "regex-syntax", + "cc", ] [[package]] @@ -818,6 +795,28 @@ version = "1.70.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "384b8ab6d37215f3c5301a95a4accb5d64aa607f1fcb26a11b5303878451b4fe" +[[package]] +name = "onig" +version = "6.5.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0cc3cbf698f9438986c11a880c90a6d04b9de27575afd28bbf45b154b6c709e2" +dependencies = [ + "bitflags", + "libc", + "once_cell", + "onig_sys", +] + +[[package]] +name = "onig_sys" +version = "69.9.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e68317604e77e53b85896388e1a803c1d21b74c899ec9e5e1112db90735edd7" +dependencies = [ + "cc", + "pkg-config", +] + [[package]] name = "parking_lot" version = "0.12.5" @@ -847,6 +846,12 @@ version = "1.0.15" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "57c0d7b74b563b49d38dae00a0c37d4d6de9b432382b2892f0574ddcae73fd0a" +[[package]] +name = "pkg-config" +version = "0.3.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "19f132c84eca552bf34cab8ec81f1c1dcc229b811638f9d283dceabe58c5569e" + [[package]] name = "potential_utf" version = "0.1.5" @@ -1311,12 +1316,12 @@ dependencies = [ "dary_heap", "derive_builder", "esaxx-rs", - "fancy-regex", "getrandom 0.3.4", "itertools 0.14.0", "log", "macro_rules_attribute", "monostate", + "onig", "paste", "rand", "rayon", diff --git a/spikes/mlx-solo/Cargo.toml b/spikes/mlx-solo/Cargo.toml index d1f9026e43..cc574dcb30 100644 --- a/spikes/mlx-solo/Cargo.toml +++ b/spikes/mlx-solo/Cargo.toml @@ -1,8 +1,10 @@ -# Standalone spike crate — deliberately its OWN workspace so it is NOT part of -# the mesh-llm cargo workspace (keeps MLX/CMake native deps off every other -# build). It path-depends on a sibling checkout of the safemlx fork: -# ../../../safemlx == /Users//Documents/code/safemlx -# It does not build in CI and is not meant to ship. +# Standalone spike crate — its OWN workspace, deliberately NOT part of the +# mesh-llm cargo workspace (keeps MLX/CMake native deps off every other build). +# +# Mirrors goose's MLX setup exactly (crates/goose-local-inference): +# - Metal backend on (Apple Silicon GPU), same feature list as goose +# - pristine safemlx fork, NO fork patches +# Path-depends on a sibling checkout of the fork at ../../../safemlx. [package] name = "mlx-solo-spike" version = "0.0.0" @@ -16,15 +18,14 @@ path = "src/main.rs" [workspace] [dependencies] -# CPU-only: accelerate backend, no metal (no Metal shader compiler on this box). +# Same feature set goose declares on safemlx. safemlx = { path = "../../../safemlx/safemlx", default-features = false, features = [ "accelerate", + "metal", "safetensors", ] } -# Pure-Rust regex engine (fancy-regex) to avoid extra C/C++ tokenizer builds. -safemlx-lm = { path = "../../../safemlx/safemlx-lm", default-features = false, features = [ - "fancy-regex", -] } +# Default features (onig + esaxx_fast), exactly as goose consumes it. +safemlx-lm = { path = "../../../safemlx/safemlx-lm" } anyhow = "1" clap = { version = "4", features = ["derive"] } serde_json = "1" diff --git a/spikes/mlx-solo/FINDINGS.md b/spikes/mlx-solo/FINDINGS.md index d82242e5ea..8af8dcad85 100644 --- a/spikes/mlx-solo/FINDINGS.md +++ b/spikes/mlx-solo/FINDINGS.md @@ -2,25 +2,54 @@ Validates the Phase-2 claim in `docs/design/MLX_STAGE_ENGINE_PLAN.md`: load an HF **safetensors** model in Rust via `safemlx-lm`, optionally **JIT-quantize on -load**, and generate tokens — with no GGUF and no ahead-of-time quant step. +load**, and generate tokens on **Metal** — with no GGUF and no ahead-of-time +quant step. Throwaway crate with its own `[workspace]`, deliberately **not** part of the -mesh-llm cargo workspace, so MLX/CMake native deps never touch other builds. +mesh-llm cargo workspace, so MLX/CMake native deps never touch other builds. It +mirrors goose's MLX setup (`../goose/crates/goose-local-inference`) exactly: +`safemlx` with `["accelerate", "metal", "safetensors"]`, `Device::Gpu` for +compute + `Device::Cpu` for weight staging, plain `LoadedModel::load`. ## TL;DR -- **Solo serving from raw safetensors works, in Rust, end to end.** ✅ - Qwen3-0.6B → coherent output at **18.1 tok/s** decode (CPU), no GGUF, no - pre-quant. -- **JIT quant is functional but has no fast CPU kernel.** 4-bit and 8-bit both - produce coherent output but collapse to **~0.4 tok/s** (~45× slower than - source precision). MLX quantized matmul is Metal-optimized; CPU is not a - viable quant path. **Metal is required to judge JIT-quant performance.** -- **safemlx-lm is young — hit (and fixed) a real loader bug** on the quant path - for tied-embedding checkpoints. -- **Feature wiring needs a fork/upstream change**: the published `safemlx-lm` - hard-enables the `metal` feature, so a Metal-less (CPU/CI) build is impossible - without a workspace-level `default-features = false`. +- **Solo serving from raw safetensors works on Metal, in Rust, end to end.** ✅ + Qwen3-0.6B, source precision (bf16): **321 tok/s** decode, coherent. +- **JIT quant on load is free at inference time.** JIT 4-bit and a pre-quantized + mlx-community 4-bit repo both run at **~604 tok/s** — identical. So "download + safetensors → quantize on load → serve" costs nothing versus shipping a + pre-quantized artifact. This is the workflow win, confirmed. +- **The goose baseline needs ZERO fork patches.** Source-precision serving runs + on a pristine `jbg/safemlx` checkout. +- **Going beyond the goose baseline hit two genuine bugs in the young fork** + (JIT quant + loading arbitrary mlx-community repos). Both are small and fixed + locally; they are **upstream-PR candidates for `jbg/safemlx`, not mesh-llm + drift** — and they confirm the plan's "safemlx-lm is young, expect per-family + papercuts" risk. + +## Results (Apple Silicon, Metal) + +Model: `Qwen/Qwen3-0.6B` (dense safetensors), and `mlx-community/Qwen3-0.6B-4bit` +(pre-quantized). `-n 128`, greedy. + +| Mode | model source | load | ttft | decode tok/s | patches needed | +| --- | --- | --- | --- | --- | --- | +| source precision (bf16) | dense safetensors | 1.36s | 0.876s | **321** | none | +| pre-quantized 4-bit (goose's path) | mlx-community repo | 0.15s | 0.158s | **603** | mode-default fix* | +| **JIT 4-bit on load** | dense safetensors | 0.27s | 0.010s | **604** | lm_head fix* | + +*The pre-quantized and JIT rows each needed one small fork fix (see below); the +source-precision row — the actual goose baseline — needed none. + +Key reading: **JIT (604) ≈ pre-quantized (603)**. Quantizing on load is not a +runtime tax; it happens during the (already fast) load. So the "no wait for a +published GGUF/quant, just serve the safetensors" story holds with no inference +penalty on Metal. + +> Historical note: an earlier CPU-only run (no Metal compiler installed) showed +> 18 tok/s source / 0.4 tok/s quantized. That 0.4 was a pure CPU-kernel artifact +> (MLX quant matmul is Metal-optimized). It is not representative and has been +> superseded by the Metal numbers above. CPU is not a serving path. ## Environment @@ -28,103 +57,93 @@ mesh-llm cargo workspace, so MLX/CMake native deps never touch other builds. | --- | --- | | Machine | Apple Silicon (arm64), macOS 26.5 | | Toolchain | rustc 1.97, cmake 4.4, Apple clang 21 | -| Metal shader compiler | **absent** (CommandLineTools only, no full Xcode) → CPU-only build | -| MLX backend | `accelerate` (CPU). MLX core `v0.32.0` via safemlx-sys FetchContent | -| safemlx fork | `jbg/safemlx` @ `0502a19` + 3 local edits (below) | -| Model | `Qwen/Qwen3-0.6B` (dense, safetensors, `tie_word_embeddings: true`) | - -## Results - -Command shape: -``` -mlx-solo --model -n 64 [--quantize 4|8] "" -``` - -| Mode | load | ttft | decode tok/s | output | -| --- | --- | --- | --- | --- | -| source precision | 0.30s | 0.174s | **18.1** | coherent ✅ | -| JIT 4-bit affine | 1.03s | 47.6s | **0.4** | coherent ✅ | -| JIT 8-bit affine | 1.13s | 48.4s | **0.4** | coherent ✅ | - -The quant slowdown is uniform across 4/8-bit and dominated by per-token decode -(ttft is essentially the first decode step), which is the signature of a missing -fast CPU quant-matmul kernel rather than a load-time cost. +| Metal | Xcode 26.6 (already installed) + `MetalToolchain` component (688 MB, pulled via `xcodebuild -downloadComponent MetalToolchain`); `metal` 32023.883. Point builds at it with `DEVELOPER_DIR=/Applications/Xcode.app/Contents/Developer` — no sudo, no re-install. | +| MLX backend | `metal` + `accelerate`; MLX core `v0.32.0` via safemlx-sys FetchContent | +| safemlx fork | `jbg/safemlx` @ `0502a19`; pristine for source precision, +2 tiny fixes for quant coverage | +| Models | `Qwen/Qwen3-0.6B`, `mlx-community/Qwen3-0.6B-4bit` | ## Build notes -- **MLX C++ builds cleanly under cmake 4.4, CPU-only** — this was the #1 feared - risk (cmake 4.x deprecates old-CMake compatibility). `libmlx.a` + `libmlxc.a` - produced, no policy errors. Cold build ~1m19s; warm rebuild after a Rust-only - change ~6.7s (MLX archives cached in the shared target dir). -- Only harmless linker warnings (`object file … has version 26.5.0, which is +- **MLX C++ builds cleanly with the Metal backend under cmake 4.4.** Cold build + (compiles MLX C++ + metallib shaders) ~2m05s; warm rebuild after a Rust-only + change ~6.7s (MLX archives cached in the target dir). +- Only harmless linker warnings (`object file … has version 26.5.x, which is newer than target minimum of 11.0.0`). +- The build needs `DEVELOPER_DIR` pointed at full Xcode so `safemlx-sys`'s + `xcrun -find metal` resolves the real compiler (CommandLineTools alone is not + enough; its `metal` is a stub that errors at runtime until the toolchain + component is installed). -## Issues found - -### 1. Published safemlx-lm hard-enables `metal` (build blocker off-Metal) +## Issues found — both upstream-PR candidates, not drift -`safemlx-lm`'s dependency on `safemlx` did not disable default features, and a -member-level `default-features = false` is **ignored** by Cargo (it warns and -requires the setting at the workspace root). Effect: any consumer inherits the -`metal` feature and the build `panic!`s when the Metal compiler is absent -(CI, Metal-less macOS, or a CPU-only lane). +The goose baseline (source-precision serving) needs no changes. The two fixes +below are only needed to go *beyond* goose's usage — JIT quant and loading +arbitrary published quantized repos — which is in-scope because that is the +"serve any safetensors model, quantized on load" workflow the plan proposes. +They are small, isolated, and should be PR'd to `jbg/safemlx`. -Fix in this spike (in the fork): -- root `Cargo.toml`: `safemlx = { …, default-features = false }` -- `safemlx-lm/Cargo.toml`: `safemlx = { workspace = true, default-features = false, features = ["accelerate", "safetensors"] }` - -**Plan implication:** the `MlxStageEngine` will need MLX backend selection as an -explicit cargo feature (`metal` / `accelerate` / `cuda`), which means either an -upstream PR to `safemlx-lm` or carrying a small fork. Real integration cost, not -a blocker. - -### 2. Tied-embedding `lm_head.weight` fails strict load on the quant path (qwen3) +### 1. Tied-embedding `lm_head.weight` fails the *quantized* strict loader (qwen3) `Qwen3-0.6B` sets `tie_word_embeddings: true` but the checkpoint still ships a -redundant `lm_head.weight`. The **dense** load uses a lenient loader and -tolerates it; the **quantized** load uses `load_safetensors_dir_quantized_strict` -with a bare `StrictLoadConfig::default()` and rejects it as an unused tensor: +redundant `lm_head.weight`. The dense loader is lenient and tolerates it; the +quantized loader uses a bare `StrictLoadConfig::default()` and rejects it: ``` strict weight-load validation failed: 0 missing, 1 unused unused: lm_head.weight ``` -Fix in this spike (in the fork, `safemlx-lm/src/models/qwen3.rs` -`load_qwen3_model_quantized`): +Fix (`safemlx-lm/src/models/qwen3.rs`, `load_qwen3_model_quantized`): ```rust let config = StrictLoadConfig::default().allow_unused_prefix("lm_head."); ``` Harmless when untied (then `lm_head.weight` is a loaded param, not unused). -**Plan implication:** confirms the "safemlx-lm is young; each model family is -bespoke and separately certified" risk. Expect per-family loader papercuts. +### 2. `WeightQuantization` requires a `mode` field many mlx-community repos omit + +`mlx-community/Qwen3-0.6B-4bit`'s `config.json` has +`"quantization": {"group_size": 64, "bits": 4}` with no `mode`. The fork's +`WeightQuantizationMetadata` makes `mode` mandatory, so the load fails with +`missing field 'mode'`. mlx-lm itself defaults a missing mode to `affine`. + +Fix (`safemlx-lm/src/quantization.rs`): +```rust +#[serde(default = "default_affine_mode_string")] +mode: String, +// ... +fn default_affine_mode_string() -> String { "affine".to_string() } +``` ## Reproduce -1. Sibling checkout of the fork at `../safemlx` (repo-relative: - `/Users//Documents/code/safemlx`), base `0502a19`, with the two Cargo - edits and the qwen3 one-liner above. -2. Model dir with `config.json`, `tokenizer*.json`, `model.safetensors` - (e.g. `Qwen/Qwen3-0.6B`). -3. Build/run (own workspace, so it will trigger a one-time MLX C++ build): +1. Install the Metal toolchain if `xcrun metal --version` errors: + `DEVELOPER_DIR=/Applications/Xcode.app/Contents/Developer \ + /Applications/Xcode.app/Contents/Developer/usr/bin/xcodebuild -downloadComponent MetalToolchain` +2. Sibling checkout of the fork at `../safemlx`, base `0502a19`, with the two + fixes above (only needed for the quant rows). +3. Model dirs (HF safetensors) for `Qwen/Qwen3-0.6B` and, optionally, + `mlx-community/Qwen3-0.6B-4bit`. +4. Build/run with Xcode selected for this shell: ``` + export DEVELOPER_DIR=/Applications/Xcode.app/Contents/Developer cd spikes/mlx-solo cargo build --release - ./target/release/mlx-solo --model -n 64 "..." - ./target/release/mlx-solo --model -n 64 --quantize 4 "..." + ./target/release/mlx-solo --model -n 128 "..." # source precision + ./target/release/mlx-solo --model -n 128 "..." # pre-quantized + ./target/release/mlx-solo --model -q 4 -n 128 "..." # JIT 4-bit ``` - Using a shared `CARGO_TARGET_DIR` avoids rebuilding MLX across crates. + A shared `CARGO_TARGET_DIR` avoids rebuilding MLX across crates. ## What this de-risks / what it does not De-risked: -- safemlx-lm as a Rust solo engine driving load + generate. -- The "no GGUF, no pre-quant, serve raw safetensors" workflow claim. -- MLX native build under this toolchain (cmake 4.4). - -Still open (needs Metal): -- Any performance judgement of JIT quant (CPU quant path is not representative). -- Metal throughput vs the llama.cpp backend. -- Everything staged/split (this spike is single-stage, whole-model) — the - partial-load and boundary-fence go/no-go spikes are unchanged. +- safemlx-lm as a Rust solo engine on **Metal** (321 tok/s bf16, ~604 tok/s 4-bit). +- The "no GGUF, no pre-quant, serve raw safetensors" workflow — including that + **JIT quant is free at inference time** (≈ pre-quantized). +- MLX native build with the Metal backend under this toolchain (cmake 4.4). + +Still open: +- Larger models + throughput vs the llama.cpp backend on the same hardware. +- Everything staged/split — this spike is single-stage, whole-model. The + partial-load and boundary-fence go/no-go spikes (plan §8) are unchanged. +- Upstreaming the two fixes to `jbg/safemlx` (or carrying a thin fork). diff --git a/spikes/mlx-solo/src/main.rs b/spikes/mlx-solo/src/main.rs index c6bb9cf963..6915ec1ada 100644 --- a/spikes/mlx-solo/src/main.rs +++ b/spikes/mlx-solo/src/main.rs @@ -52,10 +52,10 @@ struct Cli { fn main() -> Result<()> { let args = Cli::parse(); - // CPU device: no Metal compiler on this box, so the GPU backend is absent. - let device = Device::new(DeviceType::Cpu, 0); - let stream = Stream::new_with_device(&device); - let weights_stream = Stream::new_with_device(&device); + // Mirror goose (crates/goose-local-inference/src/mlx.rs): Metal GPU stream for + // compute, CPU stream for weight staging. + let stream = Stream::new_with_device(&Device::new(DeviceType::Gpu, 0)); + let weights_stream = Stream::new_with_device(&Device::new(DeviceType::Cpu, 0)); let load_options = match args.quantize { Some(bits) => { From 558be443d9246b98db7fab4746433d7a557c0170 Mon Sep 17 00:00:00 2001 From: michaelneale Date: Thu, 16 Jul 2026 14:12:43 +1000 Subject: [PATCH 04/37] docs+spike: correct MLX platform framing; prove fork-free is not viable 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. --- docs/design/MLX_STAGE_ENGINE_PLAN.md | 70 +- spikes/mlx-solo-published/Cargo.lock | 1572 +++++++++++++++++++++++++ spikes/mlx-solo-published/Cargo.toml | 25 + spikes/mlx-solo-published/src/main.rs | 145 +++ 4 files changed, 1790 insertions(+), 22 deletions(-) create mode 100644 spikes/mlx-solo-published/Cargo.lock create mode 100644 spikes/mlx-solo-published/Cargo.toml create mode 100644 spikes/mlx-solo-published/src/main.rs diff --git a/docs/design/MLX_STAGE_ENGINE_PLAN.md b/docs/design/MLX_STAGE_ENGINE_PLAN.md index 2ac7ca7b12..31710c9db8 100644 --- a/docs/design/MLX_STAGE_ENGINE_PLAN.md +++ b/docs/design/MLX_STAGE_ENGINE_PLAN.md @@ -486,9 +486,9 @@ how goose gates it. Lean by construction, for the platforms that don't use it. | iOS / tvOS / visionOS | ✅ Metal | | Linux x86_64 / aarch64 | ✅ CPU | | Linux + NVIDIA | ✅ CUDA (the `cuda`/`nccl` features **panic** on non-Linux) | -| Linux + AMD (ROCm) | ❌ | -| Vulkan (any) | ❌ | -| Windows | ❌ | +| Linux + AMD (ROCm) | ⏳ not today — large **active but unmerged** upstream experiment (see below) | +| Vulkan (any) | ❌ upstream *wishlist* only, no implementation | +| Windows | ❌ (some `if(WIN32)` scaffolding in vendored `mlx-c`, no working backend) | **Coverage is expanding, and safemlx tracks it fast.** `jbg/safemlx` is very active (103 commits, latest 2026-07-15) and pins a recent MLX core (`v0.32.0`). @@ -497,22 +497,34 @@ It wires in new backends quickly: the `Add CUDA support` commit landed a full and there is `if(WIN32)` DLL-export scaffolding in the vendored `mlx-c`. So the matrix above is a **snapshot, not a ceiling**. -**But the gaps are gated by MLX upstream, not by safemlx.** safemlx can only -expose backends `ml-explore/mlx` itself provides. MLX's backend line is -**CPU + Metal + CUDA** — which is exactly why CUDA appeared here. There is **no -ROCm or Vulkan backend in MLX upstream**, and no in-repo signal that safemlx adds -them independently. So: Windows is the most plausible next addition -(scaffolding + working Linux/CUDA); ROCm/Vulkan depend on an upstream decision -with no current signal either way. - -**Strategic consequence (unchanged).** **Today**, the ROCm / Vulkan / Windows -gaps mean **MLX cannot be Skippy's sole engine** — which reinforces (not changes) -the plan: MLX is an **additive, feature+cfg-gated second engine**, -Apple-Silicon-first (with Linux/CUDA as a real second target), while llama.cpp -stays the cross-platform default. And the most durable reason to keep llama.cpp -is **not** platform coverage (which may well close as MLX upstream grows) but its -GGUF/imatrix quant maturity and the existing patch-queue investment — those would -argue against removing it even if MLX's backend matrix later caught up. +**The gaps are gated by MLX upstream, and there are *two* gates.** safemlx does +not build backends of its own — every one of its non-`main` branches is model / +runtime / quant work, not hardware work, and `forks_count`/`network_count` are 0 +with no open PRs. A new backend must therefore (1) land in `ml-explore/mlx` +(C++), and only then (2) be wired through safemlx — exactly the sequence CUDA +followed (`Add CUDA support` was safemlx *exposing* an upstream backend, not +authoring one). So hardware coverage tracks upstream MLX, delayed by the safemlx +wiring step. + +**ROCm is real but not bankable yet.** Upstream MLX has a large, active AMD/ROCm +effort — PR **#2300 "[Experiment] ROCm backend"** (≈449 commits, +45k lines, open +~13 months, updated as of this writing) plus issue **#2556 "Add ROCm Support for +AMD GPUs"**. It is **unmerged and `mergeable_state: dirty`**, so it is genuine +momentum, not a shipped backend. Vulkan is only an upstream *wishlist* issue with +no implementation; Windows has scaffolding but no backend. Net: the matrix is +**expanding (CPU → Metal → CUDA, ROCm being actively attempted upstream)**, so +treat it as a moving target — but do not plan around ROCm/Vulkan/Windows until +they both merge upstream **and** appear in safemlx. + +**Strategic consequence.** **Today**, the ROCm / Vulkan / Windows gaps mean +**MLX cannot be Skippy's sole engine** — which reinforces (not changes) the plan: +MLX is an **additive, feature+cfg-gated second engine**, strongest on Apple +Silicon (with Linux/CUDA a real second target, and AMD plausibly later), while +llama.cpp stays the cross-platform default. Crucially, even in the optimistic +world where MLX gains ROCm/Vulkan, the durable reason to keep llama.cpp is **not** +platform coverage but its **GGUF/imatrix k-quant maturity** and the existing +**patch-queue investment** — those are the sticky arguments; hardware coverage is +the reversible one. --- @@ -669,9 +681,23 @@ Spikes 1 and 2 are more decisive than any standalone token/s benchmark. them out of early phases. - **Two artifact pipelines** add storage + certification cost; mitigate with a single canonical BF16 source and reproducible derivation. -- **safemlx maturity/maintenance** (external fork of mlx-rs) — the `MlxStageEngine` - will likely carry a small fork or need upstream PRs (backend feature exposure, - loader fixes, and eventually `forward_range`/partial-load). Pin carefully. +- **safemlx supply chain — pin to a git rev, not a crates.io version (confirmed + this session).** The published crates collide version strings with the fork + HEAD: crates.io `safemlx-lm 0.4.1` is a *different, older* codebase than the + fork's `0.4.1` (851 vs 2221 lines in `qwen3.rs`), because the fork develops on + a fixed version without bumping. A fork-free build against published crates + **compiled and ran but produced gibberish for Qwen3 source precision and + crashed on a pre-quantized repo** (`rms_norm` size mismatch) — the working + dense-Qwen3/Llama + JIT-quant code exists only in unpublished fork HEAD. So + MLX-for-Skippy must **pin a specific git commit** of `jbg/safemlx` (and carry + the small loader fixes until upstreamed), or coordinate a real published + release. This makes "track upstream + pin + possibly patch" a **standing cost**, + not a one-off. +- **Hardware coverage is a moving target with two gates.** New backends must land + in upstream `ml-explore/mlx` *then* be wired through safemlx (which authors no + backends itself). ROCm is an active-but-unmerged upstream experiment (#2300); + Vulkan is wishlist-only; Windows has scaffolding but no backend. Do not plan + around AMD/Vulkan/Windows until both gates clear. - **Compat discipline:** MLX must stay additive (feature-probe + gossip capability); homogeneous chains by default; mixed-engine only when certified. diff --git a/spikes/mlx-solo-published/Cargo.lock b/spikes/mlx-solo-published/Cargo.lock new file mode 100644 index 0000000000..47c5a538eb --- /dev/null +++ b/spikes/mlx-solo-published/Cargo.lock @@ -0,0 +1,1572 @@ +# This file is automatically @generated by Cargo. +# It is not intended for manual editing. +version = 4 + +[[package]] +name = "ahash" +version = "0.8.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5a15f179cd60c4584b8a8c596927aadc462e27f2ca70c04e0071964a73ba7a75" +dependencies = [ + "cfg-if", + "getrandom 0.3.4", + "once_cell", + "serde", + "version_check", + "zerocopy", +] + +[[package]] +name = "aho-corasick" +version = "1.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ddd31a130427c27518df266943a5308ed92d4b226cc639f5a8f1002816174301" +dependencies = [ + "memchr", +] + +[[package]] +name = "allocator-api2" +version = "0.2.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "683d7910e743518b0e34f1186f92494becacb047c7b6bf616c96772180fef923" + +[[package]] +name = "anstream" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "824a212faf96e9acacdbd09febd34438f8f711fb84e09a8916013cd7815ca28d" +dependencies = [ + "anstyle", + "anstyle-parse", + "anstyle-query", + "anstyle-wincon", + "colorchoice", + "is_terminal_polyfill", + "utf8parse", +] + +[[package]] +name = "anstyle" +version = "1.0.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "940b3a0ca603d1eade50a4846a2afffd5ef57a9feac2c0e2ec2e14f9ead76000" + +[[package]] +name = "anstyle-parse" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "52ce7f38b242319f7cabaa6813055467063ecdc9d355bbb4ce0c68908cd8130e" +dependencies = [ + "utf8parse", +] + +[[package]] +name = "anstyle-query" +version = "1.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "40c48f72fd53cd289104fc64099abca73db4166ad86ea0b4341abe65af83dadc" +dependencies = [ + "windows-sys", +] + +[[package]] +name = "anstyle-wincon" +version = "3.0.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "291e6a250ff86cd4a820112fb8898808a366d8f9f58ce16d1f538353ad55747d" +dependencies = [ + "anstyle", + "once_cell_polyfill", + "windows-sys", +] + +[[package]] +name = "anyhow" +version = "1.0.103" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2a4385e2e34eb35d6b3efe798b9eb88096925d87726c0798709bf56d9ed84af3" + +[[package]] +name = "autocfg" +version = "1.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f2032f911046de80f0a198e0901378627c33f59ea0ac00e363d481118bd70a53" + +[[package]] +name = "base64" +version = "0.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9e1b586273c5702936fe7b7d6896644d8be71e6314cfe09d3167c95f712589e8" + +[[package]] +name = "bitflags" +version = "2.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b588b76d00fde79687d7646a9b5bdf3cc0f655e0bbd080335a95d7e96f3587da" + +[[package]] +name = "bytemuck" +version = "1.25.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d6aedf8ae72766347502cf3cb4f41cf5e9cc37d28bee90f1fdaaae15f9cf9424" + +[[package]] +name = "castaway" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dec551ab6e7578819132c713a93c022a05d60159dc86e7a7050223577484c55a" +dependencies = [ + "rustversion", +] + +[[package]] +name = "cc" +version = "1.2.67" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e17dd265a7d0f31ef544e1b20e03add05d3b45b491b633b10d67145d2acc1a38" +dependencies = [ + "find-msvc-tools", + "shlex", +] + +[[package]] +name = "cfg-if" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" + +[[package]] +name = "clap" +version = "4.6.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dd059f9da4f5c36b3787f65d38ccaab1cc315f07b01f89abc8359ee6a8205011" +dependencies = [ + "clap_builder", + "clap_derive", +] + +[[package]] +name = "clap_builder" +version = "4.6.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f09628afdcc538b57f3c6341e9c8e9970f18e4a481690a64974d7023bd33548b" +dependencies = [ + "anstream", + "anstyle", + "clap_lex", + "strsim", +] + +[[package]] +name = "clap_derive" +version = "4.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f2ce8604710f6733aa641a2b3731eaa1e8b3d9973d5e3565da11800813f997a9" +dependencies = [ + "heck", + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "clap_lex" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8d4a3bb8b1e0c1050499d1815f5ab16d04f0959b233085fb31653fbfc9d98f9" + +[[package]] +name = "cmake" +version = "0.1.58" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c0f78a02292a74a88ac736019ab962ece0bc380e3f977bf72e376c5d78ff0678" +dependencies = [ + "cc", +] + +[[package]] +name = "colorchoice" +version = "1.0.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1d07550c9036bf2ae0c684c4297d503f838287c83c53686d05370d0e139ae570" + +[[package]] +name = "compact_str" +version = "0.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9dfdd1c2274d9aa354115b09dc9a901d6c5576818cdf70d14cae2bdb47df00ab" +dependencies = [ + "castaway", + "cfg-if", + "itoa", + "rustversion", + "ryu", + "serde", + "static_assertions", +] + +[[package]] +name = "crossbeam-deque" +version = "0.8.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5181e0de7b61eb03a81e347d6dd8797bae9da5146707b51077e2d71a54ec0ceb" +dependencies = [ + "crossbeam-epoch", + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-epoch" +version = "0.9.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2d6914041f254d6e9176c01941b21115dcfb7089e55135a35411081bd106ef3f" +dependencies = [ + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-utils" +version = "0.8.22" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "61803da095bee82a81bb1a452ecc25d3b2f1416d1897eb86430c6159ef717c17" + +[[package]] +name = "crunchy" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "460fbee9c2c2f33933d720630a6a0bac33ba7053db5344fac858d4b8952d77d5" + +[[package]] +name = "daachorse" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6f55d7153ba3b507595872a3874803f07a8a81d1e888abed8e5db7da0597d6e2" + +[[package]] +name = "darling" +version = "0.20.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc7f46116c46ff9ab3eb1597a45688b6715c6e628b5c133e288e709a29bcb4ee" +dependencies = [ + "darling_core 0.20.11", + "darling_macro 0.20.11", +] + +[[package]] +name = "darling" +version = "0.23.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "25ae13da2f202d56bd7f91c25fba009e7717a1e4a1cc98a76d844b65ae912e9d" +dependencies = [ + "darling_core 0.23.0", + "darling_macro 0.23.0", +] + +[[package]] +name = "darling_core" +version = "0.20.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0d00b9596d185e565c2207a0b01f8bd1a135483d02d9b7b0a54b11da8d53412e" +dependencies = [ + "fnv", + "ident_case", + "proc-macro2", + "quote", + "strsim", + "syn", +] + +[[package]] +name = "darling_core" +version = "0.23.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9865a50f7c335f53564bb694ef660825eb8610e0a53d3e11bf1b0d3df31e03b0" +dependencies = [ + "ident_case", + "proc-macro2", + "quote", + "strsim", + "syn", +] + +[[package]] +name = "darling_macro" +version = "0.20.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc34b93ccb385b40dc71c6fceac4b2ad23662c7eeb248cf10d529b7e055b6ead" +dependencies = [ + "darling_core 0.20.11", + "quote", + "syn", +] + +[[package]] +name = "darling_macro" +version = "0.23.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ac3984ec7bd6cfa798e62b4a642426a5be0e68f9401cfc2a01e3fa9ea2fcdb8d" +dependencies = [ + "darling_core 0.23.0", + "quote", + "syn", +] + +[[package]] +name = "dary_heap" +version = "0.3.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8b1e3a325bc115f096c8b77bbf027a7c2592230e70be2d985be950d3d5e60ebe" +dependencies = [ + "serde", +] + +[[package]] +name = "derive_builder" +version = "0.20.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "507dfb09ea8b7fa618fcf76e953f4f5e192547945816d5358edffe39f6f94947" +dependencies = [ + "derive_builder_macro", +] + +[[package]] +name = "derive_builder_core" +version = "0.20.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2d5bcf7b024d6835cfb3d473887cd966994907effbe9227e8c8219824d06c4e8" +dependencies = [ + "darling 0.20.11", + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "derive_builder_macro" +version = "0.20.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ab63b0e2bf4d5928aff72e83a7dace85d7bba5fe12dcc3c5a572d78caffd3f3c" +dependencies = [ + "derive_builder_core", + "syn", +] + +[[package]] +name = "displaydoc" +version = "0.2.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ac70aa55017e108007fbaf5aa0f54b021c98f92ff8af59d42eda9da96e3dd4f" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "dyn-clone" +version = "1.0.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d0881ea181b1df73ff77ffaaf9c7544ecc11e82fba9b5f27b262a3c73a332555" + +[[package]] +name = "either" +version = "1.16.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91622ff5e7162018101f2fea40d6ebf4a78bbe5a49736a2020649edf9693679e" + +[[package]] +name = "equivalent" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "877a4ace8713b0bcf2a4e7eec82529c029f1d0619886d18145fea96c3ffe5c0f" + +[[package]] +name = "errno" +version = "0.3.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" +dependencies = [ + "libc", + "windows-sys", +] + +[[package]] +name = "esaxx-rs" +version = "0.1.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d817e038c30374a4bcb22f94d0a8a0e216958d4c3dcde369b1439fec4bdda6e6" +dependencies = [ + "cc", +] + +[[package]] +name = "fastrand" +version = "2.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9f1f227452a390804cdb637b74a86990f2a7d7ba4b7d5693aac9b4dd6defd8d6" + +[[package]] +name = "find-msvc-tools" +version = "0.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5baebc0774151f905a1a2cc41989300b1e6fbb29aff0ceffa1064fdd3088d582" + +[[package]] +name = "fnv" +version = "1.0.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3f9eec918d3f24069decb9af1554cad7c880e2da24a9afd88aca000531ab82c1" + +[[package]] +name = "foldhash" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "77ce24cb58228fbb8aa041425bb1050850ac19177686ea6e0f41a70416f56fdb" + +[[package]] +name = "getrandom" +version = "0.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "899def5c37c4fd7b2664648c28120ecec138e4d395b459e5ca34f9cce2dd77fd" +dependencies = [ + "cfg-if", + "libc", + "r-efi 5.3.0", + "wasip2", +] + +[[package]] +name = "getrandom" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "300e883d756b2e4ec94e02791f39b04b522276138852cfc41d9fb7e904106099" +dependencies = [ + "cfg-if", + "libc", + "r-efi 6.0.0", +] + +[[package]] +name = "half" +version = "2.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" +dependencies = [ + "cfg-if", + "crunchy", + "zerocopy", +] + +[[package]] +name = "hashbrown" +version = "0.16.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "841d1cc9bed7f9236f321df977030373f4a4163ae1a7dbfe1a51a2c1a51d9100" +dependencies = [ + "allocator-api2", + "equivalent", + "foldhash", + "serde", + "serde_core", +] + +[[package]] +name = "hashbrown" +version = "0.17.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a" + +[[package]] +name = "heck" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" + +[[package]] +name = "icu_collections" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2984d1cd16c883d7935b9e07e44071dca8d917fd52ecc02c04d5fa0b5a3f191c" +dependencies = [ + "displaydoc", + "potential_utf", + "utf8_iter", + "yoke", + "zerofrom", + "zerovec", +] + +[[package]] +name = "icu_locale_core" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "92219b62b3e2b4d88ac5119f8904c10f8f61bf7e95b640d25ba3075e6cac2c29" +dependencies = [ + "displaydoc", + "litemap", + "tinystr", + "writeable", + "zerovec", +] + +[[package]] +name = "icu_normalizer" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c56e5ee99d6e3d33bd91c5d85458b6005a22140021cc324cea84dd0e72cff3b4" +dependencies = [ + "icu_collections", + "icu_provider", + "smallvec", + "zerovec", +] + +[[package]] +name = "icu_properties" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bee3b67d0ea5c2cca5003417989af8996f8604e34fb9ddf96208a033901e70de" +dependencies = [ + "icu_collections", + "icu_locale_core", + "icu_provider", + "zerotrie", + "zerovec", +] + +[[package]] +name = "icu_provider" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "139c4cf31c8b5f33d7e199446eff9c1e02decfc2f0eec2c8d71f65befa45b421" +dependencies = [ + "displaydoc", + "icu_locale_core", + "yoke", + "zerofrom", + "zerovec", +] + +[[package]] +name = "ident_case" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b9e0384b61958566e926dc50660321d12159025e767c18e043daf26b70104c39" + +[[package]] +name = "idna_adapter" +version = "1.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cb68373c0d6620ef8105e855e7745e18b0d00d3bdb07fb532e434244cdb9a714" +dependencies = [ + "icu_normalizer", + "icu_properties", +] + +[[package]] +name = "indexmap" +version = "2.14.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d466e9454f08e4a911e14806c24e16fba1b4c121d1ea474396f396069cf949d9" +dependencies = [ + "equivalent", + "hashbrown 0.17.1", +] + +[[package]] +name = "is_terminal_polyfill" +version = "1.70.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a6cb138bb79a146c1bd460005623e142ef0181e3d0219cb493e02f7d08a35695" + +[[package]] +name = "itertools" +version = "0.14.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2b192c782037fadd9cfa75548310488aabdbf3d2da73885b31bd0abd03351285" +dependencies = [ + "either", +] + +[[package]] +name = "itertools" +version = "0.15.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8b4baf93f58d4425749ca49a51c50ebab072c5df6994d08fed93541c331481dc" +dependencies = [ + "either", +] + +[[package]] +name = "itoa" +version = "1.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" + +[[package]] +name = "libc" +version = "0.2.186" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68ab91017fe16c622486840e4c83c9a37afeff978bd239b5293d61ece587de66" + +[[package]] +name = "linux-raw-sys" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32a66949e030da00e8c7d4434b251670a91556f4144941d37452769c25d58a53" + +[[package]] +name = "litemap" +version = "0.8.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "92daf443525c4cce67b150400bc2316076100ce0b3686209eb8cf3c31612e6f0" + +[[package]] +name = "lock_api" +version = "0.4.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "224399e74b87b5f3557511d98dff8b14089b3dadafcab6bb93eab67d3aace965" +dependencies = [ + "scopeguard", +] + +[[package]] +name = "log" +version = "0.4.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0ceec5bc11778974d1bcb055b18002eba7f4b3518b6a0081b3af5f21666da9ad" + +[[package]] +name = "mach-sys" +version = "0.5.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "48460c2e82a3a0de197152fdf8d2c2d5e43adc501501553e439bf2156e6f87c7" +dependencies = [ + "fastrand", +] + +[[package]] +name = "macro_rules_attribute" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "65049d7923698040cd0b1ddcced9b0eb14dd22c5f86ae59c3740eab64a676520" +dependencies = [ + "macro_rules_attribute-proc_macro", + "paste", +] + +[[package]] +name = "macro_rules_attribute-proc_macro" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "670fdfda89751bc4a84ac13eaa63e205cf0fd22b4c9a5fbfa085b63c1f1d3a30" + +[[package]] +name = "memchr" +version = "2.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf8baf1c55e62ffcace7a9f06f4bd9cd3f0c4beb022d3b367256b91b87513d98" + +[[package]] +name = "memo-map" +version = "0.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "38d1115007560874e373613744c6fba374c17688327a71c1476d1a5954cc857b" + +[[package]] +name = "minijinja" +version = "2.21.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cb3d648e68cea56d9858d535ee28f9538404e2dd8cb08ed0bd05dca379477f39" +dependencies = [ + "memo-map", + "serde", +] + +[[package]] +name = "minijinja-contrib" +version = "2.21.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85342f6fac0be8ccd5bd00d9066be538f34f393f577b75d81b17c8398a6b43bb" +dependencies = [ + "minijinja", + "serde", +] + +[[package]] +name = "minimal-lexical" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68354c5c6bd36d73ff3feceb05efa59b6acb7626617f4962be322a825e61f79a" + +[[package]] +name = "mlx-solo-published" +version = "0.0.0" +dependencies = [ + "anyhow", + "clap", + "safemlx", + "safemlx-lm", + "serde_json", +] + +[[package]] +name = "monostate" +version = "0.1.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3341a273f6c9d5bef1908f17b7267bbab0e95c9bf69a0d4dcf8e9e1b2c76ef67" +dependencies = [ + "monostate-impl", + "serde", + "serde_core", +] + +[[package]] +name = "monostate-impl" +version = "0.1.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e4db6d5580af57bf992f59068d4ea26fd518574ff48d7639b255a36f9de6e7e9" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "nom" +version = "7.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d273983c5a657a70a3e8f2a01329822f3b8c8172b73826411a55751e404a0a4a" +dependencies = [ + "memchr", + "minimal-lexical", +] + +[[package]] +name = "num-complex" +version = "0.4.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "73f88a1307638156682bada9d7604135552957b7818057dcef22705b4d509495" +dependencies = [ + "num-traits", +] + +[[package]] +name = "num-traits" +version = "0.2.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "071dfc062690e90b734c0b2273ce72ad0ffa95f0c74596bc250dcfd960262841" +dependencies = [ + "autocfg", +] + +[[package]] +name = "num_enum" +version = "0.7.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5d0bca838442ec211fa11de3a8b0e0e8f3a4522575b5c4c06ed722e005036f26" +dependencies = [ + "num_enum_derive", + "rustversion", +] + +[[package]] +name = "num_enum_derive" +version = "0.7.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "680998035259dcfcafe653688bf2aa6d3e2dc05e98be6ab46afb089dc84f1df8" +dependencies = [ + "proc-macro-crate", + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "once_cell" +version = "1.21.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" + +[[package]] +name = "once_cell_polyfill" +version = "1.70.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "384b8ab6d37215f3c5301a95a4accb5d64aa607f1fcb26a11b5303878451b4fe" + +[[package]] +name = "onig" +version = "6.5.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0cc3cbf698f9438986c11a880c90a6d04b9de27575afd28bbf45b154b6c709e2" +dependencies = [ + "bitflags", + "libc", + "once_cell", + "onig_sys", +] + +[[package]] +name = "onig_sys" +version = "69.9.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e68317604e77e53b85896388e1a803c1d21b74c899ec9e5e1112db90735edd7" +dependencies = [ + "cc", + "pkg-config", +] + +[[package]] +name = "parking_lot" +version = "0.12.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "93857453250e3077bd71ff98b6a65ea6621a19bb0f559a85248955ac12c45a1a" +dependencies = [ + "lock_api", + "parking_lot_core", +] + +[[package]] +name = "parking_lot_core" +version = "0.9.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2621685985a2ebf1c516881c026032ac7deafcda1a2c9b7850dc81e3dfcb64c1" +dependencies = [ + "cfg-if", + "libc", + "redox_syscall", + "smallvec", + "windows-link", +] + +[[package]] +name = "paste" +version = "1.0.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "57c0d7b74b563b49d38dae00a0c37d4d6de9b432382b2892f0574ddcae73fd0a" + +[[package]] +name = "pkg-config" +version = "0.3.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "19f132c84eca552bf34cab8ec81f1c1dcc229b811638f9d283dceabe58c5569e" + +[[package]] +name = "potential_utf" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0103b1cef7ec0cf76490e969665504990193874ea05c85ff9bab8b911d0a0564" +dependencies = [ + "zerovec", +] + +[[package]] +name = "ppv-lite86" +version = "0.2.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85eae3c4ed2f50dcfe72643da4befc30deadb458a9b590d720cde2f2b1e97da9" +dependencies = [ + "zerocopy", +] + +[[package]] +name = "proc-macro-crate" +version = "3.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e67ba7e9b2b56446f1d419b1d807906278ffa1a658a8a5d8a39dcb1f5a78614f" +dependencies = [ + "toml_edit", +] + +[[package]] +name = "proc-macro2" +version = "1.0.106" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8fd00f0bb2e90d81d1044c2b32617f68fcb9fa3bb7640c23e9c748e53fb30934" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "quote" +version = "1.0.46" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dfbc457d0c7a0759a614551b11a6409e5951f6c7537be1f1b7682b9ae9230368" +dependencies = [ + "proc-macro2", +] + +[[package]] +name = "r-efi" +version = "5.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "69cdb34c158ceb288df11e18b4bd39de994f6657d83847bdffdbd7f346754b0f" + +[[package]] +name = "r-efi" +version = "6.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8dcc9c7d52a811697d2151c701e0d08956f92b0e24136cf4cf27b57a6a0d9bf" + +[[package]] +name = "rand" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b9ef1d0d795eb7d84685bca4f72f3649f064e6641543d3a8c415898726a57b41" +dependencies = [ + "rand_chacha", + "rand_core", +] + +[[package]] +name = "rand_chacha" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d3022b5f1df60f26e1ffddd6c66e8aa15de382ae63b3a0c1bfc0e4d3e3f325cb" +dependencies = [ + "ppv-lite86", + "rand_core", +] + +[[package]] +name = "rand_core" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "76afc826de14238e6e8c374ddcc1fa19e374fd8dd986b0d2af0d02377261d83c" +dependencies = [ + "getrandom 0.3.4", +] + +[[package]] +name = "rayon" +version = "1.12.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fb39b166781f92d482534ef4b4b1b2568f42613b53e5b6c160e24cfbfa30926d" +dependencies = [ + "either", + "rayon-core", +] + +[[package]] +name = "rayon-cond" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2964d0cf57a3e7a06e8183d14a8b527195c706b7983549cd5462d5aa3747438f" +dependencies = [ + "either", + "itertools 0.14.0", + "rayon", +] + +[[package]] +name = "rayon-core" +version = "1.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "22e18b0f0062d30d4230b2e85ff77fdfe4326feb054b9783a3460d8435c8ab91" +dependencies = [ + "crossbeam-deque", + "crossbeam-utils", +] + +[[package]] +name = "redox_syscall" +version = "0.5.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed2bf2547551a7053d6fdfafda3f938979645c44812fbfcda098faae3f1a362d" +dependencies = [ + "bitflags", +] + +[[package]] +name = "regex" +version = "1.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f020237b6c8eed93db2e2cb53c00c60a8e1bc73da7d073199a1180401450218d" +dependencies = [ + "aho-corasick", + "memchr", + "regex-automata", + "regex-syntax", +] + +[[package]] +name = "regex-automata" +version = "0.4.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8fcfdb36bda0c880c5931cdc7a2bcdc8ba4556847b9d912bca70bc94708711ad" +dependencies = [ + "aho-corasick", + "memchr", + "regex-syntax", +] + +[[package]] +name = "regex-syntax" +version = "0.8.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d6f6ff9a378485b298a5286656da665ba74413d36db0979633275d2e708145d4" + +[[package]] +name = "rustix" +version = "1.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6fe4565b9518b83ef4f91bb47ce29620ca828bd32cb7e408f0062e9930ba190" +dependencies = [ + "bitflags", + "errno", + "libc", + "linux-raw-sys", + "windows-sys", +] + +[[package]] +name = "rustversion" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf54715a573b99ac80df0bc206da022bcd442c974952c7b9720069370852e21f" + +[[package]] +name = "ryu" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9774ba4a74de5f7b1c1451ed6cd5285a32eddb5cccb8cc655a4e50009e06477f" + +[[package]] +name = "safemlx" +version = "0.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ba126836f374feaa7d6842fe4666fa249b364622ba3c1806b493ce0d5461998f" +dependencies = [ + "bytemuck", + "dyn-clone", + "half", + "itertools 0.15.0", + "libc", + "mach-sys", + "num-complex", + "num-traits", + "num_enum", + "parking_lot", + "paste", + "safemlx-internal-macros", + "safemlx-macros", + "safemlx-sys", + "safetensors", + "smallvec", + "strum", + "thiserror", +] + +[[package]] +name = "safemlx-internal-macros" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6169d5cd6a9b565bc3b0db7d91a9bbfe19999839b23c92e991fd823835d8edbd" +dependencies = [ + "darling 0.23.0", + "itertools 0.15.0", + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "safemlx-lm" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "10021db1789b1707952ae7e952c81eafd5b7b2a81e34eb47719f2a182671d4db" +dependencies = [ + "anyhow", + "clap", + "idna_adapter", + "minijinja", + "safemlx", + "safemlx-lm-utils", + "serde", + "serde_json", + "thiserror", + "tokenizers", +] + +[[package]] +name = "safemlx-lm-utils" +version = "0.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1af6547ba76673530921074ad4e7608240526f30335048b7bac077dee5d392a5" +dependencies = [ + "minijinja", + "minijinja-contrib", + "serde", + "serde_json", + "thiserror", + "tokenizers", +] + +[[package]] +name = "safemlx-macros" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6deae4f49d22783e05805575e4b676e203fcc4e4e6ad57fbb5001abdb405e257" +dependencies = [ + "darling 0.23.0", + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "safemlx-sys" +version = "0.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5240fd4c15ee5c0ff377757f8643e2a555378b4fbd5d0ab1a810efc06f8fcd6c" +dependencies = [ + "cc", + "cmake", +] + +[[package]] +name = "safetensors" +version = "0.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "79b079b829cb27a1c3c374341345ed2e8b2c0c839034522cee576c140bd7f846" +dependencies = [ + "hashbrown 0.16.1", + "libc", + "serde", + "serde_json", + "tempfile", +] + +[[package]] +name = "scopeguard" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" + +[[package]] +name = "serde" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9a8e94ea7f378bd32cbbd37198a4a91436180c5bb472411e48b5ec2e2124ae9e" +dependencies = [ + "serde_core", + "serde_derive", +] + +[[package]] +name = "serde_core" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "41d385c7d4ca58e59fc732af25c3983b67ac852c1a25000afe1175de458b67ad" +dependencies = [ + "serde_derive", +] + +[[package]] +name = "serde_derive" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d540f220d3187173da220f885ab66608367b6574e925011a9353e4badda91d79" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "serde_json" +version = "1.0.150" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e8014e44b4736ed0538adeecded0fce2a272f22dc9578a7eb6b2d9993c74cfb9" +dependencies = [ + "itoa", + "memchr", + "serde", + "serde_core", + "zmij", +] + +[[package]] +name = "shlex" +version = "2.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8fadd59c855ef2080decdef8ff161eb6661b86933c9d82e5ba29dc602a55aba" + +[[package]] +name = "smallvec" +version = "1.15.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8ed6a63f02c8539c91a8685a86f4099661ba3da017932f6ebbea6de3f0fa7c90" + +[[package]] +name = "spm_precompiled" +version = "0.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5851699c4033c63636f7ea4cf7b7c1f1bf06d0cc03cfb42e711de5a5c46cf326" +dependencies = [ + "base64", + "nom", + "serde", + "unicode-segmentation", +] + +[[package]] +name = "stable_deref_trait" +version = "1.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ce2be8dc25455e1f91df71bfa12ad37d7af1092ae736f3a6cd0e37bc7810596" + +[[package]] +name = "static_assertions" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a2eb9349b6444b326872e140eb1cf5e7c522154d69e7a0ffb0fb81c06b37543f" + +[[package]] +name = "strsim" +version = "0.11.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7da8b5736845d9f2fcb837ea5d9e2628564b3b043a70948a3f0b778838c5fb4f" + +[[package]] +name = "strum" +version = "0.28.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9628de9b8791db39ceda2b119bbe13134770b56c138ec1d3af810d045c04f9bd" +dependencies = [ + "strum_macros", +] + +[[package]] +name = "strum_macros" +version = "0.28.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ab85eea0270ee17587ed4156089e10b9e6880ee688791d45a905f5b1ca36f664" +dependencies = [ + "heck", + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "syn" +version = "2.0.119" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "872831b642d1a07999a962a351ed35b955ea2cfc8f3862091e2a240a84f17297" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "synstructure" +version = "0.13.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "728a70f3dbaf5bab7f0c4b1ac8d7ae5ea60a4b5549c8a5914361c99147a709d2" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "tempfile" +version = "3.27.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32497e9a4c7b38532efcdebeef879707aa9f794296a4f0244f6f69e9bc8574bd" +dependencies = [ + "fastrand", + "getrandom 0.4.3", + "once_cell", + "rustix", + "windows-sys", +] + +[[package]] +name = "thiserror" +version = "2.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4288b5bcbc7920c07a1149a35cf9590a2aa808e0bc1eafaade0b80947865fbc4" +dependencies = [ + "thiserror-impl", +] + +[[package]] +name = "thiserror-impl" +version = "2.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ebc4ee7f67670e9b64d05fa4253e753e016c6c95ff35b89b7941d6b856dec1d5" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "tinystr" +version = "0.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8323304221c2a851516f22236c5722a72eaa19749016521d6dff0824447d96d" +dependencies = [ + "displaydoc", + "zerovec", +] + +[[package]] +name = "tokenizers" +version = "0.23.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "44e5bea67576e04b6ff8564c5d9e09c2ef0cf476502245f2f120e497769d3112" +dependencies = [ + "ahash", + "compact_str", + "daachorse", + "dary_heap", + "derive_builder", + "esaxx-rs", + "getrandom 0.3.4", + "itertools 0.14.0", + "log", + "macro_rules_attribute", + "monostate", + "onig", + "paste", + "rand", + "rayon", + "rayon-cond", + "regex", + "regex-syntax", + "serde", + "serde_json", + "spm_precompiled", + "thiserror", + "unicode-normalization-alignments", + "unicode-segmentation", + "unicode_categories", +] + +[[package]] +name = "toml_datetime" +version = "1.1.1+spec-1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3165f65f62e28e0115a00b2ebdd37eb6f3b641855f9d636d3cd4103767159ad7" +dependencies = [ + "serde_core", +] + +[[package]] +name = "toml_edit" +version = "0.25.13+spec-1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6975367e4d2ef766d86af01ffad14b622fecc8d4357a998fbc4deb6e9bacaf9b" +dependencies = [ + "indexmap", + "toml_datetime", + "toml_parser", + "winnow", +] + +[[package]] +name = "toml_parser" +version = "1.1.2+spec-1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a2abe9b86193656635d2411dc43050282ca48aa31c2451210f4202550afb7526" +dependencies = [ + "winnow", +] + +[[package]] +name = "unicode-ident" +version = "1.0.24" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75" + +[[package]] +name = "unicode-normalization-alignments" +version = "0.1.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "43f613e4fa046e69818dd287fdc4bc78175ff20331479dab6e1b0f98d57062de" +dependencies = [ + "smallvec", +] + +[[package]] +name = "unicode-segmentation" +version = "1.13.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c6f5d3c3b1bf09027a88a6bc961fc00497d651009560b5463668dc81b0fa87a8" + +[[package]] +name = "unicode_categories" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "39ec24b3121d976906ece63c9daad25b85969647682eee313cb5779fdd69e14e" + +[[package]] +name = "utf8_iter" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6c140620e7ffbb22c2dee59cafe6084a59b5ffc27a8859a5f0d494b5d52b6be" + +[[package]] +name = "utf8parse" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "06abde3611657adf66d383f00b093d7faecc7fa57071cce2578660c9f1010821" + +[[package]] +name = "version_check" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" + +[[package]] +name = "wasip2" +version = "1.0.4+wasi-0.2.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b67efb37e106e55ce722a510d6b5f9c17f083e5fc79afc2badeb12cc313d9487" +dependencies = [ + "wit-bindgen", +] + +[[package]] +name = "windows-link" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" + +[[package]] +name = "windows-sys" +version = "0.61.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ae137229bcbd6cdf0f7b80a31df61766145077ddf49416a728b02cb3921ff3fc" +dependencies = [ + "windows-link", +] + +[[package]] +name = "winnow" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "23b97319f7b8343df12cc98938e5c3eb436064524c8d2b4e30a1d3a36eecdf81" +dependencies = [ + "memchr", +] + +[[package]] +name = "wit-bindgen" +version = "0.57.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ebf944e87a7c253233ad6766e082e3cd714b5d03812acc24c318f549614536e" + +[[package]] +name = "writeable" +version = "0.6.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ffae5123b2d3fc086436f8834ae3ab053a283cfac8fe0a0b8eaae044768a4c4" + +[[package]] +name = "yoke" +version = "0.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "709fe23a0424b6a435d82152b1bd3fdfb0833487d5fa90d05d42762a9891fef5" +dependencies = [ + "stable_deref_trait", + "yoke-derive", + "zerofrom", +] + +[[package]] +name = "yoke-derive" +version = "0.8.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "de844c262c8848816172cef550288e7dc6c7b7814b4ee56b3e1553f275f1858e" +dependencies = [ + "proc-macro2", + "quote", + "syn", + "synstructure", +] + +[[package]] +name = "zerocopy" +version = "0.8.54" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b7cbbc0a705a0fd05cc3676525980d2bf5a9bc4adac6d6475209a7887cf59d19" +dependencies = [ + "zerocopy-derive", +] + +[[package]] +name = "zerocopy-derive" +version = "0.8.54" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e2e817b7b52d0c7358d3246da9d69935ebb18116b2b102b4230dac079b4862f5" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "zerofrom" +version = "0.1.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0ec05a11813ea801ff6d75110ad09cd0824ddba17dfe17128ea0d5f68e6c5272" +dependencies = [ + "zerofrom-derive", +] + +[[package]] +name = "zerofrom-derive" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "11532158c46691caf0f2593ea8358fed6bbf68a0315e80aae9bd41fbade684a1" +dependencies = [ + "proc-macro2", + "quote", + "syn", + "synstructure", +] + +[[package]] +name = "zerotrie" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0f9152d31db0792fa83f70fb2f83148effb5c1f5b8c7686c3459e361d9bc20bf" +dependencies = [ + "displaydoc", + "yoke", + "zerofrom", +] + +[[package]] +name = "zerovec" +version = "0.11.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "90f911cbc359ab6af17377d242225f4d75119aec87ea711a880987b18cd7b239" +dependencies = [ + "yoke", + "zerofrom", + "zerovec-derive", +] + +[[package]] +name = "zerovec-derive" +version = "0.11.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "625dc425cab0dca6dc3c3319506e6593dcb08a9f387ea3b284dbd52a92c40555" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "zmij" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29666d0abbfad1e3dc4dcf6144730dd3a3ab225bbbdac83319345b1b44ccfc1b" diff --git a/spikes/mlx-solo-published/Cargo.toml b/spikes/mlx-solo-published/Cargo.toml new file mode 100644 index 0000000000..030cd7534f --- /dev/null +++ b/spikes/mlx-solo-published/Cargo.toml @@ -0,0 +1,25 @@ +# Fork-free variant: depends ONLY on crates.io (no path deps), to answer +# "can we do solo MLX serving without a fork?". +# +# Uses the goose-baseline path (plain LoadedModel::load) because the PUBLISHED +# safemlx-lm 0.4.1 does NOT expose the JIT-quant-on-load API (that lives only in +# the fork's unpublished HEAD). Own [workspace], not part of mesh-llm. +[package] +name = "mlx-solo-published" +version = "0.0.0" +edition = "2021" +publish = false + +[[bin]] +name = "mlx-solo-published" +path = "src/main.rs" + +[workspace] + +[dependencies] +# Pure crates.io. safemlx 0.1.3 default features include metal + accelerate. +safemlx = { version = "0.1.3", features = ["safetensors"] } +safemlx-lm = "0.4.1" +anyhow = "1" +clap = { version = "4", features = ["derive"] } +serde_json = "1" diff --git a/spikes/mlx-solo-published/src/main.rs b/spikes/mlx-solo-published/src/main.rs new file mode 100644 index 0000000000..f03f7a63a2 --- /dev/null +++ b/spikes/mlx-solo-published/src/main.rs @@ -0,0 +1,145 @@ +//! Fork-free solo MLX serving: depends only on crates.io safemlx-lm 0.4.1. +//! +//! Proves the goose-baseline path (load + generate, no JIT quant) works with no +//! fork. The published crate exposes `LoadedModel::load` but NOT the +//! JIT-quant-on-load API (`ModelLoadOptions` / `with_quantization`), which lives +//! only in the fork's unpublished HEAD — so this binary intentionally has no +//! --quantize flag. It serves source-precision or already-quantized MLX repos. + +use std::io::Write; +use std::path::PathBuf; +use std::time::Instant; + +use anyhow::{bail, Context, Result}; +use clap::Parser; +use safemlx::transforms::async_eval; +use safemlx::{Device, DeviceType, Stream}; +use safemlx_lm::models::input::{InputPart, ModelInput}; +use safemlx_lm::models::LoadedModel; +use safemlx_lm::sampler::DefaultSampler; + +#[derive(Parser, Debug)] +#[command(about = "Fork-free solo MLX serving (crates.io only; load + generate)")] +struct Cli { + /// Model directory (HF safetensors, source-precision or pre-quantized MLX). + #[arg(short, long)] + model: PathBuf, + + /// Prompt text. + #[arg(default_value = "Explain what a mesh network is in two sentences.")] + prompt: String, + + /// Max tokens to generate. + #[arg(short = 'n', long, default_value_t = 128)] + max_tokens: usize, + + /// Skip the chat template and feed the prompt raw. + #[arg(long)] + raw: bool, +} + +fn main() -> Result<()> { + let args = Cli::parse(); + + // Mirror goose: Metal GPU stream for compute, CPU stream for weight staging. + let stream = Stream::new_with_device(&Device::new(DeviceType::Gpu, 0)); + let weights_stream = Stream::new_with_device(&Device::new(DeviceType::Cpu, 0)); + + eprintln!("[load] loading {} ...", args.model.display()); + let load_started = Instant::now(); + let mut model = LoadedModel::load(&args.model, &stream, &weights_stream) + .with_context(|| format!("failed to load model from {}", args.model.display()))?; + stream.synchronize()?; + let load_elapsed = load_started.elapsed(); + eprintln!( + "[load] ok: model_type={} in {:.2}s", + model.model_type(), + load_elapsed.as_secs_f64() + ); + + let (rendered, add_special) = if args.raw { + (args.prompt.clone(), true) + } else { + match model.apply_chat_template_json( + vec![vec![serde_json::json!({"role": "user", "content": args.prompt})]], + None, + true, + )? { + Some(rendered) => (rendered, false), + None => { + eprintln!("[prompt] no chat template; feeding raw"); + (args.prompt.clone(), true) + } + } + }; + + let tokens = model.encode_to_array(&rendered, add_special, &stream)?; + let prompt_len = tokens.shape()[1]; + if prompt_len == 0 { + bail!("prompt produced no tokens"); + } + eprintln!("[prompt] {prompt_len} tokens"); + + let eos = model.eos_token_ids().to_vec(); + let mut cache = model.new_cache(); + + let mut output_ids: Vec = Vec::with_capacity(args.max_tokens); + let gen_started = Instant::now(); + let mut ttft = None; + + { + let parts = [InputPart::text_token_ids(&tokens)]; + let input = ModelInput::new(&parts); + let mut generator = + model.generate_input_with_cache_sampler(&mut cache, 0.0, input, None, &stream, DefaultSampler); + + let mut current = generator.next().transpose()?; + for index in 0..args.max_tokens { + let Some(token) = current.take() else { break }; + + let next = if index + 1 < args.max_tokens { + let next = generator.next(); + if let Some(Ok(next_token)) = next.as_ref() { + async_eval([next_token])?; + } + next + } else { + None + }; + + let token_id = token.item::(&stream); + if ttft.is_none() { + ttft = Some(gen_started.elapsed()); + } + output_ids.push(token_id); + if eos.contains(&token_id) { + break; + } + current = next.transpose()?; + } + } + + let gen_elapsed = gen_started.elapsed(); + let text = model.decode(&output_ids, true)?; + + let mut stdout = std::io::stdout().lock(); + writeln!(stdout, "\n===== OUTPUT =====\n{text}\n==================")?; + + let decode_tokens = output_ids.len().saturating_sub(1); + let decode_elapsed = gen_elapsed.saturating_sub(ttft.unwrap_or_default()); + let decode_rate = if decode_elapsed.is_zero() { + 0.0 + } else { + decode_tokens as f64 / decode_elapsed.as_secs_f64() + }; + eprintln!( + "[stats] load={:.2}s prompt_tokens={} generated={} ttft={:.3}s decode={:.1} tok/s", + load_elapsed.as_secs_f64(), + prompt_len, + output_ids.len(), + ttft.map(|d| d.as_secs_f64()).unwrap_or(0.0), + decode_rate, + ); + + Ok(()) +} From 4632c4b0e518414e7adb20419e0c4b9a90595464 Mon Sep 17 00:00:00 2001 From: michaelneale Date: Thu, 16 Jul 2026 15:00:26 +1000 Subject: [PATCH 05/37] feat: MLX (Metal) serving engine for tensor models, goose-style 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. --- crates/skippy-engine-mlx/Cargo.lock | 2098 +++++++++++++++++ crates/skippy-engine-mlx/Cargo.toml | 62 + crates/skippy-engine-mlx/WIRING.md | 133 ++ crates/skippy-engine-mlx/src/backend.rs | 248 ++ crates/skippy-engine-mlx/src/bin/mlx-serve.rs | 104 + crates/skippy-engine-mlx/src/engine.rs | 311 +++ crates/skippy-engine-mlx/src/lib.rs | 23 + 7 files changed, 2979 insertions(+) create mode 100644 crates/skippy-engine-mlx/Cargo.lock create mode 100644 crates/skippy-engine-mlx/Cargo.toml create mode 100644 crates/skippy-engine-mlx/WIRING.md create mode 100644 crates/skippy-engine-mlx/src/backend.rs create mode 100644 crates/skippy-engine-mlx/src/bin/mlx-serve.rs create mode 100644 crates/skippy-engine-mlx/src/engine.rs create mode 100644 crates/skippy-engine-mlx/src/lib.rs diff --git a/crates/skippy-engine-mlx/Cargo.lock b/crates/skippy-engine-mlx/Cargo.lock new file mode 100644 index 0000000000..6380fc69a4 --- /dev/null +++ b/crates/skippy-engine-mlx/Cargo.lock @@ -0,0 +1,2098 @@ +# This file is automatically @generated by Cargo. +# It is not intended for manual editing. +version = 4 + +[[package]] +name = "ahash" +version = "0.8.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5a15f179cd60c4584b8a8c596927aadc462e27f2ca70c04e0071964a73ba7a75" +dependencies = [ + "cfg-if", + "getrandom 0.3.4", + "once_cell", + "serde", + "version_check", + "zerocopy", +] + +[[package]] +name = "aho-corasick" +version = "1.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ddd31a130427c27518df266943a5308ed92d4b226cc639f5a8f1002816174301" +dependencies = [ + "memchr", +] + +[[package]] +name = "allocator-api2" +version = "0.2.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "683d7910e743518b0e34f1186f92494becacb047c7b6bf616c96772180fef923" + +[[package]] +name = "anstream" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "824a212faf96e9acacdbd09febd34438f8f711fb84e09a8916013cd7815ca28d" +dependencies = [ + "anstyle", + "anstyle-parse", + "anstyle-query", + "anstyle-wincon", + "colorchoice", + "is_terminal_polyfill", + "utf8parse", +] + +[[package]] +name = "anstyle" +version = "1.0.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "940b3a0ca603d1eade50a4846a2afffd5ef57a9feac2c0e2ec2e14f9ead76000" + +[[package]] +name = "anstyle-parse" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "52ce7f38b242319f7cabaa6813055467063ecdc9d355bbb4ce0c68908cd8130e" +dependencies = [ + "utf8parse", +] + +[[package]] +name = "anstyle-query" +version = "1.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "40c48f72fd53cd289104fc64099abca73db4166ad86ea0b4341abe65af83dadc" +dependencies = [ + "windows-sys", +] + +[[package]] +name = "anstyle-wincon" +version = "3.0.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "291e6a250ff86cd4a820112fb8898808a366d8f9f58ce16d1f538353ad55747d" +dependencies = [ + "anstyle", + "once_cell_polyfill", + "windows-sys", +] + +[[package]] +name = "anyhow" +version = "1.0.103" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2a4385e2e34eb35d6b3efe798b9eb88096925d87726c0798709bf56d9ed84af3" + +[[package]] +name = "async-stream" +version = "0.3.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b5a71a6f37880a80d1d7f19efd781e4b5de42c88f0722cc13bcb6cc2cfe8476" +dependencies = [ + "async-stream-impl", + "futures-core", + "pin-project-lite", +] + +[[package]] +name = "async-stream-impl" +version = "0.3.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c7c24de15d275a1ecfd47a380fb4d5ec9bfe0933f309ed5e705b775596a3574d" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "async-trait" +version = "0.1.89" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9035ad2d096bed7955a320ee7e2230574d28fd3c3a0f186cbea1ff3c7eed5dbb" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "atomic-waker" +version = "1.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1505bd5d3d116872e7271a6d4e16d81d0c8570876c8de68093a09ac269d8aac0" + +[[package]] +name = "autocfg" +version = "1.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f2032f911046de80f0a198e0901378627c33f59ea0ac00e363d481118bd70a53" + +[[package]] +name = "axum" +version = "0.8.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "31b698c5f9a010f6573133b09e0de5408834d0c82f8d7475a89fc1867a71cd90" +dependencies = [ + "axum-core", + "bytes", + "form_urlencoded", + "futures-util", + "http", + "http-body", + "http-body-util", + "hyper", + "hyper-util", + "itoa", + "matchit", + "memchr", + "mime", + "percent-encoding", + "pin-project-lite", + "serde_core", + "serde_json", + "serde_path_to_error", + "serde_urlencoded", + "sync_wrapper", + "tokio", + "tower", + "tower-layer", + "tower-service", + "tracing", +] + +[[package]] +name = "axum-core" +version = "0.5.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "08c78f31d7b1291f7ee735c1c6780ccde7785daae9a9206026862dab7d8792d1" +dependencies = [ + "bytes", + "futures-core", + "http", + "http-body", + "http-body-util", + "mime", + "pin-project-lite", + "sync_wrapper", + "tower-layer", + "tower-service", + "tracing", +] + +[[package]] +name = "base64" +version = "0.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9e1b586273c5702936fe7b7d6896644d8be71e6314cfe09d3167c95f712589e8" + +[[package]] +name = "bitflags" +version = "2.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b588b76d00fde79687d7646a9b5bdf3cc0f655e0bbd080335a95d7e96f3587da" + +[[package]] +name = "bytemuck" +version = "1.25.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d6aedf8ae72766347502cf3cb4f41cf5e9cc37d28bee90f1fdaaae15f9cf9424" + +[[package]] +name = "bytes" +version = "1.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc652a48c352aef3ea3aed32080501cf3ef6ed5da78602a020c991775b0aff04" + +[[package]] +name = "castaway" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dec551ab6e7578819132c713a93c022a05d60159dc86e7a7050223577484c55a" +dependencies = [ + "rustversion", +] + +[[package]] +name = "cc" +version = "1.2.67" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e17dd265a7d0f31ef544e1b20e03add05d3b45b491b633b10d67145d2acc1a38" +dependencies = [ + "find-msvc-tools", + "shlex", +] + +[[package]] +name = "cfg-if" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" + +[[package]] +name = "clap" +version = "4.6.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dd059f9da4f5c36b3787f65d38ccaab1cc315f07b01f89abc8359ee6a8205011" +dependencies = [ + "clap_builder", + "clap_derive", +] + +[[package]] +name = "clap_builder" +version = "4.6.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f09628afdcc538b57f3c6341e9c8e9970f18e4a481690a64974d7023bd33548b" +dependencies = [ + "anstream", + "anstyle", + "clap_lex", + "strsim", +] + +[[package]] +name = "clap_derive" +version = "4.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f2ce8604710f6733aa641a2b3731eaa1e8b3d9973d5e3565da11800813f997a9" +dependencies = [ + "heck", + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "clap_lex" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8d4a3bb8b1e0c1050499d1815f5ab16d04f0959b233085fb31653fbfc9d98f9" + +[[package]] +name = "cmake" +version = "0.1.58" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c0f78a02292a74a88ac736019ab962ece0bc380e3f977bf72e376c5d78ff0678" +dependencies = [ + "cc", +] + +[[package]] +name = "colorchoice" +version = "1.0.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1d07550c9036bf2ae0c684c4297d503f838287c83c53686d05370d0e139ae570" + +[[package]] +name = "compact_str" +version = "0.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9dfdd1c2274d9aa354115b09dc9a901d6c5576818cdf70d14cae2bdb47df00ab" +dependencies = [ + "castaway", + "cfg-if", + "itoa", + "rustversion", + "ryu", + "serde", + "static_assertions", +] + +[[package]] +name = "crossbeam-deque" +version = "0.8.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5181e0de7b61eb03a81e347d6dd8797bae9da5146707b51077e2d71a54ec0ceb" +dependencies = [ + "crossbeam-epoch", + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-epoch" +version = "0.9.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2d6914041f254d6e9176c01941b21115dcfb7089e55135a35411081bd106ef3f" +dependencies = [ + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-utils" +version = "0.8.22" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "61803da095bee82a81bb1a452ecc25d3b2f1416d1897eb86430c6159ef717c17" + +[[package]] +name = "crunchy" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "460fbee9c2c2f33933d720630a6a0bac33ba7053db5344fac858d4b8952d77d5" + +[[package]] +name = "daachorse" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6f55d7153ba3b507595872a3874803f07a8a81d1e888abed8e5db7da0597d6e2" + +[[package]] +name = "darling" +version = "0.20.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc7f46116c46ff9ab3eb1597a45688b6715c6e628b5c133e288e709a29bcb4ee" +dependencies = [ + "darling_core 0.20.11", + "darling_macro 0.20.11", +] + +[[package]] +name = "darling" +version = "0.23.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "25ae13da2f202d56bd7f91c25fba009e7717a1e4a1cc98a76d844b65ae912e9d" +dependencies = [ + "darling_core 0.23.0", + "darling_macro 0.23.0", +] + +[[package]] +name = "darling_core" +version = "0.20.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0d00b9596d185e565c2207a0b01f8bd1a135483d02d9b7b0a54b11da8d53412e" +dependencies = [ + "fnv", + "ident_case", + "proc-macro2", + "quote", + "strsim", + "syn", +] + +[[package]] +name = "darling_core" +version = "0.23.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9865a50f7c335f53564bb694ef660825eb8610e0a53d3e11bf1b0d3df31e03b0" +dependencies = [ + "ident_case", + "proc-macro2", + "quote", + "strsim", + "syn", +] + +[[package]] +name = "darling_macro" +version = "0.20.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc34b93ccb385b40dc71c6fceac4b2ad23662c7eeb248cf10d529b7e055b6ead" +dependencies = [ + "darling_core 0.20.11", + "quote", + "syn", +] + +[[package]] +name = "darling_macro" +version = "0.23.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ac3984ec7bd6cfa798e62b4a642426a5be0e68f9401cfc2a01e3fa9ea2fcdb8d" +dependencies = [ + "darling_core 0.23.0", + "quote", + "syn", +] + +[[package]] +name = "dary_heap" +version = "0.3.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8b1e3a325bc115f096c8b77bbf027a7c2592230e70be2d985be950d3d5e60ebe" +dependencies = [ + "serde", +] + +[[package]] +name = "derive_builder" +version = "0.20.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "507dfb09ea8b7fa618fcf76e953f4f5e192547945816d5358edffe39f6f94947" +dependencies = [ + "derive_builder_macro", +] + +[[package]] +name = "derive_builder_core" +version = "0.20.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2d5bcf7b024d6835cfb3d473887cd966994907effbe9227e8c8219824d06c4e8" +dependencies = [ + "darling 0.20.11", + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "derive_builder_macro" +version = "0.20.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ab63b0e2bf4d5928aff72e83a7dace85d7bba5fe12dcc3c5a572d78caffd3f3c" +dependencies = [ + "derive_builder_core", + "syn", +] + +[[package]] +name = "displaydoc" +version = "0.2.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ac70aa55017e108007fbaf5aa0f54b021c98f92ff8af59d42eda9da96e3dd4f" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "dyn-clone" +version = "1.0.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d0881ea181b1df73ff77ffaaf9c7544ecc11e82fba9b5f27b262a3c73a332555" + +[[package]] +name = "either" +version = "1.16.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91622ff5e7162018101f2fea40d6ebf4a78bbe5a49736a2020649edf9693679e" + +[[package]] +name = "equivalent" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "877a4ace8713b0bcf2a4e7eec82529c029f1d0619886d18145fea96c3ffe5c0f" + +[[package]] +name = "errno" +version = "0.3.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" +dependencies = [ + "libc", + "windows-sys", +] + +[[package]] +name = "esaxx-rs" +version = "0.1.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d817e038c30374a4bcb22f94d0a8a0e216958d4c3dcde369b1439fec4bdda6e6" +dependencies = [ + "cc", +] + +[[package]] +name = "fastrand" +version = "2.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9f1f227452a390804cdb637b74a86990f2a7d7ba4b7d5693aac9b4dd6defd8d6" + +[[package]] +name = "find-msvc-tools" +version = "0.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5baebc0774151f905a1a2cc41989300b1e6fbb29aff0ceffa1064fdd3088d582" + +[[package]] +name = "fnv" +version = "1.0.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3f9eec918d3f24069decb9af1554cad7c880e2da24a9afd88aca000531ab82c1" + +[[package]] +name = "foldhash" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "77ce24cb58228fbb8aa041425bb1050850ac19177686ea6e0f41a70416f56fdb" + +[[package]] +name = "form_urlencoded" +version = "1.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cb4cb245038516f5f85277875cdaa4f7d2c9a0fa0468de06ed190163b1581fcf" +dependencies = [ + "percent-encoding", +] + +[[package]] +name = "futures-channel" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "07bbe89c50d7a535e539b8c17bc0b49bdb77747034daa8087407d655f3f7cc1d" +dependencies = [ + "futures-core", +] + +[[package]] +name = "futures-core" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7e3450815272ef58cec6d564423f6e755e25379b217b0bc688e295ba24df6b1d" + +[[package]] +name = "futures-macro" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e835b70203e41293343137df5c0664546da5745f82ec9b84d40be8336958447b" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "futures-task" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "037711b3d59c33004d3856fbdc83b99d4ff37a24768fa1be9ce3538a1cde4393" + +[[package]] +name = "futures-util" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "389ca41296e6190b48053de0321d02a77f32f8a5d2461dd38762c0593805c6d6" +dependencies = [ + "futures-core", + "futures-macro", + "futures-task", + "pin-project-lite", + "slab", +] + +[[package]] +name = "getrandom" +version = "0.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "899def5c37c4fd7b2664648c28120ecec138e4d395b459e5ca34f9cce2dd77fd" +dependencies = [ + "cfg-if", + "libc", + "r-efi 5.3.0", + "wasip2", +] + +[[package]] +name = "getrandom" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "300e883d756b2e4ec94e02791f39b04b522276138852cfc41d9fb7e904106099" +dependencies = [ + "cfg-if", + "libc", + "r-efi 6.0.0", +] + +[[package]] +name = "half" +version = "2.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" +dependencies = [ + "cfg-if", + "crunchy", + "zerocopy", +] + +[[package]] +name = "hashbrown" +version = "0.16.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "841d1cc9bed7f9236f321df977030373f4a4163ae1a7dbfe1a51a2c1a51d9100" +dependencies = [ + "allocator-api2", + "equivalent", + "foldhash", + "serde", + "serde_core", +] + +[[package]] +name = "hashbrown" +version = "0.17.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a" + +[[package]] +name = "heck" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" + +[[package]] +name = "http" +version = "1.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6970f50e31d6fc17d3fa27329444bfa74e196cf62e95052a3f6fee181dba6425" +dependencies = [ + "bytes", + "itoa", +] + +[[package]] +name = "http-body" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ca2a8f2913ee65f60facd6a5905613afaa448497a0230cc41ce022d93290bc2c" +dependencies = [ + "bytes", + "http", +] + +[[package]] +name = "http-body-util" +version = "0.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e9f41fd6a08e4d4ec69df65976da761afd5ad5e58a9d4acb46bd1c953a9e3ff2" +dependencies = [ + "bytes", + "futures-core", + "http", + "http-body", + "pin-project-lite", +] + +[[package]] +name = "httparse" +version = "1.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6dbf3de79e51f3d586ab4cb9d5c3e2c14aa28ed23d180cf89b4df0454a69cc87" + +[[package]] +name = "httpdate" +version = "1.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "df3b46402a9d5adb4c86a0cf463f42e19994e3ee891101b1841f30a545cb49a9" + +[[package]] +name = "hyper" +version = "1.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "55281c53a1894c864990125767da440a4e630446785086f52523b20033b74498" +dependencies = [ + "atomic-waker", + "bytes", + "futures-channel", + "futures-core", + "http", + "http-body", + "httparse", + "httpdate", + "itoa", + "pin-project-lite", + "smallvec", + "tokio", +] + +[[package]] +name = "hyper-util" +version = "0.1.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "96547c2556ec9d12fb1578c4eaf448b04993e7fb79cbaad930a656880a6bdfa0" +dependencies = [ + "bytes", + "http", + "http-body", + "hyper", + "pin-project-lite", + "tokio", + "tower-service", +] + +[[package]] +name = "icu_collections" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2984d1cd16c883d7935b9e07e44071dca8d917fd52ecc02c04d5fa0b5a3f191c" +dependencies = [ + "displaydoc", + "potential_utf", + "utf8_iter", + "yoke", + "zerofrom", + "zerovec", +] + +[[package]] +name = "icu_locale_core" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "92219b62b3e2b4d88ac5119f8904c10f8f61bf7e95b640d25ba3075e6cac2c29" +dependencies = [ + "displaydoc", + "litemap", + "tinystr", + "writeable", + "zerovec", +] + +[[package]] +name = "icu_normalizer" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c56e5ee99d6e3d33bd91c5d85458b6005a22140021cc324cea84dd0e72cff3b4" +dependencies = [ + "icu_collections", + "icu_provider", + "smallvec", + "zerovec", +] + +[[package]] +name = "icu_properties" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bee3b67d0ea5c2cca5003417989af8996f8604e34fb9ddf96208a033901e70de" +dependencies = [ + "icu_collections", + "icu_locale_core", + "icu_provider", + "zerotrie", + "zerovec", +] + +[[package]] +name = "icu_provider" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "139c4cf31c8b5f33d7e199446eff9c1e02decfc2f0eec2c8d71f65befa45b421" +dependencies = [ + "displaydoc", + "icu_locale_core", + "yoke", + "zerofrom", + "zerovec", +] + +[[package]] +name = "ident_case" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b9e0384b61958566e926dc50660321d12159025e767c18e043daf26b70104c39" + +[[package]] +name = "idna_adapter" +version = "1.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cb68373c0d6620ef8105e855e7745e18b0d00d3bdb07fb532e434244cdb9a714" +dependencies = [ + "icu_normalizer", + "icu_properties", +] + +[[package]] +name = "indexmap" +version = "2.14.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d466e9454f08e4a911e14806c24e16fba1b4c121d1ea474396f396069cf949d9" +dependencies = [ + "equivalent", + "hashbrown 0.17.1", +] + +[[package]] +name = "is_terminal_polyfill" +version = "1.70.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a6cb138bb79a146c1bd460005623e142ef0181e3d0219cb493e02f7d08a35695" + +[[package]] +name = "itertools" +version = "0.14.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2b192c782037fadd9cfa75548310488aabdbf3d2da73885b31bd0abd03351285" +dependencies = [ + "either", +] + +[[package]] +name = "itertools" +version = "0.15.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8b4baf93f58d4425749ca49a51c50ebab072c5df6994d08fed93541c331481dc" +dependencies = [ + "either", +] + +[[package]] +name = "itoa" +version = "1.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" + +[[package]] +name = "lazy_static" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe" + +[[package]] +name = "libc" +version = "0.2.186" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68ab91017fe16c622486840e4c83c9a37afeff978bd239b5293d61ece587de66" + +[[package]] +name = "linux-raw-sys" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32a66949e030da00e8c7d4434b251670a91556f4144941d37452769c25d58a53" + +[[package]] +name = "litemap" +version = "0.8.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "92daf443525c4cce67b150400bc2316076100ce0b3686209eb8cf3c31612e6f0" + +[[package]] +name = "lock_api" +version = "0.4.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "224399e74b87b5f3557511d98dff8b14089b3dadafcab6bb93eab67d3aace965" +dependencies = [ + "scopeguard", +] + +[[package]] +name = "log" +version = "0.4.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0ceec5bc11778974d1bcb055b18002eba7f4b3518b6a0081b3af5f21666da9ad" + +[[package]] +name = "macro_rules_attribute" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "65049d7923698040cd0b1ddcced9b0eb14dd22c5f86ae59c3740eab64a676520" +dependencies = [ + "macro_rules_attribute-proc_macro", + "paste", +] + +[[package]] +name = "macro_rules_attribute-proc_macro" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "670fdfda89751bc4a84ac13eaa63e205cf0fd22b4c9a5fbfa085b63c1f1d3a30" + +[[package]] +name = "matchers" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d1525a2a28c7f4fa0fc98bb91ae755d1e2d1505079e05539e35bc876b5d65ae9" +dependencies = [ + "regex-automata", +] + +[[package]] +name = "matchit" +version = "0.8.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "47e1ffaa40ddd1f3ed91f717a33c8c0ee23fff369e3aa8772b9605cc1d22f4c3" + +[[package]] +name = "memchr" +version = "2.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf8baf1c55e62ffcace7a9f06f4bd9cd3f0c4beb022d3b367256b91b87513d98" + +[[package]] +name = "memmap2" +version = "0.9.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d1219ed1b7f229ee7104d281dd01d6802fe28bb6e95d292942c4daacdeb798c0" +dependencies = [ + "libc", +] + +[[package]] +name = "memo-map" +version = "0.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "38d1115007560874e373613744c6fba374c17688327a71c1476d1a5954cc857b" + +[[package]] +name = "mesh-llm-guardrails" +version = "0.72.1" +dependencies = [ + "serde", + "serde_json", +] + +[[package]] +name = "mime" +version = "0.3.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6877bb514081ee2a7ff5ef9de3281f14a4dd4bceac4c09388074a6b5df8a139a" + +[[package]] +name = "minijinja" +version = "2.21.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cb3d648e68cea56d9858d535ee28f9538404e2dd8cb08ed0bd05dca379477f39" +dependencies = [ + "memo-map", + "serde", +] + +[[package]] +name = "minijinja-contrib" +version = "2.21.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85342f6fac0be8ccd5bd00d9066be538f34f393f577b75d81b17c8398a6b43bb" +dependencies = [ + "minijinja", + "serde", +] + +[[package]] +name = "minimal-lexical" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68354c5c6bd36d73ff3feceb05efa59b6acb7626617f4962be322a825e61f79a" + +[[package]] +name = "mio" +version = "1.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "30d65c71f1ce40ab09135ce117d742b9f8a19ff91a41a8b57ed50bc2de59c427" +dependencies = [ + "libc", + "wasi", + "windows-sys", +] + +[[package]] +name = "monostate" +version = "0.1.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3341a273f6c9d5bef1908f17b7267bbab0e95c9bf69a0d4dcf8e9e1b2c76ef67" +dependencies = [ + "monostate-impl", + "serde", + "serde_core", +] + +[[package]] +name = "monostate-impl" +version = "0.1.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e4db6d5580af57bf992f59068d4ea26fd518574ff48d7639b255a36f9de6e7e9" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "nom" +version = "7.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d273983c5a657a70a3e8f2a01329822f3b8c8172b73826411a55751e404a0a4a" +dependencies = [ + "memchr", + "minimal-lexical", +] + +[[package]] +name = "nu-ansi-term" +version = "0.50.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7957b9740744892f114936ab4a57b3f487491bbeafaf8083688b16841a4240e5" +dependencies = [ + "windows-sys", +] + +[[package]] +name = "num-complex" +version = "0.4.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "73f88a1307638156682bada9d7604135552957b7818057dcef22705b4d509495" +dependencies = [ + "num-traits", +] + +[[package]] +name = "num-traits" +version = "0.2.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "071dfc062690e90b734c0b2273ce72ad0ffa95f0c74596bc250dcfd960262841" +dependencies = [ + "autocfg", +] + +[[package]] +name = "num_enum" +version = "0.7.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5d0bca838442ec211fa11de3a8b0e0e8f3a4522575b5c4c06ed722e005036f26" +dependencies = [ + "num_enum_derive", + "rustversion", +] + +[[package]] +name = "num_enum_derive" +version = "0.7.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "680998035259dcfcafe653688bf2aa6d3e2dc05e98be6ab46afb089dc84f1df8" +dependencies = [ + "proc-macro-crate", + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "once_cell" +version = "1.21.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" + +[[package]] +name = "once_cell_polyfill" +version = "1.70.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "384b8ab6d37215f3c5301a95a4accb5d64aa607f1fcb26a11b5303878451b4fe" + +[[package]] +name = "onig" +version = "6.5.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0cc3cbf698f9438986c11a880c90a6d04b9de27575afd28bbf45b154b6c709e2" +dependencies = [ + "bitflags", + "libc", + "once_cell", + "onig_sys", +] + +[[package]] +name = "onig_sys" +version = "69.9.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e68317604e77e53b85896388e1a803c1d21b74c899ec9e5e1112db90735edd7" +dependencies = [ + "cc", + "pkg-config", +] + +[[package]] +name = "openai-frontend" +version = "0.72.1" +dependencies = [ + "async-trait", + "axum", + "futures-core", + "futures-util", + "mesh-llm-guardrails", + "serde", + "serde_json", + "tokio", + "tokio-stream", + "tracing", +] + +[[package]] +name = "parking_lot" +version = "0.12.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "93857453250e3077bd71ff98b6a65ea6621a19bb0f559a85248955ac12c45a1a" +dependencies = [ + "lock_api", + "parking_lot_core", +] + +[[package]] +name = "parking_lot_core" +version = "0.9.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2621685985a2ebf1c516881c026032ac7deafcda1a2c9b7850dc81e3dfcb64c1" +dependencies = [ + "cfg-if", + "libc", + "redox_syscall", + "smallvec", + "windows-link", +] + +[[package]] +name = "paste" +version = "1.0.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "57c0d7b74b563b49d38dae00a0c37d4d6de9b432382b2892f0574ddcae73fd0a" + +[[package]] +name = "percent-encoding" +version = "2.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9b4f627cb1b25917193a259e49bdad08f671f8d9708acfd5fe0a8c1455d87220" + +[[package]] +name = "pin-project-lite" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a89322df9ebe1c1578d689c92318e070967d1042b512afbe49518723f4e6d5cd" + +[[package]] +name = "pkg-config" +version = "0.3.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "19f132c84eca552bf34cab8ec81f1c1dcc229b811638f9d283dceabe58c5569e" + +[[package]] +name = "potential_utf" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0103b1cef7ec0cf76490e969665504990193874ea05c85ff9bab8b911d0a0564" +dependencies = [ + "zerovec", +] + +[[package]] +name = "ppv-lite86" +version = "0.2.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85eae3c4ed2f50dcfe72643da4befc30deadb458a9b590d720cde2f2b1e97da9" +dependencies = [ + "zerocopy", +] + +[[package]] +name = "proc-macro-crate" +version = "3.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e67ba7e9b2b56446f1d419b1d807906278ffa1a658a8a5d8a39dcb1f5a78614f" +dependencies = [ + "toml_edit", +] + +[[package]] +name = "proc-macro2" +version = "1.0.106" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8fd00f0bb2e90d81d1044c2b32617f68fcb9fa3bb7640c23e9c748e53fb30934" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "quote" +version = "1.0.46" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dfbc457d0c7a0759a614551b11a6409e5951f6c7537be1f1b7682b9ae9230368" +dependencies = [ + "proc-macro2", +] + +[[package]] +name = "r-efi" +version = "5.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "69cdb34c158ceb288df11e18b4bd39de994f6657d83847bdffdbd7f346754b0f" + +[[package]] +name = "r-efi" +version = "6.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8dcc9c7d52a811697d2151c701e0d08956f92b0e24136cf4cf27b57a6a0d9bf" + +[[package]] +name = "rand" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b9ef1d0d795eb7d84685bca4f72f3649f064e6641543d3a8c415898726a57b41" +dependencies = [ + "rand_chacha", + "rand_core", +] + +[[package]] +name = "rand_chacha" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d3022b5f1df60f26e1ffddd6c66e8aa15de382ae63b3a0c1bfc0e4d3e3f325cb" +dependencies = [ + "ppv-lite86", + "rand_core", +] + +[[package]] +name = "rand_core" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "76afc826de14238e6e8c374ddcc1fa19e374fd8dd986b0d2af0d02377261d83c" +dependencies = [ + "getrandom 0.3.4", +] + +[[package]] +name = "rayon" +version = "1.12.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fb39b166781f92d482534ef4b4b1b2568f42613b53e5b6c160e24cfbfa30926d" +dependencies = [ + "either", + "rayon-core", +] + +[[package]] +name = "rayon-cond" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2964d0cf57a3e7a06e8183d14a8b527195c706b7983549cd5462d5aa3747438f" +dependencies = [ + "either", + "itertools 0.14.0", + "rayon", +] + +[[package]] +name = "rayon-core" +version = "1.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "22e18b0f0062d30d4230b2e85ff77fdfe4326feb054b9783a3460d8435c8ab91" +dependencies = [ + "crossbeam-deque", + "crossbeam-utils", +] + +[[package]] +name = "redox_syscall" +version = "0.5.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed2bf2547551a7053d6fdfafda3f938979645c44812fbfcda098faae3f1a362d" +dependencies = [ + "bitflags", +] + +[[package]] +name = "regex" +version = "1.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f020237b6c8eed93db2e2cb53c00c60a8e1bc73da7d073199a1180401450218d" +dependencies = [ + "aho-corasick", + "memchr", + "regex-automata", + "regex-syntax", +] + +[[package]] +name = "regex-automata" +version = "0.4.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8fcfdb36bda0c880c5931cdc7a2bcdc8ba4556847b9d912bca70bc94708711ad" +dependencies = [ + "aho-corasick", + "memchr", + "regex-syntax", +] + +[[package]] +name = "regex-syntax" +version = "0.8.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d6f6ff9a378485b298a5286656da665ba74413d36db0979633275d2e708145d4" + +[[package]] +name = "rustix" +version = "1.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6fe4565b9518b83ef4f91bb47ce29620ca828bd32cb7e408f0062e9930ba190" +dependencies = [ + "bitflags", + "errno", + "libc", + "linux-raw-sys", + "windows-sys", +] + +[[package]] +name = "rustversion" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf54715a573b99ac80df0bc206da022bcd442c974952c7b9720069370852e21f" + +[[package]] +name = "ryu" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9774ba4a74de5f7b1c1451ed6cd5285a32eddb5cccb8cc655a4e50009e06477f" + +[[package]] +name = "safemlx" +version = "0.1.3" +dependencies = [ + "bytemuck", + "dyn-clone", + "half", + "itertools 0.15.0", + "libc", + "num-complex", + "num-traits", + "num_enum", + "parking_lot", + "paste", + "safemlx-internal-macros", + "safemlx-macros", + "safemlx-sys", + "safetensors", + "smallvec", + "strum", + "thiserror", +] + +[[package]] +name = "safemlx-internal-macros" +version = "0.1.1" +dependencies = [ + "darling 0.23.0", + "itertools 0.15.0", + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "safemlx-lm" +version = "0.4.1" +dependencies = [ + "anyhow", + "clap", + "idna_adapter", + "memmap2", + "minijinja", + "safemlx", + "safemlx-lm-utils", + "safetensors", + "serde", + "serde_json", + "thiserror", + "tokenizers", +] + +[[package]] +name = "safemlx-lm-utils" +version = "0.1.4" +dependencies = [ + "minijinja", + "minijinja-contrib", + "serde", + "serde_json", + "thiserror", + "tokenizers", +] + +[[package]] +name = "safemlx-macros" +version = "0.1.1" +dependencies = [ + "darling 0.23.0", + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "safemlx-sys" +version = "0.1.3" +dependencies = [ + "cc", + "cmake", +] + +[[package]] +name = "safetensors" +version = "0.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "79b079b829cb27a1c3c374341345ed2e8b2c0c839034522cee576c140bd7f846" +dependencies = [ + "hashbrown 0.16.1", + "libc", + "serde", + "serde_json", + "tempfile", +] + +[[package]] +name = "scopeguard" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" + +[[package]] +name = "serde" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9a8e94ea7f378bd32cbbd37198a4a91436180c5bb472411e48b5ec2e2124ae9e" +dependencies = [ + "serde_core", + "serde_derive", +] + +[[package]] +name = "serde_core" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "41d385c7d4ca58e59fc732af25c3983b67ac852c1a25000afe1175de458b67ad" +dependencies = [ + "serde_derive", +] + +[[package]] +name = "serde_derive" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d540f220d3187173da220f885ab66608367b6574e925011a9353e4badda91d79" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "serde_json" +version = "1.0.150" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e8014e44b4736ed0538adeecded0fce2a272f22dc9578a7eb6b2d9993c74cfb9" +dependencies = [ + "itoa", + "memchr", + "serde", + "serde_core", + "zmij", +] + +[[package]] +name = "serde_path_to_error" +version = "0.1.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "10a9ff822e371bb5403e391ecd83e182e0e77ba7f6fe0160b795797109d1b457" +dependencies = [ + "itoa", + "serde", + "serde_core", +] + +[[package]] +name = "serde_urlencoded" +version = "0.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d3491c14715ca2294c4d6a88f15e84739788c1d030eed8c110436aafdaa2f3fd" +dependencies = [ + "form_urlencoded", + "itoa", + "ryu", + "serde", +] + +[[package]] +name = "sharded-slab" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f40ca3c46823713e0d4209592e8d6e826aa57e928f09752619fc696c499637f6" +dependencies = [ + "lazy_static", +] + +[[package]] +name = "shlex" +version = "2.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8fadd59c855ef2080decdef8ff161eb6661b86933c9d82e5ba29dc602a55aba" + +[[package]] +name = "signal-hook-registry" +version = "1.4.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c4db69cba1110affc0e9f7bcd48bbf87b3f4fc7c61fc9155afd4c469eb3d6c1b" +dependencies = [ + "errno", + "libc", +] + +[[package]] +name = "skippy-engine-mlx" +version = "0.0.0" +dependencies = [ + "anyhow", + "async-stream", + "async-trait", + "axum", + "clap", + "futures-core", + "openai-frontend", + "safemlx", + "safemlx-lm", + "serde_json", + "tokenizers", + "tokio", + "tracing", + "tracing-subscriber", +] + +[[package]] +name = "slab" +version = "0.4.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c790de23124f9ab44544d7ac05d60440adc586479ce501c1d6d7da3cd8c9cf5" + +[[package]] +name = "smallvec" +version = "1.15.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8ed6a63f02c8539c91a8685a86f4099661ba3da017932f6ebbea6de3f0fa7c90" + +[[package]] +name = "socket2" +version = "0.6.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c3d1e2c7f27f8d4cb10542a02c49005dbd6e93095799d6f3be745fae9f8fedd4" +dependencies = [ + "libc", + "windows-sys", +] + +[[package]] +name = "spm_precompiled" +version = "0.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5851699c4033c63636f7ea4cf7b7c1f1bf06d0cc03cfb42e711de5a5c46cf326" +dependencies = [ + "base64", + "nom", + "serde", + "unicode-segmentation", +] + +[[package]] +name = "stable_deref_trait" +version = "1.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ce2be8dc25455e1f91df71bfa12ad37d7af1092ae736f3a6cd0e37bc7810596" + +[[package]] +name = "static_assertions" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a2eb9349b6444b326872e140eb1cf5e7c522154d69e7a0ffb0fb81c06b37543f" + +[[package]] +name = "strsim" +version = "0.11.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7da8b5736845d9f2fcb837ea5d9e2628564b3b043a70948a3f0b778838c5fb4f" + +[[package]] +name = "strum" +version = "0.28.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9628de9b8791db39ceda2b119bbe13134770b56c138ec1d3af810d045c04f9bd" +dependencies = [ + "strum_macros", +] + +[[package]] +name = "strum_macros" +version = "0.28.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ab85eea0270ee17587ed4156089e10b9e6880ee688791d45a905f5b1ca36f664" +dependencies = [ + "heck", + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "syn" +version = "2.0.119" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "872831b642d1a07999a962a351ed35b955ea2cfc8f3862091e2a240a84f17297" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "sync_wrapper" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0bf256ce5efdfa370213c1dabab5935a12e49f2c58d15e9eac2870d3b4f27263" + +[[package]] +name = "synstructure" +version = "0.13.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "728a70f3dbaf5bab7f0c4b1ac8d7ae5ea60a4b5549c8a5914361c99147a709d2" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "tempfile" +version = "3.27.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32497e9a4c7b38532efcdebeef879707aa9f794296a4f0244f6f69e9bc8574bd" +dependencies = [ + "fastrand", + "getrandom 0.4.3", + "once_cell", + "rustix", + "windows-sys", +] + +[[package]] +name = "thiserror" +version = "2.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4288b5bcbc7920c07a1149a35cf9590a2aa808e0bc1eafaade0b80947865fbc4" +dependencies = [ + "thiserror-impl", +] + +[[package]] +name = "thiserror-impl" +version = "2.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ebc4ee7f67670e9b64d05fa4253e753e016c6c95ff35b89b7941d6b856dec1d5" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "thread_local" +version = "1.1.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ad99c4c6d32803332c548b1af0540b357b3f5fc0be8f6c6bfe8b2e6ae784070" +dependencies = [ + "cfg-if", +] + +[[package]] +name = "tinystr" +version = "0.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8323304221c2a851516f22236c5722a72eaa19749016521d6dff0824447d96d" +dependencies = [ + "displaydoc", + "zerovec", +] + +[[package]] +name = "tokenizers" +version = "0.23.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "44e5bea67576e04b6ff8564c5d9e09c2ef0cf476502245f2f120e497769d3112" +dependencies = [ + "ahash", + "compact_str", + "daachorse", + "dary_heap", + "derive_builder", + "esaxx-rs", + "getrandom 0.3.4", + "itertools 0.14.0", + "log", + "macro_rules_attribute", + "monostate", + "onig", + "paste", + "rand", + "rayon", + "rayon-cond", + "regex", + "regex-syntax", + "serde", + "serde_json", + "spm_precompiled", + "thiserror", + "unicode-normalization-alignments", + "unicode-segmentation", + "unicode_categories", +] + +[[package]] +name = "tokio" +version = "1.52.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8fc7f01b389ac15039e4dc9531aa973a135d7a4135281b12d7c1bc79fd57fffe" +dependencies = [ + "bytes", + "libc", + "mio", + "parking_lot", + "pin-project-lite", + "signal-hook-registry", + "socket2", + "tokio-macros", + "windows-sys", +] + +[[package]] +name = "tokio-macros" +version = "2.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "385a6cb71ab9ab790c5fe8d67f1645e6c450a7ce006a33de03daa956cf70a496" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "tokio-stream" +version = "0.1.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32da49809aab5c3bc678af03902d4ccddea2a87d028d86392a4b1560c6906c70" +dependencies = [ + "futures-core", + "pin-project-lite", + "tokio", +] + +[[package]] +name = "toml_datetime" +version = "1.1.1+spec-1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3165f65f62e28e0115a00b2ebdd37eb6f3b641855f9d636d3cd4103767159ad7" +dependencies = [ + "serde_core", +] + +[[package]] +name = "toml_edit" +version = "0.25.13+spec-1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6975367e4d2ef766d86af01ffad14b622fecc8d4357a998fbc4deb6e9bacaf9b" +dependencies = [ + "indexmap", + "toml_datetime", + "toml_parser", + "winnow", +] + +[[package]] +name = "toml_parser" +version = "1.1.2+spec-1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a2abe9b86193656635d2411dc43050282ca48aa31c2451210f4202550afb7526" +dependencies = [ + "winnow", +] + +[[package]] +name = "tower" +version = "0.5.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ebe5ef63511595f1344e2d5cfa636d973292adc0eec1f0ad45fae9f0851ab1d4" +dependencies = [ + "futures-core", + "futures-util", + "pin-project-lite", + "sync_wrapper", + "tokio", + "tower-layer", + "tower-service", + "tracing", +] + +[[package]] +name = "tower-layer" +version = "0.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "121c2a6cda46980bb0fcd1647ffaf6cd3fc79a013de288782836f6df9c48780e" + +[[package]] +name = "tower-service" +version = "0.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8df9b6e13f2d32c91b9bd719c00d1958837bc7dec474d94952798cc8e69eeec3" + +[[package]] +name = "tracing" +version = "0.1.44" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "63e71662fa4b2a2c3a26f570f037eb95bb1f85397f3cd8076caed2f026a6d100" +dependencies = [ + "log", + "pin-project-lite", + "tracing-attributes", + "tracing-core", +] + +[[package]] +name = "tracing-attributes" +version = "0.1.31" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7490cfa5ec963746568740651ac6781f701c9c5ea257c58e057f3ba8cf69e8da" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "tracing-core" +version = "0.1.36" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "db97caf9d906fbde555dd62fa95ddba9eecfd14cb388e4f491a66d74cd5fb79a" +dependencies = [ + "once_cell", + "valuable", +] + +[[package]] +name = "tracing-log" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ee855f1f400bd0e5c02d150ae5de3840039a3f54b025156404e34c23c03f47c3" +dependencies = [ + "log", + "once_cell", + "tracing-core", +] + +[[package]] +name = "tracing-subscriber" +version = "0.3.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cb7f578e5945fb242538965c2d0b04418d38ec25c79d160cd279bf0731c8d319" +dependencies = [ + "matchers", + "nu-ansi-term", + "once_cell", + "regex-automata", + "sharded-slab", + "smallvec", + "thread_local", + "tracing", + "tracing-core", + "tracing-log", +] + +[[package]] +name = "unicode-ident" +version = "1.0.24" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75" + +[[package]] +name = "unicode-normalization-alignments" +version = "0.1.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "43f613e4fa046e69818dd287fdc4bc78175ff20331479dab6e1b0f98d57062de" +dependencies = [ + "smallvec", +] + +[[package]] +name = "unicode-segmentation" +version = "1.13.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c6f5d3c3b1bf09027a88a6bc961fc00497d651009560b5463668dc81b0fa87a8" + +[[package]] +name = "unicode_categories" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "39ec24b3121d976906ece63c9daad25b85969647682eee313cb5779fdd69e14e" + +[[package]] +name = "utf8_iter" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6c140620e7ffbb22c2dee59cafe6084a59b5ffc27a8859a5f0d494b5d52b6be" + +[[package]] +name = "utf8parse" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "06abde3611657adf66d383f00b093d7faecc7fa57071cce2578660c9f1010821" + +[[package]] +name = "valuable" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ba73ea9cf16a25df0c8caa16c51acb937d5712a8429db78a3ee29d5dcacd3a65" + +[[package]] +name = "version_check" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" + +[[package]] +name = "wasi" +version = "0.11.1+wasi-snapshot-preview1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ccf3ec651a847eb01de73ccad15eb7d99f80485de043efb2f370cd654f4ea44b" + +[[package]] +name = "wasip2" +version = "1.0.4+wasi-0.2.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b67efb37e106e55ce722a510d6b5f9c17f083e5fc79afc2badeb12cc313d9487" +dependencies = [ + "wit-bindgen", +] + +[[package]] +name = "windows-link" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" + +[[package]] +name = "windows-sys" +version = "0.61.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ae137229bcbd6cdf0f7b80a31df61766145077ddf49416a728b02cb3921ff3fc" +dependencies = [ + "windows-link", +] + +[[package]] +name = "winnow" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "23b97319f7b8343df12cc98938e5c3eb436064524c8d2b4e30a1d3a36eecdf81" +dependencies = [ + "memchr", +] + +[[package]] +name = "wit-bindgen" +version = "0.57.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ebf944e87a7c253233ad6766e082e3cd714b5d03812acc24c318f549614536e" + +[[package]] +name = "writeable" +version = "0.6.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ffae5123b2d3fc086436f8834ae3ab053a283cfac8fe0a0b8eaae044768a4c4" + +[[package]] +name = "yoke" +version = "0.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "709fe23a0424b6a435d82152b1bd3fdfb0833487d5fa90d05d42762a9891fef5" +dependencies = [ + "stable_deref_trait", + "yoke-derive", + "zerofrom", +] + +[[package]] +name = "yoke-derive" +version = "0.8.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "de844c262c8848816172cef550288e7dc6c7b7814b4ee56b3e1553f275f1858e" +dependencies = [ + "proc-macro2", + "quote", + "syn", + "synstructure", +] + +[[package]] +name = "zerocopy" +version = "0.8.54" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b7cbbc0a705a0fd05cc3676525980d2bf5a9bc4adac6d6475209a7887cf59d19" +dependencies = [ + "zerocopy-derive", +] + +[[package]] +name = "zerocopy-derive" +version = "0.8.54" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e2e817b7b52d0c7358d3246da9d69935ebb18116b2b102b4230dac079b4862f5" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "zerofrom" +version = "0.1.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0ec05a11813ea801ff6d75110ad09cd0824ddba17dfe17128ea0d5f68e6c5272" +dependencies = [ + "zerofrom-derive", +] + +[[package]] +name = "zerofrom-derive" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "11532158c46691caf0f2593ea8358fed6bbf68a0315e80aae9bd41fbade684a1" +dependencies = [ + "proc-macro2", + "quote", + "syn", + "synstructure", +] + +[[package]] +name = "zerotrie" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0f9152d31db0792fa83f70fb2f83148effb5c1f5b8c7686c3459e361d9bc20bf" +dependencies = [ + "displaydoc", + "yoke", + "zerofrom", +] + +[[package]] +name = "zerovec" +version = "0.11.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "90f911cbc359ab6af17377d242225f4d75119aec87ea711a880987b18cd7b239" +dependencies = [ + "yoke", + "zerofrom", + "zerovec-derive", +] + +[[package]] +name = "zerovec-derive" +version = "0.11.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "625dc425cab0dca6dc3c3319506e6593dcb08a9f387ea3b284dbd52a92c40555" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "zmij" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29666d0abbfad1e3dc4dcf6144730dd3a3ab225bbbdac83319345b1b44ccfc1b" diff --git a/crates/skippy-engine-mlx/Cargo.toml b/crates/skippy-engine-mlx/Cargo.toml new file mode 100644 index 0000000000..178d5ae313 --- /dev/null +++ b/crates/skippy-engine-mlx/Cargo.toml @@ -0,0 +1,62 @@ +# MLX (Metal) serving engine for mesh-llm. +# +# Serves HF safetensors tensor models over mesh-llm's REAL OpenAI-compatible +# frontend (`openai-frontend::OpenAiBackend`), goose-style, on Apple Silicon. +# +# This crate is intentionally its OWN cargo workspace (note the empty +# `[workspace]` table below) so it does NOT join the main mesh-llm workspace. +# Reason: it path-depends on a sibling checkout of the safemlx fork at +# ../../../safemlx, which is external to this repo. Making it a normal workspace +# member would force every `cargo build`/CI run to resolve that external path, +# breaking builds for anyone without the checkout. Promotion to a real member +# (with safemlx pinned to a git rev) is documented in WIRING.md. +[package] +name = "skippy-engine-mlx" +version = "0.0.0" +edition = "2021" +publish = false + +[workspace] + +[lib] +name = "skippy_engine_mlx" +path = "src/lib.rs" + +[[bin]] +name = "mlx-serve" +path = "src/bin/mlx-serve.rs" + +[features] +default = [] +# Enable the real MLX engine. macOS-only in practice (code is cfg-gated to +# target_os = "macos"); the deps still resolve elsewhere but compile to nothing. +mlx = ["dep:safemlx", "dep:safemlx-lm", "dep:tokenizers"] + +[dependencies] +# The real mesh-llm OpenAI frontend — this is what proves we serve over the +# same surface the shipped binary uses, not a toy. +openai-frontend = { path = "../openai-frontend" } +async-trait = "0.1" +async-stream = "0.3" +anyhow = "1" +axum = "0.8" +clap = { version = "4", features = ["derive"] } +futures-core = "0.3" +serde_json = "1" +tokio = { version = "1", features = ["full"] } +tracing = "0.1" +tracing-subscriber = { version = "0.3", features = ["env-filter"] } + +# MLX stack (pinned to the safemlx fork checkout; see WIRING.md for git-pin plan). +# Same feature set goose uses: accelerate + metal + safetensors. +safemlx = { path = "../../../safemlx/safemlx", default-features = false, features = [ + "accelerate", + "metal", + "safetensors", +], optional = true } +safemlx-lm = { path = "../../../safemlx/safemlx-lm", optional = true } +# Standalone tokenizer for incremental (streaming) detokenization without +# borrowing the model. `onig` matches safemlx-lm so the build dedupes. +tokenizers = { version = "0.23", default-features = false, features = [ + "onig", +], optional = true } diff --git a/crates/skippy-engine-mlx/WIRING.md b/crates/skippy-engine-mlx/WIRING.md new file mode 100644 index 0000000000..275e238482 --- /dev/null +++ b/crates/skippy-engine-mlx/WIRING.md @@ -0,0 +1,133 @@ +# Wiring `skippy-engine-mlx` into mesh-llm + +This crate is a **working, self-contained MLX (Metal) serving engine** that +already serves HF safetensors models over mesh-llm's real OpenAI frontend +(`openai_frontend::router_for`). It is intentionally standalone right now — its +own cargo workspace, path-depending the local safemlx fork — so it does not +perturb the main workspace or CI. This document is the concrete plan to promote +it into the shipped binary so that **on a Mac, `mesh-llm serve` can run an MLX +tensor model and users can pick one from `/v1/models`.** + +## What already works (this crate, today) + +- `MlxEngine` — a dedicated OS worker thread owns the non-`Send` MLX objects + (model, streams, arrays); the outside world talks to it with `Send` channels. +- `MlxBackend: openai_frontend::OpenAiBackend` — `models`, `chat_completion`, + `chat_completion_stream` (SSE), with usage accounting and incremental + detokenization. +- `mlx-serve` bin — `router_for(Arc)` + `axum::serve`. +- Verified on Apple Silicon (Metal): `/v1/models`, non-stream chat, and stream + chat all return real Qwen3-0.6B generations. Source precision ~321 tok/s; + JIT-4bit ~604 tok/s (see `../../spikes/mlx-solo/FINDINGS.md`). + +## Promotion plan (the actual PR) + +### 1. Make it a real workspace member with a pinned safemlx + +- Add `crates/skippy-engine-mlx` to root `Cargo.toml` `members` and remove its + local `[workspace]` table. +- Replace the `path = "../../../safemlx/..."` deps with **git-rev pins** of + `jbg/safemlx` (published crates are broken for Qwen3 — see FINDINGS §"supply + chain"). Carry the two small loader fixes as a patch/branch until upstreamed. +- Keep the crate's `mlx` feature; gate all MLX code with + `#[cfg(all(feature = "mlx", target_os = "macos"))]` (already done). +- Update `scripts/affected-crates.sh`, `scripts/plan-clippy-batches.sh`, + `scripts/publish-crates.sh` `WORKSPACE_MEMBERS`, and run + `cargo run -p xtask -- repo-consistency ci-crate-lists`. Because MLX is a heavy + native lane, add it to the CI backend-gating like the other native features + (only build/test the `mlx` feature on macOS runners). + +### 2. Depend on it from host-runtime, macOS-gated + +In `crates/mesh-llm-host-runtime/Cargo.toml`: + +```toml +[target.'cfg(target_os = "macos")'.dependencies] +skippy-engine-mlx = { path = "../skippy-engine-mlx", optional = true } + +[features] +mlx = ["dep:skippy-engine-mlx", "skippy-engine-mlx/mlx"] +``` + +Propagate a `mlx` feature up through `crates/mesh-llm/Cargo.toml`, and enable it +by default only on macOS builds in the release packaging. + +### 3. Add an `Mlx` variant to the launch enum + +`crates/mesh-llm-host-runtime/src/runtime/local.rs`: + +- `LocalRuntimeBackendHandle` currently has one variant (`Skippy { .. }`). Add + (macOS+feature gated): + + ```rust + #[cfg(all(feature = "mlx", target_os = "macos"))] + Mlx { backend: Arc, http: MlxHttpHandle, _death_tx: ... }, + ``` + +- Every `match &self.inner { LocalRuntimeBackendHandle::Skippy { .. } => ... }` + in `local.rs` (pid, ctx_used_tokens, openai_guardrails, llama_slots_snapshot, + set_openai_guardrail_mode, shutdown/http accessor) needs an `Mlx` arm. Most map + to simple/None-ish values since MLX has no llama slots or GGUF guardrail state. + +### 4. Route safetensors models to the MLX branch at launch + +`start_runtime_local_model` currently branches: + +``` +if is_layer_package_ref(..) { layer_package } else { skippy (direct GGUF) } +``` + +Add a first branch: if the resolved model is `ModelFormat::Safetensors` (the +model layer already classifies this — `crates/model-artifact` `ModelFormat`, and +`models/resolve` already detects `is_primary_mlx_weight_file`), and we're on +macOS with the `mlx` feature, start an `MlxEngine` instead: + +```rust +#[cfg(all(feature = "mlx", target_os = "macos"))] +if resolved_format == ModelFormat::Safetensors { + return start_runtime_mlx_model(spec, model_name, plan).await; +} +``` + +`start_runtime_mlx_model` mirrors `start_runtime_skippy_model`: build an +`MlxEngineConfig` from the resolved model dir + planned ctx/limits, `spawn` the +engine on a blocking task, wrap `MlxBackend` in the embedded HTTP handle +(`openai_frontend::router_for`), and return a `LocalRuntimeModelHandle` whose +`backend` string is `"mlx"`. + +### 5. Model discovery / listing already works + +The model layer already discovers, downloads, catalogs, and lists MLX +safetensors repos (`crates/mesh-llm-host-runtime/src/models/catalog.rs`, +`.../models/resolve`, `crates/model-resolver`). No change needed for a user to +*see* MLX models; the missing piece was purely the serving engine, which this +crate provides. Auto-behavior: on a Mac, a resolved safetensors model simply +routes to the MLX engine. + +### 6. Auto on Mac + user selection + +- With the `mlx` feature enabled by default on macOS builds, serving a + safetensors model "just works" with no extra flags. +- Users pick a model the same way as today: `mesh-llm serve --model ` + or from `~/.mesh-llm/config.toml`; safetensors → MLX, GGUF → llama.cpp. +- Optionally add `--serving-backend mlx|llama` to force the engine when a model + is available in both formats (parallels the existing skippy backend selector + noted in `docs/SKIPPY.md`). + +## Out of scope for the first PR + +- **Splits / staged execution.** This is single-stage, whole-model serving. The + staged `StageEngine` trait, partial-load, and activation-frame work (plan §6–§8) + are separate and remain gated on the go/no-go spikes. +- **Tool calling / reasoning parsing.** goose's `mlx.rs` has native + emulated + tool parsing and thinking-output filtering worth porting later; this crate + streams raw model text (including `` blocks) for now. +- **Draft/speculative decoding** (goose's `gemma4_mtp`). + +## Testing the promoted path + +- `just build` on macOS with the `mlx` feature. +- `mesh-llm serve --model ` → confirm the model appears in + `/v1/models` and `/v1/chat/completions` returns a generation. +- Confirm non-macOS / no-feature builds are byte-for-byte unaffected (the crate + and its deps compile out entirely). diff --git a/crates/skippy-engine-mlx/src/backend.rs b/crates/skippy-engine-mlx/src/backend.rs new file mode 100644 index 0000000000..02de047fd8 --- /dev/null +++ b/crates/skippy-engine-mlx/src/backend.rs @@ -0,0 +1,248 @@ +//! `OpenAiBackend` implementation driving the MLX worker. +//! +//! This is the adapter between mesh-llm's real OpenAI-compatible frontend and +//! the MLX engine: it converts `ChatCompletionRequest`s into engine jobs and +//! turns the worker's `TokenMsg` stream into OpenAI chat responses / SSE chunks. + +use std::sync::Arc; +use std::time::{SystemTime, UNIX_EPOCH}; + +use async_trait::async_trait; +use openai_frontend::backend::{ + ChatCompletionStream, OpenAiBackend, OpenAiRequestContext, OpenAiResult, +}; +use openai_frontend::chat::{ + message_content_to_text, AssistantMessage, ChatCompletionChoice, ChatCompletionChunk, + ChatCompletionChunkChoice, ChatCompletionDelta, ChatCompletionRequest, ChatCompletionResponse, +}; +use openai_frontend::common::{completion_id, FinishReason, Usage}; +use openai_frontend::errors::OpenAiError; +use openai_frontend::models::ModelObject; + +use crate::engine::{ChatTurn, FinishReason as EngineFinish, GenerateRequest, MlxEngine, TokenMsg}; + +pub struct MlxBackend { + engine: Arc, +} + +impl MlxBackend { + pub fn new(engine: MlxEngine) -> Self { + Self { + engine: Arc::new(engine), + } + } + + fn build_request(&self, request: &ChatCompletionRequest) -> OpenAiResult { + let messages: Vec = request + .messages + .iter() + .map(|m| ChatTurn { + role: m.role.clone(), + content: m + .content + .as_ref() + .and_then(message_content_to_text) + .unwrap_or_default(), + }) + .collect(); + if messages.is_empty() { + return Err(OpenAiError::invalid_request("no messages in request")); + } + let requested = request + .max_completion_tokens + .or(request.max_tokens) + .map(|n| n as usize); + Ok(GenerateRequest { + messages, + raw_prompt: None, + max_tokens: self.engine.clamp_max_tokens(requested), + }) + } +} + +fn now_secs() -> u64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|d| d.as_secs()) + .unwrap_or(0) +} + +fn map_finish(reason: EngineFinish) -> FinishReason { + match reason { + EngineFinish::Stop => FinishReason::Stop, + EngineFinish::Length => FinishReason::Length, + } +} + +fn role_chunk(id: &str, created: u64, model: &str) -> ChatCompletionChunk { + ChatCompletionChunk { + id: id.to_string(), + object: "chat.completion.chunk", + created, + model: model.to_string(), + choices: vec![ChatCompletionChunkChoice { + index: 0, + delta: ChatCompletionDelta { + role: Some("assistant"), + content: None, + reasoning_content: None, + tool_calls: None, + }, + logprobs: None, + finish_reason: None, + }], + usage: None, + } +} + +fn content_chunk(id: &str, created: u64, model: &str, text: String) -> ChatCompletionChunk { + ChatCompletionChunk { + id: id.to_string(), + object: "chat.completion.chunk", + created, + model: model.to_string(), + choices: vec![ChatCompletionChunkChoice { + index: 0, + delta: ChatCompletionDelta { + role: None, + content: Some(text), + reasoning_content: None, + tool_calls: None, + }, + logprobs: None, + finish_reason: None, + }], + usage: None, + } +} + +fn final_chunk( + id: &str, + created: u64, + model: &str, + finish: FinishReason, + usage: Usage, +) -> ChatCompletionChunk { + ChatCompletionChunk { + id: id.to_string(), + object: "chat.completion.chunk", + created, + model: model.to_string(), + choices: vec![ChatCompletionChunkChoice { + index: 0, + delta: ChatCompletionDelta { + role: None, + content: None, + reasoning_content: None, + tool_calls: None, + }, + logprobs: None, + finish_reason: Some(finish), + }], + usage: Some(usage), + } +} + +fn usage(prompt_tokens: u32, completion_tokens: u32) -> Usage { + Usage { + prompt_tokens, + completion_tokens, + total_tokens: prompt_tokens + completion_tokens, + prompt_tokens_details: None, + } +} + +#[async_trait] +impl OpenAiBackend for MlxBackend { + async fn models(&self) -> OpenAiResult> { + Ok(vec![ModelObject { + id: self.engine.model_id().to_string(), + object: "model", + created: now_secs(), + owned_by: "mesh-llm-mlx".to_string(), + }]) + } + + async fn chat_completion( + &self, + request: ChatCompletionRequest, + ) -> OpenAiResult { + let gen = self.build_request(&request)?; + let model = self.engine.model_id().to_string(); + let mut rx = self.engine.submit(gen); + + let mut text = String::new(); + let mut finish = FinishReason::Stop; + let mut used = usage(0, 0); + + while let Some(msg) = rx.recv().await { + match msg { + TokenMsg::Delta(delta) => text.push_str(&delta), + TokenMsg::Done { + finish_reason, + prompt_tokens, + completion_tokens, + } => { + finish = map_finish(finish_reason); + used = usage(prompt_tokens, completion_tokens); + break; + } + TokenMsg::Error(e) => return Err(OpenAiError::internal(e)), + } + } + + Ok(ChatCompletionResponse { + id: completion_id("chatcmpl"), + object: "chat.completion", + created: now_secs(), + model, + choices: vec![ChatCompletionChoice { + index: 0, + message: AssistantMessage { + role: "assistant", + content: Some(text), + reasoning_content: None, + tool_calls: None, + }, + logprobs: None, + finish_reason: Some(finish), + }], + usage: used, + timings: None, + }) + } + + async fn chat_completion_stream( + &self, + request: ChatCompletionRequest, + _context: OpenAiRequestContext, + ) -> OpenAiResult { + let gen = self.build_request(&request)?; + let model = self.engine.model_id().to_string(); + let mut rx = self.engine.submit(gen); + let id = completion_id("chatcmpl"); + let created = now_secs(); + + let stream = async_stream::stream! { + yield Ok(role_chunk(&id, created, &model)); + while let Some(msg) = rx.recv().await { + match msg { + TokenMsg::Delta(delta) => { + yield Ok(content_chunk(&id, created, &model, delta)); + } + TokenMsg::Done { finish_reason, prompt_tokens, completion_tokens } => { + let used = usage(prompt_tokens, completion_tokens); + yield Ok(final_chunk(&id, created, &model, map_finish(finish_reason), used)); + break; + } + TokenMsg::Error(e) => { + yield Err(OpenAiError::internal(e)); + break; + } + } + } + }; + + Ok(Box::pin(stream)) + } +} diff --git a/crates/skippy-engine-mlx/src/bin/mlx-serve.rs b/crates/skippy-engine-mlx/src/bin/mlx-serve.rs new file mode 100644 index 0000000000..d43a6811d7 --- /dev/null +++ b/crates/skippy-engine-mlx/src/bin/mlx-serve.rs @@ -0,0 +1,104 @@ +//! `mlx-serve` — stand up mesh-llm's real OpenAI-compatible frontend backed by +//! the MLX (Metal) engine, serving an HF safetensors model. +//! +//! This exists to prove the engine serves over the SAME `openai-frontend` +//! surface the shipped binary uses (`router_for(Arc)`), not a +//! bespoke HTTP handler. + +#[cfg(all(feature = "mlx", target_os = "macos"))] +mod real { + use std::path::PathBuf; + use std::sync::Arc; + + use anyhow::{Context, Result}; + use clap::Parser; + use skippy_engine_mlx::{MlxBackend, MlxEngine, MlxEngineConfig}; + + #[derive(Parser, Debug)] + #[command(about = "Serve an MLX safetensors model over the mesh-llm OpenAI frontend")] + struct Cli { + /// Model directory (HF safetensors: config.json + tokenizer.json + *.safetensors). + #[arg(short, long)] + model: PathBuf, + + /// Model id advertised on /v1/models (defaults to the directory name). + #[arg(long)] + model_id: Option, + + /// JIT-quantize eligible dense weights to this bit width on load (e.g. 4, 8). + #[arg(short, long)] + quantize: Option, + + #[arg(long, default_value_t = 64)] + quant_group_size: i32, + + #[arg(long, default_value_t = 512)] + default_max_tokens: usize, + + #[arg(long, default_value_t = 4096)] + max_tokens_cap: usize, + + #[arg(long, default_value = "127.0.0.1:11434")] + bind: String, + } + + pub async fn main() -> Result<()> { + tracing_subscriber::fmt() + .with_env_filter( + tracing_subscriber::EnvFilter::try_from_default_env() + .unwrap_or_else(|_| "info".into()), + ) + .init(); + + let cli = Cli::parse(); + let model_id = cli.model_id.clone().unwrap_or_else(|| { + cli.model + .file_name() + .map(|s| s.to_string_lossy().to_string()) + .unwrap_or_else(|| "mlx-model".to_string()) + }); + + let config = MlxEngineConfig { + model_dir: cli.model.clone(), + model_id: model_id.clone(), + quantize_bits: cli.quantize, + quant_group_size: cli.quant_group_size, + default_max_tokens: cli.default_max_tokens, + max_tokens_cap: cli.max_tokens_cap, + }; + + tracing::info!("loading MLX model from {} ...", cli.model.display()); + let engine = tokio::task::spawn_blocking(move || MlxEngine::spawn(config)) + .await + .context("join MLX load task")??; + + let backend = Arc::new(MlxBackend::new(engine)); + let app = openai_frontend::router::router_for(backend); + + let listener = tokio::net::TcpListener::bind(&cli.bind) + .await + .with_context(|| format!("bind {}", cli.bind))?; + tracing::info!( + "MLX serving '{}' on http://{}/v1 (try: GET /v1/models)", + model_id, + cli.bind + ); + axum::serve(listener, app).await.context("axum serve")?; + Ok(()) + } +} + +#[cfg(all(feature = "mlx", target_os = "macos"))] +#[tokio::main] +async fn main() -> anyhow::Result<()> { + real::main().await +} + +#[cfg(not(all(feature = "mlx", target_os = "macos")))] +fn main() { + eprintln!( + "mlx-serve was built without MLX support.\n\ + Rebuild on Apple Silicon with `--features mlx` to enable the MLX engine." + ); + std::process::exit(1); +} diff --git a/crates/skippy-engine-mlx/src/engine.rs b/crates/skippy-engine-mlx/src/engine.rs new file mode 100644 index 0000000000..2a4d524428 --- /dev/null +++ b/crates/skippy-engine-mlx/src/engine.rs @@ -0,0 +1,311 @@ +//! MLX generation engine backed by a dedicated OS worker thread. +//! +//! MLX arrays, streams, and the loaded model wrap raw C pointers and are neither +//! `Send` nor `Sync`. Rather than fight that, we confine every MLX object to a +//! single worker thread that owns them for its whole life, and talk to it only +//! with `Send` messages: +//! +//! - a `Send + Sync` job channel (tokio unbounded) carries generation requests; +//! - each job carries a per-request token channel the worker streams results on. +//! +//! This also naturally serializes GPU access (one generation at a time), which +//! matches how goose drives safemlx today. + +use std::path::PathBuf; +use std::thread; +use std::time::Instant; + +use anyhow::{anyhow, Result}; +use serde_json::{json, Value}; +use tokio::sync::mpsc; + +use safemlx::transforms::async_eval; +use safemlx::{Device, DeviceType, Stream}; +use safemlx_lm::models::input::{InputPart, ModelInput}; +use safemlx_lm::models::{LoadedModel, ModelLoadOptions}; +use safemlx_lm::quantization::AffineQuantization; +use safemlx_lm::sampler::DefaultSampler; + +/// How the worker should load and run a model. +#[derive(Clone, Debug)] +pub struct MlxEngineConfig { + pub model_dir: PathBuf, + pub model_id: String, + /// JIT-quantize eligible dense weights to this bit width on load (Metal only). + pub quantize_bits: Option, + pub quant_group_size: i32, + pub default_max_tokens: usize, + pub max_tokens_cap: usize, +} + +/// One chat turn, in `Send` form (no MLX types). +#[derive(Clone, Debug)] +pub struct ChatTurn { + pub role: String, + pub content: String, +} + +/// A generation request handed to the worker. +#[derive(Debug)] +pub struct GenerateRequest { + pub messages: Vec, + /// If set, skip the chat template and feed this text verbatim. + pub raw_prompt: Option, + pub max_tokens: usize, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum FinishReason { + Stop, + Length, +} + +/// Streamed output from the worker for one request. +#[derive(Debug)] +pub enum TokenMsg { + Delta(String), + Done { + finish_reason: FinishReason, + prompt_tokens: u32, + completion_tokens: u32, + }, + Error(String), +} + +struct Job { + req: GenerateRequest, + reply: mpsc::UnboundedSender, +} + +/// Handle to the MLX worker thread. `Send + Sync`, safe to share in an `Arc`. +pub struct MlxEngine { + job_tx: mpsc::UnboundedSender, + config: MlxEngineConfig, +} + +impl MlxEngine { + /// Spawns the worker and blocks until the model has finished loading. + /// Call from a blocking context (e.g. `tokio::task::spawn_blocking`). + pub fn spawn(config: MlxEngineConfig) -> Result { + let (job_tx, job_rx) = mpsc::unbounded_channel::(); + let (ready_tx, ready_rx) = std::sync::mpsc::channel::>(); + + let worker_config = config.clone(); + thread::Builder::new() + .name("mlx-engine".into()) + .spawn(move || run_worker(worker_config, job_rx, ready_tx))?; + + match ready_rx.recv() { + Ok(Ok(())) => Ok(Self { job_tx, config }), + Ok(Err(e)) => Err(anyhow!("MLX model load failed: {e}")), + Err(_) => Err(anyhow!("MLX worker exited before signalling readiness")), + } + } + + pub fn model_id(&self) -> &str { + &self.config.model_id + } + + pub fn clamp_max_tokens(&self, requested: Option) -> usize { + requested + .unwrap_or(self.config.default_max_tokens) + .clamp(1, self.config.max_tokens_cap) + } + + /// Submits a request and returns the channel its tokens will stream on. + pub fn submit(&self, req: GenerateRequest) -> mpsc::UnboundedReceiver { + let (tx, rx) = mpsc::unbounded_channel(); + if self + .job_tx + .send(Job { + req, + reply: tx.clone(), + }) + .is_err() + { + let _ = tx.send(TokenMsg::Error("MLX worker is not running".into())); + } + rx + } +} + +struct LoadedEngine { + model: LoadedModel, + stream: Stream, + tokenizer: tokenizers::Tokenizer, + eos: Vec, +} + +fn load_engine(config: &MlxEngineConfig) -> Result { + // Metal GPU stream for compute, CPU stream for weight staging (goose's split). + let stream = Stream::new_with_device(&Device::new(DeviceType::Gpu, 0)); + let weights_stream = Stream::new_with_device(&Device::new(DeviceType::Cpu, 0)); + + let options = match config.quantize_bits { + Some(bits) => ModelLoadOptions::with_quantization(AffineQuantization::new( + config.quant_group_size, + bits, + )?), + None => ModelLoadOptions::default(), + }; + + let started = Instant::now(); + let model = + LoadedModel::load_with_options(&config.model_dir, options, &stream, &weights_stream) + .map_err(|e| anyhow!("load {}: {e}", config.model_dir.display()))?; + stream.synchronize().map_err(|e| anyhow!("sync: {e}"))?; + + let tokenizer = tokenizers::Tokenizer::from_file(config.model_dir.join("tokenizer.json")) + .map_err(|e| anyhow!("tokenizer.json: {e}"))?; + let eos = model.eos_token_ids().to_vec(); + + tracing::info!( + model = %config.model_id, + kind = model.model_type(), + load_secs = started.elapsed().as_secs_f64(), + "MLX model loaded" + ); + Ok(LoadedEngine { + model, + stream, + tokenizer, + eos, + }) +} + +fn run_worker( + config: MlxEngineConfig, + mut job_rx: mpsc::UnboundedReceiver, + ready_tx: std::sync::mpsc::Sender>, +) { + let mut engine = match load_engine(&config) { + Ok(engine) => { + let _ = ready_tx.send(Ok(())); + engine + } + Err(e) => { + let _ = ready_tx.send(Err(e.to_string())); + return; + } + }; + + while let Some(job) = job_rx.blocking_recv() { + let reply = job.reply.clone(); + if let Err(e) = generate_one(&mut engine, job) { + let _ = reply.send(TokenMsg::Error(e.to_string())); + } + } +} + +fn build_prompt(model: &mut LoadedModel, req: &GenerateRequest) -> Result<(String, bool)> { + if let Some(raw) = &req.raw_prompt { + return Ok((raw.clone(), true)); + } + let messages: Vec = req + .messages + .iter() + .map(|turn| json!({"role": turn.role, "content": turn.content})) + .collect(); + let rendered = model + .apply_chat_template_json(vec![messages], None, true) + .map_err(|e| anyhow!("chat template: {e}"))?; + match rendered { + Some(prompt) => Ok((prompt, false)), + None => { + let fallback = req + .messages + .last() + .map(|turn| turn.content.clone()) + .unwrap_or_default(); + Ok((fallback, true)) + } + } +} + +fn generate_one(engine: &mut LoadedEngine, job: Job) -> Result<()> { + let LoadedEngine { + model, + stream, + tokenizer, + eos, + } = engine; + let reply = job.reply; + + let (prompt, add_special) = build_prompt(model, &job.req)?; + let tokens = model + .encode_to_array(&prompt, add_special, stream) + .map_err(|e| anyhow!("encode: {e}"))?; + let prompt_tokens = tokens.shape()[1] as u32; + + let mut cache = model.new_cache(); + let parts = [InputPart::text_token_ids(&tokens)]; + let input = ModelInput::new(&parts); + let mut generator = model.generate_input_with_cache_sampler( + &mut cache, + 0.0, + input, + None, + stream, + DefaultSampler, + ); + + let mut ids: Vec = Vec::with_capacity(job.req.max_tokens); + let mut emitted = String::new(); + let mut finish = FinishReason::Length; + + let mut current = generator.next().transpose().map_err(|e| anyhow!("{e}"))?; + for index in 0..job.req.max_tokens { + let Some(token) = current.take() else { + finish = FinishReason::Stop; + break; + }; + + // Start the next decode before reading this token back (mlx-lm's + // one-token async pipeline overlaps compute with host readback). + let next = if index + 1 < job.req.max_tokens { + let next = generator.next(); + if let Some(Ok(next_token)) = next.as_ref() { + async_eval([next_token]).map_err(|e| anyhow!("async_eval: {e}"))?; + } + next + } else { + None + }; + + let token_id = token.item::(&*stream); + if eos.contains(&token_id) { + finish = FinishReason::Stop; + break; + } + ids.push(token_id); + + // Incremental detokenization: decode the whole id sequence and emit only + // the newly appended suffix. Robust for byte-level BPE where one char can + // span multiple tokens. + let full = tokenizer + .decode(&ids, true) + .map_err(|e| anyhow!("decode: {e}"))?; + if let Some(delta) = full.strip_prefix(emitted.as_str()) { + if !delta.is_empty() { + if reply.send(TokenMsg::Delta(delta.to_string())).is_err() { + return Ok(()); // client hung up + } + emitted = full; + } + } else { + emitted = full; // rare re-render; resync silently + } + + if reply.is_closed() { + return Ok(()); + } + current = next.transpose().map_err(|e| anyhow!("{e}"))?; + } + + let _ = reply.send(TokenMsg::Done { + finish_reason: finish, + prompt_tokens, + completion_tokens: ids.len() as u32, + }); + Ok(()) +} diff --git a/crates/skippy-engine-mlx/src/lib.rs b/crates/skippy-engine-mlx/src/lib.rs new file mode 100644 index 0000000000..5dcffca700 --- /dev/null +++ b/crates/skippy-engine-mlx/src/lib.rs @@ -0,0 +1,23 @@ +//! MLX (Metal) serving engine for mesh-llm. +//! +//! Serves HF safetensors tensor models over mesh-llm's real OpenAI-compatible +//! frontend (`openai_frontend::OpenAiBackend`), goose-style, on Apple Silicon. +//! +//! All MLX-touching code is gated behind BOTH the `mlx` cargo feature AND +//! `target_os = "macos"`. On any other target, or without the feature, this +//! crate compiles to an empty shell so it never burdens non-Apple builds. + +#[cfg(all(feature = "mlx", target_os = "macos"))] +mod backend; +#[cfg(all(feature = "mlx", target_os = "macos"))] +mod engine; + +#[cfg(all(feature = "mlx", target_os = "macos"))] +pub use backend::MlxBackend; +#[cfg(all(feature = "mlx", target_os = "macos"))] +pub use engine::{ChatTurn, GenerateRequest, MlxEngine, MlxEngineConfig}; + +/// True when this build actually contains the MLX engine. +pub const fn mlx_available() -> bool { + cfg!(all(feature = "mlx", target_os = "macos")) +} From 8eded3afd01af0af34ddaf52052f4f5883b02a3b Mon Sep 17 00:00:00 2001 From: michaelneale Date: Fri, 17 Jul 2026 05:35:55 +1000 Subject: [PATCH 06/37] feat(mlx): serve safetensors models to Metal via git-pinned safemlx 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. --- crates/skippy-engine-mlx/Cargo.lock | 6 ++ crates/skippy-engine-mlx/Cargo.toml | 27 +++++--- crates/skippy-engine-mlx/WIRING.md | 69 ++++++++++++++----- crates/skippy-engine-mlx/src/bin/mlx-serve.rs | 9 --- crates/skippy-engine-mlx/src/engine.rs | 19 +---- 5 files changed, 78 insertions(+), 52 deletions(-) diff --git a/crates/skippy-engine-mlx/Cargo.lock b/crates/skippy-engine-mlx/Cargo.lock index 6380fc69a4..ce41d7739d 100644 --- a/crates/skippy-engine-mlx/Cargo.lock +++ b/crates/skippy-engine-mlx/Cargo.lock @@ -1338,6 +1338,7 @@ checksum = "9774ba4a74de5f7b1c1451ed6cd5285a32eddb5cccb8cc655a4e50009e06477f" [[package]] name = "safemlx" version = "0.1.3" +source = "git+https://github.com/jbg/safemlx?rev=4e53c5e#4e53c5ecd7cbd91c0dfd0992a3c731ca2c36e9c7" dependencies = [ "bytemuck", "dyn-clone", @@ -1361,6 +1362,7 @@ dependencies = [ [[package]] name = "safemlx-internal-macros" version = "0.1.1" +source = "git+https://github.com/jbg/safemlx?rev=4e53c5e#4e53c5ecd7cbd91c0dfd0992a3c731ca2c36e9c7" dependencies = [ "darling 0.23.0", "itertools 0.15.0", @@ -1372,6 +1374,7 @@ dependencies = [ [[package]] name = "safemlx-lm" version = "0.4.1" +source = "git+https://github.com/jbg/safemlx?rev=4e53c5e#4e53c5ecd7cbd91c0dfd0992a3c731ca2c36e9c7" dependencies = [ "anyhow", "clap", @@ -1390,6 +1393,7 @@ dependencies = [ [[package]] name = "safemlx-lm-utils" version = "0.1.4" +source = "git+https://github.com/jbg/safemlx?rev=4e53c5e#4e53c5ecd7cbd91c0dfd0992a3c731ca2c36e9c7" dependencies = [ "minijinja", "minijinja-contrib", @@ -1402,6 +1406,7 @@ dependencies = [ [[package]] name = "safemlx-macros" version = "0.1.1" +source = "git+https://github.com/jbg/safemlx?rev=4e53c5e#4e53c5ecd7cbd91c0dfd0992a3c731ca2c36e9c7" dependencies = [ "darling 0.23.0", "proc-macro2", @@ -1412,6 +1417,7 @@ dependencies = [ [[package]] name = "safemlx-sys" version = "0.1.3" +source = "git+https://github.com/jbg/safemlx?rev=4e53c5e#4e53c5ecd7cbd91c0dfd0992a3c731ca2c36e9c7" dependencies = [ "cc", "cmake", diff --git a/crates/skippy-engine-mlx/Cargo.toml b/crates/skippy-engine-mlx/Cargo.toml index 178d5ae313..d4014cca37 100644 --- a/crates/skippy-engine-mlx/Cargo.toml +++ b/crates/skippy-engine-mlx/Cargo.toml @@ -3,13 +3,17 @@ # Serves HF safetensors tensor models over mesh-llm's REAL OpenAI-compatible # frontend (`openai-frontend::OpenAiBackend`), goose-style, on Apple Silicon. # -# This crate is intentionally its OWN cargo workspace (note the empty -# `[workspace]` table below) so it does NOT join the main mesh-llm workspace. -# Reason: it path-depends on a sibling checkout of the safemlx fork at -# ../../../safemlx, which is external to this repo. Making it a normal workspace -# member would force every `cargo build`/CI run to resolve that external path, -# breaking builds for anyone without the checkout. Promotion to a real member -# (with safemlx pinned to a git rev) is documented in WIRING.md. +# It depends on the MLX-in-Rust stack (safemlx / safemlx-lm) pinned to a specific +# PUBLIC commit of github.com/jbg/safemlx. This is a deliberate git-rev pin, not +# a private fork: safemlx's crates.io releases currently produce garbage output +# for dense models (Qwen3, Llama) — verified this to be a library bug, not a +# prompting/template issue — while a recent upstream commit serves them +# correctly. We pin that commit and will swap to a normal version pin once +# safemlx cuts a working release. See WIRING.md. +# +# The crate keeps its OWN `[workspace]` table for now so the heavy MLX native +# build (CMake + full MLX C++ compile) never runs in unrelated builds/CI. +# Promotion to a real workspace member is documented in WIRING.md. [package] name = "skippy-engine-mlx" version = "0.0.0" @@ -47,14 +51,15 @@ tokio = { version = "1", features = ["full"] } tracing = "0.1" tracing-subscriber = { version = "0.3", features = ["env-filter"] } -# MLX stack (pinned to the safemlx fork checkout; see WIRING.md for git-pin plan). -# Same feature set goose uses: accelerate + metal + safetensors. -safemlx = { path = "../../../safemlx/safemlx", default-features = false, features = [ +# MLX stack — pinned to a specific PUBLIC commit of github.com/jbg/safemlx. +# See the header comment for why this is a git-rev pin rather than a crates.io +# version pin. Same feature set goose uses: accelerate + metal + safetensors. +safemlx = { git = "https://github.com/jbg/safemlx", rev = "4e53c5e", default-features = false, features = [ "accelerate", "metal", "safetensors", ], optional = true } -safemlx-lm = { path = "../../../safemlx/safemlx-lm", optional = true } +safemlx-lm = { git = "https://github.com/jbg/safemlx", rev = "4e53c5e", optional = true } # Standalone tokenizer for incremental (streaming) detokenization without # borrowing the model. `onig` matches safemlx-lm so the build dedupes. tokenizers = { version = "0.23", default-features = false, features = [ diff --git a/crates/skippy-engine-mlx/WIRING.md b/crates/skippy-engine-mlx/WIRING.md index 275e238482..96d1599d1e 100644 --- a/crates/skippy-engine-mlx/WIRING.md +++ b/crates/skippy-engine-mlx/WIRING.md @@ -3,10 +3,10 @@ This crate is a **working, self-contained MLX (Metal) serving engine** that already serves HF safetensors models over mesh-llm's real OpenAI frontend (`openai_frontend::router_for`). It is intentionally standalone right now — its -own cargo workspace, path-depending the local safemlx fork — so it does not -perturb the main workspace or CI. This document is the concrete plan to promote -it into the shipped binary so that **on a Mac, `mesh-llm serve` can run an MLX -tensor model and users can pick one from `/v1/models`.** +own cargo workspace — so the heavy MLX native build never runs in unrelated +builds/CI. This document is the concrete plan to promote it into the shipped +binary so that **on a Mac, `mesh-llm serve` can run an MLX tensor model and +users can pick one from `/v1/models`.** ## What already works (this crate, today) @@ -16,19 +16,38 @@ tensor model and users can pick one from `/v1/models`.** `chat_completion_stream` (SSE), with usage accounting and incremental detokenization. - `mlx-serve` bin — `router_for(Arc)` + `axum::serve`. -- Verified on Apple Silicon (Metal): `/v1/models`, non-stream chat, and stream - chat all return real Qwen3-0.6B generations. Source precision ~321 tok/s; - JIT-4bit ~604 tok/s (see `../../spikes/mlx-solo/FINDINGS.md`). +- Verified on Apple Silicon (Metal), source precision, over the real frontend: + Qwen3-0.6B (non-stream + streaming) and SmolLM2 (Llama arch) both generate + **coherent** output. Source precision ~321 tok/s (see + `../../spikes/mlx-solo/FINDINGS.md`). + +## Dependency: why a git-rev pin (not crates.io) + +`Cargo.toml` pins `safemlx` / `safemlx-lm` to a specific **public** commit of +`github.com/jbg/safemlx` (`rev = "4e53c5e"`), not a crates.io version. This is a +deliberate, reproducible git pin — **not a private fork and no local patches**: + +- safemlx's **published** crates (both `0.1.5` and `0.4.1`) produce **garbage + output for dense models** (Qwen3 *and* Llama both emit repeated-token gibberish + with this exact same crate code). This was verified to be a **library bug**, + not a prompting/template problem — the Qwen3 chat template is confirmed applied + correctly, and greedy sampling (`temp=0` → argmax) is used. +- The pinned upstream commit serves those families correctly. Swapping the same + crate between published and the git pin flips the output coherent↔gibberish, + which isolates the cause to the safemlx version. +- **Action item:** swap the git-rev pin for a normal version pin once safemlx + cuts a crates.io release that serves dense models correctly. +- No JIT quantization is used (plain `LoadedModel::load`), so none of the + quant-path loader quirks apply; this path needs zero source patches. ## Promotion plan (the actual PR) -### 1. Make it a real workspace member with a pinned safemlx +### 1. Make it a real workspace member - Add `crates/skippy-engine-mlx` to root `Cargo.toml` `members` and remove its - local `[workspace]` table. -- Replace the `path = "../../../safemlx/..."` deps with **git-rev pins** of - `jbg/safemlx` (published crates are broken for Qwen3 — see FINDINGS §"supply - chain"). Carry the two small loader fixes as a patch/branch until upstreamed. + local `[workspace]` table. The safemlx deps are already git-rev pinned to a + public commit (see "Dependency" above), so nothing else changes about them — + a workspace member with a git dependency is fine. - Keep the crate's `mlx` feature; gate all MLX code with `#[cfg(all(feature = "mlx", target_os = "macos"))]` (already done). - Update `scripts/affected-crates.sh`, `scripts/plan-clippy-batches.sh`, @@ -43,14 +62,16 @@ In `crates/mesh-llm-host-runtime/Cargo.toml`: ```toml [target.'cfg(target_os = "macos")'.dependencies] -skippy-engine-mlx = { path = "../skippy-engine-mlx", optional = true } +skippy-engine-mlx = { path = "../skippy-engine-mlx", features = ["mlx"], optional = true } [features] -mlx = ["dep:skippy-engine-mlx", "skippy-engine-mlx/mlx"] +mlx = ["dep:skippy-engine-mlx"] ``` -Propagate a `mlx` feature up through `crates/mesh-llm/Cargo.toml`, and enable it -by default only on macOS builds in the release packaging. +(The `path` here is the in-repo crate path once it is a workspace member; its own +safemlx deps stay git-rev pinned.) Propagate a `mlx` feature up through +`crates/mesh-llm/Cargo.toml`, and enable it by default only on macOS builds in +the release packaging. ### 3. Add an `Mlx` variant to the launch enum @@ -123,6 +144,22 @@ routes to the MLX engine. tool parsing and thinking-output filtering worth porting later; this crate streams raw model text (including `` blocks) for now. - **Draft/speculative decoding** (goose's `gemma4_mtp`). +- **JIT quantization on load.** safemlx can affine-quantize dense weights at load + time (~604 tok/s 4-bit in earlier spikes), but it is deliberately excluded here + to keep the first PR to the goose-style plain-load path. + +## Future: Linux + NVIDIA (CUDA) + +MLX is **not Apple-only** — `jbg/safemlx` supports **Linux + NVIDIA GPUs via +CUDA** (a `cuda` cargo feature, plus `nccl` for multi-GPU), gated in +`safemlx-sys/build.rs` with a hard `panic!` to Linux targets, for both +`x86_64-linux` and `sbsa-linux` (ARM). The generation code +(`LoadedModel::load` + `generate_with_cache`) is backend-independent — only the +native build backend differs (Metal vs CUDA). So a later change could serve MLX +on Linux/NVIDIA mesh nodes by widening the gate from `target_os = "macos"` to +also allow `cfg(all(target_os = "linux", feature = "cuda"))` and enabling +`safemlx/cuda`. This is **out of scope for this PR** (Mac/Metal first) and is +tracked as future research; ROCm/Vulkan/Windows are not supported upstream. ## Testing the promoted path diff --git a/crates/skippy-engine-mlx/src/bin/mlx-serve.rs b/crates/skippy-engine-mlx/src/bin/mlx-serve.rs index d43a6811d7..7a95b69877 100644 --- a/crates/skippy-engine-mlx/src/bin/mlx-serve.rs +++ b/crates/skippy-engine-mlx/src/bin/mlx-serve.rs @@ -25,13 +25,6 @@ mod real { #[arg(long)] model_id: Option, - /// JIT-quantize eligible dense weights to this bit width on load (e.g. 4, 8). - #[arg(short, long)] - quantize: Option, - - #[arg(long, default_value_t = 64)] - quant_group_size: i32, - #[arg(long, default_value_t = 512)] default_max_tokens: usize, @@ -61,8 +54,6 @@ mod real { let config = MlxEngineConfig { model_dir: cli.model.clone(), model_id: model_id.clone(), - quantize_bits: cli.quantize, - quant_group_size: cli.quant_group_size, default_max_tokens: cli.default_max_tokens, max_tokens_cap: cli.max_tokens_cap, }; diff --git a/crates/skippy-engine-mlx/src/engine.rs b/crates/skippy-engine-mlx/src/engine.rs index 2a4d524428..989ac526c2 100644 --- a/crates/skippy-engine-mlx/src/engine.rs +++ b/crates/skippy-engine-mlx/src/engine.rs @@ -22,8 +22,7 @@ use tokio::sync::mpsc; use safemlx::transforms::async_eval; use safemlx::{Device, DeviceType, Stream}; use safemlx_lm::models::input::{InputPart, ModelInput}; -use safemlx_lm::models::{LoadedModel, ModelLoadOptions}; -use safemlx_lm::quantization::AffineQuantization; +use safemlx_lm::models::LoadedModel; use safemlx_lm::sampler::DefaultSampler; /// How the worker should load and run a model. @@ -31,9 +30,6 @@ use safemlx_lm::sampler::DefaultSampler; pub struct MlxEngineConfig { pub model_dir: PathBuf, pub model_id: String, - /// JIT-quantize eligible dense weights to this bit width on load (Metal only). - pub quantize_bits: Option, - pub quant_group_size: i32, pub default_max_tokens: usize, pub max_tokens_cap: usize, } @@ -141,18 +137,9 @@ fn load_engine(config: &MlxEngineConfig) -> Result { let stream = Stream::new_with_device(&Device::new(DeviceType::Gpu, 0)); let weights_stream = Stream::new_with_device(&Device::new(DeviceType::Cpu, 0)); - let options = match config.quantize_bits { - Some(bits) => ModelLoadOptions::with_quantization(AffineQuantization::new( - config.quant_group_size, - bits, - )?), - None => ModelLoadOptions::default(), - }; - let started = Instant::now(); - let model = - LoadedModel::load_with_options(&config.model_dir, options, &stream, &weights_stream) - .map_err(|e| anyhow!("load {}: {e}", config.model_dir.display()))?; + let model = LoadedModel::load(&config.model_dir, &stream, &weights_stream) + .map_err(|e| anyhow!("load {}: {e}", config.model_dir.display()))?; stream.synchronize().map_err(|e| anyhow!("sync: {e}"))?; let tokenizer = tokenizers::Tokenizer::from_file(config.model_dir.join("tokenizer.json")) From 7448a10637f324edf3ca58990d7112a634f65159 Mon Sep 17 00:00:00 2001 From: michaelneale Date: Fri, 17 Jul 2026 07:16:58 +1000 Subject: [PATCH 07/37] wip(mlx): wire MLX engine into mesh-llm serve (safetensors routing) 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. --- Cargo.lock | 452 ++++++++++++++++-- Cargo.toml | 1 + crates/mesh-llm-host-runtime/Cargo.toml | 21 + .../src/inference/mlx.rs | 103 ++++ .../src/inference/mod.rs | 2 + .../src/runtime/local.rs | 135 +++++- crates/mesh-llm/Cargo.toml | 5 + crates/skippy-engine-mlx/Cargo.toml | 6 +- .../SERVE_INTEGRATION_STATUS.md | 113 +++++ crates/skippy-engine-mlx/src/backend.rs | 8 +- scripts/affected-crates.sh | 1 + scripts/plan-clippy-batches.sh | 1 + 12 files changed, 805 insertions(+), 43 deletions(-) create mode 100644 crates/mesh-llm-host-runtime/src/inference/mlx.rs create mode 100644 crates/skippy-engine-mlx/SERVE_INTEGRATION_STATUS.md diff --git a/Cargo.lock b/Cargo.lock index e3d2442496..1b9afdd8dd 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -43,6 +43,20 @@ dependencies = [ "subtle", ] +[[package]] +name = "ahash" +version = "0.8.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5a15f179cd60c4584b8a8c596927aadc462e27f2ca70c04e0071964a73ba7a75" +dependencies = [ + "cfg-if 1.0.4", + "getrandom 0.3.4", + "once_cell", + "serde", + "version_check", + "zerocopy", +] + [[package]] name = "aho-corasick" version = "1.1.4" @@ -388,6 +402,28 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "async-stream" +version = "0.3.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b5a71a6f37880a80d1d7f19efd781e4b5de42c88f0722cc13bcb6cc2cfe8476" +dependencies = [ + "async-stream-impl", + "futures-core", + "pin-project-lite", +] + +[[package]] +name = "async-stream-impl" +version = "0.3.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c7c24de15d275a1ecfd47a380fb4d5ec9bfe0933f309ed5e705b775596a3574d" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.118", +] + [[package]] name = "async-task" version = "4.7.1" @@ -474,7 +510,7 @@ version = "0.30.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "16e2cdb6d5ed835199484bb92bb8b3edd526effe995c61732580439c1a67e2e9" dependencies = [ - "base64", + "base64 0.22.1", "http", "log", "url", @@ -577,6 +613,12 @@ version = "1.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "fd307490d624467aa6f74b0eabb77633d1f758a7b25f12bceb0b22e08d9726f6" +[[package]] +name = "base64" +version = "0.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9e1b586273c5702936fe7b7d6896644d8be71e6314cfe09d3167c95f712589e8" + [[package]] name = "base64" version = "0.22.1" @@ -1088,6 +1130,7 @@ dependencies = [ "itoa", "rustversion", "ryu", + "serde", "static_assertions", ] @@ -1244,6 +1287,16 @@ dependencies = [ "crossbeam-utils", ] +[[package]] +name = "crossbeam-deque" +version = "0.8.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5181e0de7b61eb03a81e347d6dd8797bae9da5146707b51077e2d71a54ec0ceb" +dependencies = [ + "crossbeam-epoch", + "crossbeam-utils", +] + [[package]] name = "crossbeam-epoch" version = "0.9.18" @@ -1487,6 +1540,12 @@ dependencies = [ "syn 2.0.118", ] +[[package]] +name = "daachorse" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6f55d7153ba3b507595872a3874803f07a8a81d1e888abed8e5db7da0597d6e2" + [[package]] name = "darling" version = "0.20.11" @@ -1556,6 +1615,15 @@ dependencies = [ "syn 2.0.118", ] +[[package]] +name = "dary_heap" +version = "0.3.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8b1e3a325bc115f096c8b77bbf027a7c2592230e70be2d985be950d3d5e60ebe" +dependencies = [ + "serde", +] + [[package]] name = "dashmap" version = "6.2.1" @@ -1967,6 +2035,15 @@ version = "3.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "dea2df4cf52843e0452895c455a1a2cfbb842a1e7329671acf418fdc53ed4c59" +[[package]] +name = "esaxx-rs" +version = "0.1.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d817e038c30374a4bcb22f94d0a8a0e216958d4c3dcde369b1439fec4bdda6e6" +dependencies = [ + "cc", +] + [[package]] name = "euclid" version = "0.22.14" @@ -2507,6 +2584,8 @@ dependencies = [ "allocator-api2", "equivalent", "foldhash 0.2.0", + "serde", + "serde_core", ] [[package]] @@ -2576,7 +2655,7 @@ name = "hf-hub" version = "1.0.0" source = "git+https://github.com/Mesh-LLM/hf-hub?branch=mesh-llm#fd3bfcabba1b9b827e685649cbcc8bf45ec6b310" dependencies = [ - "base64", + "base64 0.22.1", "bon", "bytes", "futures", @@ -2855,7 +2934,7 @@ version = "0.1.20" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "96547c2556ec9d12fb1578c4eaf448b04993e7fb79cbaad930a656880a6bdfa0" dependencies = [ - "base64", + "base64 0.22.1", "bytes", "futures-channel", "futures-util", @@ -3358,6 +3437,15 @@ dependencies = [ "either", ] +[[package]] +name = "itertools" +version = "0.15.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8b4baf93f58d4425749ca49a51c50ebab072c5df6994d08fed93541c331481dc" +dependencies = [ + "either", +] + [[package]] name = "itoa" version = "1.0.18" @@ -3703,6 +3791,22 @@ dependencies = [ "winapi", ] +[[package]] +name = "macro_rules_attribute" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "65049d7923698040cd0b1ddcced9b0eb14dd22c5f86ae59c3740eab64a676520" +dependencies = [ + "macro_rules_attribute-proc_macro", + "paste", +] + +[[package]] +name = "macro_rules_attribute-proc_macro" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "670fdfda89751bc4a84ac13eaa63e205cf0fd22b4c9a5fbfa085b63c1f1d3a30" + [[package]] name = "matchers" version = "0.2.0" @@ -3739,12 +3843,27 @@ version = "2.8.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "88904434abc2901f197fe8cc55f0445e7ded921dba5911dad2e2b39b48e663c4" +[[package]] +name = "memmap2" +version = "0.9.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d1219ed1b7f229ee7104d281dd01d6802fe28bb6e95d292942c4daacdeb798c0" +dependencies = [ + "libc", +] + [[package]] name = "memmem" version = "0.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a64a92489e2744ce060c349162be1c5f33c6969234104dbd99ddb5feb08b8c15" +[[package]] +name = "memo-map" +version = "0.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "38d1115007560874e373613744c6fba374c17688327a71c1476d1a5954cc857b" + [[package]] name = "memoffset" version = "0.9.1" @@ -3822,7 +3941,7 @@ version = "0.72.1" dependencies = [ "anyhow", "async-trait", - "base64", + "base64 0.22.1", "bytes", "crypto_box", "ed25519-dalek", @@ -3980,7 +4099,7 @@ dependencies = [ "argon2", "async-trait", "axum", - "base64", + "base64 0.22.1", "bytes", "chacha20poly1305", "chrono", @@ -4046,6 +4165,7 @@ dependencies = [ "serial_test", "sha2 0.10.9", "skippy-coordinator", + "skippy-engine-mlx", "skippy-protocol", "skippy-runtime", "skippy-server", @@ -4071,7 +4191,7 @@ name = "mesh-llm-identity" version = "0.72.1" dependencies = [ "argon2", - "base64", + "base64 0.22.1", "chacha20poly1305", "chrono", "crypto_box", @@ -4336,6 +4456,26 @@ dependencies = [ "unicase", ] +[[package]] +name = "minijinja" +version = "2.21.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cb3d648e68cea56d9858d535ee28f9538404e2dd8cb08ed0bd05dca379477f39" +dependencies = [ + "memo-map", + "serde", +] + +[[package]] +name = "minijinja-contrib" +version = "2.21.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85342f6fac0be8ccd5bd00d9066be538f34f393f577b75d81b17c8398a6b43bb" +dependencies = [ + "minijinja", + "serde", +] + [[package]] name = "minimal-lexical" version = "0.2.1" @@ -4449,6 +4589,28 @@ dependencies = [ "uuid", ] +[[package]] +name = "monostate" +version = "0.1.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3341a273f6c9d5bef1908f17b7267bbab0e95c9bf69a0d4dcf8e9e1b2c76ef67" +dependencies = [ + "monostate-impl", + "serde", + "serde_core", +] + +[[package]] +name = "monostate-impl" +version = "0.1.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e4db6d5580af57bf992f59068d4ea26fd518574ff48d7639b255a36f9de6e7e9" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.118", +] + [[package]] name = "more-asserts" version = "0.3.1" @@ -4845,7 +5007,7 @@ version = "0.44.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "98cf5d15d70d1f8f4059e5f79923ac15891eb691d2843d01191e0585fb064d70" dependencies = [ - "base64", + "base64 0.22.1", "bech32", "bip39", "bitcoin_hashes", @@ -5222,6 +5384,28 @@ version = "0.1.13" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "269bca4c2591a28585d6bf10d9ed0332b7d76900a1b02bec41bdc3a2cdcda107" +[[package]] +name = "onig" +version = "6.5.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0cc3cbf698f9438986c11a880c90a6d04b9de27575afd28bbf45b154b6c709e2" +dependencies = [ + "bitflags 2.13.0", + "libc", + "once_cell", + "onig_sys", +] + +[[package]] +name = "onig_sys" +version = "69.9.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e68317604e77e53b85896388e1a803c1d21b74c899ec9e5e1112db90735edd7" +dependencies = [ + "cc", + "pkg-config", +] + [[package]] name = "opaque-debug" version = "0.3.1" @@ -5348,7 +5532,7 @@ version = "0.31.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a7175df06de5eaee9909d4805a3d07e28bb752c34cab57fa9cff549da596b30f" dependencies = [ - "base64", + "base64 0.22.1", "const-hex", "opentelemetry", "opentelemetry_sdk", @@ -5516,7 +5700,7 @@ version = "3.0.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1d30c53c26bc5b31a98cd02d20f25a7c8567146caf63ed593a9d87b2775291be" dependencies = [ - "base64", + "base64 0.22.1", "serde_core", ] @@ -5716,7 +5900,7 @@ version = "1.9.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "092791278e026273c1b65bbdcfbba3a300f2994c896bd01ab01da613c29c46f1" dependencies = [ - "base64", + "base64 0.22.1", "indexmap", "quick-xml", "serde", @@ -5788,7 +5972,7 @@ version = "0.19.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "eb3713e4977408279158444a18c1a01ac9bf2e7eaf1fbfd1a19ac9cd18d90721" dependencies = [ - "base64", + "base64 0.22.1", "bytes", "derive_more", "hyper-util", @@ -5944,7 +6128,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "03da047801ff44bb6a4d407d4860c05fd70bb81714e6b2f3812603d5b145b042" dependencies = [ "heck", - "itertools", + "itertools 0.14.0", "log", "multimap", "petgraph", @@ -5963,7 +6147,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b570b25f7617e43d59005d0990ccb79e950a423952cea19671b7a876da390adf" dependencies = [ "anyhow", - "itertools", + "itertools 0.14.0", "proc-macro2", "quote", "syn 2.0.118", @@ -6270,7 +6454,7 @@ dependencies = [ "compact_str", "critical-section", "hashbrown 0.17.1", - "itertools", + "itertools 0.14.0", "kasuari", "lru 0.18.0", "palette", @@ -6335,7 +6519,7 @@ dependencies = [ "hashbrown 0.17.1", "indoc", "instability", - "itertools", + "itertools 0.14.0", "line-clipping", "ratatui-core", "serde", @@ -6345,6 +6529,37 @@ dependencies = [ "unicode-width", ] +[[package]] +name = "rayon" +version = "1.12.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fb39b166781f92d482534ef4b4b1b2568f42613b53e5b6c160e24cfbfa30926d" +dependencies = [ + "either", + "rayon-core", +] + +[[package]] +name = "rayon-cond" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2964d0cf57a3e7a06e8183d14a8b527195c706b7983549cd5462d5aa3747438f" +dependencies = [ + "either", + "itertools 0.14.0", + "rayon", +] + +[[package]] +name = "rayon-core" +version = "1.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "22e18b0f0062d30d4230b2e85ff77fdfe4326feb054b9783a3460d8435c8ab91" +dependencies = [ + "crossbeam-deque", + "crossbeam-utils", +] + [[package]] name = "rcgen" version = "0.14.8" @@ -6466,7 +6681,7 @@ version = "0.12.28" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "eddd3ca559203180a307f12d114c268abf583f59b03cb906fd0b3ff8646c1147" dependencies = [ - "base64", + "base64 0.22.1", "bytes", "encoding_rs", "futures-channel", @@ -6514,7 +6729,7 @@ version = "0.13.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "219c5811de6525e5416c7d5d53bb656d3afdbc6c5af816e0802bcfa42dbdc1c3" dependencies = [ - "base64", + "base64 0.22.1", "bytes", "encoding_rs", "futures-core", @@ -6594,7 +6809,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1d1f571c72940a19d9532fe52dbea8bc9912bf1d766c2970bb824056b86f3f59" dependencies = [ "async-trait", - "base64", + "base64 0.22.1", "bytes", "chrono", "futures", @@ -6857,6 +7072,107 @@ version = "0.11.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3944826ff8fa8093089aba3acb4ef44b9446a99a16f3bf4e74af3f77d340ab7d" +[[package]] +name = "safemlx" +version = "0.1.3" +source = "git+https://github.com/jbg/safemlx?rev=4e53c5e#4e53c5ecd7cbd91c0dfd0992a3c731ca2c36e9c7" +dependencies = [ + "bytemuck", + "dyn-clone", + "half", + "itertools 0.15.0", + "libc", + "num-complex", + "num-traits", + "num_enum", + "parking_lot", + "paste", + "safemlx-internal-macros", + "safemlx-macros", + "safemlx-sys", + "safetensors", + "smallvec", + "strum", + "thiserror 2.0.18", +] + +[[package]] +name = "safemlx-internal-macros" +version = "0.1.1" +source = "git+https://github.com/jbg/safemlx?rev=4e53c5e#4e53c5ecd7cbd91c0dfd0992a3c731ca2c36e9c7" +dependencies = [ + "darling 0.23.0", + "itertools 0.15.0", + "proc-macro2", + "quote", + "syn 2.0.118", +] + +[[package]] +name = "safemlx-lm" +version = "0.4.1" +source = "git+https://github.com/jbg/safemlx?rev=4e53c5e#4e53c5ecd7cbd91c0dfd0992a3c731ca2c36e9c7" +dependencies = [ + "anyhow", + "clap", + "idna_adapter", + "memmap2", + "minijinja", + "safemlx", + "safemlx-lm-utils", + "safetensors", + "serde", + "serde_json", + "thiserror 2.0.18", + "tokenizers", +] + +[[package]] +name = "safemlx-lm-utils" +version = "0.1.4" +source = "git+https://github.com/jbg/safemlx?rev=4e53c5e#4e53c5ecd7cbd91c0dfd0992a3c731ca2c36e9c7" +dependencies = [ + "minijinja", + "minijinja-contrib", + "serde", + "serde_json", + "thiserror 2.0.18", + "tokenizers", +] + +[[package]] +name = "safemlx-macros" +version = "0.1.1" +source = "git+https://github.com/jbg/safemlx?rev=4e53c5e#4e53c5ecd7cbd91c0dfd0992a3c731ca2c36e9c7" +dependencies = [ + "darling 0.23.0", + "proc-macro2", + "quote", + "syn 2.0.118", +] + +[[package]] +name = "safemlx-sys" +version = "0.1.3" +source = "git+https://github.com/jbg/safemlx?rev=4e53c5e#4e53c5ecd7cbd91c0dfd0992a3c731ca2c36e9c7" +dependencies = [ + "cc", + "cmake", +] + +[[package]] +name = "safetensors" +version = "0.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "79b079b829cb27a1c3c374341345ed2e8b2c0c839034522cee576c140bd7f846" +dependencies = [ + "hashbrown 0.16.1", + "libc", + "serde", + "serde_json", + "tempfile", +] + [[package]] name = "salsa20" version = "0.10.2" @@ -7422,6 +7738,26 @@ dependencies = [ "skippy-runtime", ] +[[package]] +name = "skippy-engine-mlx" +version = "0.72.1" +dependencies = [ + "anyhow", + "async-stream", + "async-trait", + "axum", + "clap", + "futures-core", + "openai-frontend", + "safemlx", + "safemlx-lm", + "serde_json", + "tokenizers", + "tokio", + "tracing", + "tracing-subscriber", +] + [[package]] name = "skippy-ffi" version = "0.72.1" @@ -7512,7 +7848,7 @@ dependencies = [ "anyhow", "async-trait", "axum", - "base64", + "base64 0.22.1", "blake3", "clap", "futures-util", @@ -7622,6 +7958,18 @@ dependencies = [ "der", ] +[[package]] +name = "spm_precompiled" +version = "0.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5851699c4033c63636f7ea4cf7b7c1f1bf06d0cc03cfb42e711de5a5c46cf326" +dependencies = [ + "base64 0.13.1", + "nom 7.1.3", + "serde", + "unicode-segmentation", +] + [[package]] name = "sse-stream" version = "0.2.3" @@ -7853,7 +8201,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4676b37242ccbd1aabf56edb093a4827dc49086c0ffd764a5705899e0f35f8f7" dependencies = [ "anyhow", - "base64", + "base64 0.22.1", "bitflags 2.13.0", "fancy-regex", "filedescriptor", @@ -8018,6 +8366,39 @@ version = "0.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1f3ccbac311fea05f86f61904b462b55fb3df8837a366dfc601a0161d0532f20" +[[package]] +name = "tokenizers" +version = "0.23.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "44e5bea67576e04b6ff8564c5d9e09c2ef0cf476502245f2f120e497769d3112" +dependencies = [ + "ahash", + "compact_str", + "daachorse", + "dary_heap", + "derive_builder", + "esaxx-rs", + "getrandom 0.3.4", + "itertools 0.14.0", + "log", + "macro_rules_attribute", + "monostate", + "onig", + "paste", + "rand 0.9.4", + "rayon", + "rayon-cond", + "regex", + "regex-syntax", + "serde", + "serde_json", + "spm_precompiled", + "thiserror 2.0.18", + "unicode-normalization-alignments", + "unicode-segmentation", + "unicode_categories", +] + [[package]] name = "tokio" version = "1.52.3" @@ -8084,7 +8465,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1af8573b15fdad8d66da116198cd8fd8d87ff62a67c1c6c3df7f62da1170793f" dependencies = [ "async-trait", - "base64", + "base64 0.22.1", "chrono", "futures", "log", @@ -8165,7 +8546,7 @@ version = "0.13.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d52efb639344a7c6adb8e62c6f3d2c19c001ff1b79a5041ba1c6ed42e19c6aa5" dependencies = [ - "base64", + "base64 0.22.1", "bytes", "futures-core", "futures-sink", @@ -8266,7 +8647,7 @@ checksum = "ac2a5518c70fa84342385732db33fb3f44bc4cc748936eb5833d2df34d6445ef" dependencies = [ "async-trait", "axum", - "base64", + "base64 0.22.1", "bytes", "h2", "http", @@ -8522,6 +8903,15 @@ dependencies = [ "tinyvec", ] +[[package]] +name = "unicode-normalization-alignments" +version = "0.1.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "43f613e4fa046e69818dd287fdc4bc78175ff20331479dab6e1b0f98d57062de" +dependencies = [ + "smallvec", +] + [[package]] name = "unicode-segmentation" version = "1.13.3" @@ -8534,7 +8924,7 @@ version = "2.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "16b380a1238663e5f8a691f9039c73e1cdae598a30e9855f541d29b08b53e9a5" dependencies = [ - "itertools", + "itertools 0.14.0", "unicode-segmentation", "unicode-width", ] @@ -8551,6 +8941,12 @@ version = "0.2.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ebc1c04c71510c7f702b52b7c350734c9ff1295c464a03335b00bb84fc54f853" +[[package]] +name = "unicode_categories" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "39ec24b3121d976906ece63c9daad25b85969647682eee313cb5779fdd69e14e" + [[package]] name = "uniffi" version = "0.31.0" @@ -9613,7 +10009,7 @@ checksum = "3e1e496dcbe6a09017acdfaf48e1a646735e7ff5b2a49e2c7e081cca77a59bc8" dependencies = [ "anyhow", "async-trait", - "base64", + "base64 0.22.1", "bytes", "clap", "crc32fast", @@ -9650,7 +10046,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "cb838aa8eb67d730af301584cf003caad407487606058292a6750711b603fbee" dependencies = [ "async-trait", - "base64", + "base64 0.22.1", "blake3", "bytemuck", "bytes", @@ -9661,7 +10057,7 @@ dependencies = [ "futures-util", "getrandom 0.4.3", "heapify", - "itertools", + "itertools 0.14.0", "lazy_static", "lz4_flex", "more-asserts", @@ -9693,7 +10089,7 @@ dependencies = [ "clap", "gearhash", "http", - "itertools", + "itertools 0.14.0", "lazy_static", "more-asserts", "rand 0.10.1", diff --git a/Cargo.toml b/Cargo.toml index 15c37618a4..d1ea11ddd3 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -41,6 +41,7 @@ members = [ "crates/skippy-topology", "crates/skippy-cache", "crates/skippy-metrics", + "crates/skippy-engine-mlx", "crates/openai-frontend", "crates/skippy-ffi", "crates/skippy-runtime", diff --git a/crates/mesh-llm-host-runtime/Cargo.toml b/crates/mesh-llm-host-runtime/Cargo.toml index 1b921c300e..07dc74063b 100644 --- a/crates/mesh-llm-host-runtime/Cargo.toml +++ b/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" } diff --git a/crates/mesh-llm-host-runtime/src/inference/mlx.rs b/crates/mesh-llm-host-runtime/src/inference/mlx.rs new file mode 100644 index 0000000000..2afaf8ee8d --- /dev/null +++ b/crates/mesh-llm-host-runtime/src/inference/mlx.rs @@ -0,0 +1,103 @@ +//! 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, MlxEngine, MlxEngineConfig}; + +/// A loaded MLX model plus the OpenAI backend that serves it. +pub(crate) struct MlxModelHandle { + backend: Arc, +} + +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 { + let config = MlxEngineConfig { + model_dir, + model_id, + default_max_tokens: context_length.max(1) as usize, + max_tokens_cap: context_length.max(1) as usize, + }; + let engine = MlxEngine::spawn(config)?; + Ok(Self { + backend: Arc::new(MlxBackend::new(engine)), + }) + } + + /// Starts an `openai-frontend` HTTP server for this model on `port`. + pub(crate) fn start_http(&self, port: u16) -> 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 server = tokio::spawn(async move { + let listener = match tokio::net::TcpListener::bind(addr).await { + Ok(listener) => listener, + Err(error) => { + tracing::error!(%addr, %error, "MLX openai frontend failed to bind"); + return; + } + }; + 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"); + } + }); + + 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>, + server: Option>, +} + +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(); + } + } +} diff --git a/crates/mesh-llm-host-runtime/src/inference/mod.rs b/crates/mesh-llm-host-runtime/src/inference/mod.rs index 38b67c300a..3e4d803885 100644 --- a/crates/mesh-llm-host-runtime/src/inference/mod.rs +++ b/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; diff --git a/crates/mesh-llm-host-runtime/src/runtime/local.rs b/crates/mesh-llm-host-runtime/src/runtime/local.rs index 4a4edc820a..4e2060a0a8 100644 --- a/crates/mesh-llm-host-runtime/src/runtime/local.rs +++ b/crates/mesh-llm-host-runtime/src/runtime/local.rs @@ -71,10 +71,22 @@ pub(super) enum RuntimeEvent { pub(super) enum LocalRuntimeBackendHandle { Skippy { - model: skippy::SkippyModelHandle, + // Boxed: `SkippyModelHandle` is large, and boxing keeps enum variants + // similarly sized (satisfies clippy::large_enum_variant once the smaller + // `Mlx` variant exists). Created once per model load, so the extra + // indirection is negligible. + model: Box, http: skippy::SkippyHttpHandle, _death_tx: tokio::sync::oneshot::Sender<()>, }, + #[cfg(all(feature = "mlx", target_os = "macos"))] + Mlx { + // Held to keep the loaded model (and its worker thread) alive for as long + // as the HTTP server is serving it; not read directly. + _model: crate::inference::mlx::MlxModelHandle, + http: crate::inference::mlx::MlxHttpHandle, + _death_tx: tokio::sync::oneshot::Sender<()>, + }, } pub(super) struct LocalRuntimeModelHandle { @@ -88,9 +100,7 @@ pub(super) struct LocalRuntimeModelHandle { impl LocalRuntimeModelHandle { pub(super) fn pid(&self) -> u32 { - match &self.inner { - LocalRuntimeBackendHandle::Skippy { .. } => std::process::id(), - } + std::process::id() } pub(super) fn ctx_used_tokens(&self) -> Option { @@ -98,12 +108,16 @@ impl LocalRuntimeModelHandle { LocalRuntimeBackendHandle::Skippy { model, .. } => { Some(model.status().max_session_tokens) } + #[cfg(all(feature = "mlx", target_os = "macos"))] + LocalRuntimeBackendHandle::Mlx { .. } => None, } } pub(super) fn openai_guardrails(&self) -> Option { match &self.inner { LocalRuntimeBackendHandle::Skippy { model, .. } => model.openai_guardrails(), + #[cfg(all(feature = "mlx", target_os = "macos"))] + LocalRuntimeBackendHandle::Mlx { .. } => None, } } @@ -115,6 +129,8 @@ impl LocalRuntimeModelHandle { LocalRuntimeBackendHandle::Skippy { model, .. } => { model.set_openai_guardrail_mode(mode) } + #[cfg(all(feature = "mlx", target_os = "macos"))] + LocalRuntimeBackendHandle::Mlx { .. } => None, } } @@ -157,6 +173,8 @@ impl LocalRuntimeModelHandle { .collect(), }) } + #[cfg(all(feature = "mlx", target_os = "macos"))] + LocalRuntimeBackendHandle::Mlx { .. } => None, } } @@ -166,6 +184,10 @@ impl LocalRuntimeModelHandle { let _ = http.shutdown().await; model.shutdown(); } + #[cfg(all(feature = "mlx", target_os = "macos"))] + LocalRuntimeBackendHandle::Mlx { http, .. } => { + let _ = http.shutdown().await; + } } } } @@ -566,6 +588,16 @@ pub(super) async fn start_runtime_local_model( tokio::sync::oneshot::Receiver<()>, )> { let model_name = runtime_model_name.to_string(); + + // Safetensors models are served by the MLX (Metal) engine on Apple Silicon. + // This branch runs before any GGUF planning, which assumes a llama.cpp/GGUF + // artifact. On non-macOS or without the `mlx` feature, a safetensors path + // falls through to the GGUF path and fails there with a clear error. + #[cfg(all(feature = "mlx", target_os = "macos"))] + if is_safetensors_model_path(spec.model_path) { + return start_runtime_mlx_model(spec, model_name).await; + } + let package_ref = spec.model_path.to_string_lossy().to_string(); let layer_package = if skippy::is_layer_package_ref(&package_ref) { let package_ref_for_identity = package_ref.clone(); @@ -1384,7 +1416,7 @@ async fn load_split_runtime_generation_inner( slots: spec.slots, capabilities, inner: LocalRuntimeBackendHandle::Skippy { - model: handle, + model: Box::new(handle), http, _death_tx: death_tx, }, @@ -3610,6 +3642,95 @@ fn now_unix_nanos() -> i64 { .unwrap_or(0) } +/// True if the resolved model path points at a safetensors artifact (single +/// `model.safetensors`, a sharded first file, or any `*.safetensors`), which the +/// MLX engine serves. Directories are treated as safetensors when they contain a +/// primary weight file. +#[cfg(all(feature = "mlx", target_os = "macos"))] +fn is_safetensors_model_path(path: &Path) -> bool { + let name = path + .file_name() + .map(|n| n.to_string_lossy().to_string()) + .unwrap_or_default(); + if name.ends_with(".safetensors") { + return true; + } + if path.is_dir() { + return path.join("model.safetensors").exists() + || path.join("model.safetensors.index.json").exists(); + } + false +} + +/// The directory that holds a safetensors model's `config.json` / `tokenizer.json` +/// / weights — either `path` itself (a dir) or its parent (a weight file). +#[cfg(all(feature = "mlx", target_os = "macos"))] +fn mlx_model_dir(path: &Path) -> PathBuf { + if path.is_dir() { + path.to_path_buf() + } else { + path.parent() + .map(Path::to_path_buf) + .unwrap_or_else(|| path.to_path_buf()) + } +} + +#[cfg(all(feature = "mlx", target_os = "macos"))] +async fn start_runtime_mlx_model( + spec: LocalRuntimeModelStartSpec<'_>, + model_name: String, +) -> Result<( + String, + LocalRuntimeModelHandle, + tokio::sync::oneshot::Receiver<()>, +)> { + let port = alloc_local_port().await?; + let model_dir = mlx_model_dir(spec.model_path); + let context_length = spec.ctx_size_override.unwrap_or(4096); + let capabilities = models::runtime_verified_model_capabilities( + &model_name, + spec.model_path, + models::RuntimeMediaCapabilityEvidence { + vision_projector_loaded: false, + }, + ); + + let _ = emit_event(OutputEvent::ModelLoading { + model: model_name.clone(), + source: None, + }); + let load_model_name = model_name.clone(); + let mlx_model = tokio::task::spawn_blocking(move || { + crate::inference::mlx::MlxModelHandle::load(model_dir, load_model_name, context_length) + }) + .await + .context("join load MLX model task")??; + let _ = emit_event(OutputEvent::ModelLoaded { + model: model_name.clone(), + bytes: None, + }); + + let http = mlx_model.start_http(port); + let (death_tx, death_rx) = tokio::sync::oneshot::channel(); + + Ok(( + model_name, + LocalRuntimeModelHandle { + port: http.port(), + backend: "mlx".into(), + context_length, + slots: 1, + capabilities, + inner: LocalRuntimeBackendHandle::Mlx { + _model: mlx_model, + http, + _death_tx: death_tx, + }, + }, + death_rx, + )) +} + async fn start_runtime_skippy_model( spec: LocalRuntimeModelStartSpec<'_>, model_name: String, @@ -3687,7 +3808,7 @@ async fn start_runtime_skippy_model( slots: plan.slots, capabilities, inner: LocalRuntimeBackendHandle::Skippy { - model: skippy_model, + model: Box::new(skippy_model), http, _death_tx: death_tx, }, @@ -3803,7 +3924,7 @@ async fn start_runtime_layer_package_model( slots: plan.slots, capabilities, inner: LocalRuntimeBackendHandle::Skippy { - model: handle, + model: Box::new(handle), http, _death_tx: death_tx, }, diff --git a/crates/mesh-llm/Cargo.toml b/crates/mesh-llm/Cargo.toml index 850c10a5cf..eb950295fd 100644 --- a/crates/mesh-llm/Cargo.toml +++ b/crates/mesh-llm/Cargo.toml @@ -12,6 +12,11 @@ web-ui = ["mesh-llm-host-runtime/web-ui"] gpu-bench-cuda = ["mesh-llm-host-runtime/gpu-bench-cuda"] gpu-bench-hip = ["mesh-llm-host-runtime/gpu-bench-hip"] gpu-bench-intel = ["mesh-llm-host-runtime/gpu-bench-intel"] +# Serve HF safetensors models on Apple Silicon via the MLX (Metal) engine. +# Off by default; enable on macOS builds. Implies `dynamic-native-runtime` +# (see mesh-llm-host-runtime `mlx` for why: avoids a GGUF duplicate-symbol +# link failure between MLX's vendored gguflib and the patched llama.cpp). +mlx = ["mesh-llm-host-runtime/mlx", "dynamic-native-runtime"] dynamic-native-runtime = ["mesh-llm-host-runtime/dynamic-native-runtime"] [lints] diff --git a/crates/skippy-engine-mlx/Cargo.toml b/crates/skippy-engine-mlx/Cargo.toml index d4014cca37..d9509339ac 100644 --- a/crates/skippy-engine-mlx/Cargo.toml +++ b/crates/skippy-engine-mlx/Cargo.toml @@ -16,12 +16,10 @@ # Promotion to a real workspace member is documented in WIRING.md. [package] name = "skippy-engine-mlx" -version = "0.0.0" -edition = "2021" +version.workspace = true +edition.workspace = true publish = false -[workspace] - [lib] name = "skippy_engine_mlx" path = "src/lib.rs" diff --git a/crates/skippy-engine-mlx/SERVE_INTEGRATION_STATUS.md b/crates/skippy-engine-mlx/SERVE_INTEGRATION_STATUS.md new file mode 100644 index 0000000000..26f0db0fdc --- /dev/null +++ b/crates/skippy-engine-mlx/SERVE_INTEGRATION_STATUS.md @@ -0,0 +1,113 @@ +# MLX serve-integration — WIP status & blocker + +This branch (`micn/mlx-serve-wiring`) stacks the **`mesh-llm serve` integration** +on top of the standalone MLX engine crate from PR #1009 (`micn/mlx-redux`). + +Goal: on a Mac, `mesh-llm serve --model ` routes to the MLX +(Metal) engine and the model appears in `/v1/models` — with non-macOS / no-feature +builds byte-for-byte unaffected. + +**Status: functionally complete and verified locally, but blocked by one repo +publish invariant. Needs a maintainer decision before it can merge.** Captured +here so the work + reasoning aren't lost. + +## What works (verified on Apple Silicon this session) + +- `skippy-engine-mlx` is now a real workspace member (removed its private + `[workspace]`; added to root `members`, both CI crate-list scripts, and it + stays out of `default-members`). +- `mlx` feature added to `mesh-llm-host-runtime` and `mesh-llm`, wired as a + **macOS-target-gated optional dep** (`[target.'cfg(target_os = "macos")'.dependencies]`). +- `crates/mesh-llm-host-runtime/src/inference/mlx.rs`: `MlxModelHandle` + + `MlxHttpHandle` serving over the real `openai-frontend::router_for` + axum, with + graceful shutdown. +- `LocalRuntimeBackendHandle::Mlx` variant + all match arms (cfg-gated). +- `start_runtime_mlx_model` + `is_safetensors_model_path` routing: safetensors + models branch to MLX *before* the GGUF planning path. +- **Build gate PASSED:** `cargo build -p mesh-llm --features mlx` → exit 0. +- **Don't-break-default PASSED:** no-feature `cargo check`/clippy clean; + `cargo tree` confirms safemlx is absent unless `--features mlx`. +- fmt clean; clippy clean both with and without `--features mlx` + (boxed the `Skippy` enum variant to satisfy `large_enum_variant`). +- `xtask repo-consistency ci-crate-lists` PASSES. + +## Key finding: the two-native-stack link collision (RESOLVED) + +Linking MLX statically alongside the patched llama.cpp fails with: + +``` +ld64.lld: error: duplicate symbol: gguf_get_key + >>> defined in .../gguflib-src/gguflib.c (MLX's vendored GGUF parser) + >>> defined in libskippy_ffi...(gguf.cpp.o) (patched llama.cpp) +``` + +MLX statically links antirez's `gguflib` (via `MLX_BUILD_GGUF=ON`, and +`safemlx-sys` link-directs `gguflib` unconditionally); the patched llama.cpp +exports the same C symbols. Two independent GGUF parsers → duplicate symbol. + +**Fix:** the `mlx` feature now implies `dynamic-native-runtime`. That loads the +llama.cpp runtime as a dylib (as release builds already do), so its GGUF symbols +live in a separate link unit and don't collide with MLX's static `gguflib`. +`cargo build -p mesh-llm --features mlx` alone links clean after this. + +## BLOCKER: publish invariant (needs a maintainer call) + +``` +cargo run -p xtask -- repo-consistency release-targets + error: mesh-llm-host-runtime: publishable crate depends on + non-publishable workspace crate `skippy-engine-mlx` +``` + +Root-cause chain: +- `skippy-engine-mlx` git-pins `safemlx` / `safemlx-lm` to a public commit of + `github.com/jbg/safemlx` (crates.io's published safemlx produces gibberish for + dense models — verified; only the git commit serves correctly). +- crates.io forbids git dependencies → the crate must be `publish = false`. +- `mesh-llm-host-runtime` is a **published** SDK crate, and the repo invariant + (enforced by `xtask release-targets`, reflecting a real `cargo publish` + constraint) forbids a publishable crate from depending on a `publish = false` + one — even an optional, target-gated dep. + +This is **not solved elsewhere in the repo**: the other `publish = false` crates +(`skippy-quantize`, `mesh-llm-commands`) are only consumed by the non-published +binary crate `mesh-llm`, never by a published library crate. + +### Options + +- **A. `[patch.crates-io]` redirect + make the crate publishable.** Follows the + existing `hf-hub` precedent (a published-crate dep already redirected to a fork + via `[patch.crates-io]` in the root manifest). Add `skippy-engine-mlx` to + `scripts/publish-crates.sh`. Caveat that differs from hf-hub: safemlx's + *published* version is known-broken, so a hypothetically-published + `skippy-engine-mlx` with `mlx` on would reference broken upstream — acceptable + only because the feature is off by default and explicitly a stopgap until + safemlx cuts a working release. +- **B. Hold the host-runtime wiring.** Ship the standalone crate (PR #1009) as + is; keep this branch as the ready-to-go integration until safemlx releases a + working crates.io version, then flip git-pin → version-pin and merge. Most + honest; doesn't deliver "serve just works" yet. +- **C. Loosen the xtask invariant** to exempt optional/target deps. Not + recommended: it would allow a manifest that genuinely cannot `cargo publish`. + +**Recommendation: B for now** (this branch is the parked, working integration), +moving to **A** if/when we want it live before safemlx releases. The real +unlock is a working safemlx crates.io release, after which this is a trivial +git-pin → version-pin swap and the invariant is satisfied automatically. + +## How to reproduce / verify + +```bash +export DEVELOPER_DIR=/Applications/Xcode.app/Contents/Developer # Metal toolchain +cargo build -p mesh-llm --features mlx # exit 0 (links clean) +cargo check -p mesh-llm-host-runtime # no-feature: clean, no safemlx +cargo run -p xtask -- repo-consistency ci-crate-lists # PASS +cargo run -p xtask -- repo-consistency release-targets # FAILS (the blocker) +``` + +## Remaining once unblocked + +- End-to-end serve test: `mesh-llm serve --model ` → + confirm `/v1/models` + `/v1/chat/completions`. +- Fold this status into `WIRING.md` (note the `dynamic-native-runtime` + requirement and the publish resolution chosen). +- Rebase onto `main` and open/land the real PR. diff --git a/crates/skippy-engine-mlx/src/backend.rs b/crates/skippy-engine-mlx/src/backend.rs index 02de047fd8..3a14867d7b 100644 --- a/crates/skippy-engine-mlx/src/backend.rs +++ b/crates/skippy-engine-mlx/src/backend.rs @@ -167,9 +167,9 @@ impl OpenAiBackend for MlxBackend { &self, request: ChatCompletionRequest, ) -> OpenAiResult { - let gen = self.build_request(&request)?; + let job = self.build_request(&request)?; let model = self.engine.model_id().to_string(); - let mut rx = self.engine.submit(gen); + let mut rx = self.engine.submit(job); let mut text = String::new(); let mut finish = FinishReason::Stop; @@ -217,9 +217,9 @@ impl OpenAiBackend for MlxBackend { request: ChatCompletionRequest, _context: OpenAiRequestContext, ) -> OpenAiResult { - let gen = self.build_request(&request)?; + let job = self.build_request(&request)?; let model = self.engine.model_id().to_string(); - let mut rx = self.engine.submit(gen); + let mut rx = self.engine.submit(job); let id = completion_id("chatcmpl"); let created = now_secs(); diff --git a/scripts/affected-crates.sh b/scripts/affected-crates.sh index e9c25d9cf0..6f7b466ecf 100755 --- a/scripts/affected-crates.sh +++ b/scripts/affected-crates.sh @@ -50,6 +50,7 @@ WORKSPACE_MEMBERS=( "skippy-topology" "skippy-cache" "skippy-metrics" + "skippy-engine-mlx" "openai-frontend" "skippy-ffi" "skippy-runtime" diff --git a/scripts/plan-clippy-batches.sh b/scripts/plan-clippy-batches.sh index b5155c1084..abff132f92 100644 --- a/scripts/plan-clippy-batches.sh +++ b/scripts/plan-clippy-batches.sh @@ -52,6 +52,7 @@ WORKSPACE_MEMBERS=( "skippy-topology" "skippy-cache" "skippy-metrics" + "skippy-engine-mlx" "openai-frontend" "skippy-ffi" "skippy-runtime" From 99a86ec2fecb2b0744bfd6177320fddbe64d1d9e Mon Sep 17 00:00:00 2001 From: Michael Neale <14976+michaelneale@users.noreply.github.com> Date: Fri, 17 Jul 2026 15:44:06 +1000 Subject: [PATCH 08/37] feat(mlx): prove partial safetensors stage execution --- Justfile | 12 + docs/design/MLX_STAGE_ENGINE_PLAN.md | 101 +- spikes/mlx-safetensors-stages/Cargo.lock | 1402 ++++++++++++++++++++ spikes/mlx-safetensors-stages/Cargo.toml | 21 + spikes/mlx-safetensors-stages/FINDINGS.md | 321 +++++ spikes/mlx-safetensors-stages/src/main.rs | 883 ++++++++++++ spikes/mlx-solo/Cargo.lock | 182 ++- spikes/mlx-solo/Cargo.toml | 5 + spikes/mlx-solo/FINDINGS.md | 7 +- spikes/mlx-solo/src/bin/mlx-split-proof.rs | 465 +++++++ 10 files changed, 3364 insertions(+), 35 deletions(-) create mode 100644 spikes/mlx-safetensors-stages/Cargo.lock create mode 100644 spikes/mlx-safetensors-stages/Cargo.toml create mode 100644 spikes/mlx-safetensors-stages/FINDINGS.md create mode 100644 spikes/mlx-safetensors-stages/src/main.rs create mode 100644 spikes/mlx-solo/src/bin/mlx-split-proof.rs diff --git a/Justfile b/Justfile index a4f2dcc80d..b415e5d30b 100644 --- a/Justfile +++ b/Justfile @@ -239,6 +239,18 @@ 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 }} + # Generate a reproducible benchmark corpus for skippy bench tooling. bench-corpus tier="smoke" *ARGS="": scripts/generate-bench-corpus.py "{{ tier }}" {{ ARGS }} diff --git a/docs/design/MLX_STAGE_ENGINE_PLAN.md b/docs/design/MLX_STAGE_ENGINE_PLAN.md index 31710c9db8..a6f2495842 100644 --- a/docs/design/MLX_STAGE_ENGINE_PLAN.md +++ b/docs/design/MLX_STAGE_ENGINE_PLAN.md @@ -22,6 +22,21 @@ beyond it (JIT quant + loading arbitrary mlx-community repos) and are upstream-P candidates. See `spikes/mlx-solo/FINDINGS.md`; results are folded into §5.3, Phase 2, and §9. +**Update — exact stage-local SafeTensors materialization and dense split +execution are now proven** +(`spikes/mlx-safetensors-stages/`). The SafeTensors index plus per-file headers +are sufficient to map a stage to exact HTTP byte ranges, and Hugging Face honors +those range requests. This materially changes the split artifact conclusion: +nodes do not need published stage packages or even complete source shards. On +Inkling BF16, four layers contain 109.84 GiB of tensors scattered across 942.99 +GiB of shard files; exact ranges avoid 833.15 GiB. A SmolLM2-135M proof then +materialized two partial files (layers 0..15 and 15..30), loaded each directly +into MLX, and matched unsplit logits exactly for prefill plus eight decode steps +through Skippy's real F16 and F32 binary activation codec. The remaining +artifact gate is bounded-memory quantization for frontier-sized source tensors; +the remaining product gate is engine-neutral `skippy-server` integration. +See `spikes/mlx-safetensors-stages/FINDINGS.md`. + --- ## 1. Bottom line @@ -52,10 +67,16 @@ batched verify, trim/checkpoint, tokenizer/chat helpers). **Recommendation:** introduce a Rust `StageEngine` trait, keep the existing C ABI as the `LlamaStageEngine` adapter, and add an Apple-Silicon-gated `MlxStageEngine`. Do **not** extend the llama.cpp C ABI to host MLX, and do -**not** invent a separate MLX network protocol. **Lead with solo MLX + -JIT-quant serving** (immediate workflow win, minimal distributed work), then gate -the split work behind two go/no-go spikes: **partial layer loading** and -**per-token boundary fence latency**. +**not** invent a separate MLX network protocol. For split serving, treat the +immutable upstream SafeTensors checkpoint plus a small quantization profile as +the source of truth; range-fetch, optionally quantize, and cache only the local +stage. Published engine-specific layer packages become an optional prewarmed +optimization rather than a prerequisite. Gate execution behind the remaining +go/no-go work: **bounded-memory stage materialization / partial model loading** +and **per-token boundary fence latency**. The small-model proof also establishes +that a receiver must restore the model compute dtype after decoding the wire +dtype; numeric F32 residual values left as an F32 MLX array change downstream +Metal arithmetic for a BF16 model. --- @@ -377,10 +398,11 @@ JIT quant: is a *whole-model* op. A split node must instantiate only `layers[start..end]`, read only the shards overlapping its range, and quantize only those tensors. This is exactly go/no-go **Spike 1**, now with a quant step folded in. -2. **Selective shard download.** The `weight_map` makes fetching only the shards - for a layer range *possible*, but safetensors shards bundle several - consecutive layers, so nodes over-fetch at range boundaries — coarser than - GGUF layer-package parts, which Skippy slices exactly. +2. **Exact tensor-range download (proven).** The `weight_map` selects source + files, and each SafeTensors header supplies exact tensor byte offsets. HTTP + range requests therefore avoid whole-shard overfetch. This is a modest win + for layer-ordered Qwen/Nemotron/GLM checkpoints and a requirement for + Inkling, whose tensors for four layers are spread across 57 BF16 files. 3. **Deterministic cross-stage quant.** Every stage must quantize *identically* (same algo / group-size / bits / tie handling) or the split model drifts numerically from the solo model. Affine quant is deterministic given its @@ -396,14 +418,15 @@ JIT quant: #### The tension worth naming Skippy's existing chain (`skippy-quantize`, layer-package repos, BF16→GGUF) is -built around **pre-quantized, exactly-sliced GGUF parts** *specifically so split -nodes download their slice and never quantize at runtime*. -JIT-quant-from-safetensors trades that for flexible source + runtime quant + -coarser slicing. For solo there is no tension; for splits it is a genuine, but -acceptable, tradeoff. The two paths should **coexist**: - -- **JIT safetensors = fast coverage path** — try any supported HF model on the - mesh immediately, no publish step. +built around **pre-quantized, exactly-sliced GGUF parts** so split nodes never +quantize at runtime. Exact SafeTensors byte ranges remove the earlier coarser- +slicing disadvantage. The remaining trade is cold-start quantization time and +temporary source precision versus a prewarmed, published quant. The two paths +should **coexist**: + +- **JIT safetensors = flexible coverage path** — range-fetch only the stage, + adapt its precision to available hardware, cache the deterministic result, + and require no weight-republishing step. - **Pre-quantized layer packages = optimized path** — for models served seriously (exact slices, no runtime quant, tailored partial download). @@ -415,7 +438,7 @@ Use **one logical package identity, not one physical weight encoding**: model identity + source revision + tokenizer/config/chat metadata + topology variants: llama-gguf: GGUF parts + quant (existing skippy-model-package path) - mlx-jit: HF safetensors + quant spec (quantize on load; cache the result) + mlx-jit: HF tensor ranges + per-stage quant profile (quantize on load; cache) mlx-packaged: pre-quantized MLX stage shards + index (optimized split path) ``` @@ -608,12 +631,21 @@ new distributed work, and de-risks the engine before any split work. > candidates, not mesh-llm drift (see §9). CPU is not a serving path and was not > benchmarked as one. -**Phase 3 — Stage-aware partial load + activation frames.** Add `forward_range` -/ `resume_from_hidden` and the stage-aware loader to `safemlx-lm` (upstream to -the fork). Implement `prefill_chunk_frame` / `decode_step_frame` / +**Phase 3 — Streaming stage materialization + partial load + activation +frames.** Convert the proven tensor-range plan into a bounded-memory pipeline: +range-fetch one tensor, optionally quantize it, append it to a derived stage +cache, and release the source buffer. Add `forward_range` / +`resume_from_hidden` and a stage-aware model constructor to `safemlx-lm`. +Implement `prefill_chunk_frame` / `decode_step_frame` / `copy_output_activation_frame` producing Skippy `ActivationFrame`s. Two-stage single-machine parity first, then two Macs over the real network. +> **Dense single-machine spike passed.** SmolLM2-135M was split 15+15 using two +> exact-range partial SafeTensors artifacts. F16 and F32 `StageWireMessage` +> boundaries both matched unsplit MLX with zero measured logit delta across +> prompt prefill and eight decode steps. This proves the basic artifact and +> activation seams, but not bounded-memory quantization or server integration. + **Phase 4 — KV/state codec + verify + trim/checkpoint.** Implement the engine-general cache codec (§5.2), `verify_tokens_frame` for speculative decode, trim/checkpoint/reset. Add speculative (safemlx-lm already has Gemma4 MTP draft @@ -632,17 +664,20 @@ Apple-Silicon nodes. llama.cpp remains the cross-platform default. ## 8. Spike gates (go/no-go before Phase 3) -1. **Partial-loading proof (GO/NO-GO):** load only `layers[start..end]` (+ - embeddings/readout when owned) for a dense Qwen/Llama; confirm **peak RSS - contains only the selected range**. If a stage can't avoid loading the whole - model, the split story is dead for MLX. +1. **Partial-loading proof (DENSE GO, QUANT PARTIAL):** remote exact-range + selection is proven, including on 1.9 TB Inkling BF16. SmolLM2 partial files + were materialized and loaded without a complete checkpoint. Still required: + confirm peak RSS is bounded by quantized stage + one source tensor and + scratch during tensor-at-a-time load-time quantization. 2. **Boundary latency breakdown (GO/NO-GO):** measure layer compute, cast, contiguous, **eval fence**, host readback, serialize, and receive-reconstruct **independently**, at hidden widths 4096/8192/16384 and token counts 1/32/512. Decode is single-sequence latency-bound; prove the fence doesn't dominate. -3. **Two-stage dense parity:** Qwen/Llama, multiple split points, F32 + F16 - activations, chunked prefill, 128-token decode, compare logits to llama.cpp. +3. **Two-stage dense parity (INITIAL GO):** SmolLM2/Llama at split 15 passed F32 + and F16 through the real binary codec with zero measured logit delta for one + prefill plus eight decode steps. Still required: multiple split points, + chunked prefill, 128-token decode, two processes, and cross-engine comparison. 4. **Real network run:** two Macs over Wi-Fi and 1/10GbE (Thunderbolt if relevant); report end-to-end tok/s + p50/p95 inter-token latency, not local MLX throughput. @@ -705,12 +740,14 @@ Spikes 1 and 2 are more decisive than any standalone token/s benchmark. ## 10. Immediate next steps -1. Land **Phase 1** (the `StageEngine` trait refactor, llama-only) — valuable on - its own and the prerequisite for everything else. -2. Run **Spike 1 (partial load)** and **Spike 2 (boundary fence)** in - `../safemlx` against a dense Qwen/Llama. Treat both as go/no-go. -3. If both pass, proceed to Phase 2 reusing goose's `safemlx-lm` generation - patterns as the starting point. +1. Add tensor-at-a-time MLX quantization to the materializer and measure peak + RSS against the one-source-tensor memory contract. +2. Introduce `StageEngine` with the llama adapter first, then run the proven MLX + split through two real `skippy-server` processes. +3. Run **Spike 2 (boundary fence)** at frontier residual widths and keep it as a + go/no-go gate. +4. Use Nemotron-H as the first frontier-family follow-up already represented in + `safemlx-lm`; then port Inkling text from the upstream Transformers reference. --- diff --git a/spikes/mlx-safetensors-stages/Cargo.lock b/spikes/mlx-safetensors-stages/Cargo.lock new file mode 100644 index 0000000000..dc22819051 --- /dev/null +++ b/spikes/mlx-safetensors-stages/Cargo.lock @@ -0,0 +1,1402 @@ +# This file is automatically @generated by Cargo. +# It is not intended for manual editing. +version = 4 + +[[package]] +name = "anstream" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "824a212faf96e9acacdbd09febd34438f8f711fb84e09a8916013cd7815ca28d" +dependencies = [ + "anstyle", + "anstyle-parse", + "anstyle-query", + "anstyle-wincon", + "colorchoice", + "is_terminal_polyfill", + "utf8parse", +] + +[[package]] +name = "anstyle" +version = "1.0.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "940b3a0ca603d1eade50a4846a2afffd5ef57a9feac2c0e2ec2e14f9ead76000" + +[[package]] +name = "anstyle-parse" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "52ce7f38b242319f7cabaa6813055467063ecdc9d355bbb4ce0c68908cd8130e" +dependencies = [ + "utf8parse", +] + +[[package]] +name = "anstyle-query" +version = "1.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "40c48f72fd53cd289104fc64099abca73db4166ad86ea0b4341abe65af83dadc" +dependencies = [ + "windows-sys 0.61.2", +] + +[[package]] +name = "anstyle-wincon" +version = "3.0.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "291e6a250ff86cd4a820112fb8898808a366d8f9f58ce16d1f538353ad55747d" +dependencies = [ + "anstyle", + "once_cell_polyfill", + "windows-sys 0.61.2", +] + +[[package]] +name = "anyhow" +version = "1.0.103" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2a4385e2e34eb35d6b3efe798b9eb88096925d87726c0798709bf56d9ed84af3" + +[[package]] +name = "atomic-waker" +version = "1.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1505bd5d3d116872e7271a6d4e16d81d0c8570876c8de68093a09ac269d8aac0" + +[[package]] +name = "base64" +version = "0.22.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "72b3254f16251a8381aa12e40e3c4d2f0199f8c6508fbecb9d91f575e0fbb8c6" + +[[package]] +name = "bitflags" +version = "2.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b588b76d00fde79687d7646a9b5bdf3cc0f655e0bbd080335a95d7e96f3587da" + +[[package]] +name = "bumpalo" +version = "3.20.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "72f5acc6cb2ba439de613abc23857ec3d78374d8ed5ac84e9d11336e87da8649" + +[[package]] +name = "bytes" +version = "1.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc652a48c352aef3ea3aed32080501cf3ef6ed5da78602a020c991775b0aff04" + +[[package]] +name = "cc" +version = "1.2.67" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e17dd265a7d0f31ef544e1b20e03add05d3b45b491b633b10d67145d2acc1a38" +dependencies = [ + "find-msvc-tools", + "shlex", +] + +[[package]] +name = "cfg-if" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" + +[[package]] +name = "cfg_aliases" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f079e83a288787bcd14a6aea84cee5c87a67c5a3e660c30f557a3d24761b3527" + +[[package]] +name = "chacha20" +version = "0.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d524456ba66e72eb8b115ff89e01e497f8e6d11d78b70b1aa13c0fbd97540a81" +dependencies = [ + "cfg-if", + "cpufeatures", + "rand_core", +] + +[[package]] +name = "clap" +version = "4.6.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dd059f9da4f5c36b3787f65d38ccaab1cc315f07b01f89abc8359ee6a8205011" +dependencies = [ + "clap_builder", + "clap_derive", +] + +[[package]] +name = "clap_builder" +version = "4.6.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f09628afdcc538b57f3c6341e9c8e9970f18e4a481690a64974d7023bd33548b" +dependencies = [ + "anstream", + "anstyle", + "clap_lex", + "strsim", +] + +[[package]] +name = "clap_derive" +version = "4.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f2ce8604710f6733aa641a2b3731eaa1e8b3d9973d5e3565da11800813f997a9" +dependencies = [ + "heck", + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "clap_lex" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8d4a3bb8b1e0c1050499d1815f5ab16d04f0959b233085fb31653fbfc9d98f9" + +[[package]] +name = "colorchoice" +version = "1.0.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1d07550c9036bf2ae0c684c4297d503f838287c83c53686d05370d0e139ae570" + +[[package]] +name = "cpufeatures" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8b2a41393f66f16b0823bb79094d54ac5fbd34ab292ddafb9a0456ac9f87d201" +dependencies = [ + "libc", +] + +[[package]] +name = "displaydoc" +version = "0.2.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ac70aa55017e108007fbaf5aa0f54b021c98f92ff8af59d42eda9da96e3dd4f" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "find-msvc-tools" +version = "0.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5baebc0774151f905a1a2cc41989300b1e6fbb29aff0ceffa1064fdd3088d582" + +[[package]] +name = "form_urlencoded" +version = "1.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cb4cb245038516f5f85277875cdaa4f7d2c9a0fa0468de06ed190163b1581fcf" +dependencies = [ + "percent-encoding", +] + +[[package]] +name = "futures-channel" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "07bbe89c50d7a535e539b8c17bc0b49bdb77747034daa8087407d655f3f7cc1d" +dependencies = [ + "futures-core", + "futures-sink", +] + +[[package]] +name = "futures-core" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7e3450815272ef58cec6d564423f6e755e25379b217b0bc688e295ba24df6b1d" + +[[package]] +name = "futures-io" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cecba35d7ad927e23624b22ad55235f2239cfa44fd10428eecbeba6d6a717718" + +[[package]] +name = "futures-sink" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c39754e157331b013978ec91992bde1ac089843443c49cbc7f46150b0fad0893" + +[[package]] +name = "futures-task" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "037711b3d59c33004d3856fbdc83b99d4ff37a24768fa1be9ce3538a1cde4393" + +[[package]] +name = "futures-util" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "389ca41296e6190b48053de0321d02a77f32f8a5d2461dd38762c0593805c6d6" +dependencies = [ + "futures-core", + "futures-io", + "futures-sink", + "futures-task", + "memchr", + "pin-project-lite", + "slab", +] + +[[package]] +name = "getrandom" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ff2abc00be7fca6ebc474524697ae276ad847ad0a6b3faa4bcb027e9a4614ad0" +dependencies = [ + "cfg-if", + "js-sys", + "libc", + "wasi", + "wasm-bindgen", +] + +[[package]] +name = "getrandom" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "300e883d756b2e4ec94e02791f39b04b522276138852cfc41d9fb7e904106099" +dependencies = [ + "cfg-if", + "js-sys", + "libc", + "r-efi", + "rand_core", + "wasm-bindgen", +] + +[[package]] +name = "heck" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" + +[[package]] +name = "http" +version = "1.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6970f50e31d6fc17d3fa27329444bfa74e196cf62e95052a3f6fee181dba6425" +dependencies = [ + "bytes", + "itoa", +] + +[[package]] +name = "http-body" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ca2a8f2913ee65f60facd6a5905613afaa448497a0230cc41ce022d93290bc2c" +dependencies = [ + "bytes", + "http", +] + +[[package]] +name = "http-body-util" +version = "0.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e9f41fd6a08e4d4ec69df65976da761afd5ad5e58a9d4acb46bd1c953a9e3ff2" +dependencies = [ + "bytes", + "futures-core", + "http", + "http-body", + "pin-project-lite", +] + +[[package]] +name = "httparse" +version = "1.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6dbf3de79e51f3d586ab4cb9d5c3e2c14aa28ed23d180cf89b4df0454a69cc87" + +[[package]] +name = "hyper" +version = "1.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "55281c53a1894c864990125767da440a4e630446785086f52523b20033b74498" +dependencies = [ + "atomic-waker", + "bytes", + "futures-channel", + "futures-core", + "http", + "http-body", + "httparse", + "itoa", + "pin-project-lite", + "smallvec", + "tokio", + "want", +] + +[[package]] +name = "hyper-rustls" +version = "0.27.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "33ca68d021ef39cf6463ab54c1d0f5daf03377b70561305bb89a8f83aab66e0f" +dependencies = [ + "http", + "hyper", + "hyper-util", + "rustls", + "tokio", + "tokio-rustls", + "tower-service", + "webpki-roots", +] + +[[package]] +name = "hyper-util" +version = "0.1.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "96547c2556ec9d12fb1578c4eaf448b04993e7fb79cbaad930a656880a6bdfa0" +dependencies = [ + "base64", + "bytes", + "futures-channel", + "futures-util", + "http", + "http-body", + "hyper", + "ipnet", + "libc", + "percent-encoding", + "pin-project-lite", + "socket2", + "tokio", + "tower-service", + "tracing", +] + +[[package]] +name = "icu_collections" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2984d1cd16c883d7935b9e07e44071dca8d917fd52ecc02c04d5fa0b5a3f191c" +dependencies = [ + "displaydoc", + "potential_utf", + "utf8_iter", + "yoke", + "zerofrom", + "zerovec", +] + +[[package]] +name = "icu_locale_core" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "92219b62b3e2b4d88ac5119f8904c10f8f61bf7e95b640d25ba3075e6cac2c29" +dependencies = [ + "displaydoc", + "litemap", + "tinystr", + "writeable", + "zerovec", +] + +[[package]] +name = "icu_normalizer" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c56e5ee99d6e3d33bd91c5d85458b6005a22140021cc324cea84dd0e72cff3b4" +dependencies = [ + "icu_collections", + "icu_normalizer_data", + "icu_properties", + "icu_provider", + "smallvec", + "zerovec", +] + +[[package]] +name = "icu_normalizer_data" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "da3be0ae77ea334f4da67c12f149704f19f81d1adf7c51cf482943e84a2bad38" + +[[package]] +name = "icu_properties" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bee3b67d0ea5c2cca5003417989af8996f8604e34fb9ddf96208a033901e70de" +dependencies = [ + "icu_collections", + "icu_locale_core", + "icu_properties_data", + "icu_provider", + "zerotrie", + "zerovec", +] + +[[package]] +name = "icu_properties_data" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8e2bbb201e0c04f7b4b3e14382af113e17ba4f63e2c9d2ee626b720cbce54a14" + +[[package]] +name = "icu_provider" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "139c4cf31c8b5f33d7e199446eff9c1e02decfc2f0eec2c8d71f65befa45b421" +dependencies = [ + "displaydoc", + "icu_locale_core", + "writeable", + "yoke", + "zerofrom", + "zerotrie", + "zerovec", +] + +[[package]] +name = "idna" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3b0875f23caa03898994f6ddc501886a45c7d3d62d04d2d90788d47be1b1e4de" +dependencies = [ + "idna_adapter", + "smallvec", + "utf8_iter", +] + +[[package]] +name = "idna_adapter" +version = "1.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cb68373c0d6620ef8105e855e7745e18b0d00d3bdb07fb532e434244cdb9a714" +dependencies = [ + "icu_normalizer", + "icu_properties", +] + +[[package]] +name = "ipnet" +version = "2.12.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d98f6fed1fde3f8c21bc40a1abb88dd75e67924f9cffc3ef95607bad8017f8e2" + +[[package]] +name = "is_terminal_polyfill" +version = "1.70.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a6cb138bb79a146c1bd460005623e142ef0181e3d0219cb493e02f7d08a35695" + +[[package]] +name = "itoa" +version = "1.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" + +[[package]] +name = "js-sys" +version = "0.3.103" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "53b44bfcdb3f8d5837a46dae1ca9660a837176eee74a28b229bc626816589102" +dependencies = [ + "cfg-if", + "futures-util", + "wasm-bindgen", +] + +[[package]] +name = "libc" +version = "0.2.186" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68ab91017fe16c622486840e4c83c9a37afeff978bd239b5293d61ece587de66" + +[[package]] +name = "litemap" +version = "0.8.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "92daf443525c4cce67b150400bc2316076100ce0b3686209eb8cf3c31612e6f0" + +[[package]] +name = "log" +version = "0.4.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0ceec5bc11778974d1bcb055b18002eba7f4b3518b6a0081b3af5f21666da9ad" + +[[package]] +name = "lru-slab" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "112b39cec0b298b6c1999fee3e31427f74f676e4cb9879ed1a121b43661a4154" + +[[package]] +name = "memchr" +version = "2.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf8baf1c55e62ffcace7a9f06f4bd9cd3f0c4beb022d3b367256b91b87513d98" + +[[package]] +name = "mio" +version = "1.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "30d65c71f1ce40ab09135ce117d742b9f8a19ff91a41a8b57ed50bc2de59c427" +dependencies = [ + "libc", + "wasi", + "windows-sys 0.61.2", +] + +[[package]] +name = "mlx-safetensors-stage-plan" +version = "0.0.0" +dependencies = [ + "anyhow", + "clap", + "reqwest", + "serde", + "serde_json", +] + +[[package]] +name = "once_cell" +version = "1.21.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" + +[[package]] +name = "once_cell_polyfill" +version = "1.70.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "384b8ab6d37215f3c5301a95a4accb5d64aa607f1fcb26a11b5303878451b4fe" + +[[package]] +name = "percent-encoding" +version = "2.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9b4f627cb1b25917193a259e49bdad08f671f8d9708acfd5fe0a8c1455d87220" + +[[package]] +name = "pin-project-lite" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a89322df9ebe1c1578d689c92318e070967d1042b512afbe49518723f4e6d5cd" + +[[package]] +name = "potential_utf" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0103b1cef7ec0cf76490e969665504990193874ea05c85ff9bab8b911d0a0564" +dependencies = [ + "zerovec", +] + +[[package]] +name = "proc-macro2" +version = "1.0.106" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8fd00f0bb2e90d81d1044c2b32617f68fcb9fa3bb7640c23e9c748e53fb30934" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "quinn" +version = "0.11.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c1a41e437b6bbd489372cd4971de128e85c855f56c57f283d20ff016cf7c0a8" +dependencies = [ + "bytes", + "cfg_aliases", + "pin-project-lite", + "quinn-proto", + "quinn-udp", + "rustc-hash", + "rustls", + "socket2", + "thiserror", + "tokio", + "tracing", + "web-time", +] + +[[package]] +name = "quinn-proto" +version = "0.11.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2f4bfc015262b9df63c8845072ce59068853ff5872180c2ce2f13038b970e560" +dependencies = [ + "bytes", + "getrandom 0.4.3", + "lru-slab", + "rand", + "rand_pcg", + "ring", + "rustc-hash", + "rustls", + "rustls-pki-types", + "slab", + "thiserror", + "tinyvec", + "tracing", + "web-time", +] + +[[package]] +name = "quinn-udp" +version = "0.5.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "35a133f956daabe89a61a685c2649f13d82d5aa4bd5d12d1277e1072a21c0694" +dependencies = [ + "cfg_aliases", + "libc", + "once_cell", + "socket2", + "tracing", + "windows-sys 0.61.2", +] + +[[package]] +name = "quote" +version = "1.0.46" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dfbc457d0c7a0759a614551b11a6409e5951f6c7537be1f1b7682b9ae9230368" +dependencies = [ + "proc-macro2", +] + +[[package]] +name = "r-efi" +version = "6.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8dcc9c7d52a811697d2151c701e0d08956f92b0e24136cf4cf27b57a6a0d9bf" + +[[package]] +name = "rand" +version = "0.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c7f5fa3a058cd35567ef9bfa5e75732bee0f9e4c55fa90477bef2dfcdbc4be80" +dependencies = [ + "chacha20", + "getrandom 0.4.3", + "rand_core", +] + +[[package]] +name = "rand_core" +version = "0.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "63b8176103e19a2643978565ca18b50549f6101881c443590420e4dc998a3c69" + +[[package]] +name = "rand_pcg" +version = "0.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "caa0f4137e1c0a72f4c651489402276c8e8e1cf081f3b0ba156d2cbeef09e86a" +dependencies = [ + "rand_core", +] + +[[package]] +name = "reqwest" +version = "0.12.28" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "eddd3ca559203180a307f12d114c268abf583f59b03cb906fd0b3ff8646c1147" +dependencies = [ + "base64", + "bytes", + "futures-channel", + "futures-core", + "futures-util", + "http", + "http-body", + "http-body-util", + "hyper", + "hyper-rustls", + "hyper-util", + "js-sys", + "log", + "percent-encoding", + "pin-project-lite", + "quinn", + "rustls", + "rustls-pki-types", + "serde", + "serde_json", + "serde_urlencoded", + "sync_wrapper", + "tokio", + "tokio-rustls", + "tower", + "tower-http", + "tower-service", + "url", + "wasm-bindgen", + "wasm-bindgen-futures", + "web-sys", + "webpki-roots", +] + +[[package]] +name = "ring" +version = "0.17.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a4689e6c2294d81e88dc6261c768b63bc4fcdb852be6d1352498b114f61383b7" +dependencies = [ + "cc", + "cfg-if", + "getrandom 0.2.17", + "libc", + "untrusted", + "windows-sys 0.52.0", +] + +[[package]] +name = "rustc-hash" +version = "2.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6b1e7f9a428571be2dc5bc0505c13fb6bf936822b894ec87abf8a08a4e51742d" + +[[package]] +name = "rustls" +version = "0.23.42" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3c54fcab019b409d04215d3a17cb438fd7fbf192ee61461f20f4fe18704bc138" +dependencies = [ + "once_cell", + "ring", + "rustls-pki-types", + "rustls-webpki", + "subtle", + "zeroize", +] + +[[package]] +name = "rustls-pki-types" +version = "1.15.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "764899a24af3980067ee14bc143654f297b22eaebfe3c7b6b211920a5a59b046" +dependencies = [ + "web-time", + "zeroize", +] + +[[package]] +name = "rustls-webpki" +version = "0.103.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "61c429a8649f110dddef65e2a5ad240f747e85f7758a6bccc7e5777bd33f756e" +dependencies = [ + "ring", + "rustls-pki-types", + "untrusted", +] + +[[package]] +name = "rustversion" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf54715a573b99ac80df0bc206da022bcd442c974952c7b9720069370852e21f" + +[[package]] +name = "ryu" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9774ba4a74de5f7b1c1451ed6cd5285a32eddb5cccb8cc655a4e50009e06477f" + +[[package]] +name = "serde" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9a8e94ea7f378bd32cbbd37198a4a91436180c5bb472411e48b5ec2e2124ae9e" +dependencies = [ + "serde_core", + "serde_derive", +] + +[[package]] +name = "serde_core" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "41d385c7d4ca58e59fc732af25c3983b67ac852c1a25000afe1175de458b67ad" +dependencies = [ + "serde_derive", +] + +[[package]] +name = "serde_derive" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d540f220d3187173da220f885ab66608367b6574e925011a9353e4badda91d79" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "serde_json" +version = "1.0.150" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e8014e44b4736ed0538adeecded0fce2a272f22dc9578a7eb6b2d9993c74cfb9" +dependencies = [ + "itoa", + "memchr", + "serde", + "serde_core", + "zmij", +] + +[[package]] +name = "serde_urlencoded" +version = "0.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d3491c14715ca2294c4d6a88f15e84739788c1d030eed8c110436aafdaa2f3fd" +dependencies = [ + "form_urlencoded", + "itoa", + "ryu", + "serde", +] + +[[package]] +name = "shlex" +version = "2.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8fadd59c855ef2080decdef8ff161eb6661b86933c9d82e5ba29dc602a55aba" + +[[package]] +name = "slab" +version = "0.4.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c790de23124f9ab44544d7ac05d60440adc586479ce501c1d6d7da3cd8c9cf5" + +[[package]] +name = "smallvec" +version = "1.15.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8ed6a63f02c8539c91a8685a86f4099661ba3da017932f6ebbea6de3f0fa7c90" + +[[package]] +name = "socket2" +version = "0.6.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c3d1e2c7f27f8d4cb10542a02c49005dbd6e93095799d6f3be745fae9f8fedd4" +dependencies = [ + "libc", + "windows-sys 0.61.2", +] + +[[package]] +name = "stable_deref_trait" +version = "1.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ce2be8dc25455e1f91df71bfa12ad37d7af1092ae736f3a6cd0e37bc7810596" + +[[package]] +name = "strsim" +version = "0.11.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7da8b5736845d9f2fcb837ea5d9e2628564b3b043a70948a3f0b778838c5fb4f" + +[[package]] +name = "subtle" +version = "2.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "13c2bddecc57b384dee18652358fb23172facb8a2c51ccc10d74c157bdea3292" + +[[package]] +name = "syn" +version = "2.0.119" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "872831b642d1a07999a962a351ed35b955ea2cfc8f3862091e2a240a84f17297" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "sync_wrapper" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0bf256ce5efdfa370213c1dabab5935a12e49f2c58d15e9eac2870d3b4f27263" +dependencies = [ + "futures-core", +] + +[[package]] +name = "synstructure" +version = "0.13.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "728a70f3dbaf5bab7f0c4b1ac8d7ae5ea60a4b5549c8a5914361c99147a709d2" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "thiserror" +version = "2.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4288b5bcbc7920c07a1149a35cf9590a2aa808e0bc1eafaade0b80947865fbc4" +dependencies = [ + "thiserror-impl", +] + +[[package]] +name = "thiserror-impl" +version = "2.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ebc4ee7f67670e9b64d05fa4253e753e016c6c95ff35b89b7941d6b856dec1d5" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "tinystr" +version = "0.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8323304221c2a851516f22236c5722a72eaa19749016521d6dff0824447d96d" +dependencies = [ + "displaydoc", + "zerovec", +] + +[[package]] +name = "tinyvec" +version = "1.12.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bb4ebadaa0af04fab11ae01eb5f9fdb5f9c5b875506e210e71c07873528baa7f" +dependencies = [ + "tinyvec_macros", +] + +[[package]] +name = "tinyvec_macros" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1f3ccbac311fea05f86f61904b462b55fb3df8837a366dfc601a0161d0532f20" + +[[package]] +name = "tokio" +version = "1.52.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "317fafbbe3f02fc663dad00ea6186197de963cd4190e86a26d8d0fae095539af" +dependencies = [ + "bytes", + "libc", + "mio", + "pin-project-lite", + "socket2", + "windows-sys 0.61.2", +] + +[[package]] +name = "tokio-rustls" +version = "0.26.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1729aa945f29d91ba541258c8df89027d5792d85a8841fb65e8bf0f4ede4ef61" +dependencies = [ + "rustls", + "tokio", +] + +[[package]] +name = "tower" +version = "0.5.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ebe5ef63511595f1344e2d5cfa636d973292adc0eec1f0ad45fae9f0851ab1d4" +dependencies = [ + "futures-core", + "futures-util", + "pin-project-lite", + "sync_wrapper", + "tokio", + "tower-layer", + "tower-service", +] + +[[package]] +name = "tower-http" +version = "0.6.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4cfcf7e2740e6fc6d4d688b4ef00650406bb94adf4731e43c096c3a19fe40840" +dependencies = [ + "bitflags", + "bytes", + "futures-util", + "http", + "http-body", + "pin-project-lite", + "tower", + "tower-layer", + "tower-service", + "url", +] + +[[package]] +name = "tower-layer" +version = "0.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "121c2a6cda46980bb0fcd1647ffaf6cd3fc79a013de288782836f6df9c48780e" + +[[package]] +name = "tower-service" +version = "0.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8df9b6e13f2d32c91b9bd719c00d1958837bc7dec474d94952798cc8e69eeec3" + +[[package]] +name = "tracing" +version = "0.1.44" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "63e71662fa4b2a2c3a26f570f037eb95bb1f85397f3cd8076caed2f026a6d100" +dependencies = [ + "pin-project-lite", + "tracing-core", +] + +[[package]] +name = "tracing-core" +version = "0.1.36" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "db97caf9d906fbde555dd62fa95ddba9eecfd14cb388e4f491a66d74cd5fb79a" +dependencies = [ + "once_cell", +] + +[[package]] +name = "try-lock" +version = "0.2.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e421abadd41a4225275504ea4d6566923418b7f05506fbc9c0fe86ba7396114b" + +[[package]] +name = "unicode-ident" +version = "1.0.24" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75" + +[[package]] +name = "untrusted" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8ecb6da28b8a351d773b68d5825ac39017e680750f980f3a1a85cd8dd28a47c1" + +[[package]] +name = "url" +version = "2.5.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ff67a8a4397373c3ef660812acab3268222035010ab8680ec4215f38ba3d0eed" +dependencies = [ + "form_urlencoded", + "idna", + "percent-encoding", + "serde", +] + +[[package]] +name = "utf8_iter" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6c140620e7ffbb22c2dee59cafe6084a59b5ffc27a8859a5f0d494b5d52b6be" + +[[package]] +name = "utf8parse" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "06abde3611657adf66d383f00b093d7faecc7fa57071cce2578660c9f1010821" + +[[package]] +name = "want" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bfa7760aed19e106de2c7c0b581b509f2f25d3dacaf737cb82ac61bc6d760b0e" +dependencies = [ + "try-lock", +] + +[[package]] +name = "wasi" +version = "0.11.1+wasi-snapshot-preview1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ccf3ec651a847eb01de73ccad15eb7d99f80485de043efb2f370cd654f4ea44b" + +[[package]] +name = "wasm-bindgen" +version = "0.2.126" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4b067c0c11094aef6b7a801c1e34a26affafdf3d051dba08456b868789aaf9a4" +dependencies = [ + "cfg-if", + "once_cell", + "rustversion", + "wasm-bindgen-macro", + "wasm-bindgen-shared", +] + +[[package]] +name = "wasm-bindgen-futures" +version = "0.4.76" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c62df1340f32221cb9c54d6a27b030e3dba64361d4a95bed55f9aacb44da291d" +dependencies = [ + "js-sys", + "wasm-bindgen", +] + +[[package]] +name = "wasm-bindgen-macro" +version = "0.2.126" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "167ce5e579f6bcf889c4f7175a8a5a585de84e8ff93976ce393efa5f2837aab1" +dependencies = [ + "quote", + "wasm-bindgen-macro-support", +] + +[[package]] +name = "wasm-bindgen-macro-support" +version = "0.2.126" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f3997c7839262f4ef12cf90b818d6340c18e80f263f1a94bf157d0ec4420380e" +dependencies = [ + "bumpalo", + "proc-macro2", + "quote", + "syn", + "wasm-bindgen-shared", +] + +[[package]] +name = "wasm-bindgen-shared" +version = "0.2.126" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dc1b4cb0cc549fcf58d7dfc081778139b3d283a081644e833e84682ad71cea24" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "web-sys" +version = "0.3.103" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8622dcb61c0bcc9fffa6938bed81210af2da9a7e4a1a834b2e37a59b6dfb6141" +dependencies = [ + "js-sys", + "wasm-bindgen", +] + +[[package]] +name = "web-time" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5a6580f308b1fad9207618087a65c04e7a10bc77e02c8e84e9b00dd4b12fa0bb" +dependencies = [ + "js-sys", + "wasm-bindgen", +] + +[[package]] +name = "webpki-roots" +version = "1.0.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bf85cb06032201fa7c6f829d7db5a7e5aa45bcc0655327713065f6f0576731bf" +dependencies = [ + "rustls-pki-types", +] + +[[package]] +name = "windows-link" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" + +[[package]] +name = "windows-sys" +version = "0.52.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "282be5f36a8ce781fad8c8ae18fa3f9beff57ec1b52cb3de0789201425d9a33d" +dependencies = [ + "windows-targets", +] + +[[package]] +name = "windows-sys" +version = "0.61.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ae137229bcbd6cdf0f7b80a31df61766145077ddf49416a728b02cb3921ff3fc" +dependencies = [ + "windows-link", +] + +[[package]] +name = "windows-targets" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9b724f72796e036ab90c1021d4780d4d3d648aca59e491e6b98e725b84e99973" +dependencies = [ + "windows_aarch64_gnullvm", + "windows_aarch64_msvc", + "windows_i686_gnu", + "windows_i686_gnullvm", + "windows_i686_msvc", + "windows_x86_64_gnu", + "windows_x86_64_gnullvm", + "windows_x86_64_msvc", +] + +[[package]] +name = "windows_aarch64_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32a4622180e7a0ec044bb555404c800bc9fd9ec262ec147edd5989ccd0c02cd3" + +[[package]] +name = "windows_aarch64_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09ec2a7bb152e2252b53fa7803150007879548bc709c039df7627cabbd05d469" + +[[package]] +name = "windows_i686_gnu" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8e9b5ad5ab802e97eb8e295ac6720e509ee4c243f69d781394014ebfe8bbfa0b" + +[[package]] +name = "windows_i686_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0eee52d38c090b3caa76c563b86c3a4bd71ef1a819287c19d586d7334ae8ed66" + +[[package]] +name = "windows_i686_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "240948bc05c5e7c6dabba28bf89d89ffce3e303022809e73deaefe4f6ec56c66" + +[[package]] +name = "windows_x86_64_gnu" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "147a5c80aabfbf0c7d901cb5895d1de30ef2907eb21fbbab29ca94c5b08b1a78" + +[[package]] +name = "windows_x86_64_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "24d5b23dc417412679681396f2b49f3de8c1473deb516bd34410872eff51ed0d" + +[[package]] +name = "windows_x86_64_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec" + +[[package]] +name = "writeable" +version = "0.6.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ffae5123b2d3fc086436f8834ae3ab053a283cfac8fe0a0b8eaae044768a4c4" + +[[package]] +name = "yoke" +version = "0.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "709fe23a0424b6a435d82152b1bd3fdfb0833487d5fa90d05d42762a9891fef5" +dependencies = [ + "stable_deref_trait", + "yoke-derive", + "zerofrom", +] + +[[package]] +name = "yoke-derive" +version = "0.8.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "de844c262c8848816172cef550288e7dc6c7b7814b4ee56b3e1553f275f1858e" +dependencies = [ + "proc-macro2", + "quote", + "syn", + "synstructure", +] + +[[package]] +name = "zerofrom" +version = "0.1.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0ec05a11813ea801ff6d75110ad09cd0824ddba17dfe17128ea0d5f68e6c5272" +dependencies = [ + "zerofrom-derive", +] + +[[package]] +name = "zerofrom-derive" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "11532158c46691caf0f2593ea8358fed6bbf68a0315e80aae9bd41fbade684a1" +dependencies = [ + "proc-macro2", + "quote", + "syn", + "synstructure", +] + +[[package]] +name = "zeroize" +version = "1.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e13c156562582aa81c60cb29407084cdb54c4164760106ab78e6c5b0858cf64e" + +[[package]] +name = "zerotrie" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0f9152d31db0792fa83f70fb2f83148effb5c1f5b8c7686c3459e361d9bc20bf" +dependencies = [ + "displaydoc", + "yoke", + "zerofrom", +] + +[[package]] +name = "zerovec" +version = "0.11.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "90f911cbc359ab6af17377d242225f4d75119aec87ea711a880987b18cd7b239" +dependencies = [ + "yoke", + "zerofrom", + "zerovec-derive", +] + +[[package]] +name = "zerovec-derive" +version = "0.11.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "625dc425cab0dca6dc3c3319506e6593dcb08a9f387ea3b284dbd52a92c40555" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "zmij" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29666d0abbfad1e3dc4dcf6144730dd3a3ab225bbbdac83319345b1b44ccfc1b" diff --git a/spikes/mlx-safetensors-stages/Cargo.toml b/spikes/mlx-safetensors-stages/Cargo.toml new file mode 100644 index 0000000000..dd9b22e1b7 --- /dev/null +++ b/spikes/mlx-safetensors-stages/Cargo.toml @@ -0,0 +1,21 @@ +[package] +name = "mlx-safetensors-stage-plan" +version = "0.0.0" +edition = "2024" +publish = false + +[[bin]] +name = "mlx-safetensors-stage-plan" +path = "src/main.rs" + +# Standalone research spike: keep HTTP dependencies and its lockfile out of the +# mesh-llm workspace while the stage-cache design is still being validated. +[workspace] + +[dependencies] +anyhow = "1" +clap = { version = "4", features = ["derive"] } +reqwest = { version = "0.12", default-features = false, features = ["blocking", "rustls-tls"] } +serde = { version = "1", features = ["derive"] } +serde_json = "1" + diff --git a/spikes/mlx-safetensors-stages/FINDINGS.md b/spikes/mlx-safetensors-stages/FINDINGS.md new file mode 100644 index 0000000000..1b955b5d29 --- /dev/null +++ b/spikes/mlx-safetensors-stages/FINDINGS.md @@ -0,0 +1,321 @@ +# SafeTensors stage-local download and adaptive MLX quantization + +## Status + +Metadata planning, exact-range materialization, and two-stage MLX execution +proofs completed on 2026-07-17. The frontier-model measurements inspect headers +only; the SmolLM2 proof downloads and executes selected tensor payloads. + +The standalone `mlx-safetensors-stage-plan` spike proves that a layer server can: + +1. read `model.safetensors.index.json`; +2. select only tensors owned by its layer range; +3. fetch the 8-byte length and JSON header from each relevant SafeTensors file; +4. turn each tensor's `data_offsets` into absolute HTTP byte ranges; and +5. stream those ranges into a valid partial `model.safetensors` artifact. + +It refuses a response other than HTTP `206 Partial Content`, preventing an +ignored `Range` header from silently downloading a multi-gigabyte shard. + +## Bottom line + +SafeTensors supports the desired model: canonical upstream weights remain on +Hugging Face, while each layer server downloads and caches only the tensors it +owns. A separately published layer-package repository is not required. + +For normally ordered checkpoints, selecting whole shard files already gets +close to exact. For Inkling, tensors are heavily interleaved across source +shards, so exact tensor ranges are mandatory: whole-shard selection would turn +a 109.84 GiB four-layer stage into a 942.99 GiB download. + +The small dense-model path is now proven through execution. The remaining +artifact proof is bounded-memory tensor-at-a-time MLX quantization for a source +tensor too large to retain alongside a whole BF16 stage. + +## Reproduce + +```bash +just mlx-safetensors-stage-plan \ + --repo thinkingmachines/Inkling \ + --revision 86b4d430ab871652a707666b89203a866888c5e5 \ + --layer-start 30 \ + --layer-end 34 +``` + +Use `--json` for the per-shard byte ranges. Additional tensors can be assigned +with repeated `--include-prefix` arguments; for example the first stage can own +the embedding and modality towers, while the final stage owns final norm, +readout, and optional MTP tensors. + +Without `--output`, the CLI fetches only the index and SafeTensors headers. With +`--output `, it fetches the selected payload ranges, writes +`model.safetensors`, `config.json`, and a reproducible `stage-plan.json`, and +still refuses any payload response other than HTTP 206. + +## Small-model execution proof + +`HuggingFaceTB/SmolLM2-135M-Instruct` at immutable revision +`12fd25f77366fa6b3b4b768ec3050bf629380bac` was split unnecessarily at layer 15: + +| Stage | Owned tensors | Exact payload | Whole checkpoint | Avoided locally | HTTP payload spans | +| --- | ---: | ---: | ---: | ---: | ---: | +| 0: embedding + layers 0..15 | 136 | 155.28 MiB | 256.60 MiB | 101.28 MiB | 3 | +| 1: layers 15..30 + norm + tied embedding | 137 | 155.28 MiB | 256.60 MiB | 101.28 MiB | 4 | + +The tied embedding is intentionally duplicated: stage 0 uses it for token +input, while stage 1 uses it as the tied output projection. Neither stage file +contains the complete checkpoint. A strict whole-model baseline is assembled +from the union of the two partial files, so the parity test cannot silently +fall back to a full download. + +The `mlx-split-proof` harness runs layers 0..15 and 15..30 as separate MLX +stages, serializes the residual through Skippy's real `StageWireMessage` binary +codec, maintains independent per-stage KV caches, and compares against unsplit +execution. Prompt prefill plus eight greedy decode steps passed on Metal with +both F16 and F32 wire encodings: + +- identical eight-token sequence: `284, 260, 2240, 314, 1343, 327, 624, 8685`; +- worst maximum absolute logit delta: `0.0` for F16 and F32; and +- F16 total stage-wire traffic: 15,584 bytes for the tested prefill and decode. + +One important engine contract emerged: after decoding F16/F32 wire bytes, the +receiving MLX stage must cast the residual back to the model's compute dtype +(BF16 here) before its first block. Leaving the reconstructed array as F32 +changed Metal arithmetic and immediately changed the greedy token, despite +identical numeric boundary values. + +Reproduction uses the two materializer invocations followed by: + +```bash +just mlx-safetensors-split-proof \ + --stage0 /tmp/mlx-split-smol/stage0 \ + --stage1 /tmp/mlx-split-smol/stage1 \ + --split 15 \ + --steps 8 \ + --wire-dtype f16 +``` + +This proves the artifact, MLX layer-range, KV-cache, and existing binary +activation-frame seams on one Mac. It is not yet a two-process/two-node +`mesh-llm serve` implementation; `skippy-server` still binds directly to the +llama.cpp `StageModel` and needs the planned engine abstraction first. + +## Representative measurements + +All rows select four middle transformer layers. Layer types vary within hybrid +models, so the table demonstrates storage locality rather than equal compute. +Every repository is pinned to the immutable revision shown below. + +| Model / source encoding | Revision | Layers | Full tensor bytes | Selected bytes | Whole relevant shards | Avoided by exact ranges | Largest selected tensor | +| --- | --- | ---: | ---: | ---: | ---: | ---: | ---: | +| Qwen3-235B-A22B BF16 | `8efa61729e24bd65b1d152b5ab5409052aa80e65` | 40..44 | 437.89 GiB | 18.54 GiB | 22.31 GiB | 3.78 GiB | not recorded | +| Inkling BF16 | `86b4d430ab871652a707666b89203a866888c5e5` | 30..34 | 1.73 TiB | 109.84 GiB | 942.99 GiB | 833.15 GiB | 18.00 GiB | +| Inkling official NVFP4 | `d11961f515e883e37796edb9dd6ec1bf0e0e8212` | 30..34 | 551.21 GiB | 32.22 GiB | 501.59 GiB | 469.37 GiB | not recorded | +| Inkling community MLX affine-4 | `34f92fe0879faa413071c8dad23538014f0c266b` | 30..34 | 521.97 GiB | 32.22 GiB | 40.27 GiB | 8.05 GiB | not recorded | +| Nemotron 3 Ultra 550B BF16 | `624ba927cfbef0427354998700de3d51173c8c04` | 48..52 | 1.02 TiB | 41.81 GiB | 46.44 GiB | 4.63 GiB | 548 MiB | +| Kimi K2.6 published checkpoint | `7eb5002f6aadc958aed6a9177b7ed26bb94011bb` | 28..32 | 554.27 GiB | 36.54 GiB | 36.54 GiB | 0 | 112 MiB | +| GLM-5.2 BF16 | `b4734de4facf877f85769a911abafc5283eab3d9` | 36..40 | 1.37 TiB | 73.54 GiB | 79.90 GiB | 6.36 GiB | 192 MiB | +| DeepSeek V4 Pro FP8 | `b5968e9190ef611bbf34a7229255be88a0e937c1` | 28..32 | 805.32 GiB | 51.73 GiB | 51.73 GiB | 0 | 112 MiB | + +The format mechanism is stable and documented by the +[SafeTensors format](https://github.com/huggingface/safetensors#format): the +header records each tensor's dtype, shape, and byte offsets. Hugging Face's +[Xet download protocol](https://huggingface.co/docs/xet/download-protocol#range-downloads) +supports partial-file reconstruction, and the public `resolve` endpoint honored +the byte ranges in every probe above. + +## Inkling as the frontier stress test + +[Inkling](https://huggingface.co/thinkingmachines/Inkling) is a 975B-total, +41B-active, 66-layer multimodal MoE with: + +- 256 routed experts, with 6 selected per token, plus 2 shared experts; +- 6144-wide residual states; +- relative-position attention rather than RoPE; +- a 5:1 sliding-window/global-attention pattern; +- four short-convolution states per decoder layer; +- text, vision, and audio inputs; +- a 1,048,576-token maximum context; and +- eight optional MTP predictor layers. + +The source artifacts are: + +| Artifact | Exact tensor bytes | Notes | +| --- | ---: | --- | +| `thinkingmachines/Inkling` | 1,904,604,285,204 | Canonical BF16; 109 weight files plus index | +| `thinkingmachines/Inkling-NVFP4` | 591,854,374,368 | Official calibrated NVFP4; intended for Blackwell-class serving stacks | +| `mlx-community/Inkling-mlx-4bit` | 560,463,783,044 | Text-only mixed MLX affine-4 experiment | + +The community MLX artifact is useful size evidence, but it is not yet a runnable +or certified answer for mesh-llm. Its own model card says the custom Inkling +forward is not registered in upstream `mlx-lm`, logits are not numerically +verified, and the vision/audio towers are excluded. The repository contains +weights and tokenizer/config files but not the custom model implementation. + +Neither upstream `mlx-lm` nor the pinned Rust `safemlx-lm` dependency currently +ships an Inkling family. The authoritative reference implementation is now in +[Transformers' Inkling model](https://github.com/huggingface/transformers/blob/main/src/transformers/models/inkling/modular_inkling.py). + +### What an Inkling MLX stage engine must implement + +1. Text-decoder parity first: relative-logit attention, local/global masks, + query scaling, sigmoid top-k routing, routed and shared experts, and all four + SConv paths. +2. A stage constructor that creates only `layers[start..end]`, with embeddings + on the first stage and final norm/readout on the last. +3. Stage-local KV plus SConv recurrent state. Inkling cannot be treated as a + plain paged-KV Llama family. +4. An Inkling-specific weight loader and quantization predicate. +5. Logit parity against Transformers at several layer cuts before network work. +6. Vision/audio towers on the first stage after the text chain is certified. +7. MTP as a separate optional capability after ordinary decode is correct. + +The residual-stream boundary remains clean between decoder layers, so these +features make family bring-up substantial but do not invalidate pipeline +splitting. + +## Hardware-adaptive quantization + +The user's proposed model is viable: choose a quantization plan after topology +and hardware discovery, then quantize only each server's selected tensors during +cold load. MLX directly supports affine 2/3/4/5/6/8-bit quantization with group +sizes 32/64/128, plus MXFP4, MXFP8, and NVFP4. It also accepts a per-module +predicate, allowing sensitive modules and different layer ranges to retain more +precision. See [`mlx.core.quantize`](https://ml-explore.github.io/mlx/build/html/python/_autosummary/mlx.core.quantize.html) +and [`mlx.nn.quantize`](https://ml-explore.github.io/mlx/build/html/python/nn/_autosummary/mlx.nn.quantize.html). + +For Inkling, start with the community conversion's conservative policy: + +- quantize routed-expert matrices only; +- keep attention, router, shared experts, embeddings, normalization, relative + projections, and SConv weights in BF16; +- use affine 4-bit, group size 64; and +- never derive an MLX quant from the official NVFP4 artifact when BF16 is + available, because that would be a lossy requantization. + +The measured 4-bit artifact and MLX affine storage formula imply approximately +1.870 TB of BF16 source tensors are quantizable and 34.5 GB remain BF16. Holding +the same predicate constant gives these rough storage targets: + +| Routed-expert affine precision | Estimated total weights | +| --- | ---: | +| 2-bit, group 64 | 304 GiB | +| 3-bit, group 64 | 413 GiB | +| 4-bit, group 64 | 522 GiB (matches measured artifact) | +| 5-bit, group 64 | 631 GiB | +| 6-bit, group 64 | 740 GiB | +| 8-bit, group 64 | 957 GiB | + +These are capacity estimates, not quality endorsements. Two- and three-bit +profiles need evaluation, and the current community 4-bit artifact itself is +not yet logit-verified. MLX's sensitivity-based dynamic quantization can produce +mixed-bit profiles, but for a frontier model the sensitivity result should be +computed and certified once per model revision, stored as a small profile, and +then applied deterministically by every stage. Recomputing sensitivity during +every cold load would be too expensive. + +Different stages may use different precision when hardware differs. The chosen +per-stage profile must be part of the topology/model identity so that a cached +stage is reproducible and correctness evidence names the exact numeric model. + +### Cold-load memory contract + +Load-time quantization only makes small nodes viable if it is streamed. For +Inkling layers 30..33: + +- BF16 input ranges: 109.84 GiB; +- resulting mixed affine-4 stage: 32.22 GiB; and +- largest single BF16 source tensor: 18.00 GiB. + +A whole-stage loader would need BF16 input plus quantized output and fail on a +128 GB node. A tensor-streaming loader can keep the accumulated 32.22 GiB target +plus one source tensor and quantization scratch resident. The cold path should: + +1. range-fetch one tensor into a bounded temporary/mmap buffer; +2. create the MLX source array; +3. quantize according to the certified per-tensor profile; +4. evaluate and append the packed tensor/scales/biases to the derived cache; +5. release the BF16 source buffer; and +6. continue with the next tensor. + +The same rule applies to disk: do not retain an entire BF16 stage unless the +operator asks for it. A derived cache can approach `quantized stage + largest +source tensor`, rather than `BF16 stage + quantized stage`. + +### Approximate Inkling 4-bit deployment shapes + +The 521.97 GiB text-weight artifact plus Inkling's long-context cache and load +scratch makes aggregate memory, not raw model size alone, the constraint. +At full 1M context, BF16 KV is approximately 44 GiB: eleven global-attention +layers each retain about 4 GiB, while the 55 sliding layers retain only their +512-token windows (about 220 MiB combined). SConv state is comparatively small. + +Assuming balanced stages and the 18 GiB largest source tensor: + +| Topology | Weight share/node | Assessment before measured runtime overhead | +| --- | ---: | --- | +| 2 × 512 GB | ~261 GiB | Comfortable capacity; simplest plausible full-context shape | +| 3 × 256 GB | ~174 GiB | Comfortable capacity | +| 4 × 192 GB | ~131 GiB | Plausible | +| 5 × 128 GB | ~104 GiB | Too tight once 18 GiB load scratch and KV are included | +| 6 × 128 GB | ~87 GiB | Plausible but needs measured allocator/kernel headroom | +| 8 × 128 GB | ~65 GiB | Safer first 128 GB-node target | +| 12 × 64 GB | ~44 GiB | Too tight during 18 GiB source-tensor quantization | +| 16 × 64 GB | ~33 GiB | Plausible capacity; stage latency may dominate | + +Shorter context materially reduces the KV portion. These are feasibility +estimates, not throughput claims; MoE dispatch performance, per-stage latency, +and the MLX boundary fence still need measurement. + +## Frontier-family prioritization + +SafeTensors acquisition is general, but MLX execution remains family-specific. +The measured candidates suggest this order: + +1. **Qwen/Llama**: finish the partial loader and two-stage correctness proof. +2. **Nemotron 3 Ultra**: best next frontier proof because `safemlx-lm` already + has a Nemotron-H implementation, although its Mamba/recurrent blocks still + require state-boundary certification. +3. **Inkling text backbone**: high-value new family; use the upstream + Transformers modular implementation as the parity oracle. +4. **Inkling multimodal + MTP**: add first-stage towers and optional predictor + layers after text correctness. +5. **Kimi K2.6, GLM-5.2, DeepSeek V4**: all are viable range-download targets, + but each requires a new or substantially updated MLX family and native + support for its existing compressed format or a canonical BF16 source. + +Not every frontier repository should be requantized at load. Kimi K2.6's +published checkpoint is already about 554 GiB, and DeepSeek V4 Pro is already +FP8. Preserve a compatible calibrated source encoding when the local backend +supports it; use BF16-to-local-quant only when it is the cleanest compatible +source path. + +## Recommended artifact identity + +The durable identity should be: + +```text +source repo + immutable revision ++ selected tensor names and source byte ranges ++ model-family implementation revision ++ stage range / embedding / readout / modality ownership ++ per-stage quantization profile ++ activation wire dtype += derived stage cache identity +``` + +The cache is evictable derived data. The upstream checkpoint remains the source +of truth, and a small certified quantization/profile manifest replaces a large +published layer-package repository. + +## Next proof + +1. Stream BF16 tensor -> MLX affine quant -> partial SafeTensors output one + tensor at a time; prove bounded RSS and disk use. +2. Measure the MLX eval/readback/codec boundary fence independently at frontier + residual widths and prefill sizes. +3. Introduce the engine-neutral stage interface and run the same proof through + two real `skippy-server` processes. +4. Repeat the loader proof with Nemotron-H before implementing Inkling. +5. Port Inkling's text decoder to `safemlx-lm`, starting with one layer and + Transformers parity, then stage ranges, then network execution. diff --git a/spikes/mlx-safetensors-stages/src/main.rs b/spikes/mlx-safetensors-stages/src/main.rs new file mode 100644 index 0000000000..e1ce70bbd8 --- /dev/null +++ b/spikes/mlx-safetensors-stages/src/main.rs @@ -0,0 +1,883 @@ +use std::collections::{BTreeMap, BTreeSet}; +use std::env; +use std::fs::{self, File}; +use std::io::{BufWriter, Write}; +use std::ops::Range; +use std::path::{Path, PathBuf}; +use std::time::Duration; + +use anyhow::{Context, Result, anyhow, ensure}; +use clap::Parser; +use reqwest::StatusCode; +use reqwest::Url; +use reqwest::blocking::{Client, Response}; +use reqwest::header::{AUTHORIZATION, CONTENT_LENGTH, RANGE}; +use reqwest::redirect::Policy; +use serde::{Deserialize, Serialize}; + +const MAX_INDEX_BYTES: u64 = 64 * 1024 * 1024; +const MAX_HEADER_BYTES: u64 = 256 * 1024 * 1024; + +#[derive(Debug, Parser)] +#[command(about = "Plan exact SafeTensors byte ranges for one transformer layer stage")] +struct Args { + /// Hugging Face model repository, for example Qwen/Qwen3-235B-A22B. + #[arg(long)] + repo: String, + + /// Repository revision. Use an immutable commit SHA for reproducible plans. + #[arg(long, default_value = "main")] + revision: String, + + /// First transformer layer owned by this stage. + #[arg(long)] + layer_start: u32, + + /// Exclusive end of the transformer layer range. + #[arg(long)] + layer_end: u32, + + /// Include additional exact tensor-name prefixes, such as model.embed_tokens. + #[arg(long = "include-prefix")] + include_prefixes: Vec, + + /// Merge ranges separated by at most this many unneeded bytes. + #[arg(long, default_value_t = 0)] + coalesce_gap_bytes: u64, + + #[arg(long, default_value = "https://huggingface.co")] + endpoint: String, + + #[arg(long)] + json: bool, + + /// Materialize the selected tensors as output/model.safetensors. + #[arg(long)] + output: Option, +} + +#[derive(Debug, Deserialize)] +struct SafetensorIndex { + #[serde(default)] + metadata: IndexMetadata, + weight_map: BTreeMap, +} + +#[derive(Debug, Default, Deserialize)] +struct IndexMetadata { + total_size: Option, +} + +#[derive(Clone, Debug, Deserialize, Serialize)] +struct TensorHeader { + dtype: String, + shape: Vec, + data_offsets: [u64; 2], +} + +#[derive(Debug)] +struct RemoteHeader { + header_len: u64, + tensors: BTreeMap, +} + +#[derive(Debug)] +struct CheckpointLayout { + index: SafetensorIndex, + index_bytes: u64, + headers: BTreeMap, +} + +#[derive(Debug)] +struct PreparedStage { + plan: StagePlan, + tensors: Vec, +} + +#[derive(Debug)] +struct SelectedTensor { + name: String, + source_file: String, + source_range: ByteRange, + header: TensorHeader, +} + +#[derive(Debug, Serialize)] +struct StagePlan { + repo: String, + revision: String, + layer_start: u32, + layer_end: u32, + include_prefixes: Vec, + total_model_tensor_bytes: Option, + index_bytes: u64, + selected_tensor_count: usize, + selected_tensor_bytes: u64, + largest_selected_tensor_bytes: u64, + source_shard_count: usize, + source_shard_bytes: u64, + range_request_count: usize, + range_payload_bytes: u64, + header_probe_bytes: u64, + planned_download_bytes: u64, + source_shard_bytes_avoided: u64, + full_model_tensor_bytes_avoided: Option, + shards: Vec, +} + +#[derive(Debug, Serialize)] +struct ShardPlan { + file: String, + file_bytes: u64, + header_probe_bytes: u64, + selected_tensor_count: usize, + selected_tensor_bytes: u64, + largest_selected_tensor_bytes: u64, + range_request_count: usize, + range_payload_bytes: u64, + ranges: Vec, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize)] +struct ByteRange { + start: u64, + end_exclusive: u64, +} + +impl ByteRange { + fn len(&self) -> u64 { + self.end_exclusive - self.start + } +} + +fn main() -> Result<()> { + let args = Args::parse(); + ensure!( + args.layer_start < args.layer_end, + "--layer-start must be less than --layer-end" + ); + let client = build_client()?; + let prepared = prepare_stage(&client, &args)?; + let plan = &prepared.plan; + if let Some(output) = &args.output { + materialize_stage(&client, &args, &prepared, output)?; + } + if args.json { + println!("{}", serde_json::to_string_pretty(&plan)?); + } else { + print_human_plan(plan); + } + Ok(()) +} + +fn build_client() -> Result { + Client::builder() + .connect_timeout(Duration::from_secs(15)) + .timeout(Duration::from_secs(60)) + .redirect(Policy::limited(10)) + .user_agent("mesh-llm-mlx-safetensors-stage-plan/0") + .build() + .context("build HTTP client") +} + +fn prepare_stage(client: &Client, args: &Args) -> Result { + let mut layout = load_checkpoint_layout(client, args)?; + let selected = select_tensors(&layout.index.weight_map, args); + ensure!( + !selected.is_empty(), + "no tensors matched layers {}..{} or the requested prefixes", + args.layer_start, + args.layer_end + ); + + let mut by_shard: BTreeMap<&str, BTreeSet<&str>> = BTreeMap::new(); + for name in &selected { + let shard = layout + .index + .weight_map + .get(*name) + .with_context(|| format!("selected tensor {name} is absent from the weight map"))?; + by_shard.entry(shard).or_default().insert(*name); + } + + let mut shards = Vec::with_capacity(by_shard.len()); + let mut tensors = Vec::with_capacity(selected.len()); + for (file, names) in by_shard { + if !layout.headers.contains_key(file) { + let url = resolve_url(&args.endpoint, &args.repo, &args.revision, file)?; + let header = + fetch_safetensor_header(client, &url).with_context(|| format!("inspect {file}"))?; + layout.headers.insert(file.to_string(), header); + } + let header = layout + .headers + .get(file) + .with_context(|| format!("missing inspected header for {file}"))?; + shards.push(plan_shard(file, header, &names, args.coalesce_gap_bytes)?); + tensors.extend(selected_tensors(file, header, &names)?); + } + let plan = summarize_plan(args, &layout.index, layout.index_bytes, shards); + Ok(PreparedStage { plan, tensors }) +} + +fn load_checkpoint_layout(client: &Client, args: &Args) -> Result { + let index_url = resolve_url( + &args.endpoint, + &args.repo, + &args.revision, + "model.safetensors.index.json", + )?; + if let Some(index_bytes) = fetch_optional_small_file(client, index_url, MAX_INDEX_BYTES)? { + let index: SafetensorIndex = + serde_json::from_slice(&index_bytes).context("parse model.safetensors.index.json")?; + return Ok(CheckpointLayout { + index, + index_bytes: index_bytes.len() as u64, + headers: BTreeMap::new(), + }); + } + + let file = "model.safetensors"; + let url = resolve_url(&args.endpoint, &args.repo, &args.revision, file)?; + let header = fetch_safetensor_header(client, &url) + .context("inspect unsharded model.safetensors checkpoint")?; + let total_size = header + .tensors + .values() + .map(|tensor| tensor.data_offsets[1]) + .max(); + let weight_map = header + .tensors + .keys() + .map(|name| (name.clone(), file.to_string())) + .collect(); + Ok(CheckpointLayout { + index: SafetensorIndex { + metadata: IndexMetadata { total_size }, + weight_map, + }, + index_bytes: 0, + headers: BTreeMap::from([(file.to_string(), header)]), + }) +} + +fn summarize_plan( + args: &Args, + index: &SafetensorIndex, + index_bytes: u64, + shards: Vec, +) -> StagePlan { + let selected_tensor_count = shards.iter().map(|shard| shard.selected_tensor_count).sum(); + let selected_tensor_bytes = shards.iter().map(|shard| shard.selected_tensor_bytes).sum(); + let largest_selected_tensor_bytes = shards + .iter() + .map(|shard| shard.largest_selected_tensor_bytes) + .max() + .unwrap_or(0); + let source_shard_bytes = shards.iter().map(|shard| shard.file_bytes).sum(); + let range_request_count = shards.iter().map(|shard| shard.range_request_count).sum(); + let range_payload_bytes = shards.iter().map(|shard| shard.range_payload_bytes).sum(); + let header_probe_bytes = shards.iter().map(|shard| shard.header_probe_bytes).sum(); + let planned_download_bytes = index_bytes + header_probe_bytes + range_payload_bytes; + StagePlan { + repo: args.repo.clone(), + revision: args.revision.clone(), + layer_start: args.layer_start, + layer_end: args.layer_end, + include_prefixes: args.include_prefixes.clone(), + total_model_tensor_bytes: index.metadata.total_size, + index_bytes, + selected_tensor_count, + selected_tensor_bytes, + largest_selected_tensor_bytes, + source_shard_count: shards.len(), + source_shard_bytes, + range_request_count, + range_payload_bytes, + header_probe_bytes, + planned_download_bytes, + source_shard_bytes_avoided: source_shard_bytes + .saturating_sub(header_probe_bytes + range_payload_bytes), + full_model_tensor_bytes_avoided: index + .metadata + .total_size + .map(|total| total.saturating_sub(selected_tensor_bytes)), + shards, + } +} + +fn select_tensors<'a>(weight_map: &'a BTreeMap, args: &Args) -> BTreeSet<&'a str> { + weight_map + .keys() + .filter(|name| { + layer_index(name) + .is_some_and(|layer| layer >= args.layer_start && layer < args.layer_end) + || args + .include_prefixes + .iter() + .any(|prefix| name.starts_with(prefix)) + }) + .map(String::as_str) + .collect() +} + +fn layer_index(name: &str) -> Option { + let parts = name.split('.').collect::>(); + parts.windows(2).find_map(|pair| { + matches!(pair[0], "layers" | "layer" | "h") + .then(|| pair[1].parse().ok()) + .flatten() + }) +} + +fn fetch_safetensor_header(client: &Client, url: &Url) -> Result { + let len_bytes = fetch_range(client, url.clone(), 0..8)?; + let header_len = u64::from_le_bytes( + len_bytes + .as_slice() + .try_into() + .map_err(|_| anyhow!("invalid 8-byte SafeTensors header length"))?, + ); + ensure!( + header_len <= MAX_HEADER_BYTES, + "SafeTensors header is unexpectedly large: {header_len} bytes" + ); + let header_end = 8_u64 + .checked_add(header_len) + .context("SafeTensors header range overflow")?; + let header_bytes = fetch_range(client, url.clone(), 8..header_end)?; + let raw: BTreeMap = + serde_json::from_slice(&header_bytes).context("parse SafeTensors header")?; + let tensors = raw + .into_iter() + .filter(|(name, _)| name != "__metadata__") + .map(|(name, value)| { + serde_json::from_value(value) + .map(|tensor| (name, tensor)) + .context("parse SafeTensors tensor header") + }) + .collect::>()?; + Ok(RemoteHeader { + header_len, + tensors, + }) +} + +fn plan_shard( + file: &str, + header: &RemoteHeader, + selected: &BTreeSet<&str>, + coalesce_gap_bytes: u64, +) -> Result { + let data_start = 8_u64 + .checked_add(header.header_len) + .context("SafeTensors data offset overflow")?; + let file_bytes = header + .tensors + .values() + .map(|tensor| tensor.data_offsets[1]) + .max() + .unwrap_or(0) + .checked_add(data_start) + .context("SafeTensors file length overflow")?; + let mut ranges = Vec::with_capacity(selected.len()); + let mut selected_tensor_bytes = 0_u64; + let mut largest_selected_tensor_bytes = 0_u64; + for name in selected { + let tensor = header + .tensors + .get(*name) + .with_context(|| format!("weight-map tensor {name} is absent from {file}"))?; + validate_tensor(name, tensor)?; + let start = data_start + .checked_add(tensor.data_offsets[0]) + .with_context(|| format!("absolute offset overflow for {name}"))?; + let end_exclusive = data_start + .checked_add(tensor.data_offsets[1]) + .with_context(|| format!("absolute offset overflow for {name}"))?; + let tensor_bytes = end_exclusive - start; + selected_tensor_bytes += tensor_bytes; + largest_selected_tensor_bytes = largest_selected_tensor_bytes.max(tensor_bytes); + ranges.push(ByteRange { + start, + end_exclusive, + }); + } + let ranges = coalesce_ranges(ranges, coalesce_gap_bytes); + let range_payload_bytes = ranges.iter().map(ByteRange::len).sum(); + Ok(ShardPlan { + file: file.to_string(), + file_bytes, + header_probe_bytes: data_start, + selected_tensor_count: selected.len(), + selected_tensor_bytes, + largest_selected_tensor_bytes, + range_request_count: ranges.len(), + range_payload_bytes, + ranges, + }) +} + +fn selected_tensors( + file: &str, + header: &RemoteHeader, + selected: &BTreeSet<&str>, +) -> Result> { + let data_start = 8_u64 + .checked_add(header.header_len) + .context("SafeTensors data offset overflow")?; + selected + .iter() + .map(|name| { + let tensor = header + .tensors + .get(*name) + .with_context(|| format!("weight-map tensor {name} is absent from {file}"))?; + validate_tensor(name, tensor)?; + let start = data_start + .checked_add(tensor.data_offsets[0]) + .with_context(|| format!("absolute offset overflow for {name}"))?; + let end_exclusive = data_start + .checked_add(tensor.data_offsets[1]) + .with_context(|| format!("absolute offset overflow for {name}"))?; + Ok(SelectedTensor { + name: (*name).to_string(), + source_file: file.to_string(), + source_range: ByteRange { + start, + end_exclusive, + }, + header: tensor.clone(), + }) + }) + .collect() +} + +fn validate_tensor(name: &str, tensor: &TensorHeader) -> Result<()> { + ensure!( + tensor.data_offsets[0] <= tensor.data_offsets[1], + "invalid data offsets for tensor {name}" + ); + ensure!(!tensor.dtype.is_empty(), "tensor {name} has no dtype"); + let _rank = tensor.shape.len(); + Ok(()) +} + +fn coalesce_ranges(mut ranges: Vec, max_gap: u64) -> Vec { + ranges.sort_by_key(|range| range.start); + let mut merged: Vec = Vec::with_capacity(ranges.len()); + for range in ranges { + if let Some(previous) = merged.last_mut() + && range.start.saturating_sub(previous.end_exclusive) <= max_gap + { + previous.end_exclusive = previous.end_exclusive.max(range.end_exclusive); + } else { + merged.push(range); + } + } + merged +} + +fn fetch_small_file(client: &Client, url: Url, max_bytes: u64) -> Result> { + let response = authorized(client.get(url)) + .send() + .context("send HTTP request")?; + let response = response + .error_for_status() + .context("download metadata file")?; + if let Some(length) = response + .headers() + .get(CONTENT_LENGTH) + .and_then(|value| value.to_str().ok()) + .and_then(|value| value.parse::().ok()) + { + ensure!( + length <= max_bytes, + "metadata file is too large: {length} bytes" + ); + } + let bytes = response.bytes().context("read metadata response")?; + ensure!( + bytes.len() as u64 <= max_bytes, + "metadata response exceeded {max_bytes} bytes" + ); + Ok(bytes.to_vec()) +} + +fn fetch_optional_small_file(client: &Client, url: Url, max_bytes: u64) -> Result>> { + let response = authorized(client.get(url)) + .send() + .context("send HTTP request")?; + if response.status() == StatusCode::NOT_FOUND { + return Ok(None); + } + read_small_response(response, max_bytes).map(Some) +} + +fn read_small_response(response: Response, max_bytes: u64) -> Result> { + let response = response + .error_for_status() + .context("download metadata file")?; + if let Some(length) = response + .headers() + .get(CONTENT_LENGTH) + .and_then(|value| value.to_str().ok()) + .and_then(|value| value.parse::().ok()) + { + ensure!( + length <= max_bytes, + "metadata file is too large: {length} bytes" + ); + } + let bytes = response.bytes().context("read metadata response")?; + ensure!( + bytes.len() as u64 <= max_bytes, + "metadata response exceeded {max_bytes} bytes" + ); + Ok(bytes.to_vec()) +} + +fn fetch_range(client: &Client, url: Url, range: Range) -> Result> { + ensure!(range.start < range.end, "HTTP byte range must not be empty"); + let header = format!("bytes={}-{}", range.start, range.end - 1); + let response = authorized(client.get(url).header(RANGE, header.clone())) + .send() + .with_context(|| format!("request HTTP range {header}"))?; + ensure_partial_content(&response, &header)?; + let bytes = response.bytes().context("read HTTP range response")?; + ensure!( + bytes.len() as u64 == range.end - range.start, + "HTTP range {header} returned {} bytes, expected {}", + bytes.len(), + range.end - range.start + ); + Ok(bytes.to_vec()) +} + +fn ensure_partial_content(response: &Response, range: &str) -> Result<()> { + ensure!( + response.status() == StatusCode::PARTIAL_CONTENT, + "server did not honor {range}; status was {} (refusing a possible full-shard download)", + response.status() + ); + Ok(()) +} + +fn materialize_stage( + client: &Client, + args: &Args, + prepared: &PreparedStage, + output: &Path, +) -> Result<()> { + fs::create_dir_all(output) + .with_context(|| format!("create output directory {}", output.display()))?; + let mut tensors = prepared.tensors.iter().collect::>(); + tensors.sort_by(|left, right| { + (&left.source_file, left.source_range.start) + .cmp(&(&right.source_file, right.source_range.start)) + }); + + let mut output_offset = 0_u64; + let mut output_header = BTreeMap::new(); + for tensor in &tensors { + let len = tensor.source_range.len(); + let end = output_offset + .checked_add(len) + .context("partial SafeTensors output offset overflow")?; + let mut header = tensor.header.clone(); + header.data_offsets = [output_offset, end]; + output_header.insert(tensor.name.clone(), header); + output_offset = end; + } + let mut header_bytes = serde_json::to_vec(&output_header)?; + while header_bytes.len() % 8 != 0 { + header_bytes.push(b' '); + } + let header_len = u64::try_from(header_bytes.len()).context("output header is too large")?; + + let destination = output.join("model.safetensors"); + let partial = output.join("model.safetensors.partial"); + let mut writer = BufWriter::new( + File::create(&partial).with_context(|| format!("create {}", partial.display()))?, + ); + writer.write_all(&header_len.to_le_bytes())?; + writer.write_all(&header_bytes)?; + + let spans = materialization_spans(&tensors); + let mut payload_bytes = 0_u64; + for span in &spans { + let url = resolve_url( + &args.endpoint, + &args.repo, + &args.revision, + &span.source_file, + )?; + payload_bytes += fetch_range_into(client, url, span.range.clone(), &mut writer)?; + } + writer.flush()?; + drop(writer); + fs::rename(&partial, &destination).with_context(|| { + format!( + "move completed partial SafeTensors file to {}", + destination.display() + ) + })?; + + let config_url = resolve_url(&args.endpoint, &args.repo, &args.revision, "config.json")?; + let config = fetch_small_file(client, config_url, MAX_INDEX_BYTES)?; + fs::write(output.join("config.json"), config)?; + fs::write( + output.join("stage-plan.json"), + serde_json::to_vec_pretty(&prepared.plan)?, + )?; + ensure!( + payload_bytes == prepared.plan.selected_tensor_bytes, + "materialized {payload_bytes} payload bytes, planned {}", + prepared.plan.selected_tensor_bytes + ); + eprintln!( + "materialized {} tensors ({} payload) in {} HTTP spans to {}", + tensors.len(), + human_bytes(payload_bytes), + spans.len(), + destination.display() + ); + Ok(()) +} + +#[derive(Debug)] +struct MaterializationSpan { + source_file: String, + range: Range, +} + +fn materialization_spans(tensors: &[&SelectedTensor]) -> Vec { + let mut spans: Vec = Vec::new(); + for tensor in tensors { + if let Some(previous) = spans.last_mut() + && previous.source_file == tensor.source_file + && previous.range.end == tensor.source_range.start + { + previous.range.end = tensor.source_range.end_exclusive; + } else { + spans.push(MaterializationSpan { + source_file: tensor.source_file.clone(), + range: tensor.source_range.start..tensor.source_range.end_exclusive, + }); + } + } + spans +} + +fn fetch_range_into( + client: &Client, + url: Url, + range: Range, + writer: &mut impl Write, +) -> Result { + ensure!(range.start < range.end, "HTTP byte range must not be empty"); + let expected = range.end - range.start; + let header = format!("bytes={}-{}", range.start, range.end - 1); + let mut response = authorized(client.get(url).header(RANGE, header.clone())) + .send() + .with_context(|| format!("request HTTP range {header}"))?; + ensure_partial_content(&response, &header)?; + let written = std::io::copy(&mut response, writer).context("stream HTTP tensor range")?; + ensure!( + written == expected, + "HTTP range {header} returned {written} bytes, expected {expected}" + ); + Ok(written) +} + +fn authorized(builder: reqwest::blocking::RequestBuilder) -> reqwest::blocking::RequestBuilder { + if let Some(token) = hf_token() { + builder.header(AUTHORIZATION, format!("Bearer {token}")) + } else { + builder + } +} + +fn hf_token() -> Option { + ["HF_TOKEN", "HUGGING_FACE_HUB_TOKEN"] + .iter() + .find_map(|name| env::var(name).ok()) + .map(|token| token.trim().to_string()) + .filter(|token| !token.is_empty()) +} + +fn resolve_url(endpoint: &str, repo: &str, revision: &str, file: &str) -> Result { + let mut url = Url::parse(endpoint).context("parse Hugging Face endpoint")?; + { + let mut segments = url + .path_segments_mut() + .map_err(|_| anyhow!("Hugging Face endpoint cannot be a base URL"))?; + segments.pop_if_empty(); + segments.extend(repo.split('/')); + segments.push("resolve"); + segments.push(revision); + segments.extend(file.split('/')); + } + Ok(url) +} + +fn print_human_plan(plan: &StagePlan) { + println!( + "{}@{} layers {}..{}", + plan.repo, plan.revision, plan.layer_start, plan.layer_end + ); + if let Some(total) = plan.total_model_tensor_bytes { + println!("full checkpoint tensors: {}", human_bytes(total)); + } + println!( + "selected: {} tensors, {}", + plan.selected_tensor_count, + human_bytes(plan.selected_tensor_bytes) + ); + println!( + "largest selected tensor: {}", + human_bytes(plan.largest_selected_tensor_bytes) + ); + println!( + "whole source shards: {} files, {}", + plan.source_shard_count, + human_bytes(plan.source_shard_bytes) + ); + println!( + "exact ranged payload: {} requests, {}", + plan.range_request_count, + human_bytes(plan.range_payload_bytes) + ); + println!( + "planned download including index/headers: {}", + human_bytes(plan.planned_download_bytes) + ); + println!( + "avoided versus whole shards: {}", + human_bytes(plan.source_shard_bytes_avoided) + ); +} + +fn human_bytes(bytes: u64) -> String { + const UNITS: [&str; 5] = ["B", "KiB", "MiB", "GiB", "TiB"]; + let mut value = bytes as f64; + let mut unit = 0; + while value >= 1024.0 && unit + 1 < UNITS.len() { + value /= 1024.0; + unit += 1; + } + format!("{value:.2} {}", UNITS[unit]) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn recognizes_common_transformer_layer_paths() { + assert_eq!(layer_index("model.layers.42.mlp.up_proj.weight"), Some(42)); + assert_eq!(layer_index("transformer.h.7.attn.weight"), Some(7)); + assert_eq!(layer_index("model.embed_tokens.weight"), None); + } + + #[test] + fn coalesces_adjacent_and_small_gap_ranges() { + let ranges = vec![ + ByteRange { + start: 20, + end_exclusive: 30, + }, + ByteRange { + start: 0, + end_exclusive: 10, + }, + ByteRange { + start: 10, + end_exclusive: 20, + }, + ByteRange { + start: 32, + end_exclusive: 40, + }, + ]; + assert_eq!( + coalesce_ranges(ranges, 2), + vec![ByteRange { + start: 0, + end_exclusive: 40 + }] + ); + } + + #[test] + fn plans_only_selected_tensor_payloads() { + let header = RemoteHeader { + header_len: 100, + tensors: BTreeMap::from([ + ( + "model.layers.1.weight".into(), + tensor("BF16", &[2, 2], [0, 8]), + ), + ( + "model.layers.2.weight".into(), + tensor("BF16", &[2, 2], [8, 16]), + ), + ( + "model.layers.3.weight".into(), + tensor("BF16", &[2, 2], [16, 24]), + ), + ]), + }; + let selected = BTreeSet::from(["model.layers.2.weight"]); + let plan = plan_shard("model.safetensors", &header, &selected, 0).unwrap(); + assert_eq!(plan.file_bytes, 132); + assert_eq!(plan.selected_tensor_bytes, 8); + assert_eq!(plan.largest_selected_tensor_bytes, 8); + assert_eq!( + plan.ranges, + vec![ByteRange { + start: 116, + end_exclusive: 124 + }] + ); + } + + #[test] + fn materialization_spans_merge_only_contiguous_ranges_in_one_source() { + let tensors = [ + selected_tensor("a", "one.safetensors", 10, 20), + selected_tensor("b", "one.safetensors", 20, 30), + selected_tensor("c", "one.safetensors", 40, 50), + selected_tensor("d", "two.safetensors", 50, 60), + ]; + let references = tensors.iter().collect::>(); + let spans = materialization_spans(&references); + + assert_eq!(spans.len(), 3); + assert_eq!(spans[0].source_file, "one.safetensors"); + assert_eq!(spans[0].range, 10..30); + assert_eq!(spans[1].range, 40..50); + assert_eq!(spans[2].source_file, "two.safetensors"); + assert_eq!(spans[2].range, 50..60); + } + + fn selected_tensor(name: &str, file: &str, start: u64, end: u64) -> SelectedTensor { + SelectedTensor { + name: name.to_string(), + source_file: file.to_string(), + source_range: ByteRange { + start, + end_exclusive: end, + }, + header: tensor("BF16", &[1], [0, end - start]), + } + } + + fn tensor(dtype: &str, shape: &[u64], data_offsets: [u64; 2]) -> TensorHeader { + TensorHeader { + dtype: dtype.into(), + shape: shape.to_vec(), + data_offsets, + } + } +} diff --git a/spikes/mlx-solo/Cargo.lock b/spikes/mlx-solo/Cargo.lock index e9e35f0511..5d7fd4dcd4 100644 --- a/spikes/mlx-solo/Cargo.lock +++ b/spikes/mlx-solo/Cargo.lock @@ -111,6 +111,12 @@ version = "1.25.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d6aedf8ae72766347502cf3cb4f41cf5e9cc37d28bee90f1fdaaae15f9cf9424" +[[package]] +name = "bytes" +version = "1.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc652a48c352aef3ea3aed32080501cf3ef6ed5da78602a020c991775b0aff04" + [[package]] name = "castaway" version = "0.2.4" @@ -412,12 +418,24 @@ version = "0.1.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5baebc0774151f905a1a2cc41989300b1e6fbb29aff0ceffa1064fdd3088d582" +[[package]] +name = "fixedbitset" +version = "0.5.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1d674e81391d1e1ab681a28d99df07927c6d4aa5b027d7da16ba32d1d21ecd99" + [[package]] name = "fnv" version = "1.0.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3f9eec918d3f24069decb9af1554cad7c880e2da24a9afd88aca000531ab82c1" +[[package]] +name = "foldhash" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d9c4f5dac5e15c24eb999c26181a6ca40b39fe946cbe4c263c7209467bc83af2" + [[package]] name = "foldhash" version = "0.2.0" @@ -458,6 +476,15 @@ dependencies = [ "zerocopy", ] +[[package]] +name = "hashbrown" +version = "0.15.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9229cfe53dfd69f0609a49f65461bd93001ea1ef889cd5529dd176593f5338a1" +dependencies = [ + "foldhash 0.1.5", +] + [[package]] name = "hashbrown" version = "0.16.1" @@ -466,7 +493,7 @@ checksum = "841d1cc9bed7f9236f321df977030373f4a4163ae1a7dbfe1a51a2c1a51d9100" dependencies = [ "allocator-api2", "equivalent", - "foldhash", + "foldhash 0.2.0", "serde", "serde_core", ] @@ -709,6 +736,7 @@ dependencies = [ "safemlx", "safemlx-lm", "serde_json", + "skippy-protocol", ] [[package]] @@ -733,6 +761,12 @@ dependencies = [ "syn", ] +[[package]] +name = "multimap" +version = "0.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1d87ecb2933e8aeadb3e3a02b828fed80a7528047e68b4f424523a0981a3a084" + [[package]] name = "nom" version = "7.1.3" @@ -846,6 +880,17 @@ version = "1.0.15" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "57c0d7b74b563b49d38dae00a0c37d4d6de9b432382b2892f0574ddcae73fd0a" +[[package]] +name = "petgraph" +version = "0.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8701b58ea97060d5e5b155d383a69952a60943f0e6dfe30b04c287beb0b27455" +dependencies = [ + "fixedbitset", + "hashbrown 0.15.5", + "indexmap", +] + [[package]] name = "pkg-config" version = "0.3.33" @@ -870,6 +915,16 @@ dependencies = [ "zerocopy", ] +[[package]] +name = "prettyplease" +version = "0.2.37" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "479ca8adacdd7ce8f1fb39ce9ecccbfe93a3f1344b3d0d97f20bc0196208f62b" +dependencies = [ + "proc-macro2", + "syn", +] + [[package]] name = "proc-macro-crate" version = "3.5.0" @@ -888,6 +943,121 @@ dependencies = [ "unicode-ident", ] +[[package]] +name = "prost" +version = "0.14.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "528ac67416ff8646872a3c02cad9cc4ee5dc9f9540c9b10771855c95cb2e5ae1" +dependencies = [ + "bytes", + "prost-derive", +] + +[[package]] +name = "prost-build" +version = "0.14.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "03da047801ff44bb6a4d407d4860c05fd70bb81714e6b2f3812603d5b145b042" +dependencies = [ + "heck", + "itertools 0.14.0", + "log", + "multimap", + "petgraph", + "prettyplease", + "prost", + "prost-types", + "regex", + "syn", + "tempfile", +] + +[[package]] +name = "prost-derive" +version = "0.14.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b570b25f7617e43d59005d0990ccb79e950a423952cea19671b7a876da390adf" +dependencies = [ + "anyhow", + "itertools 0.14.0", + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "prost-types" +version = "0.14.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f94967dc7688f3054c7fac87473ffae4cc4c3904800e2d9f5b857246d8963b0a" +dependencies = [ + "prost", +] + +[[package]] +name = "protoc-bin-vendored" +version = "3.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d1c381df33c98266b5f08186583660090a4ffa0889e76c7e9a5e175f645a67fa" +dependencies = [ + "protoc-bin-vendored-linux-aarch_64", + "protoc-bin-vendored-linux-ppcle_64", + "protoc-bin-vendored-linux-s390_64", + "protoc-bin-vendored-linux-x86_32", + "protoc-bin-vendored-linux-x86_64", + "protoc-bin-vendored-macos-aarch_64", + "protoc-bin-vendored-macos-x86_64", + "protoc-bin-vendored-win32", +] + +[[package]] +name = "protoc-bin-vendored-linux-aarch_64" +version = "3.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c350df4d49b5b9e3ca79f7e646fde2377b199e13cfa87320308397e1f37e1a4c" + +[[package]] +name = "protoc-bin-vendored-linux-ppcle_64" +version = "3.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a55a63e6c7244f19b5c6393f025017eb5d793fd5467823a099740a7a4222440c" + +[[package]] +name = "protoc-bin-vendored-linux-s390_64" +version = "3.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1dba5565db4288e935d5330a07c264a4ee8e4a5b4a4e6f4e83fad824cc32f3b0" + +[[package]] +name = "protoc-bin-vendored-linux-x86_32" +version = "3.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8854774b24ee28b7868cd71dccaae8e02a2365e67a4a87a6cd11ee6cdbdf9cf5" + +[[package]] +name = "protoc-bin-vendored-linux-x86_64" +version = "3.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b38b07546580df720fa464ce124c4b03630a6fb83e05c336fea2a241df7e5d78" + +[[package]] +name = "protoc-bin-vendored-macos-aarch_64" +version = "3.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "89278a9926ce312e51f1d999fee8825d324d603213344a9a706daa009f1d8092" + +[[package]] +name = "protoc-bin-vendored-macos-x86_64" +version = "3.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "81745feda7ccfb9471d7a4de888f0652e806d5795b61480605d4943176299756" + +[[package]] +name = "protoc-bin-vendored-win32" +version = "3.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "95067976aca6421a523e491fce939a3e65249bac4b977adee0ee9771568e8aa3" + [[package]] name = "quote" version = "1.0.46" @@ -1182,6 +1352,16 @@ version = "2.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f8fadd59c855ef2080decdef8ff161eb6661b86933c9d82e5ba29dc602a55aba" +[[package]] +name = "skippy-protocol" +version = "0.72.1" +dependencies = [ + "prost", + "prost-build", + "protoc-bin-vendored", + "serde", +] + [[package]] name = "smallvec" version = "1.15.2" diff --git a/spikes/mlx-solo/Cargo.toml b/spikes/mlx-solo/Cargo.toml index cc574dcb30..da00a02240 100644 --- a/spikes/mlx-solo/Cargo.toml +++ b/spikes/mlx-solo/Cargo.toml @@ -15,6 +15,10 @@ publish = false name = "mlx-solo" path = "src/main.rs" +[[bin]] +name = "mlx-split-proof" +path = "src/bin/mlx-split-proof.rs" + [workspace] [dependencies] @@ -29,3 +33,4 @@ safemlx-lm = { path = "../../../safemlx/safemlx-lm" } anyhow = "1" clap = { version = "4", features = ["derive"] } serde_json = "1" +skippy-protocol = { path = "../../crates/skippy-protocol" } diff --git a/spikes/mlx-solo/FINDINGS.md b/spikes/mlx-solo/FINDINGS.md index 8af8dcad85..379b63b949 100644 --- a/spikes/mlx-solo/FINDINGS.md +++ b/spikes/mlx-solo/FINDINGS.md @@ -144,6 +144,9 @@ De-risked: Still open: - Larger models + throughput vs the llama.cpp backend on the same hardware. -- Everything staged/split — this spike is single-stage, whole-model. The - partial-load and boundary-fence go/no-go spikes (plan §8) are unchanged. +- Product staged/split integration. A follow-on SmolLM2 proof now range- + materializes two partial SafeTensors files and matches whole-model MLX through + Skippy's F16/F32 binary activation codec; see + `../mlx-safetensors-stages/FINDINGS.md`. Two `skippy-server` processes, the + boundary-fence benchmark, and bounded-memory load-time quantization remain. - Upstreaming the two fixes to `jbg/safemlx` (or carrying a thin fork). diff --git a/spikes/mlx-solo/src/bin/mlx-split-proof.rs b/spikes/mlx-solo/src/bin/mlx-split-proof.rs new file mode 100644 index 0000000000..975c6e0f2b --- /dev/null +++ b/spikes/mlx-solo/src/bin/mlx-split-proof.rs @@ -0,0 +1,465 @@ +//! Proves that two MLX stages loaded from disjoint partial SafeTensors files +//! reproduce whole-model greedy decode without a complete checkpoint file. + +use std::io::Cursor; +use std::path::{Path, PathBuf}; + +use anyhow::{bail, ensure, Context, Result}; +use clap::{Parser, ValueEnum}; +use safemlx::module::{Module, ModuleParameters, ModuleParametersExt}; +use safemlx::ops::indexing::{NewAxis, TryIndexOp}; +use safemlx::{arange, Array, Device, DeviceType, Dtype, Stream}; +use safemlx_lm::cache::{ConcatKeyValueCache, KeyValueCache}; +use safemlx_lm::models::common::linear::project_logits_maybe_quantized; +use safemlx_lm::models::llama::{self, AttentionInput, TransformerBlock}; +use safemlx_lm::weights::{ + load_safetensors_lenient, load_safetensors_strict, StrictLoadConfig, StrictLoadReport, +}; +use skippy_protocol::binary::{ + encode_f32_activation_payload, read_stage_message, write_stage_message, StageStateHeader, + StageWireMessage, WireActivationDType, WireMessageKind, +}; + +#[derive(Debug, Parser)] +#[command(about = "Compare whole-model MLX with two partial SafeTensors stages")] +struct Args { + #[arg(long)] + stage0: PathBuf, + + #[arg(long)] + stage1: PathBuf, + + /// First layer owned by stage 1. + #[arg(long, default_value_t = 15)] + split: usize, + + /// Comma-separated prompt token ids. The mesh sends token ids to stage 0. + #[arg(long, default_value = "1,1531,314,260,3575,28")] + tokens: String, + + /// Number of greedy decode steps to compare after prompt prefill. + #[arg(long, default_value_t = 8)] + steps: usize, + + /// Activation encoding used at the artificial mesh boundary. + #[arg(long, value_enum, default_value_t = WireDtype::F16)] + wire_dtype: WireDtype, +} + +#[derive(Clone, Copy, Debug, ValueEnum)] +enum WireDtype { + F16, + F32, +} + +struct Models { + baseline: llama::Model, + stage0: llama::Model, + stage1: llama::Model, +} + +struct Caches { + baseline: Vec>, + stage0: Vec>, + stage1: Vec>, +} + +impl Caches { + fn new() -> Self { + Self { + baseline: Vec::new(), + stage0: Vec::new(), + stage1: Vec::new(), + } + } +} + +fn main() -> Result<()> { + let args = Args::parse(); + ensure!(args.steps > 0, "--steps must be positive"); + let prompt = parse_tokens(&args.tokens)?; + + let stream = Stream::new_with_device(&Device::new(DeviceType::Gpu, 0)); + let weights_stream = Stream::new_with_device(&Device::new(DeviceType::Cpu, 0)); + let mut models = load_models(&args, &stream, &weights_stream)?; + stream.synchronize()?; + + let mut caches = Caches::new(); + let mut input = prompt; + let mut generated = Vec::with_capacity(args.steps); + let mut worst_logit_delta = 0.0_f32; + let mut boundary_bytes = 0_usize; + + for step in 0..args.steps { + let tokens = Array::from_slice(&input, &[1, input.len() as i32]); + let (baseline_logits, baseline_boundary, baseline_final, baseline_normed) = + baseline_forward( + &mut models.baseline, + &tokens, + &mut caches.baseline, + args.split, + &stream, + )?; + let (staged_logits, staged_boundary, staged_final, staged_normed, bytes) = staged_forward( + &mut models, + &tokens, + &mut caches, + args.split, + args.wire_dtype, + &stream, + )?; + boundary_bytes += bytes; + let boundary_delta = array_max_abs_delta(&baseline_boundary, &staged_boundary, &stream)?; + let final_delta = array_max_abs_delta(&baseline_final, &staged_final, &stream)?; + let norm_delta = array_max_abs_delta(&baseline_normed, &staged_normed, &stream)?; + + let baseline = last_logits(&baseline_logits, &stream)?; + let staged = last_logits(&staged_logits, &stream)?; + ensure!(baseline.len() == staged.len(), "logit width mismatch"); + let delta = baseline + .iter() + .zip(&staged) + .map(|(left, right)| (left - right).abs()) + .fold(0.0_f32, f32::max); + worst_logit_delta = worst_logit_delta.max(delta); + let baseline_token = argmax(&baseline)?; + let staged_token = argmax(&staged)?; + ensure!( + baseline_token == staged_token, + "greedy token diverged at step {step}: baseline={baseline_token}, staged={staged_token}, boundary_delta={boundary_delta}, final_delta={final_delta}, norm_delta={norm_delta}, max_abs_logit_delta={delta}" + ); + generated.push(staged_token); + input = vec![staged_token]; + eprintln!( + "step={step} token={staged_token} boundary_delta={boundary_delta:.6} max_abs_logit_delta={delta:.6} stage_wire_bytes={bytes}" + ); + } + + println!("PASS: whole-model and two-stage greedy tokens match"); + println!("wire_dtype={:?}", args.wire_dtype); + println!("split=0..{} | {}..30", args.split, args.split); + println!("generated_tokens={generated:?}"); + println!("worst_max_abs_logit_delta={worst_logit_delta:.6}"); + println!("stage_wire_bytes={boundary_bytes}"); + Ok(()) +} + +fn parse_tokens(value: &str) -> Result> { + let tokens = value + .split(',') + .map(|token| token.trim().parse::().context("parse token id")) + .collect::>>()?; + ensure!(!tokens.is_empty(), "--tokens must not be empty"); + Ok(tokens) +} + +fn load_models(args: &Args, stream: &Stream, weights_stream: &Stream) -> Result { + let model_args = llama::get_llama_model_args(&args.stage0)?; + let total_layers = usize::try_from(model_args.num_hidden_layers)?; + ensure!( + args.split > 0 && args.split < total_layers, + "--split must be inside 0..{total_layers}" + ); + let stage0_file = weight_file(&args.stage0); + let stage1_file = weight_file(&args.stage1); + + let mut baseline = llama::Model::new(model_args.clone(), stream)?; + let strict = StrictLoadConfig::default(); + let mut report = StrictLoadReport::default(); + load_safetensors_strict( + &mut baseline, + &stage0_file, + weights_stream, + &strict, + &mut report, + )?; + load_safetensors_strict( + &mut baseline, + &stage1_file, + weights_stream, + &strict, + &mut report, + )?; + report.finish(&baseline, &strict)?; + baseline.copy_to_stream(stream)?; + + let mut stage0 = llama::Model::new(model_args.clone(), stream)?; + load_safetensors_lenient(&mut stage0, &stage0_file, weights_stream)?; + stage0.model.layers.truncate(args.split); + stage0.model.num_hidden_layers = i32::try_from(args.split)?; + stage0.copy_to_stream(stream)?; + + let mut stage1 = llama::Model::new(model_args, stream)?; + load_safetensors_lenient(&mut stage1, &stage1_file, weights_stream)?; + for block in &mut stage1.model.layers[args.split..] { + block.copy_to_stream(stream)?; + } + stage1.model.norm.copy_to_stream(stream)?; + stage1.model.embed_tokens.copy_to_stream(stream)?; + if let Some(lm_head) = &mut stage1.lm_head { + lm_head.copy_to_stream(stream)?; + } + verify_layer_weights(&baseline, &stage1, args.split, stream)?; + + eprintln!( + "loaded baseline from union of partial files; stage layers={}+{}", + stage0.model.layers.len(), + stage1.model.layers.len() - args.split + ); + Ok(Models { + baseline, + stage0, + stage1, + }) +} + +fn verify_layer_weights( + baseline: &llama::Model, + stage1: &llama::Model, + layer: usize, + stream: &Stream, +) -> Result<()> { + let baseline_parameters = baseline.parameters().flatten(); + let stage_parameters = stage1.parameters().flatten(); + let prefix = format!("model.layers.{layer}."); + let mut checked = 0_usize; + for (name, baseline_value) in baseline_parameters + .iter() + .filter(|(name, _)| name.starts_with(&prefix)) + { + let stage_value = stage_parameters + .get(name) + .with_context(|| format!("stage 1 parameter {name} is absent"))?; + let delta = array_max_abs_delta(baseline_value, stage_value, stream)?; + ensure!(delta == 0.0, "stage 1 parameter {name} differs by {delta}"); + checked += 1; + } + ensure!(checked > 0, "no parameters found for {prefix}"); + eprintln!("verified {checked} layer-{layer} stage parameters against baseline"); + Ok(()) +} + +fn weight_file(directory: &Path) -> PathBuf { + directory.join("model.safetensors") +} + +fn baseline_forward( + model: &mut llama::Model, + tokens: &Array, + cache: &mut Vec>, + split: usize, + stream: &Stream, +) -> Result<(Array, Array, Array, Array)> { + let hidden = model.model.embed_tokens.forward(tokens, stream)?; + let mask = attention_mask(&hidden, cache, stream)?; + if cache.is_empty() { + *cache = (0..model.model.layers.len()) + .map(|_| Some(ConcatKeyValueCache::default())) + .collect(); + } + let (first_layers, last_layers) = model.model.layers.split_at_mut(split); + let (first_cache, last_cache) = cache.split_at_mut(split); + let hidden = forward_block_slice(first_layers, hidden, mask.as_ref(), first_cache, stream)?; + let boundary = hidden.clone(); + let hidden = forward_block_slice(last_layers, hidden, mask.as_ref(), last_cache, stream)?; + let final_hidden = hidden.clone(); + let hidden = model.model.norm.forward(&hidden, stream)?; + let normed = hidden.clone(); + let logits = project_logits_maybe_quantized( + &mut model.lm_head, + &mut model.model.embed_tokens, + &hidden, + stream, + )?; + Ok((logits, boundary, final_hidden, normed)) +} + +fn staged_forward( + models: &mut Models, + tokens: &Array, + caches: &mut Caches, + split: usize, + wire_dtype: WireDtype, + stream: &Stream, +) -> Result<(Array, Array, Array, Array, usize)> { + let mut hidden = models.stage0.model.embed_tokens.forward(tokens, stream)?; + let mask = attention_mask(&hidden, &caches.stage0, stream)?; + hidden = forward_blocks( + &mut models.stage0.model.layers, + hidden, + mask.as_ref(), + &mut caches.stage0, + stream, + )?; + let stage0_boundary = hidden.clone(); + let (hidden, boundary_bytes) = wire_roundtrip(&hidden, wire_dtype, stream)?; + let hidden = forward_blocks( + &mut models.stage1.model.layers[split..], + hidden, + mask.as_ref(), + &mut caches.stage1, + stream, + )?; + let final_hidden = hidden.clone(); + let hidden = models.stage1.model.norm.forward(&hidden, stream)?; + let normed = hidden.clone(); + let logits = project_logits_maybe_quantized( + &mut models.stage1.lm_head, + &mut models.stage1.model.embed_tokens, + &hidden, + stream, + )?; + Ok(( + logits, + stage0_boundary, + final_hidden, + normed, + boundary_bytes, + )) +} + +fn attention_mask( + hidden: &Array, + cache: &[Option], + stream: &Stream, +) -> Result> { + let sequence = hidden.shape()[1]; + if sequence == 1 { + return Ok(None); + } + let offset = cache + .first() + .and_then(Option::as_ref) + .map_or(0, KeyValueCache::offset); + let right = arange!(stop = offset + sequence, stream = stream)?; + let left = arange!(start = offset, stop = offset + sequence, stream = stream)?; + let left = left.try_index_device((.., NewAxis), stream)?; + let right = right.try_index_device(NewAxis, stream)?; + Ok(Some(left.ge(&right, stream)?)) +} + +fn forward_blocks( + blocks: &mut [TransformerBlock], + hidden: Array, + mask: Option<&Array>, + cache: &mut Vec>, + stream: &Stream, +) -> Result { + if cache.is_empty() { + *cache = (0..blocks.len()) + .map(|_| Some(ConcatKeyValueCache::default())) + .collect(); + } + ensure!(cache.len() == blocks.len(), "stage cache length mismatch"); + forward_block_slice(blocks, hidden, mask, cache, stream) +} + +fn forward_block_slice( + blocks: &mut [TransformerBlock], + mut hidden: Array, + mask: Option<&Array>, + cache: &mut [Option], + stream: &Stream, +) -> Result { + ensure!(cache.len() == blocks.len(), "stage cache length mismatch"); + for (block, layer_cache) in blocks.iter_mut().zip(cache.iter_mut()) { + hidden = block.forward( + AttentionInput { + x: &hidden, + mask, + cache: layer_cache.as_mut(), + generated_sliding_window: None, + }, + stream, + )?; + } + Ok(hidden) +} + +fn wire_roundtrip( + hidden: &Array, + wire_dtype: WireDtype, + stream: &Stream, +) -> Result<(Array, usize)> { + let shape = hidden.shape().to_vec(); + let compute_dtype = hidden.dtype(); + let token_count = shape[1]; + let hidden_width = shape[2]; + let f32_values = hidden + .as_dtype(Dtype::Float32, stream)? + .evaluated()? + .as_slice::() + .to_vec(); + let f32_bytes = f32_values + .iter() + .flat_map(|value| value.to_le_bytes()) + .collect::>(); + let protocol_dtype = match wire_dtype { + WireDtype::F16 => WireActivationDType::F16, + WireDtype::F32 => WireActivationDType::F32, + }; + let activation = + encode_f32_activation_payload(protocol_dtype, token_count, hidden_width, &f32_bytes)?; + let kind = if token_count > 1 { + WireMessageKind::PrefillEmbd + } else { + WireMessageKind::DecodeEmbd + }; + let mut state = StageStateHeader::new(kind, protocol_dtype); + state.source_stage_index = 0; + let message = StageWireMessage { + kind, + pos_start: 0, + token_count, + state, + request_id: 1, + session_id: 1, + sampling: None, + chat_sampling_metadata: None, + tokens: Vec::new(), + positions: Vec::new(), + activation, + raw_bytes: Vec::new(), + }; + let mut frame = Vec::new(); + write_stage_message(&mut frame, &message, protocol_dtype)?; + let decoded = read_stage_message(Cursor::new(&frame), hidden_width)?; + let decoded_f32 = decoded.activation_f32_payload(hidden_width)?; + let values = decoded_f32 + .chunks_exact(4) + .map(|chunk| f32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]])) + .collect::>(); + let restored = Array::from_slice(&values, &shape).as_dtype(compute_dtype, stream)?; + Ok((restored, frame.len())) +} + +fn last_logits(logits: &Array, stream: &Stream) -> Result> { + let row = logits + .try_index_device((0, -1, ..), stream)? + .as_dtype(Dtype::Float32, stream)?; + Ok(row.evaluated()?.as_slice::().to_vec()) +} + +fn array_max_abs_delta(left: &Array, right: &Array, stream: &Stream) -> Result { + ensure!(left.shape() == right.shape(), "activation shape mismatch"); + let left = left.as_dtype(Dtype::Float32, stream)?; + let right = right.as_dtype(Dtype::Float32, stream)?; + Ok(left + .evaluated()? + .as_slice::() + .iter() + .zip(right.evaluated()?.as_slice::()) + .map(|(left, right)| (left - right).abs()) + .fold(0.0_f32, f32::max)) +} + +fn argmax(values: &[f32]) -> Result { + let Some((index, _)) = values + .iter() + .enumerate() + .max_by(|(_, left), (_, right)| left.total_cmp(right)) + else { + bail!("cannot argmax empty logits") + }; + Ok(u32::try_from(index)?) +} From 613ddaa43e6d59f24a6041100772d583cc1d21d1 Mon Sep 17 00:00:00 2001 From: Michael Neale <14976+michaelneale@users.noreply.github.com> Date: Fri, 17 Jul 2026 16:02:45 +1000 Subject: [PATCH 09/37] feat(mlx): serve partial safetensors stages over skippy wire --- Cargo.lock | 12 + Cargo.toml | 1 + Justfile | 8 + crates/skippy-engine-mlx/Cargo.toml | 7 + crates/skippy-engine-mlx/STAGED_EXECUTION.md | 98 +++++ crates/skippy-engine-mlx/src/backend.rs | 6 +- crates/skippy-engine-mlx/src/bin/mlx-stage.rs | 265 ++++++++++++ crates/skippy-engine-mlx/src/engine.rs | 6 +- crates/skippy-engine-mlx/src/lib.rs | 4 + crates/skippy-engine-mlx/src/stage.rs | 392 ++++++++++++++++++ crates/skippy-engine/Cargo.toml | 16 + crates/skippy-engine/README.md | 12 + crates/skippy-engine/src/lib.rs | 178 ++++++++ crates/skippy-server/Cargo.toml | 1 + crates/skippy-server/src/engine_transport.rs | 360 ++++++++++++++++ crates/skippy-server/src/lib.rs | 1 + docs/design/MLX_STAGE_ENGINE_PLAN.md | 25 +- scripts/affected-crates.sh | 1 + scripts/plan-clippy-batches.sh | 1 + scripts/publish-crates.sh | 6 + 20 files changed, 1390 insertions(+), 10 deletions(-) create mode 100644 crates/skippy-engine-mlx/STAGED_EXECUTION.md create mode 100644 crates/skippy-engine-mlx/src/bin/mlx-stage.rs create mode 100644 crates/skippy-engine-mlx/src/stage.rs create mode 100644 crates/skippy-engine/Cargo.toml create mode 100644 crates/skippy-engine/README.md create mode 100644 crates/skippy-engine/src/lib.rs create mode 100644 crates/skippy-server/src/engine_transport.rs diff --git a/Cargo.lock b/Cargo.lock index 1b9afdd8dd..64e5a0ac95 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -7738,6 +7738,14 @@ dependencies = [ "skippy-runtime", ] +[[package]] +name = "skippy-engine" +version = "0.72.1" +dependencies = [ + "anyhow", + "skippy-protocol", +] + [[package]] name = "skippy-engine-mlx" version = "0.72.1" @@ -7752,6 +7760,9 @@ dependencies = [ "safemlx", "safemlx-lm", "serde_json", + "skippy-engine", + "skippy-protocol", + "skippy-server", "tokenizers", "tokio", "tracing", @@ -7859,6 +7870,7 @@ dependencies = [ "serde_json", "sha2 0.10.9", "skippy-cache", + "skippy-engine", "skippy-metrics", "skippy-protocol", "skippy-runtime", diff --git a/Cargo.toml b/Cargo.toml index d1ea11ddd3..8f5fedd0a4 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -37,6 +37,7 @@ members = [ "crates/model-hf", "crates/model-resolver", "crates/skippy-protocol", + "crates/skippy-engine", "crates/skippy-coordinator", "crates/skippy-topology", "crates/skippy-cache", diff --git a/Justfile b/Justfile index b415e5d30b..87e25d6bab 100644 --- a/Justfile +++ b/Justfile @@ -251,6 +251,14 @@ mlx-safetensors-stage-plan *ARGS: 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 production-shaped MLX stage server/client over Skippy's binary wire. +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 + +# 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 }} diff --git a/crates/skippy-engine-mlx/Cargo.toml b/crates/skippy-engine-mlx/Cargo.toml index d9509339ac..5b1feecdb6 100644 --- a/crates/skippy-engine-mlx/Cargo.toml +++ b/crates/skippy-engine-mlx/Cargo.toml @@ -28,6 +28,10 @@ path = "src/lib.rs" name = "mlx-serve" path = "src/bin/mlx-serve.rs" +[[bin]] +name = "mlx-stage" +path = "src/bin/mlx-stage.rs" + [features] default = [] # Enable the real MLX engine. macOS-only in practice (code is cfg-gated to @@ -38,6 +42,9 @@ mlx = ["dep:safemlx", "dep:safemlx-lm", "dep:tokenizers"] # The real mesh-llm OpenAI frontend — this is what proves we serve over the # same surface the shipped binary uses, not a toy. openai-frontend = { path = "../openai-frontend" } +skippy-engine = { path = "../skippy-engine" } +skippy-protocol = { path = "../skippy-protocol" } +skippy-server = { path = "../skippy-server" } async-trait = "0.1" async-stream = "0.3" anyhow = "1" diff --git a/crates/skippy-engine-mlx/STAGED_EXECUTION.md b/crates/skippy-engine-mlx/STAGED_EXECUTION.md new file mode 100644 index 0000000000..fbadeebcf0 --- /dev/null +++ b/crates/skippy-engine-mlx/STAGED_EXECUTION.md @@ -0,0 +1,98 @@ +# MLX partial-layer staged execution + +## Status + +Dense Llama-family MLX stages now run as separate OS processes from partial +SafeTensors artifacts and communicate over Skippy's existing binary stage wire. +The first proof uses `HuggingFaceTB/SmolLM2-135M-Instruct` split at layer 15. + +This is a production-shaped bridge, not yet the default mesh launch path: + +- `skippy-engine` owns the engine-neutral `StageEngine` contract and residual + buffer descriptors. +- `skippy-server::engine_transport` serves that contract using the existing + `StageWireMessage`, ready handshake, activation codec, and reply codec. +- `MlxStageEngine` loads one materialized partial SafeTensors file, owns + per-session KV caches on a dedicated MLX worker thread, and executes only its + configured layer range. +- `mlx-stage` starts a stage process or drives a chain as a proof client. + +No process in the proof has access to the complete checkpoint. The tokenizer +and config files are small shared metadata; tensor data comes only from that +process's `model.safetensors`. + +## Verified result + +On Apple Silicon Metal, using two materialized 155.28 MiB partial files: + +| Process | Layers | Tensor file available | RSS after the proof | +| --- | ---: | ---: | ---: | +| stage 0 | `0..15` | 155.28 MiB | 188,784 KiB | +| stage 1 | `15..30` | 155.28 MiB | 189,168 KiB | + +The processes exchanged F16 residual activations and generated: + +```text +[284, 260, 2240, 314, 1343, 327, 624, 8685] +``` + +That exactly matches the whole-model and in-process split reference for the +same prompt across prompt prefill and seven subsequent decode calls. Each stage +kept an independent per-layer KV cache, and `Stop` cleared the session in both +processes. + +The two partial files are the exact-range artifacts described in +`../../spikes/mlx-safetensors-stages/FINDINGS.md`. Tied input/output embeddings +are intentionally duplicated across the stages; that is why the sum of the two +files is larger than the full checkpoint even though neither process downloads +the full checkpoint. + +## Reproduce + +Build once: + +```bash +just mlx-stage-build +``` + +Start the final stage: + +```bash +just mlx-stage serve \ + --model /tmp/mlx-split-smol/stage1 \ + --model-id HuggingFaceTB/SmolLM2-135M-Instruct \ + --stage-index 1 --layer-start 15 --layer-end 30 \ + --bind 127.0.0.1:19091 --wire-dtype f16 --compute-dtype bf16 +``` + +Start the first stage in another terminal: + +```bash +just mlx-stage serve \ + --model /tmp/mlx-split-smol/stage0 \ + --model-id HuggingFaceTB/SmolLM2-135M-Instruct \ + --stage-index 0 --layer-start 0 --layer-end 15 \ + --bind 127.0.0.1:19090 --downstream 127.0.0.1:19091 \ + --wire-dtype f16 --compute-dtype bf16 +``` + +Drive the chain: + +```bash +just mlx-stage prove --connect 127.0.0.1:19090 --wire-dtype f16 +``` + +## Deliberate limitations of this checkpoint + +- Dense Llama-family checkpoints only. The engine boundary is family-neutral, + but the current MLX adapter is the smallest implementation that proves it. +- Greedy sampling only; sampling metadata is preserved in the contract and + rejected explicitly when enabled. +- No KV page import/export, cache trim/checkpoint, MTP, speculative verify, + multimodal projection, or transport batching yet. +- `engine_transport` is the reduced compatibility lane. The mature llama.cpp + binary server remains unchanged and still owns telemetry, exact-prefix cache, + batching, and OpenAI orchestration. +- Mesh topology planning does not yet launch `MlxStageEngine`; the next product + step is selecting this engine from stage config and advertising it as an + additive capability. There is no mesh protocol or Skippy ABI break here. diff --git a/crates/skippy-engine-mlx/src/backend.rs b/crates/skippy-engine-mlx/src/backend.rs index 3a14867d7b..16a1bad7fb 100644 --- a/crates/skippy-engine-mlx/src/backend.rs +++ b/crates/skippy-engine-mlx/src/backend.rs @@ -12,10 +12,10 @@ use openai_frontend::backend::{ ChatCompletionStream, OpenAiBackend, OpenAiRequestContext, OpenAiResult, }; use openai_frontend::chat::{ - message_content_to_text, AssistantMessage, ChatCompletionChoice, ChatCompletionChunk, - ChatCompletionChunkChoice, ChatCompletionDelta, ChatCompletionRequest, ChatCompletionResponse, + AssistantMessage, ChatCompletionChoice, ChatCompletionChunk, ChatCompletionChunkChoice, + ChatCompletionDelta, ChatCompletionRequest, ChatCompletionResponse, message_content_to_text, }; -use openai_frontend::common::{completion_id, FinishReason, Usage}; +use openai_frontend::common::{FinishReason, Usage, completion_id}; use openai_frontend::errors::OpenAiError; use openai_frontend::models::ModelObject; diff --git a/crates/skippy-engine-mlx/src/bin/mlx-stage.rs b/crates/skippy-engine-mlx/src/bin/mlx-stage.rs new file mode 100644 index 0000000000..db6ee4817b --- /dev/null +++ b/crates/skippy-engine-mlx/src/bin/mlx-stage.rs @@ -0,0 +1,265 @@ +//! Run or probe an MLX partial-layer engine over Skippy's binary stage wire. + +#[cfg(all(feature = "mlx", target_os = "macos"))] +mod real { + use std::{ + io::Write, + net::{SocketAddr, TcpStream}, + path::PathBuf, + sync::Arc, + }; + + use anyhow::{Context, Result, ensure}; + use clap::{Parser, Subcommand, ValueEnum}; + use skippy_engine_mlx::{MlxComputeDtype, MlxStageEngine, MlxStageEngineConfig}; + use skippy_protocol::binary::{ + StageStateHeader, StageWireMessage, WireActivationDType, WireMessageKind, WireReplyKind, + recv_ready, recv_reply, write_stage_message, + }; + use skippy_server::engine_transport::{EngineStageServerOptions, serve_stage_engine}; + + #[derive(Debug, Parser)] + #[command(about = "Serve and prove partial SafeTensors MLX stages")] + struct Cli { + #[command(subcommand)] + command: Command, + } + + #[derive(Debug, Subcommand)] + enum Command { + /// Load one partial SafeTensors artifact and serve its layer range. + Serve { + #[arg(long)] + model: PathBuf, + #[arg(long, default_value = "mlx-stage-model")] + model_id: String, + #[arg(long)] + stage_index: u32, + #[arg(long)] + layer_start: u32, + #[arg(long)] + layer_end: u32, + #[arg(long)] + bind: SocketAddr, + #[arg(long)] + downstream: Option, + #[arg(long, value_enum, default_value_t = WireDtype::F16)] + wire_dtype: WireDtype, + #[arg(long, value_enum, default_value_t = ComputeDtype::Bf16)] + compute_dtype: ComputeDtype, + }, + /// Drive a stage chain and assert its greedy token sequence. + Prove { + #[arg(long)] + connect: SocketAddr, + #[arg(long, default_value = "1,1531,314,260,3575,28")] + tokens: String, + #[arg(long, default_value = "284,260,2240,314,1343,327,624,8685")] + expected: String, + #[arg(long, value_enum, default_value_t = WireDtype::F16)] + wire_dtype: WireDtype, + }, + } + + #[derive(Clone, Copy, Debug, ValueEnum)] + enum WireDtype { + F16, + F32, + } + + impl From for WireActivationDType { + fn from(value: WireDtype) -> Self { + match value { + WireDtype::F16 => Self::F16, + WireDtype::F32 => Self::F32, + } + } + } + + #[derive(Clone, Copy, Debug, ValueEnum)] + enum ComputeDtype { + F16, + Bf16, + F32, + } + + impl From for MlxComputeDtype { + fn from(value: ComputeDtype) -> Self { + match value { + ComputeDtype::F16 => Self::F16, + ComputeDtype::Bf16 => Self::Bf16, + ComputeDtype::F32 => Self::F32, + } + } + } + + pub fn main() -> Result<()> { + match Cli::parse().command { + Command::Serve { + model, + model_id, + stage_index, + layer_start, + layer_end, + bind, + downstream, + wire_dtype, + compute_dtype, + } => serve( + MlxStageEngineConfig { + model_dir: model, + model_id, + stage_index, + layer_start, + layer_end, + compute_dtype: compute_dtype.into(), + }, + EngineStageServerOptions { + bind_addr: bind, + downstream_addr: downstream, + wire_dtype: wire_dtype.into(), + }, + ), + Command::Prove { + connect, + tokens, + expected, + wire_dtype, + } => prove(connect, &tokens, &expected, wire_dtype.into()), + } + } + + fn serve(config: MlxStageEngineConfig, options: EngineStageServerOptions) -> Result<()> { + let engine = Arc::new(MlxStageEngine::spawn(config)?); + serve_stage_engine(engine, options) + } + + fn prove( + connect: SocketAddr, + tokens: &str, + expected: &str, + wire_dtype: WireActivationDType, + ) -> Result<()> { + let prompt = parse_ids(tokens)?; + let expected = parse_ids(expected)?; + ensure!( + !expected.is_empty(), + "expected token sequence must not be empty" + ); + let mut stream = TcpStream::connect(connect) + .with_context(|| format!("connect first MLX stage at {connect}"))?; + stream.set_nodelay(true).ok(); + recv_ready(&mut stream).context("first MLX stage did not become ready")?; + + let session_id = 1; + let request_id = 1; + let prefill = execution_message( + WireMessageKind::PrefillFinalEmbd, + &prompt, + *prompt.last().context("prompt must not be empty")?, + request_id, + session_id, + wire_dtype, + 0, + ); + let mut generated = Vec::with_capacity(expected.len()); + generated.push(send_predicted(&mut stream, &prefill, wire_dtype)?); + + while generated.len() < expected.len() { + let current = *generated.last().expect("generated has first token"); + let decode = execution_message( + WireMessageKind::DecodeEmbd, + &[], + current, + request_id, + session_id, + wire_dtype, + i32::try_from(generated.len())?, + ); + generated.push(send_predicted(&mut stream, &decode, wire_dtype)?); + } + + ensure!( + generated == expected, + "two-process stage tokens diverged: expected={expected:?} actual={generated:?}" + ); + let stop = StageWireMessage::stop_with_identity(wire_dtype, request_id, session_id); + write_stage_message(&mut stream, &stop, wire_dtype)?; + stream.flush().ok(); + let reply = recv_reply(&mut stream)?; + ensure!(reply.kind == WireReplyKind::Ack, "stop did not return ACK"); + println!("PASS: two MLX stage processes matched the reference greedy tokens"); + println!("wire_dtype={wire_dtype:?}"); + println!("generated_tokens={generated:?}"); + Ok(()) + } + + fn execution_message( + kind: WireMessageKind, + tokens: &[i32], + current_token: i32, + request_id: u64, + session_id: u64, + wire_dtype: WireActivationDType, + decode_step: i32, + ) -> StageWireMessage { + let mut state = StageStateHeader::new(kind, wire_dtype); + state.current_token = current_token; + state.prompt_token_count = i32::try_from(tokens.len()).unwrap_or_default(); + state.decode_step = decode_step; + StageWireMessage { + kind, + pos_start: 0, + token_count: if kind == WireMessageKind::DecodeEmbd { + 1 + } else { + i32::try_from(tokens.len()).unwrap_or_default() + }, + state, + request_id, + session_id, + sampling: None, + chat_sampling_metadata: None, + tokens: tokens.to_vec(), + positions: Vec::new(), + activation: Vec::new(), + raw_bytes: Vec::new(), + } + } + + fn send_predicted( + stream: &mut TcpStream, + message: &StageWireMessage, + wire_dtype: WireActivationDType, + ) -> Result { + write_stage_message(&mut *stream, message, wire_dtype)?; + stream.flush().ok(); + let reply = recv_reply(&mut *stream)?; + ensure!( + matches!( + reply.kind, + WireReplyKind::PredictedToken | WireReplyKind::PredictedTokens + ), + "stage chain did not return a predicted token" + ); + Ok(reply.predicted) + } + + fn parse_ids(value: &str) -> Result> { + value + .split(',') + .map(|token| token.trim().parse().context("parse token ID")) + .collect() + } +} + +#[cfg(all(feature = "mlx", target_os = "macos"))] +fn main() -> anyhow::Result<()> { + real::main() +} + +#[cfg(not(all(feature = "mlx", target_os = "macos")))] +fn main() { + eprintln!("mlx-stage requires macOS and `--features mlx`"); + std::process::exit(1); +} diff --git a/crates/skippy-engine-mlx/src/engine.rs b/crates/skippy-engine-mlx/src/engine.rs index 989ac526c2..9378feda7b 100644 --- a/crates/skippy-engine-mlx/src/engine.rs +++ b/crates/skippy-engine-mlx/src/engine.rs @@ -15,14 +15,14 @@ use std::path::PathBuf; use std::thread; use std::time::Instant; -use anyhow::{anyhow, Result}; -use serde_json::{json, Value}; +use anyhow::{Result, anyhow}; +use serde_json::{Value, json}; use tokio::sync::mpsc; use safemlx::transforms::async_eval; use safemlx::{Device, DeviceType, Stream}; -use safemlx_lm::models::input::{InputPart, ModelInput}; use safemlx_lm::models::LoadedModel; +use safemlx_lm::models::input::{InputPart, ModelInput}; use safemlx_lm::sampler::DefaultSampler; /// How the worker should load and run a model. diff --git a/crates/skippy-engine-mlx/src/lib.rs b/crates/skippy-engine-mlx/src/lib.rs index 5dcffca700..c948045cfd 100644 --- a/crates/skippy-engine-mlx/src/lib.rs +++ b/crates/skippy-engine-mlx/src/lib.rs @@ -11,11 +11,15 @@ mod backend; #[cfg(all(feature = "mlx", target_os = "macos"))] mod engine; +#[cfg(all(feature = "mlx", target_os = "macos"))] +mod stage; #[cfg(all(feature = "mlx", target_os = "macos"))] pub use backend::MlxBackend; #[cfg(all(feature = "mlx", target_os = "macos"))] pub use engine::{ChatTurn, GenerateRequest, MlxEngine, MlxEngineConfig}; +#[cfg(all(feature = "mlx", target_os = "macos"))] +pub use stage::{MlxComputeDtype, MlxStageEngine, MlxStageEngineConfig}; /// True when this build actually contains the MLX engine. pub const fn mlx_available() -> bool { diff --git a/crates/skippy-engine-mlx/src/stage.rs b/crates/skippy-engine-mlx/src/stage.rs new file mode 100644 index 0000000000..fa78a349fb --- /dev/null +++ b/crates/skippy-engine-mlx/src/stage.rs @@ -0,0 +1,392 @@ +//! Partial-layer MLX implementation of the engine-neutral Skippy stage contract. + +use std::{ + collections::BTreeMap, + path::{Path, PathBuf}, + sync::mpsc, + thread, +}; + +use anyhow::{Context, Result, anyhow, bail, ensure}; +use safemlx::module::{Module, ModuleParameters, ModuleParametersExt}; +use safemlx::ops::indexing::{NewAxis, TryIndexOp}; +use safemlx::{Array, Device, DeviceType, Dtype, Stream, arange}; +use safemlx_lm::{ + cache::{ConcatKeyValueCache, KeyValueCache}, + models::{ + common::linear::project_logits_maybe_quantized, + llama::{self, AttentionInput, TransformerBlock}, + }, + weights::load_safetensors_lenient, +}; +use skippy_engine::{ + StageActivation, StageEngine, StageEngineInfo, StageExecutionKind, StageExecutionOutput, + StageExecutionRequest, +}; + +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub enum MlxComputeDtype { + F16, + #[default] + Bf16, + F32, +} + +impl MlxComputeDtype { + fn mlx(self) -> Dtype { + match self { + Self::F16 => Dtype::Float16, + Self::Bf16 => Dtype::Bfloat16, + Self::F32 => Dtype::Float32, + } + } +} + +#[derive(Clone, Debug)] +pub struct MlxStageEngineConfig { + pub model_dir: PathBuf, + pub model_id: String, + pub stage_index: u32, + pub layer_start: u32, + pub layer_end: u32, + pub compute_dtype: MlxComputeDtype, +} + +enum WorkerJob { + Execute { + request: StageExecutionRequest, + reply: mpsc::Sender>, + }, + Reset { + session_id: u64, + reply: mpsc::Sender>, + }, +} + +/// Send+Sync handle whose worker thread exclusively owns all MLX objects. +pub struct MlxStageEngine { + info: StageEngineInfo, + jobs: mpsc::Sender, +} + +impl MlxStageEngine { + pub fn spawn(config: MlxStageEngineConfig) -> Result { + let (jobs, job_rx) = mpsc::channel(); + let (ready_tx, ready_rx) = mpsc::channel(); + thread::Builder::new() + .name(format!("mlx-stage-{}", config.stage_index)) + .spawn(move || run_worker(config, job_rx, ready_tx))?; + match ready_rx.recv() { + Ok(Ok(info)) => Ok(Self { info, jobs }), + Ok(Err(error)) => Err(anyhow!("MLX stage load failed: {error}")), + Err(_) => Err(anyhow!("MLX stage worker exited before readiness")), + } + } + + fn request( + &self, + make_job: impl FnOnce(mpsc::Sender>) -> WorkerJob, + ) -> Result { + let (reply_tx, reply_rx) = mpsc::channel(); + self.jobs + .send(make_job(reply_tx)) + .map_err(|_| anyhow!("MLX stage worker is not running"))?; + reply_rx + .recv() + .map_err(|_| anyhow!("MLX stage worker dropped its reply"))? + .map_err(anyhow::Error::msg) + } +} + +impl StageEngine for MlxStageEngine { + fn info(&self) -> &StageEngineInfo { + &self.info + } + + fn execute(&self, request: StageExecutionRequest) -> Result { + self.request(|reply| WorkerJob::Execute { request, reply }) + } + + fn reset_session(&self, session_id: u64) -> Result<()> { + self.request(|reply| WorkerJob::Reset { session_id, reply }) + } +} + +struct LoadedStage { + model: llama::Model, + stream: Stream, + compute_dtype: Dtype, + info: StageEngineInfo, + sessions: BTreeMap>>, +} + +fn run_worker( + config: MlxStageEngineConfig, + job_rx: mpsc::Receiver, + ready_tx: mpsc::Sender>, +) { + let mut stage = match load_stage(config) { + Ok(stage) => { + let _ = ready_tx.send(Ok(stage.info.clone())); + stage + } + Err(error) => { + let _ = ready_tx.send(Err(format!("{error:#}"))); + return; + } + }; + while let Ok(job) = job_rx.recv() { + match job { + WorkerJob::Execute { request, reply } => { + let _ = reply.send(stage.execute(request).map_err(|error| format!("{error:#}"))); + } + WorkerJob::Reset { session_id, reply } => { + stage.sessions.remove(&session_id); + let _ = reply.send(Ok(())); + } + } + } +} + +fn load_stage(config: MlxStageEngineConfig) -> Result { + let stream = Stream::new_with_device(&Device::new(DeviceType::Gpu, 0)); + let weights_stream = Stream::new_with_device(&Device::new(DeviceType::Cpu, 0)); + let model_args = llama::get_llama_model_args(&config.model_dir)?; + let total_layers = u32::try_from(model_args.num_hidden_layers)?; + let info = StageEngineInfo { + engine: "mlx".to_string(), + model_id: config.model_id, + stage_index: config.stage_index, + layer_start: config.layer_start, + layer_end: config.layer_end, + total_layers, + activation_width: u32::try_from(model_args.hidden_size)?, + }; + info.validate()?; + + let mut model = llama::Model::new(model_args, &stream)?; + load_safetensors_lenient(&mut model, weight_file(&config.model_dir), &weights_stream)?; + retain_local_layers(&mut model, info.layer_start, info.layer_end)?; + copy_stage_weights_to_compute_stream(&mut model, &info, &stream)?; + stream.synchronize()?; + eprintln!( + "MLX partial stage loaded: model={} stage={} layers={}..{} tensors={}", + info.model_id, + info.stage_index, + info.layer_start, + info.layer_end, + model.parameters().flatten().len(), + ); + Ok(LoadedStage { + model, + stream, + compute_dtype: config.compute_dtype.mlx(), + info, + sessions: BTreeMap::new(), + }) +} + +fn weight_file(model_dir: &Path) -> PathBuf { + model_dir.join("model.safetensors") +} + +fn retain_local_layers(model: &mut llama::Model, start: u32, end: u32) -> Result<()> { + let start = usize::try_from(start)?; + let end = usize::try_from(end)?; + ensure!( + end <= model.model.layers.len(), + "stage layer range is out of bounds" + ); + model.model.layers = model.model.layers.drain(start..end).collect(); + model.model.num_hidden_layers = i32::try_from(model.model.layers.len())?; + Ok(()) +} + +fn copy_stage_weights_to_compute_stream( + model: &mut llama::Model, + info: &StageEngineInfo, + stream: &Stream, +) -> Result<()> { + if info.is_first() || info.is_final() { + model.model.embed_tokens.copy_to_stream(stream)?; + } + for layer in &mut model.model.layers { + layer.copy_to_stream(stream)?; + } + if info.is_final() { + model.model.norm.copy_to_stream(stream)?; + if let Some(lm_head) = &mut model.lm_head { + lm_head.copy_to_stream(stream)?; + } + } + Ok(()) +} + +impl LoadedStage { + fn execute(&mut self, request: StageExecutionRequest) -> Result { + if request.kind == StageExecutionKind::Verify { + bail!("MLX dense stage verification is not implemented yet"); + } + if request + .sampling + .as_ref() + .is_some_and(|sampling| sampling.enabled()) + { + bail!("MLX staged execution currently supports greedy sampling only"); + } + ensure!(!request.token_ids.is_empty(), "stage request has no tokens"); + let token_count = request.token_ids.len(); + if let Some(input) = request.input.as_ref() { + ensure!( + input.token_count == token_count, + "input activation token count does not match token sideband" + ); + } + + let mut hidden = self.input_hidden(&request)?; + let caches = self.sessions.entry(request.session_id).or_default(); + let mask = attention_mask(&hidden, caches, &self.stream)?; + if caches.is_empty() { + *caches = (0..self.model.model.layers.len()) + .map(|_| Some(ConcatKeyValueCache::default())) + .collect(); + } + hidden = forward_blocks( + &mut self.model.model.layers, + hidden, + mask.as_ref(), + caches, + &self.stream, + )?; + + if self.info.is_final() { + let hidden = self.model.model.norm.forward(&hidden, &self.stream)?; + let logits = project_logits_maybe_quantized( + &mut self.model.lm_head, + &mut self.model.model.embed_tokens, + &hidden, + &self.stream, + )?; + let predicted = last_argmax(&logits, &self.stream)?; + return Ok(StageExecutionOutput { + activation: None, + predicted_tokens: vec![predicted], + }); + } + + Ok(StageExecutionOutput { + activation: Some(array_activation(&hidden, &self.stream)?), + predicted_tokens: Vec::new(), + }) + } + + fn input_hidden(&mut self, request: &StageExecutionRequest) -> Result { + if self.info.is_first() { + ensure!( + request.input.is_none(), + "first stage cannot accept residual input" + ); + let tokens = request + .token_ids + .iter() + .copied() + .map(|token| u32::try_from(token).context("negative token ID")) + .collect::>>()?; + let shape = [1, i32::try_from(tokens.len())?]; + let tokens = Array::from_slice(&tokens, &shape); + return Ok(self + .model + .model + .embed_tokens + .forward(&tokens, &self.stream)?); + } + let input = request + .input + .as_ref() + .context("non-first stage requires residual input")?; + ensure!( + input.width == self.info.activation_width as usize, + "input activation width mismatch" + ); + let values = input.values(); + let hidden = Array::from_slice( + &values, + &[ + 1, + i32::try_from(input.token_count)?, + i32::try_from(input.width)?, + ], + ); + Ok(hidden.as_dtype(self.compute_dtype, &self.stream)?) + } +} + +fn attention_mask( + hidden: &Array, + cache: &[Option], + stream: &Stream, +) -> Result> { + let sequence = hidden.shape()[1]; + if sequence == 1 { + return Ok(None); + } + let offset = cache + .first() + .and_then(Option::as_ref) + .map_or(0, KeyValueCache::offset); + let right = arange!(stop = offset + sequence, stream = stream)?; + let left = arange!(start = offset, stop = offset + sequence, stream = stream)?; + let left = left.try_index_device((.., NewAxis), stream)?; + let right = right.try_index_device(NewAxis, stream)?; + Ok(Some(left.ge(&right, stream)?)) +} + +fn forward_blocks( + blocks: &mut [TransformerBlock], + mut hidden: Array, + mask: Option<&Array>, + cache: &mut [Option], + stream: &Stream, +) -> Result { + ensure!(cache.len() == blocks.len(), "stage cache length mismatch"); + for (block, layer_cache) in blocks.iter_mut().zip(cache.iter_mut()) { + hidden = block.forward( + AttentionInput { + x: &hidden, + mask, + cache: layer_cache.as_mut(), + generated_sliding_window: None, + }, + stream, + )?; + } + Ok(hidden) +} + +fn array_activation(hidden: &Array, stream: &Stream) -> Result { + let shape = hidden.shape().to_vec(); + ensure!( + shape.len() == 3 && shape[0] == 1, + "unexpected residual shape" + ); + let values = hidden + .as_dtype(Dtype::Float32, stream)? + .evaluated()? + .as_slice::() + .to_vec(); + StageActivation::from_values(shape[1] as usize, shape[2] as usize, &values) +} + +fn last_argmax(logits: &Array, stream: &Stream) -> Result { + let row = logits + .try_index_device((0, -1, ..), stream)? + .as_dtype(Dtype::Float32, stream)?; + let evaluated = row.evaluated()?; + let values = evaluated.as_slice::(); + let (index, _) = values + .iter() + .enumerate() + .max_by(|(_, left), (_, right)| left.total_cmp(right)) + .context("cannot argmax empty logits")?; + Ok(i32::try_from(index)?) +} diff --git a/crates/skippy-engine/Cargo.toml b/crates/skippy-engine/Cargo.toml new file mode 100644 index 0000000000..6039749fcb --- /dev/null +++ b/crates/skippy-engine/Cargo.toml @@ -0,0 +1,16 @@ +[package] +name = "skippy-engine" +version.workspace = true +edition.workspace = true +license.workspace = true +description = "Engine-neutral staged execution contract for Skippy" +repository = "https://github.com/Mesh-LLM/mesh-llm" +homepage = "https://github.com/Mesh-LLM/mesh-llm" +readme = "README.md" + +[lints] +workspace = true + +[dependencies] +anyhow.workspace = true +skippy-protocol = { path = "../skippy-protocol", version = "0.72.1" } diff --git a/crates/skippy-engine/README.md b/crates/skippy-engine/README.md new file mode 100644 index 0000000000..6974d5e3ed --- /dev/null +++ b/crates/skippy-engine/README.md @@ -0,0 +1,12 @@ +# skippy-engine + +Engine-neutral staged execution contract for Skippy. + +The crate contains shared stage descriptors and the `StageEngine` trait. It has +no model runtime of its own: concrete engines keep native arrays, model handles, +and KV caches private while exchanging token IDs and Skippy-owned activation +buffers with the server transport. + +The first concrete second-engine implementation is `skippy-engine-mlx`. The +existing llama.cpp implementation remains in `skippy-runtime` while its broader +cache, multimodal, MTP, and session surface is migrated behind this contract. diff --git a/crates/skippy-engine/src/lib.rs b/crates/skippy-engine/src/lib.rs new file mode 100644 index 0000000000..6fc9b4521e --- /dev/null +++ b/crates/skippy-engine/src/lib.rs @@ -0,0 +1,178 @@ +//! Engine-neutral staged execution contract. +//! +//! This crate deliberately owns no model runtime. It is the narrow boundary +//! between Skippy's stage transport and concrete engines such as llama.cpp or +//! MLX: token IDs and F32 residual bytes enter, residual bytes and optional +//! predicted token IDs leave. Native arrays and cache handles never cross it. + +use anyhow::{Result, bail, ensure}; +use skippy_protocol::binary::StageSamplingConfig; + +/// Static facts needed by the stage transport. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct StageEngineInfo { + pub engine: String, + pub model_id: String, + pub stage_index: u32, + pub layer_start: u32, + pub layer_end: u32, + pub total_layers: u32, + pub activation_width: u32, +} + +impl StageEngineInfo { + pub fn is_first(&self) -> bool { + self.layer_start == 0 + } + + pub fn is_final(&self) -> bool { + self.layer_end == self.total_layers + } + + pub fn validate(&self) -> Result<()> { + ensure!( + self.layer_start < self.layer_end, + "stage layer range must be non-empty" + ); + ensure!( + self.layer_end <= self.total_layers, + "stage layer range exceeds model layer count" + ); + ensure!( + self.activation_width > 0, + "activation width must be positive" + ); + Ok(()) + } +} + +/// One decoded residual activation tensor in token-major F32 form. +#[derive(Clone, Debug, PartialEq)] +pub struct StageActivation { + pub token_count: usize, + pub width: usize, + pub f32_le_bytes: Vec, +} + +impl StageActivation { + pub fn new(token_count: usize, width: usize, f32_le_bytes: Vec) -> Result { + let expected = token_count + .checked_mul(width) + .and_then(|elements| elements.checked_mul(size_of::())) + .ok_or_else(|| anyhow::anyhow!("stage activation size overflow"))?; + ensure!( + f32_le_bytes.len() == expected, + "stage activation has {} bytes; expected {expected}", + f32_le_bytes.len() + ); + Ok(Self { + token_count, + width, + f32_le_bytes, + }) + } + + pub fn values(&self) -> Vec { + self.f32_le_bytes + .chunks_exact(size_of::()) + .map(|chunk| f32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]])) + .collect() + } + + pub fn from_values(token_count: usize, width: usize, values: &[f32]) -> Result { + let expected = token_count + .checked_mul(width) + .ok_or_else(|| anyhow::anyhow!("stage activation element count overflow"))?; + ensure!( + values.len() == expected, + "stage activation has {} values; expected {expected}", + values.len() + ); + Self::new( + token_count, + width, + values + .iter() + .flat_map(|value| value.to_le_bytes()) + .collect(), + ) + } +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum StageExecutionKind { + Prefill, + PrefillFinal, + Decode, + Verify, +} + +/// A single stage operation after wire decoding. +#[derive(Clone, Debug, PartialEq)] +pub struct StageExecutionRequest { + pub session_id: u64, + pub kind: StageExecutionKind, + pub token_ids: Vec, + pub positions: Vec, + pub input: Option, + pub sampling: Option, +} + +/// Result of one stage operation before wire encoding. +#[derive(Clone, Debug, Default, PartialEq)] +pub struct StageExecutionOutput { + pub activation: Option, + pub predicted_tokens: Vec, +} + +impl StageExecutionOutput { + pub fn predicted(&self) -> Option { + self.predicted_tokens.first().copied() + } +} + +/// Concrete staged model execution behind a transport-neutral interface. +pub trait StageEngine: Send + Sync + 'static { + fn info(&self) -> &StageEngineInfo; + + fn execute(&self, request: StageExecutionRequest) -> Result; + + fn reset_session(&self, session_id: u64) -> Result<()>; + + fn checkpoint_session(&self, _session_id: u64) -> Result<()> { + bail!( + "{} stage engine does not support checkpoints", + self.info().engine + ) + } + + fn restore_session(&self, _session_id: u64) -> Result<()> { + bail!( + "{} stage engine does not support checkpoint restore", + self.info().engine + ) + } + + fn trim_session(&self, _session_id: u64, _token_count: u64) -> Result<()> { + bail!( + "{} stage engine does not support cache trim", + self.info().engine + ) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn activation_round_trips_values() { + let activation = StageActivation::from_values(2, 2, &[1.0, -2.0, 3.5, 4.0]).unwrap(); + assert_eq!(activation.values(), vec![1.0, -2.0, 3.5, 4.0]); + } + + #[test] + fn activation_rejects_wrong_size() { + assert!(StageActivation::new(2, 2, vec![0; 4]).is_err()); + } +} diff --git a/crates/skippy-server/Cargo.toml b/crates/skippy-server/Cargo.toml index f1abc29015..1371032482 100644 --- a/crates/skippy-server/Cargo.toml +++ b/crates/skippy-server/Cargo.toml @@ -25,6 +25,7 @@ clap.workspace = true futures-util = "0.3" skippy-runtime = { path = "../skippy-runtime", version = "0.72.1" } skippy-protocol = { path = "../skippy-protocol", version = "0.72.1" } +skippy-engine = { path = "../skippy-engine", version = "0.72.1" } skippy-cache = { path = "../skippy-cache", version = "0.72.1" } skippy-metrics = { path = "../skippy-metrics", version = "0.72.1" } openai-frontend = { path = "../openai-frontend", version = "0.72.1" } diff --git a/crates/skippy-server/src/engine_transport.rs b/crates/skippy-server/src/engine_transport.rs new file mode 100644 index 0000000000..b39bcaf433 --- /dev/null +++ b/crates/skippy-server/src/engine_transport.rs @@ -0,0 +1,360 @@ +//! Minimal engine-neutral server for the existing Skippy binary stage wire. +//! +//! The mature llama.cpp lane in [`crate::binary_transport`] still owns KV-page +//! caching, telemetry, batching, MTP, and OpenAI orchestration. This module is +//! intentionally the smaller compatibility seam: it proves a `StageEngine` +//! can participate in a real multi-process Skippy chain without introducing a +//! second wire protocol. Advanced operations stay capability-gated. + +use std::{ + io::{self, Write}, + net::{SocketAddr, TcpListener, TcpStream}, + sync::{ + Arc, + atomic::{AtomicBool, Ordering}, + }, + thread, + time::Duration, +}; + +use anyhow::{Context, Result, bail, ensure}; +use skippy_engine::{ + StageActivation, StageEngine, StageExecutionKind, StageExecutionOutput, StageExecutionRequest, +}; +use skippy_protocol::binary::{ + StageReply, StageWireMessage, WireActivationDType, WireMessageKind, WireReplyKind, + encode_f32_activation_payload, read_stage_message, recv_ready, recv_reply, send_ready, + send_reply_ack_with_stats, send_reply_predicted_tokens_with_stats, + send_reply_predicted_with_tokens_and_stats, write_stage_message, +}; + +#[derive(Clone, Debug)] +pub struct EngineStageServerOptions { + pub bind_addr: SocketAddr, + pub downstream_addr: Option, + pub wire_dtype: WireActivationDType, +} + +pub fn serve_stage_engine( + engine: Arc, + options: EngineStageServerOptions, +) -> Result<()> { + serve_stage_engine_until(engine, options, Arc::new(AtomicBool::new(false))) +} + +pub fn serve_stage_engine_until( + engine: Arc, + options: EngineStageServerOptions, + shutdown: Arc, +) -> Result<()> { + engine.info().validate()?; + validate_topology(engine.as_ref(), &options)?; + let listener = TcpListener::bind(options.bind_addr) + .with_context(|| format!("bind engine stage at {}", options.bind_addr))?; + listener.set_nonblocking(true)?; + eprintln!( + "skippy engine stage listening: engine={} model={} binary={} layers={}..{} width={} dtype={:?}", + engine.info().engine, + engine.info().model_id, + listener.local_addr()?, + engine.info().layer_start, + engine.info().layer_end, + engine.info().activation_width, + options.wire_dtype, + ); + + while !shutdown.load(Ordering::SeqCst) { + let (upstream, peer_addr) = match listener.accept() { + Ok(connection) => connection, + Err(error) if error.kind() == io::ErrorKind::WouldBlock => { + thread::sleep(Duration::from_millis(25)); + continue; + } + Err(error) => return Err(error).context("accept engine stage connection"), + }; + upstream.set_nonblocking(false)?; + upstream.set_nodelay(true).ok(); + let engine = engine.clone(); + let options = options.clone(); + thread::spawn(move || { + if let Err(error) = handle_connection(engine, options, upstream) { + eprintln!("engine stage connection from {peer_addr} failed: {error:#}"); + } + }); + } + Ok(()) +} + +fn validate_topology(engine: &dyn StageEngine, options: &EngineStageServerOptions) -> Result<()> { + ensure!( + engine.info().is_final() == options.downstream_addr.is_none(), + "only the final stage may omit a downstream address" + ); + Ok(()) +} + +fn handle_connection( + engine: Arc, + options: EngineStageServerOptions, + mut upstream: TcpStream, +) -> Result<()> { + send_ready(&mut upstream).context("send engine stage ready")?; + upstream.flush().ok(); + let mut downstream = options + .downstream_addr + .map(connect_downstream) + .transpose()?; + let activation_width = + i32::try_from(engine.info().activation_width).context("activation width exceeds i32")?; + + loop { + let message = match read_stage_message(&mut upstream, activation_width) { + Ok(message) => message, + Err(error) if error.kind() == io::ErrorKind::UnexpectedEof => return Ok(()), + Err(error) => return Err(error).context("read engine stage message"), + }; + if message.kind == WireMessageKind::Stop { + engine.reset_session(message.session_id)?; + let downstream_reply = + forward_control(downstream.as_mut(), &message, options.wire_dtype)?; + send_ack(&mut upstream, downstream_reply)?; + continue; + } + if message.kind.is_session_control() { + execute_session_control(engine.as_ref(), &message)?; + let downstream_reply = + forward_control(downstream.as_mut(), &message, options.wire_dtype)?; + send_ack(&mut upstream, downstream_reply)?; + continue; + } + + let request = execution_request(&message, activation_width)?; + let output = engine.execute(request)?; + match downstream.as_mut() { + Some(downstream) => { + let forwarded = + forwarded_message(engine.as_ref(), &message, output, options.wire_dtype)?; + write_stage_message(&mut *downstream, &forwarded, options.wire_dtype) + .context("forward engine stage message")?; + downstream.flush().ok(); + let reply = recv_reply(&mut *downstream).context("receive downstream reply")?; + send_reply(&mut upstream, reply)?; + } + None => send_final_reply(&mut upstream, &message, output)?, + } + } +} + +fn connect_downstream(addr: SocketAddr) -> Result { + let mut stream = TcpStream::connect(addr) + .with_context(|| format!("connect downstream engine stage at {addr}"))?; + stream.set_nodelay(true).ok(); + recv_ready(&mut stream).context("downstream engine stage did not become ready")?; + Ok(stream) +} + +fn execute_session_control(engine: &dyn StageEngine, message: &StageWireMessage) -> Result<()> { + match message.kind { + WireMessageKind::CheckpointSession => engine.checkpoint_session(message.session_id), + WireMessageKind::RestoreSession => engine.restore_session(message.session_id), + WireMessageKind::TrimSession => { + engine.trim_session(message.session_id, message.token_count.max(0) as u64) + } + _ => bail!("message is not session control"), + } +} + +fn forward_control( + downstream: Option<&mut TcpStream>, + message: &StageWireMessage, + wire_dtype: WireActivationDType, +) -> Result> { + let Some(downstream) = downstream else { + return Ok(None); + }; + write_stage_message(&mut *downstream, message, wire_dtype)?; + downstream.flush().ok(); + Ok(Some(recv_reply(&mut *downstream)?)) +} + +fn execution_request( + message: &StageWireMessage, + activation_width: i32, +) -> Result { + let kind = match message.kind { + WireMessageKind::PrefillEmbd => StageExecutionKind::Prefill, + WireMessageKind::PrefillFinalEmbd => StageExecutionKind::PrefillFinal, + WireMessageKind::DecodeEmbd + | WireMessageKind::DecodeReadout + | WireMessageKind::DecodeLightCtx + | WireMessageKind::DecodeReplayEmbd + | WireMessageKind::DecodeReplayFinalEmbd => StageExecutionKind::Decode, + WireMessageKind::VerifySpan => StageExecutionKind::Verify, + other => bail!("engine stage does not execute {other:?}"), + }; + let token_count = usize::try_from(message.token_count).context("negative token count")?; + let token_ids = execution_tokens(message, kind, token_count)?; + let input = if message.activation.is_empty() { + None + } else { + let bytes = message + .activation_f32_payload(activation_width) + .context("decode input activation")?; + Some(StageActivation::new( + token_count, + usize::try_from(activation_width)?, + bytes, + )?) + }; + Ok(StageExecutionRequest { + session_id: message.session_id, + kind, + token_ids, + positions: message.positions.clone(), + input, + sampling: message.sampling.clone(), + }) +} + +fn execution_tokens( + message: &StageWireMessage, + kind: StageExecutionKind, + token_count: usize, +) -> Result> { + if kind == StageExecutionKind::Decode { + ensure!(token_count == 1, "decode requires one token"); + return Ok(vec![message.state.current_token]); + } + ensure!( + message.tokens.len() == token_count, + "token sideband length does not match token count" + ); + Ok(message.tokens.clone()) +} + +fn forwarded_message( + engine: &dyn StageEngine, + incoming: &StageWireMessage, + output: StageExecutionOutput, + wire_dtype: WireActivationDType, +) -> Result { + let activation = output + .activation + .context("non-final engine stage returned no activation")?; + ensure!( + activation.width == engine.info().activation_width as usize, + "engine output activation width mismatch" + ); + let mut state = incoming.state; + state.source_stage_index = i32::try_from(engine.info().stage_index)?; + state.reserved = wire_dtype as i32; + let activation = encode_f32_activation_payload( + wire_dtype, + incoming.token_count, + i32::try_from(activation.width)?, + &activation.f32_le_bytes, + )?; + Ok(StageWireMessage { + kind: incoming.kind, + pos_start: incoming.pos_start, + token_count: incoming.token_count, + state, + request_id: incoming.request_id, + session_id: incoming.session_id, + sampling: incoming.sampling.clone(), + chat_sampling_metadata: incoming.chat_sampling_metadata.clone(), + tokens: incoming.tokens.clone(), + positions: incoming.positions.clone(), + activation, + raw_bytes: Vec::new(), + }) +} + +fn send_final_reply( + upstream: &mut TcpStream, + message: &StageWireMessage, + output: StageExecutionOutput, +) -> Result<()> { + if message.kind.requires_predicted_reply() { + ensure!( + !output.predicted_tokens.is_empty(), + "final engine stage returned no prediction" + ); + send_reply_predicted_with_tokens_and_stats( + &mut *upstream, + output.predicted().expect("checked non-empty"), + &output.predicted_tokens, + Default::default(), + )?; + } else { + send_reply_ack_with_stats(&mut *upstream, Default::default())?; + } + upstream.flush().ok(); + Ok(()) +} + +fn send_ack(upstream: &mut TcpStream, downstream: Option) -> Result<()> { + if let Some(reply) = downstream { + ensure!( + reply.kind == WireReplyKind::Ack, + "control expected downstream ACK" + ); + send_reply_ack_with_stats(&mut *upstream, reply.stats)?; + } else { + send_reply_ack_with_stats(&mut *upstream, Default::default())?; + } + upstream.flush().ok(); + Ok(()) +} + +fn send_reply(upstream: &mut TcpStream, reply: StageReply) -> Result<()> { + match reply.kind { + WireReplyKind::Ack => send_reply_ack_with_stats(&mut *upstream, reply.stats)?, + WireReplyKind::PredictedToken => send_reply_predicted_with_tokens_and_stats( + &mut *upstream, + reply.predicted, + &reply.predicted_tokens, + reply.stats, + )?, + WireReplyKind::PredictedTokens => send_reply_predicted_tokens_with_stats( + &mut *upstream, + &reply.predicted_tokens, + reply.stats, + )?, + } + upstream.flush().ok(); + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + use skippy_protocol::binary::StageStateHeader; + + fn decode_message(tokens: Vec, current_token: i32) -> StageWireMessage { + let kind = WireMessageKind::DecodeEmbd; + let mut state = StageStateHeader::new(kind, WireActivationDType::F16); + state.current_token = current_token; + StageWireMessage { + kind, + pos_start: 0, + token_count: 1, + state, + request_id: 1, + session_id: 2, + sampling: None, + chat_sampling_metadata: None, + tokens, + positions: Vec::new(), + activation: Vec::new(), + raw_bytes: Vec::new(), + } + } + + #[test] + fn decode_uses_current_token_not_prompt_sideband() { + let request = execution_request(&decode_message(vec![1, 2, 3], 7), 4).unwrap(); + assert_eq!(request.token_ids, vec![7]); + assert_eq!(request.kind, StageExecutionKind::Decode); + } +} diff --git a/crates/skippy-server/src/lib.rs b/crates/skippy-server/src/lib.rs index 7750e8a2ca..ccb10d13de 100644 --- a/crates/skippy-server/src/lib.rs +++ b/crates/skippy-server/src/lib.rs @@ -8,6 +8,7 @@ pub mod cli; pub mod config; mod decode_batch_policy; pub mod embedded; +pub mod engine_transport; pub mod frontend; pub mod http; pub mod kv_integration; diff --git a/docs/design/MLX_STAGE_ENGINE_PLAN.md b/docs/design/MLX_STAGE_ENGINE_PLAN.md index a6f2495842..25dff9bebb 100644 --- a/docs/design/MLX_STAGE_ENGINE_PLAN.md +++ b/docs/design/MLX_STAGE_ENGINE_PLAN.md @@ -33,10 +33,20 @@ GiB of shard files; exact ranges avoid 833.15 GiB. A SmolLM2-135M proof then materialized two partial files (layers 0..15 and 15..30), loaded each directly into MLX, and matched unsplit logits exactly for prefill plus eight decode steps through Skippy's real F16 and F32 binary activation codec. The remaining -artifact gate is bounded-memory quantization for frontier-sized source tensors; -the remaining product gate is engine-neutral `skippy-server` integration. +artifact gate is bounded-memory quantization for frontier-sized source tensors. See `spikes/mlx-safetensors-stages/FINDINGS.md`. +**Update — the first engine-neutral, multi-process stage chain is now proven.** +The new `skippy-engine` crate defines a runtime-neutral `StageEngine` contract; +`skippy-server::engine_transport` carries that contract over the existing +binary stage protocol; and `MlxStageEngine` runs a partial layer range with its +own KV cache. Two real processes, each given only one 155.28 MiB partial +SmolLM2 artifact, reproduced the eight-token whole-model reference exactly over +F16 residuals. Their post-proof RSS was about 189 MiB each. This closes the +dense execution and process-boundary proof. Host topology selection, advanced +cache/session operations, additional families, and bounded-memory quantization +remain. See `crates/skippy-engine-mlx/STAGED_EXECUTION.md`. + --- ## 1. Bottom line @@ -608,6 +618,12 @@ the trait in a new `skippy-engine` crate, implement it for the existing behavior change; ship this independently of MLX. Validate with existing `skippy-correctness` and `mic-lab` runs. +> **Partially implemented on this branch.** The engine-neutral crate and an +> additive reduced binary server lane now exist, and MLX uses them for the +> two-process proof. The mature llama runtime has not yet been migrated from +> concrete `RuntimeState`; that compatibility refactor remains before the +> engine-neutral lane can replace the default server internals. + **Phase 2 — Solo MLX serving + JIT quant (the workflow win; lead here).** `MlxStageEngine` as a single-stage/whole-model engine: open/load, session, prefill, decode-sampled, tokenizer/chat, final-stage sampling — plus the @@ -742,8 +758,9 @@ Spikes 1 and 2 are more decisive than any standalone token/s benchmark. 1. Add tensor-at-a-time MLX quantization to the materializer and measure peak RSS against the one-source-tensor memory contract. -2. Introduce `StageEngine` with the llama adapter first, then run the proven MLX - split through two real `skippy-server` processes. +2. Add the llama `StageEngine` adapter and route normal `skippy-server` launch + through the engine-neutral contract; the MLX two-process binary-wire proof + is complete. 3. Run **Spike 2 (boundary fence)** at frontier residual widths and keep it as a go/no-go gate. 4. Use Nemotron-H as the first frontier-family follow-up already represented in diff --git a/scripts/affected-crates.sh b/scripts/affected-crates.sh index 6f7b466ecf..6abc340c54 100755 --- a/scripts/affected-crates.sh +++ b/scripts/affected-crates.sh @@ -46,6 +46,7 @@ WORKSPACE_MEMBERS=( "model-hf" "model-resolver" "skippy-protocol" + "skippy-engine" "skippy-coordinator" "skippy-topology" "skippy-cache" diff --git a/scripts/plan-clippy-batches.sh b/scripts/plan-clippy-batches.sh index abff132f92..c80715b099 100644 --- a/scripts/plan-clippy-batches.sh +++ b/scripts/plan-clippy-batches.sh @@ -48,6 +48,7 @@ WORKSPACE_MEMBERS=( "model-hf" "model-resolver" "skippy-protocol" + "skippy-engine" "skippy-coordinator" "skippy-topology" "skippy-cache" diff --git a/scripts/publish-crates.sh b/scripts/publish-crates.sh index a5611ed758..e5097cd116 100755 --- a/scripts/publish-crates.sh +++ b/scripts/publish-crates.sh @@ -374,6 +374,10 @@ unpublished_registry_deps() { printf '%s\n' \ skippy-protocol ;; + skippy-engine) + printf '%s\n' \ + skippy-protocol + ;; skippy-runtime) printf '%s\n' \ skippy-ffi @@ -382,6 +386,7 @@ unpublished_registry_deps() { printf '%s\n' \ openai-frontend \ skippy-cache \ + skippy-engine \ skippy-metrics \ skippy-protocol \ skippy-runtime @@ -469,6 +474,7 @@ publish_crates=( mesh-llm-gpu-bench skippy-ffi skippy-protocol + skippy-engine skippy-coordinator skippy-topology skippy-metrics From 7d3852c723a2ddaec54006c3b651048a0e942a75 Mon Sep 17 00:00:00 2001 From: Michael Neale <14976+michaelneale@users.noreply.github.com> Date: Fri, 17 Jul 2026 16:07:17 +1000 Subject: [PATCH 10/37] refactor(skippy): adapt llama runtime to stage engine --- Cargo.lock | 1 + crates/skippy-engine-mlx/STAGED_EXECUTION.md | 3 + crates/skippy-server/Cargo.toml | 1 + crates/skippy-server/src/lib.rs | 1 + crates/skippy-server/src/llama_engine.rs | 282 +++++++++++++++++++ docs/design/MLX_STAGE_ENGINE_PLAN.md | 17 +- 6 files changed, 297 insertions(+), 8 deletions(-) create mode 100644 crates/skippy-server/src/llama_engine.rs diff --git a/Cargo.lock b/Cargo.lock index 64e5a0ac95..7e44eb505b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -7863,6 +7863,7 @@ dependencies = [ "blake3", "clap", "futures-util", + "half", "libc", "openai-frontend", "opentelemetry-proto", diff --git a/crates/skippy-engine-mlx/STAGED_EXECUTION.md b/crates/skippy-engine-mlx/STAGED_EXECUTION.md index fbadeebcf0..7193d16282 100644 --- a/crates/skippy-engine-mlx/STAGED_EXECUTION.md +++ b/crates/skippy-engine-mlx/STAGED_EXECUTION.md @@ -12,6 +12,9 @@ This is a production-shaped bridge, not yet the default mesh launch path: buffer descriptors. - `skippy-server::engine_transport` serves that contract using the existing `StageWireMessage`, ready handshake, activation codec, and reply codec. +- `skippy-server::llama_engine` proves the existing llama `RuntimeState` can + implement the same dense contract, including F16/BF16/F32 residual conversion + and checkpoint/restore/trim delegation, without changing the native ABI. - `MlxStageEngine` loads one materialized partial SafeTensors file, owns per-session KV caches on a dedicated MLX worker thread, and executes only its configured layer range. diff --git a/crates/skippy-server/Cargo.toml b/crates/skippy-server/Cargo.toml index 1371032482..31b9b89fb6 100644 --- a/crates/skippy-server/Cargo.toml +++ b/crates/skippy-server/Cargo.toml @@ -23,6 +23,7 @@ base64 = "0.22" blake3.workspace = true clap.workspace = true futures-util = "0.3" +half = "2" skippy-runtime = { path = "../skippy-runtime", version = "0.72.1" } skippy-protocol = { path = "../skippy-protocol", version = "0.72.1" } skippy-engine = { path = "../skippy-engine", version = "0.72.1" } diff --git a/crates/skippy-server/src/lib.rs b/crates/skippy-server/src/lib.rs index ccb10d13de..7bdbaf3058 100644 --- a/crates/skippy-server/src/lib.rs +++ b/crates/skippy-server/src/lib.rs @@ -13,6 +13,7 @@ pub mod frontend; pub mod http; pub mod kv_integration; pub mod kv_proto; +pub mod llama_engine; pub mod package; pub mod runtime_state; pub mod telemetry; diff --git a/crates/skippy-server/src/llama_engine.rs b/crates/skippy-server/src/llama_engine.rs new file mode 100644 index 0000000000..a973ae40f4 --- /dev/null +++ b/crates/skippy-server/src/llama_engine.rs @@ -0,0 +1,282 @@ +//! Compatibility adapter from the existing llama.cpp `RuntimeState` to the +//! engine-neutral dense stage contract. +//! +//! The mature binary server continues to call `RuntimeState` directly while +//! its broader cache/MTP/multimodal surface is migrated. This adapter keeps the +//! new contract honest: it is implementable by the existing runtime without +//! changing the native Skippy ABI or teaching MLX about llama types. + +use std::sync::Mutex; + +use anyhow::{Result, bail, ensure}; +use half::{bf16, f16}; +use skippy_engine::{ + StageActivation, StageEngine, StageEngineInfo, StageExecutionKind, StageExecutionOutput, + StageExecutionRequest, +}; +use skippy_runtime::{ + ActivationDesc, ActivationFrame, LogitBias, MAX_LOGIT_BIAS, RuntimeActivationDType, + RuntimeActivationLayout, SamplingConfig, +}; + +use crate::runtime_state::RuntimeState; + +pub struct LlamaStageEngine { + info: StageEngineInfo, + runtime: Mutex, +} + +impl LlamaStageEngine { + pub fn new(info: StageEngineInfo, runtime: RuntimeState) -> Result { + info.validate()?; + Ok(Self { + info, + runtime: Mutex::new(runtime), + }) + } + + pub fn into_runtime(self) -> Result { + self.runtime + .into_inner() + .map_err(|_| anyhow::anyhow!("llama stage runtime lock poisoned")) + } + + fn runtime(&self) -> Result> { + self.runtime + .lock() + .map_err(|_| anyhow::anyhow!("llama stage runtime lock poisoned")) + } + + fn activation_frame(&self, input: Option) -> Result> { + let Some(input) = input else { + return Ok(None); + }; + ensure!( + input.width == self.info.activation_width as usize, + "input activation width mismatch" + ); + let producer_stage_index = self.info.stage_index.saturating_sub(1); + Ok(Some(ActivationFrame { + desc: ActivationDesc { + version: 1, + dtype: RuntimeActivationDType::F32, + layout: RuntimeActivationLayout::TokenMajor, + producer_stage_index: i32::try_from(producer_stage_index)?, + layer_start: 0, + layer_end: i32::try_from(self.info.layer_start)?, + token_count: u32::try_from(input.token_count)?, + sequence_count: u32::from(input.token_count > 0), + payload_bytes: u64::try_from(input.f32_le_bytes.len())?, + flags: 0, + }, + payload: input.f32_le_bytes, + })) + } + + fn output( + &self, + frame: ActivationFrame, + predicted_tokens: Vec, + ) -> Result { + let activation = if self.info.is_final() { + None + } else { + Some(runtime_activation( + frame, + self.info.activation_width as usize, + )?) + }; + Ok(StageExecutionOutput { + activation, + predicted_tokens: if self.info.is_final() { + predicted_tokens + } else { + Vec::new() + }, + }) + } +} + +impl StageEngine for LlamaStageEngine { + fn info(&self) -> &StageEngineInfo { + &self.info + } + + fn execute(&self, request: StageExecutionRequest) -> Result { + let input = self.activation_frame(request.input)?; + let sampling = runtime_sampling_config(request.sampling.as_ref()); + let session_id = request.session_id.to_string(); + let mut runtime = self.runtime()?; + match request.kind { + StageExecutionKind::Prefill => { + let output = runtime.prefill_frame_with_positions( + &session_id, + &request.token_ids, + &request.positions, + input.as_ref(), + )?; + self.output(output, Vec::new()) + } + StageExecutionKind::PrefillFinal if self.info.is_final() => { + let (predicted, output) = runtime.prefill_final_frame_sampled( + &session_id, + &request.token_ids, + &request.positions, + sampling.as_ref(), + input.as_ref(), + )?; + self.output(output, vec![predicted]) + } + StageExecutionKind::PrefillFinal => { + let output = runtime.prefill_frame_with_positions( + &session_id, + &request.token_ids, + &request.positions, + input.as_ref(), + )?; + self.output(output, Vec::new()) + } + StageExecutionKind::Decode => { + let token_id = one_token(&request.token_ids, "decode")?; + let (predicted, output) = runtime.decode_frame_sampled( + &session_id, + token_id, + sampling.as_ref(), + input.as_ref(), + 0, + )?; + self.output(output, vec![predicted]) + } + StageExecutionKind::Verify => { + let (predicted, output) = runtime.verify_frame_sampled( + &session_id, + &request.token_ids, + sampling.as_ref(), + input.as_ref(), + 0, + )?; + self.output(output, predicted) + } + } + } + + fn reset_session(&self, session_id: u64) -> Result<()> { + self.runtime()? + .drop_session_timed(&session_id.to_string())?; + Ok(()) + } + + fn checkpoint_session(&self, session_id: u64) -> Result<()> { + self.runtime()?.checkpoint_session(&session_id.to_string()) + } + + fn restore_session(&self, session_id: u64) -> Result<()> { + self.runtime()?.restore_session(&session_id.to_string()) + } + + fn trim_session(&self, session_id: u64, token_count: u64) -> Result<()> { + self.runtime()? + .trim_session(&session_id.to_string(), token_count) + } +} + +fn one_token(tokens: &[i32], operation: &str) -> Result { + match tokens { + [token] => Ok(*token), + _ => bail!("{operation} requires exactly one token"), + } +} + +fn runtime_sampling_config( + sampling: Option<&skippy_protocol::binary::StageSamplingConfig>, +) -> Option { + let sampling = sampling?; + let mut config = SamplingConfig { + enabled: true, + seed: sampling.seed, + temperature: sampling.temperature, + top_p: sampling.top_p, + top_k: sampling.top_k, + min_p: sampling.min_p, + presence_penalty: sampling.presence_penalty, + frequency_penalty: sampling.frequency_penalty, + repeat_penalty: sampling.repeat_penalty, + penalty_last_n: sampling.penalty_last_n, + ..SamplingConfig::default() + }; + config.logit_bias = sampling + .logit_bias + .iter() + .take(MAX_LOGIT_BIAS) + .map(|source| LogitBias { + token_id: source.token_id, + bias: source.bias, + }) + .collect(); + sampling.enabled().then_some(config) +} + +fn runtime_activation(frame: ActivationFrame, width: usize) -> Result { + ensure!( + frame.desc.layout == RuntimeActivationLayout::TokenMajor, + "llama stage output activation is not token-major" + ); + ensure!( + frame.payload.len() as u64 == frame.desc.payload_bytes, + "llama stage output payload length mismatch" + ); + let token_count = usize::try_from(frame.desc.token_count)?; + let f32_le_bytes = match frame.desc.dtype { + RuntimeActivationDType::F32 => frame.payload, + RuntimeActivationDType::F16 => { + decode_16_bit_activation(&frame.payload, |bytes| f16::from_le_bytes(bytes).to_f32())? + } + RuntimeActivationDType::Bf16 => { + decode_16_bit_activation(&frame.payload, |bytes| bf16::from_le_bytes(bytes).to_f32())? + } + RuntimeActivationDType::Unknown => bail!("llama stage output activation dtype is unknown"), + }; + StageActivation::new(token_count, width, f32_le_bytes) +} + +fn decode_16_bit_activation(payload: &[u8], decode: impl Fn([u8; 2]) -> f32) -> Result> { + ensure!( + payload.len().is_multiple_of(2), + "16-bit activation has odd byte length" + ); + Ok(payload + .chunks_exact(2) + .flat_map(|chunk| decode([chunk[0], chunk[1]]).to_le_bytes()) + .collect()) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn f16_activation_converts_to_f32_contract_bytes() { + let values = [f16::from_f32(1.5), f16::from_f32(-2.0)]; + let payload = values + .iter() + .flat_map(|value| value.to_le_bytes()) + .collect::>(); + let frame = ActivationFrame { + desc: ActivationDesc { + version: 1, + dtype: RuntimeActivationDType::F16, + layout: RuntimeActivationLayout::TokenMajor, + producer_stage_index: 0, + layer_start: 0, + layer_end: 1, + token_count: 1, + sequence_count: 1, + payload_bytes: payload.len() as u64, + flags: 0, + }, + payload, + }; + let activation = runtime_activation(frame, 2).unwrap(); + assert_eq!(activation.values(), vec![1.5, -2.0]); + } +} diff --git a/docs/design/MLX_STAGE_ENGINE_PLAN.md b/docs/design/MLX_STAGE_ENGINE_PLAN.md index 25dff9bebb..7856d0268f 100644 --- a/docs/design/MLX_STAGE_ENGINE_PLAN.md +++ b/docs/design/MLX_STAGE_ENGINE_PLAN.md @@ -618,11 +618,12 @@ the trait in a new `skippy-engine` crate, implement it for the existing behavior change; ship this independently of MLX. Validate with existing `skippy-correctness` and `mic-lab` runs. -> **Partially implemented on this branch.** The engine-neutral crate and an -> additive reduced binary server lane now exist, and MLX uses them for the -> two-process proof. The mature llama runtime has not yet been migrated from -> concrete `RuntimeState`; that compatibility refactor remains before the -> engine-neutral lane can replace the default server internals. +> **Partially implemented on this branch.** The engine-neutral crate, an +> additive reduced binary server lane, and a dense `LlamaStageEngine` adapter +> over the existing `RuntimeState` now exist; MLX uses the same contract for the +> two-process proof. The mature llama server has not yet been switched from its +> concrete `RuntimeState` path because its broader batching, cache, MTP, and +> multimodal surface still needs capability-aware migration. **Phase 2 — Solo MLX serving + JIT quant (the workflow win; lead here).** `MlxStageEngine` as a single-stage/whole-model engine: open/load, session, @@ -758,9 +759,9 @@ Spikes 1 and 2 are more decisive than any standalone token/s benchmark. 1. Add tensor-at-a-time MLX quantization to the materializer and measure peak RSS against the one-source-tensor memory contract. -2. Add the llama `StageEngine` adapter and route normal `skippy-server` launch - through the engine-neutral contract; the MLX two-process binary-wire proof - is complete. +2. Route normal `skippy-server` launch through the engine-neutral contract while + retaining capability-gated llama-only batching/cache/MTP/multimodal paths; + the dense llama adapter and MLX two-process binary-wire proof are complete. 3. Run **Spike 2 (boundary fence)** at frontier residual widths and keep it as a go/no-go gate. 4. Use Nemotron-H as the first frontier-family follow-up already represented in From 93f5d58b3843449dec3f93d7ccca18efc1973c24 Mon Sep 17 00:00:00 2001 From: Michael Neale <14976+michaelneale@users.noreply.github.com> Date: Fri, 17 Jul 2026 18:45:21 +1000 Subject: [PATCH 11/37] feat(mlx): materialize verified safetensors stage slices --- Cargo.lock | 2 + crates/model-artifact/Cargo.toml | 1 + crates/model-artifact/src/lib.rs | 61 ++ crates/model-artifact/src/safetensors.rs | 227 ++++++ crates/model-hf/Cargo.toml | 1 + crates/model-hf/src/lib.rs | 1 + crates/model-hf/src/safetensors_stage/http.rs | 386 ++++++++++ .../model-hf/src/safetensors_stage/layout.rs | 541 ++++++++++++++ .../src/safetensors_stage/materialize.rs | 694 ++++++++++++++++++ crates/model-hf/src/safetensors_stage/mod.rs | 10 + .../model-hf/src/safetensors_stage/types.rs | 151 ++++ 11 files changed, 2075 insertions(+) create mode 100644 crates/model-artifact/src/safetensors.rs create mode 100644 crates/model-hf/src/safetensors_stage/http.rs create mode 100644 crates/model-hf/src/safetensors_stage/layout.rs create mode 100644 crates/model-hf/src/safetensors_stage/materialize.rs create mode 100644 crates/model-hf/src/safetensors_stage/mod.rs create mode 100644 crates/model-hf/src/safetensors_stage/types.rs diff --git a/Cargo.lock b/Cargo.lock index 7e44eb505b..28ecccfc95 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4512,6 +4512,7 @@ dependencies = [ "async-trait", "model-ref", "serde", + "serde_json", "tokio", ] @@ -4526,6 +4527,7 @@ dependencies = [ "hf-hub", "model-artifact", "model-ref", + "reqwest 0.12.28", "serde", "serde_json", "serial_test", diff --git a/crates/model-artifact/Cargo.toml b/crates/model-artifact/Cargo.toml index f4e5c567f6..6e8d9a1288 100644 --- a/crates/model-artifact/Cargo.toml +++ b/crates/model-artifact/Cargo.toml @@ -13,6 +13,7 @@ anyhow.workspace = true async-trait = "0.1" model-ref = { path = "../model-ref", version = "0.72.1" } serde.workspace = true +serde_json.workspace = true [dev-dependencies] tokio = { version = "1", features = ["macros", "rt"] } diff --git a/crates/model-artifact/src/lib.rs b/crates/model-artifact/src/lib.rs index 1e2016bfa0..45602aafa9 100644 --- a/crates/model-artifact/src/lib.rs +++ b/crates/model-artifact/src/lib.rs @@ -1,4 +1,5 @@ pub mod gguf; +pub mod safetensors; use std::path::Path; @@ -253,6 +254,26 @@ fn artifact_file_set(primary_file: &str, files: &[ModelArtifactFile]) -> Vec>(); + shards.sort_by(|left, right| left.path.cmp(&right.path)); + if !shards.is_empty() { + return shards; + } + } + vec![ files .iter() @@ -338,6 +359,20 @@ fn split_safetensors_shard_info(stem: &str) -> Option<(&str, &str, &str)> { Some((prefix, part, total)) } +fn split_safetensors_path_info(path: &str) -> Option<(String, String, String, String)> { + let path = Path::new(path); + let stem = path.file_name()?.to_str()?.strip_suffix(".safetensors")?; + let (prefix, part, total) = split_safetensors_shard_info(stem)?; + Some(( + path.parent() + .filter(|parent| !parent.as_os_str().is_empty()) + .map_or_else(String::new, |parent| parent.to_string_lossy().into_owned()), + prefix.to_string(), + part.to_string(), + total.to_string(), + )) +} + #[cfg(test)] mod tests { use super::*; @@ -512,6 +547,32 @@ mod tests { ); } + #[test] + fn public_artifact_set_returns_only_matching_safetensors_shards() { + let files = files(&[ + "weights/model-00002-of-00003.safetensors", + "weights/model-00001-of-00003.safetensors", + "weights/model-00003-of-00003.safetensors", + "other/model-00001-of-00003.safetensors", + "weights/adapter-00001-of-00003.safetensors", + "model.safetensors.index.json", + ]); + + let shards = artifact_files_for_primary("weights/model-00001-of-00003.safetensors", &files); + + assert_eq!( + shards + .iter() + .map(|file| file.path.as_str()) + .collect::>(), + vec![ + "weights/model-00001-of-00003.safetensors", + "weights/model-00002-of-00003.safetensors", + "weights/model-00003-of-00003.safetensors", + ] + ); + } + #[tokio::test] async fn accepts_revisioned_selector_refs() { let repository = repo(vec!["Model-Q4_K_M.gguf"]); diff --git a/crates/model-artifact/src/safetensors.rs b/crates/model-artifact/src/safetensors.rs new file mode 100644 index 0000000000..0ae8c7d8a2 --- /dev/null +++ b/crates/model-artifact/src/safetensors.rs @@ -0,0 +1,227 @@ +use std::collections::BTreeMap; + +use anyhow::{Context, Result, bail, ensure}; +use serde::{Deserialize, Serialize}; + +const MAX_TENSOR_COUNT: usize = 2_000_000; +const MAX_TENSOR_RANK: usize = 16; + +#[derive(Clone, Debug, PartialEq, Eq, Deserialize, Serialize)] +pub struct TensorHeader { + pub dtype: String, + pub shape: Vec, + pub data_offsets: [u64; 2], +} + +#[derive(Clone, Debug, Default, PartialEq, Eq, Deserialize, Serialize)] +pub struct IndexMetadata { + pub total_size: Option, +} + +#[derive(Clone, Debug, PartialEq, Eq, Deserialize, Serialize)] +pub struct SafetensorsIndex { + #[serde(default)] + pub metadata: IndexMetadata, + pub weight_map: BTreeMap, +} + +#[derive(Clone, Debug, PartialEq, Eq, Deserialize, Serialize)] +pub struct LlamaConfig { + pub model_type: String, + pub hidden_size: u64, + pub num_hidden_layers: u32, + #[serde(default)] + pub tie_word_embeddings: bool, +} + +pub fn parse_llama_config(bytes: &[u8]) -> Result { + let config: LlamaConfig = + serde_json::from_slice(bytes).context("parse SafeTensors model config")?; + ensure!( + config.model_type == "llama", + "MLX partial SafeTensors currently supports model_type=llama, got {:?}", + config.model_type + ); + ensure!(config.hidden_size > 0, "Llama hidden_size must be non-zero"); + ensure!( + config.num_hidden_layers > 0, + "Llama num_hidden_layers must be non-zero" + ); + Ok(config) +} + +pub fn parse_index(bytes: &[u8]) -> Result { + let index: SafetensorsIndex = + serde_json::from_slice(bytes).context("parse SafeTensors index")?; + ensure!( + !index.weight_map.is_empty(), + "SafeTensors index has no tensors" + ); + ensure!( + index.weight_map.len() <= MAX_TENSOR_COUNT, + "SafeTensors index has too many tensors" + ); + ensure!( + index + .weight_map + .iter() + .all(|(name, file)| !name.is_empty() && !file.is_empty()), + "SafeTensors index contains an empty tensor or shard name" + ); + Ok(index) +} + +pub fn parse_header(bytes: &[u8], data_bytes: u64) -> Result> { + let raw: BTreeMap = + serde_json::from_slice(bytes).context("parse SafeTensors header")?; + ensure!( + raw.len() <= MAX_TENSOR_COUNT.saturating_add(1), + "SafeTensors file has too many tensor entries" + ); + let tensors = raw + .into_iter() + .filter(|(name, _)| name != "__metadata__") + .map(|(name, value)| { + ensure!( + !name.is_empty(), + "SafeTensors tensor name must not be empty" + ); + serde_json::from_value(value) + .map(|tensor| (name.clone(), tensor)) + .with_context(|| format!("parse SafeTensors tensor header {name}")) + }) + .collect::>>()?; + validate_headers(&tensors, data_bytes)?; + Ok(tensors) +} + +fn validate_headers(tensors: &BTreeMap, data_bytes: u64) -> Result<()> { + ensure!( + !tensors.is_empty(), + "SafeTensors file has no tensor entries" + ); + let mut extents = Vec::with_capacity(tensors.len()); + for (name, tensor) in tensors { + validate_tensor(name, tensor, data_bytes)?; + extents.push((tensor.data_offsets[0], tensor.data_offsets[1], name)); + } + extents.sort_by_key(|(start, _, _)| *start); + let mut previous_end = 0; + for (start, end, name) in extents { + ensure!( + start == previous_end, + "SafeTensors tensor {name} leaves a gap or overlaps another tensor" + ); + previous_end = end; + } + ensure!( + previous_end == data_bytes, + "SafeTensors tensor data does not cover the declared data section" + ); + Ok(()) +} + +fn validate_tensor(name: &str, tensor: &TensorHeader, data_bytes: u64) -> Result<()> { + ensure!( + tensor.shape.len() <= MAX_TENSOR_RANK, + "SafeTensors tensor {name} exceeds maximum rank {MAX_TENSOR_RANK}" + ); + ensure!( + tensor.data_offsets[0] <= tensor.data_offsets[1], + "invalid data offsets for SafeTensors tensor {name}" + ); + ensure!( + tensor.data_offsets[1] <= data_bytes, + "SafeTensors tensor {name} exceeds the data section" + ); + let elements = tensor.shape.iter().try_fold(1_u64, |count, dimension| { + count + .checked_mul(*dimension) + .context("SafeTensors tensor element count overflow") + })?; + let expected_bytes = elements + .checked_mul(dtype_bytes(&tensor.dtype)?) + .context("SafeTensors tensor byte count overflow")?; + let actual_bytes = tensor.data_offsets[1] - tensor.data_offsets[0]; + ensure!( + expected_bytes == actual_bytes, + "SafeTensors tensor {name} has {actual_bytes} bytes but its dtype and shape require {expected_bytes}" + ); + Ok(()) +} + +fn dtype_bytes(dtype: &str) -> Result { + match dtype { + "BOOL" | "U8" | "I8" | "F8_E4M3" | "F8_E5M2" | "F8_E8M0" => Ok(1), + "U16" | "I16" | "F16" | "BF16" => Ok(2), + "U32" | "I32" | "F32" => Ok(4), + "U64" | "I64" | "F64" => Ok(8), + _ => bail!("unsupported SafeTensors dtype {dtype:?}"), + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn parses_and_validates_a_complete_header() { + let header = br#"{ + "a":{"dtype":"F16","shape":[2],"data_offsets":[0,4]}, + "b":{"dtype":"F32","shape":[1],"data_offsets":[4,8]}, + "__metadata__":{"format":"pt"} + }"#; + + let tensors = parse_header(header, 8).unwrap(); + + assert_eq!(tensors.len(), 2); + assert_eq!(tensors["a"].shape, vec![2]); + } + + #[test] + fn rejects_shape_byte_mismatch() { + let header = br#"{ + "bad":{"dtype":"F16","shape":[3],"data_offsets":[0,4]} + }"#; + + assert!(parse_header(header, 4).is_err()); + } + + #[test] + fn rejects_gaps_and_overlaps() { + let header = br#"{ + "a":{"dtype":"F16","shape":[1],"data_offsets":[0,2]}, + "b":{"dtype":"F16","shape":[1],"data_offsets":[3,5]} + }"#; + + assert!(parse_header(header, 5).is_err()); + } + + #[test] + fn parses_nonempty_index() { + let index = br#"{ + "metadata":{"total_size":4}, + "weight_map":{"model.layers.0.weight":"model-00001-of-00002.safetensors"} + }"#; + + assert_eq!(parse_index(index).unwrap().metadata.total_size, Some(4)); + } + + #[test] + fn accepts_only_supported_llama_config() { + let llama = br#"{ + "model_type":"llama", + "hidden_size":576, + "num_hidden_layers":30, + "tie_word_embeddings":true + }"#; + let qwen = br#"{ + "model_type":"qwen3", + "hidden_size":1024, + "num_hidden_layers":28 + }"#; + + assert_eq!(parse_llama_config(llama).unwrap().num_hidden_layers, 30); + assert!(parse_llama_config(qwen).is_err()); + } +} diff --git a/crates/model-hf/Cargo.toml b/crates/model-hf/Cargo.toml index 2ba5b88cee..a8a559603b 100644 --- a/crates/model-hf/Cargo.toml +++ b/crates/model-hf/Cargo.toml @@ -15,6 +15,7 @@ dirs = "6.0.0" hf_hub = { package = "hf-hub", version = "1.0.0-rc.1", default-features = false, features = ["blocking"] } model-artifact = { path = "../model-artifact", version = "0.72.1" } model-ref = { path = "../model-ref", version = "0.72.1" } +reqwest = { version = "0.12", default-features = false, features = ["blocking", "rustls-tls"] } serde.workspace = true serde_json.workspace = true sha2 = "0.10" diff --git a/crates/model-hf/src/lib.rs b/crates/model-hf/src/lib.rs index 8594db0620..a9a4811fb6 100644 --- a/crates/model-hf/src/lib.rs +++ b/crates/model-hf/src/lib.rs @@ -1,3 +1,4 @@ +pub mod safetensors_stage; pub mod store; use std::{ diff --git a/crates/model-hf/src/safetensors_stage/http.rs b/crates/model-hf/src/safetensors_stage/http.rs new file mode 100644 index 0000000000..8939d3b1f7 --- /dev/null +++ b/crates/model-hf/src/safetensors_stage/http.rs @@ -0,0 +1,386 @@ +use std::{ + io::{Read, Write}, + ops::Range, + path::{Component, Path}, + time::Duration, +}; + +use anyhow::{Context, Result, anyhow, ensure}; +use reqwest::{ + StatusCode, Url, + blocking::{Client, RequestBuilder, Response}, + header::{ACCEPT_ENCODING, AUTHORIZATION, CONTENT_LENGTH, CONTENT_RANGE, ETAG, RANGE}, + redirect::Policy, +}; + +const DEFAULT_ENDPOINT: &str = "https://huggingface.co"; + +#[derive(Clone)] +pub(crate) struct RemoteSource { + client: Client, + endpoint: Url, + token: Option, +} + +pub(crate) struct ExactRangeResponse { + response: Response, + pub total_file_bytes: u64, + expected_bytes: u64, + etag: Option, +} + +pub(crate) struct RemoteFile { + pub bytes: Vec, + pub etag: Option, +} + +impl RemoteSource { + pub fn new(endpoint: Option<&str>, token: Option) -> Result { + let endpoint = Url::parse(endpoint.unwrap_or(DEFAULT_ENDPOINT)) + .context("parse Hugging Face endpoint")?; + let client = Client::builder() + .connect_timeout(Duration::from_secs(30)) + .timeout(Duration::from_secs(120)) + .redirect(Policy::limited(10)) + .user_agent("mesh-llm-safetensors-stage/1") + .build() + .context("build SafeTensors stage HTTP client")?; + Ok(Self { + client, + endpoint, + token, + }) + } + + pub fn endpoint(&self) -> &str { + self.endpoint.as_str() + } + + pub fn url(&self, repo: &str, revision: &str, file: &str) -> Result { + validate_relative_file(file)?; + let mut url = self.endpoint.clone(); + { + let mut segments = url + .path_segments_mut() + .map_err(|_| anyhow!("Hugging Face endpoint cannot be a base URL"))?; + segments.pop_if_empty(); + segments.extend(repo.split('/')); + segments.push("resolve"); + segments.push(revision); + segments.extend(file.split('/')); + } + Ok(url) + } + + pub fn optional_small_file(&self, url: Url, max_bytes: u64) -> Result> { + let response = self + .authorized(self.client.get(url)) + .send() + .context("send Hugging Face metadata request")?; + if response.status() == StatusCode::NOT_FOUND { + return Ok(None); + } + self.read_small_response(response, max_bytes).map(Some) + } + + pub fn small_file(&self, url: Url, max_bytes: u64) -> Result { + let response = self + .authorized(self.client.get(url)) + .send() + .context("send Hugging Face metadata request")?; + self.read_small_response(response, max_bytes) + } + + pub fn exact_range(&self, url: Url, range: Range) -> Result { + ensure!(range.start < range.end, "HTTP byte range must not be empty"); + let expected_bytes = range.end - range.start; + let range_header = format!("bytes={}-{}", range.start, range.end - 1); + let response = self + .authorized( + self.client + .get(url) + .header(RANGE, range_header.clone()) + .header(ACCEPT_ENCODING, "identity"), + ) + .send() + .with_context(|| format!("request HTTP range {range_header}"))?; + ensure!( + response.status() == StatusCode::PARTIAL_CONTENT, + "server did not honor {range_header}; status was {} (refusing a possible full-shard download)", + response.status() + ); + let content_range = response + .headers() + .get(CONTENT_RANGE) + .context("206 response omitted Content-Range")? + .to_str() + .context("Content-Range is not valid ASCII")?; + let parsed = parse_content_range(content_range)?; + ensure!( + parsed.start == range.start && parsed.end_exclusive == range.end, + "Content-Range {content_range:?} did not match requested {range_header}" + ); + if let Some(length) = response + .headers() + .get(CONTENT_LENGTH) + .and_then(|value| value.to_str().ok()) + .and_then(|value| value.parse::().ok()) + { + ensure!( + length == expected_bytes, + "HTTP range Content-Length was {length}, expected {expected_bytes}" + ); + } + let etag = header_string(&response, ETAG)?; + Ok(ExactRangeResponse { + response, + total_file_bytes: parsed.total_file_bytes, + expected_bytes, + etag, + }) + } + + fn read_small_response(&self, response: Response, max_bytes: u64) -> Result { + let response = response + .error_for_status() + .context("download Hugging Face metadata file")?; + let etag = header_string(&response, ETAG)?; + if let Some(length) = response + .headers() + .get(CONTENT_LENGTH) + .and_then(|value| value.to_str().ok()) + .and_then(|value| value.parse::().ok()) + { + ensure!( + length <= max_bytes, + "metadata file is too large: {length} bytes" + ); + } + let limit = max_bytes + .checked_add(1) + .context("metadata byte limit overflow")?; + let mut bytes = Vec::new(); + response + .take(limit) + .read_to_end(&mut bytes) + .context("read metadata response")?; + ensure!( + bytes.len() as u64 <= max_bytes, + "metadata response exceeded {max_bytes} bytes" + ); + Ok(RemoteFile { bytes, etag }) + } + + fn authorized(&self, builder: RequestBuilder) -> RequestBuilder { + match &self.token { + Some(token) => builder.header(AUTHORIZATION, format!("Bearer {token}")), + None => builder, + } + } +} + +impl ExactRangeResponse { + pub fn etag(&self) -> Option<&str> { + self.etag.as_deref() + } + + pub fn into_bytes(mut self) -> Result> { + 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 { + let written = + std::io::copy(&mut self.response, writer).context("stream HTTP tensor range")?; + ensure!( + written == self.expected_bytes, + "HTTP range returned {written} bytes, expected {}", + self.expected_bytes + ); + Ok(written) + } +} + +fn header_string(response: &Response, name: reqwest::header::HeaderName) -> Result> { + response + .headers() + .get(name) + .map(|value| { + value + .to_str() + .context("HTTP identity header is not valid ASCII") + .map(str::to_string) + }) + .transpose() +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +struct ParsedContentRange { + start: u64, + end_exclusive: u64, + total_file_bytes: u64, +} + +fn parse_content_range(value: &str) -> Result { + let value = value + .strip_prefix("bytes ") + .context("Content-Range must use bytes")?; + let (range, total) = value + .split_once('/') + .context("Content-Range omitted total size")?; + let (start, end_inclusive) = range + .split_once('-') + .context("Content-Range omitted byte bounds")?; + let start = start.parse::().context("parse Content-Range start")?; + let end_inclusive = end_inclusive + .parse::() + .context("parse Content-Range end")?; + let total_file_bytes = total.parse::().context("parse Content-Range total")?; + let end_exclusive = end_inclusive + .checked_add(1) + .context("Content-Range end overflow")?; + ensure!(start < end_exclusive, "Content-Range is empty"); + ensure!( + end_exclusive <= total_file_bytes, + "Content-Range exceeds total file size" + ); + Ok(ParsedContentRange { + start, + end_exclusive, + total_file_bytes, + }) +} + +fn validate_relative_file(file: &str) -> Result<()> { + ensure!( + !file.is_empty() + && !file + .chars() + .any(|character| matches!(character, '\\' | '?' | '#')), + "unsafe Hugging Face repository file path {file:?}" + ); + let path = Path::new(file); + ensure!( + path.components() + .all(|component| matches!(component, Component::Normal(_))), + "unsafe Hugging Face repository file path {file:?}" + ); + Ok(()) +} + +#[cfg(test)] +mod tests { + use std::{ + io::{Read, Write}, + net::TcpListener, + thread, + }; + + use super::*; + + #[test] + fn parses_exact_content_range() { + assert_eq!( + parse_content_range("bytes 8-15/100").unwrap(), + ParsedContentRange { + start: 8, + end_exclusive: 16, + total_file_bytes: 100, + } + ); + } + + #[test] + fn rejects_unsafe_repository_paths() { + assert!(validate_relative_file("../secret").is_err()); + assert!(validate_relative_file("/absolute").is_err()); + assert!(validate_relative_file("weights/model.safetensors").is_ok()); + } + + #[test] + fn rejects_server_that_ignores_range() { + let endpoint = serve_once(http_response("200 OK", &[], b"full")); + let remote = RemoteSource::new(Some(&endpoint), None).unwrap(); + let url = remote + .url("org/model", &"a".repeat(40), "model.safetensors") + .unwrap(); + + let error = remote.exact_range(url, 0..4).err().unwrap(); + + assert!(error.to_string().contains("did not honor")); + } + + #[test] + fn rejects_mismatched_content_range() { + let endpoint = serve_once(http_response( + "206 Partial Content", + &[("Content-Range", "bytes 1-4/10")], + b"four", + )); + let remote = RemoteSource::new(Some(&endpoint), None).unwrap(); + let url = remote + .url("org/model", &"a".repeat(40), "model.safetensors") + .unwrap(); + + let error = remote.exact_range(url, 0..4).err().unwrap(); + + assert!(error.to_string().contains("did not match")); + } + + #[test] + fn rejects_truncated_range_body() { + let response = b"HTTP/1.1 206 Partial Content\r\n\ + Content-Length: 4\r\n\ + Content-Range: bytes 0-3/10\r\n\ + Connection: close\r\n\r\nxx" + .to_vec(); + let endpoint = serve_once(response); + let remote = RemoteSource::new(Some(&endpoint), None).unwrap(); + let url = remote + .url("org/model", &"a".repeat(40), "model.safetensors") + .unwrap(); + + let error = remote + .exact_range(url, 0..4) + .and_then(ExactRangeResponse::into_bytes) + .unwrap_err(); + + assert!(format!("{error:#}").contains("HTTP range")); + } + + fn serve_once(response: Vec) -> String { + let listener = TcpListener::bind("127.0.0.1:0").unwrap(); + let address = listener.local_addr().unwrap(); + thread::spawn(move || { + let (mut stream, _) = listener.accept().unwrap(); + let mut request = [0_u8; 4096]; + let _ = stream.read(&mut request); + let _ = stream.write_all(&response); + }); + format!("http://{address}") + } + + fn http_response(status: &str, headers: &[(&str, &str)], body: &[u8]) -> Vec { + let headers = headers + .iter() + .map(|(name, value)| format!("{name}: {value}\r\n")) + .collect::(); + format!( + "HTTP/1.1 {status}\r\nContent-Length: {}\r\n{headers}Connection: close\r\n\r\n", + body.len() + ) + .into_bytes() + .into_iter() + .chain(body.iter().copied()) + .collect() + } +} diff --git a/crates/model-hf/src/safetensors_stage/layout.rs b/crates/model-hf/src/safetensors_stage/layout.rs new file mode 100644 index 0000000000..a7af95ee98 --- /dev/null +++ b/crates/model-hf/src/safetensors_stage/layout.rs @@ -0,0 +1,541 @@ +use std::collections::{BTreeMap, BTreeSet}; + +use anyhow::{Context, Result, anyhow, ensure}; +use model_artifact::safetensors::{ + IndexMetadata, LlamaConfig, SafetensorsIndex, TensorHeader, parse_header, parse_index, + parse_llama_config, +}; +use sha2::{Digest, Sha256}; + +use super::{ + http::RemoteSource, + types::{ + ByteRange, PreparedStage, SafetensorsShardPlan, SafetensorsSourceShard, + SafetensorsStagePlan, SafetensorsStageRequest, SelectedTensor, + }, +}; + +pub(crate) const MAX_INDEX_BYTES: u64 = 64 * 1024 * 1024; +const MAX_HEADER_BYTES: u64 = 256 * 1024 * 1024; + +#[derive(Clone, Debug)] +struct RemoteHeader { + header_len: u64, + file_bytes: u64, + etag: Option, + tensors: BTreeMap, +} + +struct CheckpointLayout { + index: SafetensorsIndex, + index_bytes: u64, + index_sha256: Option, + index_etag: Option, + headers: BTreeMap, +} + +pub(crate) fn prepare( + remote: &RemoteSource, + request: &SafetensorsStageRequest, +) -> Result { + let config_url = remote.url(&request.repo, &request.revision, "config.json")?; + let config = remote + .small_file(config_url, MAX_INDEX_BYTES) + .context("download SafeTensors model config")?; + let model_config = parse_llama_config(&config.bytes)?; + validate_layer_range(request, &model_config)?; + let config_sha256 = sha256_hex(&config.bytes); + let mut selection_request = request.clone(); + add_required_prefixes(&mut selection_request, &model_config); + + let mut layout = load_checkpoint_layout(remote, request)?; + validate_layer_coverage(&layout.index, request)?; + let selected = select_tensors(&layout.index.weight_map, &selection_request); + ensure!( + !selected.is_empty(), + "no tensors matched layers {}..{} or requested prefixes", + request.layer_start, + request.layer_end + ); + validate_required_tensors(&selected, request, &model_config)?; + + let mut by_shard: BTreeMap<&str, BTreeSet<&str>> = BTreeMap::new(); + for name in &selected { + let shard = layout + .index + .weight_map + .get(*name) + .with_context(|| format!("selected tensor {name} is absent from weight map"))?; + by_shard.entry(shard).or_default().insert(*name); + } + + let mut shards = Vec::with_capacity(by_shard.len()); + let mut tensors = Vec::with_capacity(selected.len()); + for (file, names) in by_shard { + if !layout.headers.contains_key(file) { + let header = fetch_safetensor_header(remote, request, file) + .with_context(|| format!("inspect {file}"))?; + layout.headers.insert(file.to_string(), header); + } + let header = layout + .headers + .get(file) + .with_context(|| format!("missing inspected header for {file}"))?; + shards.push(plan_shard(file, header, &names)?); + tensors.extend(selected_tensors(file, header, &names)?); + } + let source_shards = shards + .iter() + .map(|shard| { + let header = layout + .headers + .get(&shard.file) + .expect("planned shard has inspected header"); + SafetensorsSourceShard { + file: shard.file.clone(), + file_bytes: header.file_bytes, + etag: header.etag.clone(), + } + }) + .collect(); + let plan = summarize_plan( + &selection_request, + &layout.index, + config.bytes.len() as u64, + layout.index_bytes, + shards, + )?; + Ok(PreparedStage { + plan, + tensors, + config: config.bytes, + config_sha256, + config_etag: config.etag, + index_sha256: layout.index_sha256, + index_etag: layout.index_etag, + source_shards, + }) +} + +fn validate_layer_range(request: &SafetensorsStageRequest, config: &LlamaConfig) -> Result<()> { + ensure!( + request.layer_end <= config.num_hidden_layers, + "stage layer end {} exceeds model layer count {}", + request.layer_end, + config.num_hidden_layers + ); + Ok(()) +} + +fn add_required_prefixes(request: &mut SafetensorsStageRequest, config: &LlamaConfig) { + if request.layer_start == 0 || request.layer_end == config.num_hidden_layers { + request + .include_prefixes + .push("model.embed_tokens.".to_string()); + } + if request.layer_end == config.num_hidden_layers { + request.include_prefixes.push("model.norm.".to_string()); + request.include_prefixes.push("lm_head.".to_string()); + } + request.include_prefixes.sort(); + request.include_prefixes.dedup(); +} + +fn validate_layer_coverage( + index: &SafetensorsIndex, + request: &SafetensorsStageRequest, +) -> Result<()> { + for layer in request.layer_start..request.layer_end { + ensure!( + index + .weight_map + .keys() + .any(|name| layer_index(name) == Some(layer)), + "SafeTensors checkpoint has no tensors for requested layer {layer}" + ); + } + Ok(()) +} + +fn validate_required_tensors( + selected: &BTreeSet<&str>, + request: &SafetensorsStageRequest, + config: &LlamaConfig, +) -> Result<()> { + let has_prefix = |prefix: &str| selected.iter().any(|name| name.starts_with(prefix)); + if request.layer_start == 0 { + ensure!( + has_prefix("model.embed_tokens."), + "first MLX stage requires model.embed_tokens tensors" + ); + } + if request.layer_end == config.num_hidden_layers { + ensure!( + has_prefix("model.norm."), + "final MLX stage requires model.norm tensors" + ); + ensure!( + has_prefix("lm_head.") || has_prefix("model.embed_tokens."), + "final MLX stage requires lm_head or tied embedding tensors" + ); + } + Ok(()) +} + +fn load_checkpoint_layout( + remote: &RemoteSource, + request: &SafetensorsStageRequest, +) -> Result { + let index_url = remote.url( + &request.repo, + &request.revision, + "model.safetensors.index.json", + )?; + if let Some(index_file) = remote.optional_small_file(index_url, MAX_INDEX_BYTES)? { + let index = parse_index(&index_file.bytes)?; + return Ok(CheckpointLayout { + index, + index_bytes: index_file.bytes.len() as u64, + index_sha256: Some(sha256_hex(&index_file.bytes)), + index_etag: index_file.etag, + headers: BTreeMap::new(), + }); + } + + let file = "model.safetensors"; + let header = fetch_safetensor_header(remote, request, file) + .context("inspect unsharded SafeTensors checkpoint")?; + let total_size = header + .tensors + .values() + .map(|tensor| tensor.data_offsets[1]) + .max(); + let weight_map = header + .tensors + .keys() + .map(|name| (name.clone(), file.to_string())) + .collect(); + Ok(CheckpointLayout { + index: SafetensorsIndex { + metadata: IndexMetadata { total_size }, + weight_map, + }, + index_bytes: 0, + index_sha256: None, + index_etag: None, + headers: BTreeMap::from([(file.to_string(), header)]), + }) +} + +fn fetch_safetensor_header( + remote: &RemoteSource, + request: &SafetensorsStageRequest, + file: &str, +) -> Result { + let url = remote.url(&request.repo, &request.revision, file)?; + let len_response = remote.exact_range(url.clone(), 0..8)?; + let file_bytes = len_response.total_file_bytes; + let len_etag = len_response.etag().map(str::to_string); + let len_bytes = len_response.into_bytes()?; + let header_len = u64::from_le_bytes( + len_bytes + .as_slice() + .try_into() + .map_err(|_| anyhow!("invalid 8-byte SafeTensors header length"))?, + ); + ensure!( + header_len <= MAX_HEADER_BYTES, + "SafeTensors header is unexpectedly large: {header_len} bytes" + ); + let header_end = 8_u64 + .checked_add(header_len) + .context("SafeTensors header range overflow")?; + ensure!( + header_end <= file_bytes, + "SafeTensors header exceeds source file length" + ); + let header_response = remote.exact_range(url, 8..header_end)?; + ensure!( + header_response.total_file_bytes == file_bytes, + "source file size changed while reading SafeTensors header" + ); + ensure_matching_etag(len_etag.as_deref(), header_response.etag(), file)?; + let etag = header_response.etag().map(str::to_string).or(len_etag); + let header_bytes = header_response.into_bytes()?; + let data_bytes = file_bytes - header_end; + let tensors = parse_header(&header_bytes, data_bytes)?; + Ok(RemoteHeader { + header_len, + file_bytes, + etag, + tensors, + }) +} + +fn ensure_matching_etag(left: Option<&str>, right: Option<&str>, file: &str) -> Result<()> { + if let (Some(left), Some(right)) = (left, right) { + ensure!( + left == right, + "source identity changed while reading SafeTensors shard {file}" + ); + } + Ok(()) +} + +fn select_tensors<'a>( + weight_map: &'a BTreeMap, + request: &SafetensorsStageRequest, +) -> BTreeSet<&'a str> { + weight_map + .keys() + .filter(|name| { + layer_index(name) + .is_some_and(|layer| layer >= request.layer_start && layer < request.layer_end) + || request + .include_prefixes + .iter() + .any(|prefix| name.starts_with(prefix)) + }) + .map(String::as_str) + .collect() +} + +fn layer_index(name: &str) -> Option { + name.strip_prefix("model.layers.")? + .split_once('.')? + .0 + .parse() + .ok() +} + +fn plan_shard( + file: &str, + header: &RemoteHeader, + selected: &BTreeSet<&str>, +) -> Result { + let data_start = 8_u64 + .checked_add(header.header_len) + .context("SafeTensors data offset overflow")?; + let mut ranges = Vec::with_capacity(selected.len()); + let mut selected_tensor_bytes = 0_u64; + let mut largest_selected_tensor_bytes = 0_u64; + for name in selected { + let tensor = header + .tensors + .get(*name) + .with_context(|| format!("weight-map tensor {name} is absent from {file}"))?; + let start = data_start + .checked_add(tensor.data_offsets[0]) + .with_context(|| format!("absolute offset overflow for {name}"))?; + let end_exclusive = data_start + .checked_add(tensor.data_offsets[1]) + .with_context(|| format!("absolute offset overflow for {name}"))?; + let tensor_bytes = end_exclusive - start; + selected_tensor_bytes = selected_tensor_bytes + .checked_add(tensor_bytes) + .context("selected tensor byte count overflow")?; + largest_selected_tensor_bytes = largest_selected_tensor_bytes.max(tensor_bytes); + ranges.push(ByteRange { + start, + end_exclusive, + }); + } + let ranges = coalesce_contiguous_ranges(ranges); + Ok(SafetensorsShardPlan { + file: file.to_string(), + file_bytes: header.file_bytes, + header_probe_bytes: data_start, + selected_tensor_count: selected.len(), + selected_tensor_bytes, + largest_selected_tensor_bytes, + ranges, + }) +} + +fn selected_tensors( + file: &str, + header: &RemoteHeader, + selected: &BTreeSet<&str>, +) -> Result> { + let data_start = 8_u64 + .checked_add(header.header_len) + .context("SafeTensors data offset overflow")?; + selected + .iter() + .map(|name| { + let tensor = header + .tensors + .get(*name) + .with_context(|| format!("weight-map tensor {name} is absent from {file}"))?; + Ok(SelectedTensor { + name: (*name).to_string(), + source_file: file.to_string(), + source_range: ByteRange { + start: data_start + .checked_add(tensor.data_offsets[0]) + .with_context(|| format!("absolute offset overflow for {name}"))?, + end_exclusive: data_start + .checked_add(tensor.data_offsets[1]) + .with_context(|| format!("absolute offset overflow for {name}"))?, + }, + header: tensor.clone(), + }) + }) + .collect() +} + +fn coalesce_contiguous_ranges(mut ranges: Vec) -> Vec { + ranges.sort_by_key(|range| range.start); + let mut merged: Vec = Vec::with_capacity(ranges.len()); + for range in ranges { + if let Some(previous) = merged.last_mut() + && range.start == previous.end_exclusive + { + previous.end_exclusive = range.end_exclusive; + } else { + merged.push(range); + } + } + merged +} + +fn summarize_plan( + request: &SafetensorsStageRequest, + index: &SafetensorsIndex, + config_bytes: u64, + index_bytes: u64, + shards: Vec, +) -> Result { + let selected_tensor_count = shards.iter().map(|shard| shard.selected_tensor_count).sum(); + let selected_tensor_bytes = + checked_sum(shards.iter().map(|shard| shard.selected_tensor_bytes))?; + let largest_selected_tensor_bytes = shards + .iter() + .map(|shard| shard.largest_selected_tensor_bytes) + .max() + .unwrap_or(0); + let source_shard_bytes = checked_sum(shards.iter().map(|shard| shard.file_bytes))?; + let range_request_count = shards.iter().map(|shard| shard.ranges.len()).sum(); + let range_payload_bytes = checked_sum( + shards + .iter() + .flat_map(|shard| shard.ranges.iter()) + .map(ByteRange::len), + )?; + let header_probe_bytes = checked_sum(shards.iter().map(|shard| shard.header_probe_bytes))?; + let planned_download_bytes = config_bytes + .checked_add(index_bytes) + .and_then(|bytes| bytes.checked_add(header_probe_bytes)) + .and_then(|bytes| bytes.checked_add(range_payload_bytes)) + .context("planned download byte count overflow")?; + Ok(SafetensorsStagePlan { + repo: request.repo.clone(), + revision: request.revision.clone(), + layer_start: request.layer_start, + layer_end: request.layer_end, + include_prefixes: request.include_prefixes.clone(), + total_model_tensor_bytes: index.metadata.total_size, + config_bytes, + index_bytes, + selected_tensor_count, + selected_tensor_bytes, + largest_selected_tensor_bytes, + source_shard_count: shards.len(), + source_shard_bytes, + range_request_count, + range_payload_bytes, + header_probe_bytes, + planned_download_bytes, + source_shard_bytes_avoided: source_shard_bytes + .saturating_sub(header_probe_bytes + range_payload_bytes), + full_model_tensor_bytes_avoided: index + .metadata + .total_size + .map(|total| total.saturating_sub(selected_tensor_bytes)), + shards, + }) +} + +fn sha256_hex(bytes: &[u8]) -> String { + format!("{:x}", Sha256::digest(bytes)) +} + +fn checked_sum(values: impl IntoIterator) -> Result { + values.into_iter().try_fold(0_u64, |total, value| { + total.checked_add(value).context("byte count overflow") + }) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn recognizes_only_llama_layer_paths() { + assert_eq!(layer_index("model.layers.42.mlp.up_proj.weight"), Some(42)); + assert_eq!(layer_index("transformer.h.7.attn.weight"), None); + assert_eq!(layer_index("model.embed_tokens.weight"), None); + } + + #[test] + fn coalesces_only_contiguous_ranges() { + assert_eq!( + coalesce_contiguous_ranges(vec![ + ByteRange { + start: 20, + end_exclusive: 30, + }, + ByteRange { + start: 0, + end_exclusive: 10, + }, + ByteRange { + start: 10, + end_exclusive: 20, + }, + ByteRange { + start: 32, + end_exclusive: 40, + }, + ]), + vec![ + ByteRange { + start: 0, + end_exclusive: 30, + }, + ByteRange { + start: 32, + end_exclusive: 40, + }, + ] + ); + } + + #[test] + fn assigns_embedding_and_readout_tensors_to_final_llama_stage() { + let mut request = SafetensorsStageRequest { + repo: "org/model".to_string(), + revision: "a".repeat(40), + layer_start: 1, + layer_end: 2, + include_prefixes: Vec::new(), + }; + let config = LlamaConfig { + model_type: "llama".to_string(), + hidden_size: 2, + num_hidden_layers: 2, + tie_word_embeddings: true, + }; + + add_required_prefixes(&mut request, &config); + + assert_eq!( + request.include_prefixes, + vec![ + "lm_head.".to_string(), + "model.embed_tokens.".to_string(), + "model.norm.".to_string(), + ] + ); + } +} diff --git a/crates/model-hf/src/safetensors_stage/materialize.rs b/crates/model-hf/src/safetensors_stage/materialize.rs new file mode 100644 index 0000000000..4b56376914 --- /dev/null +++ b/crates/model-hf/src/safetensors_stage/materialize.rs @@ -0,0 +1,694 @@ +use std::{ + collections::BTreeMap, + env, + fs::{self, File}, + io::{BufReader, BufWriter, Read, Write}, + path::{Path, PathBuf}, + sync::atomic::{AtomicU64, Ordering}, +}; + +use anyhow::{Context, Result, anyhow, ensure}; +use model_artifact::safetensors::{TensorHeader, parse_header}; +use sha2::{Digest, Sha256}; + +use super::{ + http::RemoteSource, + layout, + types::{ + MANIFEST_SCHEMA_VERSION, PreparedStage, SafetensorsSourceShard, SafetensorsStageArtifact, + SafetensorsStageManifest, SafetensorsStagePlan, SafetensorsStageRequest, SelectedTensor, + }, +}; + +const MODEL_FILE: &str = "model.safetensors"; +const CONFIG_FILE: &str = "config.json"; +const PLAN_FILE: &str = "stage-plan.json"; +const MANIFEST_FILE: &str = "stage-manifest.json"; +const MAX_LOCAL_HEADER_BYTES: u64 = 256 * 1024 * 1024; +static TEMP_SEQUENCE: AtomicU64 = AtomicU64::new(0); + +pub struct SafetensorsStageMaterializer { + remote: RemoteSource, + cache_root: PathBuf, +} + +impl SafetensorsStageMaterializer { + pub fn from_environment() -> Result { + let endpoint = env::var("HF_ENDPOINT").ok(); + Self::new( + crate::store::mesh_llm_cache_dir().join("mlx-safetensors-stages"), + endpoint.as_deref(), + hf_token(), + ) + } + + pub fn new(cache_root: PathBuf, endpoint: Option<&str>, token: Option) -> Result { + Ok(Self { + remote: RemoteSource::new(endpoint, token)?, + cache_root, + }) + } + + pub fn plan(&self, request: SafetensorsStageRequest) -> Result { + let request = request.normalized()?; + Ok(layout::prepare(&self.remote, &request)?.plan) + } + + pub fn materialize( + &self, + request: SafetensorsStageRequest, + ) -> Result { + let request = request.normalized()?; + let cache_key = cache_key(self.remote.endpoint(), &request)?; + let destination = self.cache_root.join(&cache_key); + if let Ok(Some(artifact)) = self.load_cached(&destination, &cache_key, &request) { + return Ok(artifact); + } + + let prepared = layout::prepare(&self.remote, &request)?; + fs::create_dir_all(&self.cache_root).with_context(|| { + format!( + "create SafeTensors stage cache {}", + self.cache_root.display() + ) + })?; + let temporary = self.temporary_path(&cache_key); + fs::create_dir(&temporary) + .with_context(|| format!("create temporary stage cache {}", temporary.display()))?; + if let Err(error) = self.write_stage(&temporary, &cache_key, &request, &prepared) { + let _ = fs::remove_dir_all(&temporary); + return Err(error); + } + let cache_hit = self.publish_cache(&temporary, &destination, &cache_key, &request)?; + let mut artifact = self + .load_cached(&destination, &cache_key, &request)? + .ok_or_else(|| anyhow!("published SafeTensors stage cache is missing"))?; + artifact.cache_hit = cache_hit; + Ok(artifact) + } + + fn write_stage( + &self, + directory: &Path, + cache_key: &str, + request: &SafetensorsStageRequest, + prepared: &PreparedStage, + ) -> Result<()> { + let output_path = directory.join(MODEL_FILE); + let (output_file_bytes, output_sha256) = self + .write_model(&output_path, prepared) + .context("materialize SafeTensors stage")?; + write_synced(directory.join(CONFIG_FILE), &prepared.config)?; + write_json(directory.join(PLAN_FILE), &prepared.plan)?; + let manifest = SafetensorsStageManifest { + schema_version: MANIFEST_SCHEMA_VERSION, + cache_key: cache_key.to_string(), + source_endpoint: self.remote.endpoint().to_string(), + request: request.clone(), + selected_tensor_count: prepared.plan.selected_tensor_count, + selected_tensor_bytes: prepared.plan.selected_tensor_bytes, + output_file_bytes, + output_sha256, + config_sha256: prepared.config_sha256.clone(), + config_etag: prepared.config_etag.clone(), + index_sha256: prepared.index_sha256.clone(), + index_etag: prepared.index_etag.clone(), + source_shards: prepared.source_shards.clone(), + }; + write_json(directory.join(MANIFEST_FILE), &manifest)?; + sync_directory(directory)?; + Ok(()) + } + + fn write_model(&self, path: &Path, prepared: &PreparedStage) -> Result<(u64, String)> { + let mut tensors = prepared.tensors.iter().collect::>(); + tensors.sort_by(|left, right| { + (&left.source_file, left.source_range.start) + .cmp(&(&right.source_file, right.source_range.start)) + }); + let (header, payload_bytes) = output_header(&tensors)?; + let mut header_bytes = serde_json::to_vec(&header)?; + while header_bytes.len() % 8 != 0 { + header_bytes.push(b' '); + } + let header_len = u64::try_from(header_bytes.len()).context("output header is too large")?; + + let file = File::create(path).with_context(|| format!("create {}", path.display()))?; + let mut writer = HashingWriter::new(BufWriter::new(file)); + writer.write_all(&header_len.to_le_bytes())?; + writer.write_all(&header_bytes)?; + let source_shards = prepared + .source_shards + .iter() + .map(|shard| (shard.file.as_str(), shard)) + .collect::>(); + let mut downloaded = 0_u64; + for span in materialization_spans(&tensors) { + let identity = source_shards + .get(span.source_file.as_str()) + .with_context(|| format!("missing source identity for {}", span.source_file))?; + let url = self.remote.url( + &prepared.plan.repo, + &prepared.plan.revision, + &span.source_file, + )?; + let response = self + .remote + .exact_range(url, span.start..span.end_exclusive)?; + ensure!( + response.total_file_bytes == identity.file_bytes, + "SafeTensors shard {} changed size during materialization", + span.source_file + ); + ensure_source_identity(identity, response.etag())?; + downloaded = downloaded + .checked_add(response.copy_to(&mut writer)?) + .context("materialized tensor byte count overflow")?; + } + ensure!( + downloaded == payload_bytes && downloaded == prepared.plan.selected_tensor_bytes, + "materialized {downloaded} tensor bytes but planned {}", + prepared.plan.selected_tensor_bytes + ); + writer.flush()?; + writer.inner.get_ref().sync_all()?; + let output_file_bytes = writer.bytes_written; + let output_sha256 = writer.finish_hash(); + Ok((output_file_bytes, output_sha256)) + } + + fn load_cached( + &self, + directory: &Path, + cache_key: &str, + request: &SafetensorsStageRequest, + ) -> Result> { + if !directory.exists() { + return Ok(None); + } + ensure!( + directory.is_dir(), + "SafeTensors stage cache entry is not a directory" + ); + let manifest: SafetensorsStageManifest = read_json(&directory.join(MANIFEST_FILE))?; + let plan: SafetensorsStagePlan = read_json(&directory.join(PLAN_FILE))?; + ensure!( + manifest.schema_version == MANIFEST_SCHEMA_VERSION, + "SafeTensors stage cache schema is stale" + ); + ensure!( + manifest.cache_key == cache_key, + "SafeTensors cache key mismatch" + ); + ensure!( + manifest.request == *request, + "SafeTensors cache request mismatch" + ); + ensure!( + manifest.source_endpoint == self.remote.endpoint(), + "SafeTensors cache source endpoint mismatch" + ); + ensure!( + plan.selected_tensor_count == manifest.selected_tensor_count + && plan.selected_tensor_bytes == manifest.selected_tensor_bytes, + "SafeTensors cached plan and manifest disagree" + ); + let model_path = directory.join(MODEL_FILE); + let config_path = directory.join(CONFIG_FILE); + ensure_regular_file(&model_path)?; + ensure_regular_file(&config_path)?; + ensure!( + fs::metadata(&model_path)?.len() == manifest.output_file_bytes, + "cached SafeTensors output length mismatch" + ); + ensure!( + sha256_file(&model_path)? == manifest.output_sha256, + "cached SafeTensors output hash mismatch" + ); + ensure!( + sha256_file(&config_path)? == manifest.config_sha256, + "cached SafeTensors config hash mismatch" + ); + validate_local_safetensors(&model_path)?; + Ok(Some(SafetensorsStageArtifact { + path: directory.to_path_buf(), + manifest, + plan, + cache_hit: true, + })) + } + + fn publish_cache( + &self, + temporary: &Path, + destination: &Path, + cache_key: &str, + request: &SafetensorsStageRequest, + ) -> Result { + if destination.exists() { + if matches!( + self.load_cached(destination, cache_key, request), + Ok(Some(_)) + ) { + fs::remove_dir_all(temporary)?; + return Ok(true); + } + let quarantine = self.temporary_path(&format!("{cache_key}.corrupt")); + fs::rename(destination, &quarantine).with_context(|| { + format!("quarantine corrupt stage cache {}", destination.display()) + })?; + if let Err(error) = fs::rename(temporary, destination) { + let _ = fs::rename(&quarantine, destination); + return Err(error).context("publish repaired SafeTensors stage cache"); + } + let _ = fs::remove_dir_all(quarantine); + } else { + fs::rename(temporary, destination).context("publish SafeTensors stage cache")?; + } + sync_directory(&self.cache_root)?; + Ok(false) + } + + fn temporary_path(&self, cache_key: &str) -> PathBuf { + let sequence = TEMP_SEQUENCE.fetch_add(1, Ordering::Relaxed); + self.cache_root.join(format!( + ".{cache_key}.{}.{}.partial", + std::process::id(), + sequence + )) + } +} + +struct MaterializationSpan { + source_file: String, + start: u64, + end_exclusive: u64, +} + +fn output_header(tensors: &[&SelectedTensor]) -> Result<(BTreeMap, u64)> { + let mut output_offset = 0_u64; + let mut output_header = BTreeMap::new(); + for tensor in tensors { + let end = output_offset + .checked_add(tensor.source_range.len()) + .context("partial SafeTensors output offset overflow")?; + let mut header = tensor.header.clone(); + header.data_offsets = [output_offset, end]; + output_header.insert(tensor.name.clone(), header); + output_offset = end; + } + Ok((output_header, output_offset)) +} + +fn materialization_spans(tensors: &[&SelectedTensor]) -> Vec { + let mut spans: Vec = Vec::new(); + for tensor in tensors { + if let Some(previous) = spans.last_mut() + && previous.source_file == tensor.source_file + && previous.end_exclusive == tensor.source_range.start + { + previous.end_exclusive = tensor.source_range.end_exclusive; + } else { + spans.push(MaterializationSpan { + source_file: tensor.source_file.clone(), + start: tensor.source_range.start, + end_exclusive: tensor.source_range.end_exclusive, + }); + } + } + spans +} + +fn ensure_source_identity( + identity: &SafetensorsSourceShard, + actual_etag: Option<&str>, +) -> Result<()> { + if let (Some(expected), Some(actual)) = (identity.etag.as_deref(), actual_etag) { + ensure!( + expected == actual, + "SafeTensors shard {} changed identity during materialization", + identity.file + ); + } + Ok(()) +} + +fn cache_key(endpoint: &str, request: &SafetensorsStageRequest) -> Result { + let identity = serde_json::to_vec(&(MANIFEST_SCHEMA_VERSION, endpoint, request))?; + Ok(format!("{:x}", Sha256::digest(identity))) +} + +fn hf_token() -> Option { + ["HF_TOKEN", "HUGGING_FACE_HUB_TOKEN"] + .iter() + .find_map(|name| env::var(name).ok()) + .map(|token| token.trim().to_string()) + .filter(|token| !token.is_empty()) +} + +fn write_json(path: PathBuf, value: &impl serde::Serialize) -> Result<()> { + let mut bytes = serde_json::to_vec_pretty(value)?; + bytes.push(b'\n'); + write_synced(path, &bytes) +} + +fn write_synced(path: PathBuf, bytes: &[u8]) -> Result<()> { + let mut file = File::create(&path).with_context(|| format!("create {}", path.display()))?; + file.write_all(bytes)?; + file.sync_all()?; + Ok(()) +} + +fn read_json(path: &Path) -> Result { + ensure_regular_file(path)?; + let file = File::open(path).with_context(|| format!("open {}", path.display()))?; + serde_json::from_reader(BufReader::new(file)) + .with_context(|| format!("parse {}", path.display())) +} + +fn ensure_regular_file(path: &Path) -> Result<()> { + ensure!( + fs::symlink_metadata(path)?.file_type().is_file(), + "cache path is not a regular file: {}", + path.display() + ); + Ok(()) +} + +fn sha256_file(path: &Path) -> Result { + let mut reader = BufReader::new(File::open(path)?); + let mut hasher = Sha256::new(); + std::io::copy(&mut reader, &mut HashWriter(&mut hasher))?; + Ok(format!("{:x}", hasher.finalize())) +} + +fn validate_local_safetensors(path: &Path) -> Result<()> { + let mut reader = BufReader::new(File::open(path)?); + let file_bytes = reader.get_ref().metadata()?.len(); + let mut length = [0_u8; 8]; + reader.read_exact(&mut length)?; + let header_len = u64::from_le_bytes(length); + ensure!( + header_len <= MAX_LOCAL_HEADER_BYTES, + "cached SafeTensors header is too large" + ); + let data_start = 8_u64 + .checked_add(header_len) + .context("cached SafeTensors header offset overflow")?; + ensure!( + data_start <= file_bytes, + "cached SafeTensors header is truncated" + ); + let mut header = vec![0_u8; usize::try_from(header_len)?]; + reader.read_exact(&mut header)?; + parse_header(&header, file_bytes - data_start)?; + Ok(()) +} + +fn sync_directory(path: &Path) -> Result<()> { + File::open(path)?.sync_all()?; + Ok(()) +} + +struct HashingWriter { + inner: W, + hasher: Sha256, + bytes_written: u64, +} + +impl HashingWriter { + fn new(inner: W) -> Self { + Self { + inner, + hasher: Sha256::new(), + bytes_written: 0, + } + } + + fn finish_hash(&self) -> String { + format!("{:x}", self.hasher.clone().finalize()) + } +} + +impl Write for HashingWriter { + fn write(&mut self, bytes: &[u8]) -> std::io::Result { + let written = self.inner.write(bytes)?; + self.hasher.update(&bytes[..written]); + self.bytes_written = self.bytes_written.saturating_add(written as u64); + Ok(written) + } + + fn flush(&mut self) -> std::io::Result<()> { + self.inner.flush() + } +} + +struct HashWriter<'a>(&'a mut Sha256); + +impl Write for HashWriter<'_> { + fn write(&mut self, bytes: &[u8]) -> std::io::Result { + self.0.update(bytes); + Ok(bytes.len()) + } + + fn flush(&mut self) -> std::io::Result<()> { + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use std::{ + io::{Read, Seek, SeekFrom}, + net::{TcpListener, TcpStream}, + sync::{Arc, Mutex}, + thread, + }; + + use model_artifact::safetensors::TensorHeader; + + use super::*; + use crate::safetensors_stage::types::ByteRange; + + fn selected(name: &str, file: &str, start: u64, end: u64) -> SelectedTensor { + SelectedTensor { + name: name.to_string(), + source_file: file.to_string(), + source_range: ByteRange { + start, + end_exclusive: end, + }, + header: TensorHeader { + dtype: "U8".to_string(), + shape: vec![end - start], + data_offsets: [start, end], + }, + } + } + + #[test] + fn materialization_spans_join_only_same_shard_contiguous_tensors() { + let tensors = [ + selected("a", "one", 10, 20), + selected("b", "one", 20, 30), + selected("c", "one", 40, 50), + selected("d", "two", 50, 60), + ]; + let refs = tensors.iter().collect::>(); + + let spans = materialization_spans(&refs); + + assert_eq!(spans.len(), 3); + assert_eq!(spans[0].start, 10); + assert_eq!(spans[0].end_exclusive, 30); + assert_eq!(spans[2].source_file, "two"); + } + + #[test] + fn cache_identity_changes_with_layer_range() { + let request = SafetensorsStageRequest { + repo: "org/model".to_string(), + revision: "a".repeat(40), + layer_start: 0, + layer_end: 10, + include_prefixes: Vec::new(), + }; + let mut other = request.clone(); + other.layer_start = 10; + other.layer_end = 20; + + assert_ne!( + cache_key("https://huggingface.co/", &request).unwrap(), + cache_key("https://huggingface.co/", &other).unwrap() + ); + } + + #[test] + fn materializes_exact_ranges_reuses_cache_and_repairs_corruption() { + let checkpoint = Arc::new(test_checkpoint()); + let requests = Arc::new(Mutex::new(Vec::new())); + let endpoint = start_checkpoint_server(Arc::clone(&checkpoint), Arc::clone(&requests)); + let cache = tempfile::tempdir().unwrap(); + let materializer = + SafetensorsStageMaterializer::new(cache.path().join("cache"), Some(&endpoint), None) + .unwrap(); + let request = test_request(); + + let first = materializer.materialize(request.clone()).unwrap(); + + assert!(!first.cache_hit); + assert_eq!(first.plan.selected_tensor_count, 2); + assert_eq!(first.plan.selected_tensor_bytes, 4); + assert!(first.manifest.output_file_bytes < checkpoint.len() as u64); + let request_count = requests.lock().unwrap().len(); + assert!(requests.lock().unwrap().iter().any(|request| { + request + .lines() + .any(|line| line.eq_ignore_ascii_case("range: bytes=0-7")) + })); + + let cached = materializer.materialize(request.clone()).unwrap(); + + assert!(cached.cache_hit); + assert_eq!(requests.lock().unwrap().len(), request_count); + + let model_path = cached.path.join(MODEL_FILE); + let mut model = fs::OpenOptions::new().write(true).open(model_path).unwrap(); + model.seek(SeekFrom::Start(16)).unwrap(); + model.write_all(b"broken").unwrap(); + model.sync_all().unwrap(); + + let repaired = materializer.materialize(request).unwrap(); + + assert!(!repaired.cache_hit); + assert!(requests.lock().unwrap().len() > request_count); + validate_local_safetensors(&repaired.path.join(MODEL_FILE)).unwrap(); + } + + fn test_request() -> SafetensorsStageRequest { + SafetensorsStageRequest { + repo: "org/model".to_string(), + revision: "0123456789012345678901234567890123456789".to_string(), + layer_start: 0, + layer_end: 1, + include_prefixes: Vec::new(), + } + } + + fn test_checkpoint() -> Vec { + let mut offset = 0_u64; + let mut header = BTreeMap::new(); + for name in [ + "model.embed_tokens.weight", + "model.layers.0.weight", + "model.layers.1.weight", + "model.norm.weight", + ] { + header.insert( + name, + TensorHeader { + dtype: "U8".to_string(), + shape: vec![2], + data_offsets: [offset, offset + 2], + }, + ); + offset += 2; + } + let mut header = serde_json::to_vec(&header).unwrap(); + while header.len() % 8 != 0 { + header.push(b' '); + } + let mut checkpoint = u64::try_from(header.len()).unwrap().to_le_bytes().to_vec(); + checkpoint.extend(header); + checkpoint.extend(0_u8..8); + checkpoint + } + + fn start_checkpoint_server( + checkpoint: Arc>, + requests: Arc>>, + ) -> String { + let listener = TcpListener::bind("127.0.0.1:0").unwrap(); + let address = listener.local_addr().unwrap(); + thread::spawn(move || { + for stream in listener.incoming() { + let Ok(mut stream) = stream else { + return; + }; + let checkpoint = Arc::clone(&checkpoint); + let requests = Arc::clone(&requests); + thread::spawn(move || { + handle_checkpoint_request(&mut stream, &checkpoint, &requests) + }); + } + }); + format!("http://{address}") + } + + fn handle_checkpoint_request( + stream: &mut TcpStream, + checkpoint: &[u8], + requests: &Mutex>, + ) { + let mut bytes = vec![0_u8; 8192]; + let Ok(read) = stream.read(&mut bytes) else { + return; + }; + let request = String::from_utf8_lossy(&bytes[..read]).into_owned(); + requests.lock().unwrap().push(request.clone()); + let path = request + .lines() + .next() + .and_then(|line| line.split_whitespace().nth(1)) + .unwrap_or("/"); + let response = if path.ends_with("/config.json") { + http_response( + "200 OK", + &[("ETag", "\"config-id\"")], + br#"{"model_type":"llama","hidden_size":2,"num_hidden_layers":2,"tie_word_embeddings":true}"#, + ) + } else if path.ends_with("/model.safetensors.index.json") { + http_response("404 Not Found", &[], b"") + } else if path.ends_with("/model.safetensors") { + checkpoint_range_response(&request, checkpoint) + } else { + http_response("404 Not Found", &[], b"") + }; + let _ = stream.write_all(&response); + } + + fn checkpoint_range_response(request: &str, checkpoint: &[u8]) -> Vec { + let range = request + .lines() + .find_map(|line| { + let (name, value) = line.split_once(':')?; + name.eq_ignore_ascii_case("range").then_some(value.trim()) + }) + .unwrap(); + let bounds = range.strip_prefix("bytes=").unwrap(); + let (start, end) = bounds.split_once('-').unwrap(); + let start = start.parse::().unwrap(); + let end = end.parse::().unwrap(); + let content_range = format!("bytes {start}-{end}/{}", checkpoint.len()); + http_response( + "206 Partial Content", + &[("ETag", "\"model-id\""), ("Content-Range", &content_range)], + &checkpoint[start..=end], + ) + } + + fn http_response(status: &str, headers: &[(&str, &str)], body: &[u8]) -> Vec { + let headers = headers + .iter() + .map(|(name, value)| format!("{name}: {value}\r\n")) + .collect::(); + format!( + "HTTP/1.1 {status}\r\nContent-Length: {}\r\n{headers}Connection: close\r\n\r\n", + body.len() + ) + .into_bytes() + .into_iter() + .chain(body.iter().copied()) + .collect() + } +} diff --git a/crates/model-hf/src/safetensors_stage/mod.rs b/crates/model-hf/src/safetensors_stage/mod.rs new file mode 100644 index 0000000000..642a79e398 --- /dev/null +++ b/crates/model-hf/src/safetensors_stage/mod.rs @@ -0,0 +1,10 @@ +mod http; +mod layout; +mod materialize; +mod types; + +pub use materialize::SafetensorsStageMaterializer; +pub use types::{ + ByteRange, SafetensorsShardPlan, SafetensorsSourceShard, SafetensorsStageArtifact, + SafetensorsStageManifest, SafetensorsStagePlan, SafetensorsStageRequest, +}; diff --git a/crates/model-hf/src/safetensors_stage/types.rs b/crates/model-hf/src/safetensors_stage/types.rs new file mode 100644 index 0000000000..00e6862132 --- /dev/null +++ b/crates/model-hf/src/safetensors_stage/types.rs @@ -0,0 +1,151 @@ +use std::path::PathBuf; + +use anyhow::{Result, ensure}; +use model_artifact::safetensors::TensorHeader; +use serde::{Deserialize, Serialize}; + +pub(crate) const MANIFEST_SCHEMA_VERSION: u32 = 1; + +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub struct SafetensorsStageRequest { + pub repo: String, + /// Immutable Hugging Face commit SHA, not a branch or tag. + pub revision: String, + pub layer_start: u32, + pub layer_end: u32, + #[serde(default)] + pub include_prefixes: Vec, +} + +impl SafetensorsStageRequest { + pub(crate) fn normalized(mut self) -> Result { + ensure!( + !self.repo.trim().is_empty(), + "Hugging Face repo is required" + ); + ensure!( + self.repo.split('/').count() == 2 && self.repo.split('/').all(|part| !part.is_empty()), + "Hugging Face repo must be owner/name" + ); + ensure!( + self.revision.len() == 40 && self.revision.bytes().all(|byte| byte.is_ascii_hexdigit()), + "SafeTensors stage revision must be an immutable 40-character commit SHA" + ); + ensure!( + self.layer_start < self.layer_end, + "SafeTensors stage layer range must be non-empty" + ); + self.include_prefixes = self + .include_prefixes + .into_iter() + .map(|prefix| prefix.trim().to_string()) + .filter(|prefix| !prefix.is_empty()) + .collect(); + self.include_prefixes.sort(); + self.include_prefixes.dedup(); + Ok(self) + } +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub struct SafetensorsStagePlan { + pub repo: String, + pub revision: String, + pub layer_start: u32, + pub layer_end: u32, + pub include_prefixes: Vec, + pub total_model_tensor_bytes: Option, + pub config_bytes: u64, + pub index_bytes: u64, + pub selected_tensor_count: usize, + pub selected_tensor_bytes: u64, + pub largest_selected_tensor_bytes: u64, + pub source_shard_count: usize, + pub source_shard_bytes: u64, + pub range_request_count: usize, + pub range_payload_bytes: u64, + pub header_probe_bytes: u64, + pub planned_download_bytes: u64, + pub source_shard_bytes_avoided: u64, + pub full_model_tensor_bytes_avoided: Option, + pub shards: Vec, +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub struct SafetensorsShardPlan { + pub file: String, + pub file_bytes: u64, + pub header_probe_bytes: u64, + pub selected_tensor_count: usize, + pub selected_tensor_bytes: u64, + pub largest_selected_tensor_bytes: u64, + pub ranges: Vec, +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub struct ByteRange { + pub start: u64, + pub end_exclusive: u64, +} + +impl ByteRange { + pub fn len(&self) -> u64 { + self.end_exclusive - self.start + } + + pub fn is_empty(&self) -> bool { + self.start >= self.end_exclusive + } +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub struct SafetensorsStageManifest { + pub schema_version: u32, + pub cache_key: String, + pub source_endpoint: String, + pub request: SafetensorsStageRequest, + pub selected_tensor_count: usize, + pub selected_tensor_bytes: u64, + pub output_file_bytes: u64, + pub output_sha256: String, + pub config_sha256: String, + pub config_etag: Option, + pub index_sha256: Option, + pub index_etag: Option, + pub source_shards: Vec, +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub struct SafetensorsSourceShard { + pub file: String, + pub file_bytes: u64, + pub etag: Option, +} + +#[derive(Clone, Debug)] +pub struct SafetensorsStageArtifact { + pub path: PathBuf, + pub manifest: SafetensorsStageManifest, + pub plan: SafetensorsStagePlan, + pub cache_hit: bool, +} + +#[derive(Clone, Debug)] +pub(crate) struct SelectedTensor { + pub name: String, + pub source_file: String, + pub source_range: ByteRange, + pub header: TensorHeader, +} + +#[derive(Clone, Debug)] +pub(crate) struct PreparedStage { + pub plan: SafetensorsStagePlan, + pub tensors: Vec, + pub config: Vec, + pub config_sha256: String, + pub config_etag: Option, + pub index_sha256: Option, + pub index_etag: Option, + pub source_shards: Vec, +} From 43d08d1da3dfc547d7c047c9b3f9297ad6dc2753 Mon Sep 17 00:00:00 2001 From: Michael Neale <14976+michaelneale@users.noreply.github.com> Date: Fri, 17 Jul 2026 18:56:53 +1000 Subject: [PATCH 12/37] fix(mlx): harden staged tensor identity and loading --- Cargo.lock | 1 + crates/model-hf/Cargo.toml | 3 + crates/model-hf/src/safetensors_stage/http.rs | 66 ++++- .../model-hf/src/safetensors_stage/layout.rs | 59 +++-- .../model-hf/src/safetensors_stage/locking.rs | 64 +++++ .../src/safetensors_stage/materialize.rs | 234 ++++++++++++++++-- crates/model-hf/src/safetensors_stage/mod.rs | 1 + .../model-hf/src/safetensors_stage/types.rs | 10 +- crates/skippy-engine-mlx/src/stage.rs | 31 ++- 9 files changed, 419 insertions(+), 50 deletions(-) create mode 100644 crates/model-hf/src/safetensors_stage/locking.rs diff --git a/Cargo.lock b/Cargo.lock index 28ecccfc95..a1608a15db 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4525,6 +4525,7 @@ dependencies = [ "chrono", "dirs", "hf-hub", + "libc", "model-artifact", "model-ref", "reqwest 0.12.28", diff --git a/crates/model-hf/Cargo.toml b/crates/model-hf/Cargo.toml index a8a559603b..bbcf0fba24 100644 --- a/crates/model-hf/Cargo.toml +++ b/crates/model-hf/Cargo.toml @@ -21,6 +21,9 @@ serde_json.workspace = true sha2 = "0.10" tokio = { version = "1", features = ["rt"] } +[target.'cfg(unix)'.dependencies] +libc = "0.2" + [dev-dependencies] serial_test = "3" tempfile = "3" diff --git a/crates/model-hf/src/safetensors_stage/http.rs b/crates/model-hf/src/safetensors_stage/http.rs index 8939d3b1f7..6f6a75c389 100644 --- a/crates/model-hf/src/safetensors_stage/http.rs +++ b/crates/model-hf/src/safetensors_stage/http.rs @@ -9,7 +9,9 @@ use anyhow::{Context, Result, anyhow, ensure}; use reqwest::{ StatusCode, Url, blocking::{Client, RequestBuilder, Response}, - header::{ACCEPT_ENCODING, AUTHORIZATION, CONTENT_LENGTH, CONTENT_RANGE, ETAG, RANGE}, + header::{ + ACCEPT_ENCODING, AUTHORIZATION, CONTENT_LENGTH, CONTENT_RANGE, ETAG, IF_RANGE, RANGE, + }, redirect::Policy, }; @@ -92,16 +94,38 @@ impl RemoteSource { } pub fn exact_range(&self, url: Url, range: Range) -> Result { + self.exact_range_with_identity(url, range, None) + } + + pub fn exact_range_if_range( + &self, + url: Url, + range: Range, + etag: &str, + ) -> Result { + validate_strong_etag(etag)?; + self.exact_range_with_identity(url, range, Some(etag)) + } + + fn exact_range_with_identity( + &self, + url: Url, + range: Range, + etag: Option<&str>, + ) -> Result { ensure!(range.start < range.end, "HTTP byte range must not be empty"); let expected_bytes = range.end - range.start; let range_header = format!("bytes={}-{}", range.start, range.end - 1); + let mut request = self + .client + .get(url) + .header(RANGE, range_header.clone()) + .header(ACCEPT_ENCODING, "identity"); + if let Some(etag) = etag { + request = request.header(IF_RANGE, etag); + } let response = self - .authorized( - self.client - .get(url) - .header(RANGE, range_header.clone()) - .header(ACCEPT_ENCODING, "identity"), - ) + .authorized(request) .send() .with_context(|| format!("request HTTP range {range_header}"))?; ensure!( @@ -132,6 +156,9 @@ impl RemoteSource { ); } let etag = header_string(&response, ETAG)?; + if let Some(etag) = &etag { + validate_strong_etag(etag)?; + } Ok(ExactRangeResponse { response, total_file_bytes: parsed.total_file_bytes, @@ -179,6 +206,14 @@ impl RemoteSource { } } +fn validate_strong_etag(etag: &str) -> Result<()> { + ensure!( + !etag.starts_with("W/") && etag.starts_with('"') && etag.ends_with('"'), + "SafeTensors range response requires a strong quoted ETag" + ); + Ok(()) +} + impl ExactRangeResponse { pub fn etag(&self) -> Option<&str> { self.etag.as_deref() @@ -357,6 +392,23 @@ mod tests { assert!(format!("{error:#}").contains("HTTP range")); } + #[test] + fn rejects_weak_etag_for_tensor_ranges() { + let endpoint = serve_once(http_response( + "206 Partial Content", + &[("Content-Range", "bytes 0-3/10"), ("ETag", "W/\"weak\"")], + b"four", + )); + let remote = RemoteSource::new(Some(&endpoint), None).unwrap(); + let url = remote + .url("org/model", &"a".repeat(40), "model.safetensors") + .unwrap(); + + let error = remote.exact_range(url, 0..4).err().unwrap(); + + assert!(error.to_string().contains("strong quoted ETag")); + } + fn serve_once(response: Vec) -> String { let listener = TcpListener::bind("127.0.0.1:0").unwrap(); let address = listener.local_addr().unwrap(); diff --git a/crates/model-hf/src/safetensors_stage/layout.rs b/crates/model-hf/src/safetensors_stage/layout.rs index a7af95ee98..0783e917c6 100644 --- a/crates/model-hf/src/safetensors_stage/layout.rs +++ b/crates/model-hf/src/safetensors_stage/layout.rs @@ -21,13 +21,15 @@ const MAX_HEADER_BYTES: u64 = 256 * 1024 * 1024; #[derive(Clone, Debug)] struct RemoteHeader { header_len: u64, + header_sha256: String, file_bytes: u64, - etag: Option, + etag: String, tensors: BTreeMap, } struct CheckpointLayout { index: SafetensorsIndex, + layout_sha256: String, index_bytes: u64, index_sha256: Option, index_etag: Option, @@ -49,6 +51,7 @@ pub(crate) fn prepare( add_required_prefixes(&mut selection_request, &model_config); let mut layout = load_checkpoint_layout(remote, request)?; + let checkpoint_sha256 = checkpoint_sha256(request, &config_sha256, &layout.layout_sha256)?; validate_layer_coverage(&layout.index, request)?; let selected = select_tensors(&layout.index.weight_map, &selection_request); ensure!( @@ -94,18 +97,20 @@ pub(crate) fn prepare( SafetensorsSourceShard { file: shard.file.clone(), file_bytes: header.file_bytes, - etag: header.etag.clone(), + etag: Some(header.etag.clone()), } }) .collect(); let plan = summarize_plan( &selection_request, + &checkpoint_sha256, &layout.index, config.bytes.len() as u64, layout.index_bytes, shards, )?; Ok(PreparedStage { + checkpoint_sha256, plan, tensors, config: config.bytes, @@ -193,10 +198,12 @@ fn load_checkpoint_layout( )?; if let Some(index_file) = remote.optional_small_file(index_url, MAX_INDEX_BYTES)? { let index = parse_index(&index_file.bytes)?; + let index_sha256 = sha256_hex(&index_file.bytes); return Ok(CheckpointLayout { index, + layout_sha256: index_sha256.clone(), index_bytes: index_file.bytes.len() as u64, - index_sha256: Some(sha256_hex(&index_file.bytes)), + index_sha256: Some(index_sha256), index_etag: index_file.etag, headers: BTreeMap::new(), }); @@ -220,6 +227,7 @@ fn load_checkpoint_layout( metadata: IndexMetadata { total_size }, weight_map, }, + layout_sha256: header.header_sha256.clone(), index_bytes: 0, index_sha256: None, index_etag: None, @@ -235,7 +243,10 @@ fn fetch_safetensor_header( let url = remote.url(&request.repo, &request.revision, file)?; let len_response = remote.exact_range(url.clone(), 0..8)?; let file_bytes = len_response.total_file_bytes; - let len_etag = len_response.etag().map(str::to_string); + let len_etag = len_response + .etag() + .context("SafeTensors header-length response omitted ETag")? + .to_string(); let len_bytes = len_response.into_bytes()?; let header_len = u64::from_le_bytes( len_bytes @@ -254,31 +265,34 @@ fn fetch_safetensor_header( header_end <= file_bytes, "SafeTensors header exceeds source file length" ); - let header_response = remote.exact_range(url, 8..header_end)?; + let header_response = remote.exact_range_if_range(url, 8..header_end, &len_etag)?; ensure!( header_response.total_file_bytes == file_bytes, "source file size changed while reading SafeTensors header" ); - ensure_matching_etag(len_etag.as_deref(), header_response.etag(), file)?; - let etag = header_response.etag().map(str::to_string).or(len_etag); + let header_etag = header_response + .etag() + .context("SafeTensors header response omitted ETag")?; + ensure_matching_etag(&len_etag, header_etag, file)?; + let etag = header_etag.to_string(); let header_bytes = header_response.into_bytes()?; + let header_sha256 = sha256_hex(&header_bytes); let data_bytes = file_bytes - header_end; let tensors = parse_header(&header_bytes, data_bytes)?; Ok(RemoteHeader { header_len, + header_sha256, file_bytes, etag, tensors, }) } -fn ensure_matching_etag(left: Option<&str>, right: Option<&str>, file: &str) -> Result<()> { - if let (Some(left), Some(right)) = (left, right) { - ensure!( - left == right, - "source identity changed while reading SafeTensors shard {file}" - ); - } +fn ensure_matching_etag(left: &str, right: &str, file: &str) -> Result<()> { + ensure!( + left == right, + "source identity changed while reading SafeTensors shard {file}" + ); Ok(()) } @@ -401,6 +415,7 @@ fn coalesce_contiguous_ranges(mut ranges: Vec) -> Vec { fn summarize_plan( request: &SafetensorsStageRequest, + checkpoint_sha256: &str, index: &SafetensorsIndex, config_bytes: u64, index_bytes: u64, @@ -429,6 +444,7 @@ fn summarize_plan( .and_then(|bytes| bytes.checked_add(range_payload_bytes)) .context("planned download byte count overflow")?; Ok(SafetensorsStagePlan { + checkpoint_sha256: checkpoint_sha256.to_string(), repo: request.repo.clone(), revision: request.revision.clone(), layer_start: request.layer_start, @@ -456,6 +472,21 @@ fn summarize_plan( }) } +fn checkpoint_sha256( + request: &SafetensorsStageRequest, + config_sha256: &str, + layout_sha256: &str, +) -> Result { + let identity = serde_json::to_vec(&( + "mesh-mlx-safetensors-checkpoint-v1", + &request.repo, + &request.revision, + config_sha256, + layout_sha256, + ))?; + Ok(sha256_hex(&identity)) +} + fn sha256_hex(bytes: &[u8]) -> String { format!("{:x}", Sha256::digest(bytes)) } diff --git a/crates/model-hf/src/safetensors_stage/locking.rs b/crates/model-hf/src/safetensors_stage/locking.rs new file mode 100644 index 0000000000..761b1e62d4 --- /dev/null +++ b/crates/model-hf/src/safetensors_stage/locking.rs @@ -0,0 +1,64 @@ +use std::{ + fs::{self, File, OpenOptions}, + path::Path, +}; + +use anyhow::{Context, Result}; + +pub(super) struct CacheKeyLock { + #[allow(dead_code)] + file: File, +} + +impl CacheKeyLock { + pub(super) fn acquire(cache_root: &Path, cache_key: &str) -> Result { + fs::create_dir_all(cache_root) + .with_context(|| format!("create SafeTensors stage cache {}", cache_root.display()))?; + let path = cache_root.join(format!(".{cache_key}.lock")); + let file = OpenOptions::new() + .create(true) + .truncate(false) + .read(true) + .write(true) + .open(&path) + .with_context(|| format!("open SafeTensors stage cache lock {}", path.display()))?; + lock_file(&file) + .with_context(|| format!("lock SafeTensors stage cache key {cache_key}"))?; + Ok(Self { file }) + } +} + +impl Drop for CacheKeyLock { + fn drop(&mut self) { + unlock_file(&self.file); + } +} + +#[cfg(unix)] +fn lock_file(file: &File) -> Result<()> { + use std::os::fd::AsRawFd; + + // SAFETY: `file` owns a valid descriptor for the duration of this call. + let result = unsafe { libc::flock(file.as_raw_fd(), libc::LOCK_EX) }; + if result == 0 { + Ok(()) + } else { + Err(std::io::Error::last_os_error()).context("flock failed") + } +} + +#[cfg(not(unix))] +fn lock_file(_file: &File) -> Result<()> { + Ok(()) +} + +#[cfg(unix)] +fn unlock_file(file: &File) { + use std::os::fd::AsRawFd; + + // SAFETY: `file` still owns the descriptor locked by `lock_file`. + let _ = unsafe { libc::flock(file.as_raw_fd(), libc::LOCK_UN) }; +} + +#[cfg(not(unix))] +fn unlock_file(_file: &File) {} diff --git a/crates/model-hf/src/safetensors_stage/materialize.rs b/crates/model-hf/src/safetensors_stage/materialize.rs index 4b56376914..c8802c9d86 100644 --- a/crates/model-hf/src/safetensors_stage/materialize.rs +++ b/crates/model-hf/src/safetensors_stage/materialize.rs @@ -14,6 +14,7 @@ use sha2::{Digest, Sha256}; use super::{ http::RemoteSource, layout, + locking::CacheKeyLock, types::{ MANIFEST_SCHEMA_VERSION, PreparedStage, SafetensorsSourceShard, SafetensorsStageArtifact, SafetensorsStageManifest, SafetensorsStagePlan, SafetensorsStageRequest, SelectedTensor, @@ -65,13 +66,12 @@ impl SafetensorsStageMaterializer { return Ok(artifact); } + let _cache_lock = CacheKeyLock::acquire(&self.cache_root, &cache_key)?; + if let Ok(Some(artifact)) = self.load_cached(&destination, &cache_key, &request) { + return Ok(artifact); + } + remove_stale_partials(&self.cache_root, &cache_key)?; let prepared = layout::prepare(&self.remote, &request)?; - fs::create_dir_all(&self.cache_root).with_context(|| { - format!( - "create SafeTensors stage cache {}", - self.cache_root.display() - ) - })?; let temporary = self.temporary_path(&cache_key); fs::create_dir(&temporary) .with_context(|| format!("create temporary stage cache {}", temporary.display()))?; @@ -103,6 +103,7 @@ impl SafetensorsStageMaterializer { let manifest = SafetensorsStageManifest { schema_version: MANIFEST_SCHEMA_VERSION, cache_key: cache_key.to_string(), + checkpoint_sha256: prepared.checkpoint_sha256.clone(), source_endpoint: self.remote.endpoint().to_string(), request: request.clone(), selected_tensor_count: prepared.plan.selected_tensor_count, @@ -152,9 +153,15 @@ impl SafetensorsStageMaterializer { &prepared.plan.revision, &span.source_file, )?; - let response = self - .remote - .exact_range(url, span.start..span.end_exclusive)?; + let expected_etag = identity + .etag + .as_deref() + .context("planned SafeTensors shard has no ETag")?; + let response = self.remote.exact_range_if_range( + url, + span.start..span.end_exclusive, + expected_etag, + )?; ensure!( response.total_file_bytes == identity.file_bytes, "SafeTensors shard {} changed size during materialization", @@ -173,6 +180,14 @@ impl SafetensorsStageMaterializer { writer.flush()?; writer.inner.get_ref().sync_all()?; let output_file_bytes = writer.bytes_written; + let expected_output_file_bytes = 8_u64 + .checked_add(header_len) + .and_then(|bytes| bytes.checked_add(payload_bytes)) + .context("materialized SafeTensors output length overflow")?; + ensure!( + output_file_bytes == expected_output_file_bytes, + "materialized SafeTensors output length mismatch" + ); let output_sha256 = writer.finish_hash(); Ok((output_file_bytes, output_sha256)) } @@ -213,6 +228,10 @@ impl SafetensorsStageMaterializer { && plan.selected_tensor_bytes == manifest.selected_tensor_bytes, "SafeTensors cached plan and manifest disagree" ); + ensure!( + plan.checkpoint_sha256 == manifest.checkpoint_sha256, + "SafeTensors checkpoint identity mismatch" + ); let model_path = directory.join(MODEL_FILE); let config_path = directory.join(CONFIG_FILE); ensure_regular_file(&model_path)?; @@ -323,16 +342,40 @@ fn ensure_source_identity( identity: &SafetensorsSourceShard, actual_etag: Option<&str>, ) -> Result<()> { - if let (Some(expected), Some(actual)) = (identity.etag.as_deref(), actual_etag) { - ensure!( - expected == actual, - "SafeTensors shard {} changed identity during materialization", - identity.file - ); + let expected = identity + .etag + .as_deref() + .context("planned SafeTensors shard has no ETag")?; + let actual = actual_etag.context("SafeTensors payload response omitted ETag")?; + ensure!( + expected == actual, + "SafeTensors shard {} changed identity during materialization", + identity.file + ); + Ok(()) +} + +#[cfg(unix)] +fn remove_stale_partials(cache_root: &Path, cache_key: &str) -> Result<()> { + let prefix = format!(".{cache_key}."); + for entry in fs::read_dir(cache_root)? { + let entry = entry?; + let name = entry.file_name(); + let Some(name) = name.to_str() else { + continue; + }; + if name.starts_with(&prefix) && name.ends_with(".partial") && entry.file_type()?.is_dir() { + fs::remove_dir_all(entry.path())?; + } } Ok(()) } +#[cfg(not(unix))] +fn remove_stale_partials(_cache_root: &Path, _cache_key: &str) -> Result<()> { + Ok(()) +} + fn cache_key(endpoint: &str, request: &SafetensorsStageRequest) -> Result { let identity = serde_json::to_vec(&(MANIFEST_SCHEMA_VERSION, endpoint, request))?; Ok(format!("{:x}", Sha256::digest(identity))) @@ -461,7 +504,10 @@ mod tests { use std::{ io::{Read, Seek, SeekFrom}, net::{TcpListener, TcpStream}, - sync::{Arc, Mutex}, + sync::{ + Arc, Barrier, Mutex, + atomic::{AtomicUsize, Ordering}, + }, thread, }; @@ -546,6 +592,11 @@ mod tests { .lines() .any(|line| line.eq_ignore_ascii_case("range: bytes=0-7")) })); + assert!(requests.lock().unwrap().iter().any(|request| { + request + .lines() + .any(|line| line.eq_ignore_ascii_case("if-range: \"model-id\"")) + })); let cached = materializer.materialize(request.clone()).unwrap(); @@ -565,6 +616,107 @@ mod tests { validate_local_safetensors(&repaired.path.join(MODEL_FILE)).unwrap(); } + #[test] + fn serializes_concurrent_materialization_of_the_same_cache_key() { + let checkpoint = Arc::new(test_checkpoint()); + let requests = Arc::new(Mutex::new(Vec::new())); + let endpoint = start_checkpoint_server(Arc::clone(&checkpoint), Arc::clone(&requests)); + let cache = tempfile::tempdir().unwrap(); + let cache_root = cache.path().join("cache"); + let barrier = Arc::new(Barrier::new(2)); + let handles = (0..2) + .map(|_| { + let endpoint = endpoint.clone(); + let cache_root = cache_root.clone(); + let barrier = Arc::clone(&barrier); + thread::spawn(move || { + let materializer = + SafetensorsStageMaterializer::new(cache_root, Some(&endpoint), None) + .unwrap(); + barrier.wait(); + materializer.materialize(test_request()).unwrap() + }) + }) + .collect::>(); + + let mut artifacts = handles + .into_iter() + .map(|handle| handle.join().unwrap()) + .collect::>(); + artifacts.sort_by_key(|artifact| artifact.cache_hit); + + assert!(!artifacts[0].cache_hit); + assert!(artifacts[1].cache_hit); + assert_eq!(artifacts[0].path, artifacts[1].path); + assert!(fs::read_dir(cache_root).unwrap().all(|entry| { + !entry + .unwrap() + .file_name() + .to_string_lossy() + .ends_with(".partial") + })); + } + + #[test] + fn rejects_checkpoint_shard_without_an_etag() { + let checkpoint = Arc::new(test_checkpoint()); + let requests = Arc::new(Mutex::new(Vec::new())); + let endpoint = + start_checkpoint_server_with_etag(checkpoint, requests, ModelEtagMode::Static(None)); + let cache = tempfile::tempdir().unwrap(); + let materializer = + SafetensorsStageMaterializer::new(cache.path().join("cache"), Some(&endpoint), None) + .unwrap(); + + let error = materializer.materialize(test_request()).unwrap_err(); + + assert!(format!("{error:#}").contains("omitted ETag")); + } + + #[test] + fn rejects_checkpoint_that_changes_etag_before_payload() { + let checkpoint = Arc::new(test_checkpoint()); + let requests = Arc::new(Mutex::new(Vec::new())); + let endpoint = start_checkpoint_server_with_etag( + checkpoint, + requests, + ModelEtagMode::ChangeAfterHeader, + ); + let cache = tempfile::tempdir().unwrap(); + let materializer = + SafetensorsStageMaterializer::new(cache.path().join("cache"), Some(&endpoint), None) + .unwrap(); + + let error = materializer.materialize(test_request()).unwrap_err(); + + assert!(format!("{error:#}").contains("changed identity")); + } + + #[test] + fn shares_checkpoint_identity_across_distinct_layer_ranges() { + let checkpoint = Arc::new(test_checkpoint()); + let requests = Arc::new(Mutex::new(Vec::new())); + let endpoint = start_checkpoint_server(checkpoint, requests); + let cache = tempfile::tempdir().unwrap(); + let materializer = + SafetensorsStageMaterializer::new(cache.path().join("cache"), Some(&endpoint), None) + .unwrap(); + let first_request = test_request(); + let mut final_request = first_request.clone(); + final_request.layer_start = 1; + final_request.layer_end = 2; + + let first = materializer.materialize(first_request).unwrap(); + let final_stage = materializer.materialize(final_request).unwrap(); + + assert_eq!( + first.manifest.checkpoint_sha256, + final_stage.manifest.checkpoint_sha256 + ); + assert_ne!(first.path, final_stage.path); + assert_eq!(final_stage.plan.selected_tensor_count, 3); + } + fn test_request() -> SafetensorsStageRequest { SafetensorsStageRequest { repo: "org/model".to_string(), @@ -607,9 +759,28 @@ mod tests { fn start_checkpoint_server( checkpoint: Arc>, requests: Arc>>, + ) -> String { + start_checkpoint_server_with_etag( + checkpoint, + requests, + ModelEtagMode::Static(Some("\"model-id\"")), + ) + } + + #[derive(Clone, Copy)] + enum ModelEtagMode { + Static(Option<&'static str>), + ChangeAfterHeader, + } + + fn start_checkpoint_server_with_etag( + checkpoint: Arc>, + requests: Arc>>, + etag_mode: ModelEtagMode, ) -> String { let listener = TcpListener::bind("127.0.0.1:0").unwrap(); let address = listener.local_addr().unwrap(); + let model_requests = Arc::new(AtomicUsize::new(0)); thread::spawn(move || { for stream in listener.incoming() { let Ok(mut stream) = stream else { @@ -617,8 +788,15 @@ mod tests { }; let checkpoint = Arc::clone(&checkpoint); let requests = Arc::clone(&requests); + let model_requests = Arc::clone(&model_requests); thread::spawn(move || { - handle_checkpoint_request(&mut stream, &checkpoint, &requests) + handle_checkpoint_request( + &mut stream, + &checkpoint, + &requests, + etag_mode, + &model_requests, + ) }); } }); @@ -629,6 +807,8 @@ mod tests { stream: &mut TcpStream, checkpoint: &[u8], requests: &Mutex>, + etag_mode: ModelEtagMode, + model_requests: &AtomicUsize, ) { let mut bytes = vec![0_u8; 8192]; let Ok(read) = stream.read(&mut bytes) else { @@ -650,14 +830,20 @@ mod tests { } else if path.ends_with("/model.safetensors.index.json") { http_response("404 Not Found", &[], b"") } else if path.ends_with("/model.safetensors") { - checkpoint_range_response(&request, checkpoint) + let request_index = model_requests.fetch_add(1, Ordering::SeqCst); + let etag = match etag_mode { + ModelEtagMode::Static(etag) => etag, + ModelEtagMode::ChangeAfterHeader if request_index < 2 => Some("\"model-id\""), + ModelEtagMode::ChangeAfterHeader => Some("\"changed-id\""), + }; + checkpoint_range_response(&request, checkpoint, etag) } else { http_response("404 Not Found", &[], b"") }; let _ = stream.write_all(&response); } - fn checkpoint_range_response(request: &str, checkpoint: &[u8]) -> Vec { + fn checkpoint_range_response(request: &str, checkpoint: &[u8], etag: Option<&str>) -> Vec { let range = request .lines() .find_map(|line| { @@ -670,11 +856,11 @@ mod tests { let start = start.parse::().unwrap(); let end = end.parse::().unwrap(); let content_range = format!("bytes {start}-{end}/{}", checkpoint.len()); - http_response( - "206 Partial Content", - &[("ETag", "\"model-id\""), ("Content-Range", &content_range)], - &checkpoint[start..=end], - ) + let mut headers = vec![("Content-Range", content_range.as_str())]; + if let Some(etag) = etag { + headers.push(("ETag", etag)); + } + http_response("206 Partial Content", &headers, &checkpoint[start..=end]) } fn http_response(status: &str, headers: &[(&str, &str)], body: &[u8]) -> Vec { diff --git a/crates/model-hf/src/safetensors_stage/mod.rs b/crates/model-hf/src/safetensors_stage/mod.rs index 642a79e398..afb4ed666b 100644 --- a/crates/model-hf/src/safetensors_stage/mod.rs +++ b/crates/model-hf/src/safetensors_stage/mod.rs @@ -1,5 +1,6 @@ mod http; mod layout; +mod locking; mod materialize; mod types; diff --git a/crates/model-hf/src/safetensors_stage/types.rs b/crates/model-hf/src/safetensors_stage/types.rs index 00e6862132..9f04801af1 100644 --- a/crates/model-hf/src/safetensors_stage/types.rs +++ b/crates/model-hf/src/safetensors_stage/types.rs @@ -4,12 +4,12 @@ use anyhow::{Result, ensure}; use model_artifact::safetensors::TensorHeader; use serde::{Deserialize, Serialize}; -pub(crate) const MANIFEST_SCHEMA_VERSION: u32 = 1; +pub(crate) const MANIFEST_SCHEMA_VERSION: u32 = 2; #[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] pub struct SafetensorsStageRequest { pub repo: String, - /// Immutable Hugging Face commit SHA, not a branch or tag. + /// Hugging Face commit SHA. The endpoint must honor commit-addressed immutability. pub revision: String, pub layer_start: u32, pub layer_end: u32, @@ -49,6 +49,8 @@ impl SafetensorsStageRequest { #[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] pub struct SafetensorsStagePlan { + /// Topology-wide identity shared by every layer range of this checkpoint. + pub checkpoint_sha256: String, pub repo: String, pub revision: String, pub layer_start: u32, @@ -90,7 +92,7 @@ pub struct ByteRange { impl ByteRange { pub fn len(&self) -> u64 { - self.end_exclusive - self.start + self.end_exclusive.saturating_sub(self.start) } pub fn is_empty(&self) -> bool { @@ -102,6 +104,7 @@ impl ByteRange { pub struct SafetensorsStageManifest { pub schema_version: u32, pub cache_key: String, + pub checkpoint_sha256: String, pub source_endpoint: String, pub request: SafetensorsStageRequest, pub selected_tensor_count: usize, @@ -140,6 +143,7 @@ pub(crate) struct SelectedTensor { #[derive(Clone, Debug)] pub(crate) struct PreparedStage { + pub checkpoint_sha256: String, pub plan: SafetensorsStagePlan, pub tensors: Vec, pub config: Vec, diff --git a/crates/skippy-engine-mlx/src/stage.rs b/crates/skippy-engine-mlx/src/stage.rs index fa78a349fb..c3a49ed3e1 100644 --- a/crates/skippy-engine-mlx/src/stage.rs +++ b/crates/skippy-engine-mlx/src/stage.rs @@ -17,7 +17,7 @@ use safemlx_lm::{ common::linear::project_logits_maybe_quantized, llama::{self, AttentionInput, TransformerBlock}, }, - weights::load_safetensors_lenient, + weights::{StrictLoadConfig, StrictLoadReport, load_safetensors_strict}, }; use skippy_engine::{ StageActivation, StageEngine, StageEngineInfo, StageExecutionKind, StageExecutionOutput, @@ -165,7 +165,16 @@ fn load_stage(config: MlxStageEngineConfig) -> Result { info.validate()?; let mut model = llama::Model::new(model_args, &stream)?; - load_safetensors_lenient(&mut model, weight_file(&config.model_dir), &weights_stream)?; + let load_config = partial_stage_load_config(&info); + let mut load_report = StrictLoadReport::default(); + load_safetensors_strict( + &mut model, + weight_file(&config.model_dir), + &weights_stream, + &load_config, + &mut load_report, + )?; + load_report.finish(&model, &load_config)?; retain_local_layers(&mut model, info.layer_start, info.layer_end)?; copy_stage_weights_to_compute_stream(&mut model, &info, &stream)?; stream.synchronize()?; @@ -186,6 +195,24 @@ fn load_stage(config: MlxStageEngineConfig) -> Result { }) } +fn partial_stage_load_config(info: &StageEngineInfo) -> StrictLoadConfig { + let mut config = StrictLoadConfig::default(); + for layer in 0..info.total_layers { + if layer < info.layer_start || layer >= info.layer_end { + config = config.allow_missing_contains(format!("model.layers.{layer}.")); + } + } + if !info.is_first() && !info.is_final() { + config = config.allow_missing_contains("model.embed_tokens."); + } + if !info.is_final() { + config = config + .allow_missing_contains("model.norm.") + .allow_missing_contains("lm_head."); + } + config +} + fn weight_file(model_dir: &Path) -> PathBuf { model_dir.join("model.safetensors") } From 3adc7a946682c69ca90ed23c3e1f94d660167f77 Mon Sep 17 00:00:00 2001 From: Michael Neale <14976+michaelneale@users.noreply.github.com> Date: Fri, 17 Jul 2026 19:28:21 +1000 Subject: [PATCH 13/37] feat(mlx): load safetensors stages through host control --- .../src/inference/skippy/stage/inventory.rs | 11 + .../src/inference/skippy/stage/mlx.rs | 651 ++++++++++++++++++ .../src/inference/skippy/stage/mod.rs | 226 +----- .../src/inference/skippy/stage/status.rs | 243 +++++++ .../src/inference/skippy/stage/tests.rs | 50 ++ .../src/runtime/local.rs | 54 +- crates/skippy-engine-mlx/STAGED_EXECUTION.md | 30 +- crates/skippy-engine-mlx/src/bin/mlx-stage.rs | 1 + crates/skippy-engine-mlx/src/stage.rs | 29 + crates/skippy-server/src/embedded.rs | 207 +++++- crates/skippy-server/src/engine_transport.rs | 132 +++- crates/skippy-server/src/lib.rs | 2 +- docs/design/MLX_STAGE_ENGINE_PLAN.md | 15 +- 13 files changed, 1404 insertions(+), 247 deletions(-) create mode 100644 crates/mesh-llm-host-runtime/src/inference/skippy/stage/mlx.rs create mode 100644 crates/mesh-llm-host-runtime/src/inference/skippy/stage/status.rs diff --git a/crates/mesh-llm-host-runtime/src/inference/skippy/stage/inventory.rs b/crates/mesh-llm-host-runtime/src/inference/skippy/stage/inventory.rs index f953dfa6ce..3b9a8a1ab6 100644 --- a/crates/mesh-llm-host-runtime/src/inference/skippy/stage/inventory.rs +++ b/crates/mesh-llm-host-runtime/src/inference/skippy/stage/inventory.rs @@ -188,6 +188,17 @@ struct PrepareSourceResult { } async fn prepare_stage_source(load: &StageLoadRequest) -> Result { + if load.backend == "mlx" { + #[cfg(all(feature = "mlx", target_os = "macos"))] + { + let artifact = super::mlx::prepare_stage(load).await?; + return Ok(PrepareSourceResult { + bytes_total: Some(artifact.manifest.output_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)) diff --git a/crates/mesh-llm-host-runtime/src/inference/skippy/stage/mlx.rs b/crates/mesh-llm-host-runtime/src/inference/skippy/stage/mlx.rs new file mode 100644 index 0000000000..1e1c5c6883 --- /dev/null +++ b/crates/mesh-llm-host-runtime/src/inference/skippy/stage/mlx.rs @@ -0,0 +1,651 @@ +use std::{net::SocketAddr, sync::Arc}; + +use anyhow::{Context, Result, bail, ensure}; +use model_hf::safetensors_stage::{ + SafetensorsStageArtifact, SafetensorsStageMaterializer, SafetensorsStageRequest, +}; +use skippy_engine_mlx::{MlxComputeDtype, MlxStageEngine, MlxStageEngineConfig}; +use skippy_protocol::LoadMode; +use skippy_protocol::binary::WireActivationDType; +use skippy_server::{ + EmbeddedServerHandle, EmbeddedState, engine_transport::EngineStageServerOptions, +}; + +use super::{ + RunningStage, StageControlState, StageLoadRequest, StageReadyResponse, StageStatusFilter, + StageWireDType, stage_load_failure_context, +}; + +const HF_MODEL_PREFIX: &str = "hf-model://"; + +pub(super) struct MlxStageLaunch { + pub(super) load: StageLoadRequest, + pub(super) server: EmbeddedServerHandle, + pub(super) artifact: SafetensorsStageArtifact, +} + +pub(super) async fn load_stage( + state: &mut StageControlState, + key: String, + load: StageLoadRequest, + bind_addr: SocketAddr, +) -> Result { + let launch = launch_stage(load, bind_addr).await?; + if let Err(error) = wait_for_engine_stage_ready(&launch.server, bind_addr).await { + let last_error = launch.server.status().last_error; + let context = stage_load_failure_context( + &launch.load, + "MLX engine stage did not become ready", + last_error.as_deref(), + ); + let _ = launch.server.shutdown().await; + return Err(error.context(context)); + } + let effective_load = launch.load; + state.stages.insert( + key, + RunningStage { + load: effective_load.clone(), + server: launch.server, + materialized: None, + mlx_artifact: Some(launch.artifact), + package: None, + _materialized_pin: None, + }, + ); + let status = state + .statuses(&StageStatusFilter { + topology_id: Some(effective_load.topology_id), + run_id: Some(effective_load.run_id), + stage_id: Some(effective_load.stage_id), + }) + .into_iter() + .next() + .context("MLX stage status missing after load")?; + Ok(StageReadyResponse { + accepted: true, + status, + error: None, + }) +} + +async fn wait_for_engine_stage_ready( + server: &EmbeddedServerHandle, + bind_addr: SocketAddr, +) -> Result<()> { + const STARTUP_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(30); + let deadline = tokio::time::Instant::now() + STARTUP_TIMEOUT; + loop { + let status = server.status(); + match status.state { + EmbeddedState::Failed => bail!( + "MLX engine stage startup failed: {}", + status.last_error.as_deref().unwrap_or("unknown error") + ), + EmbeddedState::Stopped => bail!("MLX engine stage stopped during startup"), + EmbeddedState::Ready => { + let ready = tokio::task::spawn_blocking(move || { + super::probe_binary_stage_ready( + bind_addr, + std::time::Duration::from_millis(500), + ) + }) + .await + .context("join MLX engine stage readiness probe")?; + if ready.is_ok() { + return Ok(()); + } + } + EmbeddedState::Starting | EmbeddedState::Stopping => {} + } + if tokio::time::Instant::now() >= deadline { + bail!( + "MLX engine stage did not become ready at {bind_addr} within {STARTUP_TIMEOUT:?}" + ); + } + tokio::time::sleep(std::time::Duration::from_millis(50)).await; + } +} + +pub(super) async fn prepare_stage(load: &StageLoadRequest) -> Result { + let load = load.clone(); + tokio::task::spawn_blocking(move || materialize_stage_blocking(&load)) + .await + .context("join MLX SafeTensors stage preparation")? +} + +pub(super) async fn launch_stage( + mut load: StageLoadRequest, + bind_addr: SocketAddr, +) -> Result { + let blocking_load = load.clone(); + let (artifact, engine) = tokio::task::spawn_blocking(move || { + let artifact = materialize_stage_blocking(&blocking_load)?; + let engine = Arc::new(MlxStageEngine::spawn(MlxStageEngineConfig { + model_dir: artifact.path.clone(), + model_id: blocking_load.model_id.clone(), + stage_index: blocking_load.stage_index, + layer_start: blocking_load.layer_start, + layer_end: blocking_load.layer_end, + compute_dtype: MlxComputeDtype::Bf16, + ctx_size: Some(blocking_load.ctx_size), + })?); + ensure!( + engine.stage_info().activation_width == blocking_load.activation_width.max(0) as u32, + "MLX stage activation width {} does not match requested {}", + engine.stage_info().activation_width, + blocking_load.activation_width + ); + Ok::<_, anyhow::Error>((artifact, engine)) + }) + .await + .context("join MLX stage load task")??; + + load.bind_addr = bind_addr.to_string(); + load.model_path = Some(artifact.path.to_string_lossy().into_owned()); + let server = skippy_server::start_stage_engine( + engine, + EngineStageServerOptions { + bind_addr, + downstream_addr: downstream_addr(&load)?, + wire_dtype: wire_dtype(load.wire_dtype)?, + }, + ); + Ok(MlxStageLaunch { + load, + server, + artifact, + }) +} + +fn materialize_stage_blocking(load: &StageLoadRequest) -> Result { + let request = request_from_load(load)?; + let artifact = SafetensorsStageMaterializer::from_environment()?.materialize(request)?; + ensure!( + artifact.manifest.checkpoint_sha256 == load.manifest_sha256, + "MLX checkpoint identity {} does not match stage claim {}", + artifact.manifest.checkpoint_sha256, + load.manifest_sha256 + ); + Ok(artifact) +} + +fn request_from_load(load: &StageLoadRequest) -> Result { + validate_load_settings(load)?; + ensure!(load.backend == "mlx", "MLX stage requires backend=mlx"); + ensure!( + load.load_mode == LoadMode::ArtifactSlice, + "MLX SafeTensors stages require load_mode=artifact-slice" + ); + ensure!( + load.lane_count == 1, + "MLX reduced stage transport currently supports lane_count=1" + ); + let model_ref = load + .package_ref + .strip_prefix(HF_MODEL_PREFIX) + .context("MLX stage package_ref must be hf-model://org/repo@commit")?; + let model_ref = model_ref::parse_model_ref(model_ref).context("parse MLX HF model ref")?; + ensure!( + model_ref.selector.is_none(), + "MLX stage HF model ref must not contain a selector" + ); + let revision = model_ref + .revision + .context("MLX stage HF model ref requires an immutable commit revision")?; + ensure!( + revision.len() == 40 && revision.bytes().all(|byte| byte.is_ascii_hexdigit()), + "MLX stage HF model ref revision must be a 40-character commit SHA" + ); + Ok(SafetensorsStageRequest { + repo: model_ref.repo, + revision, + layer_start: load.layer_start, + layer_end: load.layer_end, + include_prefixes: Vec::new(), + }) +} + +fn validate_load_settings(load: &StageLoadRequest) -> Result<()> { + ensure!(load.ctx_size > 0, "MLX stage ctx_size must be positive"); + ensure!( + load.wire_dtype != StageWireDType::Q8, + "MLX stages do not support Q8 activation wire dtype" + ); + ensure!( + load.n_batch.is_none() && load.n_ubatch.is_none(), + "MLX reduced stage transport does not support batch overrides" + ); + ensure!( + !load.native_mtp_enabled, + "MLX staged execution does not support native MTP" + ); + ensure!( + load.flash_attn_type == skippy_protocol::FlashAttentionType::Auto, + "MLX staged execution does not support flash-attention overrides" + ); + ensure!( + load.mmap.is_none() && !load.mlock, + "MLX staged execution does not support mmap/mlock overrides" + ); + ensure!( + load.projector_path.is_none(), + "MLX staged execution does not support multimodal projectors" + ); + ensure!( + matches!(load.n_gpu_layers, -1 | 0), + "MLX staged execution does not support partial GPU offload" + ); + if let Some(device) = load.selected_device.as_ref() { + ensure!( + device.backend_device.to_ascii_lowercase().contains("metal") + && device.index.is_none_or(|index| index == 0), + "MLX stage selected_device must be Metal device 0" + ); + } + Ok(()) +} + +fn downstream_addr(load: &StageLoadRequest) -> Result> { + load.downstream + .as_ref() + .map(|peer| { + peer.endpoint + .parse() + .with_context(|| format!("parse MLX downstream endpoint {}", peer.endpoint)) + }) + .transpose() +} + +fn wire_dtype(dtype: StageWireDType) -> Result { + match dtype { + StageWireDType::F32 => Ok(WireActivationDType::F32), + StageWireDType::F16 => Ok(WireActivationDType::F16), + StageWireDType::Q8 => bail!("MLX stages do not support Q8 activation wire dtype"), + } +} + +#[cfg(test)] +mod tests { + use std::{io::Write, net::TcpStream, time::Duration}; + + use skippy_protocol::FlashAttentionType; + use skippy_protocol::binary::{ + StageStateHeader, StageWireMessage, WireMessageKind, WireReplyKind, recv_ready, recv_reply, + write_stage_message, + }; + + use super::super::{ + StageControlCommand, StageControlRequest, StageControlResponse, StageInventoryRequest, + StagePeerDescriptor, StagePreparationState, StagePrepareRequest, StageStopRequest, + spawn_stage_control_loop, + }; + use super::*; + + const SMOL_REPO: &str = "HuggingFaceTB/SmolLM2-135M-Instruct"; + const SMOL_REVISION: &str = "12fd25f77366fa6b3b4b768ec3050bf629380bac"; + const SMOL_PROMPT: &[i32] = &[1, 1531, 314, 260, 3575, 28]; + const SMOL_EXPECTED: &[i32] = &[284, 260, 2240, 314, 1343, 327, 624, 8685]; + + #[test] + fn parses_commit_addressed_mlx_stage_ref() { + let request = request_from_load(&load_request()).unwrap(); + + assert_eq!(request.repo, "org/model"); + assert_eq!(request.revision, "a".repeat(40)); + assert_eq!((request.layer_start, request.layer_end), (4, 8)); + } + + #[test] + fn rejects_mutable_mlx_stage_ref() { + let mut load = load_request(); + load.package_ref = "hf-model://org/model@main".to_string(); + + assert!(request_from_load(&load).is_err()); + } + + #[test] + fn rejects_unsupported_load_mode_and_lanes() { + let mut load = load_request(); + load.load_mode = LoadMode::LayerPackage; + assert!(request_from_load(&load).is_err()); + + load.load_mode = LoadMode::ArtifactSlice; + load.lane_count = 2; + assert!(request_from_load(&load).is_err()); + } + + #[test] + fn rejects_unsupported_runtime_overrides_before_download() { + let mut load = load_request(); + load.wire_dtype = StageWireDType::Q8; + assert!(request_from_load(&load).is_err()); + + load.wire_dtype = StageWireDType::F16; + load.n_batch = Some(8); + assert!(request_from_load(&load).is_err()); + + load.n_batch = None; + load.native_mtp_enabled = true; + assert!(request_from_load(&load).is_err()); + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + #[ignore = "downloads about 310 MiB and requires Apple Silicon Metal"] + async fn real_control_plane_materializes_and_runs_two_range_only_stages() -> Result<()> { + let plan = tokio::task::spawn_blocking(|| { + SafetensorsStageMaterializer::from_environment()?.plan(SafetensorsStageRequest { + repo: SMOL_REPO.to_string(), + revision: SMOL_REVISION.to_string(), + layer_start: 0, + layer_end: 15, + include_prefixes: Vec::new(), + }) + }) + .await + .context("join SmolLM2 range plan")??; + ensure!( + plan.planned_download_bytes < plan.source_shard_bytes, + "SmolLM2 stage plan did not avoid the complete source shard" + ); + + let control = spawn_stage_control_loop(None); + let final_load = smol_load(1, 15, 30, None, &plan.checkpoint_sha256); + prepare_and_wait(&control, &final_load).await?; + let final_ready = load_through_control(&control, &final_load).await?; + ensure!(final_ready.accepted, "final MLX stage load was rejected"); + assert_range_only_status(&final_ready.status, &final_load)?; + + let downstream = StagePeerDescriptor { + stage_id: final_load.stage_id.clone(), + stage_index: final_load.stage_index, + endpoint: final_ready.status.bind_addr.clone(), + node_id: None, + }; + let first_load = smol_load(0, 0, 15, Some(downstream), &plan.checkpoint_sha256); + prepare_and_wait(&control, &first_load).await?; + let first_ready = load_through_control(&control, &first_load).await?; + ensure!(first_ready.accepted, "first MLX stage load was rejected"); + assert_range_only_status(&first_ready.status, &first_load)?; + + let first_addr = first_ready.status.bind_addr.parse()?; + let generated = tokio::task::spawn_blocking(move || prove_chain(first_addr)) + .await + .context("join MLX control-plane proof client")??; + ensure!(generated == SMOL_EXPECTED, "MLX stage tokens diverged"); + + stop_stage(&control, &first_load).await?; + stop_stage(&control, &final_load).await?; + Ok(()) + } + + async fn prepare_and_wait( + control: &tokio::sync::mpsc::UnboundedSender, + load: &StageLoadRequest, + ) -> Result<()> { + let response = send_control( + control, + StageControlRequest::Prepare(StagePrepareRequest { + load: load.clone(), + coordinator_id: None, + }), + ) + .await?; + let StageControlResponse::PrepareAccepted(accepted) = response else { + bail!("MLX prepare returned the wrong control response"); + }; + ensure!(accepted.accepted, "MLX stage prepare was rejected"); + tokio::time::timeout(Duration::from_secs(600), async { + loop { + let response = send_control( + control, + StageControlRequest::Inventory(StageInventoryRequest { + model_id: load.model_id.clone(), + package_ref: load.package_ref.clone(), + manifest_sha256: load.manifest_sha256.clone(), + }), + ) + .await?; + let StageControlResponse::Inventory(inventory) = response else { + bail!("MLX inventory returned the wrong control response"); + }; + if let Some(status) = inventory + .preparing_ranges + .into_iter() + .find(|status| status.stage_id == load.stage_id) + { + match status.state { + StagePreparationState::Available => return Ok(()), + StagePreparationState::Failed => { + bail!( + "MLX stage preparation failed: {}", + status.error.as_deref().unwrap_or("unknown error") + ); + } + _ => {} + } + } + tokio::time::sleep(Duration::from_millis(100)).await; + } + }) + .await + .context("timed out preparing MLX SafeTensors stage")? + } + + async fn load_through_control( + control: &tokio::sync::mpsc::UnboundedSender, + load: &StageLoadRequest, + ) -> Result { + let response = send_control(control, StageControlRequest::Load(load.clone())).await?; + let StageControlResponse::Ready(ready) = response else { + bail!("MLX load returned the wrong control response"); + }; + Ok(ready) + } + + fn assert_range_only_status( + status: &super::super::StageStatusSnapshot, + load: &StageLoadRequest, + ) -> Result<()> { + ensure!( + status.manifest_sha256.as_deref() == Some(load.manifest_sha256.as_str()), + "running MLX stage checkpoint identity changed" + ); + ensure!(status.source_model_sha256.is_none()); + ensure!(status.source_model_bytes.is_none()); + ensure!(!status.materialized_pinned); + let artifact_path = status + .materialized_path + .as_deref() + .context("running MLX stage has no materialized path")?; + let plan: model_hf::safetensors_stage::SafetensorsStagePlan = serde_json::from_slice( + &std::fs::read(std::path::Path::new(artifact_path).join("stage-plan.json"))?, + )?; + ensure!( + plan.planned_download_bytes < plan.source_shard_bytes, + "running MLX stage planned a complete source shard download" + ); + Ok(()) + } + + async fn stop_stage( + control: &tokio::sync::mpsc::UnboundedSender, + load: &StageLoadRequest, + ) -> Result<()> { + let response = send_control( + control, + StageControlRequest::Stop(StageStopRequest { + topology_id: load.topology_id.clone(), + run_id: load.run_id.clone(), + stage_id: load.stage_id.clone(), + shutdown_generation: load.shutdown_generation + 1, + coordinator_term: load.coordinator_term, + }), + ) + .await?; + let StageControlResponse::Ready(response) = response else { + bail!("MLX stop returned the wrong control response"); + }; + ensure!(response.accepted, "MLX stage stop was rejected"); + Ok(()) + } + + async fn send_control( + control: &tokio::sync::mpsc::UnboundedSender, + request: StageControlRequest, + ) -> Result { + let (resp, rx) = tokio::sync::oneshot::channel(); + control + .send(StageControlCommand { request, resp }) + .map_err(|_| anyhow::anyhow!("MLX stage control loop stopped"))?; + rx.await.context("MLX stage control response dropped")? + } + + fn smol_load( + stage_index: u32, + layer_start: u32, + layer_end: u32, + downstream: Option, + checkpoint_sha256: &str, + ) -> StageLoadRequest { + let mut load = load_request(); + load.topology_id = "mlx-control-proof".to_string(); + load.run_id = "smollm2-range-only".to_string(); + load.model_id = SMOL_REPO.to_string(); + load.package_ref = format!("hf-model://{SMOL_REPO}@{SMOL_REVISION}"); + load.manifest_sha256 = checkpoint_sha256.to_string(); + load.stage_id = format!("stage-{stage_index}"); + load.stage_index = stage_index; + load.layer_start = layer_start; + load.layer_end = layer_end; + load.activation_width = 576; + load.downstream = downstream; + load + } + + fn prove_chain(connect: SocketAddr) -> Result> { + let wire_dtype = WireActivationDType::F16; + let mut stream = TcpStream::connect(connect) + .with_context(|| format!("connect first MLX stage at {connect}"))?; + stream.set_nodelay(true).ok(); + recv_ready(&mut stream).context("first MLX stage did not become ready")?; + + let session_id = 1; + let request_id = 1; + let prefill = execution_message( + WireMessageKind::PrefillFinalEmbd, + SMOL_PROMPT, + *SMOL_PROMPT.last().expect("non-empty prompt"), + request_id, + session_id, + 0, + ); + let mut generated = Vec::with_capacity(SMOL_EXPECTED.len()); + generated.push(send_predicted(&mut stream, &prefill)?); + while generated.len() < SMOL_EXPECTED.len() { + let current = *generated.last().expect("generated has first token"); + let decode = execution_message( + WireMessageKind::DecodeEmbd, + &[], + current, + request_id, + session_id, + i32::try_from(generated.len())?, + ); + generated.push(send_predicted(&mut stream, &decode)?); + } + + let stop = StageWireMessage::stop_with_identity(wire_dtype, request_id, session_id); + write_stage_message(&mut stream, &stop, wire_dtype)?; + stream.flush().ok(); + ensure!(recv_reply(&mut stream)?.kind == WireReplyKind::Ack); + Ok(generated) + } + + fn execution_message( + kind: WireMessageKind, + tokens: &[i32], + current_token: i32, + request_id: u64, + session_id: u64, + decode_step: i32, + ) -> StageWireMessage { + let wire_dtype = WireActivationDType::F16; + let mut state = StageStateHeader::new(kind, wire_dtype); + state.current_token = current_token; + state.prompt_token_count = i32::try_from(tokens.len()).unwrap_or_default(); + state.decode_step = decode_step; + StageWireMessage { + kind, + pos_start: 0, + token_count: if kind == WireMessageKind::DecodeEmbd { + 1 + } else { + i32::try_from(tokens.len()).unwrap_or_default() + }, + state, + request_id, + session_id, + sampling: None, + chat_sampling_metadata: None, + tokens: tokens.to_vec(), + positions: Vec::new(), + activation: Vec::new(), + raw_bytes: Vec::new(), + } + } + + fn send_predicted(stream: &mut TcpStream, message: &StageWireMessage) -> Result { + write_stage_message(&mut *stream, message, WireActivationDType::F16)?; + stream.flush().ok(); + let reply = recv_reply(&mut *stream)?; + ensure!( + matches!( + reply.kind, + WireReplyKind::PredictedToken | WireReplyKind::PredictedTokens + ), + "MLX stage chain did not return a predicted token" + ); + Ok(reply.predicted) + } + + fn load_request() -> StageLoadRequest { + StageLoadRequest { + topology_id: "topology".to_string(), + run_id: "run".to_string(), + model_id: "org/model".to_string(), + backend: "mlx".to_string(), + package_ref: format!("hf-model://org/model@{}", "a".repeat(40)), + manifest_sha256: "b".repeat(64), + stage_id: "stage-0".to_string(), + stage_index: 0, + layer_start: 4, + layer_end: 8, + model_path: None, + source_model_bytes: None, + projector_path: None, + selected_device: None, + bind_addr: "127.0.0.1:0".to_string(), + activation_width: 2, + wire_dtype: StageWireDType::F16, + ctx_size: 128, + lane_count: 1, + n_batch: None, + n_ubatch: None, + n_gpu_layers: 0, + mmap: None, + mlock: false, + cache_type_k: "f16".to_string(), + cache_type_v: "f16".to_string(), + flash_attn_type: FlashAttentionType::Auto, + native_mtp_enabled: false, + shutdown_generation: 0, + coordinator_term: 0, + coordinator_id: None, + lease_until_unix_ms: 0, + load_mode: LoadMode::ArtifactSlice, + upstream: None, + downstream: None, + } + } +} diff --git a/crates/mesh-llm-host-runtime/src/inference/skippy/stage/mod.rs b/crates/mesh-llm-host-runtime/src/inference/skippy/stage/mod.rs index 6b37674028..61fceab6bf 100644 --- a/crates/mesh-llm-host-runtime/src/inference/skippy/stage/mod.rs +++ b/crates/mesh-llm-host-runtime/src/inference/skippy/stage/mod.rs @@ -10,7 +10,7 @@ use std::{ use anyhow::{Context, Result, anyhow}; use skippy_coordinator::{ClaimDecision, ClaimFence, LoadClaimRef}; -use skippy_protocol::{FlashAttentionType, LoadMode, PeerConfig, StageConfig}; +use skippy_protocol::{LoadMode, PeerConfig, StageConfig}; use skippy_server::{ EmbeddedServerHandle, binary_transport::{BinaryStageOptions, WireCondition}, @@ -22,17 +22,25 @@ use tokio::{ }; mod inventory; +#[cfg(all(feature = "mlx", target_os = "macos"))] +mod mlx; +mod status; #[cfg(test)] mod tests; mod types; use inventory::{resolve_inventory_source, run_stage_prepare_task}; +use status::{ + failed_status_from_load, preparation_status_from_cancel, preparation_status_from_load, + status_from_running, stopped_status, +}; pub(crate) use types::*; struct RunningStage { load: StageLoadRequest, server: EmbeddedServerHandle, materialized: Option, + mlx_artifact: Option, package: Option, _materialized_pin: Option, } @@ -329,11 +337,12 @@ impl StageControlState { } async fn load(&mut self, load: StageLoadRequest) -> Result { - anyhow::ensure!( - load.backend == "skippy", - "unsupported stage backend '{}'", - load.backend - ); + match load.backend.as_str() { + "skippy" => {} + #[cfg(all(feature = "mlx", target_os = "macos"))] + "mlx" => {} + backend => anyhow::bail!("unsupported stage backend '{backend}'"), + } if let Some(error) = self.validate_load_claim(&load) { return Ok(StageReadyResponse { accepted: false, @@ -347,6 +356,10 @@ impl StageControlState { } let bind_addr = materialize_stage_bind_addr(parse_bind_addr(&load.bind_addr)?)?; + #[cfg(all(feature = "mlx", target_os = "macos"))] + if load.backend == "mlx" { + return mlx::load_stage(self, key, load, bind_addr).await; + } let mut effective_load = load; effective_load.bind_addr = bind_addr.to_string(); super::configure_materialized_stage_cache(); @@ -399,6 +412,7 @@ impl StageControlState { load: effective_load.clone(), server, materialized: None, + mlx_artifact: None, package: resolved_package, _materialized_pin: None, }, @@ -759,206 +773,6 @@ fn empty_to_default(value: &str, default: &str) -> String { } } -fn status_from_running(stage: &RunningStage) -> StageStatusSnapshot { - let server = stage.server.status(); - let state = match server.state { - skippy_server::EmbeddedState::Starting => StageRuntimeState::Starting, - skippy_server::EmbeddedState::Ready => StageRuntimeState::Ready, - skippy_server::EmbeddedState::Stopping => StageRuntimeState::Stopping, - skippy_server::EmbeddedState::Stopped => StageRuntimeState::Stopped, - skippy_server::EmbeddedState::Failed => StageRuntimeState::Failed, - }; - StageStatusSnapshot { - topology_id: stage.load.topology_id.clone(), - run_id: stage.load.run_id.clone(), - model_id: stage.load.model_id.clone(), - backend: stage.load.backend.clone(), - package_ref: Some(stage.load.package_ref.clone()), - manifest_sha256: Some(stage.load.manifest_sha256.clone()), - source_model_path: stage - .materialized - .as_ref() - .map(|artifact| artifact.source_model_path.clone()) - .or_else(|| { - stage - .package - .as_ref() - .map(|package| package.source_model_path.clone()) - }) - .or_else(|| stage.load.model_path.clone()), - source_model_sha256: stage - .materialized - .as_ref() - .map(|artifact| artifact.source_model_sha256.clone()) - .or_else(|| { - stage - .package - .as_ref() - .map(|package| package.source_model_sha256.clone()) - }), - source_model_bytes: stage - .materialized - .as_ref() - .and_then(|artifact| artifact.source_model_bytes) - .or_else(|| { - stage - .package - .as_ref() - .and_then(|package| package.source_model_bytes) - }) - .or(stage.load.source_model_bytes), - materialized_path: stage - .materialized - .as_ref() - .map(|artifact| artifact.path.to_string_lossy().to_string()), - materialized_pinned: stage.materialized.is_some(), - projector_path: stage.load.projector_path.clone(), - stage_id: stage.load.stage_id.clone(), - stage_index: stage.load.stage_index, - layer_start: stage.load.layer_start, - layer_end: stage.load.layer_end, - state, - bind_addr: server.bind_addr.to_string(), - activation_width: stage.load.activation_width.max(0) as u32, - wire_dtype: stage.load.wire_dtype, - selected_device: stage.load.selected_device.clone(), - ctx_size: stage.load.ctx_size, - lane_count: stage.load.lane_count, - n_batch: stage.load.n_batch, - n_ubatch: stage.load.n_ubatch, - flash_attn_type: stage.load.flash_attn_type, - error: server.last_error.clone(), - shutdown_generation: stage.load.shutdown_generation, - coordinator_term: stage.load.coordinator_term, - coordinator_id: stage.load.coordinator_id, - lease_until_unix_ms: stage.load.lease_until_unix_ms, - } -} - -fn stopped_status(stop: &StageStopRequest) -> StageStatusSnapshot { - StageStatusSnapshot { - topology_id: stop.topology_id.clone(), - run_id: stop.run_id.clone(), - model_id: String::new(), - backend: "skippy".to_string(), - package_ref: None, - manifest_sha256: None, - source_model_path: None, - source_model_sha256: None, - source_model_bytes: None, - materialized_path: None, - materialized_pinned: false, - projector_path: None, - stage_id: stop.stage_id.clone(), - stage_index: 0, - layer_start: 0, - layer_end: 0, - state: StageRuntimeState::Stopped, - bind_addr: String::new(), - activation_width: 0, - wire_dtype: StageWireDType::F32, - selected_device: None, - ctx_size: 0, - lane_count: 0, - n_batch: None, - n_ubatch: None, - flash_attn_type: FlashAttentionType::Auto, - error: None, - shutdown_generation: stop.shutdown_generation, - coordinator_term: stop.coordinator_term, - coordinator_id: None, - lease_until_unix_ms: 0, - } -} - -fn failed_status_from_load(load: &StageLoadRequest, error: String) -> StageStatusSnapshot { - StageStatusSnapshot { - topology_id: load.topology_id.clone(), - run_id: load.run_id.clone(), - model_id: load.model_id.clone(), - backend: load.backend.clone(), - package_ref: Some(load.package_ref.clone()), - manifest_sha256: Some(load.manifest_sha256.clone()), - source_model_path: load.model_path.clone(), - source_model_sha256: None, - source_model_bytes: load.source_model_bytes, - materialized_path: None, - materialized_pinned: false, - projector_path: load.projector_path.clone(), - stage_id: load.stage_id.clone(), - stage_index: load.stage_index, - layer_start: load.layer_start, - layer_end: load.layer_end, - state: StageRuntimeState::Failed, - bind_addr: load.bind_addr.clone(), - activation_width: load.activation_width.max(0) as u32, - wire_dtype: load.wire_dtype, - selected_device: load.selected_device.clone(), - ctx_size: load.ctx_size, - lane_count: load.lane_count, - n_batch: load.n_batch, - n_ubatch: load.n_ubatch, - flash_attn_type: load.flash_attn_type, - error: Some(error), - shutdown_generation: load.shutdown_generation, - coordinator_term: load.coordinator_term, - coordinator_id: load.coordinator_id, - lease_until_unix_ms: load.lease_until_unix_ms, - } -} - -fn preparation_status_from_load( - load: &StageLoadRequest, - state: StagePreparationState, - error: Option, -) -> StagePreparationStatus { - StagePreparationStatus { - topology_id: load.topology_id.clone(), - run_id: load.run_id.clone(), - model_id: load.model_id.clone(), - backend: load.backend.clone(), - package_ref: load.package_ref.clone(), - manifest_sha256: load.manifest_sha256.clone(), - stage_id: load.stage_id.clone(), - stage_index: load.stage_index, - layer_start: load.layer_start, - layer_end: load.layer_end, - state, - bytes_done: None, - bytes_total: None, - bind_addr: None, - error, - shutdown_generation: load.shutdown_generation, - coordinator_term: load.coordinator_term, - coordinator_id: load.coordinator_id, - lease_until_unix_ms: load.lease_until_unix_ms, - } -} - -fn preparation_status_from_cancel(cancel: StageCancelPrepareRequest) -> StagePreparationStatus { - StagePreparationStatus { - topology_id: cancel.topology_id, - run_id: cancel.run_id, - model_id: String::new(), - backend: "skippy".to_string(), - package_ref: String::new(), - manifest_sha256: String::new(), - stage_id: cancel.stage_id, - stage_index: 0, - layer_start: 0, - layer_end: 0, - state: StagePreparationState::Cancelled, - bytes_done: None, - bytes_total: None, - bind_addr: None, - error: None, - shutdown_generation: cancel.shutdown_generation, - coordinator_term: 0, - coordinator_id: None, - lease_until_unix_ms: 0, - } -} - impl From for skippy_protocol::binary::WireActivationDType { fn from(value: StageWireDType) -> Self { match value { diff --git a/crates/mesh-llm-host-runtime/src/inference/skippy/stage/status.rs b/crates/mesh-llm-host-runtime/src/inference/skippy/stage/status.rs new file mode 100644 index 0000000000..c50e1377a7 --- /dev/null +++ b/crates/mesh-llm-host-runtime/src/inference/skippy/stage/status.rs @@ -0,0 +1,243 @@ +use skippy_protocol::FlashAttentionType; + +use super::{ + RunningStage, StageCancelPrepareRequest, StageLoadRequest, StagePreparationState, + StagePreparationStatus, StageRuntimeState, StageStatusSnapshot, StageStopRequest, + StageWireDType, +}; + +pub(super) fn status_from_running(stage: &RunningStage) -> StageStatusSnapshot { + let server = stage.server.status(); + StageStatusSnapshot { + topology_id: stage.load.topology_id.clone(), + run_id: stage.load.run_id.clone(), + model_id: stage.load.model_id.clone(), + backend: stage.load.backend.clone(), + package_ref: Some(stage.load.package_ref.clone()), + manifest_sha256: Some(stage.load.manifest_sha256.clone()), + source_model_path: source_model_path(stage), + source_model_sha256: source_model_sha256(stage), + source_model_bytes: source_model_bytes(stage), + materialized_path: materialized_path(stage), + materialized_pinned: stage.materialized.is_some(), + projector_path: stage.load.projector_path.clone(), + stage_id: stage.load.stage_id.clone(), + stage_index: stage.load.stage_index, + layer_start: stage.load.layer_start, + layer_end: stage.load.layer_end, + state: runtime_state(server.state), + bind_addr: server.bind_addr.to_string(), + activation_width: stage.load.activation_width.max(0) as u32, + wire_dtype: stage.load.wire_dtype, + selected_device: stage.load.selected_device.clone(), + ctx_size: stage.load.ctx_size, + lane_count: stage.load.lane_count, + n_batch: stage.load.n_batch, + n_ubatch: stage.load.n_ubatch, + flash_attn_type: stage.load.flash_attn_type, + error: server.last_error, + shutdown_generation: stage.load.shutdown_generation, + coordinator_term: stage.load.coordinator_term, + coordinator_id: stage.load.coordinator_id, + lease_until_unix_ms: stage.load.lease_until_unix_ms, + } +} + +fn runtime_state(state: skippy_server::EmbeddedState) -> StageRuntimeState { + match state { + skippy_server::EmbeddedState::Starting => StageRuntimeState::Starting, + skippy_server::EmbeddedState::Ready => StageRuntimeState::Ready, + skippy_server::EmbeddedState::Stopping => StageRuntimeState::Stopping, + skippy_server::EmbeddedState::Stopped => StageRuntimeState::Stopped, + skippy_server::EmbeddedState::Failed => StageRuntimeState::Failed, + } +} + +fn source_model_path(stage: &RunningStage) -> Option { + stage + .materialized + .as_ref() + .map(|artifact| artifact.source_model_path.clone()) + .or_else(|| { + stage + .package + .as_ref() + .map(|package| package.source_model_path.clone()) + }) + .or_else(|| { + stage + .mlx_artifact + .as_ref() + .map(|_| stage.load.package_ref.clone()) + }) + .or_else(|| stage.load.model_path.clone()) +} + +fn source_model_sha256(stage: &RunningStage) -> Option { + stage + .materialized + .as_ref() + .map(|artifact| artifact.source_model_sha256.clone()) + .or_else(|| { + stage + .package + .as_ref() + .map(|package| package.source_model_sha256.clone()) + }) +} + +fn source_model_bytes(stage: &RunningStage) -> Option { + stage + .materialized + .as_ref() + .and_then(|artifact| artifact.source_model_bytes) + .or_else(|| { + stage + .package + .as_ref() + .and_then(|package| package.source_model_bytes) + }) + .or(stage.load.source_model_bytes) +} + +fn materialized_path(stage: &RunningStage) -> Option { + stage + .materialized + .as_ref() + .map(|artifact| artifact.path.to_string_lossy().into_owned()) + .or_else(|| { + stage + .mlx_artifact + .as_ref() + .map(|artifact| artifact.path.to_string_lossy().into_owned()) + }) +} + +pub(super) fn stopped_status(stop: &StageStopRequest) -> StageStatusSnapshot { + StageStatusSnapshot { + topology_id: stop.topology_id.clone(), + run_id: stop.run_id.clone(), + model_id: String::new(), + backend: "skippy".to_string(), + package_ref: None, + manifest_sha256: None, + source_model_path: None, + source_model_sha256: None, + source_model_bytes: None, + materialized_path: None, + materialized_pinned: false, + projector_path: None, + stage_id: stop.stage_id.clone(), + stage_index: 0, + layer_start: 0, + layer_end: 0, + state: StageRuntimeState::Stopped, + bind_addr: String::new(), + activation_width: 0, + wire_dtype: StageWireDType::F32, + selected_device: None, + ctx_size: 0, + lane_count: 0, + n_batch: None, + n_ubatch: None, + flash_attn_type: FlashAttentionType::Auto, + error: None, + shutdown_generation: stop.shutdown_generation, + coordinator_term: stop.coordinator_term, + coordinator_id: None, + lease_until_unix_ms: 0, + } +} + +pub(super) fn failed_status_from_load( + load: &StageLoadRequest, + error: String, +) -> StageStatusSnapshot { + StageStatusSnapshot { + topology_id: load.topology_id.clone(), + run_id: load.run_id.clone(), + model_id: load.model_id.clone(), + backend: load.backend.clone(), + package_ref: Some(load.package_ref.clone()), + manifest_sha256: Some(load.manifest_sha256.clone()), + source_model_path: load.model_path.clone(), + source_model_sha256: None, + source_model_bytes: load.source_model_bytes, + materialized_path: None, + materialized_pinned: false, + projector_path: load.projector_path.clone(), + stage_id: load.stage_id.clone(), + stage_index: load.stage_index, + layer_start: load.layer_start, + layer_end: load.layer_end, + state: StageRuntimeState::Failed, + bind_addr: load.bind_addr.clone(), + activation_width: load.activation_width.max(0) as u32, + wire_dtype: load.wire_dtype, + selected_device: load.selected_device.clone(), + ctx_size: load.ctx_size, + lane_count: load.lane_count, + n_batch: load.n_batch, + n_ubatch: load.n_ubatch, + flash_attn_type: load.flash_attn_type, + error: Some(error), + shutdown_generation: load.shutdown_generation, + coordinator_term: load.coordinator_term, + coordinator_id: load.coordinator_id, + lease_until_unix_ms: load.lease_until_unix_ms, + } +} + +pub(super) fn preparation_status_from_load( + load: &StageLoadRequest, + state: StagePreparationState, + error: Option, +) -> StagePreparationStatus { + StagePreparationStatus { + topology_id: load.topology_id.clone(), + run_id: load.run_id.clone(), + model_id: load.model_id.clone(), + backend: load.backend.clone(), + package_ref: load.package_ref.clone(), + manifest_sha256: load.manifest_sha256.clone(), + stage_id: load.stage_id.clone(), + stage_index: load.stage_index, + layer_start: load.layer_start, + layer_end: load.layer_end, + state, + bytes_done: None, + bytes_total: None, + bind_addr: None, + error, + shutdown_generation: load.shutdown_generation, + coordinator_term: load.coordinator_term, + coordinator_id: load.coordinator_id, + lease_until_unix_ms: load.lease_until_unix_ms, + } +} + +pub(super) fn preparation_status_from_cancel( + cancel: StageCancelPrepareRequest, +) -> StagePreparationStatus { + StagePreparationStatus { + topology_id: cancel.topology_id, + run_id: cancel.run_id, + model_id: String::new(), + backend: "skippy".to_string(), + package_ref: String::new(), + manifest_sha256: String::new(), + stage_id: cancel.stage_id, + stage_index: 0, + layer_start: 0, + layer_end: 0, + state: StagePreparationState::Cancelled, + bytes_done: None, + bytes_total: None, + bind_addr: None, + error: None, + shutdown_generation: cancel.shutdown_generation, + coordinator_term: 0, + coordinator_id: None, + lease_until_unix_ms: 0, + } +} diff --git a/crates/mesh-llm-host-runtime/src/inference/skippy/stage/tests.rs b/crates/mesh-llm-host-runtime/src/inference/skippy/stage/tests.rs index 4a3e7c36f4..0b0cddd9ba 100644 --- a/crates/mesh-llm-host-runtime/src/inference/skippy/stage/tests.rs +++ b/crates/mesh-llm-host-runtime/src/inference/skippy/stage/tests.rs @@ -390,6 +390,56 @@ async fn prepare_stage_records_background_source_availability() { panic!("prepare did not become available, last state: {last_state:?}"); } +#[cfg(not(all(feature = "mlx", target_os = "macos")))] +#[tokio::test] +async fn prepare_mlx_fails_closed_without_downloading_on_unsupported_build() { + let file = tempfile::NamedTempFile::new().unwrap(); + let mut load = load_request(); + load.backend = "mlx".to_string(); + load.load_mode = LoadMode::ArtifactSlice; + load.package_ref = format!("hf-model://org/model@{}", "a".repeat(40)); + load.model_path = Some(file.path().to_string_lossy().into_owned()); + load.downstream = None; + let mut state = StageControlState::default(); + + let accepted = state + .prepare(StagePrepareRequest { + load: load.clone(), + coordinator_id: None, + }) + .await + .unwrap(); + assert!(accepted.accepted); + + for _ in 0..20 { + let inventory = state + .inventory(StageInventoryRequest { + model_id: load.model_id.clone(), + package_ref: load.package_ref.clone(), + manifest_sha256: load.manifest_sha256.clone(), + }) + .await; + if let Some(status) = inventory + .preparing_ranges + .iter() + .find(|status| status.stage_id == load.stage_id) + && status.state == StagePreparationState::Failed + { + assert!( + status + .error + .as_deref() + .unwrap_or_default() + .contains("unsupported stage backend 'mlx'") + ); + return; + } + tokio::time::sleep(Duration::from_millis(25)).await; + } + + panic!("unsupported MLX prepare did not fail closed"); +} + #[tokio::test] async fn prepare_layer_package_stays_downloading_while_peer_prefetch_is_pending() { let mut load = load_request(); diff --git a/crates/mesh-llm-host-runtime/src/runtime/local.rs b/crates/mesh-llm-host-runtime/src/runtime/local.rs index 4e2060a0a8..3f22988675 100644 --- a/crates/mesh-llm-host-runtime/src/runtime/local.rs +++ b/crates/mesh-llm-host-runtime/src/runtime/local.rs @@ -3549,18 +3549,12 @@ fn split_stage_source_is_ready( if ready_running_stage { return true; } - if load.load_mode != LoadMode::LayerPackage && !skippy::is_layer_package_ref(&load.package_ref) - { - return inventory - .available_ranges - .iter() - .any(|range| split_layer_range_covers(range, load)); - } - inventory.preparing_ranges.iter().any(|status| { + let prepared_stage_source = inventory.preparing_ranges.iter().any(|status| { status.topology_id == load.topology_id && status.run_id == load.run_id && status.stage_id == load.stage_id && status.model_id == load.model_id + && status.backend == load.backend && status.package_ref == load.package_ref && status.manifest_sha256 == load.manifest_sha256 && status.layer_start <= load.layer_start @@ -3569,7 +3563,18 @@ fn split_stage_source_is_ready( status.state, skippy::StagePreparationState::Available | skippy::StagePreparationState::Ready ) - }) + }); + if prepared_stage_source { + return true; + } + if load.load_mode != LoadMode::LayerPackage && !skippy::is_layer_package_ref(&load.package_ref) + { + return inventory + .available_ranges + .iter() + .any(|range| split_layer_range_covers(range, load)); + } + false } fn split_layer_range_covers(range: &skippy::LayerRange, load: &skippy::StageLoadRequest) -> bool { @@ -5043,6 +5048,37 @@ max_tokens = 222 assert!(split_stage_source_is_ready(&inventory, &load)); } + #[test] + fn mlx_artifact_slice_accepts_exact_prepare_availability() { + let mut load = stage_load_request(LoadMode::ArtifactSlice); + load.backend = "mlx".to_string(); + load.package_ref = format!( + "hf-model://HuggingFaceTB/SmolLM2-135M-Instruct@{}", + "a".repeat(40) + ); + let mut inventory = skippy::StageLayerInventory { + model_id: load.model_id.clone(), + package_ref: load.package_ref.clone(), + manifest_sha256: load.manifest_sha256.clone(), + layer_count: 0, + ready_ranges: Vec::new(), + available_ranges: Vec::new(), + missing_ranges: Vec::new(), + preparing_ranges: Vec::new(), + source_model_path: None, + source_model_bytes: None, + source_model_kind: skippy::SourceModelKind::Unknown, + }; + + assert!(!split_stage_source_is_ready(&inventory, &load)); + + inventory + .preparing_ranges + .push(test_preparation_status_from_load(&load)); + + assert!(split_stage_source_is_ready(&inventory, &load)); + } + #[test] fn runtime_slice_stage_source_accepts_inventory_availability() { let load = stage_load_request(LoadMode::RuntimeSlice); diff --git a/crates/skippy-engine-mlx/STAGED_EXECUTION.md b/crates/skippy-engine-mlx/STAGED_EXECUTION.md index 7193d16282..67319189ea 100644 --- a/crates/skippy-engine-mlx/STAGED_EXECUTION.md +++ b/crates/skippy-engine-mlx/STAGED_EXECUTION.md @@ -6,7 +6,8 @@ Dense Llama-family MLX stages now run as separate OS processes from partial SafeTensors artifacts and communicate over Skippy's existing binary stage wire. The first proof uses `HuggingFaceTB/SmolLM2-135M-Instruct` split at layer 15. -This is a production-shaped bridge, not yet the default mesh launch path: +This is a production-shaped bridge with an explicit host control path, not yet +the default automatic mesh launch path: - `skippy-engine` owns the engine-neutral `StageEngine` contract and residual buffer descriptors. @@ -19,6 +20,9 @@ This is a production-shaped bridge, not yet the default mesh launch path: per-session KV caches on a dedicated MLX worker thread, and executes only its configured layer range. - `mlx-stage` starts a stage process or drives a chain as a proof client. +- `StagePrepare` / `StageLoad` with `backend=mlx` and an immutable + `hf-model://org/repo@` reference now materialize and start the same + engine through the normal host stage-control loop. No process in the proof has access to the complete checkpoint. The tokenizer and config files are small shared metadata; tensor data comes only from that @@ -44,6 +48,22 @@ same prompt across prompt prefill and seven subsequent decode calls. Each stage kept an independent per-layer KV cache, and `Stop` cleared the session in both processes. +The host-managed proof also passed from a clean MLX stage cache. Both ranges +shared checkpoint identity +`303b5a31e5226edb03a48f6f77464736a91a404b1500f385ec43d0951ce81e87`, +but retained distinct stage cache keys: + +| Layers | Planned HTTP payload | Complete source shard | Avoided | Requests | +| --- | ---: | ---: | ---: | ---: | +| `0..15` | 162,857,381 bytes | 269,060,552 bytes | 106,204,032 bytes | 3 | +| `15..30` | 162,858,533 bytes | 269,060,552 bytes | 106,202,880 bytes | 4 | + +The test submitted Prepare, polled inventory, submitted Load, checked the +materialized status identity/path, generated the same eight reference tokens, +and submitted Stop through `spawn_stage_control_loop`. The runtime status does +not mislabel the derived slice as the full source model or claim a cache pin +that does not exist. + The two partial files are the exact-range artifacts described in `../../spikes/mlx-safetensors-stages/FINDINGS.md`. Tied input/output embeddings are intentionally duplicated across the stages; that is why the sum of the two @@ -96,6 +116,8 @@ just mlx-stage prove --connect 127.0.0.1:19090 --wire-dtype f16 - `engine_transport` is the reduced compatibility lane. The mature llama.cpp binary server remains unchanged and still owns telemetry, exact-prefix cache, batching, and OpenAI orchestration. -- Mesh topology planning does not yet launch `MlxStageEngine`; the next product - step is selecting this engine from stage config and advertising it as an - additive capability. There is no mesh protocol or Skippy ABI break here. +- Mesh topology planning does not yet produce MLX stage assignments. The host + can consume explicit `backend=mlx` Prepare/Load requests, but automatic + placement, capability advertisement, coordinator model planning, and an + OpenAI stage-0 frontend remain. There is no mesh protocol or Skippy ABI break + in the explicit consumer path. diff --git a/crates/skippy-engine-mlx/src/bin/mlx-stage.rs b/crates/skippy-engine-mlx/src/bin/mlx-stage.rs index db6ee4817b..f5a9ac5145 100644 --- a/crates/skippy-engine-mlx/src/bin/mlx-stage.rs +++ b/crates/skippy-engine-mlx/src/bin/mlx-stage.rs @@ -113,6 +113,7 @@ mod real { layer_start, layer_end, compute_dtype: compute_dtype.into(), + ctx_size: None, }, EngineStageServerOptions { bind_addr: bind, diff --git a/crates/skippy-engine-mlx/src/stage.rs b/crates/skippy-engine-mlx/src/stage.rs index c3a49ed3e1..41bf81871a 100644 --- a/crates/skippy-engine-mlx/src/stage.rs +++ b/crates/skippy-engine-mlx/src/stage.rs @@ -50,6 +50,7 @@ pub struct MlxStageEngineConfig { pub layer_start: u32, pub layer_end: u32, pub compute_dtype: MlxComputeDtype, + pub ctx_size: Option, } enum WorkerJob { @@ -83,6 +84,10 @@ impl MlxStageEngine { } } + pub fn stage_info(&self) -> &StageEngineInfo { + &self.info + } + fn request( &self, make_job: impl FnOnce(mpsc::Sender>) -> WorkerJob, @@ -116,6 +121,7 @@ struct LoadedStage { model: llama::Model, stream: Stream, compute_dtype: Dtype, + ctx_size: Option, info: StageEngineInfo, sessions: BTreeMap>>, } @@ -190,6 +196,7 @@ fn load_stage(config: MlxStageEngineConfig) -> Result { model, stream, compute_dtype: config.compute_dtype.mlx(), + ctx_size: config.ctx_size.map(usize::try_from).transpose()?, info, sessions: BTreeMap::new(), }) @@ -263,6 +270,7 @@ impl LoadedStage { } ensure!(!request.token_ids.is_empty(), "stage request has no tokens"); let token_count = request.token_ids.len(); + self.ensure_context_capacity(request.session_id, token_count)?; if let Some(input) = request.input.as_ref() { ensure!( input.token_count == token_count, @@ -307,6 +315,27 @@ impl LoadedStage { }) } + fn ensure_context_capacity(&self, session_id: u64, token_count: usize) -> Result<()> { + let Some(ctx_size) = self.ctx_size else { + return Ok(()); + }; + let offset = self + .sessions + .get(&session_id) + .and_then(|caches| caches.first()) + .and_then(Option::as_ref) + .map(KeyValueCache::offset) + .map(usize::try_from) + .transpose()? + .unwrap_or_default(); + ensure!( + offset.saturating_add(token_count) <= ctx_size, + "MLX stage context limit {ctx_size} exceeded by {} tokens", + offset.saturating_add(token_count) + ); + Ok(()) + } + fn input_hidden(&mut self, request: &StageExecutionRequest) -> Result { if self.info.is_first() { ensure!( diff --git a/crates/skippy-server/src/embedded.rs b/crates/skippy-server/src/embedded.rs index f0286f94a4..c0a348b7bb 100644 --- a/crates/skippy-server/src/embedded.rs +++ b/crates/skippy-server/src/embedded.rs @@ -1,16 +1,23 @@ use std::{ net::SocketAddr, - sync::{Arc, Mutex}, + sync::{ + Arc, Mutex, + atomic::{AtomicBool, Ordering}, + }, }; use anyhow::{Context, Result}; use openai_frontend::OpenAiBackend; +use skippy_engine::StageEngine; use skippy_protocol::{StageConfig, StageTopology}; use tokio::{sync::oneshot, task::JoinHandle}; use crate::{ binary_transport::{BinaryStageOptions, serve_binary_stage_with_shutdown}, config::validate_config, + engine_transport::{ + EngineStageServerOptions, prepare_stage_engine_listener, serve_prepared_stage_engine_until, + }, frontend::{EmbeddedOpenAiArgs, serve_embedded_openai_with_shutdown}, http::{StageHttpOptions, serve_stage_http_with_shutdown}, runtime_state::{ @@ -350,6 +357,49 @@ pub fn start_binary_stage(options: BinaryStageOptions) -> EmbeddedServerHandle { } } +pub fn start_stage_engine( + engine: Arc, + options: EngineStageServerOptions, +) -> EmbeddedServerHandle { + let bind_addr = options.bind_addr; + let status = Arc::new(Mutex::new(ServerHandleState { + name: "engine-stage", + bind_addr, + state: EmbeddedState::Starting, + started_at_unix_nanos: now_unix_nanos(), + stopped_at_unix_nanos: None, + last_error: None, + })); + let (shutdown_tx, shutdown_rx) = oneshot::channel(); + let shutdown = Arc::new(AtomicBool::new(false)); + let task_status = Arc::clone(&status); + let task_shutdown = Arc::clone(&shutdown); + let runtime = tokio::runtime::Handle::current(); + let task = tokio::task::spawn_blocking(move || { + let watcher_shutdown = Arc::clone(&task_shutdown); + let watcher = runtime.spawn(async move { + let _ = shutdown_rx.await; + watcher_shutdown.store(true, Ordering::Release); + }); + let result = + prepare_stage_engine_listener(engine.as_ref(), &options).and_then(|listener| { + { + let mut status = task_status.lock().expect("server status lock poisoned"); + status.state = EmbeddedState::Ready; + } + serve_prepared_stage_engine_until(engine, options, listener, task_shutdown) + }); + watcher.abort(); + finish_server_status(&task_status, &result); + result + }); + EmbeddedServerHandle { + status, + shutdown: Some(shutdown_tx), + task: Some(task), + } +} + fn spawn_async_server( name: &'static str, bind_addr: SocketAddr, @@ -398,3 +448,158 @@ fn finish_server_status(status: &Arc>, result: &Result< } } } + +#[cfg(test)] +mod tests { + use std::{ + io::Read, + net::{TcpListener, TcpStream}, + time::Duration, + }; + + use skippy_engine::{StageEngineInfo, StageExecutionOutput, StageExecutionRequest}; + use skippy_protocol::binary::{WireActivationDType, recv_ready}; + + use super::*; + + struct FakeStageEngine { + info: StageEngineInfo, + } + + impl StageEngine for FakeStageEngine { + fn info(&self) -> &StageEngineInfo { + &self.info + } + + fn execute(&self, _request: StageExecutionRequest) -> Result { + Ok(StageExecutionOutput::default()) + } + + fn reset_session(&self, _session_id: u64) -> Result<()> { + Ok(()) + } + } + + #[tokio::test] + async fn engine_stage_reports_bind_failure_without_false_ready() { + let occupied = TcpListener::bind("127.0.0.1:0").unwrap(); + let bind_addr = occupied.local_addr().unwrap(); + let handle = start_stage_engine(final_engine(), options(bind_addr, None)); + + let status = wait_for_state(&handle, EmbeddedState::Failed).await; + + assert!(status.last_error.unwrap().contains("bind engine stage")); + assert!(handle.shutdown().await.is_err()); + } + + #[tokio::test] + async fn engine_stage_reports_invalid_topology_without_false_ready() { + let handle = start_stage_engine(non_final_engine(), options(unused_addr(), None)); + + let status = wait_for_state(&handle, EmbeddedState::Failed).await; + + assert!(status.last_error.unwrap().contains("only the final stage")); + assert!(handle.shutdown().await.is_err()); + } + + #[tokio::test] + async fn engine_stage_preflights_unreachable_downstream() { + let handle = start_stage_engine( + non_final_engine(), + options(unused_addr(), Some(unused_addr())), + ); + + let status = wait_for_state(&handle, EmbeddedState::Failed).await; + + assert!(status.last_error.unwrap().contains("preflight downstream")); + assert!(handle.shutdown().await.is_err()); + } + + #[tokio::test] + async fn engine_stage_shutdown_closes_and_joins_live_connections() { + let bind_addr = unused_addr(); + let handle = start_stage_engine(final_engine(), options(bind_addr, None)); + wait_for_state(&handle, EmbeddedState::Ready).await; + let mut client = TcpStream::connect(bind_addr).unwrap(); + client + .set_read_timeout(Some(Duration::from_secs(2))) + .unwrap(); + recv_ready(&mut client).unwrap(); + + handle.shutdown().await.unwrap(); + + let mut byte = [0_u8; 1]; + let closed = match client.read(&mut byte) { + Ok(0) => true, + Err(error) => matches!( + error.kind(), + std::io::ErrorKind::ConnectionReset | std::io::ErrorKind::BrokenPipe + ), + _ => false, + }; + assert!(closed, "live engine-stage connection survived shutdown"); + } + + fn final_engine() -> Arc { + engine(StageEngineInfo { + engine: "fake".to_string(), + model_id: "fake-model".to_string(), + stage_index: 0, + layer_start: 0, + layer_end: 1, + total_layers: 1, + activation_width: 2, + }) + } + + fn non_final_engine() -> Arc { + engine(StageEngineInfo { + engine: "fake".to_string(), + model_id: "fake-model".to_string(), + stage_index: 0, + layer_start: 0, + layer_end: 1, + total_layers: 2, + activation_width: 2, + }) + } + + fn engine(info: StageEngineInfo) -> Arc { + Arc::new(FakeStageEngine { info }) + } + + fn options( + bind_addr: SocketAddr, + downstream_addr: Option, + ) -> EngineStageServerOptions { + EngineStageServerOptions { + bind_addr, + downstream_addr, + wire_dtype: WireActivationDType::F16, + } + } + + fn unused_addr() -> SocketAddr { + TcpListener::bind("127.0.0.1:0") + .unwrap() + .local_addr() + .unwrap() + } + + async fn wait_for_state( + handle: &EmbeddedServerHandle, + expected: EmbeddedState, + ) -> EmbeddedServerStatus { + tokio::time::timeout(Duration::from_secs(2), async { + loop { + let status = handle.status(); + if status.state == expected { + return status; + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .expect("engine stage did not reach expected state") + } +} diff --git a/crates/skippy-server/src/engine_transport.rs b/crates/skippy-server/src/engine_transport.rs index b39bcaf433..69ceff7a13 100644 --- a/crates/skippy-server/src/engine_transport.rs +++ b/crates/skippy-server/src/engine_transport.rs @@ -8,9 +8,9 @@ use std::{ io::{self, Write}, - net::{SocketAddr, TcpListener, TcpStream}, + net::{Shutdown, SocketAddr, TcpListener, TcpStream}, sync::{ - Arc, + Arc, Mutex, atomic::{AtomicBool, Ordering}, }, thread, @@ -47,11 +47,26 @@ pub fn serve_stage_engine_until( options: EngineStageServerOptions, shutdown: Arc, ) -> Result<()> { + let listener = prepare_stage_engine_listener(engine.as_ref(), &options)?; + serve_prepared_stage_engine_until(engine, options, listener, shutdown) +} + +pub(crate) fn prepare_stage_engine_listener( + engine: &dyn StageEngine, + options: &EngineStageServerOptions, +) -> Result { engine.info().validate()?; - validate_topology(engine.as_ref(), &options)?; + validate_topology(engine, options)?; let listener = TcpListener::bind(options.bind_addr) .with_context(|| format!("bind engine stage at {}", options.bind_addr))?; listener.set_nonblocking(true)?; + if let Some(downstream_addr) = options.downstream_addr { + drop( + connect_downstream(downstream_addr).with_context(|| { + format!("preflight downstream engine stage at {downstream_addr}") + })?, + ); + } eprintln!( "skippy engine stage listening: engine={} model={} binary={} layers={}..{} width={} dtype={:?}", engine.info().engine, @@ -62,27 +77,83 @@ pub fn serve_stage_engine_until( engine.info().activation_width, options.wire_dtype, ); + Ok(listener) +} - while !shutdown.load(Ordering::SeqCst) { - let (upstream, peer_addr) = match listener.accept() { - Ok(connection) => connection, - Err(error) if error.kind() == io::ErrorKind::WouldBlock => { - thread::sleep(Duration::from_millis(25)); - continue; - } - Err(error) => return Err(error).context("accept engine stage connection"), - }; - upstream.set_nonblocking(false)?; - upstream.set_nodelay(true).ok(); - let engine = engine.clone(); - let options = options.clone(); - thread::spawn(move || { - if let Err(error) = handle_connection(engine, options, upstream) { - eprintln!("engine stage connection from {peer_addr} failed: {error:#}"); - } - }); +pub(crate) fn serve_prepared_stage_engine_until( + engine: Arc, + options: EngineStageServerOptions, + listener: TcpListener, + shutdown: Arc, +) -> Result<()> { + let mut connections = Vec::new(); + let result = (|| { + while !shutdown.load(Ordering::SeqCst) { + reap_finished_connections(&mut connections); + let (upstream, peer_addr) = match listener.accept() { + Ok(connection) => connection, + Err(error) if error.kind() == io::ErrorKind::WouldBlock => { + thread::sleep(Duration::from_millis(25)); + continue; + } + Err(error) => return Err(error).context("accept engine stage connection"), + }; + upstream.set_nonblocking(false)?; + upstream.set_nodelay(true).ok(); + let control = upstream.try_clone()?; + let downstream_control = Arc::new(Mutex::new(None)); + let engine = engine.clone(); + let options = options.clone(); + let connection_downstream = Arc::clone(&downstream_control); + let connection = thread::spawn(move || { + if let Err(error) = + handle_connection(engine, options, upstream, connection_downstream) + { + eprintln!("engine stage connection from {peer_addr} failed: {error:#}"); + } + }); + connections.push(ActiveConnection { + control, + downstream_control, + thread: connection, + }); + } + Ok(()) + })(); + stop_active_connections(connections); + result +} + +struct ActiveConnection { + control: TcpStream, + downstream_control: Arc>>, + thread: thread::JoinHandle<()>, +} + +fn reap_finished_connections(connections: &mut Vec) { + let mut index = 0; + while index < connections.len() { + if connections[index].thread.is_finished() { + let connection = connections.swap_remove(index); + let _ = connection.thread.join(); + } else { + index += 1; + } + } +} + +fn stop_active_connections(connections: Vec) { + for connection in &connections { + let _ = connection.control.shutdown(Shutdown::Both); + if let Ok(downstream) = connection.downstream_control.lock() + && let Some(downstream) = downstream.as_ref() + { + let _ = downstream.shutdown(Shutdown::Both); + } + } + for connection in connections { + let _ = connection.thread.join(); } - Ok(()) } fn validate_topology(engine: &dyn StageEngine, options: &EngineStageServerOptions) -> Result<()> { @@ -97,13 +168,19 @@ fn handle_connection( engine: Arc, options: EngineStageServerOptions, mut upstream: TcpStream, + downstream_control: Arc>>, ) -> Result<()> { - send_ready(&mut upstream).context("send engine stage ready")?; - upstream.flush().ok(); let mut downstream = options .downstream_addr .map(connect_downstream) .transpose()?; + if let Some(stream) = downstream.as_ref() { + *downstream_control + .lock() + .expect("downstream control lock poisoned") = Some(stream.try_clone()?); + } + send_ready(&mut upstream).context("send engine stage ready")?; + upstream.flush().ok(); let activation_width = i32::try_from(engine.info().activation_width).context("activation width exceeds i32")?; @@ -146,10 +223,15 @@ fn handle_connection( } fn connect_downstream(addr: SocketAddr) -> Result { - let mut stream = TcpStream::connect(addr) + const CONNECT_TIMEOUT: Duration = Duration::from_secs(5); + let mut stream = TcpStream::connect_timeout(&addr, CONNECT_TIMEOUT) .with_context(|| format!("connect downstream engine stage at {addr}"))?; stream.set_nodelay(true).ok(); + stream.set_read_timeout(Some(CONNECT_TIMEOUT))?; + stream.set_write_timeout(Some(CONNECT_TIMEOUT))?; recv_ready(&mut stream).context("downstream engine stage did not become ready")?; + stream.set_read_timeout(None)?; + stream.set_write_timeout(None)?; Ok(stream) } diff --git a/crates/skippy-server/src/lib.rs b/crates/skippy-server/src/lib.rs index 7bdbaf3058..6ff0ddc138 100644 --- a/crates/skippy-server/src/lib.rs +++ b/crates/skippy-server/src/lib.rs @@ -24,7 +24,7 @@ pub use cli::ServeBinaryArgs; pub use embedded::{ EmbeddedRuntimeOptions, EmbeddedRuntimeStatus, EmbeddedServerHandle, EmbeddedServerStatus, EmbeddedState, SkippyRuntimeHandle, start_binary_stage, start_embedded_openai, - start_openai_backend, start_stage_http, + start_openai_backend, start_stage_engine, start_stage_http, }; pub use frontend::{ CONTEXT_BUDGET_MAX_TOKENS, DEFAULT_EMBEDDED_MAX_TOKENS, EmbeddedOpenAiArgs, diff --git a/docs/design/MLX_STAGE_ENGINE_PLAN.md b/docs/design/MLX_STAGE_ENGINE_PLAN.md index 7856d0268f..db0cb13101 100644 --- a/docs/design/MLX_STAGE_ENGINE_PLAN.md +++ b/docs/design/MLX_STAGE_ENGINE_PLAN.md @@ -1,6 +1,6 @@ # MLX as a Skippy Stage Engine — Deep Dive and Plan -## Status: exploratory design proposal +## Status: implemented prototype and remaining research plan This document evaluates using Apple **MLX** (via the Rust `safemlx` / `safemlx-lm` crates) as an alternative inference engine behind Skippy's staged execution @@ -47,6 +47,19 @@ dense execution and process-boundary proof. Host topology selection, advanced cache/session operations, additional families, and bounded-memory quantization remain. See `crates/skippy-engine-mlx/STAGED_EXECUTION.md`. +**Update — host `StagePrepare` / `StageLoad` now consumes range-only MLX +stages.** An immutable `hf-model://org/repo@` request is validated +before network work, materialized on a blocking worker, checked against a +topology-wide checkpoint identity, strict-loaded into `MlxStageEngine`, and +served through the existing Skippy binary wire. A clean-cache SmolLM2 proof ran +both 15-layer ranges through `spawn_stage_control_loop`, reproduced the same +eight reference tokens, and stopped both stages. Each node-side range plan read +about 162.86 MB of a 269.06 MB source shard and avoided about 106.2 MB. Startup +now fails closed on bind/topology/downstream errors, and Stop closes and joins +active connections. Automatic MLX topology production, capability +advertisement, remote two-node proof, and bounded-memory load-time quantization +remain; this checkpoint proves the host consumer path, not automatic placement. + --- ## 1. Bottom line From 6e4693418a509eac79c0ad7e8f65adb441d0186c Mon Sep 17 00:00:00 2001 From: Michael Neale <14976+michaelneale@users.noreply.github.com> Date: Fri, 17 Jul 2026 19:45:22 +1000 Subject: [PATCH 14/37] feat(mlx): quantize partial stages on load --- .../src/inference/skippy/stage/mlx.rs | 1 + crates/skippy-engine-mlx/STAGED_EXECUTION.md | 33 +++++- crates/skippy-engine-mlx/src/bin/mlx-stage.rs | 39 +++++- crates/skippy-engine-mlx/src/lib.rs | 2 +- crates/skippy-engine-mlx/src/stage.rs | 103 ++++++++++++++-- docs/design/MLX_STAGE_ENGINE_PLAN.md | 61 +++++++--- spikes/mlx-safetensors-stages/FINDINGS.md | 111 +++++++++++++----- 7 files changed, 284 insertions(+), 66 deletions(-) diff --git a/crates/mesh-llm-host-runtime/src/inference/skippy/stage/mlx.rs b/crates/mesh-llm-host-runtime/src/inference/skippy/stage/mlx.rs index 1e1c5c6883..f0b3c548bd 100644 --- a/crates/mesh-llm-host-runtime/src/inference/skippy/stage/mlx.rs +++ b/crates/mesh-llm-host-runtime/src/inference/skippy/stage/mlx.rs @@ -128,6 +128,7 @@ pub(super) async fn launch_stage( layer_start: blocking_load.layer_start, layer_end: blocking_load.layer_end, compute_dtype: MlxComputeDtype::Bf16, + weight_quantization: None, ctx_size: Some(blocking_load.ctx_size), })?); ensure!( diff --git a/crates/skippy-engine-mlx/STAGED_EXECUTION.md b/crates/skippy-engine-mlx/STAGED_EXECUTION.md index 67319189ea..461e6423f7 100644 --- a/crates/skippy-engine-mlx/STAGED_EXECUTION.md +++ b/crates/skippy-engine-mlx/STAGED_EXECUTION.md @@ -64,6 +64,24 @@ and submitted Stop through `spawn_stage_control_loop`. The runtime status does not mislabel the derived slice as the full source model or claim a cache pin that does not exist. +The next engine-level proof enabled tensor-at-a-time JIT weight quantization. +The pinned safemlx loader visits one dense tensor at a time, quantizes it, and +eagerly evaluates and synchronizes the packed weight/scales/biases before +visiting the next tensor. This bounds the lazy graph, but the TensorView and +stream-copy path can temporarily hold more than one physical source copy. With +affine 4-bit, group size 64, the whole 30-layer reference and the two +independently loaded 15-layer stages generated the same quantized-model tokens: + +```text +[260, 2240, 314, 253, 1379, 282, 25801, 28] +``` + +The two-stage processes retained 349 MLX parameters each and had post-proof RSS +of 87,392 KiB and 87,952 KiB, versus roughly 189 MiB each at source precision. +This proves deterministic per-stage quantization and quantized stage execution; +it does not yet prove peak RSS, direct range-to-quantized-cache disk bounds, or +host/topology selection of the quantization profile. + The two partial files are the exact-range artifacts described in `../../spikes/mlx-safetensors-stages/FINDINGS.md`. Tied input/output embeddings are intentionally duplicated across the stages; that is why the sum of the two @@ -88,6 +106,10 @@ just mlx-stage serve \ --bind 127.0.0.1:19091 --wire-dtype f16 --compute-dtype bf16 ``` +Add `--weight-quantization affine4` to both stage commands to reproduce the +JIT-quantized proof, then pass +`--expected 260,2240,314,253,1379,282,25801,28` to `mlx-stage prove`. + Start the first stage in another terminal: ```bash @@ -107,8 +129,9 @@ just mlx-stage prove --connect 127.0.0.1:19090 --wire-dtype f16 ## Deliberate limitations of this checkpoint -- Dense Llama-family checkpoints only. The engine boundary is family-neutral, - but the current MLX adapter is the smallest implementation that proves it. +- Dense Llama-family checkpoints only in `MlxStageEngine`. The pinned safemlx + revision has whole-model Inkling and Nemotron-H implementations, but neither + is exposed through this partial-stage adapter yet. - Greedy sampling only; sampling metadata is preserved in the contract and rejected explicitly when enabled. - No KV page import/export, cache trim/checkpoint, MTP, speculative verify, @@ -119,5 +142,7 @@ just mlx-stage prove --connect 127.0.0.1:19090 --wire-dtype f16 - Mesh topology planning does not yet produce MLX stage assignments. The host can consume explicit `backend=mlx` Prepare/Load requests, but automatic placement, capability advertisement, coordinator model planning, and an - OpenAI stage-0 frontend remain. There is no mesh protocol or Skippy ABI break - in the explicit consumer path. + OpenAI stage-0 frontend remain. Explicit host requests still use source + precision; quantization selection currently exists only in the engine config + and `mlx-stage` proof CLI. There is no mesh protocol or Skippy ABI break in + the explicit consumer path. diff --git a/crates/skippy-engine-mlx/src/bin/mlx-stage.rs b/crates/skippy-engine-mlx/src/bin/mlx-stage.rs index f5a9ac5145..e1025e838c 100644 --- a/crates/skippy-engine-mlx/src/bin/mlx-stage.rs +++ b/crates/skippy-engine-mlx/src/bin/mlx-stage.rs @@ -11,7 +11,9 @@ mod real { use anyhow::{Context, Result, ensure}; use clap::{Parser, Subcommand, ValueEnum}; - use skippy_engine_mlx::{MlxComputeDtype, MlxStageEngine, MlxStageEngineConfig}; + use skippy_engine_mlx::{ + MlxComputeDtype, MlxStageEngine, MlxStageEngineConfig, MlxWeightQuantization, + }; use skippy_protocol::binary::{ StageStateHeader, StageWireMessage, WireActivationDType, WireMessageKind, WireReplyKind, recv_ready, recv_reply, write_stage_message, @@ -47,6 +49,8 @@ mod real { wire_dtype: WireDtype, #[arg(long, value_enum, default_value_t = ComputeDtype::Bf16)] compute_dtype: ComputeDtype, + #[arg(long, value_enum)] + weight_quantization: Option, }, /// Drive a stage chain and assert its greedy token sequence. Prove { @@ -93,6 +97,29 @@ mod real { } } + #[derive(Clone, Copy, Debug, ValueEnum)] + enum WeightQuantization { + Affine4, + Affine8, + Mxfp4, + } + + impl From for MlxWeightQuantization { + fn from(value: WeightQuantization) -> Self { + match value { + WeightQuantization::Affine4 => Self::Affine { + group_size: 64, + bits: 4, + }, + WeightQuantization::Affine8 => Self::Affine { + group_size: 64, + bits: 8, + }, + WeightQuantization::Mxfp4 => Self::MxFp4, + } + } + } + pub fn main() -> Result<()> { match Cli::parse().command { Command::Serve { @@ -105,6 +132,7 @@ mod real { downstream, wire_dtype, compute_dtype, + weight_quantization, } => serve( MlxStageEngineConfig { model_dir: model, @@ -113,6 +141,7 @@ mod real { layer_start, layer_end, compute_dtype: compute_dtype.into(), + weight_quantization: weight_quantization.map(Into::into), ctx_size: None, }, EngineStageServerOptions { @@ -180,15 +209,15 @@ mod real { generated.push(send_predicted(&mut stream, &decode, wire_dtype)?); } - ensure!( - generated == expected, - "two-process stage tokens diverged: expected={expected:?} actual={generated:?}" - ); let stop = StageWireMessage::stop_with_identity(wire_dtype, request_id, session_id); write_stage_message(&mut stream, &stop, wire_dtype)?; stream.flush().ok(); let reply = recv_reply(&mut stream)?; ensure!(reply.kind == WireReplyKind::Ack, "stop did not return ACK"); + ensure!( + generated == expected, + "two-process stage tokens diverged: expected={expected:?} actual={generated:?}" + ); println!("PASS: two MLX stage processes matched the reference greedy tokens"); println!("wire_dtype={wire_dtype:?}"); println!("generated_tokens={generated:?}"); diff --git a/crates/skippy-engine-mlx/src/lib.rs b/crates/skippy-engine-mlx/src/lib.rs index c948045cfd..16ad2e15bf 100644 --- a/crates/skippy-engine-mlx/src/lib.rs +++ b/crates/skippy-engine-mlx/src/lib.rs @@ -19,7 +19,7 @@ pub use backend::MlxBackend; #[cfg(all(feature = "mlx", target_os = "macos"))] pub use engine::{ChatTurn, GenerateRequest, MlxEngine, MlxEngineConfig}; #[cfg(all(feature = "mlx", target_os = "macos"))] -pub use stage::{MlxComputeDtype, MlxStageEngine, MlxStageEngineConfig}; +pub use stage::{MlxComputeDtype, MlxStageEngine, MlxStageEngineConfig, MlxWeightQuantization}; /// True when this build actually contains the MLX engine. pub const fn mlx_available() -> bool { diff --git a/crates/skippy-engine-mlx/src/stage.rs b/crates/skippy-engine-mlx/src/stage.rs index 41bf81871a..a6cc6f7130 100644 --- a/crates/skippy-engine-mlx/src/stage.rs +++ b/crates/skippy-engine-mlx/src/stage.rs @@ -17,7 +17,11 @@ use safemlx_lm::{ common::linear::project_logits_maybe_quantized, llama::{self, AttentionInput, TransformerBlock}, }, - weights::{StrictLoadConfig, StrictLoadReport, load_safetensors_strict}, + quantization::{AffineQuantization, WeightQuantization}, + weights::{ + StrictLoadConfig, StrictLoadReport, load_safetensors_quantized_strict, + load_safetensors_strict, + }, }; use skippy_engine::{ StageActivation, StageEngine, StageEngineInfo, StageExecutionKind, StageExecutionOutput, @@ -42,6 +46,30 @@ impl MlxComputeDtype { } } +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum MlxWeightQuantization { + Affine { group_size: i32, bits: i32 }, + MxFp4, +} + +impl MlxWeightQuantization { + fn safemlx(self) -> Result { + match self { + Self::Affine { group_size, bits } => { + Ok(AffineQuantization::new(group_size, bits)?.into()) + } + Self::MxFp4 => Ok(WeightQuantization::MxFp4), + } + } + + fn label(self) -> String { + match self { + Self::Affine { group_size, bits } => format!("affine-{bits}bit-g{group_size}"), + Self::MxFp4 => "mxfp4".to_string(), + } + } +} + #[derive(Clone, Debug)] pub struct MlxStageEngineConfig { pub model_dir: PathBuf, @@ -50,6 +78,7 @@ pub struct MlxStageEngineConfig { pub layer_start: u32, pub layer_end: u32, pub compute_dtype: MlxComputeDtype, + pub weight_quantization: Option, pub ctx_size: Option, } @@ -157,7 +186,7 @@ fn run_worker( fn load_stage(config: MlxStageEngineConfig) -> Result { let stream = Stream::new_with_device(&Device::new(DeviceType::Gpu, 0)); let weights_stream = Stream::new_with_device(&Device::new(DeviceType::Cpu, 0)); - let model_args = llama::get_llama_model_args(&config.model_dir)?; + let mut model_args = llama::get_llama_model_args(&config.model_dir)?; let total_layers = u32::try_from(model_args.num_hidden_layers)?; let info = StageEngineInfo { engine: "mlx".to_string(), @@ -170,27 +199,52 @@ fn load_stage(config: MlxStageEngineConfig) -> Result { }; info.validate()?; + let quantization = config + .weight_quantization + .map(MlxWeightQuantization::safemlx) + .transpose()?; + if let Some(quantization) = quantization { + ensure!( + model_args.quantization.is_none() && model_args.quantization_config.is_none(), + "MLX stage load-time quantization requires a dense source checkpoint" + ); + model_args.quantization = Some(quantization); + } let mut model = llama::Model::new(model_args, &stream)?; let load_config = partial_stage_load_config(&info); let mut load_report = StrictLoadReport::default(); - load_safetensors_strict( - &mut model, - weight_file(&config.model_dir), - &weights_stream, - &load_config, - &mut load_report, - )?; + match quantization { + Some(quantization) => load_safetensors_quantized_strict( + &mut model, + weight_file(&config.model_dir), + &weights_stream, + &stream, + quantization, + &load_config, + &mut load_report, + )?, + None => load_safetensors_strict( + &mut model, + weight_file(&config.model_dir), + &weights_stream, + &load_config, + &mut load_report, + )?, + } load_report.finish(&model, &load_config)?; retain_local_layers(&mut model, info.layer_start, info.layer_end)?; copy_stage_weights_to_compute_stream(&mut model, &info, &stream)?; stream.synchronize()?; eprintln!( - "MLX partial stage loaded: model={} stage={} layers={}..{} tensors={}", + "MLX partial stage loaded: model={} stage={} layers={}..{} tensors={} weight_quantization={}", info.model_id, info.stage_index, info.layer_start, info.layer_end, model.parameters().flatten().len(), + config + .weight_quantization + .map_or_else(|| "none".to_string(), MlxWeightQuantization::label), ); Ok(LoadedStage { model, @@ -446,3 +500,32 @@ fn last_argmax(logits: &Array, stream: &Stream) -> Result { .context("cannot argmax empty logits")?; Ok(i32::try_from(index)?) } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn validates_requested_weight_quantization_before_model_load() { + let affine = MlxWeightQuantization::Affine { + group_size: 64, + bits: 4, + } + .safemlx() + .unwrap(); + assert_eq!(affine.group_size(), 64); + assert_eq!(affine.bits(), 4); + assert!( + MlxWeightQuantization::Affine { + group_size: 48, + bits: 4, + } + .safemlx() + .is_err() + ); + assert_eq!( + MlxWeightQuantization::MxFp4.safemlx().unwrap(), + WeightQuantization::MxFp4 + ); + } +} diff --git a/docs/design/MLX_STAGE_ENGINE_PLAN.md b/docs/design/MLX_STAGE_ENGINE_PLAN.md index db0cb13101..8aac8d6cf6 100644 --- a/docs/design/MLX_STAGE_ENGINE_PLAN.md +++ b/docs/design/MLX_STAGE_ENGINE_PLAN.md @@ -33,7 +33,9 @@ GiB of shard files; exact ranges avoid 833.15 GiB. A SmolLM2-135M proof then materialized two partial files (layers 0..15 and 15..30), loaded each directly into MLX, and matched unsplit logits exactly for prefill plus eight decode steps through Skippy's real F16 and F32 binary activation codec. The remaining -artifact gate is bounded-memory quantization for frontier-sized source tensors. +artifact gate is direct range-to-derived-cache quantization for frontier-sized +source tensors; tensor-at-a-time quantization into a live partial model is now +proven on SmolLM2. See `spikes/mlx-safetensors-stages/FINDINGS.md`. **Update — the first engine-neutral, multi-process stage chain is now proven.** @@ -44,8 +46,9 @@ own KV cache. Two real processes, each given only one 155.28 MiB partial SmolLM2 artifact, reproduced the eight-token whole-model reference exactly over F16 residuals. Their post-proof RSS was about 189 MiB each. This closes the dense execution and process-boundary proof. Host topology selection, advanced -cache/session operations, additional families, and bounded-memory quantization -remain. See `crates/skippy-engine-mlx/STAGED_EXECUTION.md`. +cache/session operations, additional staged families, and bounded +range-to-derived-cache quantization remain. See +`crates/skippy-engine-mlx/STAGED_EXECUTION.md`. **Update — host `StagePrepare` / `StageLoad` now consumes range-only MLX stages.** An immutable `hf-model://org/repo@` request is validated @@ -57,8 +60,25 @@ eight reference tokens, and stopped both stages. Each node-side range plan read about 162.86 MB of a 269.06 MB source shard and avoided about 106.2 MB. Startup now fails closed on bind/topology/downstream errors, and Stop closes and joins active connections. Automatic MLX topology production, capability -advertisement, remote two-node proof, and bounded-memory load-time quantization -remain; this checkpoint proves the host consumer path, not automatic placement. +advertisement, remote two-node proof, host quantization selection, and bounded +range-to-derived-cache materialization remain; this checkpoint proves the host +consumer path, not automatic placement. + +**Update — partial MLX stages now JIT-quantize one tensor at a time.** The +pinned safemlx strict loader already contains the required bounded lazy-graph +seam: visit one tensor, quantize, `eval`, synchronize, install packed parameters, +repeat. A whole-model affine-4 reference and two independently quantized +SmolLM2 stages generated the same eight tokens over F16 residuals. Post-proof +RSS was about 85.3 and 85.9 MiB per stage. This bounds the lazy graph, not +physical copies: TensorView conversion, stream copies, mmap pages, and MLX +scratch all require explicit high-water measurement. The host does not select +this profile yet, and the BF16 partial stage still exists on disk first. + +The pinned safemlx revision also already includes whole-model Inkling text, +vision, and audio execution. Earlier notes that called for porting Inkling were +stale. The remaining Inkling work is a partial-stage API plus wiring its +existing quantized grouped-expert runtime into the constructor and transformed +loader, not a model-family implementation from scratch. --- @@ -670,11 +690,14 @@ Implement `prefill_chunk_frame` / `decode_step_frame` / `copy_output_activation_frame` producing Skippy `ActivationFrame`s. Two-stage single-machine parity first, then two Macs over the real network. -> **Dense single-machine spike passed.** SmolLM2-135M was split 15+15 using two -> exact-range partial SafeTensors artifacts. F16 and F32 `StageWireMessage` -> boundaries both matched unsplit MLX with zero measured logit delta across -> prompt prefill and eight decode steps. This proves the basic artifact and -> activation seams, but not bounded-memory quantization or server integration. +> **Dense and JIT-quantized process proofs passed.** SmolLM2-135M was split +> 15+15 using two exact-range partial SafeTensors artifacts. F16 and F32 +> `StageWireMessage` boundaries matched unsplit MLX with zero measured dense +> logit delta; two real F16-wire processes also matched the whole-model +> affine-4 token reference after tensor-wise on-load quantization. Explicit host +> stage control now starts the source-precision path. Direct +> range-to-quantized-cache materialization, profile-bearing topology requests, +> and remote two-node execution remain. **Phase 4 — KV/state codec + verify + trim/checkpoint.** Implement the engine-general cache codec (§5.2), `verify_tokens_frame` for speculative decode, @@ -697,8 +720,10 @@ Apple-Silicon nodes. llama.cpp remains the cross-platform default. 1. **Partial-loading proof (DENSE GO, QUANT PARTIAL):** remote exact-range selection is proven, including on 1.9 TB Inkling BF16. SmolLM2 partial files were materialized and loaded without a complete checkpoint. Still required: - confirm peak RSS is bounded by quantized stage + one source tensor and - scratch during tensor-at-a-time load-time quantization. + The live-model loader now quantizes/evaluates one tensor at a time and split + output matches a whole quantized reference. Still required: measure cold-load + high-water RSS and stream HTTP ranges directly into a derived quantized cache + without retaining the whole BF16 stage slice on disk. 2. **Boundary latency breakdown (GO/NO-GO):** measure layer compute, cast, contiguous, **eval fence**, host readback, serialize, and receive-reconstruct **independently**, at hidden widths 4096/8192/16384 and token counts @@ -770,15 +795,19 @@ Spikes 1 and 2 are more decisive than any standalone token/s benchmark. ## 10. Immediate next steps -1. Add tensor-at-a-time MLX quantization to the materializer and measure peak - RSS against the one-source-tensor memory contract. +1. Expose a sequential selected-range callback/temp-artifact seam, add direct + range-to-quantized-cache shards, and measure peak RSS/disk use with OS and + MLX counters. The live model loader's tensor-wise quantization and split + correctness are proven, but physical source copies remain unbounded evidence. 2. Route normal `skippy-server` launch through the engine-neutral contract while retaining capability-gated llama-only batching/cache/MTP/multimodal paths; the dense llama adapter and MLX two-process binary-wire proof are complete. 3. Run **Spike 2 (boundary fence)** at frontier residual widths and keep it as a go/no-go gate. -4. Use Nemotron-H as the first frontier-family follow-up already represented in - `safemlx-lm`; then port Inkling text from the upstream Transformers reference. +4. Use Nemotron-H as the first frontier-family follow-up: wire its existing + public layer/cache structures and affine expert runtime to public-checkpoint + on-load packing. Then expose safemlx's existing Inkling implementation as a + staged text decoder and use Transformers as the parity oracle. --- diff --git a/spikes/mlx-safetensors-stages/FINDINGS.md b/spikes/mlx-safetensors-stages/FINDINGS.md index 1b955b5d29..6669d312d1 100644 --- a/spikes/mlx-safetensors-stages/FINDINGS.md +++ b/spikes/mlx-safetensors-stages/FINDINGS.md @@ -28,9 +28,11 @@ close to exact. For Inkling, tensors are heavily interleaved across source shards, so exact tensor ranges are mandatory: whole-shard selection would turn a 109.84 GiB four-layer stage into a 942.99 GiB download. -The small dense-model path is now proven through execution. The remaining -artifact proof is bounded-memory tensor-at-a-time MLX quantization for a source -tensor too large to retain alongside a whole BF16 stage. +The small dense-model path is now proven through execution, including +tensor-at-a-time affine-4 quantization into the live MLX model. The remaining +artifact proof is direct range-to-quantized-cache materialization with measured +peak RSS and disk use: the current path first retains the complete BF16 stage +slice on disk, then quantizes its tensors one at a time during model load. ## Reproduce @@ -95,10 +97,36 @@ just mlx-safetensors-split-proof \ --wire-dtype f16 ``` -This proves the artifact, MLX layer-range, KV-cache, and existing binary -activation-frame seams on one Mac. It is not yet a two-process/two-node -`mesh-llm serve` implementation; `skippy-server` still binds directly to the -llama.cpp `StageModel` and needs the planned engine abstraction first. +This original harness proved the artifact, MLX layer-range, KV-cache, and +existing binary activation-frame seams on one Mac. Subsequent commits added the +engine abstraction, two real stage processes, and explicit host +`StagePrepare`/`StageLoad` consumption. Automatic topology production and a +remote two-node MLX run remain. + +## Small-model load-time quantization proof + +`MlxStageEngine` now optionally constructs quantized Llama modules and calls +safemlx's strict tensor-streaming loader. For every selected dense tensor, that +loader produces the packed weight/scales/biases, calls `eval`, synchronizes the +quantization stream, installs the result, and then continues. It does not build +a BF16-stage-sized lazy quantization graph. This is not yet a one-source-copy +guarantee: `Array::try_from(TensorView)` and the subsequent stream copy can both +contribute to the physical high-water mark, and mmap pages are outside MLX +allocator counters. + +On the same SmolLM2 split, affine 4-bit with group size 64 produced this +whole-model reference: + +```text +[260, 2240, 314, 253, 1379, 282, 25801, 28] +``` + +The two separately quantized `0..15` and `15..30` processes reproduced all +eight tokens over F16 stage residuals. Each process retained 349 MLX parameters; +post-proof RSS was 87,392 KiB and 87,952 KiB. This is correctness and steady-RSS +evidence, not a peak-memory claim. The next memory gate must sample the cold +load high-water mark and remove the requirement to keep the complete BF16 stage +slice on disk. ## Representative measurements @@ -152,20 +180,33 @@ forward is not registered in upstream `mlx-lm`, logits are not numerically verified, and the vision/audio towers are excluded. The repository contains weights and tokenizer/config files but not the custom model implementation. -Neither upstream `mlx-lm` nor the pinned Rust `safemlx-lm` dependency currently -ships an Inkling family. The authoritative reference implementation is now in +The pinned Rust dependency is commit +`4e53c5ecd7cbd91c0dfd0992a3c731ca2c36e9c7` ("Add Thinking Machines Inkling +support"). Its `safemlx-lm` Inkling family implements the text decoder, dMel +audio and hMLP vision towers, native SafeTensors key transforms, heterogeneous +KV/SConv cache, and ordinary generation; its loader intentionally skips MTP +weights. The authoritative parity oracle remains [Transformers' Inkling model](https://github.com/huggingface/transformers/blob/main/src/transformers/models/inkling/modular_inkling.py). +What is missing is narrower than a family port but still substantial: safemlx's +Inkling model internals are not exposed as a layer-range stage. Its common +`PackedSwiGluExperts` runtime already supports affine/MXFP4 grouped execution, +but Inkling constructs those expert banks with no quantization and its custom +weight-transform loader only installs dense arrays. The high-level loader's +claim that grouped quantized execution is absent is stale; the real gap is +constructor metadata plus transformed rank-3 packing/loading. + ### What an Inkling MLX stage engine must implement -1. Text-decoder parity first: relative-logit attention, local/global masks, - query scaling, sigmoid top-k routing, routed and shared experts, and all four - SConv paths. -2. A stage constructor that creates only `layers[start..end]`, with embeddings +1. Confirm the existing whole-model text decoder against Transformers on a + tractable fixture/reduced config before changing its visibility. +2. Expose a stage constructor that creates only `layers[start..end]`, with embeddings on the first stage and final norm/readout on the last. -3. Stage-local KV plus SConv recurrent state. Inkling cannot be treated as a - plain paged-KV Llama family. -4. An Inkling-specific weight loader and quantization predicate. +3. Split the existing heterogeneous KV plus SConv recurrent cache by stage; + Inkling cannot be treated as a plain paged-KV Llama family. +4. Wire the existing packed affine/MXFP4 expert runtime into Inkling's + constructor and transformed loader/profile. Keep sensitive tensors dense + initially. 5. Logit parity against Transformers at several layer cuts before network work. 6. Vision/audio towers on the first stage after the text chain is certified. 7. MTP as a separate optional capability after ordinary decode is correct. @@ -228,10 +269,13 @@ Inkling layers 30..33: - largest single BF16 source tensor: 18.00 GiB. A whole-stage loader would need BF16 input plus quantized output and fail on a -128 GB node. A tensor-streaming loader can keep the accumulated 32.22 GiB target -plus one source tensor and quantization scratch resident. The cold path should: +128 GB node. A carefully bounded loader targets the accumulated 32.22 GiB +output plus one source unit and quantization scratch, but the current +full-tensor copy path does not yet meet that contract for an 18 GiB Inkling +expert bank. Inkling may need expert/row slabs or a zero-copy managed-array +seam. The cold path should: -1. range-fetch one tensor into a bounded temporary/mmap buffer; +1. range-fetch one bounded tensor or expert slab into a temporary/mmap buffer; 2. create the MLX source array; 3. quantize according to the certified per-tensor profile; 4. evaluate and append the packed tensor/scales/biases to the derived cache; @@ -274,10 +318,12 @@ The measured candidates suggest this order: 1. **Qwen/Llama**: finish the partial loader and two-stage correctness proof. 2. **Nemotron 3 Ultra**: best next frontier proof because `safemlx-lm` already - has a Nemotron-H implementation, although its Mamba/recurrent blocks still - require state-boundary certification. -3. **Inkling text backbone**: high-value new family; use the upstream - Transformers modular implementation as the parity oracle. + has a Nemotron-H implementation with public layer/cache structures and an + affine rank-3 expert runtime. Its public SafeTensors loader still needs the + matching on-load packing path and Mamba/recurrent stage-boundary certification. +3. **Inkling text backbone**: safemlx already has whole-model text/multimodal + execution; expose a partial-stage surface and use Transformers as the parity + oracle rather than porting the family again. 4. **Inkling multimodal + MTP**: add first-stage towers and optional predictor layers after text correctness. 5. **Kimi K2.6, GLM-5.2, DeepSeek V4**: all are viable range-download targets, @@ -300,22 +346,27 @@ source repo + immutable revision + model-family implementation revision + stage range / embedding / readout / modality ownership + per-stage quantization profile -+ activation wire dtype = derived stage cache identity ``` The cache is evictable derived data. The upstream checkpoint remains the source of truth, and a small certified quantization/profile manifest replaces a large -published layer-package repository. +published layer-package repository. Activation wire dtype belongs in topology / +deployment identity and correctness evidence, not in the weight-cache key, +because it does not alter the derived packed weights. ## Next proof -1. Stream BF16 tensor -> MLX affine quant -> partial SafeTensors output one - tensor at a time; prove bounded RSS and disk use. +1. Expose a sequential selected-range callback/temp-artifact seam from + `model-hf`, then replace BF16-stage-on-disk + live quantization with direct + range -> MLX affine -> bounded derived-cache shards. Prove peak RSS and disk + bounds using both OS physical footprint and MLX allocator counters. 2. Measure the MLX eval/readback/codec boundary fence independently at frontier residual widths and prefill sizes. 3. Introduce the engine-neutral stage interface and run the same proof through two real `skippy-server` processes. -4. Repeat the loader proof with Nemotron-H before implementing Inkling. -5. Port Inkling's text decoder to `safemlx-lm`, starting with one layer and - Transformers parity, then stage ranges, then network execution. +4. First quantize one real Nemotron-H BF16 matrix reproducibly, then implement + one complete split-expert bank without accumulating every dense expert. + Prove a small family member before attempting Ultra-scale ranges. +5. Expose the existing safemlx Inkling text decoder as one stage, prove + Transformers parity, then add stage ranges and network execution. From 54de5f45cf1ef404f08dedb525e8eb6f2e435c6c Mon Sep 17 00:00:00 2001 From: Michael Neale <14976+michaelneale@users.noreply.github.com> Date: Fri, 17 Jul 2026 19:58:31 +1000 Subject: [PATCH 15/37] feat(model-hf): stream exact stage tensors sequentially --- .../model-hf/src/safetensors_stage/locking.rs | 64 +++ .../src/safetensors_stage/materialize.rs | 78 +++- crates/model-hf/src/safetensors_stage/mod.rs | 3 + .../src/safetensors_stage/tensor_stream.rs | 392 ++++++++++++++++++ .../model-hf/src/safetensors_stage/types.rs | 30 ++ docs/design/MLX_STAGE_ENGINE_PLAN.md | 19 +- spikes/mlx-safetensors-stages/FINDINGS.md | 26 +- 7 files changed, 598 insertions(+), 14 deletions(-) create mode 100644 crates/model-hf/src/safetensors_stage/tensor_stream.rs diff --git a/crates/model-hf/src/safetensors_stage/locking.rs b/crates/model-hf/src/safetensors_stage/locking.rs index 761b1e62d4..40336af8d7 100644 --- a/crates/model-hf/src/safetensors_stage/locking.rs +++ b/crates/model-hf/src/safetensors_stage/locking.rs @@ -34,6 +34,43 @@ impl Drop for CacheKeyLock { } } +pub(super) struct AdvisoryFileLock { + file: File, +} + +impl AdvisoryFileLock { + pub(super) fn acquire(path: &Path) -> Result { + let file = open_lock_file(path)?; + lock_file(&file).with_context(|| format!("lock {}", path.display()))?; + Ok(Self { file }) + } + + pub(super) fn try_acquire(path: &Path) -> Result> { + let file = open_lock_file(path)?; + if try_lock_file(&file).with_context(|| format!("try lock {}", path.display()))? { + Ok(Some(Self { file })) + } else { + Ok(None) + } + } +} + +impl Drop for AdvisoryFileLock { + fn drop(&mut self) { + unlock_file(&self.file); + } +} + +fn open_lock_file(path: &Path) -> Result { + OpenOptions::new() + .create(true) + .truncate(false) + .read(true) + .write(true) + .open(path) + .with_context(|| format!("open lock file {}", path.display())) +} + #[cfg(unix)] fn lock_file(file: &File) -> Result<()> { use std::os::fd::AsRawFd; @@ -47,11 +84,38 @@ fn lock_file(file: &File) -> Result<()> { } } +#[cfg(unix)] +fn try_lock_file(file: &File) -> Result { + use std::os::fd::AsRawFd; + + // SAFETY: `file` owns a valid descriptor for the duration of this call. + let result = unsafe { libc::flock(file.as_raw_fd(), libc::LOCK_EX | libc::LOCK_NB) }; + if result == 0 { + return Ok(true); + } + let error = std::io::Error::last_os_error(); + if error + .raw_os_error() + .is_some_and(|code| code == libc::EWOULDBLOCK || code == libc::EAGAIN) + { + Ok(false) + } else { + Err(error).context("nonblocking flock failed") + } +} + #[cfg(not(unix))] fn lock_file(_file: &File) -> Result<()> { Ok(()) } +#[cfg(not(unix))] +fn try_lock_file(_file: &File) -> Result { + // No portable advisory lock is available here, so never treat a visit as + // abandoned and risk removing another process's active tensor. + Ok(false) +} + #[cfg(unix)] fn unlock_file(file: &File) { use std::os::fd::AsRawFd; diff --git a/crates/model-hf/src/safetensors_stage/materialize.rs b/crates/model-hf/src/safetensors_stage/materialize.rs index c8802c9d86..bac73433ea 100644 --- a/crates/model-hf/src/safetensors_stage/materialize.rs +++ b/crates/model-hf/src/safetensors_stage/materialize.rs @@ -29,8 +29,8 @@ const MAX_LOCAL_HEADER_BYTES: u64 = 256 * 1024 * 1024; static TEMP_SEQUENCE: AtomicU64 = AtomicU64::new(0); pub struct SafetensorsStageMaterializer { - remote: RemoteSource, - cache_root: PathBuf, + pub(crate) remote: RemoteSource, + pub(crate) cache_root: PathBuf, } impl SafetensorsStageMaterializer { @@ -338,7 +338,7 @@ fn materialization_spans(tensors: &[&SelectedTensor]) -> Vec, ) -> Result<()> { @@ -616,6 +616,78 @@ mod tests { validate_local_safetensors(&repaired.path.join(MODEL_FILE)).unwrap(); } + #[test] + fn visits_selected_tensors_as_one_ephemeral_file_at_a_time() { + let checkpoint = Arc::new(test_checkpoint()); + let requests = Arc::new(Mutex::new(Vec::new())); + let endpoint = start_checkpoint_server(checkpoint, Arc::clone(&requests)); + let cache = tempfile::tempdir().unwrap(); + let cache_root = cache.path().join("cache"); + let materializer = + SafetensorsStageMaterializer::new(cache_root.clone(), Some(&endpoint), None).unwrap(); + let mut visited = Vec::new(); + + let visit = materializer.prepare_tensor_visit(test_request()).unwrap(); + assert_eq!(visit.checkpoint_sha256(), visit.plan().checkpoint_sha256); + assert_eq!( + format!("{:x}", Sha256::digest(visit.config())), + visit.config_sha256() + ); + let metadata_request_count = requests.lock().unwrap().len(); + let report = visit + .visit_tensor_files(|tensor| { + assert!(tensor.path.is_file()); + assert_eq!( + fs::read_dir(tensor.path.parent().unwrap())?.count(), + 1, + "more than one ephemeral source tensor was retained" + ); + validate_local_safetensors(&tensor.path)?; + let mut reader = BufReader::new(File::open(&tensor.path)?); + let mut header_len = [0_u8; 8]; + reader.read_exact(&mut header_len)?; + let header_len = u64::from_le_bytes(header_len); + let mut header = vec![0_u8; usize::try_from(header_len)?]; + reader.read_exact(&mut header)?; + let parsed = parse_header(&header, tensor.source_range.len())?; + assert_eq!(parsed.keys().collect::>(), vec![&tensor.name]); + visited.push(tensor.name.clone()); + Ok(()) + }) + .unwrap(); + + assert_eq!(visited.len(), report.plan.selected_tensor_count); + assert_eq!(report.visited_tensor_count, visited.len()); + assert_eq!(report.source_range_request_count, visited.len()); + assert_eq!( + report.visited_tensor_bytes, + report.plan.selected_tensor_bytes + ); + assert!(report.temporary_file_peak_bytes > 0); + let payload_request_count = requests.lock().unwrap().len() - metadata_request_count; + assert_eq!(payload_request_count, report.source_range_request_count); + assert!(fs::read_dir(cache_root).unwrap().next().is_none()); + } + + #[test] + fn removes_ephemeral_tensor_files_when_the_visitor_fails() { + let checkpoint = Arc::new(test_checkpoint()); + let endpoint = start_checkpoint_server(checkpoint, Arc::new(Mutex::new(Vec::new()))); + let cache = tempfile::tempdir().unwrap(); + let cache_root = cache.path().join("cache"); + let materializer = + SafetensorsStageMaterializer::new(cache_root.clone(), Some(&endpoint), None).unwrap(); + + let error = materializer + .prepare_tensor_visit(test_request()) + .unwrap() + .visit_tensor_files(|_| anyhow::bail!("stop after first tensor")) + .unwrap_err(); + + assert!(error.to_string().contains("stop after first tensor")); + assert!(fs::read_dir(cache_root).unwrap().next().is_none()); + } + #[test] fn serializes_concurrent_materialization_of_the_same_cache_key() { let checkpoint = Arc::new(test_checkpoint()); diff --git a/crates/model-hf/src/safetensors_stage/mod.rs b/crates/model-hf/src/safetensors_stage/mod.rs index afb4ed666b..687d36d39a 100644 --- a/crates/model-hf/src/safetensors_stage/mod.rs +++ b/crates/model-hf/src/safetensors_stage/mod.rs @@ -2,10 +2,13 @@ mod http; mod layout; mod locking; mod materialize; +mod tensor_stream; mod types; pub use materialize::SafetensorsStageMaterializer; +pub use tensor_stream::SafetensorsStageTensorVisit; pub use types::{ ByteRange, SafetensorsShardPlan, SafetensorsSourceShard, SafetensorsStageArtifact, SafetensorsStageManifest, SafetensorsStagePlan, SafetensorsStageRequest, + SafetensorsStageTensorFile, SafetensorsStageTensorVisitReport, }; diff --git a/crates/model-hf/src/safetensors_stage/tensor_stream.rs b/crates/model-hf/src/safetensors_stage/tensor_stream.rs new file mode 100644 index 0000000000..90de225444 --- /dev/null +++ b/crates/model-hf/src/safetensors_stage/tensor_stream.rs @@ -0,0 +1,392 @@ +use std::{ + collections::BTreeMap, + fs::{self, File}, + io::{BufWriter, Write}, + path::{Path, PathBuf}, + sync::atomic::{AtomicU64, Ordering}, +}; + +use anyhow::{Context, Result, ensure}; + +use super::{ + layout, + locking::AdvisoryFileLock, + materialize::{SafetensorsStageMaterializer, ensure_source_identity}, + types::{ + PreparedStage, SafetensorsSourceShard, SafetensorsStagePlan, SafetensorsStageRequest, + SafetensorsStageTensorFile, SafetensorsStageTensorVisitReport, SelectedTensor, + }, +}; + +static TENSOR_VISIT_SEQUENCE: AtomicU64 = AtomicU64::new(0); + +impl SafetensorsStageMaterializer { + /// Prepares a verified sequential visit before any tensor payload is fetched. + pub fn prepare_tensor_visit( + &self, + request: SafetensorsStageRequest, + ) -> Result> { + let request = request.normalized()?; + let prepared = layout::prepare(&self.remote, &request)?; + Ok(SafetensorsStageTensorVisit { + materializer: self, + prepared, + }) + } +} + +/// A verified stage selection ready for sequential tensor-range consumption. +/// +/// The plan, checkpoint identity, and source config are available before +/// `visit_tensor_files` starts fetching tensor payloads, so a consumer can +/// construct its destination model or derived-cache metadata first. +pub struct SafetensorsStageTensorVisit<'a> { + materializer: &'a SafetensorsStageMaterializer, + prepared: PreparedStage, +} + +impl SafetensorsStageTensorVisit<'_> { + pub fn plan(&self) -> &SafetensorsStagePlan { + &self.prepared.plan + } + + pub fn checkpoint_sha256(&self) -> &str { + &self.prepared.checkpoint_sha256 + } + + pub fn config(&self) -> &[u8] { + &self.prepared.config + } + + pub fn config_sha256(&self) -> &str { + &self.prepared.config_sha256 + } + + /// Visits selected tensors as ephemeral one-tensor SafeTensors files. + /// + /// Each file is removed immediately after `visitor` returns. Consumers + /// using mmap or lazy device graphs must evaluate and synchronize all work + /// that reads the file before returning from the callback. This lets a + /// backend quantize or transform exact HTTP ranges sequentially without + /// retaining the complete BF16 stage artifact on disk. + pub fn visit_tensor_files(self, mut visitor: F) -> Result + where + F: FnMut(&SafetensorsStageTensorFile) -> Result<()>, + { + let directory = EphemeralTensorDirectory::create(&self.materializer.cache_root)?; + let source_shards = self + .prepared + .source_shards + .iter() + .map(|shard| (shard.file.as_str(), shard)) + .collect::>(); + let mut tensors = self.prepared.tensors.iter().collect::>(); + tensors.sort_by(|left, right| { + (&left.source_file, left.source_range.start) + .cmp(&(&right.source_file, right.source_range.start)) + }); + + let mut visited_tensor_bytes = 0_u64; + let mut temporary_file_peak_bytes = 0_u64; + for (index, tensor) in tensors.iter().enumerate() { + let source = source_shards + .get(tensor.source_file.as_str()) + .with_context(|| format!("missing source identity for {}", tensor.source_file))?; + let path = directory + .path() + .join(format!("tensor-{index:06}.safetensors")); + let file_bytes = + self.materializer + .write_tensor_file(&path, &self.prepared.plan, tensor, source)?; + temporary_file_peak_bytes = temporary_file_peak_bytes.max(file_bytes); + visitor(&tensor_file(tensor, path.clone(), file_bytes))?; + fs::remove_file(&path) + .with_context(|| format!("remove ephemeral tensor file {}", path.display()))?; + visited_tensor_bytes = visited_tensor_bytes + .checked_add(tensor.source_range.len()) + .context("visited tensor byte count overflow")?; + } + + ensure!( + tensors.len() == self.prepared.plan.selected_tensor_count + && visited_tensor_bytes == self.prepared.plan.selected_tensor_bytes, + "sequential tensor visit did not cover the selected stage" + ); + Ok(SafetensorsStageTensorVisitReport { + plan: self.prepared.plan, + visited_tensor_count: tensors.len(), + visited_tensor_bytes, + source_range_request_count: tensors.len(), + temporary_file_peak_bytes, + }) + } +} + +impl SafetensorsStageMaterializer { + fn write_tensor_file( + &self, + path: &Path, + plan: &super::types::SafetensorsStagePlan, + tensor: &SelectedTensor, + source: &SafetensorsSourceShard, + ) -> Result { + let (header, header_len) = one_tensor_header(tensor)?; + let file = File::create(path).with_context(|| format!("create {}", path.display()))?; + let mut writer = BufWriter::new(file); + writer.write_all(&header_len.to_le_bytes())?; + writer.write_all(&header)?; + + let url = self + .remote + .url(&plan.repo, &plan.revision, &tensor.source_file)?; + let expected_etag = source + .etag + .as_deref() + .context("planned SafeTensors shard has no ETag")?; + let response = self.remote.exact_range_if_range( + url, + tensor.source_range.start..tensor.source_range.end_exclusive, + expected_etag, + )?; + ensure!( + response.total_file_bytes == source.file_bytes, + "SafeTensors shard {} changed size during tensor visit", + tensor.source_file + ); + ensure_source_identity(source, response.etag())?; + let copied = response.copy_to(&mut writer)?; + ensure!( + copied == tensor.source_range.len(), + "ephemeral SafeTensors tensor payload length mismatch" + ); + writer.flush()?; + let file_bytes = 8_u64 + .checked_add(header_len) + .and_then(|bytes| bytes.checked_add(copied)) + .context("ephemeral SafeTensors file length overflow")?; + ensure!( + fs::metadata(path)?.len() == file_bytes, + "ephemeral SafeTensors file length mismatch" + ); + Ok(file_bytes) + } +} + +fn one_tensor_header(tensor: &SelectedTensor) -> Result<(Vec, u64)> { + let tensor_bytes = tensor.source_range.len(); + let mut header = tensor.header.clone(); + header.data_offsets = [0, tensor_bytes]; + let mut header = serde_json::to_vec(&BTreeMap::from([(tensor.name.clone(), header)]))?; + while header.len() % 8 != 0 { + header.push(b' '); + } + let header_len = u64::try_from(header.len()).context("one-tensor header is too large")?; + Ok((header, header_len)) +} + +fn tensor_file( + tensor: &SelectedTensor, + path: PathBuf, + file_bytes: u64, +) -> SafetensorsStageTensorFile { + SafetensorsStageTensorFile { + name: tensor.name.clone(), + dtype: tensor.header.dtype.clone(), + shape: tensor.header.shape.clone(), + source_file: tensor.source_file.clone(), + source_range: tensor.source_range.clone(), + path, + file_bytes, + } +} + +struct EphemeralTensorDirectory { + path: PathBuf, + lock_path: PathBuf, + _lock: AdvisoryFileLock, +} + +impl EphemeralTensorDirectory { + fn create(cache_root: &Path) -> Result { + fs::create_dir_all(cache_root) + .with_context(|| format!("create SafeTensors cache root {}", cache_root.display()))?; + remove_abandoned_tensor_directories(cache_root)?; + for _ in 0..100 { + let sequence = TENSOR_VISIT_SEQUENCE.fetch_add(1, Ordering::Relaxed); + let base = format!(".tensor-visit.{}.{}", std::process::id(), sequence); + let path = cache_root.join(format!("{base}.partial")); + let lock_path = cache_root.join(format!("{base}.lock")); + let lock = AdvisoryFileLock::acquire(&lock_path)?; + match fs::create_dir(&path) { + Ok(()) => { + return Ok(Self { + path, + lock_path, + _lock: lock, + }); + } + Err(error) if error.kind() == std::io::ErrorKind::AlreadyExists => { + drop(lock); + let _ = fs::remove_file(&lock_path); + } + Err(error) => { + drop(lock); + let _ = fs::remove_file(&lock_path); + return Err(error) + .with_context(|| format!("create tensor visit dir {}", path.display())); + } + } + } + anyhow::bail!("could not allocate a unique SafeTensors tensor visit directory") + } + + fn path(&self) -> &Path { + &self.path + } +} + +impl Drop for EphemeralTensorDirectory { + fn drop(&mut self) { + let _ = fs::remove_dir_all(&self.path); + let _ = fs::remove_file(&self.lock_path); + } +} + +#[cfg(unix)] +fn remove_abandoned_tensor_directories(cache_root: &Path) -> Result<()> { + for entry in fs::read_dir(cache_root)? { + let entry = entry?; + let name = entry.file_name(); + let Some(name) = name.to_str() else { + continue; + }; + let Some(base) = name.strip_suffix(".partial") else { + continue; + }; + if !base.starts_with(".tensor-visit.") || !entry.file_type()?.is_dir() { + continue; + } + let lock_path = cache_root.join(format!("{base}.lock")); + let Some(lock) = AdvisoryFileLock::try_acquire(&lock_path)? else { + continue; + }; + fs::remove_dir_all(entry.path())?; + drop(lock); + fs::remove_file(lock_path)?; + } + Ok(()) +} + +#[cfg(not(unix))] +fn remove_abandoned_tensor_directories(_cache_root: &Path) -> Result<()> { + Ok(()) +} + +#[cfg(test)] +mod tests { + use model_artifact::safetensors::TensorHeader; + + use super::*; + + #[test] + fn one_tensor_header_rebases_payload_offsets() { + let tensor = SelectedTensor { + name: "model.layers.2.weight".to_string(), + source_file: "model.safetensors".to_string(), + source_range: super::super::types::ByteRange { + start: 100, + end_exclusive: 108, + }, + header: TensorHeader { + dtype: "F32".to_string(), + shape: vec![2], + data_offsets: [92, 100], + }, + }; + + let (header, _) = one_tensor_header(&tensor).unwrap(); + let parsed: BTreeMap = serde_json::from_slice(&header).unwrap(); + + assert_eq!(parsed[&tensor.name].data_offsets, [0, 8]); + } + + #[test] + #[ignore = "downloads one layer from a pinned Hugging Face SafeTensors checkpoint"] + fn visits_real_smollm2_layer_without_retaining_source_shard() { + let cache = tempfile::tempdir().unwrap(); + let cache_root = cache.path().join("cache"); + let materializer = + SafetensorsStageMaterializer::new(cache_root.clone(), None, None).unwrap(); + let mut visited = Vec::new(); + let request = SafetensorsStageRequest { + repo: "HuggingFaceTB/SmolLM2-135M-Instruct".to_string(), + revision: "12fd25f77366fa6b3b4b768ec3050bf629380bac".to_string(), + layer_start: 14, + layer_end: 15, + include_prefixes: Vec::new(), + }; + + let visit = materializer.prepare_tensor_visit(request).unwrap(); + assert_eq!(visit.checkpoint_sha256(), visit.plan().checkpoint_sha256); + assert_eq!(visit.config_sha256().len(), 64); + assert!(!visit.config().is_empty()); + let report = visit + .visit_tensor_files(|tensor| { + assert!(tensor.name.starts_with("model.layers.14.")); + assert!(tensor.path.is_file()); + assert_eq!(fs::read_dir(tensor.path.parent().unwrap())?.count(), 1); + visited.push((tensor.name.clone(), tensor.source_range.len())); + Ok(()) + }) + .unwrap(); + + assert_eq!(report.visited_tensor_count, visited.len()); + assert_eq!(report.source_range_request_count, visited.len()); + assert_eq!( + report.visited_tensor_bytes, + visited.iter().map(|(_, bytes)| bytes).sum::() + ); + assert!(report.visited_tensor_bytes < report.plan.source_shard_bytes); + assert!(fs::read_dir(cache_root).unwrap().next().is_none()); + eprintln!( + "visited {} tensors / {} bytes; source shard {} bytes; peak temporary file {} bytes", + report.visited_tensor_count, + report.visited_tensor_bytes, + report.plan.source_shard_bytes, + report.temporary_file_peak_bytes + ); + } + + #[cfg(unix)] + #[test] + fn removes_abandoned_tensor_visit_directory() { + let cache = tempfile::tempdir().unwrap(); + let cache_root = cache.path().join("cache"); + let abandoned = cache_root.join(".tensor-visit.999.0.partial"); + fs::create_dir_all(&abandoned).unwrap(); + fs::write(abandoned.join("tensor-000000.safetensors"), b"abandoned").unwrap(); + + let active = EphemeralTensorDirectory::create(&cache_root).unwrap(); + + assert!(!abandoned.exists()); + drop(active); + assert!(fs::read_dir(cache_root).unwrap().next().is_none()); + } + + #[cfg(unix)] + #[test] + fn keeps_concurrent_tensor_visit_directory_locked() { + let cache = tempfile::tempdir().unwrap(); + let cache_root = cache.path().join("cache"); + let first = EphemeralTensorDirectory::create(&cache_root).unwrap(); + let first_path = first.path().to_path_buf(); + + let second = EphemeralTensorDirectory::create(&cache_root).unwrap(); + + assert!(first_path.is_dir()); + drop(second); + assert!(first_path.is_dir()); + drop(first); + assert!(fs::read_dir(cache_root).unwrap().next().is_none()); + } +} diff --git a/crates/model-hf/src/safetensors_stage/types.rs b/crates/model-hf/src/safetensors_stage/types.rs index 9f04801af1..5aab3f08d3 100644 --- a/crates/model-hf/src/safetensors_stage/types.rs +++ b/crates/model-hf/src/safetensors_stage/types.rs @@ -133,6 +133,36 @@ pub struct SafetensorsStageArtifact { pub cache_hit: bool, } +/// One selected tensor materialized as an ephemeral, valid SafeTensors file. +/// +/// The file exists only for the duration of the visitor callback that receives +/// this value. Callers must consume it before returning from that callback. +#[derive(Debug)] +pub struct SafetensorsStageTensorFile { + pub name: String, + pub dtype: String, + pub shape: Vec, + pub source_file: String, + pub source_range: ByteRange, + pub path: PathBuf, + pub file_bytes: u64, +} + +/// Summary of a sequential selected-tensor visit. +#[derive(Clone, Debug)] +pub struct SafetensorsStageTensorVisitReport { + /// The artifact-oriented range plan used to select and verify tensors. + /// Its `range_request_count` describes coalesced materialization spans; + /// `source_range_request_count` is the visitor's actual request count. + pub plan: SafetensorsStagePlan, + pub visited_tensor_count: usize, + pub visited_tensor_bytes: u64, + pub source_range_request_count: usize, + /// Largest ephemeral source file produced by `model-hf` during this visit. + /// This excludes consumer output files, filesystem overhead, and RSS. + pub temporary_file_peak_bytes: u64, +} + #[derive(Clone, Debug)] pub(crate) struct SelectedTensor { pub name: String, diff --git a/docs/design/MLX_STAGE_ENGINE_PLAN.md b/docs/design/MLX_STAGE_ENGINE_PLAN.md index 8aac8d6cf6..f3fc0aa2c8 100644 --- a/docs/design/MLX_STAGE_ENGINE_PLAN.md +++ b/docs/design/MLX_STAGE_ENGINE_PLAN.md @@ -74,6 +74,18 @@ physical copies: TensorView conversion, stream copies, mmap pages, and MLX scratch all require explicit high-water measurement. The host does not select this profile yet, and the BF16 partial stage still exists on disk first. +**Update — exact ranges can now be consumed sequentially without a BF16 stage +artifact.** `model-hf` exposes each selected tensor as an ephemeral, valid +one-tensor SafeTensors file, verifies the pinned source identity, and removes +that file before downloading the next tensor. This is an engine-neutral source +seam, not yet a quantized derived cache: the MLX consumer, bounded output +shards, cache identity, and cold-load high-water measurements remain. On the +pinned SmolLM2 layer-14 proof, it fetched 7,080,192 tensor bytes from a +269,060,552-byte source shard while the largest temporary file was 1,769,584 +bytes; the temporary directory was empty at completion. A prepared visit makes +the verified config and checkpoint identity available before tensor callbacks, +and macOS/Unix advisory locks safely scavenge crash-abandoned visits. + The pinned safemlx revision also already includes whole-model Inkling text, vision, and audio execution. Earlier notes that called for porting Inkling were stale. The remaining Inkling work is a partial-stage API plus wiring its @@ -695,9 +707,10 @@ single-machine parity first, then two Macs over the real network. > `StageWireMessage` boundaries matched unsplit MLX with zero measured dense > logit delta; two real F16-wire processes also matched the whole-model > affine-4 token reference after tensor-wise on-load quantization. Explicit host -> stage control now starts the source-precision path. Direct -> range-to-quantized-cache materialization, profile-bearing topology requests, -> and remote two-node execution remain. +> stage control now starts the source-precision path. A sequential exact-range +> tensor visitor now removes the need to create a BF16 stage artifact in the +> next cache builder. Direct range-to-quantized-cache materialization, +> profile-bearing topology requests, and remote two-node execution remain. **Phase 4 — KV/state codec + verify + trim/checkpoint.** Implement the engine-general cache codec (§5.2), `verify_tokens_frame` for speculative decode, diff --git a/spikes/mlx-safetensors-stages/FINDINGS.md b/spikes/mlx-safetensors-stages/FINDINGS.md index 6669d312d1..fd8155026e 100644 --- a/spikes/mlx-safetensors-stages/FINDINGS.md +++ b/spikes/mlx-safetensors-stages/FINDINGS.md @@ -29,10 +29,20 @@ shards, so exact tensor ranges are mandatory: whole-shard selection would turn a 109.84 GiB four-layer stage into a 942.99 GiB download. The small dense-model path is now proven through execution, including -tensor-at-a-time affine-4 quantization into the live MLX model. The remaining -artifact proof is direct range-to-quantized-cache materialization with measured -peak RSS and disk use: the current path first retains the complete BF16 stage -slice on disk, then quantizes its tensors one at a time during model load. +tensor-at-a-time affine-4 quantization into the live MLX model. `model-hf` now +also exposes a backend-neutral sequential visitor that downloads each selected +range as an ephemeral, valid one-tensor SafeTensors file and deletes it before +fetching the next tensor. The remaining artifact proof is consuming that seam +to build bounded quantized-cache shards with measured peak RSS and disk use; +the serving path still first retains the complete BF16 stage slice on disk. +Its prepared session exposes the verified config, config hash, checkpoint +identity, and range plan before payload callbacks. On macOS/Unix, advisory +locks also scavenge crash-abandoned visits without removing concurrent ones. + +The pinned SmolLM2 layer-14 visitor proof fetched 9 tensors totaling 7,080,192 +bytes from a 269,060,552-byte source shard. Its largest temporary one-tensor +file was 1,769,584 bytes, and the visitor cache directory was empty afterward. +Those figures bound temporary disk use for this fixture, not process memory. ## Reproduce @@ -357,10 +367,10 @@ because it does not alter the derived packed weights. ## Next proof -1. Expose a sequential selected-range callback/temp-artifact seam from - `model-hf`, then replace BF16-stage-on-disk + live quantization with direct - range -> MLX affine -> bounded derived-cache shards. Prove peak RSS and disk - bounds using both OS physical footprint and MLX allocator counters. +1. Consume the sequential selected-range visitor from `model-hf` to replace + BF16-stage-on-disk + live quantization with direct range -> MLX affine -> + bounded derived-cache shards. Prove peak RSS and disk bounds using both OS + physical footprint and MLX allocator counters. 2. Measure the MLX eval/readback/codec boundary fence independently at frontier residual widths and prefill sizes. 3. Introduce the engine-neutral stage interface and run the same proof through From 884395eef41f76e993e5c8b0704a2542459386bc Mon Sep 17 00:00:00 2001 From: Michael Neale <14976+michaelneale@users.noreply.github.com> Date: Fri, 17 Jul 2026 20:32:27 +1000 Subject: [PATCH 16/37] feat(mlx): derive quantized stages from exact ranges --- Cargo.lock | 24 + crates/skippy-engine-mlx/Cargo.toml | 25 +- crates/skippy-engine-mlx/STAGED_EXECUTION.md | 66 +- crates/skippy-engine-mlx/src/bin/mlx-stage.rs | 69 +- crates/skippy-engine-mlx/src/derived.rs | 910 ++++++++++++++++++ crates/skippy-engine-mlx/src/lib.rs | 6 + crates/skippy-engine-mlx/src/stage.rs | 43 +- docs/design/MLX_STAGE_ENGINE_PLAN.md | 30 +- spikes/mlx-safetensors-stages/FINDINGS.md | 56 +- 9 files changed, 1180 insertions(+), 49 deletions(-) create mode 100644 crates/skippy-engine-mlx/src/derived.rs diff --git a/Cargo.lock b/Cargo.lock index a1608a15db..e1b4677fe7 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -856,6 +856,20 @@ name = "bytemuck" version = "1.25.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f65693059b6b9c588b9f62fed1cedbf0a8b805631457ea162d68f0de186f3de5" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.118", +] [[package]] name = "byteorder" @@ -2555,6 +2569,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if 1.0.4", "crunchy", "zerocopy", @@ -7757,15 +7772,24 @@ dependencies = [ "async-stream", "async-trait", "axum", + "bytemuck", "clap", "futures-core", + "half", + "libc", + "memmap2", + "model-hf", "openai-frontend", "safemlx", "safemlx-lm", + "safetensors", + "serde", "serde_json", + "sha2 0.10.9", "skippy-engine", "skippy-protocol", "skippy-server", + "tempfile", "tokenizers", "tokio", "tracing", diff --git a/crates/skippy-engine-mlx/Cargo.toml b/crates/skippy-engine-mlx/Cargo.toml index 5b1feecdb6..425b1c4ba0 100644 --- a/crates/skippy-engine-mlx/Cargo.toml +++ b/crates/skippy-engine-mlx/Cargo.toml @@ -36,12 +36,25 @@ path = "src/bin/mlx-stage.rs" default = [] # Enable the real MLX engine. macOS-only in practice (code is cfg-gated to # target_os = "macos"); the deps still resolve elsewhere but compile to nothing. -mlx = ["dep:safemlx", "dep:safemlx-lm", "dep:tokenizers"] +mlx = [ + "dep:bytemuck", + "dep:half", + "dep:libc", + "dep:memmap2", + "dep:model-hf", + "dep:safetensors", + "dep:safemlx", + "dep:safemlx-lm", + "dep:serde", + "dep:sha2", + "dep:tokenizers", +] [dependencies] # The real mesh-llm OpenAI frontend — this is what proves we serve over the # same surface the shipped binary uses, not a toy. openai-frontend = { path = "../openai-frontend" } +model-hf = { path = "../model-hf", optional = true } skippy-engine = { path = "../skippy-engine" } skippy-protocol = { path = "../skippy-protocol" } skippy-server = { path = "../skippy-server" } @@ -49,9 +62,16 @@ async-trait = "0.1" async-stream = "0.3" anyhow = "1" axum = "0.8" +bytemuck = { version = "1", optional = true } clap = { version = "4", features = ["derive"] } futures-core = "0.3" +half = { version = "2", features = ["bytemuck"], optional = true } +libc = { version = "0.2", optional = true } +memmap2 = { version = "0.9", optional = true } serde_json = "1" +serde = { version = "1", features = ["derive"], optional = true } +sha2 = { version = "0.10", optional = true } +safetensors = { version = "0.8", optional = true } tokio = { version = "1", features = ["full"] } tracing = "0.1" tracing-subscriber = { version = "0.3", features = ["env-filter"] } @@ -70,3 +90,6 @@ safemlx-lm = { git = "https://github.com/jbg/safemlx", rev = "4e53c5e", optional tokenizers = { version = "0.23", default-features = false, features = [ "onig", ], optional = true } + +[dev-dependencies] +tempfile = "3" diff --git a/crates/skippy-engine-mlx/STAGED_EXECUTION.md b/crates/skippy-engine-mlx/STAGED_EXECUTION.md index 461e6423f7..6fde8f3963 100644 --- a/crates/skippy-engine-mlx/STAGED_EXECUTION.md +++ b/crates/skippy-engine-mlx/STAGED_EXECUTION.md @@ -16,7 +16,8 @@ the default automatic mesh launch path: - `skippy-server::llama_engine` proves the existing llama `RuntimeState` can implement the same dense contract, including F16/BF16/F32 residual conversion and checkpoint/restore/trim delegation, without changing the native ABI. -- `MlxStageEngine` loads one materialized partial SafeTensors file, owns +- `MlxStageEngine` loads one materialized partial SafeTensors file or sharded + directory, owns per-session KV caches on a dedicated MLX worker thread, and executes only its configured layer range. - `mlx-stage` starts a stage process or drives a chain as a proof client. @@ -79,8 +80,39 @@ independently loaded 15-layer stages generated the same quantized-model tokens: The two-stage processes retained 349 MLX parameters each and had post-proof RSS of 87,392 KiB and 87,952 KiB, versus roughly 189 MiB each at source precision. This proves deterministic per-stage quantization and quantized stage execution; -it does not yet prove peak RSS, direct range-to-quantized-cache disk bounds, or -host/topology selection of the quantization profile. +it does not by itself prove peak RSS or remove the dense partial artifact. + +The next proof removed that dense partial artifact. `mlx-stage derive` consumes +the sequential exact-range session from `model-hf`, quantizes and synchronizes +one matrix at a time, copies packed results into a bounded host-side output +shard, and deletes each dense source tensor before fetching the next. Pure-Rust +SafeTensors I/O avoids linking MLX's bundled GGUF symbols into the existing +Skippy/llama.cpp binary. + +With 16 MiB output shards, the two SmolLM2 halves produced: + +| Layers | Dense ranges fetched | Quantized artifact | Shards | Largest source temp | MLX peak active | Process max RSS | macOS peak footprint | +| --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | +| `0..15` | 162,825,984 B | 45,887,848 B | 3 | 56,623,208 B | 72,548,352 B | 140,853,248 B | 240,157,440 B | +| `15..30` | 162,827,136 B | 45,889,726 B | 3 | 56,623,208 B | 72,548,352 B | 140,722,176 B | 241,550,080 B | + +Both derived directories loaded without a quantization request because their +config records the affine-4/group-64 encoding. They again produced exactly: + +```text +[260, 2240, 314, 253, 1379, 282, 25801, 28] +``` + +No complete dense stage or source shard was written. The report separates the +checkpoint/plan/quantizer recipe hash from an output-content digest and records +every output shard hash. Repeating the one-layer derivation produced a +byte-identical weight shard; the whole directories intentionally differ because +reports include local paths and runtime memory evidence. The shard-size option +is a soft bundle target, so one packed tensor may exceed it. This is a bounded +`model_type=llama` artifact builder, not yet the host's reusable cache and not +evidence that frontier expert-bank transforms fit the same bound. Artifact byte +counts and the measured source-plus-output working-disk high-water mark exclude +the report, lock files, and filesystem allocation overhead. The two partial files are the exact-range artifacts described in `../../spikes/mlx-safetensors-stages/FINDINGS.md`. Tied input/output embeddings @@ -96,6 +128,27 @@ Build once: just mlx-stage-build ``` +Derive both quantized stage directories directly from immutable source ranges: + +```bash +just mlx-stage derive \ + --repo HuggingFaceTB/SmolLM2-135M-Instruct \ + --revision 12fd25f77366fa6b3b4b768ec3050bf629380bac \ + --layer-start 0 --layer-end 15 \ + --output /tmp/mlx-derived-smol-stage0 \ + --weight-quantization affine4 --shard-size-mib 16 + +just mlx-stage derive \ + --repo HuggingFaceTB/SmolLM2-135M-Instruct \ + --revision 12fd25f77366fa6b3b4b768ec3050bf629380bac \ + --layer-start 15 --layer-end 30 \ + --output /tmp/mlx-derived-smol-stage1 \ + --weight-quantization affine4 --shard-size-mib 16 +``` + +The derived directories are already quantized; do not pass +`--weight-quantization` when serving them. + Start the final stage: ```bash @@ -132,6 +185,10 @@ just mlx-stage prove --connect 127.0.0.1:19090 --wire-dtype f16 - Dense Llama-family checkpoints only in `MlxStageEngine`. The pinned safemlx revision has whole-model Inkling and Nemotron-H implementations, but neither is exposed through this partial-stage adapter yet. +- The derived builder handles ordinary rank-2 Llama weights. Nemotron split + experts need per-layer expert-bank assembly before quantization; Inkling needs + its transformed rank-3 grouped-expert loader. Neither is silently treated as + Llama. - Greedy sampling only; sampling metadata is preserved in the contract and rejected explicitly when enabled. - No KV page import/export, cache trim/checkpoint, MTP, speculative verify, @@ -144,5 +201,6 @@ just mlx-stage prove --connect 127.0.0.1:19090 --wire-dtype f16 placement, capability advertisement, coordinator model planning, and an OpenAI stage-0 frontend remain. Explicit host requests still use source precision; quantization selection currently exists only in the engine config - and `mlx-stage` proof CLI. There is no mesh protocol or Skippy ABI break in + and `mlx-stage` proof/derive CLI. The explicit derived artifact is not yet an + automatic cache hit path. There is no mesh protocol or Skippy ABI break in the explicit consumer path. diff --git a/crates/skippy-engine-mlx/src/bin/mlx-stage.rs b/crates/skippy-engine-mlx/src/bin/mlx-stage.rs index e1025e838c..2e815f070f 100644 --- a/crates/skippy-engine-mlx/src/bin/mlx-stage.rs +++ b/crates/skippy-engine-mlx/src/bin/mlx-stage.rs @@ -11,8 +11,10 @@ mod real { use anyhow::{Context, Result, ensure}; use clap::{Parser, Subcommand, ValueEnum}; + use model_hf::safetensors_stage::{SafetensorsStageMaterializer, SafetensorsStageRequest}; use skippy_engine_mlx::{ - MlxComputeDtype, MlxStageEngine, MlxStageEngineConfig, MlxWeightQuantization, + MlxComputeDtype, MlxDerivedStageConfig, MlxStageEngine, MlxStageEngineConfig, + MlxWeightQuantization, derive_quantized_stage, }; use skippy_protocol::binary::{ StageStateHeader, StageWireMessage, WireActivationDType, WireMessageKind, WireReplyKind, @@ -52,6 +54,27 @@ mod real { #[arg(long, value_enum)] weight_quantization: Option, }, + /// Download exact stage tensors and write bounded MLX-quantized shards. + Derive { + #[arg(long)] + repo: String, + /// Immutable 40-character Hugging Face commit SHA. + #[arg(long)] + revision: String, + #[arg(long)] + layer_start: u32, + #[arg(long)] + layer_end: u32, + #[arg(long = "include-prefix")] + include_prefixes: Vec, + #[arg(long)] + output: PathBuf, + #[arg(long, value_enum, default_value_t = WeightQuantization::Affine4)] + weight_quantization: WeightQuantization, + /// Soft output shard target; one packed tensor may exceed it. + #[arg(long, default_value_t = 256)] + shard_size_mib: usize, + }, /// Drive a stage chain and assert its greedy token sequence. Prove { #[arg(long)] @@ -150,6 +173,27 @@ mod real { wire_dtype: wire_dtype.into(), }, ), + Command::Derive { + repo, + revision, + layer_start, + layer_end, + include_prefixes, + output, + weight_quantization, + shard_size_mib, + } => derive( + SafetensorsStageRequest { + repo, + revision, + layer_start, + layer_end, + include_prefixes, + }, + output, + weight_quantization.into(), + shard_size_mib, + ), Command::Prove { connect, tokens, @@ -164,6 +208,29 @@ mod real { serve_stage_engine(engine, options) } + fn derive( + source: SafetensorsStageRequest, + output_dir: PathBuf, + quantization: MlxWeightQuantization, + shard_size_mib: usize, + ) -> Result<()> { + let shard_size_bytes = shard_size_mib + .checked_mul(1024 * 1024) + .context("derived shard size overflow")?; + let materializer = SafetensorsStageMaterializer::from_environment()?; + let report = derive_quantized_stage( + &materializer, + &MlxDerivedStageConfig { + source, + output_dir, + quantization, + shard_size_bytes, + }, + )?; + println!("{}", serde_json::to_string_pretty(&report)?); + Ok(()) + } + fn prove( connect: SocketAddr, tokens: &str, diff --git a/crates/skippy-engine-mlx/src/derived.rs b/crates/skippy-engine-mlx/src/derived.rs new file mode 100644 index 0000000000..c155f8d8a4 --- /dev/null +++ b/crates/skippy-engine-mlx/src/derived.rs @@ -0,0 +1,910 @@ +//! Bounded exact-range to MLX-quantized stage artifact conversion. + +use std::{ + borrow::Cow, + collections::BTreeMap, + fs::{self, File}, + io::{BufReader, Read}, + path::{Path, PathBuf}, + sync::atomic::{AtomicU64, Ordering}, +}; + +use anyhow::{Context, Result, ensure}; +use half::{bf16, f16}; +use memmap2::MmapOptions; +use model_hf::safetensors_stage::{SafetensorsStageMaterializer, SafetensorsStageRequest}; +use safemlx::{ + Array, Device, DeviceType, Dtype as MlxDtype, Stream, + memory::{active_memory, cache_memory, peak_memory, reset_peak_memory}, + transforms::eval, +}; +use safemlx_lm::quantization::{WeightQuantization, quantize_tensor}; +use safetensors::tensor::{Dtype as SafeDtype, SafeTensors, View, serialize_to_file}; +use serde::Serialize; +use serde_json::{Value, json}; +use sha2::{Digest, Sha256}; + +use crate::stage::MlxWeightQuantization; + +const DERIVED_STAGE_SCHEMA_VERSION: u32 = 1; +const DERIVED_STAGE_IMPLEMENTATION: &str = "mesh-mlx-range-derived-v1"; +const SAFEMLX_REVISION: &str = "4e53c5ecd7cbd91c0dfd0992a3c731ca2c36e9c7"; +const PLAN_FILE: &str = "stage-plan.json"; +const REPORT_FILE: &str = "derived-stage.json"; +static DERIVED_SEQUENCE: AtomicU64 = AtomicU64::new(0); + +/// Configuration for producing one MLX-quantized partial stage. +#[derive(Clone, Debug)] +pub struct MlxDerivedStageConfig { + pub source: SafetensorsStageRequest, + pub output_dir: PathBuf, + pub quantization: MlxWeightQuantization, + /// Soft output bundle target. A single packed tensor may exceed this size. + pub shard_size_bytes: usize, +} + +/// One finalized SafeTensors shard in a derived stage artifact. +#[derive(Clone, Debug, Serialize)] +pub struct MlxDerivedStageShard { + pub file: String, + pub file_bytes: u64, + pub sha256: String, +} + +/// Evidence and identity for a bounded quantized stage derivation. +#[derive(Clone, Debug, Serialize)] +pub struct MlxDerivedStageReport { + pub schema_version: u32, + /// Identity of the source checkpoint, stage plan, quantization, implementation, and sharding. + pub derivation_recipe_sha256: String, + /// Tamper-evident digest of every published artifact file except this report. + pub output_content_sha256: String, + pub checkpoint_sha256: String, + pub plan_sha256: String, + pub repo: String, + pub revision: String, + pub layer_start: u32, + pub layer_end: u32, + pub quantization: Value, + pub quantization_label: String, + pub safemlx_revision: String, + pub output_dir: PathBuf, + pub source_tensor_count: usize, + pub source_tensor_bytes: u64, + pub source_range_request_count: usize, + pub source_temporary_file_peak_bytes: u64, + pub quantized_tensor_count: usize, + pub copied_tensor_count: usize, + pub output_tensor_bytes: u64, + /// Bytes in the published artifact files, excluding this report. + pub artifact_file_bytes: u64, + /// Measured source-tensor plus artifact payload high-water mark. + /// + /// Filesystem allocation, lock files, and this report are intentionally excluded. + pub working_disk_peak_bytes: u64, + pub mlx_active_memory_bytes: usize, + pub mlx_cache_memory_bytes: usize, + pub mlx_peak_memory_bytes: usize, + pub shards: Vec, +} + +struct PendingShard { + tensors: BTreeMap, + bytes: usize, +} + +impl PendingShard { + fn new() -> Self { + Self { + tensors: BTreeMap::new(), + bytes: 0, + } + } + + fn insert(&mut self, name: String, tensor: OwnedTensor) { + self.bytes = self.bytes.saturating_add(tensor.data.len()); + self.tensors.insert(name, tensor); + } +} + +struct OwnedTensor { + dtype: SafeDtype, + shape: Vec, + data: Vec, +} + +impl View for &OwnedTensor { + fn dtype(&self) -> SafeDtype { + self.dtype + } + + fn shape(&self) -> &[usize] { + &self.shape + } + + fn data(&self) -> Cow<'_, [u8]> { + Cow::Borrowed(&self.data) + } + + fn data_len(&self) -> usize { + self.data.len() + } +} + +struct BuildState { + pending: PendingShard, + temporary_shards: Vec, + locations: BTreeMap, + quantized_tensor_count: usize, + copied_tensor_count: usize, + output_tensor_bytes: u64, + written_output_file_bytes: u64, + working_disk_peak_bytes: u64, +} + +impl BuildState { + fn new(initial_output_file_bytes: u64) -> Self { + Self { + pending: PendingShard::new(), + temporary_shards: Vec::new(), + locations: BTreeMap::new(), + quantized_tensor_count: 0, + copied_tensor_count: 0, + output_tensor_bytes: 0, + written_output_file_bytes: initial_output_file_bytes, + working_disk_peak_bytes: initial_output_file_bytes, + } + } + + fn observe_source_file(&mut self, bytes: u64) { + self.working_disk_peak_bytes = self + .working_disk_peak_bytes + .max(self.written_output_file_bytes.saturating_add(bytes)); + } +} + +/// Derives an MLX-compatible quantized partial stage without retaining a dense stage artifact. +pub fn derive_quantized_stage( + materializer: &SafetensorsStageMaterializer, + config: &MlxDerivedStageConfig, +) -> Result { + ensure!( + config.shard_size_bytes > 0, + "derived shard size must be non-zero" + ); + remove_abandoned_outputs_for_destination(&config.output_dir)?; + ensure!( + !config.output_dir.exists(), + "derived output already exists: {}", + config.output_dir.display() + ); + let visit = materializer.prepare_tensor_visit(config.source.clone())?; + let plan = visit.plan().clone(); + let source_config = visit.config().to_vec(); + ensure_dense_source_config(&source_config)?; + let quantization = config.quantization.safemlx()?; + let quantization_value = serde_json::to_value(quantization)?; + let plan_bytes = serde_json::to_vec(&plan)?; + let plan_sha256 = sha256_bytes(&plan_bytes); + let derivation_recipe_sha256 = derived_identity( + &plan, + &plan_sha256, + &quantization_value, + config.shard_size_bytes, + )?; + let temporary = TemporaryOutput::create(&config.output_dir)?; + write_quantized_config( + temporary.path().join("config.json"), + &source_config, + &quantization_value, + )?; + write_json(temporary.path().join(PLAN_FILE), &plan)?; + + reset_peak_memory()?; + let weights_stream = Stream::new_with_device(&Device::new(DeviceType::Cpu, 0)); + let quantization_stream = Stream::new_with_device(&Device::new(DeviceType::Gpu, 0)); + let initial_output_file_bytes = directory_file_bytes(temporary.path())?; + let mut state = BuildState::new(initial_output_file_bytes); + let visit_report = visit.visit_tensor_files(|tensor| { + state.observe_source_file(tensor.file_bytes); + let arrays = convert_tensor( + tensor, + quantization, + &weights_stream, + &quantization_stream, + &mut state, + )?; + append_arrays( + arrays, + config.shard_size_bytes, + temporary.path(), + tensor.file_bytes, + &mut state, + ) + })?; + if !state.pending.tensors.is_empty() { + flush_shard(temporary.path(), &mut state)?; + } + ensure!( + !state.temporary_shards.is_empty(), + "derived stage contains no tensors" + ); + let finalized = finalize_shards(temporary.path(), &state)?; + let shards = finalized + .iter() + .map(|path| derived_shard(path)) + .collect::>>()?; + let output_content_sha256 = output_content_sha256(temporary.path())?; + let artifact_file_bytes = directory_file_bytes(temporary.path())?; + state.working_disk_peak_bytes = state.working_disk_peak_bytes.max(artifact_file_bytes); + let report = MlxDerivedStageReport { + schema_version: DERIVED_STAGE_SCHEMA_VERSION, + derivation_recipe_sha256, + output_content_sha256, + checkpoint_sha256: plan.checkpoint_sha256.clone(), + plan_sha256, + repo: plan.repo.clone(), + revision: plan.revision.clone(), + layer_start: plan.layer_start, + layer_end: plan.layer_end, + quantization: quantization_value, + quantization_label: config.quantization.label(), + safemlx_revision: SAFEMLX_REVISION.to_string(), + output_dir: config.output_dir.clone(), + source_tensor_count: visit_report.visited_tensor_count, + source_tensor_bytes: visit_report.visited_tensor_bytes, + source_range_request_count: visit_report.source_range_request_count, + source_temporary_file_peak_bytes: visit_report.temporary_file_peak_bytes, + quantized_tensor_count: state.quantized_tensor_count, + copied_tensor_count: state.copied_tensor_count, + output_tensor_bytes: state.output_tensor_bytes, + artifact_file_bytes, + working_disk_peak_bytes: state.working_disk_peak_bytes, + mlx_active_memory_bytes: active_memory()?, + mlx_cache_memory_bytes: cache_memory()?, + mlx_peak_memory_bytes: peak_memory()?, + shards, + }; + write_json(temporary.path().join(REPORT_FILE), &report)?; + temporary.publish(&config.output_dir)?; + Ok(report) +} + +fn convert_tensor( + tensor: &model_hf::safetensors_stage::SafetensorsStageTensorFile, + quantization: WeightQuantization, + weights_stream: &Stream, + quantization_stream: &Stream, + state: &mut BuildState, +) -> Result> { + let file = File::open(&tensor.path)?; + // SAFETY: the mapping is read-only and remains alive until all MLX work + // derived from its sole TensorView is evaluated and synchronized below. + let mmap = unsafe { MmapOptions::new().map(&file)? }; + let tensors = SafeTensors::deserialize(&mmap)?; + ensure!( + tensors.len() == 1, + "ephemeral SafeTensors file did not contain exactly {}", + tensor.name + ); + let dense = Array::try_from(tensors.tensor(&tensor.name)?)?.copy(weights_stream)?; + let arrays = if should_quantize_source_weight(&tensor.name, &dense, quantization)? { + state.quantized_tensor_count += 1; + quantize_tensor(&dense, quantization, quantization_stream)? + .into_named_arrays(&tensor.name)? + } else { + state.copied_tensor_count += 1; + vec![(tensor.name.clone(), dense)] + }; + eval(arrays.iter().map(|(_, array)| array))?; + weights_stream.synchronize()?; + quantization_stream.synchronize()?; + arrays + .into_iter() + .map(|(name, array)| Ok((name, owned_tensor(&array)?))) + .collect() +} + +fn should_quantize_source_weight( + name: &str, + tensor: &Array, + quantization: WeightQuantization, +) -> Result { + ensure!( + !name.ends_with(".scales") + && !name.ends_with(".biases") + && !name.ends_with("_scales") + && !name.ends_with("_biases"), + "source checkpoint already contains packed quantization companion {name}" + ); + if !name.ends_with(".weight") || tensor.ndim() < 2 { + return Ok(false); + } + ensure!( + tensor.ndim() == 2, + "derived stage v1 only supports dense rank-2 Llama weights; {name} has rank {}", + tensor.ndim() + ); + ensure!( + tensor.dtype().is_float(), + "source weight {name} is already packed or uses unsupported dtype {:?}", + tensor.dtype() + ); + ensure!( + tensor.dim(1) % quantization.group_size() == 0 && tensor.dim(1) % 32 == 0, + "source weight {name} input dimension {} is incompatible with {}", + tensor.dim(1), + quantization_label(quantization) + ); + Ok(true) +} + +fn quantization_label(quantization: WeightQuantization) -> String { + format!( + "{:?}-{}bit-g{}", + quantization.mode(), + quantization.bits(), + quantization.group_size() + ) +} + +fn owned_tensor(array: &Array) -> Result { + let evaluated = array.evaluated()?; + let (dtype, data) = match array.dtype() { + MlxDtype::Bool => ( + SafeDtype::BOOL, + evaluated + .as_slice::() + .iter() + .map(|value| u8::from(*value)) + .collect(), + ), + MlxDtype::Uint8 => (SafeDtype::U8, evaluated.as_slice::().to_vec()), + MlxDtype::Uint16 => ( + SafeDtype::U16, + bytemuck::cast_slice(evaluated.as_slice::()).to_vec(), + ), + MlxDtype::Uint32 => ( + SafeDtype::U32, + bytemuck::cast_slice(evaluated.as_slice::()).to_vec(), + ), + MlxDtype::Uint64 => ( + SafeDtype::U64, + bytemuck::cast_slice(evaluated.as_slice::()).to_vec(), + ), + MlxDtype::Int8 => ( + SafeDtype::I8, + bytemuck::cast_slice(evaluated.as_slice::()).to_vec(), + ), + MlxDtype::Int16 => ( + SafeDtype::I16, + bytemuck::cast_slice(evaluated.as_slice::()).to_vec(), + ), + MlxDtype::Int32 => ( + SafeDtype::I32, + bytemuck::cast_slice(evaluated.as_slice::()).to_vec(), + ), + MlxDtype::Int64 => ( + SafeDtype::I64, + bytemuck::cast_slice(evaluated.as_slice::()).to_vec(), + ), + MlxDtype::Float16 => ( + SafeDtype::F16, + bytemuck::cast_slice(evaluated.as_slice::()).to_vec(), + ), + MlxDtype::Bfloat16 => ( + SafeDtype::BF16, + bytemuck::cast_slice(evaluated.as_slice::()).to_vec(), + ), + MlxDtype::Float32 => ( + SafeDtype::F32, + bytemuck::cast_slice(evaluated.as_slice::()).to_vec(), + ), + MlxDtype::Float64 => ( + SafeDtype::F64, + bytemuck::cast_slice(evaluated.as_slice::()).to_vec(), + ), + MlxDtype::Complex64 => { + anyhow::bail!("complex MLX tensors cannot be saved as stage weights") + } + }; + let shape = array + .shape() + .iter() + .copied() + .map(usize::try_from) + .collect::, _>>()?; + Ok(OwnedTensor { dtype, shape, data }) +} + +fn append_arrays( + arrays: Vec<(String, OwnedTensor)>, + shard_size_bytes: usize, + output_dir: &Path, + source_file_bytes: u64, + state: &mut BuildState, +) -> Result<()> { + let incoming_bytes = arrays + .iter() + .map(|(_, tensor)| tensor.data.len()) + .sum::(); + if !state.pending.tensors.is_empty() + && state.pending.bytes.saturating_add(incoming_bytes) > shard_size_bytes + { + flush_shard(output_dir, state)?; + state.observe_source_file(source_file_bytes); + } + for (name, tensor) in arrays { + state.output_tensor_bytes = state + .output_tensor_bytes + .checked_add(u64::try_from(tensor.data.len())?) + .context("derived output tensor byte count overflow")?; + state.pending.insert(name, tensor); + } + Ok(()) +} + +fn flush_shard(output_dir: &Path, state: &mut BuildState) -> Result<()> { + let index = state.temporary_shards.len(); + let path = output_dir.join(format!(".derived-{index:05}.safetensors")); + serialize_to_file( + state + .pending + .tensors + .iter() + .map(|(name, tensor)| (name.as_str(), tensor)), + None, + &path, + )?; + state.written_output_file_bytes = state + .written_output_file_bytes + .checked_add(fs::metadata(&path)?.len()) + .context("derived output file byte count overflow")?; + for name in state.pending.tensors.keys() { + state.locations.insert(name.clone(), index); + } + state.pending.tensors.clear(); + state.pending.bytes = 0; + state.temporary_shards.push(path); + Ok(()) +} + +fn finalize_shards(output_dir: &Path, state: &BuildState) -> Result> { + if state.temporary_shards.len() == 1 { + let output = output_dir.join("model.safetensors"); + fs::rename(&state.temporary_shards[0], &output)?; + return Ok(vec![output]); + } + let count = state.temporary_shards.len(); + let mut outputs = Vec::with_capacity(count); + for (index, temporary) in state.temporary_shards.iter().enumerate() { + let output = output_dir.join(format!("model-{:05}-of-{count:05}.safetensors", index + 1)); + fs::rename(temporary, &output)?; + outputs.push(output); + } + let weight_map = state + .locations + .iter() + .map(|(name, index)| { + ( + name.clone(), + Value::String( + outputs[*index] + .file_name() + .expect("derived shard has a file name") + .to_string_lossy() + .into_owned(), + ), + ) + }) + .collect::>(); + write_json( + output_dir.join("model.safetensors.index.json"), + &json!({ + "metadata": { "total_size": state.output_tensor_bytes }, + "weight_map": weight_map, + }), + )?; + Ok(outputs) +} + +fn write_quantized_config(path: PathBuf, source: &[u8], quantization: &Value) -> Result<()> { + let mut config: Value = serde_json::from_slice(source).context("parse source config.json")?; + let object = config + .as_object_mut() + .context("source config.json must contain an object")?; + object.insert("quantization".to_string(), quantization.clone()); + object.insert("quantization_config".to_string(), quantization.clone()); + write_json(path, &config) +} + +fn ensure_dense_source_config(source: &[u8]) -> Result<()> { + let config: Value = serde_json::from_slice(source).context("parse source config.json")?; + let object = config + .as_object() + .context("source config.json must contain an object")?; + ensure!( + object.get("model_type").and_then(Value::as_str) == Some("llama"), + "derived stage v1 only supports model_type=llama" + ); + for key in ["quantization", "quantization_config", "compression_config"] { + ensure!( + object.get(key).is_none_or(Value::is_null), + "source checkpoint declares {key}; implicit dequantization/requantization is unsupported" + ); + } + Ok(()) +} + +fn derived_identity( + plan: &model_hf::safetensors_stage::SafetensorsStagePlan, + plan_sha256: &str, + quantization: &Value, + shard_size_bytes: usize, +) -> Result { + let bytes = serde_json::to_vec(&( + DERIVED_STAGE_IMPLEMENTATION, + DERIVED_STAGE_SCHEMA_VERSION, + SAFEMLX_REVISION, + &plan.checkpoint_sha256, + plan_sha256, + plan.layer_start, + plan.layer_end, + &plan.include_prefixes, + quantization, + shard_size_bytes, + ))?; + Ok(sha256_bytes(&bytes)) +} + +fn derived_shard(path: &Path) -> Result { + Ok(MlxDerivedStageShard { + file: path + .file_name() + .context("derived shard has no file name")? + .to_string_lossy() + .into_owned(), + file_bytes: fs::metadata(path)?.len(), + sha256: sha256_file(path)?, + }) +} + +fn output_content_sha256(path: &Path) -> Result { + let mut files = Vec::new(); + for entry in fs::read_dir(path)? { + let entry = entry?; + if entry.file_type()?.is_file() && entry.file_name() != REPORT_FILE { + files.push(entry); + } + } + files.sort_by_key(|entry| entry.file_name()); + let mut hasher = Sha256::new(); + hasher.update(b"mesh-mlx-derived-output-v1"); + for entry in files { + let name = entry.file_name(); + let name = name.to_string_lossy(); + hasher.update(u64::try_from(name.len())?.to_le_bytes()); + hasher.update(name.as_bytes()); + let mut reader = BufReader::new(File::open(entry.path())?); + let file_bytes = reader.get_ref().metadata()?.len(); + hasher.update(file_bytes.to_le_bytes()); + let mut buffer = vec![0_u8; 1024 * 1024]; + loop { + let read = reader.read(&mut buffer)?; + if read == 0 { + break; + } + hasher.update(&buffer[..read]); + } + } + Ok(format!("{:x}", hasher.finalize())) +} + +fn directory_file_bytes(path: &Path) -> Result { + fs::read_dir(path)?.try_fold(0_u64, |total, entry| { + let entry = entry?; + let bytes = if entry.file_type()?.is_file() { + entry.metadata()?.len() + } else { + 0 + }; + total + .checked_add(bytes) + .context("derived directory byte count overflow") + }) +} + +fn write_json(path: PathBuf, value: &impl Serialize) -> Result<()> { + let mut bytes = serde_json::to_vec_pretty(value)?; + bytes.push(b'\n'); + fs::write(path, bytes)?; + Ok(()) +} + +fn sha256_bytes(bytes: &[u8]) -> String { + format!("{:x}", Sha256::digest(bytes)) +} + +fn sha256_file(path: &Path) -> Result { + let mut reader = BufReader::new(File::open(path)?); + let mut hasher = Sha256::new(); + let mut buffer = vec![0_u8; 1024 * 1024]; + loop { + let read = reader.read(&mut buffer)?; + if read == 0 { + break; + } + hasher.update(&buffer[..read]); + } + Ok(format!("{:x}", hasher.finalize())) +} + +struct TemporaryOutput { + path: Option, + lock_path: PathBuf, + _lock: File, +} + +impl TemporaryOutput { + fn create(destination: &Path) -> Result { + let parent = destination + .parent() + .filter(|path| !path.as_os_str().is_empty()) + .unwrap_or_else(|| Path::new(".")); + fs::create_dir_all(parent)?; + let name = destination + .file_name() + .context("derived output path has no file name")? + .to_string_lossy(); + remove_abandoned_outputs(parent, &name)?; + for _ in 0..100 { + let sequence = DERIVED_SEQUENCE.fetch_add(1, Ordering::Relaxed); + let base = format!(".{name}.{}.{}", std::process::id(), sequence); + let path = parent.join(format!("{base}.partial")); + let lock_path = parent.join(format!("{base}.lock")); + let lock = open_locked(&lock_path, false)?.expect("blocking lock is acquired"); + match fs::create_dir(&path) { + Ok(()) => { + return Ok(Self { + path: Some(path), + lock_path, + _lock: lock, + }); + } + Err(error) if error.kind() == std::io::ErrorKind::AlreadyExists => { + drop(lock); + let _ = fs::remove_file(lock_path); + } + Err(error) => { + drop(lock); + let _ = fs::remove_file(lock_path); + return Err(error).context("create derived stage temporary output"); + } + } + } + anyhow::bail!("could not allocate a unique derived stage temporary output") + } + + fn path(&self) -> &Path { + self.path.as_deref().expect("temporary output is active") + } + + fn publish(mut self, destination: &Path) -> Result<()> { + let path = self.path.as_ref().expect("temporary output is active"); + fs::rename(path, destination).context("publish derived stage output")?; + self.path = None; + let _ = fs::remove_file(&self.lock_path); + Ok(()) + } +} + +fn remove_abandoned_outputs_for_destination(destination: &Path) -> Result<()> { + let parent = destination + .parent() + .filter(|path| !path.as_os_str().is_empty()) + .unwrap_or_else(|| Path::new(".")); + fs::create_dir_all(parent)?; + let name = destination + .file_name() + .context("derived output path has no file name")? + .to_string_lossy(); + remove_abandoned_outputs(parent, &name) +} + +impl Drop for TemporaryOutput { + fn drop(&mut self) { + if let Some(path) = &self.path { + let _ = fs::remove_dir_all(path); + } + let _ = fs::remove_file(&self.lock_path); + } +} + +fn remove_abandoned_outputs(parent: &Path, destination_name: &str) -> Result<()> { + let prefix = format!(".{destination_name}."); + for entry in fs::read_dir(parent)? { + let entry = entry?; + let name = entry.file_name(); + let Some(name) = name.to_str() else { + continue; + }; + let Some(base) = name.strip_suffix(".partial") else { + continue; + }; + if !base.starts_with(&prefix) || !entry.file_type()?.is_dir() { + continue; + } + let lock_path = parent.join(format!("{base}.lock")); + let Some(lock) = open_locked(&lock_path, true)? else { + continue; + }; + fs::remove_dir_all(entry.path())?; + drop(lock); + fs::remove_file(lock_path)?; + } + Ok(()) +} + +fn open_locked(path: &Path, nonblocking: bool) -> Result> { + use std::fs::OpenOptions; + use std::os::fd::AsRawFd; + + let file = OpenOptions::new() + .create(true) + .truncate(false) + .read(true) + .write(true) + .open(path)?; + let operation = libc::LOCK_EX | if nonblocking { libc::LOCK_NB } else { 0 }; + // SAFETY: `file` owns a valid descriptor for the duration of this call. + let result = unsafe { libc::flock(file.as_raw_fd(), operation) }; + if result == 0 { + return Ok(Some(file)); + } + let error = std::io::Error::last_os_error(); + if nonblocking + && error + .raw_os_error() + .is_some_and(|code| code == libc::EWOULDBLOCK || code == libc::EAGAIN) + { + Ok(None) + } else { + Err(error).context("lock derived stage temporary output") + } +} + +#[cfg(test)] +mod tests { + use safemlx_lm::quantization::AffineQuantization; + + use super::*; + + #[test] + fn quantized_config_preserves_source_and_adds_both_metadata_keys() { + let directory = tempfile::tempdir().unwrap(); + let path = directory.path().join("config.json"); + let quantization = serde_json::to_value(WeightQuantization::Affine( + AffineQuantization::new(64, 4).unwrap(), + )) + .unwrap(); + + write_quantized_config(path.clone(), br#"{"model_type":"llama"}"#, &quantization).unwrap(); + + let config: Value = serde_json::from_slice(&fs::read(path).unwrap()).unwrap(); + assert_eq!(config["model_type"], "llama"); + assert_eq!(config["quantization"], quantization); + assert_eq!(config["quantization_config"], quantization); + } + + #[test] + fn rejects_prequantized_source_metadata_and_ineligible_weights() { + let error = ensure_dense_source_config( + br#"{"model_type":"llama","quantization_config":{"quant_method":"fp8"}}"#, + ) + .unwrap_err(); + assert!(error.to_string().contains("implicit dequantization")); + + let error = ensure_dense_source_config(br#"{"model_type":"nemotron_h"}"#).unwrap_err(); + assert!(error.to_string().contains("model_type=llama")); + + let quantization: WeightQuantization = AffineQuantization::new(64, 4).unwrap().into(); + let packed = Array::from_slice(&vec![0_u32; 128], &[2, 64]); + assert!( + should_quantize_source_weight("model.layers.0.q_proj.weight", &packed, quantization) + .unwrap_err() + .to_string() + .contains("already packed") + ); + let incompatible = Array::from_slice(&vec![0_f32; 126], &[2, 63]); + assert!( + should_quantize_source_weight( + "model.layers.0.q_proj.weight", + &incompatible, + quantization + ) + .unwrap_err() + .to_string() + .contains("incompatible") + ); + let expert_bank = Array::from_slice(&vec![0_f32; 256], &[2, 2, 64]); + assert!( + should_quantize_source_weight( + "model.layers.0.experts.weight", + &expert_bank, + quantization + ) + .unwrap_err() + .to_string() + .contains("rank-2 Llama") + ); + } + + #[test] + fn pure_rust_safetensors_output_preserves_bfloat16_and_packed_u32_bytes() { + let directory = tempfile::tempdir().unwrap(); + let path = directory.path().join("model.safetensors"); + let source = [bf16::from_f32(1.5), bf16::from_f32(-2.25)]; + let array = Array::from_slice(&source, &[2]); + let tensor = owned_tensor(&array).unwrap(); + + serialize_to_file([("weight", &tensor)], None, &path).unwrap(); + + let bytes = fs::read(path).unwrap(); + let saved = SafeTensors::deserialize(&bytes).unwrap(); + let saved = saved.tensor("weight").unwrap(); + assert_eq!(saved.dtype(), SafeDtype::BF16); + assert_eq!(saved.shape(), &[2]); + assert_eq!(saved.data(), bytemuck::cast_slice::(&source)); + + let packed_path = directory.path().join("packed.safetensors"); + let packed_source = [0x0123_4567_u32, 0x89ab_cdef_u32]; + let packed_array = Array::from_slice(&packed_source, &[1, 2]); + let packed_tensor = owned_tensor(&packed_array).unwrap(); + serialize_to_file([("weight", &packed_tensor)], None, &packed_path).unwrap(); + + let packed_bytes = fs::read(packed_path).unwrap(); + let packed_saved = SafeTensors::deserialize(&packed_bytes).unwrap(); + let packed_saved = packed_saved.tensor("weight").unwrap(); + assert_eq!(packed_saved.dtype(), SafeDtype::U32); + assert_eq!(packed_saved.shape(), &[1, 2]); + assert_eq!( + packed_saved.data(), + bytemuck::cast_slice::(&packed_source) + ); + } + + #[test] + fn removes_abandoned_output_but_preserves_concurrent_output() { + let directory = tempfile::tempdir().unwrap(); + let destination = directory.path().join("stage"); + let abandoned = directory.path().join(".stage.999.0.partial"); + fs::create_dir(&abandoned).unwrap(); + fs::write(abandoned.join("large.safetensors"), b"stale").unwrap(); + + let first = TemporaryOutput::create(&destination).unwrap(); + let first_path = first.path().to_path_buf(); + let second = TemporaryOutput::create(&destination).unwrap(); + + assert!(!abandoned.exists()); + assert!(first_path.is_dir()); + drop(second); + assert!(first_path.is_dir()); + drop(first); + assert!(fs::read_dir(directory.path()).unwrap().next().is_none()); + } + + #[test] + fn removes_abandoned_output_even_when_destination_exists() { + let directory = tempfile::tempdir().unwrap(); + let destination = directory.path().join("stage"); + fs::create_dir(&destination).unwrap(); + let abandoned = directory.path().join(".stage.999.0.partial"); + fs::create_dir(&abandoned).unwrap(); + fs::write(abandoned.join("large.safetensors"), b"stale").unwrap(); + + remove_abandoned_outputs_for_destination(&destination).unwrap(); + + assert!(destination.is_dir()); + assert!(!abandoned.exists()); + } +} diff --git a/crates/skippy-engine-mlx/src/lib.rs b/crates/skippy-engine-mlx/src/lib.rs index 16ad2e15bf..099b4560e4 100644 --- a/crates/skippy-engine-mlx/src/lib.rs +++ b/crates/skippy-engine-mlx/src/lib.rs @@ -10,6 +10,8 @@ #[cfg(all(feature = "mlx", target_os = "macos"))] mod backend; #[cfg(all(feature = "mlx", target_os = "macos"))] +mod derived; +#[cfg(all(feature = "mlx", target_os = "macos"))] mod engine; #[cfg(all(feature = "mlx", target_os = "macos"))] mod stage; @@ -17,6 +19,10 @@ mod stage; #[cfg(all(feature = "mlx", target_os = "macos"))] pub use backend::MlxBackend; #[cfg(all(feature = "mlx", target_os = "macos"))] +pub use derived::{ + MlxDerivedStageConfig, MlxDerivedStageReport, MlxDerivedStageShard, derive_quantized_stage, +}; +#[cfg(all(feature = "mlx", target_os = "macos"))] pub use engine::{ChatTurn, GenerateRequest, MlxEngine, MlxEngineConfig}; #[cfg(all(feature = "mlx", target_os = "macos"))] pub use stage::{MlxComputeDtype, MlxStageEngine, MlxStageEngineConfig, MlxWeightQuantization}; diff --git a/crates/skippy-engine-mlx/src/stage.rs b/crates/skippy-engine-mlx/src/stage.rs index a6cc6f7130..de007db717 100644 --- a/crates/skippy-engine-mlx/src/stage.rs +++ b/crates/skippy-engine-mlx/src/stage.rs @@ -1,11 +1,6 @@ //! Partial-layer MLX implementation of the engine-neutral Skippy stage contract. -use std::{ - collections::BTreeMap, - path::{Path, PathBuf}, - sync::mpsc, - thread, -}; +use std::{collections::BTreeMap, path::PathBuf, sync::mpsc, thread}; use anyhow::{Context, Result, anyhow, bail, ensure}; use safemlx::module::{Module, ModuleParameters, ModuleParametersExt}; @@ -19,8 +14,8 @@ use safemlx_lm::{ }, quantization::{AffineQuantization, WeightQuantization}, weights::{ - StrictLoadConfig, StrictLoadReport, load_safetensors_quantized_strict, - load_safetensors_strict, + StrictLoadConfig, StrictLoadReport, load_safetensors_dir_quantized_strict, + load_safetensors_dir_strict, }, }; use skippy_engine::{ @@ -53,7 +48,7 @@ pub enum MlxWeightQuantization { } impl MlxWeightQuantization { - fn safemlx(self) -> Result { + pub(crate) fn safemlx(self) -> Result { match self { Self::Affine { group_size, bits } => { Ok(AffineQuantization::new(group_size, bits)?.into()) @@ -62,7 +57,7 @@ impl MlxWeightQuantization { } } - fn label(self) -> String { + pub(crate) fn label(self) -> String { match self { Self::Affine { group_size, bits } => format!("affine-{bits}bit-g{group_size}"), Self::MxFp4 => "mxfp4".to_string(), @@ -203,6 +198,18 @@ fn load_stage(config: MlxStageEngineConfig) -> Result { .weight_quantization .map(MlxWeightQuantization::safemlx) .transpose()?; + let weight_quantization_label = config.weight_quantization.map_or_else( + || { + model_args + .quantization + .or(model_args.quantization_config) + .map_or_else( + || "none".to_string(), + |value| format!("checkpoint-{value:?}"), + ) + }, + MlxWeightQuantization::label, + ); if let Some(quantization) = quantization { ensure!( model_args.quantization.is_none() && model_args.quantization_config.is_none(), @@ -214,18 +221,18 @@ fn load_stage(config: MlxStageEngineConfig) -> Result { let load_config = partial_stage_load_config(&info); let mut load_report = StrictLoadReport::default(); match quantization { - Some(quantization) => load_safetensors_quantized_strict( + Some(quantization) => load_safetensors_dir_quantized_strict( &mut model, - weight_file(&config.model_dir), + &config.model_dir, &weights_stream, &stream, quantization, &load_config, &mut load_report, )?, - None => load_safetensors_strict( + None => load_safetensors_dir_strict( &mut model, - weight_file(&config.model_dir), + &config.model_dir, &weights_stream, &load_config, &mut load_report, @@ -242,9 +249,7 @@ fn load_stage(config: MlxStageEngineConfig) -> Result { info.layer_start, info.layer_end, model.parameters().flatten().len(), - config - .weight_quantization - .map_or_else(|| "none".to_string(), MlxWeightQuantization::label), + weight_quantization_label, ); Ok(LoadedStage { model, @@ -274,10 +279,6 @@ fn partial_stage_load_config(info: &StageEngineInfo) -> StrictLoadConfig { config } -fn weight_file(model_dir: &Path) -> PathBuf { - model_dir.join("model.safetensors") -} - fn retain_local_layers(model: &mut llama::Model, start: u32, end: u32) -> Result<()> { let start = usize::try_from(start)?; let end = usize::try_from(end)?; diff --git a/docs/design/MLX_STAGE_ENGINE_PLAN.md b/docs/design/MLX_STAGE_ENGINE_PLAN.md index f3fc0aa2c8..154af31b61 100644 --- a/docs/design/MLX_STAGE_ENGINE_PLAN.md +++ b/docs/design/MLX_STAGE_ENGINE_PLAN.md @@ -32,10 +32,10 @@ Inkling BF16, four layers contain 109.84 GiB of tensors scattered across 942.99 GiB of shard files; exact ranges avoid 833.15 GiB. A SmolLM2-135M proof then materialized two partial files (layers 0..15 and 15..30), loaded each directly into MLX, and matched unsplit logits exactly for prefill plus eight decode steps -through Skippy's real F16 and F32 binary activation codec. The remaining -artifact gate is direct range-to-derived-cache quantization for frontier-sized -source tensors; tensor-at-a-time quantization into a live partial model is now -proven on SmolLM2. +through Skippy's real F16 and F32 binary activation codec. Direct range-to- +quantized artifacts is now also proven for dense Llama below. The remaining +frontier artifact gate is bounded family-specific expert transformation and a +managed cache rather than the basic SafeTensors range mechanism. See `spikes/mlx-safetensors-stages/FINDINGS.md`. **Update — the first engine-neutral, multi-process stage chain is now proven.** @@ -77,15 +77,29 @@ this profile yet, and the BF16 partial stage still exists on disk first. **Update — exact ranges can now be consumed sequentially without a BF16 stage artifact.** `model-hf` exposes each selected tensor as an ephemeral, valid one-tensor SafeTensors file, verifies the pinned source identity, and removes -that file before downloading the next tensor. This is an engine-neutral source -seam, not yet a quantized derived cache: the MLX consumer, bounded output -shards, cache identity, and cold-load high-water measurements remain. On the -pinned SmolLM2 layer-14 proof, it fetched 7,080,192 tensor bytes from a +that file before downloading the next tensor. This engine-neutral source seam +is consumed by the Llama builder below; managed cache hits and frontier-family +transforms remain. On the pinned SmolLM2 layer-14 proof, it fetched 7,080,192 +tensor bytes from a 269,060,552-byte source shard while the largest temporary file was 1,769,584 bytes; the temporary directory was empty at completion. A prepared visit makes the verified config and checkpoint identity available before tensor callbacks, and macOS/Unix advisory locks safely scavenge crash-abandoned visits. +**Update — direct exact-range to bounded affine stage artifacts is proven for +Llama.** `mlx-stage derive` now consumes that prepared visit, quantizes/evaluates +one rank-2 matrix at a time, serializes packed results into bounded SafeTensors +shards, and records the checkpoint/plan/quantizer identity and output hashes. +Two direct-derived SmolLM2 halves were about 45.89 MB each in three shards, +versus about 162.83 MB of dense ranges fetched per half. Largest source temp was +56.62 MB, MLX peak-active memory was 72.55 MB, max RSS was about 140.8 MB, and +macOS peak footprint was about 241 MB. The two derived stages reproduced the +established affine-4 tokens exactly. The v1 builder deliberately requires +`model_type=llama` and rejects pre-quantized sources, rank-3 weights, and +incompatible matrix dimensions. This closes the dense-Llama artifact-builder +proof, not reusable host cache hits, Nemotron expert packing, or Inkling's +transformed rank-3 path. + The pinned safemlx revision also already includes whole-model Inkling text, vision, and audio execution. Earlier notes that called for porting Inkling were stale. The remaining Inkling work is a partial-stage API plus wiring its diff --git a/spikes/mlx-safetensors-stages/FINDINGS.md b/spikes/mlx-safetensors-stages/FINDINGS.md index fd8155026e..1bda31a97c 100644 --- a/spikes/mlx-safetensors-stages/FINDINGS.md +++ b/spikes/mlx-safetensors-stages/FINDINGS.md @@ -134,9 +134,40 @@ whole-model reference: The two separately quantized `0..15` and `15..30` processes reproduced all eight tokens over F16 stage residuals. Each process retained 349 MLX parameters; post-proof RSS was 87,392 KiB and 87,952 KiB. This is correctness and steady-RSS -evidence, not a peak-memory claim. The next memory gate must sample the cold -load high-water mark and remove the requirement to keep the complete BF16 stage -slice on disk. +evidence, not a peak-memory claim. + +## Direct range-to-quantized artifact proof + +`mlx-stage derive` now removes the complete BF16 stage slice from the workflow. +It consumes one verified range file at a time, quantizes eligible rank-2 Llama +weights, evaluates and synchronizes MLX work, copies the packed arrays into a +bounded host-side shard, and returns before `model-hf` deletes the dense source +file. Output uses pure-Rust SafeTensors serialization to coexist with the +Skippy/llama.cpp native link. + +For the two 15-layer SmolLM2 halves at affine-4/group-64 and 16 MiB output +shards: + +| Layers | Dense source payload | Derived artifact | Largest source temp | MLX peak active | Max RSS | macOS peak footprint | +| --- | ---: | ---: | ---: | ---: | ---: | ---: | +| `0..15` | 162,825,984 B | 45,887,848 B | 56,623,208 B | 72,548,352 B | 140,853,248 B | 240,157,440 B | +| `15..30` | 162,827,136 B | 45,889,726 B | 56,623,208 B | 72,548,352 B | 140,722,176 B | 241,550,080 B | + +Each artifact contained three shards and an index. Both loaded directly as +pre-quantized partial stages and reproduced the same eight-token affine-4 +reference over F16 residuals. Repeating a one-layer derivation produced a +byte-identical weight shard. The report separates its input recipe hash from an +output-content digest and per-shard hashes; runtime memory evidence is +deliberately not part of either. Whole directories are not byte-identical +because reports include the local path and measurements. Shard size is a soft +bundle target and a single packed tensor may exceed it. The artifact byte count +excludes the report. The working-disk metric is a measured source-tensor plus +artifact-payload high-water mark, not allocated filesystem blocks or lock/report +overhead. The v1 builder fails closed unless `model_type` is exactly `llama`. + +This is still an explicit artifact output, not the host's evictable cache. The +next cache step must map the recorded identity to a managed destination, +validate shard hashes on hit, and avoid all tensor-range requests on reuse. ## Representative measurements @@ -367,16 +398,13 @@ because it does not alter the derived packed weights. ## Next proof -1. Consume the sequential selected-range visitor from `model-hf` to replace - BF16-stage-on-disk + live quantization with direct range -> MLX affine -> - bounded derived-cache shards. Prove peak RSS and disk bounds using both OS - physical footprint and MLX allocator counters. -2. Measure the MLX eval/readback/codec boundary fence independently at frontier +1. Use the derivation recipe hash as a managed host cache key; validate the + output-content digest and shard hashes on hit, then prove a warm load makes + no range requests. +2. Quantize one real Nemotron-H BF16 matrix reproducibly, then implement one + complete split-expert bank without accumulating every dense expert. Prove a + small family member before attempting Ultra-scale ranges. +3. Measure the MLX eval/readback/codec boundary fence independently at frontier residual widths and prefill sizes. -3. Introduce the engine-neutral stage interface and run the same proof through - two real `skippy-server` processes. -4. First quantize one real Nemotron-H BF16 matrix reproducibly, then implement - one complete split-expert bank without accumulating every dense expert. - Prove a small family member before attempting Ultra-scale ranges. -5. Expose the existing safemlx Inkling text decoder as one stage, prove +4. Expose the existing safemlx Inkling text decoder as one stage, prove Transformers parity, then add stage ranges and network execution. From e71a926e536f6e576e494102fba13059df54ff84 Mon Sep 17 00:00:00 2001 From: Michael Neale <14976+michaelneale@users.noreply.github.com> Date: Fri, 17 Jul 2026 20:40:42 +1000 Subject: [PATCH 17/37] feat(mlx): cache derived stages by recipe --- crates/skippy-engine-mlx/STAGED_EXECUTION.md | 23 +- crates/skippy-engine-mlx/src/bin/mlx-stage.rs | 73 ++++- crates/skippy-engine-mlx/src/derived.rs | 54 +++- crates/skippy-engine-mlx/src/derived/cache.rs | 259 ++++++++++++++++++ crates/skippy-engine-mlx/src/lib.rs | 4 +- docs/design/MLX_STAGE_ENGINE_PLAN.md | 11 +- spikes/mlx-safetensors-stages/FINDINGS.md | 17 +- 7 files changed, 416 insertions(+), 25 deletions(-) create mode 100644 crates/skippy-engine-mlx/src/derived/cache.rs diff --git a/crates/skippy-engine-mlx/STAGED_EXECUTION.md b/crates/skippy-engine-mlx/STAGED_EXECUTION.md index 6fde8f3963..adc4cbfcfc 100644 --- a/crates/skippy-engine-mlx/STAGED_EXECUTION.md +++ b/crates/skippy-engine-mlx/STAGED_EXECUTION.md @@ -109,10 +109,20 @@ every output shard hash. Repeating the one-layer derivation produced a byte-identical weight shard; the whole directories intentionally differ because reports include local paths and runtime memory evidence. The shard-size option is a soft bundle target, so one packed tensor may exceed it. This is a bounded -`model_type=llama` artifact builder, not yet the host's reusable cache and not -evidence that frontier expert-bank transforms fit the same bound. Artifact byte -counts and the measured source-plus-output working-disk high-water mark exclude -the report, lock files, and filesystem allocation overhead. +`model_type=llama` artifact builder, not evidence that frontier expert-bank +transforms fit the same bound. Artifact byte counts and the measured +source-plus-output working-disk high-water mark exclude the report, lock files, +and filesystem allocation overhead. + +`mlx-stage derive-cached` then proved the reusable cache seam. It maps the +strong recipe identity to a locked managed directory and validates schema, +recipe, aggregate artifact bytes, output-content digest, and every shard hash +before accepting a hit. On the same pinned layer-14 slice, the cold call made 9 +tensor-payload range requests; the warm call returned the identical recipe and +content hashes with `cache_hit=true`, made 0 tensor-payload range requests, and +used 17,809,408 B max RSS. It still re-plans lightweight config/index/header +metadata to reconstruct the strong recipe key. Host lifecycle/eviction wiring +is not yet connected to this library/CLI cache. The two partial files are the exact-range artifacts described in `../../spikes/mlx-safetensors-stages/FINDINGS.md`. Tied input/output embeddings @@ -146,6 +156,11 @@ just mlx-stage derive \ --weight-quantization affine4 --shard-size-mib 16 ``` +To use the identity-bound cache instead of an explicit output path, replace +`derive` with `derive-cached`, omit `--output`, and optionally pass +`--cache-root`. Repeating the command reports `cache_hit=true` and +`source_range_request_count=0`. + The derived directories are already quantized; do not pass `--weight-quantization` when serving them. diff --git a/crates/skippy-engine-mlx/src/bin/mlx-stage.rs b/crates/skippy-engine-mlx/src/bin/mlx-stage.rs index 2e815f070f..a88b0f4102 100644 --- a/crates/skippy-engine-mlx/src/bin/mlx-stage.rs +++ b/crates/skippy-engine-mlx/src/bin/mlx-stage.rs @@ -13,8 +13,9 @@ mod real { use clap::{Parser, Subcommand, ValueEnum}; use model_hf::safetensors_stage::{SafetensorsStageMaterializer, SafetensorsStageRequest}; use skippy_engine_mlx::{ - MlxComputeDtype, MlxDerivedStageConfig, MlxStageEngine, MlxStageEngineConfig, - MlxWeightQuantization, derive_quantized_stage, + MlxComputeDtype, MlxDerivedStageCacheConfig, MlxDerivedStageConfig, MlxStageEngine, + MlxStageEngineConfig, MlxWeightQuantization, derive_quantized_stage, + derive_quantized_stage_cached, }; use skippy_protocol::binary::{ StageStateHeader, StageWireMessage, WireActivationDType, WireMessageKind, WireReplyKind, @@ -75,6 +76,28 @@ mod real { #[arg(long, default_value_t = 256)] shard_size_mib: usize, }, + /// Reuse or build an identity-bound quantized stage cache entry. + DeriveCached { + #[arg(long)] + repo: String, + /// Immutable 40-character Hugging Face commit SHA. + #[arg(long)] + revision: String, + #[arg(long)] + layer_start: u32, + #[arg(long)] + layer_end: u32, + #[arg(long = "include-prefix")] + include_prefixes: Vec, + /// Defaults to the mesh-llm cache directory. + #[arg(long)] + cache_root: Option, + #[arg(long, value_enum, default_value_t = WeightQuantization::Affine4)] + weight_quantization: WeightQuantization, + /// Soft output shard target; one packed tensor may exceed it. + #[arg(long, default_value_t = 256)] + shard_size_mib: usize, + }, /// Drive a stage chain and assert its greedy token sequence. Prove { #[arg(long)] @@ -194,6 +217,27 @@ mod real { weight_quantization.into(), shard_size_mib, ), + Command::DeriveCached { + repo, + revision, + layer_start, + layer_end, + include_prefixes, + cache_root, + weight_quantization, + shard_size_mib, + } => derive_cached( + SafetensorsStageRequest { + repo, + revision, + layer_start, + layer_end, + include_prefixes, + }, + cache_root, + weight_quantization.into(), + shard_size_mib, + ), Command::Prove { connect, tokens, @@ -231,6 +275,31 @@ mod real { Ok(()) } + fn derive_cached( + source: SafetensorsStageRequest, + cache_root: Option, + quantization: MlxWeightQuantization, + shard_size_mib: usize, + ) -> Result<()> { + let shard_size_bytes = shard_size_mib + .checked_mul(1024 * 1024) + .context("derived shard size overflow")?; + let cache_root = cache_root + .unwrap_or_else(|| model_hf::store::mesh_llm_cache_dir().join("mlx-derived-stages")); + let materializer = SafetensorsStageMaterializer::from_environment()?; + let result = derive_quantized_stage_cached( + &materializer, + &MlxDerivedStageCacheConfig { + source, + cache_root, + quantization, + shard_size_bytes, + }, + )?; + println!("{}", serde_json::to_string_pretty(&result)?); + Ok(()) + } + fn prove( connect: SocketAddr, tokens: &str, diff --git a/crates/skippy-engine-mlx/src/derived.rs b/crates/skippy-engine-mlx/src/derived.rs index c155f8d8a4..9e69185ca7 100644 --- a/crates/skippy-engine-mlx/src/derived.rs +++ b/crates/skippy-engine-mlx/src/derived.rs @@ -20,17 +20,23 @@ use safemlx::{ }; use safemlx_lm::quantization::{WeightQuantization, quantize_tensor}; use safetensors::tensor::{Dtype as SafeDtype, SafeTensors, View, serialize_to_file}; -use serde::Serialize; +use serde::{Deserialize, Serialize}; use serde_json::{Value, json}; use sha2::{Digest, Sha256}; use crate::stage::MlxWeightQuantization; -const DERIVED_STAGE_SCHEMA_VERSION: u32 = 1; +mod cache; + +pub use cache::{ + MlxDerivedStageCacheConfig, MlxDerivedStageCacheResult, derive_quantized_stage_cached, +}; + +pub(super) const DERIVED_STAGE_SCHEMA_VERSION: u32 = 1; const DERIVED_STAGE_IMPLEMENTATION: &str = "mesh-mlx-range-derived-v1"; const SAFEMLX_REVISION: &str = "4e53c5ecd7cbd91c0dfd0992a3c731ca2c36e9c7"; const PLAN_FILE: &str = "stage-plan.json"; -const REPORT_FILE: &str = "derived-stage.json"; +pub(super) const REPORT_FILE: &str = "derived-stage.json"; static DERIVED_SEQUENCE: AtomicU64 = AtomicU64::new(0); /// Configuration for producing one MLX-quantized partial stage. @@ -44,7 +50,7 @@ pub struct MlxDerivedStageConfig { } /// One finalized SafeTensors shard in a derived stage artifact. -#[derive(Clone, Debug, Serialize)] +#[derive(Clone, Debug, Deserialize, Serialize)] pub struct MlxDerivedStageShard { pub file: String, pub file_bytes: u64, @@ -52,12 +58,12 @@ pub struct MlxDerivedStageShard { } /// Evidence and identity for a bounded quantized stage derivation. -#[derive(Clone, Debug, Serialize)] +#[derive(Clone, Debug, Deserialize, Serialize)] pub struct MlxDerivedStageReport { pub schema_version: u32, /// Identity of the source checkpoint, stage plan, quantization, implementation, and sharding. pub derivation_recipe_sha256: String, - /// Tamper-evident digest of every published artifact file except this report. + /// Corruption-detecting digest of every published artifact file except this report. pub output_content_sha256: String, pub checkpoint_sha256: String, pub plan_sha256: String, @@ -235,7 +241,7 @@ pub fn derive_quantized_stage( .map(|path| derived_shard(path)) .collect::>>()?; let output_content_sha256 = output_content_sha256(temporary.path())?; - let artifact_file_bytes = directory_file_bytes(temporary.path())?; + let artifact_file_bytes = artifact_file_bytes(temporary.path())?; state.working_disk_peak_bytes = state.working_disk_peak_bytes.max(artifact_file_bytes); let report = MlxDerivedStageReport { schema_version: DERIVED_STAGE_SCHEMA_VERSION, @@ -536,6 +542,20 @@ fn ensure_dense_source_config(source: &[u8]) -> Result<()> { Ok(()) } +pub(super) fn prepare_derivation_recipe( + materializer: &SafetensorsStageMaterializer, + source: SafetensorsStageRequest, + quantization: MlxWeightQuantization, + shard_size_bytes: usize, +) -> Result { + let visit = materializer.prepare_tensor_visit(source)?; + ensure_dense_source_config(visit.config())?; + let quantization = serde_json::to_value(quantization.safemlx()?)?; + let plan_bytes = serde_json::to_vec(visit.plan())?; + let plan_sha256 = sha256_bytes(&plan_bytes); + derived_identity(visit.plan(), &plan_sha256, &quantization, shard_size_bytes) +} + fn derived_identity( plan: &model_hf::safetensors_stage::SafetensorsStagePlan, plan_sha256: &str, @@ -569,7 +589,7 @@ fn derived_shard(path: &Path) -> Result { }) } -fn output_content_sha256(path: &Path) -> Result { +pub(super) fn output_content_sha256(path: &Path) -> Result { let mut files = Vec::new(); for entry in fs::read_dir(path)? { let entry = entry?; @@ -614,6 +634,20 @@ fn directory_file_bytes(path: &Path) -> Result { }) } +pub(super) fn artifact_file_bytes(path: &Path) -> Result { + fs::read_dir(path)?.try_fold(0_u64, |total, entry| { + let entry = entry?; + let bytes = if entry.file_type()?.is_file() && entry.file_name() != REPORT_FILE { + entry.metadata()?.len() + } else { + 0 + }; + total + .checked_add(bytes) + .context("derived artifact byte count overflow") + }) +} + fn write_json(path: PathBuf, value: &impl Serialize) -> Result<()> { let mut bytes = serde_json::to_vec_pretty(value)?; bytes.push(b'\n'); @@ -625,7 +659,7 @@ fn sha256_bytes(bytes: &[u8]) -> String { format!("{:x}", Sha256::digest(bytes)) } -fn sha256_file(path: &Path) -> Result { +pub(super) fn sha256_file(path: &Path) -> Result { let mut reader = BufReader::new(File::open(path)?); let mut hasher = Sha256::new(); let mut buffer = vec![0_u8; 1024 * 1024]; @@ -745,7 +779,7 @@ fn remove_abandoned_outputs(parent: &Path, destination_name: &str) -> Result<()> Ok(()) } -fn open_locked(path: &Path, nonblocking: bool) -> Result> { +pub(super) fn open_locked(path: &Path, nonblocking: bool) -> Result> { use std::fs::OpenOptions; use std::os::fd::AsRawFd; diff --git a/crates/skippy-engine-mlx/src/derived/cache.rs b/crates/skippy-engine-mlx/src/derived/cache.rs new file mode 100644 index 0000000000..c2fa4ceb80 --- /dev/null +++ b/crates/skippy-engine-mlx/src/derived/cache.rs @@ -0,0 +1,259 @@ +//! Identity-bound reusable cache for derived MLX stage artifacts. + +use std::{ + fs, + path::{Component, Path, PathBuf}, +}; + +use anyhow::{Context, Result, ensure}; +use model_hf::safetensors_stage::{SafetensorsStageMaterializer, SafetensorsStageRequest}; +use serde::Serialize; + +use super::{ + DERIVED_STAGE_SCHEMA_VERSION, MlxDerivedStageConfig, MlxDerivedStageReport, REPORT_FILE, + artifact_file_bytes, derive_quantized_stage, open_locked, output_content_sha256, + prepare_derivation_recipe, sha256_file, +}; +use crate::stage::MlxWeightQuantization; + +/// Configuration for an identity-bound, reusable derived-stage cache entry. +#[derive(Clone, Debug)] +pub struct MlxDerivedStageCacheConfig { + pub source: SafetensorsStageRequest, + pub cache_root: PathBuf, + pub quantization: MlxWeightQuantization, + /// Soft output bundle target. A single packed tensor may exceed this size. + pub shard_size_bytes: usize, +} + +/// Result of a managed derived-stage lookup or build. +#[derive(Clone, Debug, Serialize)] +pub struct MlxDerivedStageCacheResult { + pub cache_hit: bool, + /// Tensor payload range requests made by this invocation. + pub source_range_request_count: usize, + pub output_dir: PathBuf, + pub report: MlxDerivedStageReport, +} + +/// Loads a verified derived stage from cache or builds it from exact tensor ranges. +pub fn derive_quantized_stage_cached( + materializer: &SafetensorsStageMaterializer, + config: &MlxDerivedStageCacheConfig, +) -> Result { + ensure!( + config.shard_size_bytes > 0, + "derived shard size must be non-zero" + ); + let recipe = prepare_derivation_recipe( + materializer, + config.source.clone(), + config.quantization, + config.shard_size_bytes, + )?; + fs::create_dir_all(&config.cache_root).with_context(|| { + format!( + "create MLX derived stage cache {}", + config.cache_root.display() + ) + })?; + let output_dir = config.cache_root.join(&recipe); + let lock_path = config.cache_root.join(format!(".{recipe}.lock")); + // Keep this pathname stable across invocations. Removing an advisory-lock + // file after unlock can split waiters between the unlinked and new inodes. + let _lock = open_locked(&lock_path, false)?.expect("blocking cache lock is acquired"); + + if let Some(report) = load_cached(&output_dir, &recipe)? { + return Ok(MlxDerivedStageCacheResult { + cache_hit: true, + source_range_request_count: 0, + output_dir, + report, + }); + } + remove_invalid_cache_entry(&output_dir)?; + let report = derive_quantized_stage( + materializer, + &MlxDerivedStageConfig { + source: config.source.clone(), + output_dir: output_dir.clone(), + quantization: config.quantization, + shard_size_bytes: config.shard_size_bytes, + }, + )?; + Ok(MlxDerivedStageCacheResult { + cache_hit: false, + source_range_request_count: report.source_range_request_count, + output_dir, + report, + }) +} + +fn load_cached(output_dir: &Path, recipe: &str) -> Result> { + if !output_dir.is_dir() { + return Ok(None); + } + if !contains_only_regular_files(output_dir)? { + return Ok(None); + } + let bytes = match fs::read(output_dir.join(REPORT_FILE)) { + Ok(bytes) => bytes, + Err(_) => return Ok(None), + }; + let mut report = match serde_json::from_slice::(&bytes) { + Ok(report) => report, + Err(_) => return Ok(None), + }; + if report.schema_version != DERIVED_STAGE_SCHEMA_VERSION + || report.derivation_recipe_sha256 != recipe + || report.artifact_file_bytes != artifact_file_bytes(output_dir)? + || report.output_content_sha256 != output_content_sha256(output_dir)? + || !shards_match(output_dir, &report)? + { + return Ok(None); + } + // `output_dir` is diagnostic, not artifact identity. Refresh it so a + // relocated cache remains reusable without exposing its stale build path. + report.output_dir = output_dir.to_path_buf(); + Ok(Some(report)) +} + +fn contains_only_regular_files(output_dir: &Path) -> Result { + for entry in fs::read_dir(output_dir)? { + if !entry?.file_type()?.is_file() { + return Ok(false); + } + } + Ok(true) +} + +fn shards_match(output_dir: &Path, report: &MlxDerivedStageReport) -> Result { + for shard in &report.shards { + let relative = Path::new(&shard.file); + if relative.components().count() != 1 + || !matches!(relative.components().next(), Some(Component::Normal(_))) + { + return Ok(false); + } + let path = output_dir.join(relative); + let metadata = match fs::metadata(&path) { + Ok(metadata) if metadata.is_file() => metadata, + _ => return Ok(false), + }; + if metadata.len() != shard.file_bytes || sha256_file(&path)? != shard.sha256 { + return Ok(false); + } + } + Ok(!report.shards.is_empty()) +} + +fn remove_invalid_cache_entry(path: &Path) -> Result<()> { + if path.is_dir() { + fs::remove_dir_all(path) + .with_context(|| format!("remove invalid derived stage cache {}", path.display()))?; + } else if path.exists() { + fs::remove_file(path) + .with_context(|| format!("remove invalid derived stage cache {}", path.display()))?; + } + Ok(()) +} + +#[cfg(test)] +mod tests { + use serde_json::json; + + use super::*; + use crate::derived::{MlxDerivedStageShard, write_json}; + + fn write_test_cache(output_dir: &Path, recipe: &str) -> MlxDerivedStageReport { + fs::create_dir_all(output_dir).unwrap(); + fs::write(output_dir.join("config.json"), b"{}\n").unwrap(); + fs::write(output_dir.join("model.safetensors"), b"packed").unwrap(); + let shard_path = output_dir.join("model.safetensors"); + let mut report = MlxDerivedStageReport { + schema_version: DERIVED_STAGE_SCHEMA_VERSION, + derivation_recipe_sha256: recipe.to_string(), + output_content_sha256: output_content_sha256(output_dir).unwrap(), + checkpoint_sha256: "checkpoint".to_string(), + plan_sha256: "plan".to_string(), + repo: "owner/model".to_string(), + revision: "0".repeat(40), + layer_start: 0, + layer_end: 1, + quantization: json!({"mode":"affine","bits":4,"group_size":64}), + quantization_label: "affine-4bit-g64".to_string(), + safemlx_revision: "revision".to_string(), + output_dir: output_dir.to_path_buf(), + source_tensor_count: 1, + source_tensor_bytes: 16, + source_range_request_count: 1, + source_temporary_file_peak_bytes: 16, + quantized_tensor_count: 1, + copied_tensor_count: 0, + output_tensor_bytes: 6, + artifact_file_bytes: artifact_file_bytes(output_dir).unwrap(), + working_disk_peak_bytes: 22, + mlx_active_memory_bytes: 0, + mlx_cache_memory_bytes: 0, + mlx_peak_memory_bytes: 0, + shards: vec![MlxDerivedStageShard { + file: "model.safetensors".to_string(), + file_bytes: 6, + sha256: sha256_file(&shard_path).unwrap(), + }], + }; + write_json(output_dir.join(REPORT_FILE), &report).unwrap(); + report.output_content_sha256 = output_content_sha256(output_dir).unwrap(); + report + } + + #[test] + fn validates_cache_content_and_rejects_tampering() { + let directory = tempfile::tempdir().unwrap(); + let output = directory.path().join("recipe"); + let expected = write_test_cache(&output, "recipe"); + + let cached = load_cached(&output, "recipe").unwrap().unwrap(); + assert_eq!(cached.output_content_sha256, expected.output_content_sha256); + + fs::write(output.join("model.safetensors"), b"broken").unwrap(); + assert!(load_cached(&output, "recipe").unwrap().is_none()); + } + + #[test] + fn rejects_traversal_in_cached_shard_name() { + let directory = tempfile::tempdir().unwrap(); + let output = directory.path().join("recipe"); + let mut report = write_test_cache(&output, "recipe"); + report.shards[0].file = "../model.safetensors".to_string(); + assert!(!shards_match(&output, &report).unwrap()); + } + + #[test] + fn accepts_relocated_cache_and_refreshes_diagnostic_path() { + let directory = tempfile::tempdir().unwrap(); + let original_root = directory.path().join("original"); + let original = original_root.join("recipe"); + write_test_cache(&original, "recipe"); + let relocated_root = directory.path().join("relocated"); + fs::rename(&original_root, &relocated_root).unwrap(); + let relocated = relocated_root.join("recipe"); + + let cached = load_cached(&relocated, "recipe").unwrap().unwrap(); + + assert_eq!(cached.output_dir, relocated); + } + + #[cfg(unix)] + #[test] + fn rejects_non_regular_cache_entries() { + use std::os::unix::fs::symlink; + + let directory = tempfile::tempdir().unwrap(); + let output = directory.path().join("recipe"); + write_test_cache(&output, "recipe"); + symlink("config.json", output.join("unexpected-index.json")).unwrap(); + + assert!(load_cached(&output, "recipe").unwrap().is_none()); + } +} diff --git a/crates/skippy-engine-mlx/src/lib.rs b/crates/skippy-engine-mlx/src/lib.rs index 099b4560e4..5d8cdccb07 100644 --- a/crates/skippy-engine-mlx/src/lib.rs +++ b/crates/skippy-engine-mlx/src/lib.rs @@ -20,7 +20,9 @@ mod stage; pub use backend::MlxBackend; #[cfg(all(feature = "mlx", target_os = "macos"))] pub use derived::{ - MlxDerivedStageConfig, MlxDerivedStageReport, MlxDerivedStageShard, derive_quantized_stage, + MlxDerivedStageCacheConfig, MlxDerivedStageCacheResult, MlxDerivedStageConfig, + MlxDerivedStageReport, MlxDerivedStageShard, derive_quantized_stage, + derive_quantized_stage_cached, }; #[cfg(all(feature = "mlx", target_os = "macos"))] pub use engine::{ChatTurn, GenerateRequest, MlxEngine, MlxEngineConfig}; diff --git a/docs/design/MLX_STAGE_ENGINE_PLAN.md b/docs/design/MLX_STAGE_ENGINE_PLAN.md index 154af31b61..f4aa2d43d5 100644 --- a/docs/design/MLX_STAGE_ENGINE_PLAN.md +++ b/docs/design/MLX_STAGE_ENGINE_PLAN.md @@ -97,8 +97,15 @@ macOS peak footprint was about 241 MB. The two derived stages reproduced the established affine-4 tokens exactly. The v1 builder deliberately requires `model_type=llama` and rejects pre-quantized sources, rank-3 weights, and incompatible matrix dimensions. This closes the dense-Llama artifact-builder -proof, not reusable host cache hits, Nemotron expert packing, or Inkling's -transformed rank-3 path. +proof, not Nemotron expert packing or Inkling's transformed rank-3 path. + +**Update — the derived-stage cache seam is proven.** `mlx-stage derive-cached` +uses the derivation recipe as a destination identity, serializes competing +builders with an advisory lock, and validates the output-content plus per-shard +hashes on hits. A cold pinned layer-14 run made 9 tensor payload requests; the +warm run made 0, skipped quantization, and used about 17.8 MB max RSS. The warm +path still fetches lightweight config/index/header metadata to reconstruct the +strong recipe. Host prepare/load integration and eviction ownership remain. The pinned safemlx revision also already includes whole-model Inkling text, vision, and audio execution. Earlier notes that called for porting Inkling were diff --git a/spikes/mlx-safetensors-stages/FINDINGS.md b/spikes/mlx-safetensors-stages/FINDINGS.md index 1bda31a97c..7f3dd1d147 100644 --- a/spikes/mlx-safetensors-stages/FINDINGS.md +++ b/spikes/mlx-safetensors-stages/FINDINGS.md @@ -165,9 +165,14 @@ excludes the report. The working-disk metric is a measured source-tensor plus artifact-payload high-water mark, not allocated filesystem blocks or lock/report overhead. The v1 builder fails closed unless `model_type` is exactly `llama`. -This is still an explicit artifact output, not the host's evictable cache. The -next cache step must map the recorded identity to a managed destination, -validate shard hashes on hit, and avoid all tensor-range requests on reuse. +The follow-on `derive-cached` path maps that recipe to a locked managed +directory and validates the report schema, recipe, aggregate artifact bytes, +output-content digest, and every shard hash before accepting a hit. On the same +pinned layer-14 slice, a cold call made 9 tensor-payload requests and the warm +call made 0, returned the identical recipe/content/shard hashes, and used +17,809,408 B max RSS. The warm path still reads lightweight config/index/header +metadata to reconstruct the strong key. This is the reusable library/CLI seam; +host lifecycle and eviction integration remain. ## Representative measurements @@ -398,9 +403,9 @@ because it does not alter the derived packed weights. ## Next proof -1. Use the derivation recipe hash as a managed host cache key; validate the - output-content digest and shard hashes on hit, then prove a warm load makes - no range requests. +1. Wire the proven derived-stage cache into host prepare/load and cache eviction, + then decide whether a local request-to-recipe locator should remove even the + warm metadata probes. 2. Quantize one real Nemotron-H BF16 matrix reproducibly, then implement one complete split-expert bank without accumulating every dense expert. Prove a small family member before attempting Ultra-scale ranges. From fb731b9e45141fbbd468919942f4135a1ef32d54 Mon Sep 17 00:00:00 2001 From: Michael Neale <14976+michaelneale@users.noreply.github.com> Date: Fri, 17 Jul 2026 21:30:03 +1000 Subject: [PATCH 18/37] feat(mlx): prepare quantized stages through host cache --- crates/mesh-llm-host-runtime/src/api/tests.rs | 1 + .../src/inference/skippy/deployment.rs | 1 + .../src/inference/skippy/materialization.rs | 3 + .../src/inference/skippy/mod.rs | 4 +- .../src/inference/skippy/stage/inventory.rs | 13 +- .../src/inference/skippy/stage/mlx.rs | 134 ++++++++++-- .../src/inference/skippy/stage/mod.rs | 9 +- .../src/inference/skippy/stage/status.rs | 5 + .../src/inference/skippy/stage/tests.rs | 41 ++++ .../src/inference/skippy/stage/types.rs | 14 ++ crates/mesh-llm-host-runtime/src/mesh/mod.rs | 1 + .../src/mesh/stage_proto.rs | 87 +++++++- .../mesh-llm-host-runtime/src/mesh/tests.rs | 85 ++++++++ .../src/runtime/local.rs | 31 +++ .../src/safetensors_stage/materialize.rs | 24 +++ .../src/safetensors_stage/tensor_stream.rs | 24 ++- crates/skippy-engine-mlx/STAGED_EXECUTION.md | 55 ++++- crates/skippy-engine-mlx/src/bin/mlx-stage.rs | 11 +- crates/skippy-engine-mlx/src/derived.rs | 115 +++++++++-- crates/skippy-engine-mlx/src/derived/cache.rs | 195 +++++++++++++++--- crates/skippy-engine-mlx/src/lib.rs | 6 +- crates/skippy-protocol/proto/stage.proto | 13 ++ crates/skippy-protocol/src/lib.rs | 1 + docs/design/MLX_STAGE_ENGINE_PLAN.md | 74 ++++--- spikes/mlx-safetensors-stages/FINDINGS.md | 54 +++-- 25 files changed, 864 insertions(+), 137 deletions(-) diff --git a/crates/mesh-llm-host-runtime/src/api/tests.rs b/crates/mesh-llm-host-runtime/src/api/tests.rs index ef7dc2aa49..b5f2acf1e9 100644 --- a/crates/mesh-llm-host-runtime/src/api/tests.rs +++ b/crates/mesh-llm-host-runtime/src/api/tests.rs @@ -2400,6 +2400,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, diff --git a/crates/mesh-llm-host-runtime/src/inference/skippy/deployment.rs b/crates/mesh-llm-host-runtime/src/inference/skippy/deployment.rs index 572d36c14a..487980db1a 100644 --- a/crates/mesh-llm-host-runtime/src/inference/skippy/deployment.rs +++ b/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, diff --git a/crates/mesh-llm-host-runtime/src/inference/skippy/materialization.rs b/crates/mesh-llm-host-runtime/src/inference/skippy/materialization.rs index 7d93ecf05b..cc65c0bacc 100644 --- a/crates/mesh-llm-host-runtime/src/inference/skippy/materialization.rs +++ b/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, diff --git a/crates/mesh-llm-host-runtime/src/inference/skippy/mod.rs b/crates/mesh-llm-host-runtime/src/inference/skippy/mod.rs index 3e3f7edf94..2bb725149e 100644 --- a/crates/mesh-llm-host-runtime/src/inference/skippy/mod.rs +++ b/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}; diff --git a/crates/mesh-llm-host-runtime/src/inference/skippy/stage/inventory.rs b/crates/mesh-llm-host-runtime/src/inference/skippy/stage/inventory.rs index 3b9a8a1ab6..5a03773254 100644 --- a/crates/mesh-llm-host-runtime/src/inference/skippy/stage/inventory.rs +++ b/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,13 +187,18 @@ struct PrepareSourceResult { bytes_total: Option, } -async fn prepare_stage_source(load: &StageLoadRequest) -> Result { +async fn prepare_stage_source( + load: &StageLoadRequest, + cancelled: Arc, +) -> Result { + #[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).await?; + let artifact = super::mlx::prepare_stage(load, cancelled).await?; return Ok(PrepareSourceResult { - bytes_total: Some(artifact.manifest.output_file_bytes), + bytes_total: Some(artifact.report.artifact_file_bytes), }); } #[cfg(not(all(feature = "mlx", target_os = "macos")))] diff --git a/crates/mesh-llm-host-runtime/src/inference/skippy/stage/mlx.rs b/crates/mesh-llm-host-runtime/src/inference/skippy/stage/mlx.rs index f0b3c548bd..4f9c8da20c 100644 --- a/crates/mesh-llm-host-runtime/src/inference/skippy/stage/mlx.rs +++ b/crates/mesh-llm-host-runtime/src/inference/skippy/stage/mlx.rs @@ -1,10 +1,15 @@ -use std::{net::SocketAddr, sync::Arc}; +use std::{ + net::SocketAddr, + sync::{Arc, atomic::AtomicBool}, +}; use anyhow::{Context, Result, bail, ensure}; -use model_hf::safetensors_stage::{ - SafetensorsStageArtifact, SafetensorsStageMaterializer, SafetensorsStageRequest, +use model_hf::safetensors_stage::{SafetensorsStageMaterializer, SafetensorsStageRequest}; +use skippy_engine_mlx::{ + MlxComputeDtype, MlxDerivationControl, MlxDerivedStageCacheConfig, MlxDerivedStageCacheResult, + MlxStageEngine, MlxStageEngineConfig, MlxWeightQuantization, derive_quantized_stage_cached, + load_prepared_quantized_stage, mlx_derived_stage_cache_root, }; -use skippy_engine_mlx::{MlxComputeDtype, MlxStageEngine, MlxStageEngineConfig}; use skippy_protocol::LoadMode; use skippy_protocol::binary::WireActivationDType; use skippy_server::{ @@ -12,16 +17,17 @@ use skippy_server::{ }; use super::{ - RunningStage, StageControlState, StageLoadRequest, StageReadyResponse, StageStatusFilter, - StageWireDType, stage_load_failure_context, + MlxStageArtifact, RunningStage, StageControlState, StageLoadRequest, StageReadyResponse, + StageStatusFilter, StageWeightQuantization, StageWireDType, stage_load_failure_context, }; const HF_MODEL_PREFIX: &str = "hf-model://"; +const DERIVED_SHARD_SIZE_BYTES: usize = 256 * 1024 * 1024; pub(super) struct MlxStageLaunch { pub(super) load: StageLoadRequest, pub(super) server: EmbeddedServerHandle, - pub(super) artifact: SafetensorsStageArtifact, + pub(super) artifact: MlxDerivedStageCacheResult, } pub(super) async fn load_stage( @@ -42,13 +48,16 @@ pub(super) async fn load_stage( return Err(error.context(context)); } let effective_load = launch.load; + let artifact = MlxStageArtifact { + path: launch.artifact.output_dir, + }; state.stages.insert( key, RunningStage { load: effective_load.clone(), server: launch.server, materialized: None, - mlx_artifact: Some(launch.artifact), + mlx_artifact: Some(artifact), package: None, _materialized_pin: None, }, @@ -107,11 +116,14 @@ async fn wait_for_engine_stage_ready( } } -pub(super) async fn prepare_stage(load: &StageLoadRequest) -> Result { +pub(super) async fn prepare_stage( + load: &StageLoadRequest, + cancelled: Arc, +) -> Result { let load = load.clone(); - tokio::task::spawn_blocking(move || materialize_stage_blocking(&load)) + tokio::task::spawn_blocking(move || derive_stage_blocking(&load, Some(cancelled))) .await - .context("join MLX SafeTensors stage preparation")? + .context("join MLX derived stage preparation")? } pub(super) async fn launch_stage( @@ -120,9 +132,9 @@ pub(super) async fn launch_stage( ) -> Result { let blocking_load = load.clone(); let (artifact, engine) = tokio::task::spawn_blocking(move || { - let artifact = materialize_stage_blocking(&blocking_load)?; + let artifact = resolve_stage_blocking(&blocking_load, None, false)?; let engine = Arc::new(MlxStageEngine::spawn(MlxStageEngineConfig { - model_dir: artifact.path.clone(), + model_dir: artifact.output_dir.clone(), model_id: blocking_load.model_id.clone(), stage_index: blocking_load.stage_index, layer_start: blocking_load.layer_start, @@ -143,7 +155,7 @@ pub(super) async fn launch_stage( .context("join MLX stage load task")??; load.bind_addr = bind_addr.to_string(); - load.model_path = Some(artifact.path.to_string_lossy().into_owned()); + load.model_path = Some(artifact.output_dir.to_string_lossy().into_owned()); let server = skippy_server::start_stage_engine( engine, EngineStageServerOptions { @@ -159,18 +171,64 @@ pub(super) async fn launch_stage( }) } -fn materialize_stage_blocking(load: &StageLoadRequest) -> Result { +fn derive_stage_blocking( + load: &StageLoadRequest, + cancelled: Option>, +) -> Result { + resolve_stage_blocking(load, cancelled, true) +} + +fn resolve_stage_blocking( + load: &StageLoadRequest, + cancelled: Option>, + build_on_miss: bool, +) -> Result { let request = request_from_load(load)?; - let artifact = SafetensorsStageMaterializer::from_environment()?.materialize(request)?; + let materializer = SafetensorsStageMaterializer::from_environment()?; + let config = MlxDerivedStageCacheConfig { + source: request, + cache_root: mlx_derived_stage_cache_root(), + quantization: mlx_weight_quantization(load.weight_quantization), + control: MlxDerivationControl::new(Some(load.manifest_sha256.clone()), cancelled), + shard_size_bytes: DERIVED_SHARD_SIZE_BYTES, + }; + let artifact = if build_on_miss { + derive_quantized_stage_cached(&materializer, &config)? + } else { + load_prepared_quantized_stage(&materializer, &config)? + }; ensure!( - artifact.manifest.checkpoint_sha256 == load.manifest_sha256, + artifact.report.checkpoint_sha256 == load.manifest_sha256, "MLX checkpoint identity {} does not match stage claim {}", - artifact.manifest.checkpoint_sha256, + artifact.report.checkpoint_sha256, load.manifest_sha256 ); + tracing::info!( + cache_hit = artifact.cache_hit, + source_range_request_count = artifact.source_range_request_count, + derivation_recipe_sha256 = %artifact.report.derivation_recipe_sha256, + stage_id = %load.stage_id, + "MLX derived stage cache resolved" + ); Ok(artifact) } +fn mlx_weight_quantization(quantization: StageWeightQuantization) -> MlxWeightQuantization { + match quantization { + StageWeightQuantization::Auto | StageWeightQuantization::Affine4 => { + MlxWeightQuantization::Affine { + group_size: 64, + bits: 4, + } + } + StageWeightQuantization::Affine8 => MlxWeightQuantization::Affine { + group_size: 64, + bits: 8, + }, + StageWeightQuantization::MxFp4 => MlxWeightQuantization::MxFp4, + } +} + fn request_from_load(load: &StageLoadRequest) -> Result { validate_load_settings(load)?; ensure!(load.backend == "mlx", "MLX stage requires backend=mlx"); @@ -270,6 +328,7 @@ fn wire_dtype(dtype: StageWireDType) -> Result { mod tests { use std::{io::Write, net::TcpStream, time::Duration}; + use skippy_engine_mlx::MlxDerivedStageReport; use skippy_protocol::FlashAttentionType; use skippy_protocol::binary::{ StageStateHeader, StageWireMessage, WireMessageKind, WireReplyKind, recv_ready, recv_reply, @@ -286,7 +345,7 @@ mod tests { const SMOL_REPO: &str = "HuggingFaceTB/SmolLM2-135M-Instruct"; const SMOL_REVISION: &str = "12fd25f77366fa6b3b4b768ec3050bf629380bac"; const SMOL_PROMPT: &[i32] = &[1, 1531, 314, 260, 3575, 28]; - const SMOL_EXPECTED: &[i32] = &[284, 260, 2240, 314, 1343, 327, 624, 8685]; + const SMOL_EXPECTED: &[i32] = &[260, 2240, 314, 253, 1379, 282, 25801, 28]; #[test] fn parses_commit_addressed_mlx_stage_ref() { @@ -331,9 +390,31 @@ mod tests { assert!(request_from_load(&load).is_err()); } + #[test] + fn maps_explicit_and_hardware_auto_weight_quantization() { + assert_eq!( + mlx_weight_quantization(StageWeightQuantization::Auto), + MlxWeightQuantization::Affine { + group_size: 64, + bits: 4 + } + ); + assert_eq!( + mlx_weight_quantization(StageWeightQuantization::Affine8), + MlxWeightQuantization::Affine { + group_size: 64, + bits: 8 + } + ); + assert_eq!( + mlx_weight_quantization(StageWeightQuantization::MxFp4), + MlxWeightQuantization::MxFp4 + ); + } + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] #[ignore = "downloads about 310 MiB and requires Apple Silicon Metal"] - async fn real_control_plane_materializes_and_runs_two_range_only_stages() -> Result<()> { + async fn real_control_plane_derives_and_runs_two_quantized_range_stages() -> Result<()> { let plan = tokio::task::spawn_blocking(|| { SafetensorsStageMaterializer::from_environment()?.plan(SafetensorsStageRequest { repo: SMOL_REPO.to_string(), @@ -404,6 +485,7 @@ mod tests { model_id: load.model_id.clone(), package_ref: load.package_ref.clone(), manifest_sha256: load.manifest_sha256.clone(), + weight_quantization: load.weight_quantization, }), ) .await?; @@ -462,10 +544,21 @@ mod tests { let plan: model_hf::safetensors_stage::SafetensorsStagePlan = serde_json::from_slice( &std::fs::read(std::path::Path::new(artifact_path).join("stage-plan.json"))?, )?; + let report: MlxDerivedStageReport = serde_json::from_slice(&std::fs::read( + std::path::Path::new(artifact_path).join("derived-stage.json"), + )?)?; ensure!( plan.planned_download_bytes < plan.source_shard_bytes, "running MLX stage planned a complete source shard download" ); + ensure!( + report.checkpoint_sha256 == load.manifest_sha256, + "running MLX stage changed source checkpoint identity" + ); + ensure!( + report.quantization_label == "affine-4bit-g64", + "running MLX stage did not use the requested quantization" + ); Ok(()) } @@ -636,6 +729,7 @@ mod tests { n_gpu_layers: 0, mmap: None, mlock: false, + weight_quantization: StageWeightQuantization::Affine4, cache_type_k: "f16".to_string(), cache_type_v: "f16".to_string(), flash_attn_type: FlashAttentionType::Auto, diff --git a/crates/mesh-llm-host-runtime/src/inference/skippy/stage/mod.rs b/crates/mesh-llm-host-runtime/src/inference/skippy/stage/mod.rs index 61fceab6bf..3b30704bd6 100644 --- a/crates/mesh-llm-host-runtime/src/inference/skippy/stage/mod.rs +++ b/crates/mesh-llm-host-runtime/src/inference/skippy/stage/mod.rs @@ -40,11 +40,15 @@ struct RunningStage { load: StageLoadRequest, server: EmbeddedServerHandle, materialized: Option, - mlx_artifact: Option, + mlx_artifact: Option, package: Option, _materialized_pin: Option, } +struct MlxStageArtifact { + path: std::path::PathBuf, +} + #[derive(Default)] struct StageControlState { stages: HashMap, @@ -152,6 +156,7 @@ impl StageControlState { status.model_id == request.model_id && status.package_ref == request.package_ref && status.manifest_sha256 == request.manifest_sha256 + && status.weight_quantization == request.weight_quantization }) .cloned() .collect::>(); @@ -175,6 +180,7 @@ impl StageControlState { stage.load.model_id == request.model_id && stage.load.package_ref == request.package_ref && stage.load.manifest_sha256 == request.manifest_sha256 + && stage.load.weight_quantization == request.weight_quantization }) .map(|stage| LayerRange { layer_start: stage.load.layer_start, @@ -206,6 +212,7 @@ impl StageControlState { .as_ref() .map(|source| source.kind) .unwrap_or(SourceModelKind::Unknown), + weight_quantization: request.weight_quantization, } } diff --git a/crates/mesh-llm-host-runtime/src/inference/skippy/stage/status.rs b/crates/mesh-llm-host-runtime/src/inference/skippy/stage/status.rs index c50e1377a7..cf597c9b73 100644 --- a/crates/mesh-llm-host-runtime/src/inference/skippy/stage/status.rs +++ b/crates/mesh-llm-host-runtime/src/inference/skippy/stage/status.rs @@ -35,6 +35,7 @@ pub(super) fn status_from_running(stage: &RunningStage) -> StageStatusSnapshot { n_batch: stage.load.n_batch, n_ubatch: stage.load.n_ubatch, flash_attn_type: stage.load.flash_attn_type, + weight_quantization: stage.load.weight_quantization, error: server.last_error, shutdown_generation: stage.load.shutdown_generation, coordinator_term: stage.load.coordinator_term, @@ -141,6 +142,7 @@ pub(super) fn stopped_status(stop: &StageStopRequest) -> StageStatusSnapshot { n_batch: None, n_ubatch: None, flash_attn_type: FlashAttentionType::Auto, + weight_quantization: super::StageWeightQuantization::Auto, error: None, shutdown_generation: stop.shutdown_generation, coordinator_term: stop.coordinator_term, @@ -180,6 +182,7 @@ pub(super) fn failed_status_from_load( n_batch: load.n_batch, n_ubatch: load.n_ubatch, flash_attn_type: load.flash_attn_type, + weight_quantization: load.weight_quantization, error: Some(error), shutdown_generation: load.shutdown_generation, coordinator_term: load.coordinator_term, @@ -204,6 +207,7 @@ pub(super) fn preparation_status_from_load( stage_index: load.stage_index, layer_start: load.layer_start, layer_end: load.layer_end, + weight_quantization: load.weight_quantization, state, bytes_done: None, bytes_total: None, @@ -230,6 +234,7 @@ pub(super) fn preparation_status_from_cancel( stage_index: 0, layer_start: 0, layer_end: 0, + weight_quantization: super::StageWeightQuantization::Auto, state: StagePreparationState::Cancelled, bytes_done: None, bytes_total: None, diff --git a/crates/mesh-llm-host-runtime/src/inference/skippy/stage/tests.rs b/crates/mesh-llm-host-runtime/src/inference/skippy/stage/tests.rs index 0b0cddd9ba..940da99259 100644 --- a/crates/mesh-llm-host-runtime/src/inference/skippy/stage/tests.rs +++ b/crates/mesh-llm-host-runtime/src/inference/skippy/stage/tests.rs @@ -40,6 +40,7 @@ fn load_request() -> StageLoadRequest { n_gpu_layers: -1, mmap: Some(false), mlock: true, + weight_quantization: StageWeightQuantization::Auto, cache_type_k: "f16".to_string(), cache_type_v: "q8_0".to_string(), flash_attn_type: FlashAttentionType::Enabled, @@ -370,6 +371,7 @@ async fn prepare_stage_records_background_source_availability() { model_id: load.model_id.clone(), package_ref: load.package_ref.clone(), manifest_sha256: load.manifest_sha256.clone(), + weight_quantization: load.weight_quantization, }) .await; if let Some(status) = inventory @@ -417,6 +419,7 @@ async fn prepare_mlx_fails_closed_without_downloading_on_unsupported_build() { model_id: load.model_id.clone(), package_ref: load.package_ref.clone(), manifest_sha256: load.manifest_sha256.clone(), + weight_quantization: load.weight_quantization, }) .await; if let Some(status) = inventory @@ -471,6 +474,7 @@ async fn prepare_layer_package_stays_downloading_while_peer_prefetch_is_pending( model_id: load.model_id.clone(), package_ref: load.package_ref.clone(), manifest_sha256: load.manifest_sha256.clone(), + weight_quantization: load.weight_quantization, }) .await; let status = inventory @@ -513,6 +517,7 @@ async fn prepare_layer_package_fails_only_after_peer_prefetch_and_local_resoluti model_id: load.model_id.clone(), package_ref: load.package_ref.clone(), manifest_sha256: load.manifest_sha256.clone(), + weight_quantization: load.weight_quantization, }) .await; if let Some(status) = inventory @@ -579,6 +584,7 @@ async fn cancel_prepare_persists_cancelled_status_and_blocks_late_prefetch_resul model_id: load.model_id.clone(), package_ref: load.package_ref.clone(), manifest_sha256: load.manifest_sha256.clone(), + weight_quantization: load.weight_quantization, }) .await; let status = inventory @@ -649,6 +655,7 @@ async fn prepare_preserves_equal_or_newer_cancelled_status() { model_id: load.model_id.clone(), package_ref: load.package_ref.clone(), manifest_sha256: load.manifest_sha256.clone(), + weight_quantization: load.weight_quantization, }) .await; let stored = inventory @@ -700,6 +707,7 @@ async fn status_update_upserts_preparation_status_and_rejects_stale_generation() model_id: load.model_id.clone(), package_ref: load.package_ref.clone(), manifest_sha256: load.manifest_sha256.clone(), + weight_quantization: load.weight_quantization, }) .await; let status = inventory @@ -724,6 +732,7 @@ async fn status_update_upserts_preparation_status_and_rejects_stale_generation() model_id: load.model_id.clone(), package_ref: load.package_ref.clone(), manifest_sha256: load.manifest_sha256.clone(), + weight_quantization: load.weight_quantization, }) .await; let status = inventory @@ -749,6 +758,7 @@ async fn inventory_retains_failed_prepare_status() { model_id: load.model_id.clone(), package_ref: load.package_ref.clone(), manifest_sha256: load.manifest_sha256.clone(), + weight_quantization: load.weight_quantization, }) .await; @@ -761,12 +771,43 @@ async fn inventory_retains_failed_prepare_status() { assert_eq!(status.error.as_deref(), Some("source unavailable")); } +#[tokio::test] +async fn inventory_does_not_report_a_different_weight_profile_as_prepared() { + let mut load = load_request(); + load.backend = "mlx".to_string(); + load.weight_quantization = StageWeightQuantization::Affine4; + let key = stage_key(&load.topology_id, &load.run_id, &load.stage_id); + let state = StageControlState::default(); + state.preparations.lock().await.insert( + key, + preparation_status_from_load(&load, StagePreparationState::Available, None), + ); + + let mut request = StageInventoryRequest { + model_id: load.model_id.clone(), + package_ref: load.package_ref.clone(), + manifest_sha256: load.manifest_sha256.clone(), + weight_quantization: StageWeightQuantization::Affine8, + }; + assert!( + state + .inventory(request.clone()) + .await + .preparing_ranges + .is_empty() + ); + + request.weight_quantization = StageWeightQuantization::Affine4; + assert_eq!(state.inventory(request).await.preparing_ranges.len(), 1); +} + #[test] fn inventory_source_candidates_prefer_explicit_gguf_ref() { let request = StageInventoryRequest { model_id: "catalog-model".to_string(), package_ref: "gguf:///tmp/source-model.gguf".to_string(), manifest_sha256: "sha256".to_string(), + weight_quantization: StageWeightQuantization::Auto, }; let candidates = inventory_source_candidates(&request); diff --git a/crates/mesh-llm-host-runtime/src/inference/skippy/stage/types.rs b/crates/mesh-llm-host-runtime/src/inference/skippy/stage/types.rs index e61be71034..260d2c7dfb 100644 --- a/crates/mesh-llm-host-runtime/src/inference/skippy/stage/types.rs +++ b/crates/mesh-llm-host-runtime/src/inference/skippy/stage/types.rs @@ -68,6 +68,7 @@ pub(crate) struct StageLoadRequest { pub(crate) n_gpu_layers: i32, pub(crate) mmap: Option, pub(crate) mlock: bool, + pub(crate) weight_quantization: StageWeightQuantization, pub(crate) cache_type_k: String, pub(crate) cache_type_v: String, pub(crate) flash_attn_type: FlashAttentionType, @@ -102,6 +103,7 @@ pub(crate) struct StageInventoryRequest { pub(crate) model_id: String, pub(crate) package_ref: String, pub(crate) manifest_sha256: String, + pub(crate) weight_quantization: StageWeightQuantization, } #[derive(Clone, Debug)] @@ -137,6 +139,7 @@ pub(crate) struct StageLayerInventory { pub(crate) source_model_path: Option, pub(crate) source_model_bytes: Option, pub(crate) source_model_kind: SourceModelKind, + pub(crate) weight_quantization: StageWeightQuantization, } #[derive(Clone, Copy, Debug, Eq, PartialEq)] @@ -162,6 +165,15 @@ pub(crate) enum StageWireDType { Q8, } +#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] +pub(crate) enum StageWeightQuantization { + #[default] + Auto, + Affine4, + Affine8, + MxFp4, +} + #[derive(Clone, Copy, Debug, Eq, PartialEq)] pub(crate) enum StageRuntimeState { Starting, @@ -218,6 +230,7 @@ pub(crate) struct StageStatusSnapshot { pub(crate) n_batch: Option, pub(crate) n_ubatch: Option, pub(crate) flash_attn_type: FlashAttentionType, + pub(crate) weight_quantization: StageWeightQuantization, pub(crate) error: Option, pub(crate) shutdown_generation: u64, pub(crate) coordinator_term: u64, @@ -237,6 +250,7 @@ pub(crate) struct StagePreparationStatus { pub(crate) stage_index: u32, pub(crate) layer_start: u32, pub(crate) layer_end: u32, + pub(crate) weight_quantization: StageWeightQuantization, pub(crate) state: StagePreparationState, pub(crate) bytes_done: Option, pub(crate) bytes_total: Option, diff --git a/crates/mesh-llm-host-runtime/src/mesh/mod.rs b/crates/mesh-llm-host-runtime/src/mesh/mod.rs index 375d7ccda8..347c95722b 100644 --- a/crates/mesh-llm-host-runtime/src/mesh/mod.rs +++ b/crates/mesh-llm-host-runtime/src/mesh/mod.rs @@ -2699,6 +2699,7 @@ pub struct StageRuntimeStatus { pub n_batch: Option, pub n_ubatch: Option, pub flash_attn_type: skippy_protocol::FlashAttentionType, + pub weight_quantization: crate::inference::skippy::StageWeightQuantization, pub error: Option, pub shutdown_generation: u64, } diff --git a/crates/mesh-llm-host-runtime/src/mesh/stage_proto.rs b/crates/mesh-llm-host-runtime/src/mesh/stage_proto.rs index 7f765a2047..ffd9d8268e 100644 --- a/crates/mesh-llm-host-runtime/src/mesh/stage_proto.rs +++ b/crates/mesh-llm-host-runtime/src/mesh/stage_proto.rs @@ -65,6 +65,7 @@ pub(super) fn stage_runtime_status_from_snapshot( n_batch: status.n_batch, n_ubatch: status.n_ubatch, flash_attn_type: status.flash_attn_type, + weight_quantization: status.weight_quantization, error: status.error, shutdown_generation: status.shutdown_generation, } @@ -102,6 +103,7 @@ pub(super) fn stage_snapshot_from_runtime_status( n_batch: status.n_batch, n_ubatch: status.n_ubatch, flash_attn_type: status.flash_attn_type, + weight_quantization: status.weight_quantization, error, shutdown_generation: status.shutdown_generation, coordinator_term: 0, @@ -167,6 +169,9 @@ pub(super) fn stage_control_request_to_proto( model_id: inventory.model_id, package_ref: inventory.package_ref, manifest_sha256: inventory.manifest_sha256, + weight_quantization: stage_weight_quantization_to_proto( + inventory.weight_quantization, + ) as i32, }) } crate::inference::skippy::StageControlRequest::Prepare(prepare) => { @@ -225,6 +230,7 @@ pub(super) fn stage_load_to_proto( n_gpu_layers: load.n_gpu_layers, mmap: load.mmap, mlock: Some(load.mlock), + weight_quantization: stage_weight_quantization_to_proto(load.weight_quantization) as i32, cache_type_k: load.cache_type_k, cache_type_v: load.cache_type_v, flash_attn_type: stage_flash_attn_type_to_proto(load.flash_attn_type) as i32, @@ -329,6 +335,9 @@ pub(super) fn stage_control_request_from_proto( model_id: inventory.model_id, package_ref: inventory.package_ref, manifest_sha256: inventory.manifest_sha256, + weight_quantization: stage_weight_quantization_from_proto( + inventory.weight_quantization, + )?, }, )) } @@ -362,7 +371,7 @@ pub(super) fn stage_control_request_from_proto( .status .ok_or_else(|| anyhow::anyhow!("stage status update missing status"))?; Ok(crate::inference::skippy::StageControlRequest::StatusUpdate( - stage_preparation_status_from_proto(status), + stage_preparation_status_from_proto(status)?, )) } } @@ -404,6 +413,7 @@ pub(super) fn stage_load_from_proto( n_gpu_layers: load.n_gpu_layers, mmap: load.mmap, mlock: load.mlock.unwrap_or(false), + weight_quantization: stage_weight_quantization_from_proto(load.weight_quantization)?, cache_type_k: load.cache_type_k, cache_type_v: load.cache_type_v, flash_attn_type: stage_flash_attn_type_from_proto(load.flash_attn_type), @@ -501,6 +511,30 @@ pub(super) fn stage_wire_dtype_from_proto(value: i32) -> crate::inference::skipp } } +pub(super) fn stage_weight_quantization_from_proto( + value: i32, +) -> anyhow::Result { + Ok( + match skippy_stage_proto::StageWeightQuantization::try_from(value) + .map_err(|_| anyhow::anyhow!("unsupported stage weight quantization value {value}"))? + { + skippy_stage_proto::StageWeightQuantization::Unspecified + | skippy_stage_proto::StageWeightQuantization::Auto => { + crate::inference::skippy::StageWeightQuantization::Auto + } + skippy_stage_proto::StageWeightQuantization::Affine4 => { + crate::inference::skippy::StageWeightQuantization::Affine4 + } + skippy_stage_proto::StageWeightQuantization::Affine8 => { + crate::inference::skippy::StageWeightQuantization::Affine8 + } + skippy_stage_proto::StageWeightQuantization::Mxfp4 => { + crate::inference::skippy::StageWeightQuantization::MxFp4 + } + }, + ) +} + pub(super) fn stage_control_unavailable_response( request: crate::inference::skippy::StageControlRequest, ) -> crate::inference::skippy::StageControlResponse { @@ -545,6 +579,7 @@ pub(super) fn stage_control_unavailable_response( n_batch: None, n_ubatch: None, flash_attn_type: skippy_protocol::FlashAttentionType::Auto, + weight_quantization: crate::inference::skippy::StageWeightQuantization::Auto, error: Some("stage control is not available".to_string()), shutdown_generation: stop.shutdown_generation, coordinator_term: stop.coordinator_term, @@ -569,6 +604,7 @@ pub(super) fn stage_control_unavailable_response( source_model_path: None, source_model_bytes: None, source_model_kind: crate::inference::skippy::SourceModelKind::Unknown, + weight_quantization: inventory.weight_quantization, }, ); } @@ -643,6 +679,7 @@ pub(super) fn stage_status_from_load( n_batch: load.n_batch, n_ubatch: load.n_ubatch, flash_attn_type: load.flash_attn_type, + weight_quantization: load.weight_quantization, error: Some("stage control is not available".to_string()), shutdown_generation: load.shutdown_generation, coordinator_term: load.coordinator_term, @@ -667,6 +704,7 @@ pub(super) fn stage_preparation_status_from_load( stage_index: load.stage_index, layer_start: load.layer_start, layer_end: load.layer_end, + weight_quantization: load.weight_quantization, state, bytes_done: None, bytes_total: None, @@ -695,6 +733,7 @@ pub(super) fn stage_preparation_status_from_cancel( stage_index: 0, layer_start: 0, layer_end: 0, + weight_quantization: crate::inference::skippy::StageWeightQuantization::Auto, state, bytes_done: None, bytes_total: None, @@ -821,7 +860,7 @@ pub(super) fn stage_control_response_from_proto( } Response::LayerInventory(inventory) => { Ok(crate::inference::skippy::StageControlResponse::Inventory( - layer_inventory_from_proto(inventory), + layer_inventory_from_proto(inventory)?, )) } Response::PrepareStageAccepted(accepted) => { @@ -832,7 +871,7 @@ pub(super) fn stage_control_response_from_proto( crate::inference::skippy::StageControlResponse::PrepareAccepted( crate::inference::skippy::StagePrepareAcceptedResponse { accepted: accepted.accepted, - status: stage_preparation_status_from_proto(status), + status: stage_preparation_status_from_proto(status)?, error: accepted.error, }, ), @@ -840,7 +879,7 @@ pub(super) fn stage_control_response_from_proto( } Response::StagePreparationStatus(status) => Ok( crate::inference::skippy::StageControlResponse::PreparationStatus( - stage_preparation_status_from_proto(status), + stage_preparation_status_from_proto(status)?, ), ), Response::StageStatusAck(ack) => { @@ -885,13 +924,15 @@ pub(super) fn layer_inventory_to_proto( source_model_path: inventory.source_model_path, source_model_bytes: inventory.source_model_bytes, source_model_kind: source_model_kind_to_proto(inventory.source_model_kind) as i32, + weight_quantization: stage_weight_quantization_to_proto(inventory.weight_quantization) + as i32, } } pub(super) fn layer_inventory_from_proto( inventory: skippy_stage_proto::LayerInventory, -) -> crate::inference::skippy::StageLayerInventory { - crate::inference::skippy::StageLayerInventory { +) -> anyhow::Result { + Ok(crate::inference::skippy::StageLayerInventory { model_id: inventory.model_id, package_ref: inventory.package_ref, manifest_sha256: inventory.manifest_sha256, @@ -915,11 +956,12 @@ pub(super) fn layer_inventory_from_proto( .preparing_ranges .into_iter() .map(stage_preparation_status_from_proto) - .collect(), + .collect::>>()?, source_model_path: inventory.source_model_path, source_model_bytes: inventory.source_model_bytes, source_model_kind: source_model_kind_from_proto(inventory.source_model_kind), - } + weight_quantization: stage_weight_quantization_from_proto(inventory.weight_quantization)?, + }) } pub(super) fn layer_range_to_proto( @@ -1003,12 +1045,13 @@ pub(super) fn stage_preparation_status_to_proto( coordinator_term: status.coordinator_term, coordinator_id: status.coordinator_id.map(|id| id.to_string()), lease_until_unix_ms: status.lease_until_unix_ms, + weight_quantization: stage_weight_quantization_to_proto(status.weight_quantization) as i32, } } pub(super) fn stage_preparation_status_from_proto( status: skippy_stage_proto::StagePreparationStatus, -) -> crate::inference::skippy::StagePreparationStatus { +) -> anyhow::Result { let coordinator_id = status.coordinator_id.and_then(|id| match id.parse() { Ok(id) => Some(id), Err(error) => { @@ -1020,7 +1063,7 @@ pub(super) fn stage_preparation_status_from_proto( None } }); - crate::inference::skippy::StagePreparationStatus { + Ok(crate::inference::skippy::StagePreparationStatus { topology_id: status.topology_id, run_id: status.run_id, model_id: status.model_id, @@ -1031,6 +1074,7 @@ pub(super) fn stage_preparation_status_from_proto( stage_index: status.stage_index, layer_start: status.layer_start, layer_end: status.layer_end, + weight_quantization: stage_weight_quantization_from_proto(status.weight_quantization)?, state: stage_preparation_state_from_proto(status.state), bytes_done: status.bytes_done, bytes_total: status.bytes_total, @@ -1040,7 +1084,7 @@ pub(super) fn stage_preparation_status_from_proto( coordinator_term: status.coordinator_term, coordinator_id, lease_until_unix_ms: status.lease_until_unix_ms, - } + }) } pub(super) fn stage_status_to_proto( @@ -1078,6 +1122,7 @@ pub(super) fn stage_status_to_proto( coordinator_term: status.coordinator_term, coordinator_id: status.coordinator_id.map(|id| id.to_string()), lease_until_unix_ms: status.lease_until_unix_ms, + weight_quantization: stage_weight_quantization_to_proto(status.weight_quantization) as i32, } } @@ -1118,6 +1163,7 @@ pub(super) fn stage_status_from_proto( materialized_pinned: status.materialized_pinned.unwrap_or(false), projector_path: status.projector_path, flash_attn_type: stage_flash_attn_type_from_proto(status.flash_attn_type), + weight_quantization: stage_weight_quantization_from_proto(status.weight_quantization)?, error: status.error, shutdown_generation: status.shutdown_generation, coordinator_term: status.coordinator_term, @@ -1286,3 +1332,22 @@ pub(super) fn stage_wire_dtype_to_proto( } } } + +pub(super) fn stage_weight_quantization_to_proto( + quantization: crate::inference::skippy::StageWeightQuantization, +) -> skippy_stage_proto::StageWeightQuantization { + match quantization { + crate::inference::skippy::StageWeightQuantization::Auto => { + skippy_stage_proto::StageWeightQuantization::Auto + } + crate::inference::skippy::StageWeightQuantization::Affine4 => { + skippy_stage_proto::StageWeightQuantization::Affine4 + } + crate::inference::skippy::StageWeightQuantization::Affine8 => { + skippy_stage_proto::StageWeightQuantization::Affine8 + } + crate::inference::skippy::StageWeightQuantization::MxFp4 => { + skippy_stage_proto::StageWeightQuantization::Mxfp4 + } + } +} diff --git a/crates/mesh-llm-host-runtime/src/mesh/tests.rs b/crates/mesh-llm-host-runtime/src/mesh/tests.rs index b7c64bc70f..6f06d24abc 100644 --- a/crates/mesh-llm-host-runtime/src/mesh/tests.rs +++ b/crates/mesh-llm-host-runtime/src/mesh/tests.rs @@ -308,6 +308,7 @@ fn stage_load_request() -> crate::inference::skippy::StageLoadRequest { n_gpu_layers: -1, mmap: Some(false), mlock: true, + weight_quantization: crate::inference::skippy::StageWeightQuantization::Affine8, cache_type_k: "f16".to_string(), cache_type_v: "q8_0".to_string(), flash_attn_type: skippy_protocol::FlashAttentionType::Auto, @@ -651,12 +652,75 @@ fn stage_load_proto_roundtrip_preserves_source_model_bytes() { assert_eq!(proto.source_model_bytes, Some(123_456_789)); assert_eq!(proto.mmap, Some(false)); assert_eq!(proto.mlock, Some(true)); + assert_eq!( + proto.weight_quantization, + skippy_stage_proto::StageWeightQuantization::Affine8 as i32 + ); let decoded = stage_load_from_proto(proto).unwrap(); assert_eq!(decoded.source_model_bytes, Some(123_456_789)); assert_eq!(decoded.model_path.as_deref(), Some("/models/demo.gguf")); assert_eq!(decoded.mmap, Some(false)); assert!(decoded.mlock); + assert_eq!( + decoded.weight_quantization, + crate::inference::skippy::StageWeightQuantization::Affine8 + ); +} + +#[test] +fn stage_load_proto_defaults_missing_weight_quantization_and_rejects_unknown_values() { + let mut proto = stage_load_to_proto(stage_load_request()); + proto.weight_quantization = 0; + assert_eq!( + stage_load_from_proto(proto.clone()) + .unwrap() + .weight_quantization, + crate::inference::skippy::StageWeightQuantization::Auto + ); + + proto.weight_quantization = 99; + assert!(stage_load_from_proto(proto).is_err()); +} + +#[test] +fn stage_status_proto_defaults_missing_weight_quantization_and_rejects_unknown_values() { + let mut preparation = stage_preparation_status_to_proto(test_preparation_status( + crate::inference::skippy::StagePreparationState::Available, + )); + preparation.weight_quantization = 0; + assert_eq!( + stage_preparation_status_from_proto(preparation.clone()) + .unwrap() + .weight_quantization, + crate::inference::skippy::StageWeightQuantization::Auto + ); + preparation.weight_quantization = 99; + assert!(stage_preparation_status_from_proto(preparation).is_err()); + + let mut status = stage_status_to_proto(stage_status_from_load( + &stage_load_request(), + crate::inference::skippy::StageRuntimeState::Ready, + )); + status.weight_quantization = 0; + assert_eq!( + stage_status_from_proto(status.clone()) + .unwrap() + .weight_quantization, + crate::inference::skippy::StageWeightQuantization::Auto + ); + status.weight_quantization = 99; + assert!(stage_status_from_proto(status).is_err()); + + let mut inventory = skippy_stage_proto::LayerInventory::default(); + assert_eq!( + layer_inventory_from_proto(inventory.clone()) + .unwrap() + .weight_quantization, + crate::inference::skippy::StageWeightQuantization::Auto + ); + inventory.weight_quantization = 99; + assert!(layer_inventory_from_proto(inventory).is_err()); } #[test] @@ -7535,6 +7599,7 @@ fn test_stage_status( n_batch: None, n_ubatch: None, flash_attn_type: skippy_protocol::FlashAttentionType::Auto, + weight_quantization: crate::inference::skippy::StageWeightQuantization::Auto, error: None, shutdown_generation: 1, } @@ -7566,6 +7631,7 @@ fn test_stage_load_request() -> crate::inference::skippy::StageLoadRequest { n_gpu_layers: -1, mmap: None, mlock: false, + weight_quantization: crate::inference::skippy::StageWeightQuantization::Affine8, cache_type_k: "f16".to_string(), cache_type_v: "f16".to_string(), flash_attn_type: skippy_protocol::FlashAttentionType::Auto, @@ -7599,6 +7665,7 @@ fn test_preparation_status( stage_index: 1, layer_start: 12, layer_end: 24, + weight_quantization: crate::inference::skippy::StageWeightQuantization::Affine8, state, bytes_done: Some(1024), bytes_total: Some(4096), @@ -7619,6 +7686,7 @@ fn stage_control_inventory_request_round_trips_proto() { model_id: "model-a".to_string(), package_ref: "gguf:///model.gguf".to_string(), manifest_sha256: "direct-gguf:1:model.gguf".to_string(), + weight_quantization: crate::inference::skippy::StageWeightQuantization::Affine8, }, ); @@ -7632,6 +7700,10 @@ fn stage_control_inventory_request_round_trips_proto() { assert_eq!(inventory.model_id, "model-a"); assert_eq!(inventory.package_ref, "gguf:///model.gguf"); assert_eq!(inventory.manifest_sha256, "direct-gguf:1:model.gguf"); + assert_eq!( + inventory.weight_quantization, + crate::inference::skippy::StageWeightQuantization::Affine8 + ); } #[test] @@ -7687,6 +7759,10 @@ fn stage_control_status_update_request_round_trips_proto() { assert_eq!(status.bind_addr.as_deref(), Some("127.0.0.1:51234")); assert_eq!(status.bytes_done, Some(1024)); assert_eq!(status.bytes_total, Some(4096)); + assert_eq!( + status.weight_quantization, + crate::inference::skippy::StageWeightQuantization::Affine8 + ); } #[test] @@ -7712,6 +7788,7 @@ fn stage_control_inventory_response_round_trips_plain_gguf_source() { source_model_path: Some("/model.gguf".to_string()), source_model_bytes: Some(4_096), source_model_kind: crate::inference::skippy::SourceModelKind::PlainGguf, + weight_quantization: crate::inference::skippy::StageWeightQuantization::Affine8, }, ); @@ -7729,10 +7806,18 @@ fn stage_control_inventory_response_round_trips_plain_gguf_source() { assert_eq!(inventory.source_model_path.as_deref(), Some("/model.gguf")); assert_eq!(inventory.available_ranges[0].layer_start, 0); assert_eq!(inventory.available_ranges[0].layer_end, 32); + assert_eq!( + inventory.weight_quantization, + crate::inference::skippy::StageWeightQuantization::Affine8 + ); assert_eq!( inventory.preparing_ranges[0].state, crate::inference::skippy::StagePreparationState::Resolving ); + assert_eq!( + inventory.preparing_ranges[0].weight_quantization, + crate::inference::skippy::StageWeightQuantization::Affine8 + ); } #[test] diff --git a/crates/mesh-llm-host-runtime/src/runtime/local.rs b/crates/mesh-llm-host-runtime/src/runtime/local.rs index 3f22988675..9c5fd47fcc 100644 --- a/crates/mesh-llm-host-runtime/src/runtime/local.rs +++ b/crates/mesh-llm-host-runtime/src/runtime/local.rs @@ -1546,6 +1546,7 @@ fn split_runtime_stage_load_request( n_gpu_layers: resolved_config.n_gpu_layers, mmap: resolved_config.mmap, mlock: resolved_config.mlock, + weight_quantization: skippy::StageWeightQuantization::Auto, cache_type_k: resolved_config.cache_type_k.clone(), cache_type_v: resolved_config.cache_type_v.clone(), flash_attn_type: resolved_config.flash_attn_type, @@ -3172,6 +3173,7 @@ async fn split_peer_package_signal( model_id: model_ref.to_string(), package_ref: package.package_ref.clone(), manifest_sha256: package.manifest_sha256.clone(), + weight_quantization: skippy::StageWeightQuantization::Auto, }; let result = node .send_stage_control(peer_id, skippy::StageControlRequest::Inventory(request)) @@ -3542,6 +3544,9 @@ fn split_stage_source_is_ready( inventory: &skippy::StageLayerInventory, load: &skippy::StageLoadRequest, ) -> bool { + if inventory.weight_quantization != load.weight_quantization { + return false; + } let ready_running_stage = inventory .ready_ranges .iter() @@ -3557,6 +3562,7 @@ fn split_stage_source_is_ready( && status.backend == load.backend && status.package_ref == load.package_ref && status.manifest_sha256 == load.manifest_sha256 + && status.weight_quantization == load.weight_quantization && status.layer_start <= load.layer_start && status.layer_end >= load.layer_end && matches!( @@ -3590,6 +3596,7 @@ async fn query_stage_inventory( model_id: load.model_id.clone(), package_ref: load.package_ref.clone(), manifest_sha256: load.manifest_sha256.clone(), + weight_quantization: load.weight_quantization, }; let response = if stage_node_id == node.id() { node.send_local_stage_control(skippy::StageControlRequest::Inventory(request)) @@ -4058,6 +4065,7 @@ mod tests { n_gpu_layers: -1, mmap: None, mlock: false, + weight_quantization: skippy::StageWeightQuantization::Auto, cache_type_k: "f16".to_string(), cache_type_v: "f16".to_string(), flash_attn_type: FlashAttentionType::Auto, @@ -4295,6 +4303,7 @@ mod tests { n_batch: None, n_ubatch: None, flash_attn_type: FlashAttentionType::Auto, + weight_quantization: skippy::StageWeightQuantization::Auto, error: None, shutdown_generation: generation.generation, } @@ -4621,6 +4630,7 @@ max_tokens = 222 n_batch: load.n_batch, n_ubatch: load.n_ubatch, flash_attn_type: load.flash_attn_type, + weight_quantization: load.weight_quantization, error: None, shutdown_generation: load.shutdown_generation, coordinator_term: load.coordinator_term, @@ -4657,6 +4667,7 @@ max_tokens = 222 n_batch: None, n_ubatch: None, flash_attn_type: FlashAttentionType::Auto, + weight_quantization: skippy::StageWeightQuantization::Auto, error: None, shutdown_generation: stop.shutdown_generation, coordinator_term: stop.coordinator_term, @@ -4679,6 +4690,7 @@ max_tokens = 222 stage_index: load.stage_index, layer_start: load.layer_start, layer_end: load.layer_end, + weight_quantization: load.weight_quantization, state: skippy::StagePreparationState::Available, bytes_done: load.source_model_bytes, bytes_total: load.source_model_bytes, @@ -4709,6 +4721,7 @@ max_tokens = 222 source_model_path: Some("/models/qwen.gguf".to_string()), source_model_bytes: Some(40_000_000), source_model_kind: skippy::SourceModelKind::LayerPackage, + weight_quantization: request.weight_quantization, } } @@ -4825,6 +4838,7 @@ max_tokens = 222 source_model_path: None, source_model_bytes: None, source_model_kind: skippy::SourceModelKind::LayerPackage, + weight_quantization: skippy::StageWeightQuantization::Auto, }; let signal = split_inventory_package_signal(&inventory, &package); @@ -5037,6 +5051,7 @@ max_tokens = 222 ), source_model_bytes: Some(4_900_000_000), source_model_kind: skippy::SourceModelKind::LayerPackage, + weight_quantization: load.weight_quantization, }; assert!(!split_stage_source_is_ready(&inventory, &load)); @@ -5052,6 +5067,7 @@ max_tokens = 222 fn mlx_artifact_slice_accepts_exact_prepare_availability() { let mut load = stage_load_request(LoadMode::ArtifactSlice); load.backend = "mlx".to_string(); + load.weight_quantization = skippy::StageWeightQuantization::Affine8; load.package_ref = format!( "hf-model://HuggingFaceTB/SmolLM2-135M-Instruct@{}", "a".repeat(40) @@ -5068,10 +5084,20 @@ max_tokens = 222 source_model_path: None, source_model_bytes: None, source_model_kind: skippy::SourceModelKind::Unknown, + weight_quantization: skippy::StageWeightQuantization::Auto, }; assert!(!split_stage_source_is_ready(&inventory, &load)); + inventory.weight_quantization = load.weight_quantization; + + let mut wrong_profile = test_preparation_status_from_load(&load); + wrong_profile.weight_quantization = skippy::StageWeightQuantization::Affine4; + inventory.preparing_ranges.push(wrong_profile); + + assert!(!split_stage_source_is_ready(&inventory, &load)); + + inventory.preparing_ranges.clear(); inventory .preparing_ranges .push(test_preparation_status_from_load(&load)); @@ -5097,6 +5123,7 @@ max_tokens = 222 source_model_path: Some("/models/qwen.gguf".to_string()), source_model_bytes: Some(4_900_000_000), source_model_kind: skippy::SourceModelKind::PlainGguf, + weight_quantization: load.weight_quantization, }; assert!(split_stage_source_is_ready(&inventory, &load)); @@ -5121,6 +5148,7 @@ max_tokens = 222 source_model_path: None, source_model_bytes: None, source_model_kind: skippy::SourceModelKind::Unknown, + weight_quantization: skippy::StageWeightQuantization::Auto, }; let signal = split_inventory_package_signal(&inventory, &package); @@ -5154,6 +5182,7 @@ max_tokens = 222 source_model_path: None, source_model_bytes: None, source_model_kind: skippy::SourceModelKind::Unknown, + weight_quantization: skippy::StageWeightQuantization::Auto, }; assert_eq!( @@ -5184,6 +5213,7 @@ max_tokens = 222 source_model_path: Some("/cache/layer-package".to_string()), source_model_bytes: Some(1_000), source_model_kind: skippy::SourceModelKind::LayerPackage, + weight_quantization: skippy::StageWeightQuantization::Auto, }; inventory.manifest_sha256 = "other-manifest".to_string(); @@ -5218,6 +5248,7 @@ max_tokens = 222 source_model_path: Some("/cache/layer-package".to_string()), source_model_bytes: Some(1_000), source_model_kind: skippy::SourceModelKind::LayerPackage, + weight_quantization: skippy::StageWeightQuantization::Auto, }; assert_eq!( diff --git a/crates/model-hf/src/safetensors_stage/materialize.rs b/crates/model-hf/src/safetensors_stage/materialize.rs index bac73433ea..5aaa4b66ca 100644 --- a/crates/model-hf/src/safetensors_stage/materialize.rs +++ b/crates/model-hf/src/safetensors_stage/materialize.rs @@ -688,6 +688,30 @@ mod tests { assert!(fs::read_dir(cache_root).unwrap().next().is_none()); } + #[test] + fn cancelled_tensor_visit_stops_before_the_first_payload_request() { + let checkpoint = Arc::new(test_checkpoint()); + let requests = Arc::new(Mutex::new(Vec::new())); + let endpoint = start_checkpoint_server(checkpoint, Arc::clone(&requests)); + let cache = tempfile::tempdir().unwrap(); + let cache_root = cache.path().join("cache"); + let materializer = + SafetensorsStageMaterializer::new(cache_root.clone(), Some(&endpoint), None).unwrap(); + + let visit = materializer.prepare_tensor_visit(test_request()).unwrap(); + let metadata_request_count = requests.lock().unwrap().len(); + let error = visit + .visit_tensor_files_cancellable(|| true, |_| Ok(())) + .unwrap_err(); + + assert!(error.to_string().contains("tensor visit cancelled")); + assert_eq!(requests.lock().unwrap().len(), metadata_request_count); + assert!( + !cache_root.exists() || fs::read_dir(cache_root).unwrap().next().is_none(), + "cancelled visit retained an ephemeral tensor directory" + ); + } + #[test] fn serializes_concurrent_materialization_of_the_same_cache_key() { let checkpoint = Arc::new(test_checkpoint()); diff --git a/crates/model-hf/src/safetensors_stage/tensor_stream.rs b/crates/model-hf/src/safetensors_stage/tensor_stream.rs index 90de225444..fbbe01a1fb 100644 --- a/crates/model-hf/src/safetensors_stage/tensor_stream.rs +++ b/crates/model-hf/src/safetensors_stage/tensor_stream.rs @@ -69,10 +69,29 @@ impl SafetensorsStageTensorVisit<'_> { /// that reads the file before returning from the callback. This lets a /// backend quantize or transform exact HTTP ranges sequentially without /// retaining the complete BF16 stage artifact on disk. - pub fn visit_tensor_files(self, mut visitor: F) -> Result + pub fn visit_tensor_files(self, visitor: F) -> Result where F: FnMut(&SafetensorsStageTensorFile) -> Result<()>, { + self.visit_tensor_files_cancellable(|| false, visitor) + } + + /// Visits selected tensors while cooperatively checking for cancellation. + /// + /// Cancellation is checked before each payload request and immediately + /// before and after the visitor. An in-flight HTTP response or visitor call + /// is allowed to finish; its ephemeral file is removed when this method + /// returns. + pub fn visit_tensor_files_cancellable( + self, + mut is_cancelled: C, + mut visitor: F, + ) -> Result + where + F: FnMut(&SafetensorsStageTensorFile) -> Result<()>, + C: FnMut() -> bool, + { + ensure!(!is_cancelled(), "SafeTensors tensor visit cancelled"); let directory = EphemeralTensorDirectory::create(&self.materializer.cache_root)?; let source_shards = self .prepared @@ -89,6 +108,7 @@ impl SafetensorsStageTensorVisit<'_> { let mut visited_tensor_bytes = 0_u64; let mut temporary_file_peak_bytes = 0_u64; for (index, tensor) in tensors.iter().enumerate() { + ensure!(!is_cancelled(), "SafeTensors tensor visit cancelled"); let source = source_shards .get(tensor.source_file.as_str()) .with_context(|| format!("missing source identity for {}", tensor.source_file))?; @@ -99,7 +119,9 @@ impl SafetensorsStageTensorVisit<'_> { self.materializer .write_tensor_file(&path, &self.prepared.plan, tensor, source)?; temporary_file_peak_bytes = temporary_file_peak_bytes.max(file_bytes); + ensure!(!is_cancelled(), "SafeTensors tensor visit cancelled"); visitor(&tensor_file(tensor, path.clone(), file_bytes))?; + ensure!(!is_cancelled(), "SafeTensors tensor visit cancelled"); fs::remove_file(&path) .with_context(|| format!("remove ephemeral tensor file {}", path.display()))?; visited_tensor_bytes = visited_tensor_bytes diff --git a/crates/skippy-engine-mlx/STAGED_EXECUTION.md b/crates/skippy-engine-mlx/STAGED_EXECUTION.md index adc4cbfcfc..0808a64eda 100644 --- a/crates/skippy-engine-mlx/STAGED_EXECUTION.md +++ b/crates/skippy-engine-mlx/STAGED_EXECUTION.md @@ -22,8 +22,9 @@ the default automatic mesh launch path: configured layer range. - `mlx-stage` starts a stage process or drives a chain as a proof client. - `StagePrepare` / `StageLoad` with `backend=mlx` and an immutable - `hf-model://org/repo@` reference now materialize and start the same - engine through the normal host stage-control loop. + `hf-model://org/repo@` reference now derive or reuse a validated + quantized stage and start the same engine through the normal host + stage-control loop. No process in the proof has access to the complete checkpoint. The tokenizer and config files are small shared metadata; tensor data comes only from that @@ -121,8 +122,42 @@ before accepting a hit. On the same pinned layer-14 slice, the cold call made 9 tensor-payload range requests; the warm call returned the identical recipe and content hashes with `cache_hit=true`, made 0 tensor-payload range requests, and used 17,809,408 B max RSS. It still re-plans lightweight config/index/header -metadata to reconstruct the strong recipe key. Host lifecycle/eviction wiring -is not yet connected to this library/CLI cache. +metadata to reconstruct the strong recipe key. + +The host control path now consumes this cache directly. `StagePrepare` maps the +load request to a derivation recipe and builds or validates it on a blocking +worker; `StageLoad` validates the same entry and loads MLX from the derived +directory. It fails on a cache miss instead of downloading or quantizing tensor +payloads during Load. The load request carries an additive quantization profile: +`auto`, affine 4-bit, affine 8-bit, or MXFP4. An absent profile from an older +peer means `auto`; an unknown value fails closed. On the current Apple Metal +backend, `auto` selects affine 4-bit with group size 64. The chosen profile is +part of the recipe identity and is carried through inventory, preparation, and +running status. Inventory responses echo the requested profile, so one profile +cannot satisfy readiness for another, including across mixed-version peers. + +The host's claimed checkpoint identity is verified from the lightweight +metadata plan before the first tensor payload request. Prepare cancellation is +also threaded into cache-lock waits and the sequential visitor. It is checked +before every payload request and before and after each quantization callback; +an HTTP transfer or MLX operation already in flight finishes cooperatively +before its temporary file is removed. + +A clean host-control run built both halves without retaining a dense stage: + +| Layers | Exact source tensor bytes | Derived artifact | Payload requests | +| --- | ---: | ---: | ---: | +| `0..15` | 162,825,984 B | 45,859,713 B | 136 | +| `15..30` | 162,827,136 B | 45,861,308 B | 137 | + +It completed Prepare, Load, Start, generation, and Stop in 120.74 seconds and +produced the established affine-4 token reference. An immediate identical run +hit both validated entries, performed the same lifecycle in 8.87 seconds, and +used 258,162,688 B max RSS. `MESH_MLX_DERIVED_CACHE_DIR` can isolate or relocate +the host cache for testing and operations. Cache capacity and eviction still +need an owner; warm lookup also still probes lightweight upstream metadata to +reconstruct the strong recipe, then streams each cached shard once to verify +both its shard hash and the aggregate content digest. The two partial files are the exact-range artifacts described in `../../spikes/mlx-safetensors-stages/FINDINGS.md`. Tied input/output embeddings @@ -214,8 +249,10 @@ just mlx-stage prove --connect 127.0.0.1:19090 --wire-dtype f16 - Mesh topology planning does not yet produce MLX stage assignments. The host can consume explicit `backend=mlx` Prepare/Load requests, but automatic placement, capability advertisement, coordinator model planning, and an - OpenAI stage-0 frontend remain. Explicit host requests still use source - precision; quantization selection currently exists only in the engine config - and `mlx-stage` proof/derive CLI. The explicit derived artifact is not yet an - automatic cache hit path. There is no mesh protocol or Skippy ABI break in - the explicit consumer path. + OpenAI stage-0 frontend remain. Explicit host requests now derive and reuse + quantized artifacts, but cache eviction and an optional local + request-to-recipe locator remain. The quantization field is an additive mesh + protocol change; old peers omit it and therefore mean `auto`, while unknown + values fail closed on new peers. Automatic placement must capability-gate + explicit non-default profiles before mixed-version deployment. No Skippy ABI + changed. diff --git a/crates/skippy-engine-mlx/src/bin/mlx-stage.rs b/crates/skippy-engine-mlx/src/bin/mlx-stage.rs index a88b0f4102..ce84bb0153 100644 --- a/crates/skippy-engine-mlx/src/bin/mlx-stage.rs +++ b/crates/skippy-engine-mlx/src/bin/mlx-stage.rs @@ -13,9 +13,9 @@ mod real { use clap::{Parser, Subcommand, ValueEnum}; use model_hf::safetensors_stage::{SafetensorsStageMaterializer, SafetensorsStageRequest}; use skippy_engine_mlx::{ - MlxComputeDtype, MlxDerivedStageCacheConfig, MlxDerivedStageConfig, MlxStageEngine, - MlxStageEngineConfig, MlxWeightQuantization, derive_quantized_stage, - derive_quantized_stage_cached, + MlxComputeDtype, MlxDerivationControl, MlxDerivedStageCacheConfig, MlxDerivedStageConfig, + MlxStageEngine, MlxStageEngineConfig, MlxWeightQuantization, derive_quantized_stage, + derive_quantized_stage_cached, mlx_derived_stage_cache_root, }; use skippy_protocol::binary::{ StageStateHeader, StageWireMessage, WireActivationDType, WireMessageKind, WireReplyKind, @@ -268,6 +268,7 @@ mod real { source, output_dir, quantization, + control: MlxDerivationControl::default(), shard_size_bytes, }, )?; @@ -284,8 +285,7 @@ mod real { let shard_size_bytes = shard_size_mib .checked_mul(1024 * 1024) .context("derived shard size overflow")?; - let cache_root = cache_root - .unwrap_or_else(|| model_hf::store::mesh_llm_cache_dir().join("mlx-derived-stages")); + let cache_root = cache_root.unwrap_or_else(mlx_derived_stage_cache_root); let materializer = SafetensorsStageMaterializer::from_environment()?; let result = derive_quantized_stage_cached( &materializer, @@ -293,6 +293,7 @@ mod real { source, cache_root, quantization, + control: MlxDerivationControl::default(), shard_size_bytes, }, )?; diff --git a/crates/skippy-engine-mlx/src/derived.rs b/crates/skippy-engine-mlx/src/derived.rs index 9e69185ca7..9df256443a 100644 --- a/crates/skippy-engine-mlx/src/derived.rs +++ b/crates/skippy-engine-mlx/src/derived.rs @@ -6,7 +6,10 @@ use std::{ fs::{self, File}, io::{BufReader, Read}, path::{Path, PathBuf}, - sync::atomic::{AtomicU64, Ordering}, + sync::{ + Arc, + atomic::{AtomicBool, AtomicU64, Ordering}, + }, }; use anyhow::{Context, Result, ensure}; @@ -30,6 +33,7 @@ mod cache; pub use cache::{ MlxDerivedStageCacheConfig, MlxDerivedStageCacheResult, derive_quantized_stage_cached, + load_prepared_quantized_stage, mlx_derived_stage_cache_root, }; pub(super) const DERIVED_STAGE_SCHEMA_VERSION: u32 = 1; @@ -39,12 +43,53 @@ const PLAN_FILE: &str = "stage-plan.json"; pub(super) const REPORT_FILE: &str = "derived-stage.json"; static DERIVED_SEQUENCE: AtomicU64 = AtomicU64::new(0); +/// Expected source identity and cooperative cancellation for a derivation. +#[derive(Clone, Debug, Default)] +pub struct MlxDerivationControl { + expected_checkpoint_sha256: Option, + cancelled: Option>, +} + +impl MlxDerivationControl { + pub fn new( + expected_checkpoint_sha256: Option, + cancelled: Option>, + ) -> Self { + Self { + expected_checkpoint_sha256, + cancelled, + } + } + + fn is_cancelled(&self) -> bool { + self.cancelled + .as_ref() + .is_some_and(|cancelled| cancelled.load(Ordering::Acquire)) + } + + fn ensure_active(&self) -> Result<()> { + ensure!(!self.is_cancelled(), "MLX stage derivation cancelled"); + Ok(()) + } + + fn verify_checkpoint(&self, actual: &str) -> Result<()> { + if let Some(expected) = &self.expected_checkpoint_sha256 { + ensure!( + actual == expected, + "MLX checkpoint identity {actual} does not match stage claim {expected}" + ); + } + Ok(()) + } +} + /// Configuration for producing one MLX-quantized partial stage. #[derive(Clone, Debug)] pub struct MlxDerivedStageConfig { pub source: SafetensorsStageRequest, pub output_dir: PathBuf, pub quantization: MlxWeightQuantization, + pub control: MlxDerivationControl, /// Soft output bundle target. A single packed tensor may exceed this size. pub shard_size_bytes: usize, } @@ -174,6 +219,7 @@ pub fn derive_quantized_stage( materializer: &SafetensorsStageMaterializer, config: &MlxDerivedStageConfig, ) -> Result { + config.control.ensure_active()?; ensure!( config.shard_size_bytes > 0, "derived shard size must be non-zero" @@ -185,6 +231,10 @@ pub fn derive_quantized_stage( config.output_dir.display() ); let visit = materializer.prepare_tensor_visit(config.source.clone())?; + config.control.ensure_active()?; + config + .control + .verify_checkpoint(visit.checkpoint_sha256())?; let plan = visit.plan().clone(); let source_config = visit.config().to_vec(); ensure_dense_source_config(&source_config)?; @@ -211,23 +261,30 @@ pub fn derive_quantized_stage( let quantization_stream = Stream::new_with_device(&Device::new(DeviceType::Gpu, 0)); let initial_output_file_bytes = directory_file_bytes(temporary.path())?; let mut state = BuildState::new(initial_output_file_bytes); - let visit_report = visit.visit_tensor_files(|tensor| { - state.observe_source_file(tensor.file_bytes); - let arrays = convert_tensor( - tensor, - quantization, - &weights_stream, - &quantization_stream, - &mut state, - )?; - append_arrays( - arrays, - config.shard_size_bytes, - temporary.path(), - tensor.file_bytes, - &mut state, - ) - })?; + let visit_report = visit.visit_tensor_files_cancellable( + || config.control.is_cancelled(), + |tensor| { + config.control.ensure_active()?; + state.observe_source_file(tensor.file_bytes); + let arrays = convert_tensor( + tensor, + quantization, + &weights_stream, + &quantization_stream, + &mut state, + )?; + config.control.ensure_active()?; + append_arrays( + arrays, + config.shard_size_bytes, + temporary.path(), + tensor.file_bytes, + &mut state, + )?; + config.control.ensure_active() + }, + )?; + config.control.ensure_active()?; if !state.pending.tensors.is_empty() { flush_shard(temporary.path(), &mut state)?; } @@ -271,7 +328,9 @@ pub fn derive_quantized_stage( mlx_peak_memory_bytes: peak_memory()?, shards, }; + config.control.ensure_active()?; write_json(temporary.path().join(REPORT_FILE), &report)?; + config.control.ensure_active()?; temporary.publish(&config.output_dir)?; Ok(report) } @@ -547,8 +606,12 @@ pub(super) fn prepare_derivation_recipe( source: SafetensorsStageRequest, quantization: MlxWeightQuantization, shard_size_bytes: usize, + control: &MlxDerivationControl, ) -> Result { + control.ensure_active()?; let visit = materializer.prepare_tensor_visit(source)?; + control.ensure_active()?; + control.verify_checkpoint(visit.checkpoint_sha256())?; ensure_dense_source_config(visit.config())?; let quantization = serde_json::to_value(quantization.safemlx()?)?; let plan_bytes = serde_json::to_vec(visit.plan())?; @@ -813,6 +876,22 @@ mod tests { use super::*; + #[test] + fn derivation_control_rejects_cancellation_and_wrong_checkpoint() { + let cancelled = Arc::new(AtomicBool::new(false)); + let control = MlxDerivationControl::new( + Some("expected-checkpoint".to_string()), + Some(Arc::clone(&cancelled)), + ); + + assert!(control.ensure_active().is_ok()); + assert!(control.verify_checkpoint("expected-checkpoint").is_ok()); + assert!(control.verify_checkpoint("other-checkpoint").is_err()); + + cancelled.store(true, Ordering::Release); + assert!(control.ensure_active().is_err()); + } + #[test] fn quantized_config_preserves_source_and_adds_both_metadata_keys() { let directory = tempfile::tempdir().unwrap(); diff --git a/crates/skippy-engine-mlx/src/derived/cache.rs b/crates/skippy-engine-mlx/src/derived/cache.rs index c2fa4ceb80..8c904ee52f 100644 --- a/crates/skippy-engine-mlx/src/derived/cache.rs +++ b/crates/skippy-engine-mlx/src/derived/cache.rs @@ -1,27 +1,43 @@ //! Identity-bound reusable cache for derived MLX stage artifacts. use std::{ - fs, + collections::BTreeMap, + fs::{self, File}, + io::{BufReader, Read}, path::{Component, Path, PathBuf}, + thread, + time::Duration, }; -use anyhow::{Context, Result, ensure}; +use anyhow::{Context, Result, bail, ensure}; use model_hf::safetensors_stage::{SafetensorsStageMaterializer, SafetensorsStageRequest}; use serde::Serialize; +use sha2::{Digest, Sha256}; use super::{ - DERIVED_STAGE_SCHEMA_VERSION, MlxDerivedStageConfig, MlxDerivedStageReport, REPORT_FILE, - artifact_file_bytes, derive_quantized_stage, open_locked, output_content_sha256, - prepare_derivation_recipe, sha256_file, + DERIVED_STAGE_SCHEMA_VERSION, MlxDerivationControl, MlxDerivedStageConfig, + MlxDerivedStageReport, REPORT_FILE, artifact_file_bytes, derive_quantized_stage, open_locked, + prepare_derivation_recipe, }; use crate::stage::MlxWeightQuantization; +const DERIVED_CACHE_ROOT_ENV: &str = "MESH_MLX_DERIVED_CACHE_DIR"; + +/// Returns the managed MLX derived-stage cache root for this process. +pub fn mlx_derived_stage_cache_root() -> PathBuf { + std::env::var_os(DERIVED_CACHE_ROOT_ENV).map_or_else( + || model_hf::store::mesh_llm_cache_dir().join("mlx-derived-stages"), + PathBuf::from, + ) +} + /// Configuration for an identity-bound, reusable derived-stage cache entry. #[derive(Clone, Debug)] pub struct MlxDerivedStageCacheConfig { pub source: SafetensorsStageRequest, pub cache_root: PathBuf, pub quantization: MlxWeightQuantization, + pub control: MlxDerivationControl, /// Soft output bundle target. A single packed tensor may exceed this size. pub shard_size_bytes: usize, } @@ -41,6 +57,26 @@ pub fn derive_quantized_stage_cached( materializer: &SafetensorsStageMaterializer, config: &MlxDerivedStageCacheConfig, ) -> Result { + resolve_quantized_stage_cache(materializer, config, true) +} + +/// Loads a verified stage that was already produced by the prepare lifecycle. +/// +/// This performs metadata planning and full cache validation, but never fetches +/// tensor payloads or derives a replacement on a miss. +pub fn load_prepared_quantized_stage( + materializer: &SafetensorsStageMaterializer, + config: &MlxDerivedStageCacheConfig, +) -> Result { + resolve_quantized_stage_cache(materializer, config, false) +} + +fn resolve_quantized_stage_cache( + materializer: &SafetensorsStageMaterializer, + config: &MlxDerivedStageCacheConfig, + build_on_miss: bool, +) -> Result { + config.control.ensure_active()?; ensure!( config.shard_size_bytes > 0, "derived shard size must be non-zero" @@ -50,6 +86,7 @@ pub fn derive_quantized_stage_cached( config.source.clone(), config.quantization, config.shard_size_bytes, + &config.control, )?; fs::create_dir_all(&config.cache_root).with_context(|| { format!( @@ -61,9 +98,11 @@ pub fn derive_quantized_stage_cached( let lock_path = config.cache_root.join(format!(".{recipe}.lock")); // Keep this pathname stable across invocations. Removing an advisory-lock // file after unlock can split waiters between the unlinked and new inodes. - let _lock = open_locked(&lock_path, false)?.expect("blocking cache lock is acquired"); + let _lock = open_cache_lock(&lock_path, &config.control)?; + config.control.ensure_active()?; - if let Some(report) = load_cached(&output_dir, &recipe)? { + if let Some(report) = load_cached(&output_dir, &recipe, &config.control)? { + config.control.ensure_active()?; return Ok(MlxDerivedStageCacheResult { cache_hit: true, source_range_request_count: 0, @@ -71,6 +110,12 @@ pub fn derive_quantized_stage_cached( report, }); } + if !build_on_miss { + bail!( + "prepared MLX derived stage cache entry {recipe} is missing or invalid; run StagePrepare before StageLoad" + ); + } + config.control.ensure_active()?; remove_invalid_cache_entry(&output_dir)?; let report = derive_quantized_stage( materializer, @@ -78,6 +123,7 @@ pub fn derive_quantized_stage_cached( source: config.source.clone(), output_dir: output_dir.clone(), quantization: config.quantization, + control: config.control.clone(), shard_size_bytes: config.shard_size_bytes, }, )?; @@ -89,7 +135,21 @@ pub fn derive_quantized_stage_cached( }) } -fn load_cached(output_dir: &Path, recipe: &str) -> Result> { +fn open_cache_lock(path: &Path, control: &MlxDerivationControl) -> Result { + loop { + control.ensure_active()?; + if let Some(lock) = open_locked(path, true)? { + return Ok(lock); + } + thread::sleep(Duration::from_millis(25)); + } +} + +fn load_cached( + output_dir: &Path, + recipe: &str, + control: &MlxDerivationControl, +) -> Result> { if !output_dir.is_dir() { return Ok(None); } @@ -107,8 +167,7 @@ fn load_cached(output_dir: &Path, recipe: &str) -> Result Result { Ok(true) } -fn shards_match(output_dir: &Path, report: &MlxDerivedStageReport) -> Result { +fn validate_cache_files( + output_dir: &Path, + report: &MlxDerivedStageReport, + control: &MlxDerivationControl, +) -> Result { + let mut expected_shards = BTreeMap::new(); for shard in &report.shards { let relative = Path::new(&shard.file); if relative.components().count() != 1 @@ -135,16 +199,64 @@ fn shards_match(output_dir: &Path, report: &MlxDerivedStageReport) -> Result metadata, - _ => return Ok(false), - }; - if metadata.len() != shard.file_bytes || sha256_file(&path)? != shard.sha256 { + if expected_shards.insert(shard.file.as_str(), shard).is_some() { + return Ok(false); + } + } + if expected_shards.is_empty() { + return Ok(false); + } + + let mut files = fs::read_dir(output_dir)? + .filter_map(|entry| match entry { + Ok(entry) if entry.file_name() != REPORT_FILE => Some(Ok(entry)), + Ok(_) => None, + Err(error) => Some(Err(error)), + }) + .collect::>>()?; + files.sort_by_key(|entry| entry.file_name()); + + let mut aggregate = Sha256::new(); + aggregate.update(b"mesh-mlx-derived-output-v1"); + for entry in files { + control.ensure_active()?; + if !entry.file_type()?.is_file() { + return Ok(false); + } + let name = entry.file_name(); + let name = name.to_string_lossy(); + let file_bytes = entry.metadata()?.len(); + aggregate.update(u64::try_from(name.len())?.to_le_bytes()); + aggregate.update(name.as_bytes()); + aggregate.update(file_bytes.to_le_bytes()); + + let expected_shard = expected_shards.remove(name.as_ref()); + if expected_shard.is_some_and(|shard| shard.file_bytes != file_bytes) { + return Ok(false); + } + let mut shard_hasher = expected_shard.map(|_| Sha256::new()); + let mut reader = BufReader::new(File::open(entry.path())?); + let mut buffer = vec![0_u8; 1024 * 1024]; + loop { + control.ensure_active()?; + let read = reader.read(&mut buffer)?; + if read == 0 { + break; + } + aggregate.update(&buffer[..read]); + if let Some(hasher) = &mut shard_hasher { + hasher.update(&buffer[..read]); + } + } + if let (Some(expected), Some(hasher)) = (expected_shard, shard_hasher) + && format!("{:x}", hasher.finalize()) != expected.sha256 + { return Ok(false); } } - Ok(!report.shards.is_empty()) + control.ensure_active()?; + Ok(expected_shards.is_empty() + && format!("{:x}", aggregate.finalize()) == report.output_content_sha256) } fn remove_invalid_cache_entry(path: &Path) -> Result<()> { @@ -160,10 +272,33 @@ fn remove_invalid_cache_entry(path: &Path) -> Result<()> { #[cfg(test)] mod tests { + use std::sync::{ + Arc, + atomic::{AtomicBool, Ordering}, + }; + use serde_json::json; use super::*; - use crate::derived::{MlxDerivedStageShard, write_json}; + use crate::derived::{MlxDerivedStageShard, output_content_sha256, sha256_file, write_json}; + + #[test] + fn cancelled_cache_waiter_stops_without_acquiring_the_recipe_lock() { + let directory = tempfile::tempdir().unwrap(); + let lock_path = directory.path().join("recipe.lock"); + let held = open_locked(&lock_path, false).unwrap().unwrap(); + let cancelled = Arc::new(AtomicBool::new(false)); + let control = MlxDerivationControl::new(None, Some(Arc::clone(&cancelled))); + let waiter_path = lock_path.clone(); + let waiter = std::thread::spawn(move || open_cache_lock(&waiter_path, &control)); + + std::thread::sleep(Duration::from_millis(75)); + cancelled.store(true, Ordering::Release); + let error = waiter.join().unwrap().unwrap_err(); + + assert!(error.to_string().contains("derivation cancelled")); + drop(held); + } fn write_test_cache(output_dir: &Path, recipe: &str) -> MlxDerivedStageReport { fs::create_dir_all(output_dir).unwrap(); @@ -213,11 +348,17 @@ mod tests { let output = directory.path().join("recipe"); let expected = write_test_cache(&output, "recipe"); - let cached = load_cached(&output, "recipe").unwrap().unwrap(); + let cached = load_cached(&output, "recipe", &MlxDerivationControl::default()) + .unwrap() + .unwrap(); assert_eq!(cached.output_content_sha256, expected.output_content_sha256); fs::write(output.join("model.safetensors"), b"broken").unwrap(); - assert!(load_cached(&output, "recipe").unwrap().is_none()); + assert!( + load_cached(&output, "recipe", &MlxDerivationControl::default()) + .unwrap() + .is_none() + ); } #[test] @@ -226,7 +367,7 @@ mod tests { let output = directory.path().join("recipe"); let mut report = write_test_cache(&output, "recipe"); report.shards[0].file = "../model.safetensors".to_string(); - assert!(!shards_match(&output, &report).unwrap()); + assert!(!validate_cache_files(&output, &report, &MlxDerivationControl::default()).unwrap()); } #[test] @@ -239,7 +380,9 @@ mod tests { fs::rename(&original_root, &relocated_root).unwrap(); let relocated = relocated_root.join("recipe"); - let cached = load_cached(&relocated, "recipe").unwrap().unwrap(); + let cached = load_cached(&relocated, "recipe", &MlxDerivationControl::default()) + .unwrap() + .unwrap(); assert_eq!(cached.output_dir, relocated); } @@ -254,6 +397,10 @@ mod tests { write_test_cache(&output, "recipe"); symlink("config.json", output.join("unexpected-index.json")).unwrap(); - assert!(load_cached(&output, "recipe").unwrap().is_none()); + assert!( + load_cached(&output, "recipe", &MlxDerivationControl::default()) + .unwrap() + .is_none() + ); } } diff --git a/crates/skippy-engine-mlx/src/lib.rs b/crates/skippy-engine-mlx/src/lib.rs index 5d8cdccb07..0a5bec45c7 100644 --- a/crates/skippy-engine-mlx/src/lib.rs +++ b/crates/skippy-engine-mlx/src/lib.rs @@ -20,9 +20,9 @@ mod stage; pub use backend::MlxBackend; #[cfg(all(feature = "mlx", target_os = "macos"))] pub use derived::{ - MlxDerivedStageCacheConfig, MlxDerivedStageCacheResult, MlxDerivedStageConfig, - MlxDerivedStageReport, MlxDerivedStageShard, derive_quantized_stage, - derive_quantized_stage_cached, + MlxDerivationControl, MlxDerivedStageCacheConfig, MlxDerivedStageCacheResult, + MlxDerivedStageConfig, MlxDerivedStageReport, MlxDerivedStageShard, derive_quantized_stage, + derive_quantized_stage_cached, load_prepared_quantized_stage, mlx_derived_stage_cache_root, }; #[cfg(all(feature = "mlx", target_os = "macos"))] pub use engine::{ChatTurn, GenerateRequest, MlxEngine, MlxEngineConfig}; diff --git a/crates/skippy-protocol/proto/stage.proto b/crates/skippy-protocol/proto/stage.proto index 1481e2cd91..8bc4ff77a9 100644 --- a/crates/skippy-protocol/proto/stage.proto +++ b/crates/skippy-protocol/proto/stage.proto @@ -94,6 +94,7 @@ message LoadStage { optional bool native_mtp_enabled = 33; optional bool mmap = 34; optional bool mlock = 35; + StageWeightQuantization weight_quantization = 36; } message StopStage { @@ -114,6 +115,7 @@ message GetLayerInventory { string model_id = 1; string package_ref = 2; string manifest_sha256 = 3; + StageWeightQuantization weight_quantization = 4; } message PrepareStage { @@ -200,6 +202,7 @@ message StageStatus { uint64 coordinator_term = 29; optional string coordinator_id = 30; uint64 lease_until_unix_ms = 31; + StageWeightQuantization weight_quantization = 32; } message StageStatusList { @@ -218,6 +221,7 @@ message LayerInventory { optional string source_model_path = 9; optional uint64 source_model_bytes = 10; SourceModelKind source_model_kind = 11; + StageWeightQuantization weight_quantization = 12; } message LayerRange { @@ -245,6 +249,7 @@ message StagePreparationStatus { uint64 coordinator_term = 17; optional string coordinator_id = 18; uint64 lease_until_unix_ms = 19; + StageWeightQuantization weight_quantization = 20; } enum SourceModelKind { @@ -297,6 +302,14 @@ enum StageWireDType { STAGE_WIRE_DTYPE_Q8 = 3; } +enum StageWeightQuantization { + STAGE_WEIGHT_QUANTIZATION_UNSPECIFIED = 0; + STAGE_WEIGHT_QUANTIZATION_AUTO = 1; + STAGE_WEIGHT_QUANTIZATION_AFFINE4 = 2; + STAGE_WEIGHT_QUANTIZATION_AFFINE8 = 3; + STAGE_WEIGHT_QUANTIZATION_MXFP4 = 4; +} + enum StageRuntimeState { STAGE_RUNTIME_STATE_UNSPECIFIED = 0; STAGE_RUNTIME_STATE_STARTING = 1; diff --git a/crates/skippy-protocol/src/lib.rs b/crates/skippy-protocol/src/lib.rs index b2aeec0ba2..3053ea0718 100644 --- a/crates/skippy-protocol/src/lib.rs +++ b/crates/skippy-protocol/src/lib.rs @@ -733,6 +733,7 @@ mod tests { model_id: "qwen".to_string(), package_ref: "hf://repo/model".to_string(), manifest_sha256: "a5".repeat(32), + weight_quantization: 0, }, )), ..frame.clone() diff --git a/docs/design/MLX_STAGE_ENGINE_PLAN.md b/docs/design/MLX_STAGE_ENGINE_PLAN.md index f4aa2d43d5..a4a01f8f43 100644 --- a/docs/design/MLX_STAGE_ENGINE_PLAN.md +++ b/docs/design/MLX_STAGE_ENGINE_PLAN.md @@ -52,17 +52,17 @@ range-to-derived-cache quantization remain. See **Update — host `StagePrepare` / `StageLoad` now consumes range-only MLX stages.** An immutable `hf-model://org/repo@` request is validated -before network work, materialized on a blocking worker, checked against a -topology-wide checkpoint identity, strict-loaded into `MlxStageEngine`, and -served through the existing Skippy binary wire. A clean-cache SmolLM2 proof ran -both 15-layer ranges through `spawn_stage_control_loop`, reproduced the same -eight reference tokens, and stopped both stages. Each node-side range plan read -about 162.86 MB of a 269.06 MB source shard and avoided about 106.2 MB. Startup -now fails closed on bind/topology/downstream errors, and Stop closes and joins -active connections. Automatic MLX topology production, capability -advertisement, remote two-node proof, host quantization selection, and bounded -range-to-derived-cache materialization remain; this checkpoint proves the host -consumer path, not automatic placement. +before network work, checked against a topology-wide checkpoint identity, +derived into or reused from the recipe-keyed quantized cache on a blocking +worker, strict-loaded into `MlxStageEngine`, and served through the existing +Skippy binary wire. `StageLoad` requires that prepared entry and cannot derive +or download tensor payloads on a miss. A clean-cache SmolLM2 proof ran both +15-layer ranges through +`spawn_stage_control_loop`, reproduced the affine-4 eight-token reference, and +stopped both stages without retaining dense stage artifacts. Automatic MLX +topology production, capability advertisement, remote two-node proof, and +cache eviction remain; this checkpoint proves the host consumer path, not +automatic placement. **Update — partial MLX stages now JIT-quantize one tensor at a time.** The pinned safemlx strict loader already contains the required bounded lazy-graph @@ -71,8 +71,8 @@ repeat. A whole-model affine-4 reference and two independently quantized SmolLM2 stages generated the same eight tokens over F16 residuals. Post-proof RSS was about 85.3 and 85.9 MiB per stage. This bounds the lazy graph, not physical copies: TensorView conversion, stream copies, mmap pages, and MLX -scratch all require explicit high-water measurement. The host does not select -this profile yet, and the BF16 partial stage still exists on disk first. +scratch all require explicit high-water measurement. The newer direct-derived +host path below removes the intermediate BF16 stage. **Update — exact ranges can now be consumed sequentially without a BF16 stage artifact.** `model-hf` exposes each selected tensor as an ephemeral, valid @@ -105,7 +105,32 @@ builders with an advisory lock, and validates the output-content plus per-shard hashes on hits. A cold pinned layer-14 run made 9 tensor payload requests; the warm run made 0, skipped quantization, and used about 17.8 MB max RSS. The warm path still fetches lightweight config/index/header metadata to reconstruct the -strong recipe. Host prepare/load integration and eviction ownership remain. +strong recipe. + +**Update — the host now owns the direct-derived load path and carries the +quantization choice.** The additive stage-load field supports `auto`, affine +4-bit, affine 8-bit, and MXFP4. Missing values from older peers map to `auto`; +unknown values fail closed. Apple Metal currently maps `auto` to affine +4-bit/group-64, and the resolved recipe remains identity-bound. A cold two-half +host lifecycle fetched 162,825,984 and 162,827,136 source tensor bytes and wrote +45,859,713 and 45,861,308-byte derived artifacts. It passed in 120.74 seconds; +the immediate validated-cache run passed in 8.87 seconds with 258,162,688 B max +RSS. `MESH_MLX_DERIVED_CACHE_DIR` overrides the host cache root. Eviction, +automatic topology selection, and the remote two-node proof remain. + +Prepare is the only lifecycle operation allowed to build a missing entry; +Load validates and consumes the prepared artifact or fails closed. + +The host additionally verifies the topology checkpoint claim from metadata +before any tensor payload is fetched. Prepare cancellation reaches recipe-lock +waits and the range visitor, with checks between payload requests and around +each quantization callback. In-flight HTTP or MLX work is cooperative rather +than preemptive. Inventory responses echo the requested profile, and +preparation plus running status carry it, preventing (for example) an affine-4 +preparation from satisfying an affine-8 load. Old peers omit these additive +fields and therefore mean `auto`; automatic non-default placement must require +a peer that advertises the new semantics. Cache-hit validation streams each +shard once while checking both its shard hash and the aggregate content digest. The pinned safemlx revision also already includes whole-model Inkling text, vision, and audio execution. Earlier notes that called for porting Inkling were @@ -728,10 +753,11 @@ single-machine parity first, then two Macs over the real network. > `StageWireMessage` boundaries matched unsplit MLX with zero measured dense > logit delta; two real F16-wire processes also matched the whole-model > affine-4 token reference after tensor-wise on-load quantization. Explicit host -> stage control now starts the source-precision path. A sequential exact-range -> tensor visitor now removes the need to create a BF16 stage artifact in the -> next cache builder. Direct range-to-quantized-cache materialization, -> profile-bearing topology requests, and remote two-node execution remain. +> stage control now derives or reuses a quantized artifact directly from exact +> tensor ranges. The additive load request carries the quantization profile, +> and clean plus warm-cache host lifecycles reproduce the affine-4 reference. +> Automatic profile-bearing MLX topology production, cache eviction, and remote +> two-node execution remain. **Phase 4 — KV/state codec + verify + trim/checkpoint.** Implement the engine-general cache codec (§5.2), `verify_tokens_frame` for speculative decode, @@ -751,13 +777,13 @@ Apple-Silicon nodes. llama.cpp remains the cross-platform default. ## 8. Spike gates (go/no-go before Phase 3) -1. **Partial-loading proof (DENSE GO, QUANT PARTIAL):** remote exact-range +1. **Partial-loading proof (DENSE + LLAMA QUANT GO, FRONTIER PARTIAL):** remote exact-range selection is proven, including on 1.9 TB Inkling BF16. SmolLM2 partial files - were materialized and loaded without a complete checkpoint. Still required: - The live-model loader now quantizes/evaluates one tensor at a time and split - output matches a whole quantized reference. Still required: measure cold-load - high-water RSS and stream HTTP ranges directly into a derived quantized cache - without retaining the whole BF16 stage slice on disk. + were materialized and loaded without a complete checkpoint. The live-model + loader quantizes/evaluates one tensor at a time, and the host now streams + exact ranges into a derived quantized cache without retaining the BF16 stage + slice. Still required: bounded expert-bank/slab transforms and high-water + evidence for Inkling and Nemotron-family frontier tensors. 2. **Boundary latency breakdown (GO/NO-GO):** measure layer compute, cast, contiguous, **eval fence**, host readback, serialize, and receive-reconstruct **independently**, at hidden widths 4096/8192/16384 and token counts diff --git a/spikes/mlx-safetensors-stages/FINDINGS.md b/spikes/mlx-safetensors-stages/FINDINGS.md index 7f3dd1d147..ee6f27d4a1 100644 --- a/spikes/mlx-safetensors-stages/FINDINGS.md +++ b/spikes/mlx-safetensors-stages/FINDINGS.md @@ -28,16 +28,15 @@ close to exact. For Inkling, tensors are heavily interleaved across source shards, so exact tensor ranges are mandatory: whole-shard selection would turn a 109.84 GiB four-layer stage into a 942.99 GiB download. -The small dense-model path is now proven through execution, including -tensor-at-a-time affine-4 quantization into the live MLX model. `model-hf` now -also exposes a backend-neutral sequential visitor that downloads each selected -range as an ephemeral, valid one-tensor SafeTensors file and deletes it before -fetching the next tensor. The remaining artifact proof is consuming that seam -to build bounded quantized-cache shards with measured peak RSS and disk use; -the serving path still first retains the complete BF16 stage slice on disk. -Its prepared session exposes the verified config, config hash, checkpoint -identity, and range plan before payload callbacks. On macOS/Unix, advisory -locks also scavenge crash-abandoned visits without removing concurrent ones. +The small dense-model path is now proven through execution, including a host +path that streams selected tensor ranges into bounded affine-4 cache shards, +then loads the derived stage without retaining a complete BF16 stage slice. +`model-hf` exposes the backend-neutral sequential visitor that downloads each +selected range as an ephemeral, valid one-tensor SafeTensors file and deletes it +before fetching the next tensor. Its prepared session exposes the verified +config, config hash, checkpoint identity, and range plan before payload +callbacks. On macOS/Unix, advisory locks also scavenge crash-abandoned visits +without removing concurrent ones. The pinned SmolLM2 layer-14 visitor proof fetched 9 tensors totaling 7,080,192 bytes from a 269,060,552-byte source shard. Its largest temporary one-tensor @@ -171,8 +170,33 @@ output-content digest, and every shard hash before accepting a hit. On the same pinned layer-14 slice, a cold call made 9 tensor-payload requests and the warm call made 0, returned the identical recipe/content/shard hashes, and used 17,809,408 B max RSS. The warm path still reads lightweight config/index/header -metadata to reconstruct the strong key. This is the reusable library/CLI seam; -host lifecycle and eviction integration remain. +metadata to reconstruct the strong key. + +The host `StagePrepare` and `StageLoad` lifecycle now uses that cache. The +stage-load wire carries an additive `auto`/affine-4/affine-8/MXFP4 profile; +older peers that omit it select `auto`, while unknown values fail closed. On +Apple Metal, `auto` currently resolves to affine-4/group-64. A clean two-half +host test made 136 and 137 tensor-payload requests, fetched 162,825,984 and +162,827,136 bytes, and produced 45,859,713 and 45,861,308-byte artifacts. The +complete Prepare/Load/Start/generate/Stop lifecycle passed in 120.74 seconds +with the established affine-4 tokens. An immediate validated-cache run passed +in 8.87 seconds with 258,162,688 B max RSS. Cache eviction and an optional +local request-to-recipe locator remain; the host cache root can be overridden +with `MESH_MLX_DERIVED_CACHE_DIR`. + +Only Prepare may fetch tensor payloads and build a missing entry. Load performs +metadata planning and full cache validation, then fails rather than deriving if +the prepared entry is absent or corrupt. + +Checkpoint claims are compared with the metadata-derived identity before the +first tensor payload request. Cancellation is cooperative through cache-lock +waits and the sequential range/quantization loop; an already-running HTTP or +MLX operation finishes before cleanup. Inventory responses echo the requested +quantization profile, and preparation plus running status carry it, so readiness +for one recipe cannot silently satisfy another. Cache-hit validation reads each +shard once while checking its own hash and the aggregate content digest. +Mixed-version automatic placement must still capability-gate explicit +non-default profiles because an old receiver ignores new additive fields. ## Representative measurements @@ -403,9 +427,9 @@ because it does not alter the derived packed weights. ## Next proof -1. Wire the proven derived-stage cache into host prepare/load and cache eviction, - then decide whether a local request-to-recipe locator should remove even the - warm metadata probes. +1. Add capacity/eviction ownership to the host derived-stage cache, then decide + whether a local request-to-recipe locator should remove even the warm + metadata probes. 2. Quantize one real Nemotron-H BF16 matrix reproducibly, then implement one complete split-expert bank without accumulating every dense expert. Prove a small family member before attempting Ultra-scale ranges. From 24c4e0de6fb67d9849cde3370f03bbcacb58ebd5 Mon Sep 17 00:00:00 2001 From: Michael Neale <14976+michaelneale@users.noreply.github.com> Date: Fri, 17 Jul 2026 21:42:10 +1000 Subject: [PATCH 19/37] feat(model-hf): plan Nemotron-H tensor ranges --- Cargo.lock | 1 + crates/model-hf/Cargo.toml | 1 + .../model-hf/src/safetensors_stage/layout.rs | 236 +++++++++++++++--- .../src/safetensors_stage/tensor_stream.rs | 27 ++ crates/skippy-engine-mlx/STAGED_EXECUTION.md | 21 ++ docs/design/MLX_STAGE_ENGINE_PLAN.md | 25 +- spikes/mlx-safetensors-stages/FINDINGS.md | 28 ++- 7 files changed, 292 insertions(+), 47 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index e1b4677fe7..378920d59d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4540,6 +4540,7 @@ dependencies = [ "chrono", "dirs", "hf-hub", + "json5", "libc", "model-artifact", "model-ref", diff --git a/crates/model-hf/Cargo.toml b/crates/model-hf/Cargo.toml index bbcf0fba24..6a710467ff 100644 --- a/crates/model-hf/Cargo.toml +++ b/crates/model-hf/Cargo.toml @@ -13,6 +13,7 @@ async-trait = "0.1" chrono = { version = "0.4", features = ["serde"] } dirs = "6.0.0" hf_hub = { package = "hf-hub", version = "1.0.0-rc.1", default-features = false, features = ["blocking"] } +json5 = "1.3.1" model-artifact = { path = "../model-artifact", version = "0.72.1" } model-ref = { path = "../model-ref", version = "0.72.1" } reqwest = { version = "0.12", default-features = false, features = ["blocking", "rustls-tls"] } diff --git a/crates/model-hf/src/safetensors_stage/layout.rs b/crates/model-hf/src/safetensors_stage/layout.rs index 0783e917c6..39f1ae21c5 100644 --- a/crates/model-hf/src/safetensors_stage/layout.rs +++ b/crates/model-hf/src/safetensors_stage/layout.rs @@ -2,9 +2,9 @@ use std::collections::{BTreeMap, BTreeSet}; use anyhow::{Context, Result, anyhow, ensure}; use model_artifact::safetensors::{ - IndexMetadata, LlamaConfig, SafetensorsIndex, TensorHeader, parse_header, parse_index, - parse_llama_config, + IndexMetadata, SafetensorsIndex, TensorHeader, parse_header, parse_index, }; +use serde::Deserialize; use sha2::{Digest, Sha256}; use super::{ @@ -18,6 +18,57 @@ use super::{ pub(crate) const MAX_INDEX_BYTES: u64 = 64 * 1024 * 1024; const MAX_HEADER_BYTES: u64 = 256 * 1024 * 1024; +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +enum StageModelFamily { + Llama, + NemotronH, +} + +#[derive(Debug, Deserialize)] +struct RawStageModelConfig { + model_type: String, + num_hidden_layers: u32, +} + +#[derive(Clone, Copy, Debug)] +struct StageModelLayout { + family: StageModelFamily, + num_hidden_layers: u32, +} + +impl StageModelLayout { + fn layer_index(self, name: &str) -> Option { + self.layer_prefixes() + .iter() + .find_map(|prefix| name.strip_prefix(prefix)?.split_once('.')?.0.parse().ok()) + } + + fn layer_prefixes(self) -> &'static [&'static str] { + match self.family { + StageModelFamily::Llama => &["model.layers."], + StageModelFamily::NemotronH => &["backbone.layers.", "model.backbone.layers."], + } + } + + fn embedding_prefix(self) -> &'static str { + match self.family { + StageModelFamily::Llama => "model.embed_tokens.", + StageModelFamily::NemotronH => "backbone.embeddings.", + } + } + + fn final_norm_prefix(self) -> &'static str { + match self.family { + StageModelFamily::Llama => "model.norm.", + StageModelFamily::NemotronH => "backbone.norm_f.", + } + } + + fn readout_prefix(self) -> &'static str { + "lm_head." + } +} + #[derive(Clone, Debug)] struct RemoteHeader { header_len: u64, @@ -44,7 +95,7 @@ pub(crate) fn prepare( let config = remote .small_file(config_url, MAX_INDEX_BYTES) .context("download SafeTensors model config")?; - let model_config = parse_llama_config(&config.bytes)?; + let model_config = parse_stage_model_layout(&config.bytes)?; validate_layer_range(request, &model_config)?; let config_sha256 = sha256_hex(&config.bytes); let mut selection_request = request.clone(); @@ -52,8 +103,8 @@ pub(crate) fn prepare( let mut layout = load_checkpoint_layout(remote, request)?; let checkpoint_sha256 = checkpoint_sha256(request, &config_sha256, &layout.layout_sha256)?; - validate_layer_coverage(&layout.index, request)?; - let selected = select_tensors(&layout.index.weight_map, &selection_request); + validate_layer_coverage(&layout.index, request, model_config)?; + let selected = select_tensors(&layout.index.weight_map, &selection_request, model_config); ensure!( !selected.is_empty(), "no tensors matched layers {}..{} or requested prefixes", @@ -122,7 +173,38 @@ pub(crate) fn prepare( }) } -fn validate_layer_range(request: &SafetensorsStageRequest, config: &LlamaConfig) -> Result<()> { +fn parse_stage_model_layout(bytes: &[u8]) -> Result { + let config: RawStageModelConfig = match serde_json::from_slice(bytes) { + Ok(config) => config, + Err(strict_error) => { + let text = + std::str::from_utf8(bytes).context("SafeTensors model config is not UTF-8")?; + json5::from_str(text).with_context(|| { + format!("parse SafeTensors model config as strict JSON ({strict_error}) or JSON5") + })? + } + }; + ensure!( + config.num_hidden_layers > 0, + "SafeTensors model num_hidden_layers must be non-zero" + ); + let family = match config.model_type.as_str() { + "llama" => StageModelFamily::Llama, + "nemotron_h" => StageModelFamily::NemotronH, + model_type => anyhow::bail!( + "MLX partial SafeTensors currently supports model_type=llama or nemotron_h, got {model_type:?}" + ), + }; + Ok(StageModelLayout { + family, + num_hidden_layers: config.num_hidden_layers, + }) +} + +fn validate_layer_range( + request: &SafetensorsStageRequest, + config: &StageModelLayout, +) -> Result<()> { ensure!( request.layer_end <= config.num_hidden_layers, "stage layer end {} exceeds model layer count {}", @@ -132,15 +214,22 @@ fn validate_layer_range(request: &SafetensorsStageRequest, config: &LlamaConfig) Ok(()) } -fn add_required_prefixes(request: &mut SafetensorsStageRequest, config: &LlamaConfig) { - if request.layer_start == 0 || request.layer_end == config.num_hidden_layers { +fn add_required_prefixes(request: &mut SafetensorsStageRequest, config: &StageModelLayout) { + if request.layer_start == 0 + || (config.family == StageModelFamily::Llama + && request.layer_end == config.num_hidden_layers) + { request .include_prefixes - .push("model.embed_tokens.".to_string()); + .push(config.embedding_prefix().to_string()); } if request.layer_end == config.num_hidden_layers { - request.include_prefixes.push("model.norm.".to_string()); - request.include_prefixes.push("lm_head.".to_string()); + request + .include_prefixes + .push(config.final_norm_prefix().to_string()); + request + .include_prefixes + .push(config.readout_prefix().to_string()); } request.include_prefixes.sort(); request.include_prefixes.dedup(); @@ -149,13 +238,14 @@ fn add_required_prefixes(request: &mut SafetensorsStageRequest, config: &LlamaCo fn validate_layer_coverage( index: &SafetensorsIndex, request: &SafetensorsStageRequest, + config: StageModelLayout, ) -> Result<()> { for layer in request.layer_start..request.layer_end { ensure!( index .weight_map .keys() - .any(|name| layer_index(name) == Some(layer)), + .any(|name| config.layer_index(name) == Some(layer)), "SafeTensors checkpoint has no tensors for requested layer {layer}" ); } @@ -165,23 +255,33 @@ fn validate_layer_coverage( fn validate_required_tensors( selected: &BTreeSet<&str>, request: &SafetensorsStageRequest, - config: &LlamaConfig, + config: &StageModelLayout, ) -> Result<()> { let has_prefix = |prefix: &str| selected.iter().any(|name| name.starts_with(prefix)); if request.layer_start == 0 { ensure!( - has_prefix("model.embed_tokens."), - "first MLX stage requires model.embed_tokens tensors" + has_prefix(config.embedding_prefix()), + "first MLX stage requires {} tensors", + config.embedding_prefix() ); } if request.layer_end == config.num_hidden_layers { ensure!( - has_prefix("model.norm."), - "final MLX stage requires model.norm tensors" + has_prefix(config.final_norm_prefix()), + "final MLX stage requires {} tensors", + config.final_norm_prefix() ); + let has_readout = has_prefix(config.readout_prefix()) + || (config.family == StageModelFamily::Llama && has_prefix(config.embedding_prefix())); ensure!( - has_prefix("lm_head.") || has_prefix("model.embed_tokens."), - "final MLX stage requires lm_head or tied embedding tensors" + has_readout, + "final MLX stage requires {} tensors{}", + config.readout_prefix(), + if config.family == StageModelFamily::Llama { + " or tied embeddings" + } else { + "" + } ); } Ok(()) @@ -299,11 +399,13 @@ fn ensure_matching_etag(left: &str, right: &str, file: &str) -> Result<()> { fn select_tensors<'a>( weight_map: &'a BTreeMap, request: &SafetensorsStageRequest, + config: StageModelLayout, ) -> BTreeSet<&'a str> { weight_map .keys() .filter(|name| { - layer_index(name) + config + .layer_index(name) .is_some_and(|layer| layer >= request.layer_start && layer < request.layer_end) || request .include_prefixes @@ -314,14 +416,6 @@ fn select_tensors<'a>( .collect() } -fn layer_index(name: &str) -> Option { - name.strip_prefix("model.layers.")? - .split_once('.')? - .0 - .parse() - .ok() -} - fn plan_shard( file: &str, header: &RemoteHeader, @@ -502,10 +596,50 @@ mod tests { use super::*; #[test] - fn recognizes_only_llama_layer_paths() { - assert_eq!(layer_index("model.layers.42.mlp.up_proj.weight"), Some(42)); - assert_eq!(layer_index("transformer.h.7.attn.weight"), None); - assert_eq!(layer_index("model.embed_tokens.weight"), None); + fn recognizes_family_specific_layer_paths() { + let llama = parse_stage_model_layout( + br#"{"model_type":"llama","hidden_size":64,"num_hidden_layers":48}"#, + ) + .unwrap(); + assert_eq!( + llama.layer_index("model.layers.42.mlp.up_proj.weight"), + Some(42) + ); + assert_eq!(llama.layer_index("backbone.layers.7.mixer.weight"), None); + assert_eq!(llama.layer_index("model.embed_tokens.weight"), None); + + let nemotron = parse_stage_model_layout( + br#"{"model_type":"nemotron_h","hidden_size":64,"num_hidden_layers":52}"#, + ) + .unwrap(); + assert_eq!( + nemotron.layer_index("backbone.layers.7.mixer.up_proj.weight"), + Some(7) + ); + assert_eq!( + nemotron.layer_index("model.backbone.layers.8.mixer.weight"), + Some(8) + ); + assert_eq!( + nemotron.layer_index("model.layers.42.mlp.up_proj.weight"), + None + ); + } + + #[test] + fn parses_hugging_face_nonfinite_config_values_without_using_them() { + let layout = parse_stage_model_layout( + br#"{ + "model_type":"nemotron_h", + "hidden_size":2688, + "num_hidden_layers":52, + "time_step_limit":[0.0, Infinity] + }"#, + ) + .unwrap(); + + assert_eq!(layout.family, StageModelFamily::NemotronH); + assert_eq!(layout.num_hidden_layers, 52); } #[test] @@ -551,11 +685,9 @@ mod tests { layer_end: 2, include_prefixes: Vec::new(), }; - let config = LlamaConfig { - model_type: "llama".to_string(), - hidden_size: 2, + let config = StageModelLayout { + family: StageModelFamily::Llama, num_hidden_layers: 2, - tie_word_embeddings: true, }; add_required_prefixes(&mut request, &config); @@ -569,4 +701,36 @@ mod tests { ] ); } + + #[test] + fn assigns_nemotron_embeddings_and_readout_to_their_own_boundaries() { + let config = StageModelLayout { + family: StageModelFamily::NemotronH, + num_hidden_layers: 2, + }; + let mut first = SafetensorsStageRequest { + repo: "org/model".to_string(), + revision: "a".repeat(40), + layer_start: 0, + layer_end: 1, + include_prefixes: Vec::new(), + }; + add_required_prefixes(&mut first, &config); + assert_eq!( + first.include_prefixes, + vec!["backbone.embeddings.".to_string()] + ); + + let mut final_stage = SafetensorsStageRequest { + layer_start: 1, + layer_end: 2, + ..first + }; + final_stage.include_prefixes.clear(); + add_required_prefixes(&mut final_stage, &config); + assert_eq!( + final_stage.include_prefixes, + vec!["backbone.norm_f.".to_string(), "lm_head.".to_string()] + ); + } } diff --git a/crates/model-hf/src/safetensors_stage/tensor_stream.rs b/crates/model-hf/src/safetensors_stage/tensor_stream.rs index fbbe01a1fb..d6a146518f 100644 --- a/crates/model-hf/src/safetensors_stage/tensor_stream.rs +++ b/crates/model-hf/src/safetensors_stage/tensor_stream.rs @@ -332,6 +332,33 @@ mod tests { assert_eq!(parsed[&tensor.name].data_offsets, [0, 8]); } + #[test] + #[ignore = "reads pinned Nemotron-H config, index, and one shard header"] + fn plans_real_nemotron_h_moe_layer_without_tensor_payloads() { + let cache = tempfile::tempdir().unwrap(); + let materializer = + SafetensorsStageMaterializer::new(cache.path().join("cache"), None, None).unwrap(); + let visit = materializer + .prepare_tensor_visit(SafetensorsStageRequest { + repo: "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-Base-BF16".to_string(), + revision: "97ab8012882a655dc38df4fee47422aca9caca07".to_string(), + layer_start: 1, + layer_end: 2, + include_prefixes: Vec::new(), + }) + .unwrap(); + let plan = visit.plan(); + + assert_eq!(plan.selected_tensor_count, 261); + assert_eq!(plan.selected_tensor_bytes, 2_594_936_576); + assert_eq!(plan.largest_selected_tensor_bytes, 19_955_712); + assert_eq!(plan.source_shard_count, 1); + assert_eq!(plan.source_shard_bytes, 4_991_210_024); + assert_eq!(plan.range_request_count, 2); + assert!(plan.source_shard_bytes_avoided > 2_396_000_000); + assert_eq!(visit.checkpoint_sha256().len(), 64); + } + #[test] #[ignore = "downloads one layer from a pinned Hugging Face SafeTensors checkpoint"] fn visits_real_smollm2_layer_without_retaining_source_shard() { diff --git a/crates/skippy-engine-mlx/STAGED_EXECUTION.md b/crates/skippy-engine-mlx/STAGED_EXECUTION.md index 0808a64eda..bda5eaaf92 100644 --- a/crates/skippy-engine-mlx/STAGED_EXECUTION.md +++ b/crates/skippy-engine-mlx/STAGED_EXECUTION.md @@ -165,6 +165,24 @@ are intentionally duplicated across the stages; that is why the sum of the two files is larger than the full checkpoint even though neither process downloads the full checkpoint. +The production `model-hf` planner now also understands the `nemotron_h` +architecture used by Nemotron 3 Nano, including its `backbone.layers.*` layout +and first/final boundary tensors. Against pinned +`nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-Base-BF16` layer `1`, it selected +2,594,936,576 bytes in 261 tensors from one 4,991,210,024-byte shard; the +largest individual tensor was 19,955,712 bytes. This is metadata/range-planning +evidence only. The derived builder and stage engine remain fail-closed for +Nemotron-H until split ReLU2 experts and recurrent state are implemented. + +Reproduce the metadata-only proof (it downloads the pinned config, index, and +one SafeTensors header, but no tensor payloads): + +```bash +cargo test -p model-hf --lib \ + plans_real_nemotron_h_moe_layer_without_tensor_payloads -- \ + --ignored --nocapture +``` + ## Reproduce Build once: @@ -239,6 +257,9 @@ just mlx-stage prove --connect 127.0.0.1:19090 --wire-dtype f16 experts need per-layer expert-bank assembly before quantization; Inkling needs its transformed rank-3 grouped-expert loader. Neither is silently treated as Llama. +- The pinned safemlx Nemotron-H implementation matches the 52-layer Nano + schema, not Nemotron 3 Ultra's 108-layer latent-MoE schema. Ultra range plans + are storage-locality evidence, not executable-family support. - Greedy sampling only; sampling metadata is preserved in the contract and rejected explicitly when enabled. - No KV page import/export, cache trim/checkpoint, MTP, speculative verify, diff --git a/docs/design/MLX_STAGE_ENGINE_PLAN.md b/docs/design/MLX_STAGE_ENGINE_PLAN.md index a4a01f8f43..0b0b207dd4 100644 --- a/docs/design/MLX_STAGE_ENGINE_PLAN.md +++ b/docs/design/MLX_STAGE_ENGINE_PLAN.md @@ -86,6 +86,22 @@ bytes; the temporary directory was empty at completion. A prepared visit makes the verified config and checkpoint identity available before tensor callbacks, and macOS/Unix advisory locks safely scavenge crash-abandoned visits. +**Update — production range planning now covers the `nemotron_h` architecture +used by Nemotron 3 Nano.** The planner recognizes its `backbone.layers.*`, +embedding, final-norm, and readout paths and falls back to JSON5 for Hugging +Face configs containing bare non-finite values. +For pinned NVIDIA 30B-A3B Base BF16 layer `1`, it selected 261 tensors totaling +2,594,936,576 bytes from a 4,991,210,024-byte shard; the largest tensor was +19,955,712 bytes. No tensor payload was fetched by this proof. This removes the +acquisition blocker for the next bounded expert-pack experiment. + +Nemotron 3 Ultra is not that next executable target. Its current public config +uses a 108-layer `layers_block_type` latent-MoE design with 512 experts and +`moe_latent_size=2048`, while the pinned safemlx `nemotron_h` implementation is +the 52-layer Nano schema driven by `num_hidden_layers` and +`hybrid_override_pattern`. Ultra range measurements remain valid storage +evidence, but execution needs separate model-family work. + **Update — direct exact-range to bounded affine stage artifacts is proven for Llama.** `mlx-stage derive` now consumes that prepared visit, quantizes/evaluates one rank-2 matrix at a time, serializes packed results into bounded SafeTensors @@ -864,10 +880,11 @@ Spikes 1 and 2 are more decisive than any standalone token/s benchmark. the dense llama adapter and MLX two-process binary-wire proof are complete. 3. Run **Spike 2 (boundary fence)** at frontier residual widths and keep it as a go/no-go gate. -4. Use Nemotron-H as the first frontier-family follow-up: wire its existing - public layer/cache structures and affine expert runtime to public-checkpoint - on-load packing. Then expose safemlx's existing Inkling implementation as a - staged text decoder and use Transformers as the parity oracle. +4. Use Nemotron-H **Nano** as the first frontier-family follow-up: wire its + existing public layer/cache structures and affine expert runtime to + public-checkpoint on-load packing. Do not treat Ultra as the same runtime + family. Then expose safemlx's existing Inkling implementation as a staged + text decoder and use Transformers as the parity oracle. --- diff --git a/spikes/mlx-safetensors-stages/FINDINGS.md b/spikes/mlx-safetensors-stages/FINDINGS.md index ee6f27d4a1..d5bfd4c885 100644 --- a/spikes/mlx-safetensors-stages/FINDINGS.md +++ b/spikes/mlx-safetensors-stages/FINDINGS.md @@ -215,6 +215,15 @@ Every repository is pinned to the immutable revision shown below. | GLM-5.2 BF16 | `b4734de4facf877f85769a911abafc5283eab3d9` | 36..40 | 1.37 TiB | 73.54 GiB | 79.90 GiB | 6.36 GiB | 192 MiB | | DeepSeek V4 Pro FP8 | `b5968e9190ef611bbf34a7229255be88a0e937c1` | 28..32 | 805.32 GiB | 51.73 GiB | 51.73 GiB | 0 | 112 MiB | +The production `model-hf` planner now reproduces a tractable `nemotron_h` +architecture case from Nemotron 3 Nano rather than relying only on this spike. +At immutable revision +`97ab8012882a655dc38df4fee47422aca9caca07`, layer `1` of NVIDIA's 30B-A3B Base +BF16 checkpoint selects 261 tensors / 2,594,936,576 bytes from a single +4,991,210,024-byte shard, with a 19,955,712-byte largest tensor and two +coalesced payload ranges. The config contains bare `Infinity`, so production +metadata parsing uses strict JSON first and a JSON5 fallback. + The format mechanism is stable and documented by the [SafeTensors format](https://github.com/huggingface/safetensors#format): the header records each tensor's dtype, shape, and byte offsets. Hugging Face's @@ -387,10 +396,10 @@ SafeTensors acquisition is general, but MLX execution remains family-specific. The measured candidates suggest this order: 1. **Qwen/Llama**: finish the partial loader and two-stage correctness proof. -2. **Nemotron 3 Ultra**: best next frontier proof because `safemlx-lm` already - has a Nemotron-H implementation with public layer/cache structures and an - affine rank-3 expert runtime. Its public SafeTensors loader still needs the - matching on-load packing path and Mamba/recurrent stage-boundary certification. +2. **Nemotron-H Nano 30B-A3B**: best next frontier proof because + `safemlx-lm` has matching public layer/cache structures and an affine rank-3 + expert runtime. Its public SafeTensors path still needs bounded quantized + ReLU2 expert-bank assembly and Mamba/recurrent boundary certification. 3. **Inkling text backbone**: safemlx already has whole-model text/multimodal execution; expose a partial-stage surface and use Transformers as the parity oracle rather than porting the family again. @@ -406,6 +415,11 @@ FP8. Preserve a compatible calibrated source encoding when the local backend supports it; use BF16-to-local-quant only when it is the cleanest compatible source path. +Nemotron 3 Ultra is a distinct follow-up, not a larger drop-in Nano checkpoint. +Its public config describes 108 layers through `layers_block_type`, 512 experts, +and `moe_latent_size=2048`; it lacks the Nano fields consumed by the pinned +safemlx implementation. The Ultra row above proves range locality only. + ## Recommended artifact identity The durable identity should be: @@ -430,9 +444,9 @@ because it does not alter the derived packed weights. 1. Add capacity/eviction ownership to the host derived-stage cache, then decide whether a local request-to-recipe locator should remove even the warm metadata probes. -2. Quantize one real Nemotron-H BF16 matrix reproducibly, then implement one - complete split-expert bank without accumulating every dense expert. Prove a - small family member before attempting Ultra-scale ranges. +2. Quantize one real Nemotron-H Nano BF16 matrix reproducibly, then implement + one complete split-expert bank without accumulating every dense expert. Keep + Ultra gated behind its separate latent-MoE family implementation. 3. Measure the MLX eval/readback/codec boundary fence independently at frontier residual widths and prefill sizes. 4. Expose the existing safemlx Inkling text decoder as one stage, prove From 16447d17e19257225a82a781f8e169aea5473846 Mon Sep 17 00:00:00 2001 From: Michael Neale <14976+michaelneale@users.noreply.github.com> Date: Fri, 17 Jul 2026 22:01:48 +1000 Subject: [PATCH 20/37] feat(mlx): derive bounded Nemotron-H expert layers --- crates/skippy-engine-mlx/STAGED_EXECUTION.md | 32 +- crates/skippy-engine-mlx/src/bin/mlx-stage.rs | 18 +- crates/skippy-engine-mlx/src/derived.rs | 281 ++++++++++-- .../src/derived/expert_bank.rs | 290 ++++++++++++ .../src/derived/nemotron_h.rs | 414 ++++++++++++++++++ crates/skippy-engine-mlx/src/lib.rs | 5 +- docs/design/MLX_STAGE_ENGINE_PLAN.md | 30 +- spikes/mlx-safetensors-stages/FINDINGS.md | 21 +- 8 files changed, 1048 insertions(+), 43 deletions(-) create mode 100644 crates/skippy-engine-mlx/src/derived/expert_bank.rs create mode 100644 crates/skippy-engine-mlx/src/derived/nemotron_h.rs diff --git a/crates/skippy-engine-mlx/STAGED_EXECUTION.md b/crates/skippy-engine-mlx/STAGED_EXECUTION.md index bda5eaaf92..dc8cc7f4e5 100644 --- a/crates/skippy-engine-mlx/STAGED_EXECUTION.md +++ b/crates/skippy-engine-mlx/STAGED_EXECUTION.md @@ -171,8 +171,9 @@ and first/final boundary tensors. Against pinned `nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-Base-BF16` layer `1`, it selected 2,594,936,576 bytes in 261 tensors from one 4,991,210,024-byte shard; the largest individual tensor was 19,955,712 bytes. This is metadata/range-planning -evidence only. The derived builder and stage engine remain fail-closed for -Nemotron-H until split ReLU2 experts and recurrent state are implemented. +evidence only for the general family layout. The derived builder now supports +exactly one Nemotron-H Nano MoE layer at a time; the serving stage engine +remains fail-closed until hybrid recurrent/attention boundaries are implemented. Reproduce the metadata-only proof (it downloads the pinned config, index, and one SafeTensors header, but no tensor payloads): @@ -183,6 +184,30 @@ cargo test -p model-hf --lib \ --ignored --nocapture ``` +The bounded affine4 implementation has also been exercised against that exact +pinned layer. It streamed 2,594,936,576 BF16 bytes through 261 individual range +requests, quantized 258 matrices while retaining three dense tensors, and wrote +730,324,736 tensor bytes. Maximum process RSS was 822,165,504 bytes and the +largest ephemeral source tensor file was 19,955,848 bytes. The resulting +artifact strict-loaded into safemlx's real layer-1 `TransformerBlock` and +produced a finite `[1, 1, 2688]` output for a deterministic nonzero input. +Here, bounded memory means bounded by the final packed layer: the six routed +bank buffers total 718,405,632 bytes. It does not mean derivation stays at the +one-expert (~20 MB source tensor) footprint. The forward is an executable smoke +test, not a dense-versus-quantized numerical parity result. + +```bash +just mlx-stage-build +just mlx-stage derive \ + --repo nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-Base-BF16 \ + --revision 97ab8012882a655dc38df4fee47422aca9caca07 \ + --layer-start 1 --layer-end 2 \ + --output /tmp/nemotron-nano-layer1-affine4 \ + --weight-quantization affine4 +just mlx-stage validate-nemotron-h \ + --model /tmp/nemotron-nano-layer1-affine4 --layer 1 +``` + ## Reproduce Build once: @@ -260,6 +285,9 @@ just mlx-stage prove --connect 127.0.0.1:19090 --wire-dtype f16 - The pinned safemlx Nemotron-H implementation matches the 52-layer Nano schema, not Nemotron 3 Ultra's 108-layer latent-MoE schema. Ultra range plans are storage-locality evidence, not executable-family support. +- Bounded Nemotron-H derivation currently accepts exactly one `E`/MoE layer. + It does not yet expose a hybrid multi-layer stage or recurrent state on the + wire. - Greedy sampling only; sampling metadata is preserved in the contract and rejected explicitly when enabled. - No KV page import/export, cache trim/checkpoint, MTP, speculative verify, diff --git a/crates/skippy-engine-mlx/src/bin/mlx-stage.rs b/crates/skippy-engine-mlx/src/bin/mlx-stage.rs index ce84bb0153..4e96aeea63 100644 --- a/crates/skippy-engine-mlx/src/bin/mlx-stage.rs +++ b/crates/skippy-engine-mlx/src/bin/mlx-stage.rs @@ -15,7 +15,7 @@ mod real { use skippy_engine_mlx::{ MlxComputeDtype, MlxDerivationControl, MlxDerivedStageCacheConfig, MlxDerivedStageConfig, MlxStageEngine, MlxStageEngineConfig, MlxWeightQuantization, derive_quantized_stage, - derive_quantized_stage_cached, mlx_derived_stage_cache_root, + derive_quantized_stage_cached, mlx_derived_stage_cache_root, validate_nemotron_h_moe_stage, }; use skippy_protocol::binary::{ StageStateHeader, StageWireMessage, WireActivationDType, WireMessageKind, WireReplyKind, @@ -72,7 +72,7 @@ mod real { output: PathBuf, #[arg(long, value_enum, default_value_t = WeightQuantization::Affine4)] weight_quantization: WeightQuantization, - /// Soft output shard target; one packed tensor may exceed it. + /// Soft output shard target; one converted tensor bundle may exceed it. #[arg(long, default_value_t = 256)] shard_size_mib: usize, }, @@ -94,10 +94,17 @@ mod real { cache_root: Option, #[arg(long, value_enum, default_value_t = WeightQuantization::Affine4)] weight_quantization: WeightQuantization, - /// Soft output shard target; one packed tensor may exceed it. + /// Soft output shard target; one converted tensor bundle may exceed it. #[arg(long, default_value_t = 256)] shard_size_mib: usize, }, + /// Strict-load and execute one derived Nemotron-H MoE layer. + ValidateNemotronH { + #[arg(long)] + model: PathBuf, + #[arg(long)] + layer: usize, + }, /// Drive a stage chain and assert its greedy token sequence. Prove { #[arg(long)] @@ -238,6 +245,11 @@ mod real { weight_quantization.into(), shard_size_mib, ), + Command::ValidateNemotronH { model, layer } => { + let report = validate_nemotron_h_moe_stage(model, layer)?; + println!("{}", serde_json::to_string_pretty(&report)?); + Ok(()) + } Command::Prove { connect, tokens, diff --git a/crates/skippy-engine-mlx/src/derived.rs b/crates/skippy-engine-mlx/src/derived.rs index 9df256443a..2cd436304c 100644 --- a/crates/skippy-engine-mlx/src/derived.rs +++ b/crates/skippy-engine-mlx/src/derived.rs @@ -28,13 +28,17 @@ use serde_json::{Value, json}; use sha2::{Digest, Sha256}; use crate::stage::MlxWeightQuantization; +use nemotron_h::NemotronHDerivation; mod cache; +mod expert_bank; +mod nemotron_h; pub use cache::{ MlxDerivedStageCacheConfig, MlxDerivedStageCacheResult, derive_quantized_stage_cached, load_prepared_quantized_stage, mlx_derived_stage_cache_root, }; +pub use nemotron_h::{MlxNemotronHValidationReport, validate_nemotron_h_moe_stage}; pub(super) const DERIVED_STAGE_SCHEMA_VERSION: u32 = 1; const DERIVED_STAGE_IMPLEMENTATION: &str = "mesh-mlx-range-derived-v1"; @@ -152,9 +156,17 @@ impl PendingShard { } } - fn insert(&mut self, name: String, tensor: OwnedTensor) { - self.bytes = self.bytes.saturating_add(tensor.data.len()); + fn insert(&mut self, name: String, tensor: OwnedTensor) -> Result<()> { + ensure!( + !self.tensors.contains_key(&name), + "derived tensor name {name:?} was produced more than once" + ); + self.bytes = self + .bytes + .checked_add(tensor.data.len()) + .context("pending derived shard byte count overflow")?; self.tensors.insert(name, tensor); + Ok(()) } } @@ -193,6 +205,45 @@ struct BuildState { working_disk_peak_bytes: u64, } +enum TensorDerivation { + Llama, + NemotronH(NemotronHDerivation), +} + +impl TensorDerivation { + fn new( + config: &Value, + layer_start: u32, + layer_end: u32, + quantization: WeightQuantization, + ) -> Result { + match config.get("model_type").and_then(Value::as_str) { + Some("llama") => Ok(Self::Llama), + Some("nemotron_h") => { + ensure!( + matches!(quantization, WeightQuantization::Affine(_)), + "bounded Nemotron-H expert banks currently require affine quantization" + ); + Ok(Self::NemotronH(NemotronHDerivation::new( + config, + layer_start, + layer_end, + )?)) + } + model_type => anyhow::bail!( + "derived MLX stages support model_type=llama or nemotron_h, got {model_type:?}" + ), + } + } + + fn finish(self) -> Result> { + match self { + Self::Llama => Ok(Vec::new()), + Self::NemotronH(derivation) => derivation.finish(), + } + } +} + impl BuildState { fn new(initial_output_file_bytes: u64) -> Self { Self { @@ -236,9 +287,15 @@ pub fn derive_quantized_stage( .control .verify_checkpoint(visit.checkpoint_sha256())?; let plan = visit.plan().clone(); - let source_config = visit.config().to_vec(); - ensure_dense_source_config(&source_config)?; + let source_config = parse_source_config(visit.config())?; + ensure_unquantized_source_config(&source_config)?; let quantization = config.quantization.safemlx()?; + let mut tensor_derivation = TensorDerivation::new( + &source_config, + plan.layer_start, + plan.layer_end, + quantization, + )?; let quantization_value = serde_json::to_value(quantization)?; let plan_bytes = serde_json::to_vec(&plan)?; let plan_sha256 = sha256_bytes(&plan_bytes); @@ -271,6 +328,7 @@ pub fn derive_quantized_stage( quantization, &weights_stream, &quantization_stream, + &mut tensor_derivation, &mut state, )?; config.control.ensure_active()?; @@ -285,19 +343,39 @@ pub fn derive_quantized_stage( }, )?; config.control.ensure_active()?; + let final_arrays = tensor_derivation.finish()?; + if !final_arrays.is_empty() { + append_arrays( + final_arrays, + config.shard_size_bytes, + temporary.path(), + 0, + &mut state, + )?; + } + config.control.ensure_active()?; if !state.pending.tensors.is_empty() { flush_shard(temporary.path(), &mut state)?; } + config.control.ensure_active()?; ensure!( !state.temporary_shards.is_empty(), "derived stage contains no tensors" ); let finalized = finalize_shards(temporary.path(), &state)?; + config.control.ensure_active()?; let shards = finalized .iter() - .map(|path| derived_shard(path)) + .map(|path| { + config.control.ensure_active()?; + let shard = derived_shard(path)?; + config.control.ensure_active()?; + Ok(shard) + }) .collect::>>()?; + config.control.ensure_active()?; let output_content_sha256 = output_content_sha256(temporary.path())?; + config.control.ensure_active()?; let artifact_file_bytes = artifact_file_bytes(temporary.path())?; state.working_disk_peak_bytes = state.working_disk_peak_bytes.max(artifact_file_bytes); let report = MlxDerivedStageReport { @@ -340,6 +418,7 @@ fn convert_tensor( quantization: WeightQuantization, weights_stream: &Stream, quantization_stream: &Stream, + derivation: &mut TensorDerivation, state: &mut BuildState, ) -> Result> { let file = File::open(&tensor.path)?; @@ -353,14 +432,31 @@ fn convert_tensor( tensor.name ); let dense = Array::try_from(tensors.tensor(&tensor.name)?)?.copy(weights_stream)?; - let arrays = if should_quantize_source_weight(&tensor.name, &dense, quantization)? { + if let TensorDerivation::NemotronH(nemotron) = derivation + && nemotron.consume_expert(&tensor.name, &dense, quantization, quantization_stream)? + { state.quantized_tensor_count += 1; - quantize_tensor(&dense, quantization, quantization_stream)? - .into_named_arrays(&tensor.name)? - } else { - state.copied_tensor_count += 1; - vec![(tensor.name.clone(), dense)] + weights_stream.synchronize()?; + quantization_stream.synchronize()?; + return Ok(Vec::new()); + } + let output_name = match derivation { + TensorDerivation::Llama => tensor.name.clone(), + TensorDerivation::NemotronH(nemotron) => nemotron.rewrite_name(&tensor.name)?, }; + let keep_dense = matches!( + derivation, + TensorDerivation::NemotronH(_) if NemotronHDerivation::keep_dense(&output_name) + ); + let arrays = + if !keep_dense && should_quantize_source_weight(&output_name, &dense, quantization)? { + state.quantized_tensor_count += 1; + quantize_tensor(&dense, quantization, quantization_stream)? + .into_named_arrays(&output_name)? + } else { + state.copied_tensor_count += 1; + vec![(output_name, dense)] + }; eval(arrays.iter().map(|(_, array)| array))?; weights_stream.synchronize()?; quantization_stream.synchronize()?; @@ -387,7 +483,7 @@ fn should_quantize_source_weight( } ensure!( tensor.ndim() == 2, - "derived stage v1 only supports dense rank-2 Llama weights; {name} has rank {}", + "derived stage only supports dense rank-2 matrix weights; {name} has rank {}", tensor.ndim() ); ensure!( @@ -504,7 +600,7 @@ fn append_arrays( .output_tensor_bytes .checked_add(u64::try_from(tensor.data.len())?) .context("derived output tensor byte count overflow")?; - state.pending.insert(name, tensor); + state.pending.insert(name, tensor)?; } Ok(()) } @@ -573,8 +669,8 @@ fn finalize_shards(output_dir: &Path, state: &BuildState) -> Result Ok(outputs) } -fn write_quantized_config(path: PathBuf, source: &[u8], quantization: &Value) -> Result<()> { - let mut config: Value = serde_json::from_slice(source).context("parse source config.json")?; +fn write_quantized_config(path: PathBuf, source: &Value, quantization: &Value) -> Result<()> { + let mut config = source.clone(); let object = config .as_object_mut() .context("source config.json must contain an object")?; @@ -583,14 +679,16 @@ fn write_quantized_config(path: PathBuf, source: &[u8], quantization: &Value) -> write_json(path, &config) } -fn ensure_dense_source_config(source: &[u8]) -> Result<()> { - let config: Value = serde_json::from_slice(source).context("parse source config.json")?; +fn ensure_unquantized_source_config(config: &Value) -> Result<()> { let object = config .as_object() .context("source config.json must contain an object")?; ensure!( - object.get("model_type").and_then(Value::as_str) == Some("llama"), - "derived stage v1 only supports model_type=llama" + matches!( + object.get("model_type").and_then(Value::as_str), + Some("llama" | "nemotron_h") + ), + "derived MLX stages support model_type=llama or nemotron_h" ); for key in ["quantization", "quantization_config", "compression_config"] { ensure!( @@ -601,6 +699,70 @@ fn ensure_dense_source_config(source: &[u8]) -> Result<()> { Ok(()) } +fn parse_source_config(source: &[u8]) -> Result { + match serde_json::from_slice(source) { + Ok(config) => Ok(config), + Err(strict_error) => { + let normalized = normalize_nonfinite_json_tokens(source)?; + serde_json::from_slice(&normalized).with_context(|| { + format!( + "parse source config.json after normalizing non-finite values; strict JSON error: {strict_error}" + ) + }) + } + } +} + +fn normalize_nonfinite_json_tokens(source: &[u8]) -> Result> { + let source = std::str::from_utf8(source).context("source config.json is not UTF-8")?; + let bytes = source.as_bytes(); + let mut output = Vec::with_capacity(bytes.len()); + let mut index = 0; + let mut in_string = false; + let mut escaped = false; + while index < bytes.len() { + let byte = bytes[index]; + if in_string { + output.push(byte); + if escaped { + escaped = false; + } else if byte == b'\\' { + escaped = true; + } else if byte == b'"' { + in_string = false; + } + index += 1; + continue; + } + if byte == b'"' { + in_string = true; + output.push(byte); + index += 1; + continue; + } + let replacement = ["-Infinity", "+Infinity", "Infinity", "NaN"] + .into_iter() + .find(|token| source[index..].starts_with(token)); + if let Some(token) = replacement { + let end = index + token.len(); + let leading_boundary = index == 0 + || bytes[index - 1].is_ascii_whitespace() + || matches!(bytes[index - 1], b':' | b',' | b'['); + let trailing_boundary = bytes.get(end).is_none_or(|next| { + next.is_ascii_whitespace() || matches!(next, b',' | b']' | b'}') + }); + if leading_boundary && trailing_boundary { + output.extend_from_slice(b"null"); + index = end; + continue; + } + } + output.push(byte); + index += 1; + } + Ok(output) +} + pub(super) fn prepare_derivation_recipe( materializer: &SafetensorsStageMaterializer, source: SafetensorsStageRequest, @@ -612,8 +774,16 @@ pub(super) fn prepare_derivation_recipe( let visit = materializer.prepare_tensor_visit(source)?; control.ensure_active()?; control.verify_checkpoint(visit.checkpoint_sha256())?; - ensure_dense_source_config(visit.config())?; - let quantization = serde_json::to_value(quantization.safemlx()?)?; + let source_config = parse_source_config(visit.config())?; + ensure_unquantized_source_config(&source_config)?; + let quantization = quantization.safemlx()?; + TensorDerivation::new( + &source_config, + visit.plan().layer_start, + visit.plan().layer_end, + quantization, + )?; + let quantization = serde_json::to_value(quantization)?; let plan_bytes = serde_json::to_vec(visit.plan())?; let plan_sha256 = sha256_bytes(&plan_bytes); derived_identity(visit.plan(), &plan_sha256, &quantization, shard_size_bytes) @@ -901,7 +1071,8 @@ mod tests { )) .unwrap(); - write_quantized_config(path.clone(), br#"{"model_type":"llama"}"#, &quantization).unwrap(); + write_quantized_config(path.clone(), &json!({"model_type": "llama"}), &quantization) + .unwrap(); let config: Value = serde_json::from_slice(&fs::read(path).unwrap()).unwrap(); assert_eq!(config["model_type"], "llama"); @@ -911,14 +1082,33 @@ mod tests { #[test] fn rejects_prequantized_source_metadata_and_ineligible_weights() { - let error = ensure_dense_source_config( - br#"{"model_type":"llama","quantization_config":{"quant_method":"fp8"}}"#, - ) + let error = ensure_unquantized_source_config(&json!({ + "model_type": "llama", + "quantization_config": {"quant_method": "fp8"} + })) .unwrap_err(); assert!(error.to_string().contains("implicit dequantization")); - let error = ensure_dense_source_config(br#"{"model_type":"nemotron_h"}"#).unwrap_err(); - assert!(error.to_string().contains("model_type=llama")); + let error = + ensure_unquantized_source_config(&json!({"model_type": "unknown"})).unwrap_err(); + assert!(error.to_string().contains("llama or nemotron_h")); + + let error = TensorDerivation::new( + &json!({ + "model_type": "nemotron_h", + "hidden_size": 64, + "num_hidden_layers": 1, + "hybrid_override_pattern": "E", + "n_routed_experts": 2, + "moe_intermediate_size": 32 + }), + 0, + 1, + WeightQuantization::MxFp4, + ) + .err() + .expect("Nemotron-H MXFP4 should fail before tensor payloads"); + assert!(error.to_string().contains("require affine")); let quantization: WeightQuantization = AffineQuantization::new(64, 4).unwrap().into(); let packed = Array::from_slice(&vec![0_u32; 128], &[2, 64]); @@ -948,10 +1138,45 @@ mod tests { ) .unwrap_err() .to_string() - .contains("rank-2 Llama") + .contains("rank-2 matrix") ); } + #[test] + fn normalizes_nonfinite_config_values_but_not_strings() { + let config = parse_source_config( + br#"{ + "model_type":"nemotron_h", + "label":"Infinity and NaN", + "limits":[Infinity,-Infinity,+Infinity,NaN] + }"#, + ) + .unwrap(); + + assert_eq!(config["label"], "Infinity and NaN"); + assert_eq!(config["limits"], json!([null, null, null, null])); + assert!(parse_source_config(br#"{"bad":123Infinity}"#).is_err()); + } + + #[test] + fn pending_shard_rejects_canonical_name_collisions() { + let tensor = || OwnedTensor { + dtype: SafeDtype::F32, + shape: vec![1], + data: 1_f32.to_le_bytes().to_vec(), + }; + let mut pending = PendingShard::new(); + + pending + .insert("model.weight".to_string(), tensor()) + .unwrap(); + let error = pending + .insert("model.weight".to_string(), tensor()) + .unwrap_err(); + + assert!(error.to_string().contains("more than once")); + } + #[test] fn pure_rust_safetensors_output_preserves_bfloat16_and_packed_u32_bytes() { let directory = tempfile::tempdir().unwrap(); diff --git a/crates/skippy-engine-mlx/src/derived/expert_bank.rs b/crates/skippy-engine-mlx/src/derived/expert_bank.rs new file mode 100644 index 0000000000..382e88f289 --- /dev/null +++ b/crates/skippy-engine-mlx/src/derived/expert_bank.rs @@ -0,0 +1,290 @@ +//! Incremental affine expert-bank assembly without retaining dense experts. + +use anyhow::{Context, Result, ensure}; +use safemlx::{Stream, transforms::eval}; +use safemlx_lm::quantization::QuantizedTensor; + +use super::{OwnedTensor, owned_tensor}; + +const MAX_EXPERTS: usize = u128::BITS as usize; + +/// Builds one rank-3 affine expert projection from independently quantized +/// rank-2 experts. +/// +/// The final packed payloads are allocated once. Each expert is copied into its +/// own leading-dimension slice, so callers never need to retain every dense or +/// MLX expert array at the same time. +pub(super) struct AffineExpertBankAssembler { + output_prefix: String, + expert_count: usize, + completed: u128, + weights: Option, + scales: Option, + biases: Option, +} + +impl AffineExpertBankAssembler { + pub(super) fn new(output_prefix: impl Into, expert_count: usize) -> Result { + ensure!( + expert_count > 0, + "expert bank must contain at least one expert" + ); + ensure!( + expert_count <= MAX_EXPERTS, + "expert bank completion bitmap supports at most {MAX_EXPERTS} experts" + ); + Ok(Self { + output_prefix: output_prefix.into(), + expert_count, + completed: 0, + weights: None, + scales: None, + biases: None, + }) + } + + pub(super) fn insert( + &mut self, + expert_index: usize, + quantized: QuantizedTensor, + stream: &Stream, + ) -> Result<()> { + ensure!( + expert_index < self.expert_count, + "expert index {expert_index} exceeds bank size {}", + self.expert_count + ); + let bit = 1_u128 << expert_index; + ensure!( + self.completed & bit == 0, + "expert {expert_index} was inserted more than once" + ); + let biases = quantized + .biases + .context("affine expert quantization did not produce biases")?; + eval([&quantized.weight, &quantized.scales, &biases])?; + stream.synchronize()?; + + insert_tensor( + &mut self.weights, + expert_index, + self.expert_count, + owned_tensor(&quantized.weight)?, + )?; + insert_tensor( + &mut self.scales, + expert_index, + self.expert_count, + owned_tensor(&quantized.scales)?, + )?; + insert_tensor( + &mut self.biases, + expert_index, + self.expert_count, + owned_tensor(&biases)?, + )?; + self.completed |= bit; + Ok(()) + } + + pub(super) fn finish(self) -> Result> { + ensure!( + self.completed == completion_mask(self.expert_count), + "expert bank is incomplete: received {}/{} experts", + self.completed.count_ones(), + self.expert_count + ); + Ok(vec![ + ( + self.output_prefix.clone(), + self.weights + .context("expert bank has no packed weights")? + .finish(), + ), + ( + format!("{}_scales", self.output_prefix), + self.scales.context("expert bank has no scales")?.finish(), + ), + ( + format!("{}_biases", self.output_prefix), + self.biases.context("expert bank has no biases")?.finish(), + ), + ]) + } +} + +struct BankTensor { + tensor: OwnedTensor, + expert_shape: Vec, + expert_bytes: usize, +} + +impl BankTensor { + fn new(expert_count: usize, expert: &OwnedTensor) -> Result { + ensure!( + expert.shape.len() == 2, + "expert projection output must be rank 2, got shape {:?}", + expert.shape + ); + let total_bytes = expert + .data + .len() + .checked_mul(expert_count) + .context("expert bank byte count overflow")?; + let mut shape = Vec::with_capacity(3); + shape.push(expert_count); + shape.extend_from_slice(&expert.shape); + Ok(Self { + tensor: OwnedTensor { + dtype: expert.dtype, + shape, + data: vec![0_u8; total_bytes], + }, + expert_shape: expert.shape.clone(), + expert_bytes: expert.data.len(), + }) + } + + fn insert(&mut self, expert_index: usize, expert: OwnedTensor) -> Result<()> { + ensure!( + expert.dtype == self.tensor.dtype, + "expert projection dtype changed from {:?} to {:?}", + self.tensor.dtype, + expert.dtype + ); + ensure!( + expert.shape == self.expert_shape, + "expert projection shape changed from {:?} to {:?}", + self.expert_shape, + expert.shape + ); + ensure!( + expert.data.len() == self.expert_bytes, + "expert projection byte count changed from {} to {}", + self.expert_bytes, + expert.data.len() + ); + let start = expert_index + .checked_mul(self.expert_bytes) + .context("expert bank slice offset overflow")?; + let end = start + .checked_add(self.expert_bytes) + .context("expert bank slice end overflow")?; + self.tensor.data[start..end].copy_from_slice(&expert.data); + Ok(()) + } + + fn finish(self) -> OwnedTensor { + self.tensor + } +} + +fn insert_tensor( + bank: &mut Option, + expert_index: usize, + expert_count: usize, + expert: OwnedTensor, +) -> Result<()> { + if bank.is_none() { + *bank = Some(BankTensor::new(expert_count, &expert)?); + } + bank.as_mut() + .expect("expert bank initialized above") + .insert(expert_index, expert) +} + +const fn completion_mask(expert_count: usize) -> u128 { + if expert_count == MAX_EXPERTS { + u128::MAX + } else { + (1_u128 << expert_count) - 1 + } +} + +#[cfg(test)] +mod tests { + use std::collections::BTreeMap; + + use half::bf16; + use safemlx::{Array, Device, DeviceType, Stream, ops::stack_axis, transforms::eval}; + use safemlx_lm::quantization::{AffineQuantization, quantize_tensor}; + + use super::*; + + #[test] + fn incremental_bank_matches_quantizing_the_stacked_experts() { + let stream = Stream::new_with_device(&Device::new(DeviceType::Gpu, 0)); + let quantization = AffineQuantization::new(64, 4).unwrap(); + let experts = (0..2) + .map(|expert| { + let values = (0..(7 * 64)) + .map(|index| { + let value = (expert * 449 + index) as f32; + bf16::from_f32(value.mul_add(0.03125, -9.0)) + }) + .collect::>(); + Array::from_slice(&values, &[7, 64]) + }) + .collect::>(); + + let mut actual = AffineExpertBankAssembler::new("experts.up_proj", 2).unwrap(); + for expert_index in [1, 0] { + let quantized = quantize_tensor(&experts[expert_index], quantization, &stream).unwrap(); + actual.insert(expert_index, quantized, &stream).unwrap(); + } + let actual = actual + .finish() + .unwrap() + .into_iter() + .collect::>(); + + let stacked = stack_axis(&experts, 0, &stream).unwrap(); + let expected = quantize_tensor(&stacked, quantization, &stream).unwrap(); + let biases = expected.biases.as_ref().unwrap(); + eval([&expected.weight, &expected.scales, biases]).unwrap(); + stream.synchronize().unwrap(); + let expected = BTreeMap::from([ + ("experts.up_proj", owned_tensor(&expected.weight).unwrap()), + ( + "experts.up_proj_scales", + owned_tensor(&expected.scales).unwrap(), + ), + ("experts.up_proj_biases", owned_tensor(biases).unwrap()), + ]); + + assert_eq!( + actual.keys().collect::>(), + expected.keys().collect::>() + ); + for (name, expected) in expected { + let actual = &actual[name]; + assert_eq!(actual.dtype, expected.dtype, "{name} dtype"); + assert_eq!(actual.shape, expected.shape, "{name} shape"); + assert_eq!(actual.data, expected.data, "{name} bytes"); + } + } + + #[test] + fn incomplete_and_duplicate_banks_fail_closed() { + let stream = Stream::new_with_device(&Device::new(DeviceType::Gpu, 0)); + let quantization = AffineQuantization::new(64, 4).unwrap(); + let dense = Array::from_slice(&vec![bf16::from_f32(0.5); 3 * 64], &[3, 64]); + let mut bank = AffineExpertBankAssembler::new("experts.down_proj", 2).unwrap(); + bank.insert( + 0, + quantize_tensor(&dense, quantization, &stream).unwrap(), + &stream, + ) + .unwrap(); + let duplicate = bank + .insert( + 0, + quantize_tensor(&dense, quantization, &stream).unwrap(), + &stream, + ) + .unwrap_err(); + assert!(duplicate.to_string().contains("more than once")); + let incomplete = bank.finish().err().expect("bank should be incomplete"); + assert!(incomplete.to_string().contains("incomplete")); + } +} diff --git a/crates/skippy-engine-mlx/src/derived/nemotron_h.rs b/crates/skippy-engine-mlx/src/derived/nemotron_h.rs new file mode 100644 index 0000000000..4132bbdb82 --- /dev/null +++ b/crates/skippy-engine-mlx/src/derived/nemotron_h.rs @@ -0,0 +1,414 @@ +//! Nemotron-H checkpoint rewriting and bounded routed-expert conversion. + +use std::{collections::BTreeMap, path::Path}; + +use anyhow::{Context, Result, ensure}; +use half::bf16; +use safemlx::{ + Array, Device, DeviceType, Stream, + memory::{active_memory, cache_memory, peak_memory, reset_peak_memory}, + module::Module, + transforms::eval, +}; +use safemlx_lm::{ + models::nemotron_h::{BlockInput, TransformerBlock, get_nemotron_h_model_args}, + quantization::{WeightQuantization, quantize_tensor}, + weights::{StrictLoadConfig, StrictLoadReport, load_safetensors_dir_strict}, +}; +use serde::{Deserialize, Serialize}; +use serde_json::Value; + +use super::{OwnedTensor, expert_bank::AffineExpertBankAssembler}; + +#[derive(Clone, Copy, Debug, Eq, Ord, PartialEq, PartialOrd)] +enum Projection { + Up, + Down, +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +enum LayerKind { + Mamba, + Attention, + Mlp, + Moe, +} + +impl LayerKind { + fn from_marker(marker: char) -> Result { + match marker { + 'M' => Ok(Self::Mamba), + '*' => Ok(Self::Attention), + '-' => Ok(Self::Mlp), + 'E' => Ok(Self::Moe), + other => { + anyhow::bail!("Nemotron-H layer pattern contains unsupported marker {other:?}") + } + } + } + + const fn parameter_field(self) -> &'static str { + match self { + Self::Mamba => "mamba", + Self::Attention => "attention", + Self::Mlp => "mlp", + Self::Moe => "moe", + } + } +} + +#[derive(Debug, Deserialize)] +struct SourceConfig { + hidden_size: i32, + num_hidden_layers: usize, + hybrid_override_pattern: String, + n_routed_experts: i32, + moe_intermediate_size: i32, +} + +pub(super) struct NemotronHDerivation { + hidden_size: i32, + intermediate_size: i32, + expert_count: usize, + layer_kinds: Vec, + banks: BTreeMap<(usize, Projection), AffineExpertBankAssembler>, +} + +/// Evidence that one derived Nemotron-H MoE layer strict-loads and executes. +#[derive(Clone, Debug, Serialize)] +pub struct MlxNemotronHValidationReport { + pub model_dir: std::path::PathBuf, + pub layer: usize, + pub input_shape: Vec, + pub output_shape: Vec, + pub output_is_finite: bool, + pub mlx_active_memory_bytes: usize, + pub mlx_cache_memory_bytes: usize, + pub mlx_peak_memory_bytes: usize, +} + +/// Strict-loads a derived affine Nemotron-H MoE block and runs a deterministic +/// nonzero hidden state through it. +pub fn validate_nemotron_h_moe_stage( + model_dir: impl AsRef, + layer: usize, +) -> Result { + let model_dir = model_dir.as_ref(); + let args = get_nemotron_h_model_args(model_dir)?; + ensure!( + args.hybrid_override_pattern + .chars() + .nth(layer) + .is_some_and(|marker| marker == 'E'), + "Nemotron-H layer {layer} is not an MoE layer" + ); + reset_peak_memory()?; + let stream = Stream::new_with_device(&Device::new(DeviceType::Gpu, 0)); + let mut block = TransformerBlock::new(&args, layer, &stream)?; + let load_config = StrictLoadConfig::default().strip_prefix(format!("model.layers.{layer}.")); + let mut load_report = StrictLoadReport::default(); + load_safetensors_dir_strict( + &mut block, + model_dir, + &stream, + &load_config, + &mut load_report, + )?; + load_report.finish(&block, &load_config)?; + + let values = (0..args.hidden_size) + .map(|index| bf16::from_f32(((index % 31) as f32 - 15.0) / 32.0)) + .collect::>(); + let input = Array::from_slice(&values, &[1, 1, args.hidden_size]); + let output = block.forward( + BlockInput { + x: &input, + mask: None, + cache: None, + }, + &stream, + )?; + let finite = output.is_finite(&stream)?.all(None, &stream)?; + eval([&output, &finite])?; + stream.synchronize()?; + let output_is_finite = finite.try_item::(&stream)?; + ensure!( + output_is_finite, + "Nemotron-H layer output contains non-finite values" + ); + + Ok(MlxNemotronHValidationReport { + model_dir: model_dir.to_path_buf(), + layer, + input_shape: input.shape().to_vec(), + output_shape: output.shape().to_vec(), + output_is_finite, + mlx_active_memory_bytes: active_memory()?, + mlx_cache_memory_bytes: cache_memory()?, + mlx_peak_memory_bytes: peak_memory()?, + }) +} + +impl NemotronHDerivation { + pub(super) fn new(config: &Value, layer_start: u32, layer_end: u32) -> Result { + ensure!( + config.get("layers_block_type").is_none() && config.get("moe_latent_size").is_none(), + "Nemotron-H latent-MoE/Ultra configs require a separate runtime family" + ); + let source: SourceConfig = + serde_json::from_value(config.clone()).context("parse Nemotron-H derivation config")?; + ensure!( + source.hidden_size > 0, + "Nemotron-H hidden_size must be non-zero" + ); + ensure!( + source.moe_intermediate_size > 0, + "Nemotron-H moe_intermediate_size must be non-zero" + ); + ensure!( + source.n_routed_experts > 0 && source.n_routed_experts <= i32::try_from(u128::BITS)?, + "Nemotron-H n_routed_experts must be in 1..={} for bounded assembly", + u128::BITS + ); + let layer_kinds = source + .hybrid_override_pattern + .chars() + .map(LayerKind::from_marker) + .collect::>>()?; + ensure!( + layer_kinds.len() == source.num_hidden_layers, + "Nemotron-H layer pattern has {} entries, expected {}", + layer_kinds.len(), + source.num_hidden_layers + ); + let start = usize::try_from(layer_start)?; + let end = usize::try_from(layer_end)?; + ensure!( + end <= layer_kinds.len(), + "Nemotron-H stage exceeds layer pattern" + ); + ensure!( + end == start + 1 && layer_kinds[start] == LayerKind::Moe, + "bounded Nemotron-H derivation currently requires exactly one MoE layer" + ); + let expert_count = usize::try_from(source.n_routed_experts)?; + + let mut banks = BTreeMap::new(); + for (layer, kind) in layer_kinds.iter().enumerate().take(end).skip(start) { + if *kind != LayerKind::Moe { + continue; + } + for projection in [Projection::Up, Projection::Down] { + let suffix = match projection { + Projection::Up => "up_proj", + Projection::Down => "down_proj", + }; + banks.insert( + (layer, projection), + AffineExpertBankAssembler::new( + format!("model.layers.{layer}.moe.experts.{suffix}"), + expert_count, + )?, + ); + } + } + Ok(Self { + hidden_size: source.hidden_size, + intermediate_size: source.moe_intermediate_size, + expert_count, + layer_kinds, + banks, + }) + } + + /// Consumes a split routed-expert matrix when `source_name` names one. + pub(super) fn consume_expert( + &mut self, + source_name: &str, + dense: &Array, + quantization: WeightQuantization, + stream: &Stream, + ) -> Result { + let Some(expert) = parse_expert_name(source_name)? else { + return Ok(false); + }; + ensure!( + expert.expert < self.expert_count, + "Nemotron-H expert {} exceeds configured count {}", + expert.expert, + self.expert_count + ); + let expected_shape = match expert.projection { + Projection::Up => [self.intermediate_size, self.hidden_size], + Projection::Down => [self.hidden_size, self.intermediate_size], + }; + ensure!( + dense.shape() == expected_shape, + "Nemotron-H expert source {source_name} has shape {:?}, expected {:?}", + dense.shape(), + expected_shape + ); + let bank = self + .banks + .get_mut(&(expert.layer, expert.projection)) + .with_context(|| { + format!("Nemotron-H expert {source_name} does not belong to a selected MoE layer") + })?; + bank.insert( + expert.expert, + quantize_tensor(dense, quantization, stream)?, + stream, + )?; + Ok(true) + } + + pub(super) fn rewrite_name(&self, source_name: &str) -> Result { + let layer_root = source_name + .strip_prefix("backbone.layers.") + .or_else(|| source_name.strip_prefix("model.backbone.layers.")); + if let Some(rest) = layer_root { + let (layer, suffix) = rest + .split_once('.') + .with_context(|| format!("invalid Nemotron-H layer tensor {source_name}"))?; + let layer = layer.parse::().with_context(|| { + format!("invalid Nemotron-H layer index in tensor {source_name}") + })?; + let kind = self + .layer_kinds + .get(layer) + .with_context(|| format!("Nemotron-H tensor layer {layer} is out of range"))?; + if let Some(mixer_suffix) = suffix.strip_prefix("mixer.") { + return Ok(format!( + "model.layers.{layer}.{}.{mixer_suffix}", + kind.parameter_field() + )); + } + return Ok(format!("model.layers.{layer}.{suffix}")); + } + if let Some(rest) = source_name.strip_prefix("model.backbone.") { + return Ok(format!("model.{rest}")); + } + if let Some(rest) = source_name.strip_prefix("backbone.") { + return Ok(format!("model.{rest}")); + } + Ok(source_name.to_string()) + } + + pub(super) fn keep_dense(output_name: &str) -> bool { + output_name.ends_with(".moe.gate.weight") + } + + pub(super) fn finish(self) -> Result> { + self.banks + .into_values() + .try_fold(Vec::new(), |mut output, bank| { + output.extend(bank.finish()?); + Ok(output) + }) + } +} + +struct ExpertSource { + layer: usize, + expert: usize, + projection: Projection, +} + +fn parse_expert_name(name: &str) -> Result> { + let root = name + .strip_prefix("backbone.layers.") + .or_else(|| name.strip_prefix("model.backbone.layers.")); + let Some(rest) = root else { + return Ok(None); + }; + let Some((layer, rest)) = rest.split_once(".mixer.experts.") else { + return Ok(None); + }; + let (expert, projection) = rest + .split_once('.') + .with_context(|| format!("invalid split Nemotron-H expert tensor {name}"))?; + let projection = match projection { + "up_proj.weight" => Projection::Up, + "down_proj.weight" => Projection::Down, + other => anyhow::bail!("unsupported split Nemotron-H expert projection {other:?}"), + }; + Ok(Some(ExpertSource { + layer: layer + .parse() + .with_context(|| format!("invalid Nemotron-H expert layer in {name}"))?, + expert: expert + .parse() + .with_context(|| format!("invalid Nemotron-H expert index in {name}"))?, + projection, + })) +} + +#[cfg(test)] +mod tests { + use serde_json::json; + + use super::*; + + fn config() -> Value { + json!({ + "model_type": "nemotron_h", + "hidden_size": 64, + "num_hidden_layers": 4, + "hybrid_override_pattern": "ME-*", + "n_routed_experts": 2, + "moe_intermediate_size": 32 + }) + } + + #[test] + fn rewrites_public_layer_keys_to_runtime_fields() { + let derivation = NemotronHDerivation::new(&config(), 1, 2).unwrap(); + assert_eq!( + derivation + .rewrite_name("backbone.layers.0.mixer.in_proj.weight") + .unwrap(), + "model.layers.0.mamba.in_proj.weight" + ); + assert_eq!( + derivation + .rewrite_name("backbone.layers.1.mixer.gate.weight") + .unwrap(), + "model.layers.1.moe.gate.weight" + ); + assert_eq!( + derivation + .rewrite_name("backbone.layers.2.mixer.up_proj.weight") + .unwrap(), + "model.layers.2.mlp.up_proj.weight" + ); + assert_eq!( + derivation + .rewrite_name("backbone.layers.3.mixer.q_proj.weight") + .unwrap(), + "model.layers.3.attention.q_proj.weight" + ); + assert_eq!( + derivation + .rewrite_name("backbone.embeddings.weight") + .unwrap(), + "model.embeddings.weight" + ); + } + + #[test] + fn parses_split_experts_and_rejects_ultra_schema() { + let expert = parse_expert_name("backbone.layers.1.mixer.experts.127.down_proj.weight") + .unwrap() + .unwrap(); + assert_eq!(expert.layer, 1); + assert_eq!(expert.expert, 127); + assert_eq!(expert.projection, Projection::Down); + + let mut ultra = config(); + ultra["layers_block_type"] = json!(["mamba"]); + let error = NemotronHDerivation::new(&ultra, 0, 1) + .err() + .expect("Ultra schema should fail"); + assert!(error.to_string().contains("separate runtime family")); + } +} diff --git a/crates/skippy-engine-mlx/src/lib.rs b/crates/skippy-engine-mlx/src/lib.rs index 0a5bec45c7..77db2a2500 100644 --- a/crates/skippy-engine-mlx/src/lib.rs +++ b/crates/skippy-engine-mlx/src/lib.rs @@ -21,8 +21,9 @@ pub use backend::MlxBackend; #[cfg(all(feature = "mlx", target_os = "macos"))] pub use derived::{ MlxDerivationControl, MlxDerivedStageCacheConfig, MlxDerivedStageCacheResult, - MlxDerivedStageConfig, MlxDerivedStageReport, MlxDerivedStageShard, derive_quantized_stage, - derive_quantized_stage_cached, load_prepared_quantized_stage, mlx_derived_stage_cache_root, + MlxDerivedStageConfig, MlxDerivedStageReport, MlxDerivedStageShard, + MlxNemotronHValidationReport, derive_quantized_stage, derive_quantized_stage_cached, + load_prepared_quantized_stage, mlx_derived_stage_cache_root, validate_nemotron_h_moe_stage, }; #[cfg(all(feature = "mlx", target_os = "macos"))] pub use engine::{ChatTurn, GenerateRequest, MlxEngine, MlxEngineConfig}; diff --git a/docs/design/MLX_STAGE_ENGINE_PLAN.md b/docs/design/MLX_STAGE_ENGINE_PLAN.md index 0b0b207dd4..d77e647195 100644 --- a/docs/design/MLX_STAGE_ENGINE_PLAN.md +++ b/docs/design/MLX_STAGE_ENGINE_PLAN.md @@ -95,6 +95,26 @@ For pinned NVIDIA 30B-A3B Base BF16 layer `1`, it selected 261 tensors totaling 19,955,712 bytes. No tensor payload was fetched by this proof. This removes the acquisition blocker for the next bounded expert-pack experiment. +**Update — one real Nano MoE layer now derives and executes.** The production +builder recognizes split ReLU2 experts, quantizes one expert matrix at a time, +and writes each result into preallocated rank-3 affine banks. A two-expert test +proves this incremental layout is byte-identical to quantizing a stacked bank. +For pinned layer `1`, 2,594,936,576 source bytes became 730,324,736 affine4 +tensor bytes (258 source matrices quantized, three tensors dense) in 199.70 +seconds. Maximum RSS was 822,165,504 bytes; the only source payload retained at +any instant was one tensor, at most 19,955,848 bytes. A validation command then +strict-loaded safemlx's actual layer-1 `TransformerBlock` and executed a finite +nonzero `[1, 1, 2688] -> [1, 1, 2688]` forward pass with 811,696,812 bytes of +reported MLX peak memory. General hybrid-stage execution is still gated on +recurrent/attention state and boundary work. + +The derivation memory bound is the final packed routed bank, not one expert: +six preallocated payload buffers total 718,405,632 bytes. Moving those buffers +to a disk-backed random-write spool is the next step if preparation RSS must +stay near the largest individual source tensor. The finite forward plus strict +parameter coverage proves artifact assembly and executability; quantization +quality still needs a dense/reference parity oracle. + Nemotron 3 Ultra is not that next executable target. Its current public config uses a 108-layer `layers_block_type` latent-MoE design with 512 experts and `moe_latent_size=2048`, while the pinned safemlx `nemotron_h` implementation is @@ -880,11 +900,11 @@ Spikes 1 and 2 are more decisive than any standalone token/s benchmark. the dense llama adapter and MLX two-process binary-wire proof are complete. 3. Run **Spike 2 (boundary fence)** at frontier residual widths and keep it as a go/no-go gate. -4. Use Nemotron-H **Nano** as the first frontier-family follow-up: wire its - existing public layer/cache structures and affine expert runtime to - public-checkpoint on-load packing. Do not treat Ultra as the same runtime - family. Then expose safemlx's existing Inkling implementation as a staged - text decoder and use Transformers as the parity oracle. +4. Extend the proven single-layer Nemotron-H **Nano** affine expert path into a + hybrid staged runtime with explicit recurrent/attention boundary state. Do + not treat Ultra as the same runtime family. Then expose safemlx's existing + Inkling implementation as a staged text decoder and use Transformers as the + parity oracle. --- diff --git a/spikes/mlx-safetensors-stages/FINDINGS.md b/spikes/mlx-safetensors-stages/FINDINGS.md index d5bfd4c885..2170582fc0 100644 --- a/spikes/mlx-safetensors-stages/FINDINGS.md +++ b/spikes/mlx-safetensors-stages/FINDINGS.md @@ -224,6 +224,21 @@ BF16 checkpoint selects 261 tensors / 2,594,936,576 bytes from a single coalesced payload ranges. The config contains bare `Infinity`, so production metadata parsing uses strict JSON first and a JSON5 fallback. +The next production proof now exists too. The exact layer ranges were consumed +one tensor at a time and converted into an affine4/g64 artifact without ever +constructing the dense 128-expert bank. The run quantized 258 matrices, copied +three dense tensors, and reduced 2,594,936,576 source bytes to 730,324,736 +tensor bytes in 199.70 seconds. Maximum RSS was 822,165,504 bytes; the largest +ephemeral source tensor file was 19,955,848 bytes. The six routed bank tensors +use underscore companions required by safemlx, while shared expert matrices use +normal dotted affine companions. The artifact strict-loaded into the actual +Nano layer-1 block and produced finite `[1, 1, 2688]` output. + +This is bounded by the final packed layer rather than one-expert RAM: the six +routed buffers total 718,405,632 bytes, consistent with the measured RSS. A +disk-backed spool would lower derivation memory further. Strict loading and a +finite forward prove executable assembly, not dense-versus-affine accuracy. + The format mechanism is stable and documented by the [SafeTensors format](https://github.com/huggingface/safetensors#format): the header records each tensor's dtype, shape, and byte offsets. Hugging Face's @@ -444,9 +459,9 @@ because it does not alter the derived packed weights. 1. Add capacity/eviction ownership to the host derived-stage cache, then decide whether a local request-to-recipe locator should remove even the warm metadata probes. -2. Quantize one real Nemotron-H Nano BF16 matrix reproducibly, then implement - one complete split-expert bank without accumulating every dense expert. Keep - Ultra gated behind its separate latent-MoE family implementation. +2. Extend the proven single-layer Nemotron-H Nano artifact into a hybrid stage + while making recurrent/attention state explicit on the wire. Keep Ultra + gated behind its separate latent-MoE family implementation. 3. Measure the MLX eval/readback/codec boundary fence independently at frontier residual widths and prefill sizes. 4. Expose the existing safemlx Inkling text decoder as one stage, prove From 9a5e9ace0bdd4e6882ead14b483aeeb29d641adf Mon Sep 17 00:00:00 2001 From: Michael Neale <14976+michaelneale@users.noreply.github.com> Date: Fri, 17 Jul 2026 22:15:15 +1000 Subject: [PATCH 21/37] feat(mlx): execute Nemotron-H MoE stages --- crates/skippy-engine-mlx/STAGED_EXECUTION.md | 45 +- crates/skippy-engine-mlx/src/bin/mlx-stage.rs | 13 + crates/skippy-engine-mlx/src/derived.rs | 1 + .../src/derived/nemotron_h.rs | 62 ++- crates/skippy-engine-mlx/src/lib.rs | 5 +- crates/skippy-engine-mlx/src/stage.rs | 81 +++- .../skippy-engine-mlx/src/stage/nemotron_h.rs | 388 ++++++++++++++++++ docs/design/MLX_STAGE_ENGINE_PLAN.md | 39 +- spikes/mlx-safetensors-stages/FINDINGS.md | 21 +- 9 files changed, 596 insertions(+), 59 deletions(-) create mode 100644 crates/skippy-engine-mlx/src/stage/nemotron_h.rs diff --git a/crates/skippy-engine-mlx/STAGED_EXECUTION.md b/crates/skippy-engine-mlx/STAGED_EXECUTION.md index dc8cc7f4e5..a67ff4d800 100644 --- a/crates/skippy-engine-mlx/STAGED_EXECUTION.md +++ b/crates/skippy-engine-mlx/STAGED_EXECUTION.md @@ -16,10 +16,10 @@ the default automatic mesh launch path: - `skippy-server::llama_engine` proves the existing llama `RuntimeState` can implement the same dense contract, including F16/BF16/F32 residual conversion and checkpoint/restore/trim delegation, without changing the native ABI. -- `MlxStageEngine` loads one materialized partial SafeTensors file or sharded - directory, owns - per-session KV caches on a dedicated MLX worker thread, and executes only its - configured layer range. +- `MlxStageEngine` auto-detects the materialized SafeTensors family. Dense + Llama stages own per-session KV caches; the first frontier adapter executes + one internal, stateless Nemotron-H Nano MoE layer. MLX objects remain on a + dedicated worker thread. - `mlx-stage` starts a stage process or drives a chain as a proof client. - `StagePrepare` / `StageLoad` with `backend=mlx` and an immutable `hf-model://org/repo@` reference now derive or reuse a validated @@ -171,9 +171,10 @@ and first/final boundary tensors. Against pinned `nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-Base-BF16` layer `1`, it selected 2,594,936,576 bytes in 261 tensors from one 4,991,210,024-byte shard; the largest individual tensor was 19,955,712 bytes. This is metadata/range-planning -evidence only for the general family layout. The derived builder now supports -exactly one Nemotron-H Nano MoE layer at a time; the serving stage engine -remains fail-closed until hybrid recurrent/attention boundaries are implemented. +evidence only for the general family layout. The derived builder and stage +engine now support exactly one internal Nemotron-H Nano MoE layer at a time. +Mamba, attention, first/final boundaries, and multi-layer hybrid stages remain +fail-closed until their state and boundary semantics are implemented. Reproduce the metadata-only proof (it downloads the pinned config, index, and one SafeTensors header, but no tensor payloads): @@ -191,6 +192,14 @@ requests, quantized 258 matrices while retaining three dense tensors, and wrote largest ephemeral source tensor file was 19,955,848 bytes. The resulting artifact strict-loaded into safemlx's real layer-1 `TransformerBlock` and produced a finite `[1, 1, 2688]` output for a deterministic nonzero input. +The same artifact then loaded through `MlxStageEngine`; execution through the +shared F32 `StageActivation` contract matched direct block execution within +`atol=1e-4`, `rtol=1e-4` (across repeated validation runs, worst observed max +absolute difference `1.1920929e-7`, max relative difference `1.8225228e-5` for +reference values above `atol`). It +compared two session IDs, reset session 1, and independently compared its +repeated output too. Separate sparse executions were not bit-identical, so the +validator records both hashes and enforces the declared numerical tolerance. Here, bounded memory means bounded by the final packed layer: the six routed bank buffers total 718,405,632 bytes. It does not mean derivation stays at the one-expert (~20 MB source tensor) footprint. The forward is an executable smoke @@ -206,6 +215,8 @@ just mlx-stage derive \ --weight-quantization affine4 just mlx-stage validate-nemotron-h \ --model /tmp/nemotron-nano-layer1-affine4 --layer 1 +just mlx-stage validate-nemotron-h-stage \ + --model /tmp/nemotron-nano-layer1-affine4 --layer 1 ``` ## Reproduce @@ -275,19 +286,19 @@ just mlx-stage prove --connect 127.0.0.1:19090 --wire-dtype f16 ## Deliberate limitations of this checkpoint -- Dense Llama-family checkpoints only in `MlxStageEngine`. The pinned safemlx - revision has whole-model Inkling and Nemotron-H implementations, but neither - is exposed through this partial-stage adapter yet. -- The derived builder handles ordinary rank-2 Llama weights. Nemotron split - experts need per-layer expert-bank assembly before quantization; Inkling needs - its transformed rank-3 grouped-expert loader. Neither is silently treated as - Llama. +- `MlxStageEngine` supports dense Llama ranges and exactly one internal, + stateless Nemotron-H Nano `E`/MoE layer. It rejects Nemotron Mamba, attention, + dense-MLP, first/final, and multi-layer ranges. Inkling is not exposed through + the partial-stage adapter. +- The derived builder handles ordinary rank-2 Llama weights and one Nano split + expert bank. Inkling still needs its transformed rank-3 grouped-expert loader; + unsupported families are not silently treated as Llama. - The pinned safemlx Nemotron-H implementation matches the 52-layer Nano schema, not Nemotron 3 Ultra's 108-layer latent-MoE schema. Ultra range plans are storage-locality evidence, not executable-family support. -- Bounded Nemotron-H derivation currently accepts exactly one `E`/MoE layer. - It does not yet expose a hybrid multi-layer stage or recurrent state on the - wire. +- Bounded Nemotron-H derivation and execution currently accept exactly one + internal `E`/MoE layer. They do not expose a hybrid multi-layer stage or + recurrent state on the wire. - Greedy sampling only; sampling metadata is preserved in the contract and rejected explicitly when enabled. - No KV page import/export, cache trim/checkpoint, MTP, speculative verify, diff --git a/crates/skippy-engine-mlx/src/bin/mlx-stage.rs b/crates/skippy-engine-mlx/src/bin/mlx-stage.rs index 4e96aeea63..81d9367339 100644 --- a/crates/skippy-engine-mlx/src/bin/mlx-stage.rs +++ b/crates/skippy-engine-mlx/src/bin/mlx-stage.rs @@ -16,6 +16,7 @@ mod real { MlxComputeDtype, MlxDerivationControl, MlxDerivedStageCacheConfig, MlxDerivedStageConfig, MlxStageEngine, MlxStageEngineConfig, MlxWeightQuantization, derive_quantized_stage, derive_quantized_stage_cached, mlx_derived_stage_cache_root, validate_nemotron_h_moe_stage, + validate_nemotron_h_stage_engine, }; use skippy_protocol::binary::{ StageStateHeader, StageWireMessage, WireActivationDType, WireMessageKind, WireReplyKind, @@ -105,6 +106,13 @@ mod real { #[arg(long)] layer: usize, }, + /// Prove StageEngine output numerically matches direct execution. + ValidateNemotronHStage { + #[arg(long)] + model: PathBuf, + #[arg(long)] + layer: usize, + }, /// Drive a stage chain and assert its greedy token sequence. Prove { #[arg(long)] @@ -250,6 +258,11 @@ mod real { println!("{}", serde_json::to_string_pretty(&report)?); Ok(()) } + Command::ValidateNemotronHStage { model, layer } => { + let report = validate_nemotron_h_stage_engine(model, layer)?; + println!("{}", serde_json::to_string_pretty(&report)?); + Ok(()) + } Command::Prove { connect, tokens, diff --git a/crates/skippy-engine-mlx/src/derived.rs b/crates/skippy-engine-mlx/src/derived.rs index 2cd436304c..0a87fcb79e 100644 --- a/crates/skippy-engine-mlx/src/derived.rs +++ b/crates/skippy-engine-mlx/src/derived.rs @@ -29,6 +29,7 @@ use sha2::{Digest, Sha256}; use crate::stage::MlxWeightQuantization; use nemotron_h::NemotronHDerivation; +pub(crate) use nemotron_h::{nemotron_h_validation_values, validate_nemotron_h_moe_stage_output}; mod cache; mod expert_bank; diff --git a/crates/skippy-engine-mlx/src/derived/nemotron_h.rs b/crates/skippy-engine-mlx/src/derived/nemotron_h.rs index 4132bbdb82..a30c402f9d 100644 --- a/crates/skippy-engine-mlx/src/derived/nemotron_h.rs +++ b/crates/skippy-engine-mlx/src/derived/nemotron_h.rs @@ -17,6 +17,7 @@ use safemlx_lm::{ }; use serde::{Deserialize, Serialize}; use serde_json::Value; +use sha2::{Digest, Sha256}; use super::{OwnedTensor, expert_bank::AffineExpertBankAssembler}; @@ -82,6 +83,7 @@ pub struct MlxNemotronHValidationReport { pub input_shape: Vec, pub output_shape: Vec, pub output_is_finite: bool, + pub output_f32_sha256: String, pub mlx_active_memory_bytes: usize, pub mlx_cache_memory_bytes: usize, pub mlx_peak_memory_bytes: usize, @@ -93,6 +95,13 @@ pub fn validate_nemotron_h_moe_stage( model_dir: impl AsRef, layer: usize, ) -> Result { + Ok(validate_nemotron_h_moe_stage_output(model_dir, layer)?.0) +} + +pub(crate) fn validate_nemotron_h_moe_stage_output( + model_dir: impl AsRef, + layer: usize, +) -> Result<(MlxNemotronHValidationReport, Vec)> { let model_dir = model_dir.as_ref(); let args = get_nemotron_h_model_args(model_dir)?; ensure!( @@ -116,10 +125,9 @@ pub fn validate_nemotron_h_moe_stage( )?; load_report.finish(&block, &load_config)?; - let values = (0..args.hidden_size) - .map(|index| bf16::from_f32(((index % 31) as f32 - 15.0) / 32.0)) - .collect::>(); - let input = Array::from_slice(&values, &[1, 1, args.hidden_size]); + let values = nemotron_h_validation_values(args.hidden_size); + let input = Array::from_slice(&values, &[1, 1, args.hidden_size]) + .as_dtype(safemlx::Dtype::Bfloat16, &stream)?; let output = block.forward( BlockInput { x: &input, @@ -128,25 +136,47 @@ pub fn validate_nemotron_h_moe_stage( }, &stream, )?; - let finite = output.is_finite(&stream)?.all(None, &stream)?; - eval([&output, &finite])?; + let output_f32 = output.as_dtype(safemlx::Dtype::Float32, &stream)?; + let finite = output_f32.is_finite(&stream)?.all(None, &stream)?; + eval([&output_f32, &finite])?; stream.synchronize()?; let output_is_finite = finite.try_item::(&stream)?; ensure!( output_is_finite, "Nemotron-H layer output contains non-finite values" ); + let evaluated_output = output_f32.evaluated()?; + let output_values = evaluated_output.as_slice::().to_vec(); + let output_f32_sha256 = f32_sha256(&output_values); - Ok(MlxNemotronHValidationReport { - model_dir: model_dir.to_path_buf(), - layer, - input_shape: input.shape().to_vec(), - output_shape: output.shape().to_vec(), - output_is_finite, - mlx_active_memory_bytes: active_memory()?, - mlx_cache_memory_bytes: cache_memory()?, - mlx_peak_memory_bytes: peak_memory()?, - }) + Ok(( + MlxNemotronHValidationReport { + model_dir: model_dir.to_path_buf(), + layer, + input_shape: input.shape().to_vec(), + output_shape: output.shape().to_vec(), + output_is_finite, + output_f32_sha256, + mlx_active_memory_bytes: active_memory()?, + mlx_cache_memory_bytes: cache_memory()?, + mlx_peak_memory_bytes: peak_memory()?, + }, + output_values, + )) +} + +pub(crate) fn nemotron_h_validation_values(hidden_size: i32) -> Vec { + (0..hidden_size) + .map(|index| bf16::from_f32(((index % 31) as f32 - 15.0) / 32.0).to_f32()) + .collect() +} + +fn f32_sha256(values: &[f32]) -> String { + let mut hasher = Sha256::new(); + for value in values { + hasher.update(value.to_le_bytes()); + } + format!("{:x}", hasher.finalize()) } impl NemotronHDerivation { diff --git a/crates/skippy-engine-mlx/src/lib.rs b/crates/skippy-engine-mlx/src/lib.rs index 77db2a2500..d8603c4325 100644 --- a/crates/skippy-engine-mlx/src/lib.rs +++ b/crates/skippy-engine-mlx/src/lib.rs @@ -28,7 +28,10 @@ pub use derived::{ #[cfg(all(feature = "mlx", target_os = "macos"))] pub use engine::{ChatTurn, GenerateRequest, MlxEngine, MlxEngineConfig}; #[cfg(all(feature = "mlx", target_os = "macos"))] -pub use stage::{MlxComputeDtype, MlxStageEngine, MlxStageEngineConfig, MlxWeightQuantization}; +pub use stage::{ + MlxComputeDtype, MlxNemotronHStageValidationReport, MlxStageEngine, MlxStageEngineConfig, + MlxWeightQuantization, validate_nemotron_h_stage_engine, +}; /// True when this build actually contains the MLX engine. pub const fn mlx_available() -> bool { diff --git a/crates/skippy-engine-mlx/src/stage.rs b/crates/skippy-engine-mlx/src/stage.rs index de007db717..7de9a14271 100644 --- a/crates/skippy-engine-mlx/src/stage.rs +++ b/crates/skippy-engine-mlx/src/stage.rs @@ -1,5 +1,7 @@ //! Partial-layer MLX implementation of the engine-neutral Skippy stage contract. +mod nemotron_h; + use std::{collections::BTreeMap, path::PathBuf, sync::mpsc, thread}; use anyhow::{Context, Result, anyhow, bail, ensure}; @@ -23,6 +25,9 @@ use skippy_engine::{ StageExecutionRequest, }; +use self::nemotron_h::NemotronHMoeStage; +pub use self::nemotron_h::{MlxNemotronHStageValidationReport, validate_nemotron_h_stage_engine}; + #[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] pub enum MlxComputeDtype { F16, @@ -141,7 +146,7 @@ impl StageEngine for MlxStageEngine { } } -struct LoadedStage { +struct LlamaLoadedStage { model: llama::Model, stream: Stream, compute_dtype: Dtype, @@ -150,6 +155,36 @@ struct LoadedStage { sessions: BTreeMap>>, } +enum LoadedStage { + Llama(Box), + NemotronH(Box), +} + +impl LoadedStage { + fn info(&self) -> &StageEngineInfo { + match self { + Self::Llama(stage) => &stage.info, + Self::NemotronH(stage) => stage.info(), + } + } + + fn execute(&mut self, request: StageExecutionRequest) -> Result { + match self { + Self::Llama(stage) => stage.execute(request), + Self::NemotronH(stage) => stage.execute(request), + } + } + + fn reset_session(&mut self, session_id: u64) { + match self { + Self::Llama(stage) => { + stage.sessions.remove(&session_id); + } + Self::NemotronH(stage) => stage.reset_session(session_id), + } + } +} + fn run_worker( config: MlxStageEngineConfig, job_rx: mpsc::Receiver, @@ -157,7 +192,7 @@ fn run_worker( ) { let mut stage = match load_stage(config) { Ok(stage) => { - let _ = ready_tx.send(Ok(stage.info.clone())); + let _ = ready_tx.send(Ok(stage.info().clone())); stage } Err(error) => { @@ -171,7 +206,7 @@ fn run_worker( let _ = reply.send(stage.execute(request).map_err(|error| format!("{error:#}"))); } WorkerJob::Reset { session_id, reply } => { - stage.sessions.remove(&session_id); + stage.reset_session(session_id); let _ = reply.send(Ok(())); } } @@ -179,6 +214,28 @@ fn run_worker( } fn load_stage(config: MlxStageEngineConfig) -> Result { + if model_type(&config.model_dir)?.as_deref() == Some("nemotron_h") { + return Ok(LoadedStage::NemotronH(Box::new(NemotronHMoeStage::load( + config, + )?))); + } + Ok(LoadedStage::Llama(Box::new(load_llama_stage(config)?))) +} + +fn model_type(model_dir: &std::path::Path) -> Result> { + let config_path = model_dir.join("config.json"); + let config: serde_json::Value = serde_json::from_slice( + &std::fs::read(&config_path) + .with_context(|| format!("read MLX stage config {}", config_path.display()))?, + ) + .with_context(|| format!("parse MLX stage config {}", config_path.display()))?; + Ok(config + .get("model_type") + .and_then(serde_json::Value::as_str) + .map(str::to_owned)) +} + +fn load_llama_stage(config: MlxStageEngineConfig) -> Result { let stream = Stream::new_with_device(&Device::new(DeviceType::Gpu, 0)); let weights_stream = Stream::new_with_device(&Device::new(DeviceType::Cpu, 0)); let mut model_args = llama::get_llama_model_args(&config.model_dir)?; @@ -251,7 +308,7 @@ fn load_stage(config: MlxStageEngineConfig) -> Result { model.parameters().flatten().len(), weight_quantization_label, ); - Ok(LoadedStage { + Ok(LlamaLoadedStage { model, stream, compute_dtype: config.compute_dtype.mlx(), @@ -311,7 +368,7 @@ fn copy_stage_weights_to_compute_stream( Ok(()) } -impl LoadedStage { +impl LlamaLoadedStage { fn execute(&mut self, request: StageExecutionRequest) -> Result { if request.kind == StageExecutionKind::Verify { bail!("MLX dense stage verification is not implemented yet"); @@ -506,6 +563,20 @@ fn last_argmax(logits: &Array, stream: &Stream) -> Result { mod tests { use super::*; + #[test] + fn detects_nemotron_h_stage_family_from_config() { + let dir = tempfile::tempdir().unwrap(); + std::fs::write( + dir.path().join("config.json"), + br#"{"model_type":"nemotron_h"}"#, + ) + .unwrap(); + assert_eq!( + model_type(dir.path()).unwrap().as_deref(), + Some("nemotron_h") + ); + } + #[test] fn validates_requested_weight_quantization_before_model_load() { let affine = MlxWeightQuantization::Affine { diff --git a/crates/skippy-engine-mlx/src/stage/nemotron_h.rs b/crates/skippy-engine-mlx/src/stage/nemotron_h.rs new file mode 100644 index 0000000000..cb2d286aa2 --- /dev/null +++ b/crates/skippy-engine-mlx/src/stage/nemotron_h.rs @@ -0,0 +1,388 @@ +//! Stateless Nemotron-H MoE blocks behind the shared stage-engine contract. + +use std::path::{Path, PathBuf}; + +use anyhow::{Context, Result, bail, ensure}; +use safemlx::{ + Array, Device, DeviceType, Dtype, Stream, + memory::{active_memory, cache_memory, peak_memory, reset_peak_memory}, + module::{Module, ModuleParameters, ModuleParametersExt}, +}; +use safemlx_lm::{ + models::nemotron_h::{BlockInput, LayerBlockType, TransformerBlock, get_nemotron_h_model_args}, + weights::{StrictLoadConfig, StrictLoadReport, load_safetensors_dir_strict}, +}; +use serde::Serialize; +use sha2::{Digest, Sha256}; +use skippy_engine::{ + StageActivation, StageEngine, StageEngineInfo, StageExecutionKind, StageExecutionOutput, + StageExecutionRequest, +}; + +use crate::derived::{nemotron_h_validation_values, validate_nemotron_h_moe_stage_output}; + +use super::{MlxComputeDtype, MlxStageEngine, MlxStageEngineConfig, array_activation}; + +/// Tolerance-aware direct-block versus shared-stage-contract evidence. +#[derive(Clone, Debug, Serialize)] +pub struct MlxNemotronHStageValidationReport { + pub model_dir: PathBuf, + pub layer: usize, + pub input_shape: Vec, + pub output_shape: Vec, + pub output_is_finite: bool, + pub direct_output_f32_sha256: String, + pub stage_output_f32_sha256: String, + pub output_within_tolerance: bool, + pub cross_session_stable: bool, + pub session_reset_stable: bool, + pub executions_compared: usize, + pub max_abs_diff: f32, + pub max_relative_diff_for_reference_magnitude_above_atol: f32, + pub cross_session_max_abs_diff: f32, + pub cross_session_max_relative_diff_for_reference_magnitude_above_atol: f32, + pub reset_max_abs_diff: f32, + pub reset_max_relative_diff_for_reference_magnitude_above_atol: f32, + pub comparison_atol: f32, + pub comparison_rtol: f32, + pub mlx_active_memory_bytes: usize, + pub mlx_cache_memory_bytes: usize, + pub mlx_peak_memory_bytes: usize, +} + +/// Compares direct safemlx execution with execution through `StageEngine`. +pub fn validate_nemotron_h_stage_engine( + model_dir: impl AsRef, + layer: usize, +) -> Result { + let model_dir = model_dir.as_ref(); + const ATOL: f32 = 1.0e-4; + const RTOL: f32 = 1.0e-4; + let (direct, direct_values) = validate_nemotron_h_moe_stage_output(model_dir, layer)?; + reset_peak_memory()?; + let engine = MlxStageEngine::spawn(MlxStageEngineConfig { + model_dir: model_dir.to_path_buf(), + model_id: "nemotron-h-stage-validation".to_string(), + stage_index: 1, + layer_start: u32::try_from(layer)?, + layer_end: u32::try_from(layer.checked_add(1).context("layer index overflow")?)?, + compute_dtype: MlxComputeDtype::Bf16, + weight_quantization: None, + ctx_size: Some(1), + })?; + let width = usize::try_from(engine.info().activation_width)?; + let values = nemotron_h_validation_values(i32::try_from(width)?); + let first = execute_validation_input(&engine, 1, width, &values)?; + let second_session = execute_validation_input(&engine, 2, width, &values)?; + engine.reset_session(2)?; + engine.reset_session(1)?; + let after_reset = execute_validation_input(&engine, 1, width, &values)?; + let outputs = [first, second_session, after_reset]; + let output_is_finite = outputs + .iter() + .flat_map(StageActivation::values) + .all(f32::is_finite); + ensure!( + output_is_finite, + "Nemotron-H stage-engine output contains non-finite values" + ); + let stage_output_f32_sha256 = bytes_sha256(&outputs[0].f32_le_bytes); + let mut comparison = OutputComparison::accumulator(); + for output in &outputs { + comparison.include(compare_outputs( + &direct_values, + &output.values(), + ATOL, + RTOL, + )?); + } + let cross_session_comparison = + compare_outputs(&outputs[0].values(), &outputs[1].values(), ATOL, RTOL)?; + let cross_session_stable = cross_session_comparison.all_close; + ensure!( + cross_session_stable, + "Nemotron-H output changed across session IDs: max_abs={} max_relative_for_reference_magnitude_above_atol={} atol={ATOL} rtol={RTOL}", + cross_session_comparison.max_abs, + cross_session_comparison.max_relative_for_reference_magnitude_above_atol, + ); + let reset_comparison = compare_outputs(&outputs[0].values(), &outputs[2].values(), ATOL, RTOL)?; + let session_reset_stable = reset_comparison.all_close; + ensure!( + session_reset_stable, + "Nemotron-H output changed after session reset: max_abs={} max_relative_for_reference_magnitude_above_atol={} atol={ATOL} rtol={RTOL}", + reset_comparison.max_abs, + reset_comparison.max_relative_for_reference_magnitude_above_atol, + ); + let output_within_tolerance = comparison.all_close; + ensure!( + output_within_tolerance, + "Nemotron-H stage-engine output differs from direct block execution: max_abs={} max_relative_for_reference_magnitude_above_atol={} atol={ATOL} rtol={RTOL}", + comparison.max_abs, + comparison.max_relative_for_reference_magnitude_above_atol, + ); + Ok(MlxNemotronHStageValidationReport { + model_dir: model_dir.to_path_buf(), + layer, + input_shape: vec![1, 1, width], + output_shape: vec![1, outputs[0].token_count, outputs[0].width], + output_is_finite, + direct_output_f32_sha256: direct.output_f32_sha256, + stage_output_f32_sha256, + output_within_tolerance, + cross_session_stable, + session_reset_stable, + executions_compared: outputs.len(), + max_abs_diff: comparison.max_abs, + max_relative_diff_for_reference_magnitude_above_atol: comparison + .max_relative_for_reference_magnitude_above_atol, + cross_session_max_abs_diff: cross_session_comparison.max_abs, + cross_session_max_relative_diff_for_reference_magnitude_above_atol: + cross_session_comparison.max_relative_for_reference_magnitude_above_atol, + reset_max_abs_diff: reset_comparison.max_abs, + reset_max_relative_diff_for_reference_magnitude_above_atol: reset_comparison + .max_relative_for_reference_magnitude_above_atol, + comparison_atol: ATOL, + comparison_rtol: RTOL, + mlx_active_memory_bytes: active_memory()?, + mlx_cache_memory_bytes: cache_memory()?, + mlx_peak_memory_bytes: peak_memory()?, + }) +} + +fn execute_validation_input( + engine: &MlxStageEngine, + session_id: u64, + width: usize, + values: &[f32], +) -> Result { + engine + .execute(StageExecutionRequest { + session_id, + kind: StageExecutionKind::Prefill, + token_ids: vec![0], + positions: vec![0], + input: Some(StageActivation::from_values(1, width, values)?), + sampling: None, + })? + .activation + .context("Nemotron-H internal stage returned no activation") +} + +struct OutputComparison { + all_close: bool, + max_abs: f32, + max_relative_for_reference_magnitude_above_atol: f32, +} + +impl OutputComparison { + const fn accumulator() -> Self { + Self { + all_close: true, + max_abs: 0.0, + max_relative_for_reference_magnitude_above_atol: 0.0, + } + } + + fn include(&mut self, next: Self) { + self.all_close = self.all_close && next.all_close; + self.max_abs = self.max_abs.max(next.max_abs); + self.max_relative_for_reference_magnitude_above_atol = self + .max_relative_for_reference_magnitude_above_atol + .max(next.max_relative_for_reference_magnitude_above_atol); + } +} + +fn compare_outputs( + direct: &[f32], + staged: &[f32], + atol: f32, + rtol: f32, +) -> Result { + ensure!( + direct.len() == staged.len(), + "direct and staged output lengths differ" + ); + let mut comparison = OutputComparison::accumulator(); + for (&expected, &actual) in direct.iter().zip(staged) { + let abs = (expected - actual).abs(); + let relative = if expected.abs() > atol { + abs / expected.abs() + } else { + 0.0 + }; + comparison.max_abs = comparison.max_abs.max(abs); + comparison.max_relative_for_reference_magnitude_above_atol = comparison + .max_relative_for_reference_magnitude_above_atol + .max(relative); + comparison.all_close = comparison.all_close && abs <= atol + rtol * expected.abs(); + } + Ok(comparison) +} + +fn bytes_sha256(bytes: &[u8]) -> String { + format!("{:x}", Sha256::digest(bytes)) +} + +pub(super) struct NemotronHMoeStage { + block: TransformerBlock, + stream: Stream, + compute_dtype: Dtype, + ctx_size: Option, + info: StageEngineInfo, +} + +impl NemotronHMoeStage { + pub(super) fn load(config: MlxStageEngineConfig) -> Result { + ensure!( + config.compute_dtype == MlxComputeDtype::Bf16, + "Nemotron-H staged execution is currently validated only with BF16 compute" + ); + ensure!( + config.weight_quantization.is_none(), + "Nemotron-H stages must be loaded from an already-derived checkpoint" + ); + ensure!( + config.layer_end == config.layer_start.saturating_add(1), + "Nemotron-H staged execution currently requires exactly one layer" + ); + let args = get_nemotron_h_model_args(&config.model_dir)?; + let layer = usize::try_from(config.layer_start)?; + ensure!( + args.layer_block_types()?.get(layer) == Some(&LayerBlockType::Moe), + "Nemotron-H staged execution currently supports only stateless MoE layers" + ); + let info = StageEngineInfo { + engine: "mlx".to_string(), + model_id: config.model_id, + stage_index: config.stage_index, + layer_start: config.layer_start, + layer_end: config.layer_end, + total_layers: u32::try_from(args.num_hidden_layers)?, + activation_width: u32::try_from(args.hidden_size)?, + }; + info.validate()?; + ensure!( + !info.is_first() && !info.is_final(), + "Nemotron-H MoE proof stages must be internal residual stages" + ); + + let stream = Stream::new_with_device(&Device::new(DeviceType::Gpu, 0)); + let weights_stream = Stream::new_with_device(&Device::new(DeviceType::Cpu, 0)); + let mut block = TransformerBlock::new(&args, layer, &stream)?; + let load_config = + StrictLoadConfig::default().strip_prefix(format!("model.layers.{layer}.")); + let mut load_report = StrictLoadReport::default(); + load_safetensors_dir_strict( + &mut block, + &config.model_dir, + &weights_stream, + &load_config, + &mut load_report, + )?; + load_report.finish(&block, &load_config)?; + block.copy_to_stream(&stream)?; + stream.synchronize()?; + tracing::info!( + model = %info.model_id, + stage = info.stage_index, + layer = info.layer_start, + tensors = block.parameters().flatten().len(), + weight_quantization = "checkpoint", + "MLX Nemotron-H MoE stage loaded", + ); + Ok(Self { + block, + stream, + compute_dtype: config.compute_dtype.mlx(), + ctx_size: config.ctx_size.map(usize::try_from).transpose()?, + info, + }) + } + + pub(super) const fn info(&self) -> &StageEngineInfo { + &self.info + } + + pub(super) fn execute( + &mut self, + request: StageExecutionRequest, + ) -> Result { + if request.kind == StageExecutionKind::Verify { + bail!("MLX Nemotron-H stage verification is not implemented yet"); + } + if request + .sampling + .as_ref() + .is_some_and(|sampling| sampling.enabled()) + { + bail!("MLX staged execution currently supports greedy sampling only"); + } + ensure!(!request.token_ids.is_empty(), "stage request has no tokens"); + let input = request + .input + .as_ref() + .context("Nemotron-H MoE stage requires residual input")?; + ensure!( + input.token_count == request.token_ids.len(), + "input activation token count does not match token sideband" + ); + ensure!( + input.width == self.info.activation_width as usize, + "input activation width mismatch" + ); + if let Some(ctx_size) = self.ctx_size { + ensure!( + input.token_count <= ctx_size, + "MLX stage context limit {ctx_size} exceeded by {} tokens", + input.token_count + ); + } + + let hidden = Array::from_slice( + &input.values(), + &[ + 1, + i32::try_from(input.token_count)?, + i32::try_from(input.width)?, + ], + ) + .as_dtype(self.compute_dtype, &self.stream)?; + let output = self.block.forward( + BlockInput { + x: &hidden, + mask: None, + cache: None, + }, + &self.stream, + )?; + Ok(StageExecutionOutput { + activation: Some(array_activation(&output, &self.stream)?), + predicted_tokens: Vec::new(), + }) + } + + pub(super) fn reset_session(&mut self, _session_id: u64) {} +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn output_comparison_accepts_roundoff_and_rejects_drift() { + let close = compare_outputs( + &[1.0, -0.5, 0.0], + &[1.000_000_1, -0.5, 1.0e-5], + 1.0e-4, + 1.0e-4, + ) + .unwrap(); + assert!(close.all_close); + assert!(close.max_abs <= 1.0e-5); + assert!(close.max_relative_for_reference_magnitude_above_atol < 1.0e-6); + + let drift = compare_outputs(&[1.0, -0.5], &[1.01, -0.5], 1.0e-4, 1.0e-4).unwrap(); + assert!(!drift.all_close); + assert!(drift.max_abs > 0.009); + } +} diff --git a/docs/design/MLX_STAGE_ENGINE_PLAN.md b/docs/design/MLX_STAGE_ENGINE_PLAN.md index d77e647195..fd12489d48 100644 --- a/docs/design/MLX_STAGE_ENGINE_PLAN.md +++ b/docs/design/MLX_STAGE_ENGINE_PLAN.md @@ -95,7 +95,8 @@ For pinned NVIDIA 30B-A3B Base BF16 layer `1`, it selected 261 tensors totaling 19,955,712 bytes. No tensor payload was fetched by this proof. This removes the acquisition blocker for the next bounded expert-pack experiment. -**Update — one real Nano MoE layer now derives and executes.** The production +**Update — one real Nano MoE layer now derives and executes through the shared +stage contract.** The production builder recognizes split ReLU2 experts, quantizes one expert matrix at a time, and writes each result into preallocated rank-3 affine banks. A two-expert test proves this incremental layout is byte-identical to quantizing a stacked bank. @@ -104,16 +105,28 @@ tensor bytes (258 source matrices quantized, three tensors dense) in 199.70 seconds. Maximum RSS was 822,165,504 bytes; the only source payload retained at any instant was one tensor, at most 19,955,848 bytes. A validation command then strict-loaded safemlx's actual layer-1 `TransformerBlock` and executed a finite -nonzero `[1, 1, 2688] -> [1, 1, 2688]` forward pass with 811,696,812 bytes of -reported MLX peak memory. General hybrid-stage execution is still gated on -recurrent/attention state and boundary work. +nonzero `[1, 1, 2688] -> [1, 1, 2688]` forward pass. `MlxStageEngine` now +auto-detects that artifact, strict-loads only the selected block, accepts the +normal F32 residual activation, computes in BF16, and returns F32 residuals. +Its result matched a direct block execution within `atol=1e-4`, `rtol=1e-4`; +across repeated validation runs, worst observed max absolute and relative +differences were `1.1920929e-7` and `1.8225228e-5` (the relative metric excludes +reference magnitudes at or below `atol`). Two session IDs plus an independent +reset/reuse comparison of session 1 all passed. Runtime MLX active/peak memory +was 730,404,608 / 811,763,256 bytes. +The comparison command loads a direct reference and then the stage, so its +process RSS high-water is not representative of a single serving stage. +General hybrid-stage execution is still gated on recurrent/attention state and +boundary work. The derivation memory bound is the final packed routed bank, not one expert: six preallocated payload buffers total 718,405,632 bytes. Moving those buffers to a disk-backed random-write spool is the next step if preparation RSS must stay near the largest individual source tensor. The finite forward plus strict parameter coverage proves artifact assembly and executability; quantization -quality still needs a dense/reference parity oracle. +quality still needs a dense-BF16 or Transformers reference parity oracle; this +new gate proves the stage wrapper agrees with direct execution of the same +affine artifact. Nemotron 3 Ultra is not that next executable target. Its current public config uses a 108-layer `layers_block_type` latent-MoE design with 512 experts and @@ -891,17 +904,15 @@ Spikes 1 and 2 are more decisive than any standalone token/s benchmark. ## 10. Immediate next steps -1. Expose a sequential selected-range callback/temp-artifact seam, add direct - range-to-quantized-cache shards, and measure peak RSS/disk use with OS and - MLX counters. The live model loader's tensor-wise quantization and split - correctness are proven, but physical source copies remain unbounded evidence. -2. Route normal `skippy-server` launch through the engine-neutral contract while - retaining capability-gated llama-only batching/cache/MTP/multimodal paths; - the dense llama adapter and MLX two-process binary-wire proof are complete. +1. Add capacity and eviction ownership to the host derived-stage cache, and + decide whether a local request-to-recipe locator should eliminate warm-path + metadata probes. +2. Teach topology planning and capability advertisement to select MLX stages; + explicit host Prepare/Load and the engine-neutral server lane are proven. 3. Run **Spike 2 (boundary fence)** at frontier residual widths and keep it as a go/no-go gate. -4. Extend the proven single-layer Nemotron-H **Nano** affine expert path into a - hybrid staged runtime with explicit recurrent/attention boundary state. Do +4. Extend the proven single-layer Nemotron-H **Nano** `StageEngine` adapter into + a hybrid staged runtime with explicit recurrent/attention boundary state. Do not treat Ultra as the same runtime family. Then expose safemlx's existing Inkling implementation as a staged text decoder and use Transformers as the parity oracle. diff --git a/spikes/mlx-safetensors-stages/FINDINGS.md b/spikes/mlx-safetensors-stages/FINDINGS.md index 2170582fc0..9b89e92ca3 100644 --- a/spikes/mlx-safetensors-stages/FINDINGS.md +++ b/spikes/mlx-safetensors-stages/FINDINGS.md @@ -232,12 +232,21 @@ tensor bytes in 199.70 seconds. Maximum RSS was 822,165,504 bytes; the largest ephemeral source tensor file was 19,955,848 bytes. The six routed bank tensors use underscore companions required by safemlx, while shared expert matrices use normal dotted affine companions. The artifact strict-loaded into the actual -Nano layer-1 block and produced finite `[1, 1, 2688]` output. +Nano layer-1 block and produced finite `[1, 1, 2688]` output. The shared +`MlxStageEngine` adapter now also auto-detects and loads this one internal +stateless MoE layer. Its F32 residual output matched direct execution of the +same affine block within `atol=1e-4`, `rtol=1e-4` (max absolute +`1.1920929e-7`, max relative `1.8225228e-5` above the `atol` +reference-magnitude floor across repeated runs). Two session IDs and an +independent reset/reuse comparison of session 1 all passed. The different output +hashes show that sparse MLX executions are not bit-identical; the numerical gate +is explicit rather than claiming exactness. This is bounded by the final packed layer rather than one-expert RAM: the six routed buffers total 718,405,632 bytes, consistent with the measured RSS. A -disk-backed spool would lower derivation memory further. Strict loading and a -finite forward prove executable assembly, not dense-versus-affine accuracy. +disk-backed spool would lower derivation memory further. Strict loading, a +finite forward, and stage-wrapper parity prove executable assembly and the mesh +engine seam, not dense-versus-affine accuracy. The format mechanism is stable and documented by the [SafeTensors format](https://github.com/huggingface/safetensors#format): the @@ -459,9 +468,9 @@ because it does not alter the derived packed weights. 1. Add capacity/eviction ownership to the host derived-stage cache, then decide whether a local request-to-recipe locator should remove even the warm metadata probes. -2. Extend the proven single-layer Nemotron-H Nano artifact into a hybrid stage - while making recurrent/attention state explicit on the wire. Keep Ultra - gated behind its separate latent-MoE family implementation. +2. Extend the proven single-layer Nemotron-H Nano `StageEngine` adapter into a + hybrid stage while making recurrent/attention state explicit on the wire. + Keep Ultra gated behind its separate latent-MoE family implementation. 3. Measure the MLX eval/readback/codec boundary fence independently at frontier residual widths and prefill sizes. 4. Expose the existing safemlx Inkling text decoder as one stage, prove From 5493b7707ab8c064c95fb4dc061aa884e555eb06 Mon Sep 17 00:00:00 2001 From: Michael Neale <14976+michaelneale@users.noreply.github.com> Date: Fri, 17 Jul 2026 22:28:28 +1000 Subject: [PATCH 22/37] feat(mlx): prove Nemotron stage binary wire --- crates/skippy-engine-mlx/Cargo.toml | 3 + crates/skippy-engine-mlx/STAGED_EXECUTION.md | 27 ++ crates/skippy-engine-mlx/src/bin/mlx-stage.rs | 21 +- crates/skippy-engine-mlx/src/lib.rs | 5 +- crates/skippy-engine-mlx/src/stage.rs | 5 +- .../skippy-engine-mlx/src/stage/nemotron_h.rs | 384 +++++++++++++++++- docs/design/MLX_STAGE_ENGINE_PLAN.md | 20 + 7 files changed, 459 insertions(+), 6 deletions(-) diff --git a/crates/skippy-engine-mlx/Cargo.toml b/crates/skippy-engine-mlx/Cargo.toml index 425b1c4ba0..0495bf2f4e 100644 --- a/crates/skippy-engine-mlx/Cargo.toml +++ b/crates/skippy-engine-mlx/Cargo.toml @@ -48,6 +48,9 @@ mlx = [ "dep:serde", "dep:sha2", "dep:tokenizers", + # SafeMLX and static llama.cpp both export GGUF C symbols such as + # `gguf_get_key`; keep the native llama runtime in a separate link unit. + "skippy-server/dynamic-native-runtime", ] [dependencies] diff --git a/crates/skippy-engine-mlx/STAGED_EXECUTION.md b/crates/skippy-engine-mlx/STAGED_EXECUTION.md index a67ff4d800..dfc193a732 100644 --- a/crates/skippy-engine-mlx/STAGED_EXECUTION.md +++ b/crates/skippy-engine-mlx/STAGED_EXECUTION.md @@ -217,8 +217,31 @@ just mlx-stage validate-nemotron-h \ --model /tmp/nemotron-nano-layer1-affine4 --layer 1 just mlx-stage validate-nemotron-h-stage \ --model /tmp/nemotron-nano-layer1-affine4 --layer 1 +just mlx-stage validate-nemotron-h-wire \ + --model /tmp/nemotron-nano-layer1-affine4 --layer 1 --wire-dtype f32 +just mlx-stage validate-nemotron-h-wire \ + --model /tmp/nemotron-nano-layer1-affine4 --layer 1 --wire-dtype f16 ``` +The last two commands deliberately put the real layer-1 engine in an +unnecessary two-stage loopback chain. The downstream stage is a synthetic +capture/final engine, not another Nemotron layer. It asserts the forwarded +`PrefillFinal` kind, session, token, position, and `[1, 1, 2688]` residual; it +returns a sentinel prediction; and it records the session reset before the +upstream Stop/ACK completes. The F32 boundary matched direct block execution +with maximum absolute error `1.1920929e-7` under `atol=1e-4`, `rtol=1e-4`. +The F16 boundary had maximum absolute error `0.00062298775` and maximum +relative error `0.00048053052` under `atol=5e-4`, `rtol=1e-3`. Those thresholds +are empirical evidence for this layer and deterministic one-token input, not a +family certification. The input values are multiples of 1/32, so the F16 result +mostly exercises output-boundary rounding rather than difficult input +rounding. + +This proves the real Skippy TCP framing, activation codec, sideband forwarding, +predicted reply propagation, and chained Stop/ACK around one real MLX frontier +layer. It does not prove a second real model stage, multi-token prefill, decode, +Nemotron recurrent state, host/QUIC orchestration, or end-to-end token logits. + ## Reproduce Build once: @@ -299,6 +322,10 @@ just mlx-stage prove --connect 127.0.0.1:19090 --wire-dtype f16 - Bounded Nemotron-H derivation and execution currently accept exactly one internal `E`/MoE layer. They do not expose a hybrid multi-layer stage or recurrent state on the wire. +- The Nemotron binary-wire validator uses a synthetic adjacent final stage and + a one-token loopback request. Its three-layer synthetic topology exists only + to exercise the transport harness; it is not a deployable 52-layer model + topology. - Greedy sampling only; sampling metadata is preserved in the contract and rejected explicitly when enabled. - No KV page import/export, cache trim/checkpoint, MTP, speculative verify, diff --git a/crates/skippy-engine-mlx/src/bin/mlx-stage.rs b/crates/skippy-engine-mlx/src/bin/mlx-stage.rs index 81d9367339..960f71d6fc 100644 --- a/crates/skippy-engine-mlx/src/bin/mlx-stage.rs +++ b/crates/skippy-engine-mlx/src/bin/mlx-stage.rs @@ -15,7 +15,8 @@ mod real { use skippy_engine_mlx::{ MlxComputeDtype, MlxDerivationControl, MlxDerivedStageCacheConfig, MlxDerivedStageConfig, MlxStageEngine, MlxStageEngineConfig, MlxWeightQuantization, derive_quantized_stage, - derive_quantized_stage_cached, mlx_derived_stage_cache_root, validate_nemotron_h_moe_stage, + derive_quantized_stage_cached, mlx_derived_stage_cache_root, + validate_nemotron_h_binary_wire, validate_nemotron_h_moe_stage, validate_nemotron_h_stage_engine, }; use skippy_protocol::binary::{ @@ -113,6 +114,15 @@ mod real { #[arg(long)] layer: usize, }, + /// Prove a Nemotron-H layer over the real binary stage wire. + ValidateNemotronHWire { + #[arg(long)] + model: PathBuf, + #[arg(long)] + layer: usize, + #[arg(long, value_enum, default_value_t = WireDtype::F16)] + wire_dtype: WireDtype, + }, /// Drive a stage chain and assert its greedy token sequence. Prove { #[arg(long)] @@ -263,6 +273,15 @@ mod real { println!("{}", serde_json::to_string_pretty(&report)?); Ok(()) } + Command::ValidateNemotronHWire { + model, + layer, + wire_dtype, + } => { + let report = validate_nemotron_h_binary_wire(model, layer, wire_dtype.into())?; + println!("{}", serde_json::to_string_pretty(&report)?); + Ok(()) + } Command::Prove { connect, tokens, diff --git a/crates/skippy-engine-mlx/src/lib.rs b/crates/skippy-engine-mlx/src/lib.rs index d8603c4325..8aa8a3d4a1 100644 --- a/crates/skippy-engine-mlx/src/lib.rs +++ b/crates/skippy-engine-mlx/src/lib.rs @@ -29,8 +29,9 @@ pub use derived::{ pub use engine::{ChatTurn, GenerateRequest, MlxEngine, MlxEngineConfig}; #[cfg(all(feature = "mlx", target_os = "macos"))] pub use stage::{ - MlxComputeDtype, MlxNemotronHStageValidationReport, MlxStageEngine, MlxStageEngineConfig, - MlxWeightQuantization, validate_nemotron_h_stage_engine, + MlxComputeDtype, MlxNemotronHStageValidationReport, MlxNemotronHWireValidationReport, + MlxStageEngine, MlxStageEngineConfig, MlxWeightQuantization, validate_nemotron_h_binary_wire, + validate_nemotron_h_stage_engine, }; /// True when this build actually contains the MLX engine. diff --git a/crates/skippy-engine-mlx/src/stage.rs b/crates/skippy-engine-mlx/src/stage.rs index 7de9a14271..52db36b1d6 100644 --- a/crates/skippy-engine-mlx/src/stage.rs +++ b/crates/skippy-engine-mlx/src/stage.rs @@ -26,7 +26,10 @@ use skippy_engine::{ }; use self::nemotron_h::NemotronHMoeStage; -pub use self::nemotron_h::{MlxNemotronHStageValidationReport, validate_nemotron_h_stage_engine}; +pub use self::nemotron_h::{ + MlxNemotronHStageValidationReport, MlxNemotronHWireValidationReport, + validate_nemotron_h_binary_wire, validate_nemotron_h_stage_engine, +}; #[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] pub enum MlxComputeDtype { diff --git a/crates/skippy-engine-mlx/src/stage/nemotron_h.rs b/crates/skippy-engine-mlx/src/stage/nemotron_h.rs index cb2d286aa2..b2bb6f7ec1 100644 --- a/crates/skippy-engine-mlx/src/stage/nemotron_h.rs +++ b/crates/skippy-engine-mlx/src/stage/nemotron_h.rs @@ -1,8 +1,19 @@ //! Stateless Nemotron-H MoE blocks behind the shared stage-engine contract. -use std::path::{Path, PathBuf}; +use std::{ + io::Write, + net::{SocketAddr, TcpListener, TcpStream}, + path::{Path, PathBuf}, + sync::{ + Arc, + atomic::{AtomicBool, Ordering}, + mpsc, + }, + thread, + time::{Duration, Instant}, +}; -use anyhow::{Context, Result, bail, ensure}; +use anyhow::{Context, Result, anyhow, bail, ensure}; use safemlx::{ Array, Device, DeviceType, Dtype, Stream, memory::{active_memory, cache_memory, peak_memory, reset_peak_memory}, @@ -18,6 +29,11 @@ use skippy_engine::{ StageActivation, StageEngine, StageEngineInfo, StageExecutionKind, StageExecutionOutput, StageExecutionRequest, }; +use skippy_protocol::binary::{ + StageStateHeader, StageWireMessage, WireActivationDType, WireMessageKind, WireReplyKind, + encode_f32_activation_payload, recv_ready, recv_reply, write_stage_message, +}; +use skippy_server::engine_transport::{EngineStageServerOptions, serve_stage_engine_until}; use crate::derived::{nemotron_h_validation_values, validate_nemotron_h_moe_stage_output}; @@ -50,6 +66,29 @@ pub struct MlxNemotronHStageValidationReport { pub mlx_peak_memory_bytes: usize, } +/// Evidence from an intentionally unnecessary two-stage binary-wire chain. +#[derive(Clone, Debug, Serialize)] +pub struct MlxNemotronHWireValidationReport { + pub model_dir: PathBuf, + pub layer: usize, + pub wire_dtype: String, + pub input_shape: Vec, + pub captured_shape: Vec, + pub output_is_finite: bool, + pub output_within_tolerance: bool, + pub max_abs_diff: f32, + pub max_relative_diff_for_reference_magnitude_above_atol: f32, + pub comparison_atol: f32, + pub comparison_rtol: f32, + pub predicted_sentinel: i32, + pub forwarded_kind: String, + pub forwarded_session_id: u64, + pub downstream_reset_session_id: u64, + pub mlx_active_memory_bytes: usize, + pub mlx_cache_memory_bytes: usize, + pub mlx_peak_memory_bytes: usize, +} + /// Compares direct safemlx execution with execution through `StageEngine`. pub fn validate_nemotron_h_stage_engine( model_dir: impl AsRef, @@ -149,6 +188,169 @@ pub fn validate_nemotron_h_stage_engine( }) } +/// Sends one deterministic residual through the real binary stage wire and a +/// synthetic final capture stage, then compares the captured residual with a +/// direct safemlx block execution. +pub fn validate_nemotron_h_binary_wire( + model_dir: impl AsRef, + layer: usize, + wire_dtype: WireActivationDType, +) -> Result { + const PREDICTED_SENTINEL: i32 = 424_242; + let model_dir = model_dir.as_ref(); + let (atol, rtol) = wire_tolerances(wire_dtype)?; + let (_direct, direct_values) = validate_nemotron_h_moe_stage_output(model_dir, layer)?; + reset_peak_memory()?; + let engine = Arc::new(MlxStageEngine::spawn(MlxStageEngineConfig { + model_dir: model_dir.to_path_buf(), + model_id: "nemotron-h-wire-validation".to_string(), + stage_index: 1, + layer_start: u32::try_from(layer)?, + layer_end: u32::try_from(layer.checked_add(1).context("layer index overflow")?)?, + compute_dtype: MlxComputeDtype::Bf16, + weight_quantization: None, + ctx_size: Some(1), + })?); + let width = usize::try_from(engine.info().activation_width)?; + let values = nemotron_h_validation_values(i32::try_from(width)?); + let (captured_tx, captured_rx) = mpsc::channel(); + let (reset_tx, reset_rx) = mpsc::channel(); + let capture = Arc::new(CaptureStageEngine::new( + engine.info(), + PREDICTED_SENTINEL, + captured_tx, + reset_tx, + )?); + let (capture_server, capture_addr, capture_ready) = WireServer::spawn_ready( + capture, + EngineStageServerOptions { + bind_addr: "127.0.0.1:0".parse()?, + downstream_addr: None, + wire_dtype, + }, + )?; + drop(capture_ready); + let (stage_server, _stage_addr, mut client) = WireServer::spawn_ready( + engine, + EngineStageServerOptions { + bind_addr: "127.0.0.1:0".parse()?, + downstream_addr: Some(capture_addr), + wire_dtype, + }, + )?; + let input = StageActivation::from_values(1, width, &values)?; + let message = wire_validation_message(&input, wire_dtype)?; + write_stage_message(&mut client, &message, wire_dtype)?; + client.flush().ok(); + let reply = recv_reply(&mut client)?; + ensure!( + reply.kind == WireReplyKind::PredictedToken && reply.predicted == PREDICTED_SENTINEL, + "binary wire chain returned the wrong final reply" + ); + let captured = captured_rx + .recv_timeout(Duration::from_secs(5)) + .context("capture stage did not receive the Nemotron-H residual")?; + let stop = StageWireMessage::stop_with_identity(wire_dtype, 1, 1); + write_stage_message(&mut client, &stop, wire_dtype)?; + client.flush().ok(); + ensure!( + recv_reply(&mut client)?.kind == WireReplyKind::Ack, + "binary wire chain stop did not return ACK" + ); + let downstream_reset_session_id = reset_rx + .recv_timeout(Duration::from_secs(5)) + .context("capture stage did not receive the forwarded session reset")?; + ensure!( + downstream_reset_session_id == 1, + "capture stage reset the wrong session" + ); + let mlx_active_memory_bytes = active_memory()?; + let mlx_cache_memory_bytes = cache_memory()?; + let mlx_peak_memory_bytes = peak_memory()?; + drop(client); + stage_server.stop()?; + capture_server.stop()?; + + let captured_values = captured.values(); + let output_is_finite = captured_values.iter().copied().all(f32::is_finite); + ensure!(output_is_finite, "binary wire output is not finite"); + let comparison = compare_outputs(&direct_values, &captured_values, atol, rtol)?; + ensure!( + comparison.all_close, + "binary wire output differs from direct execution: dtype={} max_abs={} max_relative_for_reference_magnitude_above_atol={} atol={atol} rtol={rtol}", + wire_dtype_label(wire_dtype)?, + comparison.max_abs, + comparison.max_relative_for_reference_magnitude_above_atol, + ); + Ok(MlxNemotronHWireValidationReport { + model_dir: model_dir.to_path_buf(), + layer, + wire_dtype: wire_dtype_label(wire_dtype)?.to_string(), + input_shape: vec![1, 1, width], + captured_shape: vec![1, captured.token_count, captured.width], + output_is_finite, + output_within_tolerance: comparison.all_close, + max_abs_diff: comparison.max_abs, + max_relative_diff_for_reference_magnitude_above_atol: comparison + .max_relative_for_reference_magnitude_above_atol, + comparison_atol: atol, + comparison_rtol: rtol, + predicted_sentinel: PREDICTED_SENTINEL, + forwarded_kind: "prefill_final".to_string(), + forwarded_session_id: 1, + downstream_reset_session_id, + mlx_active_memory_bytes, + mlx_cache_memory_bytes, + mlx_peak_memory_bytes, + }) +} + +fn wire_tolerances(wire_dtype: WireActivationDType) -> Result<(f32, f32)> { + match wire_dtype { + WireActivationDType::F32 => Ok((1.0e-4, 1.0e-4)), + WireActivationDType::F16 => Ok((5.0e-4, 1.0e-3)), + other => bail!("Nemotron-H binary-wire validation does not support {other:?}"), + } +} + +fn wire_dtype_label(wire_dtype: WireActivationDType) -> Result<&'static str> { + match wire_dtype { + WireActivationDType::F32 => Ok("f32"), + WireActivationDType::F16 => Ok("f16"), + other => bail!("Nemotron-H binary-wire validation does not support {other:?}"), + } +} + +fn wire_validation_message( + input: &StageActivation, + wire_dtype: WireActivationDType, +) -> Result { + let kind = WireMessageKind::PrefillFinalEmbd; + let mut state = StageStateHeader::new(kind, wire_dtype); + state.current_token = 0; + state.prompt_token_count = 1; + state.source_stage_index = 0; + Ok(StageWireMessage { + kind, + pos_start: 0, + token_count: 1, + state, + request_id: 1, + session_id: 1, + sampling: None, + chat_sampling_metadata: None, + tokens: vec![0], + positions: vec![0], + activation: encode_f32_activation_payload( + wire_dtype, + 1, + i32::try_from(input.width)?, + &input.f32_le_bytes, + )?, + raw_bytes: Vec::new(), + }) +} + fn execute_validation_input( engine: &MlxStageEngine, session_id: u64, @@ -168,6 +370,184 @@ fn execute_validation_input( .context("Nemotron-H internal stage returned no activation") } +struct CaptureStageEngine { + info: StageEngineInfo, + predicted_sentinel: i32, + captured: mpsc::Sender, + reset: mpsc::Sender, +} + +impl CaptureStageEngine { + fn new( + source: &StageEngineInfo, + predicted_sentinel: i32, + captured: mpsc::Sender, + reset: mpsc::Sender, + ) -> Result { + let stage_index = source + .stage_index + .checked_add(1) + .context("capture stage index overflow")?; + let total_layers = source + .layer_end + .checked_add(1) + .context("capture layer index overflow")?; + let info = StageEngineInfo { + engine: "capture".to_string(), + model_id: source.model_id.clone(), + stage_index, + layer_start: source.layer_end, + layer_end: total_layers, + total_layers, + activation_width: source.activation_width, + }; + info.validate()?; + Ok(Self { + info, + predicted_sentinel, + captured, + reset, + }) + } +} + +impl StageEngine for CaptureStageEngine { + fn info(&self) -> &StageEngineInfo { + &self.info + } + + fn execute(&self, request: StageExecutionRequest) -> Result { + ensure!( + request.kind == StageExecutionKind::PrefillFinal, + "capture stage received the wrong execution kind" + ); + ensure!( + request.session_id == 1, + "capture stage received the wrong session" + ); + ensure!( + request.token_ids == [0], + "capture stage received the wrong token sideband" + ); + ensure!( + request.positions == [0], + "capture stage received the wrong position sideband" + ); + let input = request + .input + .context("capture stage requires a residual activation")?; + self.captured + .send(input) + .map_err(|_| anyhow!("wire validation capture receiver was dropped"))?; + Ok(StageExecutionOutput { + activation: None, + predicted_tokens: vec![self.predicted_sentinel], + }) + } + + fn reset_session(&self, session_id: u64) -> Result<()> { + self.reset + .send(session_id) + .map_err(|_| anyhow!("wire validation reset receiver was dropped")) + } +} + +struct WireServer { + shutdown: Arc, + join: Option>>, +} + +impl WireServer { + fn spawn(engine: Arc, options: EngineStageServerOptions) -> Self { + let shutdown = Arc::new(AtomicBool::new(false)); + let thread_shutdown = Arc::clone(&shutdown); + let join = + thread::spawn(move || serve_stage_engine_until(engine, options, thread_shutdown)); + Self { + shutdown, + join: Some(join), + } + } + + fn spawn_ready( + engine: Arc, + options: EngineStageServerOptions, + ) -> Result<(Self, SocketAddr, TcpStream)> { + const BIND_ATTEMPTS: usize = 3; + let mut last_error = None; + for _ in 0..BIND_ATTEMPTS { + let addr = reserve_loopback_addr()?; + let mut attempt_options = options.clone(); + attempt_options.bind_addr = addr; + let server = Self::spawn(Arc::clone(&engine), attempt_options); + match connect_ready(addr) { + Ok(client) => return Ok((server, addr, client)), + Err(connect_error) => { + let server_error = server.stop().err(); + last_error = Some(match server_error { + Some(server_error) => anyhow!( + "connect wire stage at {addr}: {connect_error:#}; server failed: {server_error:#}" + ), + None => connect_error, + }); + } + } + } + Err(last_error.unwrap_or_else(|| anyhow!("wire stage did not start"))) + .context("start binary wire validation server") + } + + fn stop(mut self) -> Result<()> { + self.finish() + } + + fn finish(&mut self) -> Result<()> { + self.shutdown.store(true, Ordering::SeqCst); + let Some(join) = self.join.take() else { + return Ok(()); + }; + join.join() + .map_err(|_| anyhow!("binary wire validation server panicked"))? + } +} + +impl Drop for WireServer { + fn drop(&mut self) { + let _ = self.finish(); + } +} + +fn reserve_loopback_addr() -> Result { + let listener = TcpListener::bind("127.0.0.1:0")?; + Ok(listener.local_addr()?) +} + +fn connect_ready(addr: SocketAddr) -> Result { + let deadline = Instant::now() + Duration::from_secs(5); + loop { + let error = match TcpStream::connect(addr) { + Ok(mut stream) => { + stream.set_nodelay(true).ok(); + stream + .set_read_timeout(Some(Duration::from_millis(250))) + .ok(); + match recv_ready(&mut stream) { + Ok(()) => { + stream.set_read_timeout(None).ok(); + return Ok(stream); + } + Err(error) => anyhow!(error).context("receive wire ready handshake"), + } + } + Err(error) => anyhow!(error).context("connect TCP socket"), + }; + if Instant::now() >= deadline { + return Err(error).with_context(|| format!("connect wire stage at {addr}")); + } + thread::sleep(Duration::from_millis(25)); + } +} + struct OutputComparison { all_close: bool, max_abs: f32, diff --git a/docs/design/MLX_STAGE_ENGINE_PLAN.md b/docs/design/MLX_STAGE_ENGINE_PLAN.md index fd12489d48..a4bb698d67 100644 --- a/docs/design/MLX_STAGE_ENGINE_PLAN.md +++ b/docs/design/MLX_STAGE_ENGINE_PLAN.md @@ -119,6 +119,26 @@ process RSS high-water is not representative of a single serving stage. General hybrid-stage execution is still gated on recurrent/attention state and boundary work. +**Update — that real Nano layer now crosses the Skippy binary wire.** An +intentionally unnecessary loopback proof sends one deterministic +`PrefillFinalEmbd` residual through the real affine-4 layer-1 `MlxStageEngine` +and then into a fabricated capture/final engine. The capture asserts the +forwarded execution kind, session, token, position, and activation shape; +returns a sentinel prediction; and observes the forwarded session reset before +Stop/ACK completes. F32 matched the direct block with maximum absolute error +`1.1920929e-7` at `atol=1e-4`, `rtol=1e-4`. F16 produced maximum absolute and +relative errors `0.00062298775` and `0.00048053052` at `atol=5e-4`, +`rtol=1e-3`. Runtime MLX active/peak memory remained 730,404,608 / 811,763,256 +bytes. These are empirical thresholds for one layer and one input, not family +certification. The deterministic input is exactly representable in F16, so the +F16 delta primarily measures the output boundary. + +This is concrete codec, forwarding, reply, and control-chain evidence around +one real frontier layer. The adjacent final stage and its three-layer topology +are synthetic harness devices. It is not evidence for two real Nemotron +stages, multi-token prefill, decode, recurrent state, host/QUIC placement, or +full-model logits. + The derivation memory bound is the final packed routed bank, not one expert: six preallocated payload buffers total 718,405,632 bytes. Moving those buffers to a disk-backed random-write spool is the next step if preparation RSS must From 4a261f6d5b52445dc9c8456428a47fec2dd193ae Mon Sep 17 00:00:00 2001 From: Michael Neale <14976+michaelneale@users.noreply.github.com> Date: Fri, 17 Jul 2026 22:33:22 +1000 Subject: [PATCH 23/37] test(mlx): prove multi-token Nemotron wire --- crates/skippy-engine-mlx/STAGED_EXECUTION.md | 37 +++++++------ crates/skippy-engine-mlx/src/bin/mlx-stage.rs | 12 ++++- crates/skippy-engine-mlx/src/derived.rs | 5 +- .../src/derived/nemotron_h.rs | 34 ++++++++++-- crates/skippy-engine-mlx/src/lib.rs | 2 +- crates/skippy-engine-mlx/src/stage.rs | 3 +- .../skippy-engine-mlx/src/stage/nemotron_h.rs | 54 ++++++++++++++----- docs/design/MLX_STAGE_ENGINE_PLAN.md | 23 ++++---- 8 files changed, 120 insertions(+), 50 deletions(-) diff --git a/crates/skippy-engine-mlx/STAGED_EXECUTION.md b/crates/skippy-engine-mlx/STAGED_EXECUTION.md index dfc193a732..17c9c97355 100644 --- a/crates/skippy-engine-mlx/STAGED_EXECUTION.md +++ b/crates/skippy-engine-mlx/STAGED_EXECUTION.md @@ -218,29 +218,36 @@ just mlx-stage validate-nemotron-h \ just mlx-stage validate-nemotron-h-stage \ --model /tmp/nemotron-nano-layer1-affine4 --layer 1 just mlx-stage validate-nemotron-h-wire \ - --model /tmp/nemotron-nano-layer1-affine4 --layer 1 --wire-dtype f32 + --model /tmp/nemotron-nano-layer1-affine4 --layer 1 --tokens 32 \ + --wire-dtype f32 just mlx-stage validate-nemotron-h-wire \ - --model /tmp/nemotron-nano-layer1-affine4 --layer 1 --wire-dtype f16 + --model /tmp/nemotron-nano-layer1-affine4 --layer 1 --tokens 32 \ + --wire-dtype f16 ``` The last two commands deliberately put the real layer-1 engine in an unnecessary two-stage loopback chain. The downstream stage is a synthetic capture/final engine, not another Nemotron layer. It asserts the forwarded -`PrefillFinal` kind, session, token, position, and `[1, 1, 2688]` residual; it +`PrefillFinal` kind, session, all token/position sidebands, and +`[1, 32, 2688]` residual; it returns a sentinel prediction; and it records the session reset before the upstream Stop/ACK completes. The F32 boundary matched direct block execution -with maximum absolute error `1.1920929e-7` under `atol=1e-4`, `rtol=1e-4`. -The F16 boundary had maximum absolute error `0.00062298775` and maximum -relative error `0.00048053052` under `atol=5e-4`, `rtol=1e-3`. Those thresholds -are empirical evidence for this layer and deterministic one-token input, not a -family certification. The input values are multiples of 1/32, so the F16 result -mostly exercises output-boundary rounding rather than difficult input -rounding. +with maximum absolute error `2.3841858e-7` under `atol=1e-4`, `rtol=1e-4`. +The F16 boundary had maximum absolute error `0.000923872` and maximum relative +error `0.00048756658` under `atol=5e-4`, `rtol=1e-3`. The corresponding +activation payloads were 344,064 F32 bytes and 172,032 F16 bytes per boundary. +Runtime active memory stayed at 730,404,608 bytes and peak MLX memory was +820,697,688 bytes. Those thresholds are empirical evidence for this layer and +deterministic 32-token input, not a family certification. The input values are +multiples of 1/32, so the F16 result mostly exercises output-boundary rounding +rather than difficult input rounding. The validator defaults to one token and +accepts `--tokens` for larger prefill checks. This proves the real Skippy TCP framing, activation codec, sideband forwarding, predicted reply propagation, and chained Stop/ACK around one real MLX frontier -layer. It does not prove a second real model stage, multi-token prefill, decode, -Nemotron recurrent state, host/QUIC orchestration, or end-to-end token logits. +layer, including a 32-token prefill. It does not prove a second real model +stage, decode, Nemotron recurrent state, host/QUIC orchestration, or end-to-end +token logits. ## Reproduce @@ -323,9 +330,9 @@ just mlx-stage prove --connect 127.0.0.1:19090 --wire-dtype f16 internal `E`/MoE layer. They do not expose a hybrid multi-layer stage or recurrent state on the wire. - The Nemotron binary-wire validator uses a synthetic adjacent final stage and - a one-token loopback request. Its three-layer synthetic topology exists only - to exercise the transport harness; it is not a deployable 52-layer model - topology. + configurable loopback prefill (one and 32 tokens have been exercised). Its + three-layer synthetic topology exists only to exercise the transport harness; + it is not a deployable 52-layer model topology. - Greedy sampling only; sampling metadata is preserved in the contract and rejected explicitly when enabled. - No KV page import/export, cache trim/checkpoint, MTP, speculative verify, diff --git a/crates/skippy-engine-mlx/src/bin/mlx-stage.rs b/crates/skippy-engine-mlx/src/bin/mlx-stage.rs index 960f71d6fc..259e9ec466 100644 --- a/crates/skippy-engine-mlx/src/bin/mlx-stage.rs +++ b/crates/skippy-engine-mlx/src/bin/mlx-stage.rs @@ -16,7 +16,7 @@ mod real { MlxComputeDtype, MlxDerivationControl, MlxDerivedStageCacheConfig, MlxDerivedStageConfig, MlxStageEngine, MlxStageEngineConfig, MlxWeightQuantization, derive_quantized_stage, derive_quantized_stage_cached, mlx_derived_stage_cache_root, - validate_nemotron_h_binary_wire, validate_nemotron_h_moe_stage, + validate_nemotron_h_binary_wire_tokens, validate_nemotron_h_moe_stage, validate_nemotron_h_stage_engine, }; use skippy_protocol::binary::{ @@ -120,6 +120,8 @@ mod real { model: PathBuf, #[arg(long)] layer: usize, + #[arg(long, default_value_t = 1)] + tokens: usize, #[arg(long, value_enum, default_value_t = WireDtype::F16)] wire_dtype: WireDtype, }, @@ -276,9 +278,15 @@ mod real { Command::ValidateNemotronHWire { model, layer, + tokens, wire_dtype, } => { - let report = validate_nemotron_h_binary_wire(model, layer, wire_dtype.into())?; + let report = validate_nemotron_h_binary_wire_tokens( + model, + layer, + wire_dtype.into(), + tokens, + )?; println!("{}", serde_json::to_string_pretty(&report)?); Ok(()) } diff --git a/crates/skippy-engine-mlx/src/derived.rs b/crates/skippy-engine-mlx/src/derived.rs index 0a87fcb79e..3754d22584 100644 --- a/crates/skippy-engine-mlx/src/derived.rs +++ b/crates/skippy-engine-mlx/src/derived.rs @@ -29,7 +29,10 @@ use sha2::{Digest, Sha256}; use crate::stage::MlxWeightQuantization; use nemotron_h::NemotronHDerivation; -pub(crate) use nemotron_h::{nemotron_h_validation_values, validate_nemotron_h_moe_stage_output}; +pub(crate) use nemotron_h::{ + nemotron_h_validation_values, validate_nemotron_h_moe_stage_output, + validate_nemotron_h_moe_stage_output_for_tokens, +}; mod cache; mod expert_bank; diff --git a/crates/skippy-engine-mlx/src/derived/nemotron_h.rs b/crates/skippy-engine-mlx/src/derived/nemotron_h.rs index a30c402f9d..a069c68d62 100644 --- a/crates/skippy-engine-mlx/src/derived/nemotron_h.rs +++ b/crates/skippy-engine-mlx/src/derived/nemotron_h.rs @@ -101,8 +101,18 @@ pub fn validate_nemotron_h_moe_stage( pub(crate) fn validate_nemotron_h_moe_stage_output( model_dir: impl AsRef, layer: usize, +) -> Result<(MlxNemotronHValidationReport, Vec)> { + validate_nemotron_h_moe_stage_output_for_tokens(model_dir, layer, 1) +} + +pub(crate) fn validate_nemotron_h_moe_stage_output_for_tokens( + model_dir: impl AsRef, + layer: usize, + token_count: usize, ) -> Result<(MlxNemotronHValidationReport, Vec)> { let model_dir = model_dir.as_ref(); + ensure!(token_count > 0, "Nemotron-H validation needs tokens"); + let token_count_i32 = i32::try_from(token_count).context("token count exceeds i32")?; let args = get_nemotron_h_model_args(model_dir)?; ensure!( args.hybrid_override_pattern @@ -125,8 +135,8 @@ pub(crate) fn validate_nemotron_h_moe_stage_output( )?; load_report.finish(&block, &load_config)?; - let values = nemotron_h_validation_values(args.hidden_size); - let input = Array::from_slice(&values, &[1, 1, args.hidden_size]) + let values = nemotron_h_validation_values(args.hidden_size, token_count)?; + let input = Array::from_slice(&values, &[1, token_count_i32, args.hidden_size]) .as_dtype(safemlx::Dtype::Bfloat16, &stream)?; let output = block.forward( BlockInput { @@ -165,10 +175,16 @@ pub(crate) fn validate_nemotron_h_moe_stage_output( )) } -pub(crate) fn nemotron_h_validation_values(hidden_size: i32) -> Vec { - (0..hidden_size) +pub(crate) fn nemotron_h_validation_values( + hidden_size: i32, + token_count: usize, +) -> Result> { + let element_count = usize::try_from(hidden_size)? + .checked_mul(token_count) + .context("Nemotron-H validation input size overflow")?; + Ok((0..element_count) .map(|index| bf16::from_f32(((index % 31) as f32 - 15.0) / 32.0).to_f32()) - .collect() + .collect()) } fn f32_sha256(values: &[f32]) -> String { @@ -390,6 +406,14 @@ mod tests { }) } + #[test] + fn validation_input_covers_every_token() { + let values = nemotron_h_validation_values(3, 2).unwrap(); + assert_eq!(values.len(), 6); + assert_eq!(values[0], -15.0 / 32.0); + assert_eq!(values[3], -12.0 / 32.0); + } + #[test] fn rewrites_public_layer_keys_to_runtime_fields() { let derivation = NemotronHDerivation::new(&config(), 1, 2).unwrap(); diff --git a/crates/skippy-engine-mlx/src/lib.rs b/crates/skippy-engine-mlx/src/lib.rs index 8aa8a3d4a1..dc5f65dbbf 100644 --- a/crates/skippy-engine-mlx/src/lib.rs +++ b/crates/skippy-engine-mlx/src/lib.rs @@ -31,7 +31,7 @@ pub use engine::{ChatTurn, GenerateRequest, MlxEngine, MlxEngineConfig}; pub use stage::{ MlxComputeDtype, MlxNemotronHStageValidationReport, MlxNemotronHWireValidationReport, MlxStageEngine, MlxStageEngineConfig, MlxWeightQuantization, validate_nemotron_h_binary_wire, - validate_nemotron_h_stage_engine, + validate_nemotron_h_binary_wire_tokens, validate_nemotron_h_stage_engine, }; /// True when this build actually contains the MLX engine. diff --git a/crates/skippy-engine-mlx/src/stage.rs b/crates/skippy-engine-mlx/src/stage.rs index 52db36b1d6..218b72002b 100644 --- a/crates/skippy-engine-mlx/src/stage.rs +++ b/crates/skippy-engine-mlx/src/stage.rs @@ -28,7 +28,8 @@ use skippy_engine::{ use self::nemotron_h::NemotronHMoeStage; pub use self::nemotron_h::{ MlxNemotronHStageValidationReport, MlxNemotronHWireValidationReport, - validate_nemotron_h_binary_wire, validate_nemotron_h_stage_engine, + validate_nemotron_h_binary_wire, validate_nemotron_h_binary_wire_tokens, + validate_nemotron_h_stage_engine, }; #[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] diff --git a/crates/skippy-engine-mlx/src/stage/nemotron_h.rs b/crates/skippy-engine-mlx/src/stage/nemotron_h.rs index b2bb6f7ec1..c7634fca62 100644 --- a/crates/skippy-engine-mlx/src/stage/nemotron_h.rs +++ b/crates/skippy-engine-mlx/src/stage/nemotron_h.rs @@ -35,7 +35,10 @@ use skippy_protocol::binary::{ }; use skippy_server::engine_transport::{EngineStageServerOptions, serve_stage_engine_until}; -use crate::derived::{nemotron_h_validation_values, validate_nemotron_h_moe_stage_output}; +use crate::derived::{ + nemotron_h_validation_values, validate_nemotron_h_moe_stage_output, + validate_nemotron_h_moe_stage_output_for_tokens, +}; use super::{MlxComputeDtype, MlxStageEngine, MlxStageEngineConfig, array_activation}; @@ -71,6 +74,7 @@ pub struct MlxNemotronHStageValidationReport { pub struct MlxNemotronHWireValidationReport { pub model_dir: PathBuf, pub layer: usize, + pub token_count: usize, pub wire_dtype: String, pub input_shape: Vec, pub captured_shape: Vec, @@ -110,7 +114,7 @@ pub fn validate_nemotron_h_stage_engine( ctx_size: Some(1), })?; let width = usize::try_from(engine.info().activation_width)?; - let values = nemotron_h_validation_values(i32::try_from(width)?); + let values = nemotron_h_validation_values(i32::try_from(width)?, 1)?; let first = execute_validation_input(&engine, 1, width, &values)?; let second_session = execute_validation_input(&engine, 2, width, &values)?; engine.reset_session(2)?; @@ -195,11 +199,24 @@ pub fn validate_nemotron_h_binary_wire( model_dir: impl AsRef, layer: usize, wire_dtype: WireActivationDType, +) -> Result { + validate_nemotron_h_binary_wire_tokens(model_dir, layer, wire_dtype, 1) +} + +/// Runs [`validate_nemotron_h_binary_wire`] with a configurable prefill size. +pub fn validate_nemotron_h_binary_wire_tokens( + model_dir: impl AsRef, + layer: usize, + wire_dtype: WireActivationDType, + token_count: usize, ) -> Result { const PREDICTED_SENTINEL: i32 = 424_242; let model_dir = model_dir.as_ref(); + ensure!(token_count > 0, "binary wire validation needs tokens"); + let token_count_u32 = u32::try_from(token_count).context("token count exceeds u32")?; let (atol, rtol) = wire_tolerances(wire_dtype)?; - let (_direct, direct_values) = validate_nemotron_h_moe_stage_output(model_dir, layer)?; + let (_direct, direct_values) = + validate_nemotron_h_moe_stage_output_for_tokens(model_dir, layer, token_count)?; reset_peak_memory()?; let engine = Arc::new(MlxStageEngine::spawn(MlxStageEngineConfig { model_dir: model_dir.to_path_buf(), @@ -209,15 +226,16 @@ pub fn validate_nemotron_h_binary_wire( layer_end: u32::try_from(layer.checked_add(1).context("layer index overflow")?)?, compute_dtype: MlxComputeDtype::Bf16, weight_quantization: None, - ctx_size: Some(1), + ctx_size: Some(token_count_u32), })?); let width = usize::try_from(engine.info().activation_width)?; - let values = nemotron_h_validation_values(i32::try_from(width)?); + let values = nemotron_h_validation_values(i32::try_from(width)?, token_count)?; let (captured_tx, captured_rx) = mpsc::channel(); let (reset_tx, reset_rx) = mpsc::channel(); let capture = Arc::new(CaptureStageEngine::new( engine.info(), PREDICTED_SENTINEL, + token_count, captured_tx, reset_tx, )?); @@ -238,7 +256,7 @@ pub fn validate_nemotron_h_binary_wire( wire_dtype, }, )?; - let input = StageActivation::from_values(1, width, &values)?; + let input = StageActivation::from_values(token_count, width, &values)?; let message = wire_validation_message(&input, wire_dtype)?; write_stage_message(&mut client, &message, wire_dtype)?; client.flush().ok(); @@ -285,8 +303,9 @@ pub fn validate_nemotron_h_binary_wire( Ok(MlxNemotronHWireValidationReport { model_dir: model_dir.to_path_buf(), layer, + token_count, wire_dtype: wire_dtype_label(wire_dtype)?.to_string(), - input_shape: vec![1, 1, width], + input_shape: vec![1, token_count, width], captured_shape: vec![1, captured.token_count, captured.width], output_is_finite, output_within_tolerance: comparison.all_close, @@ -326,24 +345,25 @@ fn wire_validation_message( wire_dtype: WireActivationDType, ) -> Result { let kind = WireMessageKind::PrefillFinalEmbd; + let token_count = i32::try_from(input.token_count).context("token count exceeds i32")?; let mut state = StageStateHeader::new(kind, wire_dtype); state.current_token = 0; - state.prompt_token_count = 1; + state.prompt_token_count = token_count; state.source_stage_index = 0; Ok(StageWireMessage { kind, pos_start: 0, - token_count: 1, + token_count, state, request_id: 1, session_id: 1, sampling: None, chat_sampling_metadata: None, - tokens: vec![0], - positions: vec![0], + tokens: vec![0; input.token_count], + positions: (0..token_count).collect(), activation: encode_f32_activation_payload( wire_dtype, - 1, + token_count, i32::try_from(input.width)?, &input.f32_le_bytes, )?, @@ -373,6 +393,8 @@ fn execute_validation_input( struct CaptureStageEngine { info: StageEngineInfo, predicted_sentinel: i32, + expected_token_ids: Vec, + expected_positions: Vec, captured: mpsc::Sender, reset: mpsc::Sender, } @@ -381,6 +403,7 @@ impl CaptureStageEngine { fn new( source: &StageEngineInfo, predicted_sentinel: i32, + token_count: usize, captured: mpsc::Sender, reset: mpsc::Sender, ) -> Result { @@ -402,9 +425,12 @@ impl CaptureStageEngine { activation_width: source.activation_width, }; info.validate()?; + let token_count = i32::try_from(token_count).context("token count exceeds i32")?; Ok(Self { info, predicted_sentinel, + expected_token_ids: vec![0; usize::try_from(token_count)?], + expected_positions: (0..token_count).collect(), captured, reset, }) @@ -426,11 +452,11 @@ impl StageEngine for CaptureStageEngine { "capture stage received the wrong session" ); ensure!( - request.token_ids == [0], + request.token_ids == self.expected_token_ids, "capture stage received the wrong token sideband" ); ensure!( - request.positions == [0], + request.positions == self.expected_positions, "capture stage received the wrong position sideband" ); let input = request diff --git a/docs/design/MLX_STAGE_ENGINE_PLAN.md b/docs/design/MLX_STAGE_ENGINE_PLAN.md index a4bb698d67..4cb9e839b4 100644 --- a/docs/design/MLX_STAGE_ENGINE_PLAN.md +++ b/docs/design/MLX_STAGE_ENGINE_PLAN.md @@ -120,24 +120,25 @@ General hybrid-stage execution is still gated on recurrent/attention state and boundary work. **Update — that real Nano layer now crosses the Skippy binary wire.** An -intentionally unnecessary loopback proof sends one deterministic +intentionally unnecessary loopback proof sends a deterministic `PrefillFinalEmbd` residual through the real affine-4 layer-1 `MlxStageEngine` and then into a fabricated capture/final engine. The capture asserts the -forwarded execution kind, session, token, position, and activation shape; +forwarded execution kind, session, complete token/position sidebands, and +activation shape; returns a sentinel prediction; and observes the forwarded session reset before -Stop/ACK completes. F32 matched the direct block with maximum absolute error -`1.1920929e-7` at `atol=1e-4`, `rtol=1e-4`. F16 produced maximum absolute and -relative errors `0.00062298775` and `0.00048053052` at `atol=5e-4`, -`rtol=1e-3`. Runtime MLX active/peak memory remained 730,404,608 / 811,763,256 -bytes. These are empirical thresholds for one layer and one input, not family -certification. The deterministic input is exactly representable in F16, so the -F16 delta primarily measures the output boundary. +Stop/ACK completes. At 32 tokens, F32 matched the direct block with maximum +absolute error `2.3841858e-7` at `atol=1e-4`, `rtol=1e-4`. F16 produced maximum +absolute and relative errors `0.000923872` and `0.00048756658` at +`atol=5e-4`, `rtol=1e-3`. The boundary payloads were 344,064 and 172,032 bytes. +Runtime MLX active/peak memory was 730,404,608 / 820,697,688 bytes. These are +empirical thresholds for one layer and one deterministic prefill, not family +certification. The input is exactly representable in F16, so the F16 delta +primarily measures the output boundary. This is concrete codec, forwarding, reply, and control-chain evidence around one real frontier layer. The adjacent final stage and its three-layer topology are synthetic harness devices. It is not evidence for two real Nemotron -stages, multi-token prefill, decode, recurrent state, host/QUIC placement, or -full-model logits. +stages, decode, recurrent state, host/QUIC placement, or full-model logits. The derivation memory bound is the final packed routed bank, not one expert: six preallocated payload buffers total 718,405,632 bytes. Moving those buffers From d381bbd3500361a0d75d2f31251093620db8684c Mon Sep 17 00:00:00 2001 From: Michael Neale <14976+michaelneale@users.noreply.github.com> Date: Fri, 17 Jul 2026 22:54:28 +1000 Subject: [PATCH 24/37] bench(mlx): instrument activation boundary fence --- Cargo.lock | 2 + crates/metrics-server/README.md | 14 + crates/skippy-engine-mlx/Cargo.toml | 4 + crates/skippy-engine-mlx/STAGED_EXECUTION.md | 54 ++ crates/skippy-engine-mlx/src/bin/mlx-stage.rs | 81 +- .../skippy-engine-mlx/src/boundary_bench.rs | 737 ++++++++++++++++++ crates/skippy-engine-mlx/src/lib.rs | 4 + crates/skippy-metrics/src/lib.rs | 7 + docs/design/MLX_STAGE_ENGINE_PLAN.md | 19 +- docs/plugins/telemetry.md | 13 + 10 files changed, 928 insertions(+), 7 deletions(-) create mode 100644 crates/skippy-engine-mlx/src/boundary_bench.rs diff --git a/Cargo.lock b/Cargo.lock index 378920d59d..0742d5970e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -7781,6 +7781,7 @@ dependencies = [ "memmap2", "model-hf", "openai-frontend", + "reqwest 0.12.28", "safemlx", "safemlx-lm", "safetensors", @@ -7788,6 +7789,7 @@ dependencies = [ "serde_json", "sha2 0.10.9", "skippy-engine", + "skippy-metrics", "skippy-protocol", "skippy-server", "tempfile", diff --git a/crates/metrics-server/README.md b/crates/metrics-server/README.md index d7da19fe1b..71fda7174e 100644 --- a/crates/metrics-server/README.md +++ b/crates/metrics-server/README.md @@ -26,6 +26,20 @@ activation byte counts, credit counters, and lifecycle/queue counters. Experimental feature runs should emit enough run, request, session, topology, stage, and model identifiers for report export and post-run debugging. +The opt-in MLX boundary benchmark emits four shared span names: + +- `stage.mlx_boundary_eval_fence` +- `stage.mlx_boundary_host_copy` +- `stage.mlx_boundary_encode` +- `stage.mlx_boundary_decode` + +Its run config and spans contain bounded benchmark shape, dtype, byte-count, +schema/revision, and duration data only. The activation-byte metric is attached +only to the encode span. The benchmark requires explicit HTTP and OTLP +collector arguments; those transport targets are not copied into the run +config or span attributes. It does not export activation values, prompts, +paths, endpoint URLs, hardware identifiers, or model contents. + ## Commands ```bash diff --git a/crates/skippy-engine-mlx/Cargo.toml b/crates/skippy-engine-mlx/Cargo.toml index 0495bf2f4e..bcc25dbff9 100644 --- a/crates/skippy-engine-mlx/Cargo.toml +++ b/crates/skippy-engine-mlx/Cargo.toml @@ -42,11 +42,13 @@ mlx = [ "dep:libc", "dep:memmap2", "dep:model-hf", + "dep:reqwest", "dep:safetensors", "dep:safemlx", "dep:safemlx-lm", "dep:serde", "dep:sha2", + "dep:skippy-metrics", "dep:tokenizers", # SafeMLX and static llama.cpp both export GGUF C symbols such as # `gguf_get_key`; keep the native llama runtime in a separate link unit. @@ -59,6 +61,7 @@ mlx = [ openai-frontend = { path = "../openai-frontend" } model-hf = { path = "../model-hf", optional = true } skippy-engine = { path = "../skippy-engine" } +skippy-metrics = { path = "../skippy-metrics", optional = true } skippy-protocol = { path = "../skippy-protocol" } skippy-server = { path = "../skippy-server" } async-trait = "0.1" @@ -71,6 +74,7 @@ futures-core = "0.3" half = { version = "2", features = ["bytemuck"], optional = true } libc = { version = "0.2", optional = true } memmap2 = { version = "0.9", optional = true } +reqwest = { version = "0.12", default-features = false, features = ["blocking", "json", "rustls-tls"], optional = true } serde_json = "1" serde = { version = "1", features = ["derive"], optional = true } sha2 = { version = "0.10", optional = true } diff --git a/crates/skippy-engine-mlx/STAGED_EXECUTION.md b/crates/skippy-engine-mlx/STAGED_EXECUTION.md index 17c9c97355..2695b87ef1 100644 --- a/crates/skippy-engine-mlx/STAGED_EXECUTION.md +++ b/crates/skippy-engine-mlx/STAGED_EXECUTION.md @@ -249,6 +249,60 @@ layer, including a 32-token prefill. It does not prove a second real model stage, decode, Nemotron recurrent state, host/QUIC orchestration, or end-to-end token logits. +## Boundary-fence benchmark + +`mlx-stage bench-boundary` is the first instrumented pass over the independent +boundary cost. It creates one evaluated F32 MLX array, applies a synthetic lazy +F32 add, and times four release-mode phases separately: + +1. completion of the synthetic add through MLX eval/synchronize, with graph + construction outside the timer; +2. the evaluated host view plus allocation/copy into the F32 byte buffer; +3. the production Skippy F32 or F16 activation-payload encoder; and +4. post-receive activation-payload reconstruction into F32. + +The host-copy phase includes the MLX evaluated-view call. The decode phase does +not include message framing, socket reads, TCP, QUIC, or receive-buffer +allocation. This is not a model-layer benchmark. + +All samples are collected before the telemetry exporter starts, then their +original timestamps are emitted as bounded OTLP spans to an explicitly +configured metrics-server. The benchmark fails on codec drift, non-finite +values, telemetry loss, or a canonical span-count mismatch; it finalizes the +run and saves metrics-server's canonical `report.json`. The report's +`eval_and_host_copy_total` and `codec_total` percentiles are calculated from +paired per-iteration sums, not by adding independent phase percentiles. + +F32 encoding and decoding are straight byte copies. F16 includes numeric +conversion in both directions, so paired F32/F16 results come from separate +per-dtype runs and are not a within-run equivalence comparison. MLX memory +counters do not include the Rust host buffers. + +No prompt, activation values, local paths, collector endpoints, hardware IDs, +or model contents enter telemetry. The required HTTP and OTLP endpoints are +transport targets and are not copied into spans or run config. The operator run +label is validated to a bounded URL-safe character set; a local output report +may contain its explicitly requested report path. + +Reproduce one matrix cell with metrics-server running in another terminal: + +```bash +just metrics-server \ + db=/tmp/mlx-boundary.sqlite \ + http_addr=127.0.0.1:18081 \ + otlp_addr=127.0.0.1:14317 + +just mlx-stage-build +just mlx-stage bench-boundary \ + --width 16384 --tokens 512 --wire-dtype f16 \ + --warmup-iterations 3 --measured-iterations 20 \ + --metrics-http http://127.0.0.1:18081 \ + --metrics-otlp-grpc http://127.0.0.1:14317 \ + --metrics-run-id mlx-boundary-w16384-t512-f16-v2 \ + --metrics-report /tmp/mlx-boundary-metrics.json \ + --output /tmp/mlx-boundary-local.json +``` + ## Reproduce Build once: diff --git a/crates/skippy-engine-mlx/src/bin/mlx-stage.rs b/crates/skippy-engine-mlx/src/bin/mlx-stage.rs index 259e9ec466..10540d8a19 100644 --- a/crates/skippy-engine-mlx/src/bin/mlx-stage.rs +++ b/crates/skippy-engine-mlx/src/bin/mlx-stage.rs @@ -13,11 +13,11 @@ mod real { use clap::{Parser, Subcommand, ValueEnum}; use model_hf::safetensors_stage::{SafetensorsStageMaterializer, SafetensorsStageRequest}; use skippy_engine_mlx::{ - MlxComputeDtype, MlxDerivationControl, MlxDerivedStageCacheConfig, MlxDerivedStageConfig, - MlxStageEngine, MlxStageEngineConfig, MlxWeightQuantization, derive_quantized_stage, - derive_quantized_stage_cached, mlx_derived_stage_cache_root, - validate_nemotron_h_binary_wire_tokens, validate_nemotron_h_moe_stage, - validate_nemotron_h_stage_engine, + MlxBoundaryBenchConfig, MlxComputeDtype, MlxDerivationControl, MlxDerivedStageCacheConfig, + MlxDerivedStageConfig, MlxStageEngine, MlxStageEngineConfig, MlxWeightQuantization, + benchmark_mlx_boundary, derive_quantized_stage, derive_quantized_stage_cached, + mlx_derived_stage_cache_root, validate_nemotron_h_binary_wire_tokens, + validate_nemotron_h_moe_stage, validate_nemotron_h_stage_engine, }; use skippy_protocol::binary::{ StageStateHeader, StageWireMessage, WireActivationDType, WireMessageKind, WireReplyKind, @@ -125,6 +125,29 @@ mod real { #[arg(long, value_enum, default_value_t = WireDtype::F16)] wire_dtype: WireDtype, }, + /// Measure the MLX completion, host-copy, and activation-codec boundary. + BenchBoundary { + #[arg(long)] + width: usize, + #[arg(long)] + tokens: usize, + #[arg(long, value_enum)] + wire_dtype: WireDtype, + #[arg(long, default_value_t = 3)] + warmup_iterations: usize, + #[arg(long, default_value_t = 20)] + measured_iterations: usize, + #[arg(long)] + metrics_http: String, + #[arg(long)] + metrics_otlp_grpc: String, + #[arg(long)] + metrics_run_id: Option, + #[arg(long)] + metrics_report: PathBuf, + #[arg(long)] + output: Option, + }, /// Drive a stage chain and assert its greedy token sequence. Prove { #[arg(long)] @@ -290,6 +313,46 @@ mod real { println!("{}", serde_json::to_string_pretty(&report)?); Ok(()) } + Command::BenchBoundary { + width, + tokens, + wire_dtype, + warmup_iterations, + measured_iterations, + metrics_http, + metrics_otlp_grpc, + metrics_run_id, + metrics_report, + output, + } => { + let metrics_run_id = match metrics_run_id { + Some(metrics_run_id) => metrics_run_id, + None => default_boundary_run_id()?, + }; + let report = benchmark_mlx_boundary(&MlxBoundaryBenchConfig { + width, + token_count: tokens, + wire_dtype: wire_dtype.into(), + warmup_iterations, + measured_iterations, + metrics_http, + metrics_otlp_grpc, + metrics_run_id, + metrics_report_path: metrics_report, + })?; + let json = serde_json::to_vec_pretty(&report)?; + if let Some(output) = output { + if let Some(parent) = output + .parent() + .filter(|parent| !parent.as_os_str().is_empty()) + { + std::fs::create_dir_all(parent)?; + } + std::fs::write(output, &json)?; + } + println!("{}", String::from_utf8(json)?); + Ok(()) + } Command::Prove { connect, tokens, @@ -470,6 +533,14 @@ mod real { .map(|token| token.trim().parse().context("parse token ID")) .collect() } + + fn default_boundary_run_id() -> Result { + let nanos = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .context("system clock is before Unix epoch")? + .as_nanos(); + Ok(format!("mlx-boundary-{nanos}")) + } } #[cfg(all(feature = "mlx", target_os = "macos"))] diff --git a/crates/skippy-engine-mlx/src/boundary_bench.rs b/crates/skippy-engine-mlx/src/boundary_bench.rs new file mode 100644 index 0000000000..1a6064da20 --- /dev/null +++ b/crates/skippy-engine-mlx/src/boundary_bench.rs @@ -0,0 +1,737 @@ +//! Metrics-server-backed measurements for the MLX activation boundary fence. + +use std::{ + collections::BTreeMap, + fs, + mem::size_of, + path::PathBuf, + thread, + time::{Duration, Instant}, +}; + +use anyhow::{Context, Result, anyhow, bail, ensure}; +use safemlx::{ + Array, Device, DeviceType, Dtype, Stream, + memory::{active_memory, cache_memory, peak_memory, reset_peak_memory}, + transforms::eval, +}; +use serde::Serialize; +use serde_json::{Value, json}; +use skippy_metrics::{attr, metric, span}; +use skippy_protocol::{ + StageConfig, + binary::{ + StageStateHeader, StageWireMessage, WireActivationDType, WireMessageKind, + encode_f32_activation_payload, + }, +}; +use skippy_server::telemetry::{ + Telemetry, TelemetryLevel, TelemetryStats, lifecycle_attrs, now_unix_nanos, +}; + +const BENCHMARK_SCHEMA: &str = "mlx-boundary-fence-v2"; + +/// One explicit, reproducible boundary-fence benchmark run. +#[derive(Clone, Debug)] +pub struct MlxBoundaryBenchConfig { + pub width: usize, + pub token_count: usize, + pub wire_dtype: WireActivationDType, + pub warmup_iterations: usize, + pub measured_iterations: usize, + pub metrics_http: String, + pub metrics_otlp_grpc: String, + pub metrics_run_id: String, + pub metrics_report_path: PathBuf, +} + +/// Duration distribution in microseconds. +#[derive(Clone, Debug, Serialize)] +pub struct MlxBoundaryDurationSummary { + pub samples: usize, + pub min_us: f64, + pub mean_us: f64, + pub p50_us: f64, + pub p95_us: f64, + pub max_us: f64, +} + +/// Local measurements paired with a canonical metrics-server report. +#[derive(Clone, Debug, Serialize)] +pub struct MlxBoundaryBenchReport { + pub benchmark: &'static str, + pub code_revision: String, + pub metrics_run_id: String, + pub metrics_report_path: PathBuf, + pub width: usize, + pub token_count: usize, + pub wire_dtype: String, + pub warmup_iterations: usize, + pub measured_iterations: usize, + pub f32_boundary_bytes: usize, + pub wire_activation_payload_bytes: usize, + pub eval_fence: MlxBoundaryDurationSummary, + pub host_copy: MlxBoundaryDurationSummary, + pub eval_and_host_copy_total: MlxBoundaryDurationSummary, + pub encode: MlxBoundaryDurationSummary, + pub decode: MlxBoundaryDurationSummary, + pub codec_total: MlxBoundaryDurationSummary, + pub max_roundtrip_abs_diff: f32, + pub mlx_active_memory_bytes: usize, + pub mlx_cache_memory_bytes: usize, + pub mlx_peak_memory_bytes: usize, + pub telemetry: TelemetryStats, + pub canonical_span_count: u64, +} + +struct BoundarySample { + eval_fence: PhaseTiming, + host_copy: PhaseTiming, + encode: PhaseTiming, + decode: PhaseTiming, + wire_bytes: usize, + roundtrip_max_abs_diff: f32, +} + +struct PhaseTiming { + start_unix_nanos: u64, + end_unix_nanos: u64, + elapsed: Duration, +} + +struct TelemetryRuntime { + telemetry: Telemetry, + _runtime: tokio::runtime::Runtime, +} + +struct MetricsRunGuard<'a> { + config: &'a MlxBoundaryBenchConfig, + finalized: bool, +} + +/// Measures the lazy MLX eval/readback fence and the existing Skippy activation +/// codecs, emits bounded OTLP spans, and exports the canonical metrics report. +pub fn benchmark_mlx_boundary(config: &MlxBoundaryBenchConfig) -> Result { + let config = config.clone(); + thread::spawn(move || benchmark_mlx_boundary_inner(&config)) + .join() + .map_err(|_| anyhow!("MLX boundary benchmark thread panicked"))? +} + +fn benchmark_mlx_boundary_inner(config: &MlxBoundaryBenchConfig) -> Result { + validate_config(config)?; + let stream = Stream::new_with_device(&Device::new(DeviceType::Gpu, 0)); + let (source, offset) = prepared_source(config, &stream)?; + let attrs = benchmark_attrs(config)?; + + for _ in 0..config.warmup_iterations { + run_sample(config, &source, &offset, &stream)?; + } + reset_peak_memory()?; + create_metrics_run(config)?; + let mut metrics_run = MetricsRunGuard::new(config); + + let samples = measure_samples(config, &source, &offset, &stream)?; + let mlx_active_memory_bytes = active_memory()?; + let mlx_cache_memory_bytes = cache_memory()?; + let mlx_peak_memory_bytes = peak_memory()?; + + // Start the exporter only after all timed work has completed. The spans + // retain their original wall-clock timestamps without perturbing samples. + let telemetry_runtime = TelemetryRuntime::new(config)?; + for sample in &samples { + emit_sample(&telemetry_runtime.telemetry, &attrs, sample); + } + let expected_spans = config + .measured_iterations + .checked_mul(4) + .context("telemetry span count overflow")?; + let telemetry = wait_for_telemetry(&telemetry_runtime.telemetry, expected_spans)?; + let canonical_report = finalize_metrics_run(config)?; + metrics_run.mark_finalized(); + let canonical_span_count = canonical_report + .get("counts") + .and_then(|counts| counts.get("spans")) + .and_then(Value::as_u64) + .context("metrics-server report has no span count")?; + ensure!( + canonical_span_count == u64::try_from(expected_spans)?, + "metrics-server stored {canonical_span_count} spans; expected exactly {expected_spans}" + ); + write_metrics_report(config, &canonical_report)?; + + let wire_activation_payload_bytes = samples + .first() + .context("boundary benchmark produced no samples")? + .wire_bytes; + ensure!( + samples + .iter() + .all(|sample| sample.wire_bytes == wire_activation_payload_bytes), + "wire payload size changed between samples" + ); + Ok(MlxBoundaryBenchReport { + benchmark: BENCHMARK_SCHEMA, + code_revision: code_revision(), + metrics_run_id: config.metrics_run_id.clone(), + metrics_report_path: config.metrics_report_path.clone(), + width: config.width, + token_count: config.token_count, + wire_dtype: wire_dtype_label(config.wire_dtype)?.to_string(), + warmup_iterations: config.warmup_iterations, + measured_iterations: config.measured_iterations, + f32_boundary_bytes: boundary_bytes(config, size_of::())?, + wire_activation_payload_bytes, + eval_fence: summarize(samples.iter().map(|sample| sample.eval_fence.elapsed))?, + host_copy: summarize(samples.iter().map(|sample| sample.host_copy.elapsed))?, + eval_and_host_copy_total: summarize( + samples + .iter() + .map(|sample| sample.eval_fence.elapsed + sample.host_copy.elapsed), + )?, + encode: summarize(samples.iter().map(|sample| sample.encode.elapsed))?, + decode: summarize(samples.iter().map(|sample| sample.decode.elapsed))?, + codec_total: summarize( + samples + .iter() + .map(|sample| sample.encode.elapsed + sample.decode.elapsed), + )?, + max_roundtrip_abs_diff: samples + .iter() + .map(|sample| sample.roundtrip_max_abs_diff) + .fold(0.0_f32, f32::max), + mlx_active_memory_bytes, + mlx_cache_memory_bytes, + mlx_peak_memory_bytes, + telemetry, + canonical_span_count, + }) +} + +fn measure_samples( + config: &MlxBoundaryBenchConfig, + source: &Array, + offset: &Array, + stream: &Stream, +) -> Result> { + (0..config.measured_iterations) + .map(|_| run_sample(config, source, offset, stream)) + .collect() +} + +fn validate_config(config: &MlxBoundaryBenchConfig) -> Result<()> { + ensure!(config.width > 0, "boundary width must be non-zero"); + ensure!( + config.token_count > 0, + "boundary token count must be non-zero" + ); + ensure!( + config.measured_iterations > 0, + "boundary benchmark needs measured iterations" + ); + ensure!( + !config.metrics_http.trim().is_empty(), + "metrics-server HTTP endpoint is required" + ); + ensure!( + !config.metrics_otlp_grpc.trim().is_empty(), + "metrics-server OTLP endpoint is required" + ); + validate_metrics_run_id(&config.metrics_run_id)?; + wire_dtype_label(config.wire_dtype)?; + boundary_bytes(config, size_of::())?; + Ok(()) +} + +fn prepared_source(config: &MlxBoundaryBenchConfig, stream: &Stream) -> Result<(Array, Array)> { + let element_count = config + .width + .checked_mul(config.token_count) + .context("boundary element count overflow")?; + let values = (0..element_count) + .map(|index| ((index % 257) as f32 - 128.0) / 127.0) + .collect::>(); + let shape = [ + 1, + i32::try_from(config.token_count)?, + i32::try_from(config.width)?, + ]; + let source = Array::from_slice(&values, &shape).as_dtype(Dtype::Float32, stream)?; + let offset = Array::from_slice(&[0.1_f32], &[1]); + eval([&source, &offset])?; + stream.synchronize()?; + Ok((source, offset)) +} + +fn run_sample( + config: &MlxBoundaryBenchConfig, + source: &Array, + offset: &Array, + stream: &Stream, +) -> Result { + // Graph construction is intentionally outside the completion fence. + let output = source + .add(offset, stream)? + .as_dtype(Dtype::Float32, stream)?; + let ((), eval_fence) = time_phase(|| { + eval([&output])?; + stream.synchronize()?; + Ok(()) + })?; + let (f32_bytes, host_copy) = time_phase(|| { + let evaluated = output.evaluated()?; + Ok(bytemuck::cast_slice(evaluated.as_slice::()).to_vec()) + })?; + let token_count = i32::try_from(config.token_count)?; + let width = i32::try_from(config.width)?; + let (wire_payload, encode) = time_phase(|| { + Ok(encode_f32_activation_payload( + config.wire_dtype, + token_count, + width, + &f32_bytes, + )?) + })?; + let wire_bytes = wire_payload.len(); + let (decoded, decode) = + time_phase(|| Ok(wire_message(config, wire_payload)?.activation_f32_payload(width)?))?; + let roundtrip_max_abs_diff = max_abs_diff(&f32_bytes, &decoded)?; + validate_roundtrip(config.wire_dtype, roundtrip_max_abs_diff)?; + Ok(BoundarySample { + eval_fence, + host_copy, + encode, + decode, + wire_bytes, + roundtrip_max_abs_diff, + }) +} + +fn validate_roundtrip(wire_dtype: WireActivationDType, max_abs_diff: f32) -> Result<()> { + ensure!( + max_abs_diff.is_finite(), + "activation codec produced a non-finite difference" + ); + match wire_dtype { + WireActivationDType::F32 => { + ensure!( + max_abs_diff == 0.0, + "F32 activation codec was not exact: max abs diff {max_abs_diff}" + ); + } + WireActivationDType::F16 => { + ensure!( + max_abs_diff <= 0.001, + "F16 activation codec exceeded synthetic-range error bound: {max_abs_diff} > 0.001" + ); + } + _ => bail!("unsupported activation codec dtype"), + } + Ok(()) +} + +fn wire_message(config: &MlxBoundaryBenchConfig, activation: Vec) -> Result { + let kind = WireMessageKind::PrefillFinalEmbd; + let mut state = StageStateHeader::new(kind, config.wire_dtype); + state.prompt_token_count = i32::try_from(config.token_count)?; + state.source_stage_index = 0; + Ok(StageWireMessage { + kind, + pos_start: 0, + token_count: i32::try_from(config.token_count)?, + state, + request_id: 1, + session_id: 1, + sampling: None, + chat_sampling_metadata: None, + tokens: Vec::new(), + positions: Vec::new(), + activation, + raw_bytes: Vec::new(), + }) +} + +fn time_phase(work: impl FnOnce() -> Result) -> Result<(T, PhaseTiming)> { + let start_unix_nanos = u64::try_from(now_unix_nanos())?; + let start = Instant::now(); + let value = work()?; + let elapsed = start.elapsed(); + let end_unix_nanos = start_unix_nanos + .checked_add(u64::try_from(elapsed.as_nanos())?) + .context("boundary phase timestamp overflow")?; + Ok(( + value, + PhaseTiming { + start_unix_nanos, + end_unix_nanos, + elapsed, + }, + )) +} + +fn emit_sample(telemetry: &Telemetry, attrs: &BTreeMap, sample: &BoundarySample) { + emit_phase( + telemetry, + span::MLX_BOUNDARY_EVAL_FENCE, + attrs, + &sample.eval_fence, + ); + emit_phase( + telemetry, + span::MLX_BOUNDARY_HOST_COPY, + attrs, + &sample.host_copy, + ); + let mut encode_attrs = attrs.clone(); + encode_attrs.insert( + metric::ACTIVATION_BYTES_SENT.to_string(), + json!(sample.wire_bytes), + ); + emit_phase( + telemetry, + span::MLX_BOUNDARY_ENCODE, + &encode_attrs, + &sample.encode, + ); + emit_phase(telemetry, span::MLX_BOUNDARY_DECODE, attrs, &sample.decode); +} + +fn emit_phase( + telemetry: &Telemetry, + name: &str, + attrs: &BTreeMap, + timing: &PhaseTiming, +) { + telemetry.emit_span( + name, + attrs.clone(), + timing.start_unix_nanos, + timing.end_unix_nanos, + ); +} + +fn benchmark_attrs(config: &MlxBoundaryBenchConfig) -> Result> { + let mut attrs = lifecycle_attrs(&telemetry_stage_config(config)?); + attrs.insert(attr::TOKEN_COUNT.to_string(), json!(config.token_count)); + attrs.insert(attr::MESSAGE_KIND.to_string(), json!("mlx_boundary_fence")); + Ok(attrs) +} + +fn validate_metrics_run_id(run_id: &str) -> Result<()> { + ensure!(!run_id.is_empty(), "metrics run ID is required"); + ensure!(run_id.len() <= 128, "metrics run ID exceeds 128 bytes"); + ensure!( + run_id + .bytes() + .all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'.' | b'_' | b'-')), + "metrics run ID may contain only ASCII letters, digits, '.', '_', and '-'" + ); + Ok(()) +} + +fn telemetry_stage_config(config: &MlxBoundaryBenchConfig) -> Result { + Ok(serde_json::from_value(json!({ + "run_id": config.metrics_run_id, + "topology_id": BENCHMARK_SCHEMA, + "model_id": "synthetic/mlx-boundary-fence", + "stage_id": "mlx-boundary-fence", + "stage_index": 0, + "layer_start": 0, + "layer_end": 1, + "ctx_size": config.token_count, + "load_mode": "artifact-slice", + "bind_addr": "127.0.0.1:0" + }))?) +} + +impl TelemetryRuntime { + fn new(config: &MlxBoundaryBenchConfig) -> Result { + let runtime = tokio::runtime::Builder::new_multi_thread() + .worker_threads(2) + .enable_all() + .build() + .context("build boundary telemetry runtime")?; + let telemetry = { + let _guard = runtime.enter(); + Telemetry::new( + Some(config.metrics_otlp_grpc.clone()), + 4096, + telemetry_stage_config(config)?, + TelemetryLevel::Summary, + ) + }; + Ok(Self { + telemetry, + _runtime: runtime, + }) + } +} + +impl<'a> MetricsRunGuard<'a> { + fn new(config: &'a MlxBoundaryBenchConfig) -> Self { + Self { + config, + finalized: false, + } + } + + fn mark_finalized(&mut self) { + self.finalized = true; + } +} + +impl Drop for MetricsRunGuard<'_> { + fn drop(&mut self) { + if !self.finalized { + finalize_metrics_run_best_effort(self.config); + } + } +} + +fn wait_for_telemetry(telemetry: &Telemetry, expected_spans: usize) -> Result { + let deadline = Instant::now() + Duration::from_secs(10); + loop { + let stats = telemetry.stats(); + if stats.sent >= u64::try_from(expected_spans)? { + ensure!(stats.dropped == 0, "boundary telemetry dropped events"); + ensure!(stats.export_errors == 0, "boundary telemetry export failed"); + return Ok(stats); + } + ensure!( + Instant::now() < deadline, + "metrics-server did not ingest {expected_spans} boundary spans; queued={} sent={} dropped={} errors={}", + stats.queued, + stats.sent, + stats.dropped, + stats.export_errors, + ); + thread::sleep(Duration::from_millis(25)); + } +} + +fn create_metrics_run(config: &MlxBoundaryBenchConfig) -> Result<()> { + let response = metrics_client()? + .post(format!( + "{}/v1/runs", + config.metrics_http.trim_end_matches('/') + )) + .json(&metrics_run_body(config)?) + .send() + .context("create metrics-server boundary run")?; + ensure_http_success(response, "create metrics-server boundary run")?; + Ok(()) +} + +fn metrics_run_body(config: &MlxBoundaryBenchConfig) -> Result { + Ok(json!({ + "run_id": config.metrics_run_id, + "benchmark": BENCHMARK_SCHEMA, + "code_revision": code_revision(), + "width": config.width, + "token_count": config.token_count, + "wire_dtype": wire_dtype_label(config.wire_dtype)?, + "warmup_iterations": config.warmup_iterations, + "measured_iterations": config.measured_iterations, + "stages": [{ + "stage_id": "mlx-boundary-fence", + "engine": "mlx", + "model_id": "synthetic/mlx-boundary-fence" + }] + })) +} + +fn code_revision() -> String { + std::env::var("MESH_LLM_BUILD_REVISION") + .ok() + .filter(|revision| validate_metrics_run_id(revision).is_ok()) + .unwrap_or_else(|| env!("CARGO_PKG_VERSION").to_string()) +} + +fn finalize_metrics_run(config: &MlxBoundaryBenchConfig) -> Result { + let client = metrics_client()?; + let base = config.metrics_http.trim_end_matches('/'); + let response = client + .post(format!("{base}/v1/runs/{}/finalize", config.metrics_run_id)) + .send() + .context("finalize metrics-server boundary run")?; + ensure_http_success(response, "finalize metrics-server boundary run")?; + let response = client + .get(format!( + "{base}/v1/runs/{}/report.json", + config.metrics_run_id + )) + .send() + .context("fetch metrics-server boundary report")?; + ensure_http_success(response, "fetch metrics-server boundary report")? + .json() + .context("decode metrics-server boundary report") +} + +fn finalize_metrics_run_best_effort(config: &MlxBoundaryBenchConfig) { + let Ok(client) = metrics_client() else { + return; + }; + let base = config.metrics_http.trim_end_matches('/'); + let _ = client + .post(format!("{base}/v1/runs/{}/finalize", config.metrics_run_id)) + .send(); +} + +fn metrics_client() -> Result { + Ok(reqwest::blocking::Client::builder() + .connect_timeout(Duration::from_secs(5)) + .timeout(Duration::from_secs(30)) + .build()?) +} + +fn ensure_http_success( + response: reqwest::blocking::Response, + operation: &str, +) -> Result { + if response.status().is_success() { + return Ok(response); + } + let status = response.status(); + let body = response.text().unwrap_or_default(); + Err(anyhow!("{operation} failed with HTTP {status}: {body}")) +} + +fn write_metrics_report(config: &MlxBoundaryBenchConfig, report: &Value) -> Result<()> { + if let Some(parent) = config + .metrics_report_path + .parent() + .filter(|parent| !parent.as_os_str().is_empty()) + { + fs::create_dir_all(parent).with_context(|| format!("create {}", parent.display()))?; + } + fs::write( + &config.metrics_report_path, + serde_json::to_vec_pretty(report)?, + ) + .with_context(|| format!("write {}", config.metrics_report_path.display())) +} + +fn max_abs_diff(expected_bytes: &[u8], actual_bytes: &[u8]) -> Result { + ensure!( + expected_bytes.len() == actual_bytes.len(), + "roundtrip activation byte counts differ" + ); + ensure!( + expected_bytes.len().is_multiple_of(size_of::()), + "activation bytes are not aligned to F32 values" + ); + expected_bytes + .chunks_exact(size_of::()) + .zip(actual_bytes.chunks_exact(size_of::())) + .try_fold(0.0_f32, |max_diff, (left, right)| { + let left = f32::from_le_bytes(left.try_into().expect("four-byte chunk")); + let right = f32::from_le_bytes(right.try_into().expect("four-byte chunk")); + ensure!( + left.is_finite() && right.is_finite(), + "activation codec produced a non-finite value" + ); + Ok(max_diff.max((left - right).abs())) + }) +} + +fn summarize(durations: impl Iterator) -> Result { + let mut values = durations + .map(|duration| duration.as_secs_f64() * 1_000_000.0) + .collect::>(); + ensure!(!values.is_empty(), "cannot summarize zero samples"); + values.sort_by(f64::total_cmp); + let mean_us = values.iter().sum::() / values.len() as f64; + Ok(MlxBoundaryDurationSummary { + samples: values.len(), + min_us: values[0], + mean_us, + p50_us: percentile(&values, 0.50), + p95_us: percentile(&values, 0.95), + max_us: values[values.len() - 1], + }) +} + +fn percentile(sorted: &[f64], quantile: f64) -> f64 { + let rank = (quantile * sorted.len() as f64).ceil() as usize; + sorted[rank.saturating_sub(1).min(sorted.len() - 1)] +} + +fn boundary_bytes(config: &MlxBoundaryBenchConfig, element_bytes: usize) -> Result { + config + .width + .checked_mul(config.token_count) + .and_then(|elements| elements.checked_mul(element_bytes)) + .context("boundary byte count overflow") +} + +fn wire_dtype_label(wire_dtype: WireActivationDType) -> Result<&'static str> { + match wire_dtype { + WireActivationDType::F32 => Ok("f32"), + WireActivationDType::F16 => Ok("f16"), + other => bail!("MLX boundary benchmark does not support {other:?}"), + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn duration_summary_uses_nearest_rank_percentiles() { + let summary = + summarize([1_u64, 2, 3, 4, 100].into_iter().map(Duration::from_micros)).unwrap(); + assert_eq!(summary.samples, 5); + assert_eq!(summary.p50_us, 3.0); + assert_eq!(summary.p95_us, 100.0); + } + + #[test] + fn telemetry_config_contains_no_external_or_local_identifiers() { + let config = MlxBoundaryBenchConfig { + width: 4096, + token_count: 32, + wire_dtype: WireActivationDType::F16, + warmup_iterations: 1, + measured_iterations: 2, + metrics_http: "https://collector.invalid/private".to_string(), + metrics_otlp_grpc: "https://otlp.invalid/private".to_string(), + metrics_run_id: "test-run".to_string(), + metrics_report_path: PathBuf::from("/private/report.json"), + }; + let exported = json!({ + "stage_config": telemetry_stage_config(&config).unwrap(), + "span_attributes": benchmark_attrs(&config).unwrap(), + "run_create": metrics_run_body(&config).unwrap(), + }); + let serialized = serde_json::to_string(&exported).unwrap(); + assert!(!serialized.contains("collector.invalid")); + assert!(!serialized.contains("otlp.invalid")); + assert!(!serialized.contains("/private/")); + } + + #[test] + fn metrics_run_id_is_safe_for_export_and_url_paths() { + for valid in ["test-run", "mlx.boundary_01", "A1"] { + validate_metrics_run_id(valid).unwrap(); + } + for invalid in [ + "", + "has/slash", + "has?query", + "https://collector", + "has space", + ] { + assert!(validate_metrics_run_id(invalid).is_err(), "{invalid}"); + } + assert!(validate_metrics_run_id(&"x".repeat(129)).is_err()); + } + + #[test] + fn activation_roundtrip_error_is_gated() { + validate_roundtrip(WireActivationDType::F32, 0.0).unwrap(); + assert!(validate_roundtrip(WireActivationDType::F32, f32::EPSILON).is_err()); + validate_roundtrip(WireActivationDType::F16, 0.001).unwrap(); + assert!(validate_roundtrip(WireActivationDType::F16, 0.001_1).is_err()); + assert!(validate_roundtrip(WireActivationDType::F16, f32::NAN).is_err()); + } +} diff --git a/crates/skippy-engine-mlx/src/lib.rs b/crates/skippy-engine-mlx/src/lib.rs index dc5f65dbbf..406cb7fd14 100644 --- a/crates/skippy-engine-mlx/src/lib.rs +++ b/crates/skippy-engine-mlx/src/lib.rs @@ -10,6 +10,8 @@ #[cfg(all(feature = "mlx", target_os = "macos"))] mod backend; #[cfg(all(feature = "mlx", target_os = "macos"))] +mod boundary_bench; +#[cfg(all(feature = "mlx", target_os = "macos"))] mod derived; #[cfg(all(feature = "mlx", target_os = "macos"))] mod engine; @@ -19,6 +21,8 @@ mod stage; #[cfg(all(feature = "mlx", target_os = "macos"))] pub use backend::MlxBackend; #[cfg(all(feature = "mlx", target_os = "macos"))] +pub use boundary_bench::{MlxBoundaryBenchConfig, MlxBoundaryBenchReport, benchmark_mlx_boundary}; +#[cfg(all(feature = "mlx", target_os = "macos"))] pub use derived::{ MlxDerivationControl, MlxDerivedStageCacheConfig, MlxDerivedStageCacheResult, MlxDerivedStageConfig, MlxDerivedStageReport, MlxDerivedStageShard, diff --git a/crates/skippy-metrics/src/lib.rs b/crates/skippy-metrics/src/lib.rs index e208751b39..8d2f7b9ab6 100644 --- a/crates/skippy-metrics/src/lib.rs +++ b/crates/skippy-metrics/src/lib.rs @@ -57,3 +57,10 @@ pub mod metric { pub const KV_PEER_TRANSFER_ATTEMPTS: &str = "skippy.kv.peer_transfer_attempts"; pub const KV_PEER_TRANSFER_ERRORS: &str = "skippy.kv.peer_transfer_errors"; } + +pub mod span { + pub const MLX_BOUNDARY_EVAL_FENCE: &str = "stage.mlx_boundary_eval_fence"; + pub const MLX_BOUNDARY_HOST_COPY: &str = "stage.mlx_boundary_host_copy"; + pub const MLX_BOUNDARY_ENCODE: &str = "stage.mlx_boundary_encode"; + pub const MLX_BOUNDARY_DECODE: &str = "stage.mlx_boundary_decode"; +} diff --git a/docs/design/MLX_STAGE_ENGINE_PLAN.md b/docs/design/MLX_STAGE_ENGINE_PLAN.md index 4cb9e839b4..f2a63ce94a 100644 --- a/docs/design/MLX_STAGE_ENGINE_PLAN.md +++ b/docs/design/MLX_STAGE_ENGINE_PLAN.md @@ -140,6 +140,19 @@ one real frontier layer. The adjacent final stage and its three-layer topology are synthetic harness devices. It is not evidence for two real Nemotron stages, decode, recurrent state, host/QUIC placement, or full-model logits. +**Update — a metrics-backed frontier-width boundary-fence runner is ready.** +The release `mlx-stage bench-boundary` command separately times synthetic MLX +add completion/eval, host readback and buffer copy, production F32 or F16 +activation encode, and post-receive reconstruction. It calculates real paired +per-iteration boundary and encode-plus-decode distributions, gates codec +correctness, and does not start telemetry export until timed work is complete. +It requires explicit metrics-server HTTP and OTLP endpoints, emits bounded +nonblocking spans, fails on loss or canonical count mismatch, finalizes the +run, and writes the canonical metrics-server report. This is an independent +synthetic fence and codec measurement, not model compute, message framing, +TCP, QUIC, or a complete network gate. The reproduction command and evidence +contract are in `crates/skippy-engine-mlx/STAGED_EXECUTION.md`. + The derivation memory bound is the final packed routed bank, not one expert: six preallocated payload buffers total 718,405,632 bytes. Moving those buffers to a disk-backed random-write spool is the next step if preparation RSS must @@ -930,8 +943,10 @@ Spikes 1 and 2 are more decisive than any standalone token/s benchmark. metadata probes. 2. Teach topology planning and capability advertisement to select MLX stages; explicit host Prepare/Load and the engine-neutral server lane are proven. -3. Run **Spike 2 (boundary fence)** at frontier residual widths and keep it as a - go/no-go gate. +3. Extend the initial metrics-backed synthetic **Spike 2 (boundary fence)** + matrix to real model outputs and TCP/QUIC links. Preserve the separate + eval/synchronize, host copy, codec, and network phase evidence; do not assume + F16 wins when CPU conversion can exceed the bytes saved on a fast link. 4. Extend the proven single-layer Nemotron-H **Nano** `StageEngine` adapter into a hybrid staged runtime with explicit recurrent/attention boundary state. Do not treat Ultra as the same runtime family. Then expose safemlx's existing diff --git a/docs/plugins/telemetry.md b/docs/plugins/telemetry.md index d80a1e8e0a..a9459cb961 100644 --- a/docs/plugins/telemetry.md +++ b/docs/plugins/telemetry.md @@ -96,6 +96,19 @@ capability only. Neither path exports prompts, completions, logs, traces, hostnames, mesh gossip, relay messages, raw node IDs, raw GPU stable IDs, endpoint URLs, or prompt hashes. +Benchmark-only staged-runtime traces are a separate, explicit operator flow. +For example, `mlx-stage bench-boundary` requires metrics-server HTTP and OTLP +arguments on every run and exports only bounded synthetic shape, dtype, +byte-count, duration, validated run-label, schema/revision, and stage +attributes under the shared +`stage.mlx_boundary_*` span names. It does not consume ambient collector +configuration. The collector arguments are explicit transport targets but are +not copied into span attributes or run config. It does not export +prompt/completion text, activation values, paths, collector URLs, raw hardware +identifiers, or model contents. These benchmark traces are not emitted by +normal mesh runtime telemetry. The separately written local report may contain +the output path chosen by the operator; that field is not telemetry. + Guardrail telemetry follows the same boundary. It exports only bounded labels for guardrail mode, contract kind, decision, bypass reason, parser stage, and retry bucket. It does not export prompt text, completion text, schemas, tool From 08c6d74d98a99fd071f799a13e7317f1f6e99c96 Mon Sep 17 00:00:00 2001 From: Michael Neale <14976+michaelneale@users.noreply.github.com> Date: Fri, 17 Jul 2026 22:59:24 +1000 Subject: [PATCH 25/37] docs(mlx): record boundary fence evidence --- crates/skippy-engine-mlx/STAGED_EXECUTION.md | 36 ++++++++++++++++++++ docs/design/MLX_STAGE_ENGINE_PLAN.md | 10 ++++++ 2 files changed, 46 insertions(+) diff --git a/crates/skippy-engine-mlx/STAGED_EXECUTION.md b/crates/skippy-engine-mlx/STAGED_EXECUTION.md index 2695b87ef1..3c874fabba 100644 --- a/crates/skippy-engine-mlx/STAGED_EXECUTION.md +++ b/crates/skippy-engine-mlx/STAGED_EXECUTION.md @@ -284,6 +284,42 @@ transport targets and are not copied into spans or run config. The operator run label is validated to a bounded URL-safe character set; a local output report may contain its explicitly requested report path. +### V2 evidence + +The commit `d381bbd3` matrix ran serially on an Apple M5 Max with 128 GB +unified memory and macOS 26.5.2. Each release-mode process used three warmups +and 20 measured iterations. Boundary p50 is the paired +`eval_and_host_copy_total`; codec p50 is the paired `codec_total`. + +| Width | Tokens | F32 payload | F32 boundary | F32 codec | F16 payload | F16 boundary | F16 codec | +| ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | +| 2,688 | 32 | 0.328 MiB | 0.629 ms | 0.010 ms | 0.164 MiB | 0.236 ms | 0.259 ms | +| 4,096 | 1 | 0.016 MiB | 0.244 ms | 0.001 ms | 0.008 MiB | 0.234 ms | 0.010 ms | +| 4,096 | 32 | 0.500 MiB | 0.580 ms | 0.014 ms | 0.250 MiB | 0.570 ms | 0.306 ms | +| 4,096 | 512 | 8 MiB | 0.832 ms | 0.217 ms | 4 MiB | 0.920 ms | 4.986 ms | +| 8,192 | 1 | 0.031 MiB | 0.255 ms | 0.001 ms | 0.016 MiB | 0.640 ms | 0.020 ms | +| 8,192 | 32 | 1 MiB | 0.631 ms | 0.031 ms | 0.500 MiB | 0.604 ms | 0.596 ms | +| 8,192 | 512 | 16 MiB | 1.212 ms | 0.437 ms | 8 MiB | 1.340 ms | 10.036 ms | +| 16,384 | 1 | 0.063 MiB | 0.613 ms | 0.003 ms | 0.031 MiB | 0.716 ms | 0.042 ms | +| 16,384 | 32 | 2 MiB | 0.609 ms | 0.052 ms | 1 MiB | 0.680 ms | 1.211 ms | +| 16,384 | 512 | 32 MiB | 1.572 ms | 0.927 ms | 16 MiB | 1.651 ms | 20.187 ms | + +Do not over-interpret the one-token or independent F32/F16 boundary timings: +at that scale dispatch, allocator, and process-level noise are material. The +large-prefill codec result is much more stable and nearly linear. At widths +4K, 8K, and 16K, F16 saved 4, 8, and 16 MiB while adding 4.77, 9.60, and +19.26 ms over the F32 codec. With serialized conversion and transfer, that is +an approximately 0.81 GiB/s (7.0 Gbit/s) effective-payload break-even point: +below it F16 should recover its conversion cost from bytes saved; above it F32 +should be faster. Actual selection must be measured per host/link because +conversion can be optimized or overlapped and TCP/QUIC costs are absent here. + +All F32 round trips were exact. All F16 runs stayed finite with maximum absolute +error `0.00045216084`, below the declared `0.001` synthetic-range gate. The 20 +canonical reports contain 20 completed runs and exactly 1,600 spans: 400 for +each phase, zero drops, and zero export errors. Only the 400 encode spans carry +`skippy.activation_bytes_sent`. + Reproduce one matrix cell with metrics-server running in another terminal: ```bash diff --git a/docs/design/MLX_STAGE_ENGINE_PLAN.md b/docs/design/MLX_STAGE_ENGINE_PLAN.md index f2a63ce94a..6314afc43e 100644 --- a/docs/design/MLX_STAGE_ENGINE_PLAN.md +++ b/docs/design/MLX_STAGE_ENGINE_PLAN.md @@ -153,6 +153,16 @@ synthetic fence and codec measurement, not model compute, message framing, TCP, QUIC, or a complete network gate. The reproduction command and evidence contract are in `crates/skippy-engine-mlx/STAGED_EXECUTION.md`. +**V2 result.** Commit `d381bbd3` produced 20 completed canonical runs / 1,600 +spans with no telemetry loss. For 512-token boundaries at widths 4K, 8K, and +16K, F16 halved payloads while adding 4.77, 9.60, and 19.26 ms over F32 codec +work. With conversion and transfer serialized, those three cells independently +place the effective-payload break-even near 0.81 GiB/s (7.0 Gbit/s): use F16 +below that measured link rate and F32 above it, subject to a real TCP/QUIC test +and any conversion/transfer overlap. The 16K cell measured 32 MiB / 0.927 ms +for F32 versus 16 MiB / 20.187 ms for F16. This makes wire dtype a hardware and +link policy choice, not a model-format constant. + The derivation memory bound is the final packed routed bank, not one expert: six preallocated payload buffers total 718,405,632 bytes. Moving those buffers to a disk-backed random-write spool is the next step if preparation RSS must From 6350e3a98f0f856c0f01224e1b4cc0cbf8b2e3c7 Mon Sep 17 00:00:00 2001 From: Michael Neale <14976+michaelneale@users.noreply.github.com> Date: Fri, 17 Jul 2026 23:09:31 +1000 Subject: [PATCH 26/37] bench(mlx): measure production TCP boundary --- crates/metrics-server/README.md | 15 +- crates/skippy-engine-mlx/STAGED_EXECUTION.md | 35 + crates/skippy-engine-mlx/src/bin/mlx-stage.rs | 72 +- .../skippy-engine-mlx/src/boundary_bench.rs | 4 + .../src/boundary_bench/tcp.rs | 709 ++++++++++++++++++ crates/skippy-engine-mlx/src/lib.rs | 5 +- crates/skippy-metrics/src/lib.rs | 1 + docs/design/MLX_STAGE_ENGINE_PLAN.md | 7 + 8 files changed, 840 insertions(+), 8 deletions(-) create mode 100644 crates/skippy-engine-mlx/src/boundary_bench/tcp.rs diff --git a/crates/metrics-server/README.md b/crates/metrics-server/README.md index 71fda7174e..3795cf7e50 100644 --- a/crates/metrics-server/README.md +++ b/crates/metrics-server/README.md @@ -33,10 +33,19 @@ The opt-in MLX boundary benchmark emits four shared span names: - `stage.mlx_boundary_encode` - `stage.mlx_boundary_decode` +The follow-on production loopback-TCP benchmark emits: + +- `stage.mlx_boundary_tcp_roundtrip` + +That span covers sender-side activation encoding, Skippy message framing and +loopback TCP, engine-transport activation reconstruction, the synthetic sink +adapter, and the predicted-token reply. + Its run config and spans contain bounded benchmark shape, dtype, byte-count, -schema/revision, and duration data only. The activation-byte metric is attached -only to the encode span. The benchmark requires explicit HTTP and OTLP -collector arguments; those transport targets are not copied into the run +schema/revision, and duration data only. In the four-phase benchmark, the +activation-byte metric is attached only to the encode span; the TCP benchmark +attaches it once to each round-trip span. The benchmark requires explicit HTTP +and OTLP collector arguments; those transport targets are not copied into the run config or span attributes. It does not export activation values, prompts, paths, endpoint URLs, hardware identifiers, or model contents. diff --git a/crates/skippy-engine-mlx/STAGED_EXECUTION.md b/crates/skippy-engine-mlx/STAGED_EXECUTION.md index 3c874fabba..bc9d9523ea 100644 --- a/crates/skippy-engine-mlx/STAGED_EXECUTION.md +++ b/crates/skippy-engine-mlx/STAGED_EXECUTION.md @@ -339,6 +339,41 @@ just mlx-stage bench-boundary \ --output /tmp/mlx-boundary-local.json ``` +### Production loopback-TCP follow-on + +`mlx-stage bench-tcp-boundary` moves the same synthetic activation through the +production engine-neutral Skippy TCP server. Its paired round-trip timer starts +before F32/F16 activation encoding and ends after the predicted-token reply, so +it includes sender encoding, binary framing/write, loopback TCP, server +read/framing, F32 reconstruction, the synthetic final `StageEngine` adapter, +and reply framing/read. It also includes construction and destruction of the +message plus its token/position sidebands and the sink/reply assertions. The +reported `wire_activation_payload_bytes` excludes the fixed frame, eight bytes +per token of sidebands, and the reply. It excludes MLX/model compute, QUIC, +remote links, and the outbound activation encoding of a non-final stage. + +Connection bind/connect/READY and teardown are outside the timer. Samples run +sequentially over one warmed persistent connection with client and server as +threads in the same process. This is steady-state loopback latency, not +connection startup, multi-process behavior, concurrent throughput, or pipeline +overlap. + +The first warmup validates the complete decoded tensor against the source using +the same exact-F32 / bounded-F16 gate. Measured samples finish before telemetry +starts. Each sample then becomes one +`stage.mlx_boundary_tcp_roundtrip` span in a canonical metrics-server run. + +```bash +just mlx-stage bench-tcp-boundary \ + --width 16384 --tokens 512 --wire-dtype f16 \ + --warmup-iterations 3 --measured-iterations 20 \ + --metrics-http http://127.0.0.1:18081 \ + --metrics-otlp-grpc http://127.0.0.1:14317 \ + --metrics-run-id mlx-tcp-boundary-w16384-t512-f16-v1 \ + --metrics-report /tmp/mlx-tcp-boundary-metrics.json \ + --output /tmp/mlx-tcp-boundary-local.json +``` + ## Reproduce Build once: diff --git a/crates/skippy-engine-mlx/src/bin/mlx-stage.rs b/crates/skippy-engine-mlx/src/bin/mlx-stage.rs index 10540d8a19..eed1fefcc3 100644 --- a/crates/skippy-engine-mlx/src/bin/mlx-stage.rs +++ b/crates/skippy-engine-mlx/src/bin/mlx-stage.rs @@ -14,10 +14,11 @@ mod real { use model_hf::safetensors_stage::{SafetensorsStageMaterializer, SafetensorsStageRequest}; use skippy_engine_mlx::{ MlxBoundaryBenchConfig, MlxComputeDtype, MlxDerivationControl, MlxDerivedStageCacheConfig, - MlxDerivedStageConfig, MlxStageEngine, MlxStageEngineConfig, MlxWeightQuantization, - benchmark_mlx_boundary, derive_quantized_stage, derive_quantized_stage_cached, - mlx_derived_stage_cache_root, validate_nemotron_h_binary_wire_tokens, - validate_nemotron_h_moe_stage, validate_nemotron_h_stage_engine, + MlxDerivedStageConfig, MlxStageEngine, MlxStageEngineConfig, MlxTcpBoundaryBenchConfig, + MlxWeightQuantization, benchmark_mlx_boundary, benchmark_mlx_tcp_boundary, + derive_quantized_stage, derive_quantized_stage_cached, mlx_derived_stage_cache_root, + validate_nemotron_h_binary_wire_tokens, validate_nemotron_h_moe_stage, + validate_nemotron_h_stage_engine, }; use skippy_protocol::binary::{ StageStateHeader, StageWireMessage, WireActivationDType, WireMessageKind, WireReplyKind, @@ -148,6 +149,29 @@ mod real { #[arg(long)] output: Option, }, + /// Measure one activation through production Skippy loopback TCP. + BenchTcpBoundary { + #[arg(long)] + width: usize, + #[arg(long)] + tokens: usize, + #[arg(long, value_enum)] + wire_dtype: WireDtype, + #[arg(long, default_value_t = 3)] + warmup_iterations: usize, + #[arg(long, default_value_t = 20)] + measured_iterations: usize, + #[arg(long)] + metrics_http: String, + #[arg(long)] + metrics_otlp_grpc: String, + #[arg(long)] + metrics_run_id: Option, + #[arg(long)] + metrics_report: PathBuf, + #[arg(long)] + output: Option, + }, /// Drive a stage chain and assert its greedy token sequence. Prove { #[arg(long)] @@ -353,6 +377,46 @@ mod real { println!("{}", String::from_utf8(json)?); Ok(()) } + Command::BenchTcpBoundary { + width, + tokens, + wire_dtype, + warmup_iterations, + measured_iterations, + metrics_http, + metrics_otlp_grpc, + metrics_run_id, + metrics_report, + output, + } => { + let metrics_run_id = match metrics_run_id { + Some(metrics_run_id) => metrics_run_id, + None => default_boundary_run_id()?, + }; + let report = benchmark_mlx_tcp_boundary(&MlxTcpBoundaryBenchConfig { + width, + token_count: tokens, + wire_dtype: wire_dtype.into(), + warmup_iterations, + measured_iterations, + metrics_http, + metrics_otlp_grpc, + metrics_run_id, + metrics_report_path: metrics_report, + })?; + let json = serde_json::to_vec_pretty(&report)?; + if let Some(output) = output { + if let Some(parent) = output + .parent() + .filter(|parent| !parent.as_os_str().is_empty()) + { + std::fs::create_dir_all(parent)?; + } + std::fs::write(output, &json)?; + } + println!("{}", String::from_utf8(json)?); + Ok(()) + } Command::Prove { connect, tokens, diff --git a/crates/skippy-engine-mlx/src/boundary_bench.rs b/crates/skippy-engine-mlx/src/boundary_bench.rs index 1a6064da20..b35d67c56a 100644 --- a/crates/skippy-engine-mlx/src/boundary_bench.rs +++ b/crates/skippy-engine-mlx/src/boundary_bench.rs @@ -1,5 +1,9 @@ //! Metrics-server-backed measurements for the MLX activation boundary fence. +mod tcp; + +pub use tcp::{MlxTcpBoundaryBenchConfig, MlxTcpBoundaryBenchReport, benchmark_mlx_tcp_boundary}; + use std::{ collections::BTreeMap, fs, diff --git a/crates/skippy-engine-mlx/src/boundary_bench/tcp.rs b/crates/skippy-engine-mlx/src/boundary_bench/tcp.rs new file mode 100644 index 0000000000..c9f48f359d --- /dev/null +++ b/crates/skippy-engine-mlx/src/boundary_bench/tcp.rs @@ -0,0 +1,709 @@ +//! Production Skippy TCP framing measurements for synthetic activation boundaries. + +use std::{ + collections::BTreeMap, + fs, + net::{SocketAddr, TcpListener, TcpStream}, + path::PathBuf, + sync::{ + Arc, Mutex, + atomic::{AtomicBool, Ordering}, + }, + thread, + time::{Duration, Instant}, +}; + +use anyhow::{Context, Result, anyhow, ensure}; +use serde::Serialize; +use serde_json::{Value, json}; +use skippy_engine::{ + StageEngine, StageEngineInfo, StageExecutionKind, StageExecutionOutput, StageExecutionRequest, +}; +use skippy_metrics::{attr, metric, span}; +use skippy_protocol::{ + StageConfig, + binary::{ + MAX_STAGE_ACTIVATION_BYTES, MAX_STAGE_DECODED_ACTIVATION_BYTES, MAX_STAGE_SIDEBAND_VALUES, + StageStateHeader, StageWireMessage, WireActivationDType, WireMessageKind, WireReplyKind, + activation_wire_bytes, encode_f32_activation_payload, recv_ready, recv_reply, + write_stage_message, + }, +}; +use skippy_server::{ + engine_transport::{EngineStageServerOptions, serve_stage_engine_until}, + telemetry::{Telemetry, TelemetryLevel, TelemetryStats, lifecycle_attrs}, +}; + +use super::{ + MlxBoundaryDurationSummary, PhaseTiming, code_revision, ensure_http_success, max_abs_diff, + metrics_client, summarize, time_phase, validate_metrics_run_id, validate_roundtrip, + wait_for_telemetry, wire_dtype_label, +}; + +const BENCHMARK_SCHEMA: &str = "mlx-tcp-boundary-v1"; +const PREDICTED_SENTINEL: i32 = 424_243; + +/// One explicit, reproducible production-TCP boundary benchmark run. +#[derive(Clone, Debug)] +pub struct MlxTcpBoundaryBenchConfig { + pub width: usize, + pub token_count: usize, + pub wire_dtype: WireActivationDType, + pub warmup_iterations: usize, + pub measured_iterations: usize, + pub metrics_http: String, + pub metrics_otlp_grpc: String, + pub metrics_run_id: String, + pub metrics_report_path: PathBuf, +} + +/// Sender encode through production server decode and predicted reply. +#[derive(Clone, Debug, Serialize)] +pub struct MlxTcpBoundaryBenchReport { + pub benchmark: &'static str, + pub code_revision: String, + pub metrics_run_id: String, + pub metrics_report_path: PathBuf, + pub width: usize, + pub token_count: usize, + pub wire_dtype: String, + pub warmup_iterations: usize, + pub measured_iterations: usize, + pub f32_boundary_bytes: usize, + pub wire_activation_payload_bytes: usize, + pub tcp_roundtrip: MlxBoundaryDurationSummary, + pub warmup_roundtrip_max_abs_diff: f32, + pub telemetry: TelemetryStats, + pub canonical_span_count: u64, +} + +/// Measures production activation encoding, Skippy framing, loopback TCP, +/// engine-transport decoding, and the predicted-token reply as one fence. +pub fn benchmark_mlx_tcp_boundary( + config: &MlxTcpBoundaryBenchConfig, +) -> Result { + let config = config.clone(); + thread::spawn(move || benchmark_mlx_tcp_boundary_inner(&config)) + .join() + .map_err(|_| anyhow!("MLX TCP boundary benchmark thread panicked"))? +} + +fn benchmark_mlx_tcp_boundary_inner( + config: &MlxTcpBoundaryBenchConfig, +) -> Result { + validate_config(config)?; + let source = source_bytes(config)?; + let engine = Arc::new(TcpBoundarySink::new(config, Arc::clone(&source))?); + let (server, mut client) = TcpStageServer::spawn_ready( + engine.clone(), + EngineStageServerOptions { + bind_addr: "127.0.0.1:0".parse()?, + downstream_addr: None, + wire_dtype: config.wire_dtype, + }, + )?; + + for iteration in 0..config.warmup_iterations { + run_roundtrip(config, &source, &mut client, u64::try_from(iteration)?)?; + } + let warmup_roundtrip_max_abs_diff = engine + .validation_diff()? + .context("TCP boundary warmup did not reach the sink engine")?; + + let samples = measure_roundtrips(config, &source, &mut client)?; + drop(client); + server.stop()?; + + create_metrics_run(config)?; + let mut metrics_run = MetricsRunGuard::new(config); + let telemetry_runtime = TcpTelemetryRuntime::new(config)?; + let attrs = benchmark_attrs(config)?; + for sample in &samples { + telemetry_runtime.telemetry.emit_span( + span::MLX_BOUNDARY_TCP_ROUNDTRIP, + attrs.clone(), + sample.start_unix_nanos, + sample.end_unix_nanos, + ); + } + let telemetry = wait_for_telemetry(&telemetry_runtime.telemetry, samples.len())?; + let canonical_report = finalize_metrics_run(config)?; + metrics_run.mark_finalized(); + let canonical_span_count = canonical_report + .get("counts") + .and_then(|counts| counts.get("spans")) + .and_then(Value::as_u64) + .context("metrics-server report has no span count")?; + ensure!( + canonical_span_count == u64::try_from(samples.len())?, + "metrics-server stored {canonical_span_count} spans; expected exactly {}", + samples.len() + ); + write_metrics_report(config, &canonical_report)?; + + Ok(MlxTcpBoundaryBenchReport { + benchmark: BENCHMARK_SCHEMA, + code_revision: code_revision(), + metrics_run_id: config.metrics_run_id.clone(), + metrics_report_path: config.metrics_report_path.clone(), + width: config.width, + token_count: config.token_count, + wire_dtype: wire_dtype_label(config.wire_dtype)?.to_string(), + warmup_iterations: config.warmup_iterations, + measured_iterations: config.measured_iterations, + f32_boundary_bytes: activation_bytes(config, size_of::())?, + wire_activation_payload_bytes: activation_bytes( + config, + match config.wire_dtype { + WireActivationDType::F32 => size_of::(), + WireActivationDType::F16 => size_of::(), + _ => unreachable!("validated wire dtype"), + }, + )?, + tcp_roundtrip: summarize(samples.iter().map(|sample| sample.elapsed))?, + warmup_roundtrip_max_abs_diff, + telemetry, + canonical_span_count, + }) +} + +fn validate_config(config: &MlxTcpBoundaryBenchConfig) -> Result<()> { + ensure!(config.width > 0, "TCP boundary width must be non-zero"); + ensure!( + config.token_count > 0, + "TCP boundary token count must be non-zero" + ); + ensure!( + config.warmup_iterations > 0, + "TCP boundary benchmark needs a correctness warmup" + ); + ensure!( + config.measured_iterations > 0, + "TCP boundary benchmark needs measured iterations" + ); + ensure!( + !config.metrics_http.trim().is_empty(), + "metrics-server HTTP endpoint is required" + ); + ensure!( + !config.metrics_otlp_grpc.trim().is_empty(), + "metrics-server OTLP endpoint is required" + ); + validate_metrics_run_id(&config.metrics_run_id)?; + wire_dtype_label(config.wire_dtype)?; + let token_count = i32::try_from(config.token_count).context("token count exceeds i32")?; + let width = i32::try_from(config.width).context("activation width exceeds i32")?; + u32::try_from(config.width).context("activation width exceeds u32")?; + ensure!( + config.token_count <= MAX_STAGE_SIDEBAND_VALUES, + "token sideband count exceeds protocol maximum" + ); + let wire_bytes = activation_wire_bytes(config.wire_dtype, token_count, width)?; + ensure!( + wire_bytes <= MAX_STAGE_ACTIVATION_BYTES, + "wire activation exceeds protocol maximum" + ); + let decoded_bytes = activation_wire_bytes(WireActivationDType::F32, token_count, width)?; + ensure!( + decoded_bytes <= MAX_STAGE_DECODED_ACTIVATION_BYTES, + "decoded activation exceeds protocol maximum" + ); + Ok(()) +} + +fn source_bytes(config: &MlxTcpBoundaryBenchConfig) -> Result>> { + let elements = config + .width + .checked_mul(config.token_count) + .context("TCP boundary element count overflow")?; + Ok(Arc::new( + (0..elements) + .flat_map(|index| { + let value = ((index % 257) as f32 - 128.0) / 127.0 + 0.1; + value.to_le_bytes() + }) + .collect(), + )) +} + +fn measure_roundtrips( + config: &MlxTcpBoundaryBenchConfig, + source: &[u8], + client: &mut TcpStream, +) -> Result> { + (0..config.measured_iterations) + .map(|iteration| { + let request_id = u64::try_from(config.warmup_iterations + iteration)?; + let ((), timing) = time_phase(|| run_roundtrip(config, source, client, request_id))?; + Ok(timing) + }) + .collect() +} + +fn run_roundtrip( + config: &MlxTcpBoundaryBenchConfig, + source: &[u8], + client: &mut TcpStream, + request_id: u64, +) -> Result<()> { + let token_count = i32::try_from(config.token_count)?; + let width = i32::try_from(config.width)?; + let kind = WireMessageKind::PrefillFinalEmbd; + let mut state = StageStateHeader::new(kind, config.wire_dtype); + state.prompt_token_count = token_count; + state.source_stage_index = 0; + let message = StageWireMessage { + kind, + pos_start: 0, + token_count, + state, + request_id, + session_id: 1, + sampling: None, + chat_sampling_metadata: None, + tokens: vec![0; config.token_count], + positions: (0..token_count).collect(), + activation: encode_f32_activation_payload(config.wire_dtype, token_count, width, source)?, + raw_bytes: Vec::new(), + }; + write_stage_message(&mut *client, &message, config.wire_dtype)?; + use std::io::Write as _; + client.flush()?; + let reply = recv_reply(&mut *client)?; + ensure!( + reply.kind == WireReplyKind::PredictedToken + && reply.predicted == PREDICTED_SENTINEL + && reply.predicted_tokens == [PREDICTED_SENTINEL], + "TCP boundary sink returned the wrong prediction" + ); + Ok(()) +} + +struct TcpBoundarySink { + info: StageEngineInfo, + token_count: usize, + wire_dtype: WireActivationDType, + source: Arc>, + validated: AtomicBool, + validation_diff: Mutex>, +} + +impl TcpBoundarySink { + fn new(config: &MlxTcpBoundaryBenchConfig, source: Arc>) -> Result { + let info = StageEngineInfo { + engine: "mlx-tcp-boundary-sink".to_string(), + model_id: "synthetic/mlx-tcp-boundary".to_string(), + stage_index: 1, + layer_start: 1, + layer_end: 2, + total_layers: 2, + activation_width: u32::try_from(config.width)?, + }; + info.validate()?; + Ok(Self { + info, + token_count: config.token_count, + wire_dtype: config.wire_dtype, + source, + validated: AtomicBool::new(false), + validation_diff: Mutex::new(None), + }) + } + + fn validation_diff(&self) -> Result> { + Ok(*self + .validation_diff + .lock() + .map_err(|_| anyhow!("TCP boundary validation lock poisoned"))?) + } +} + +impl StageEngine for TcpBoundarySink { + fn info(&self) -> &StageEngineInfo { + &self.info + } + + fn execute(&self, request: StageExecutionRequest) -> Result { + ensure!( + request.kind == StageExecutionKind::PrefillFinal, + "TCP boundary sink received the wrong execution kind" + ); + let activation = request + .input + .context("TCP boundary sink requires an activation")?; + ensure!( + activation.token_count == self.token_count + && activation.width == self.info.activation_width as usize, + "TCP boundary sink received the wrong activation shape" + ); + if !self.validated.swap(true, Ordering::SeqCst) { + let diff = max_abs_diff(&self.source, &activation.f32_le_bytes)?; + validate_roundtrip(self.wire_dtype, diff)?; + *self + .validation_diff + .lock() + .map_err(|_| anyhow!("TCP boundary validation lock poisoned"))? = Some(diff); + } + Ok(StageExecutionOutput { + activation: None, + predicted_tokens: vec![PREDICTED_SENTINEL], + }) + } + + fn reset_session(&self, _session_id: u64) -> Result<()> { + Ok(()) + } +} + +struct TcpStageServer { + shutdown: Arc, + join: Option>>, +} + +impl TcpStageServer { + fn spawn(engine: Arc, options: EngineStageServerOptions) -> Self { + let shutdown = Arc::new(AtomicBool::new(false)); + let thread_shutdown = Arc::clone(&shutdown); + let join = + thread::spawn(move || serve_stage_engine_until(engine, options, thread_shutdown)); + Self { + shutdown, + join: Some(join), + } + } + + fn spawn_ready( + engine: Arc, + options: EngineStageServerOptions, + ) -> Result<(Self, TcpStream)> { + const BIND_ATTEMPTS: usize = 3; + let mut last_error = None; + for _ in 0..BIND_ATTEMPTS { + let addr = reserve_loopback_addr()?; + let mut attempt_options = options.clone(); + attempt_options.bind_addr = addr; + let server = Self::spawn(Arc::clone(&engine), attempt_options); + match connect_ready(addr) { + Ok(client) => return Ok((server, client)), + Err(error) => { + let server_error = server.stop().err(); + last_error = Some(match server_error { + Some(server_error) => anyhow!( + "connect TCP boundary stage at {addr}: {error:#}; server failed: {server_error:#}" + ), + None => error, + }); + } + } + } + Err(last_error.unwrap_or_else(|| anyhow!("TCP boundary stage did not start"))) + } + + fn stop(mut self) -> Result<()> { + self.finish() + } + + fn finish(&mut self) -> Result<()> { + self.shutdown.store(true, Ordering::SeqCst); + let Some(join) = self.join.take() else { + return Ok(()); + }; + join.join() + .map_err(|_| anyhow!("TCP boundary server panicked"))? + } +} + +impl Drop for TcpStageServer { + fn drop(&mut self) { + let _ = self.finish(); + } +} + +fn reserve_loopback_addr() -> Result { + let listener = TcpListener::bind("127.0.0.1:0")?; + Ok(listener.local_addr()?) +} + +fn connect_ready(addr: SocketAddr) -> Result { + let deadline = Instant::now() + Duration::from_secs(5); + loop { + let error = match TcpStream::connect(addr) { + Ok(mut stream) => { + stream.set_nodelay(true).ok(); + stream.set_read_timeout(Some(Duration::from_millis(250)))?; + stream.set_write_timeout(Some(Duration::from_secs(5)))?; + match recv_ready(&mut stream) { + Ok(()) => { + stream.set_read_timeout(Some(Duration::from_secs(10)))?; + return Ok(stream); + } + Err(error) => anyhow!(error).context("receive TCP boundary ready handshake"), + } + } + Err(error) => anyhow!(error).context("connect TCP boundary socket"), + }; + if Instant::now() >= deadline { + return Err(error).with_context(|| format!("connect TCP boundary stage at {addr}")); + } + thread::sleep(Duration::from_millis(25)); + } +} + +struct TcpTelemetryRuntime { + telemetry: Telemetry, + _runtime: tokio::runtime::Runtime, +} + +impl TcpTelemetryRuntime { + fn new(config: &MlxTcpBoundaryBenchConfig) -> Result { + let runtime = tokio::runtime::Builder::new_multi_thread() + .worker_threads(2) + .enable_all() + .build() + .context("build TCP boundary telemetry runtime")?; + let telemetry = { + let _guard = runtime.enter(); + Telemetry::new( + Some(config.metrics_otlp_grpc.clone()), + 4096, + telemetry_stage_config(config)?, + TelemetryLevel::Summary, + ) + }; + Ok(Self { + telemetry, + _runtime: runtime, + }) + } +} + +fn benchmark_attrs(config: &MlxTcpBoundaryBenchConfig) -> Result> { + let mut attrs = lifecycle_attrs(&telemetry_stage_config(config)?); + attrs.insert(attr::TOKEN_COUNT.to_string(), json!(config.token_count)); + attrs.insert( + attr::MESSAGE_KIND.to_string(), + json!("mlx_boundary_tcp_roundtrip"), + ); + attrs.insert( + metric::ACTIVATION_BYTES_SENT.to_string(), + json!(activation_bytes( + config, + match config.wire_dtype { + WireActivationDType::F32 => size_of::(), + WireActivationDType::F16 => size_of::(), + _ => unreachable!("validated wire dtype"), + } + )?), + ); + Ok(attrs) +} + +fn telemetry_stage_config(config: &MlxTcpBoundaryBenchConfig) -> Result { + Ok(serde_json::from_value(json!({ + "run_id": config.metrics_run_id, + "topology_id": BENCHMARK_SCHEMA, + "model_id": "synthetic/mlx-tcp-boundary", + "stage_id": "mlx-tcp-boundary", + "stage_index": 1, + "layer_start": 1, + "layer_end": 2, + "ctx_size": config.token_count, + "load_mode": "artifact-slice", + "bind_addr": "127.0.0.1:0" + }))?) +} + +fn create_metrics_run(config: &MlxTcpBoundaryBenchConfig) -> Result<()> { + let response = metrics_client()? + .post(format!( + "{}/v1/runs", + config.metrics_http.trim_end_matches('/') + )) + .json(&metrics_run_body(config)?) + .send() + .context("create metrics-server TCP boundary run")?; + ensure_http_success(response, "create metrics-server TCP boundary run")?; + Ok(()) +} + +fn metrics_run_body(config: &MlxTcpBoundaryBenchConfig) -> Result { + Ok(json!({ + "run_id": config.metrics_run_id, + "benchmark": BENCHMARK_SCHEMA, + "code_revision": code_revision(), + "width": config.width, + "token_count": config.token_count, + "wire_dtype": wire_dtype_label(config.wire_dtype)?, + "warmup_iterations": config.warmup_iterations, + "measured_iterations": config.measured_iterations, + "stages": [{ + "stage_id": "mlx-tcp-boundary", + "engine": "skippy-engine-transport", + "model_id": "synthetic/mlx-tcp-boundary" + }] + })) +} + +fn finalize_metrics_run(config: &MlxTcpBoundaryBenchConfig) -> Result { + let client = metrics_client()?; + let base = config.metrics_http.trim_end_matches('/'); + let response = client + .post(format!("{base}/v1/runs/{}/finalize", config.metrics_run_id)) + .send() + .context("finalize metrics-server TCP boundary run")?; + ensure_http_success(response, "finalize metrics-server TCP boundary run")?; + let response = client + .get(format!( + "{base}/v1/runs/{}/report.json", + config.metrics_run_id + )) + .send() + .context("fetch metrics-server TCP boundary report")?; + ensure_http_success(response, "fetch metrics-server TCP boundary report")? + .json() + .context("decode metrics-server TCP boundary report") +} + +fn finalize_metrics_run_best_effort(config: &MlxTcpBoundaryBenchConfig) { + let Ok(client) = metrics_client() else { + return; + }; + let base = config.metrics_http.trim_end_matches('/'); + let _ = client + .post(format!("{base}/v1/runs/{}/finalize", config.metrics_run_id)) + .send(); +} + +struct MetricsRunGuard<'a> { + config: &'a MlxTcpBoundaryBenchConfig, + finalized: bool, +} + +impl<'a> MetricsRunGuard<'a> { + fn new(config: &'a MlxTcpBoundaryBenchConfig) -> Self { + Self { + config, + finalized: false, + } + } + + fn mark_finalized(&mut self) { + self.finalized = true; + } +} + +impl Drop for MetricsRunGuard<'_> { + fn drop(&mut self) { + if !self.finalized { + finalize_metrics_run_best_effort(self.config); + } + } +} + +fn write_metrics_report(config: &MlxTcpBoundaryBenchConfig, report: &Value) -> Result<()> { + if let Some(parent) = config + .metrics_report_path + .parent() + .filter(|parent| !parent.as_os_str().is_empty()) + { + fs::create_dir_all(parent).with_context(|| format!("create {}", parent.display()))?; + } + fs::write( + &config.metrics_report_path, + serde_json::to_vec_pretty(report)?, + ) + .with_context(|| format!("write {}", config.metrics_report_path.display())) +} + +fn activation_bytes(config: &MlxTcpBoundaryBenchConfig, element_bytes: usize) -> Result { + config + .width + .checked_mul(config.token_count) + .and_then(|elements| elements.checked_mul(element_bytes)) + .context("TCP boundary byte count overflow") +} + +#[cfg(test)] +mod tests { + use super::*; + + fn config() -> MlxTcpBoundaryBenchConfig { + MlxTcpBoundaryBenchConfig { + width: 4096, + token_count: 32, + wire_dtype: WireActivationDType::F16, + warmup_iterations: 1, + measured_iterations: 2, + metrics_http: "https://collector.invalid/private".to_string(), + metrics_otlp_grpc: "https://otlp.invalid/private".to_string(), + metrics_run_id: "tcp-test-run".to_string(), + metrics_report_path: PathBuf::from("/private/report.json"), + } + } + + #[test] + fn telemetry_contains_no_transport_targets_or_local_paths() { + let config = config(); + let exported = json!({ + "stage_config": telemetry_stage_config(&config).unwrap(), + "span_attributes": benchmark_attrs(&config).unwrap(), + "run_create": metrics_run_body(&config).unwrap(), + }); + let serialized = serde_json::to_string(&exported).unwrap(); + assert!(!serialized.contains("collector.invalid")); + assert!(!serialized.contains("otlp.invalid")); + assert!(!serialized.contains("/private/")); + } + + #[test] + fn source_values_are_finite_and_in_the_codec_gate_range() { + let config = config(); + let source = source_bytes(&config).unwrap(); + assert_eq!(source.len(), 4096 * 32 * size_of::()); + assert!( + source + .chunks_exact(4) + .all(|chunk| { f32::from_le_bytes(chunk.try_into().unwrap()).is_finite() }) + ); + } + + #[test] + fn production_tcp_roundtrip_reaches_and_validates_sink() { + for (wire_dtype, expected_diff) in [ + (WireActivationDType::F32, 0.0), + (WireActivationDType::F16, 0.001), + ] { + let mut config = config(); + config.width = 8; + config.token_count = 2; + config.wire_dtype = wire_dtype; + let source = source_bytes(&config).unwrap(); + let engine = Arc::new(TcpBoundarySink::new(&config, Arc::clone(&source)).unwrap()); + let (server, mut client) = TcpStageServer::spawn_ready( + engine.clone(), + EngineStageServerOptions { + bind_addr: "127.0.0.1:0".parse().unwrap(), + downstream_addr: None, + wire_dtype: config.wire_dtype, + }, + ) + .unwrap(); + + run_roundtrip(&config, &source, &mut client, 1).unwrap(); + assert!(engine.validation_diff().unwrap().unwrap() <= expected_diff); + drop(client); + server.stop().unwrap(); + } + } + + #[test] + fn oversized_protocol_shapes_fail_before_source_allocation() { + let mut config = config(); + config.token_count = MAX_STAGE_SIDEBAND_VALUES + 1; + assert!(validate_config(&config).is_err()); + + config.token_count = 1; + config.width = MAX_STAGE_DECODED_ACTIVATION_BYTES / 4 + 1; + assert!(validate_config(&config).is_err()); + } +} diff --git a/crates/skippy-engine-mlx/src/lib.rs b/crates/skippy-engine-mlx/src/lib.rs index 406cb7fd14..1e93b3f42a 100644 --- a/crates/skippy-engine-mlx/src/lib.rs +++ b/crates/skippy-engine-mlx/src/lib.rs @@ -21,7 +21,10 @@ mod stage; #[cfg(all(feature = "mlx", target_os = "macos"))] pub use backend::MlxBackend; #[cfg(all(feature = "mlx", target_os = "macos"))] -pub use boundary_bench::{MlxBoundaryBenchConfig, MlxBoundaryBenchReport, benchmark_mlx_boundary}; +pub use boundary_bench::{ + MlxBoundaryBenchConfig, MlxBoundaryBenchReport, MlxTcpBoundaryBenchConfig, + MlxTcpBoundaryBenchReport, benchmark_mlx_boundary, benchmark_mlx_tcp_boundary, +}; #[cfg(all(feature = "mlx", target_os = "macos"))] pub use derived::{ MlxDerivationControl, MlxDerivedStageCacheConfig, MlxDerivedStageCacheResult, diff --git a/crates/skippy-metrics/src/lib.rs b/crates/skippy-metrics/src/lib.rs index 8d2f7b9ab6..22bff971fa 100644 --- a/crates/skippy-metrics/src/lib.rs +++ b/crates/skippy-metrics/src/lib.rs @@ -63,4 +63,5 @@ pub mod span { pub const MLX_BOUNDARY_HOST_COPY: &str = "stage.mlx_boundary_host_copy"; pub const MLX_BOUNDARY_ENCODE: &str = "stage.mlx_boundary_encode"; pub const MLX_BOUNDARY_DECODE: &str = "stage.mlx_boundary_decode"; + pub const MLX_BOUNDARY_TCP_ROUNDTRIP: &str = "stage.mlx_boundary_tcp_roundtrip"; } diff --git a/docs/design/MLX_STAGE_ENGINE_PLAN.md b/docs/design/MLX_STAGE_ENGINE_PLAN.md index 6314afc43e..f60803c3d9 100644 --- a/docs/design/MLX_STAGE_ENGINE_PLAN.md +++ b/docs/design/MLX_STAGE_ENGINE_PLAN.md @@ -163,6 +163,13 @@ and any conversion/transfer overlap. The 16K cell measured 32 MiB / 0.927 ms for F32 versus 16 MiB / 20.187 ms for F16. This makes wire dtype a hardware and link policy choice, not a model-format constant. +The next runner, `mlx-stage bench-tcp-boundary`, now wraps the production +engine-neutral Skippy TCP server around the synthetic sink. Its end-to-end +sample covers sender activation encode through framing, loopback TCP, +engine-transport reconstruction, and predicted reply. This tests whether the +codec-only policy survives actual Skippy framing and host copies; it still does +not stand in for QUIC or a remote link. + The derivation memory bound is the final packed routed bank, not one expert: six preallocated payload buffers total 718,405,632 bytes. Moving those buffers to a disk-backed random-write spool is the next step if preparation RSS must From 3854f27a418a6e55e93e5c2314fa57235e4a7222 Mon Sep 17 00:00:00 2001 From: Michael Neale <14976+michaelneale@users.noreply.github.com> Date: Fri, 17 Jul 2026 23:12:19 +1000 Subject: [PATCH 27/37] docs(mlx): record TCP boundary evidence --- crates/skippy-engine-mlx/STAGED_EXECUTION.md | 24 ++++++++++++++++++++ docs/design/MLX_STAGE_ENGINE_PLAN.md | 9 ++++++++ 2 files changed, 33 insertions(+) diff --git a/crates/skippy-engine-mlx/STAGED_EXECUTION.md b/crates/skippy-engine-mlx/STAGED_EXECUTION.md index bc9d9523ea..2bc13086b6 100644 --- a/crates/skippy-engine-mlx/STAGED_EXECUTION.md +++ b/crates/skippy-engine-mlx/STAGED_EXECUTION.md @@ -363,6 +363,30 @@ the same exact-F32 / bounded-F16 gate. Measured samples finish before telemetry starts. Each sample then becomes one `stage.mlx_boundary_tcp_roundtrip` span in a canonical metrics-server run. +#### TCP V1 evidence + +Commit `6350e3a9` ran release-mode F32/F16 pairs on the same M5 Max host as the +codec matrix, using three warmups and 20 sequential samples over one connection. + +| Width | Tokens | F32 payload | F32 p50 / p95 | F16 payload | F16 p50 / p95 | +| ---: | ---: | ---: | ---: | ---: | ---: | +| 2,688 | 32 | 0.328 MiB | 0.396 / 0.428 ms | 0.164 MiB | 0.627 / 0.677 ms | +| 4,096 | 512 | 8 MiB | 2.764 / 4.336 ms | 4 MiB | 8.086 / 8.231 ms | +| 8,192 | 512 | 16 MiB | 5.655 / 6.002 ms | 8 MiB | 13.601 / 13.814 ms | +| 16,384 | 512 | 32 MiB | 4.936 / 9.361 ms | 16 MiB | 25.398 / 25.622 ms | + +F32 wins on this high-bandwidth loopback path in every pair, consistent with +the codec-only prediction that F32 wins above the roughly 7.0 Gbit/s effective +payload break-even. The non-monotonic 16K F32 p50 and its wider p95 tail also +show why these numbers must not be converted into a remote-link bandwidth +claim: same-process allocation, kernel buffering, scheduling, and host copies +are part of this steady-state round trip. + +The eight canonical runs contain exactly 160 round-trip spans with zero drops +or export errors. F32 warmup reconstruction was exact; F16 maximum absolute +error was `0.00045216084`. This validates the production TCP framing direction, +but the next policy gate remains a controlled remote TCP/QUIC sweep. + ```bash just mlx-stage bench-tcp-boundary \ --width 16384 --tokens 512 --wire-dtype f16 \ diff --git a/docs/design/MLX_STAGE_ENGINE_PLAN.md b/docs/design/MLX_STAGE_ENGINE_PLAN.md index f60803c3d9..314a97b61c 100644 --- a/docs/design/MLX_STAGE_ENGINE_PLAN.md +++ b/docs/design/MLX_STAGE_ENGINE_PLAN.md @@ -170,6 +170,15 @@ engine-transport reconstruction, and predicted reply. This tests whether the codec-only policy survives actual Skippy framing and host copies; it still does not stand in for QUIC or a remote link. +**Loopback result.** Commit `6350e3a9` produced eight completed canonical TCP +runs / 160 spans with zero telemetry loss. At 512 tokens, F32/F16 round-trip +p50 was 2.764/8.086 ms at width 4K, 5.655/13.601 ms at 8K, and 4.936/25.398 ms +at 16K. F32 won every pair on this high-bandwidth loopback path, which is the +direction predicted by the codec-only ~7.0 Gbit/s break-even. The non-monotonic +F32 16K p50 and broad p95 make this evidence a framing/host-copy validation, +not a link-throughput estimate. A controlled remote TCP/QUIC sweep is still +required before automatic wire-dtype selection. + The derivation memory bound is the final packed routed bank, not one expert: six preallocated payload buffers total 718,405,632 bytes. Moving those buffers to a disk-backed random-write spool is the next step if preparation RSS must From 27bd588087b8186ccc902b000e79a90cc3b39d43 Mon Sep 17 00:00:00 2001 From: Michael Neale <14976+michaelneale@users.noreply.github.com> Date: Fri, 17 Jul 2026 23:26:33 +1000 Subject: [PATCH 28/37] bench(mlx): support external TCP boundary --- crates/skippy-engine-mlx/STAGED_EXECUTION.md | 46 +++ crates/skippy-engine-mlx/src/bin/mlx-stage.rs | 34 +- .../skippy-engine-mlx/src/boundary_bench.rs | 5 +- .../src/boundary_bench/tcp.rs | 333 ++++++++++++++---- crates/skippy-engine-mlx/src/lib.rs | 3 +- docs/design/MLX_STAGE_ENGINE_PLAN.md | 11 + 6 files changed, 350 insertions(+), 82 deletions(-) diff --git a/crates/skippy-engine-mlx/STAGED_EXECUTION.md b/crates/skippy-engine-mlx/STAGED_EXECUTION.md index 2bc13086b6..d1b33188c5 100644 --- a/crates/skippy-engine-mlx/STAGED_EXECUTION.md +++ b/crates/skippy-engine-mlx/STAGED_EXECUTION.md @@ -398,6 +398,52 @@ just mlx-stage bench-tcp-boundary \ --output /tmp/mlx-tcp-boundary-local.json ``` +### Separate-process and two-host TCP fence + +TCP v2 can move the validating sink into a separate process or host while +retaining the same production `engine_transport` framing and reconstruction +path. The sink is intentionally a benchmark tool: it is unauthenticated, +unencrypted TCP and must be bound only on a trusted private network or behind a +firewall. Its `width`, `tokens`, and `wire-dtype` must exactly match the sender. + +Start the sink in the foreground on the receiving host: + +```bash +just mlx-stage serve-tcp-boundary-sink \ + --bind 0.0.0.0:19090 \ + --width 16384 --tokens 512 --wire-dtype f16 +``` + +Then add `--connect :19090` to +`bench-tcp-boundary` on the sending host and use a fresh metrics run ID. The +runner allows 10 seconds for connect/READY and 30 seconds for each write and +reply. It uses one warmed persistent connection per invocation; the foreground +sink can accept later invocations without a restart. Run only one benchmark +client at a time per sink: its deliberately bounded validation cache retains +the most recent session, so concurrent clients can force revalidation into a +measured sample. + +Each run derives a distinct wire session from its metrics run ID. The sink +validates the first activation for that session and returns the observed +maximum absolute error as an explicit acknowledgement. It records the session +only after the exact-F32 / bounded-F16 gate succeeds, so a failed attempt cannot +poison a retry. The sender requires and independently gates that acknowledgement +before recording samples. Reports use the neutral `external_tcp` label because +an address supplied with `--connect` may still be localhost; the address itself +is excluded from the local/canonical telemetry payload. + +The current READY handshake does not carry sink build identity. For controlled +two-host evidence, copy the exact same release `mlx-stage` artifact to the sink +host and compare its SHA-256 on both hosts. Record that out-of-band checksum +alongside the client `code_revision`; do not infer sink provenance from the +client revision alone. + +The acknowledgement is a claim made by the sink, not cryptographic remote +attestation. `warmup_validation_ack_received` means the expected structured +reply arrived, and `warmup_sink_acknowledged_max_abs_diff` is the value reported +by that sink. The identical-artifact checksum procedure above is therefore part +of the controlled evidence, not an optional provenance detail. + ## Reproduce Build once: diff --git a/crates/skippy-engine-mlx/src/bin/mlx-stage.rs b/crates/skippy-engine-mlx/src/bin/mlx-stage.rs index eed1fefcc3..e137ed6ec3 100644 --- a/crates/skippy-engine-mlx/src/bin/mlx-stage.rs +++ b/crates/skippy-engine-mlx/src/bin/mlx-stage.rs @@ -15,8 +15,9 @@ mod real { use skippy_engine_mlx::{ MlxBoundaryBenchConfig, MlxComputeDtype, MlxDerivationControl, MlxDerivedStageCacheConfig, MlxDerivedStageConfig, MlxStageEngine, MlxStageEngineConfig, MlxTcpBoundaryBenchConfig, - MlxWeightQuantization, benchmark_mlx_boundary, benchmark_mlx_tcp_boundary, - derive_quantized_stage, derive_quantized_stage_cached, mlx_derived_stage_cache_root, + MlxTcpBoundarySinkConfig, MlxWeightQuantization, benchmark_mlx_boundary, + benchmark_mlx_tcp_boundary, derive_quantized_stage, derive_quantized_stage_cached, + mlx_derived_stage_cache_root, serve_mlx_tcp_boundary_sink, validate_nemotron_h_binary_wire_tokens, validate_nemotron_h_moe_stage, validate_nemotron_h_stage_engine, }; @@ -149,7 +150,7 @@ mod real { #[arg(long)] output: Option, }, - /// Measure one activation through production Skippy loopback TCP. + /// Measure one activation through local or remote production Skippy TCP. BenchTcpBoundary { #[arg(long)] width: usize, @@ -169,9 +170,23 @@ mod real { metrics_run_id: Option, #[arg(long)] metrics_report: PathBuf, + /// Connect to a separately running validating sink instead of loopback. + #[arg(long)] + connect: Option, #[arg(long)] output: Option, }, + /// Run a trusted-network-only validating sink for the TCP boundary benchmark. + ServeTcpBoundarySink { + #[arg(long)] + bind: SocketAddr, + #[arg(long)] + width: usize, + #[arg(long)] + tokens: usize, + #[arg(long, value_enum)] + wire_dtype: WireDtype, + }, /// Drive a stage chain and assert its greedy token sequence. Prove { #[arg(long)] @@ -387,6 +402,7 @@ mod real { metrics_otlp_grpc, metrics_run_id, metrics_report, + connect, output, } => { let metrics_run_id = match metrics_run_id { @@ -403,6 +419,7 @@ mod real { metrics_otlp_grpc, metrics_run_id, metrics_report_path: metrics_report, + connect_addr: connect, })?; let json = serde_json::to_vec_pretty(&report)?; if let Some(output) = output { @@ -417,6 +434,17 @@ mod real { println!("{}", String::from_utf8(json)?); Ok(()) } + Command::ServeTcpBoundarySink { + bind, + width, + tokens, + wire_dtype, + } => serve_mlx_tcp_boundary_sink(&MlxTcpBoundarySinkConfig { + bind_addr: bind, + width, + token_count: tokens, + wire_dtype: wire_dtype.into(), + }), Command::Prove { connect, tokens, diff --git a/crates/skippy-engine-mlx/src/boundary_bench.rs b/crates/skippy-engine-mlx/src/boundary_bench.rs index b35d67c56a..4ac3cceb8c 100644 --- a/crates/skippy-engine-mlx/src/boundary_bench.rs +++ b/crates/skippy-engine-mlx/src/boundary_bench.rs @@ -2,7 +2,10 @@ mod tcp; -pub use tcp::{MlxTcpBoundaryBenchConfig, MlxTcpBoundaryBenchReport, benchmark_mlx_tcp_boundary}; +pub use tcp::{ + MlxTcpBoundaryBenchConfig, MlxTcpBoundaryBenchReport, MlxTcpBoundarySinkConfig, + benchmark_mlx_tcp_boundary, serve_mlx_tcp_boundary_sink, +}; use std::{ collections::BTreeMap, diff --git a/crates/skippy-engine-mlx/src/boundary_bench/tcp.rs b/crates/skippy-engine-mlx/src/boundary_bench/tcp.rs index c9f48f359d..fec1c17ff1 100644 --- a/crates/skippy-engine-mlx/src/boundary_bench/tcp.rs +++ b/crates/skippy-engine-mlx/src/boundary_bench/tcp.rs @@ -16,6 +16,7 @@ use std::{ use anyhow::{Context, Result, anyhow, ensure}; use serde::Serialize; use serde_json::{Value, json}; +use sha2::{Digest, Sha256}; use skippy_engine::{ StageEngine, StageEngineInfo, StageExecutionKind, StageExecutionOutput, StageExecutionRequest, }; @@ -30,7 +31,7 @@ use skippy_protocol::{ }, }; use skippy_server::{ - engine_transport::{EngineStageServerOptions, serve_stage_engine_until}, + engine_transport::{EngineStageServerOptions, serve_stage_engine, serve_stage_engine_until}, telemetry::{Telemetry, TelemetryLevel, TelemetryStats, lifecycle_attrs}, }; @@ -40,8 +41,12 @@ use super::{ wait_for_telemetry, wire_dtype_label, }; -const BENCHMARK_SCHEMA: &str = "mlx-tcp-boundary-v1"; +const BENCHMARK_SCHEMA: &str = "mlx-tcp-boundary-v2"; const PREDICTED_SENTINEL: i32 = 424_243; +const CONNECT_DEADLINE: Duration = Duration::from_secs(10); +const CONNECT_ATTEMPT_TIMEOUT: Duration = Duration::from_secs(1); +const READY_TIMEOUT: Duration = Duration::from_secs(5); +const ROUNDTRIP_TIMEOUT: Duration = Duration::from_secs(30); /// One explicit, reproducible production-TCP boundary benchmark run. #[derive(Clone, Debug)] @@ -55,6 +60,16 @@ pub struct MlxTcpBoundaryBenchConfig { pub metrics_otlp_grpc: String, pub metrics_run_id: String, pub metrics_report_path: PathBuf, + pub connect_addr: Option, +} + +/// Foreground production sink for a separate-process or remote benchmark. +#[derive(Clone, Debug)] +pub struct MlxTcpBoundarySinkConfig { + pub bind_addr: SocketAddr, + pub width: usize, + pub token_count: usize, + pub wire_dtype: WireActivationDType, } /// Sender encode through production server decode and predicted reply. @@ -67,12 +82,14 @@ pub struct MlxTcpBoundaryBenchReport { pub width: usize, pub token_count: usize, pub wire_dtype: String, + pub transport: &'static str, pub warmup_iterations: usize, pub measured_iterations: usize, pub f32_boundary_bytes: usize, pub wire_activation_payload_bytes: usize, pub tcp_roundtrip: MlxBoundaryDurationSummary, - pub warmup_roundtrip_max_abs_diff: f32, + pub warmup_validation_ack_received: bool, + pub warmup_sink_acknowledged_max_abs_diff: f32, pub telemetry: TelemetryStats, pub canonical_span_count: u64, } @@ -88,34 +105,82 @@ pub fn benchmark_mlx_tcp_boundary( .map_err(|_| anyhow!("MLX TCP boundary benchmark thread panicked"))? } -fn benchmark_mlx_tcp_boundary_inner( - config: &MlxTcpBoundaryBenchConfig, -) -> Result { - validate_config(config)?; - let source = source_bytes(config)?; - let engine = Arc::new(TcpBoundarySink::new(config, Arc::clone(&source))?); - let (server, mut client) = TcpStageServer::spawn_ready( - engine.clone(), +/// Runs the validating final-stage sink in the foreground until interrupted. +pub fn serve_mlx_tcp_boundary_sink(config: &MlxTcpBoundarySinkConfig) -> Result<()> { + validate_protocol_shape(config.width, config.token_count, config.wire_dtype)?; + let source = source_bytes(config.width, config.token_count)?; + let engine = Arc::new(TcpBoundarySink::new( + config.width, + config.token_count, + config.wire_dtype, + source, + )?); + serve_stage_engine( + engine, EngineStageServerOptions { - bind_addr: "127.0.0.1:0".parse()?, + bind_addr: config.bind_addr, downstream_addr: None, wire_dtype: config.wire_dtype, }, - )?; + ) +} +fn benchmark_mlx_tcp_boundary_inner( + config: &MlxTcpBoundaryBenchConfig, +) -> Result { + validate_config(config)?; + let source = source_bytes(config.width, config.token_count)?; + let (server, engine, mut client) = match config.connect_addr { + Some(connect_addr) => (None, None, connect_ready(connect_addr)?), + None => { + let engine = Arc::new(TcpBoundarySink::new( + config.width, + config.token_count, + config.wire_dtype, + Arc::clone(&source), + )?); + let (server, client) = TcpStageServer::spawn_ready( + engine.clone(), + EngineStageServerOptions { + bind_addr: "127.0.0.1:0".parse()?, + downstream_addr: None, + wire_dtype: config.wire_dtype, + }, + )?; + (Some(server), Some(engine), client) + } + }; + + create_metrics_run(config)?; + let mut metrics_run = MetricsRunGuard::new(config); + let session_id = benchmark_session_id(&config.metrics_run_id); + let mut warmup_roundtrip_max_abs_diff = 0.0_f32; for iteration in 0..config.warmup_iterations { - run_roundtrip(config, &source, &mut client, u64::try_from(iteration)?)?; + let diff = run_roundtrip( + config, + &source, + &mut client, + session_id, + u64::try_from(iteration)?, + )?; + warmup_roundtrip_max_abs_diff = warmup_roundtrip_max_abs_diff.max(diff); + } + if let Some(engine) = engine.as_ref() { + let local_diff = engine + .validation_diff(session_id)? + .context("TCP boundary warmup acknowledgement has no matching local sink validation")?; + ensure!( + local_diff.to_bits() == warmup_roundtrip_max_abs_diff.to_bits(), + "TCP boundary warmup acknowledgement differs from local sink validation" + ); } - let warmup_roundtrip_max_abs_diff = engine - .validation_diff()? - .context("TCP boundary warmup did not reach the sink engine")?; - let samples = measure_roundtrips(config, &source, &mut client)?; + let samples = measure_roundtrips(config, &source, &mut client, session_id)?; drop(client); - server.stop()?; + if let Some(server) = server { + server.stop()?; + } - create_metrics_run(config)?; - let mut metrics_run = MetricsRunGuard::new(config); let telemetry_runtime = TcpTelemetryRuntime::new(config)?; let attrs = benchmark_attrs(config)?; for sample in &samples { @@ -149,6 +214,7 @@ fn benchmark_mlx_tcp_boundary_inner( width: config.width, token_count: config.token_count, wire_dtype: wire_dtype_label(config.wire_dtype)?.to_string(), + transport: transport_label(config), warmup_iterations: config.warmup_iterations, measured_iterations: config.measured_iterations, f32_boundary_bytes: activation_bytes(config, size_of::())?, @@ -161,18 +227,15 @@ fn benchmark_mlx_tcp_boundary_inner( }, )?, tcp_roundtrip: summarize(samples.iter().map(|sample| sample.elapsed))?, - warmup_roundtrip_max_abs_diff, + warmup_validation_ack_received: true, + warmup_sink_acknowledged_max_abs_diff: warmup_roundtrip_max_abs_diff, telemetry, canonical_span_count, }) } fn validate_config(config: &MlxTcpBoundaryBenchConfig) -> Result<()> { - ensure!(config.width > 0, "TCP boundary width must be non-zero"); - ensure!( - config.token_count > 0, - "TCP boundary token count must be non-zero" - ); + validate_protocol_shape(config.width, config.token_count, config.wire_dtype)?; ensure!( config.warmup_iterations > 0, "TCP boundary benchmark needs a correctness warmup" @@ -190,20 +253,31 @@ fn validate_config(config: &MlxTcpBoundaryBenchConfig) -> Result<()> { "metrics-server OTLP endpoint is required" ); validate_metrics_run_id(&config.metrics_run_id)?; - wire_dtype_label(config.wire_dtype)?; - let token_count = i32::try_from(config.token_count).context("token count exceeds i32")?; - let width = i32::try_from(config.width).context("activation width exceeds i32")?; - u32::try_from(config.width).context("activation width exceeds u32")?; + Ok(()) +} + +fn validate_protocol_shape( + width: usize, + token_count: usize, + wire_dtype: WireActivationDType, +) -> Result<()> { + ensure!(width > 0, "TCP boundary width must be non-zero"); + ensure!(token_count > 0, "TCP boundary token count must be non-zero"); + wire_dtype_label(wire_dtype)?; + let token_count_i32 = i32::try_from(token_count).context("token count exceeds i32")?; + let width_i32 = i32::try_from(width).context("activation width exceeds i32")?; + u32::try_from(width).context("activation width exceeds u32")?; ensure!( - config.token_count <= MAX_STAGE_SIDEBAND_VALUES, + token_count <= MAX_STAGE_SIDEBAND_VALUES, "token sideband count exceeds protocol maximum" ); - let wire_bytes = activation_wire_bytes(config.wire_dtype, token_count, width)?; + let wire_bytes = activation_wire_bytes(wire_dtype, token_count_i32, width_i32)?; ensure!( wire_bytes <= MAX_STAGE_ACTIVATION_BYTES, "wire activation exceeds protocol maximum" ); - let decoded_bytes = activation_wire_bytes(WireActivationDType::F32, token_count, width)?; + let decoded_bytes = + activation_wire_bytes(WireActivationDType::F32, token_count_i32, width_i32)?; ensure!( decoded_bytes <= MAX_STAGE_DECODED_ACTIVATION_BYTES, "decoded activation exceeds protocol maximum" @@ -211,10 +285,9 @@ fn validate_config(config: &MlxTcpBoundaryBenchConfig) -> Result<()> { Ok(()) } -fn source_bytes(config: &MlxTcpBoundaryBenchConfig) -> Result>> { - let elements = config - .width - .checked_mul(config.token_count) +fn source_bytes(width: usize, token_count: usize) -> Result>> { + let elements = width + .checked_mul(token_count) .context("TCP boundary element count overflow")?; Ok(Arc::new( (0..elements) @@ -230,11 +303,13 @@ fn measure_roundtrips( config: &MlxTcpBoundaryBenchConfig, source: &[u8], client: &mut TcpStream, + session_id: u64, ) -> Result> { (0..config.measured_iterations) .map(|iteration| { let request_id = u64::try_from(config.warmup_iterations + iteration)?; - let ((), timing) = time_phase(|| run_roundtrip(config, source, client, request_id))?; + let (_, timing) = + time_phase(|| run_roundtrip(config, source, client, session_id, request_id))?; Ok(timing) }) .collect() @@ -244,8 +319,9 @@ fn run_roundtrip( config: &MlxTcpBoundaryBenchConfig, source: &[u8], client: &mut TcpStream, + session_id: u64, request_id: u64, -) -> Result<()> { +) -> Result { let token_count = i32::try_from(config.token_count)?; let width = i32::try_from(config.width)?; let kind = WireMessageKind::PrefillFinalEmbd; @@ -258,7 +334,7 @@ fn run_roundtrip( token_count, state, request_id, - session_id: 1, + session_id, sampling: None, chat_sampling_metadata: None, tokens: vec![0; config.token_count], @@ -273,10 +349,14 @@ fn run_roundtrip( ensure!( reply.kind == WireReplyKind::PredictedToken && reply.predicted == PREDICTED_SENTINEL - && reply.predicted_tokens == [PREDICTED_SENTINEL], - "TCP boundary sink returned the wrong prediction" + && reply.predicted_tokens.len() == 2 + && reply.predicted_tokens[0] == PREDICTED_SENTINEL, + "TCP boundary sink returned no validation acknowledgement" ); - Ok(()) + let diff = validation_diff_from_ack(reply.predicted_tokens[1]); + validate_roundtrip(config.wire_dtype, diff) + .context("TCP boundary sink acknowledged an invalid activation")?; + Ok(diff) } struct TcpBoundarySink { @@ -284,12 +364,16 @@ struct TcpBoundarySink { token_count: usize, wire_dtype: WireActivationDType, source: Arc>, - validated: AtomicBool, - validation_diff: Mutex>, + validation: Mutex>, } impl TcpBoundarySink { - fn new(config: &MlxTcpBoundaryBenchConfig, source: Arc>) -> Result { + fn new( + width: usize, + token_count: usize, + wire_dtype: WireActivationDType, + source: Arc>, + ) -> Result { let info = StageEngineInfo { engine: "mlx-tcp-boundary-sink".to_string(), model_id: "synthetic/mlx-tcp-boundary".to_string(), @@ -297,24 +381,39 @@ impl TcpBoundarySink { layer_start: 1, layer_end: 2, total_layers: 2, - activation_width: u32::try_from(config.width)?, + activation_width: u32::try_from(width)?, }; info.validate()?; Ok(Self { info, - token_count: config.token_count, - wire_dtype: config.wire_dtype, + token_count, + wire_dtype, source, - validated: AtomicBool::new(false), - validation_diff: Mutex::new(None), + validation: Mutex::new(None), }) } - fn validation_diff(&self) -> Result> { - Ok(*self - .validation_diff + fn validation_diff(&self, session_id: u64) -> Result> { + Ok(self + .validation .lock() - .map_err(|_| anyhow!("TCP boundary validation lock poisoned"))?) + .map_err(|_| anyhow!("TCP boundary validation lock poisoned"))? + .filter(|(validated_session, _)| *validated_session == session_id) + .map(|(_, diff)| diff)) + } + + fn validate_session(&self, session_id: u64, activation: &[u8]) -> Result { + if let Some(diff) = self.validation_diff(session_id)? { + return Ok(diff); + } + let diff = max_abs_diff(&self.source, activation)?; + validate_roundtrip(self.wire_dtype, diff)?; + *self + .validation + .lock() + .map_err(|_| anyhow!("TCP boundary validation lock poisoned"))? = + Some((session_id, diff)); + Ok(diff) } } @@ -336,21 +435,24 @@ impl StageEngine for TcpBoundarySink { && activation.width == self.info.activation_width as usize, "TCP boundary sink received the wrong activation shape" ); - if !self.validated.swap(true, Ordering::SeqCst) { - let diff = max_abs_diff(&self.source, &activation.f32_le_bytes)?; - validate_roundtrip(self.wire_dtype, diff)?; - *self - .validation_diff - .lock() - .map_err(|_| anyhow!("TCP boundary validation lock poisoned"))? = Some(diff); - } + let diff = self.validate_session(request.session_id, &activation.f32_le_bytes)?; Ok(StageExecutionOutput { activation: None, - predicted_tokens: vec![PREDICTED_SENTINEL], + predicted_tokens: vec![PREDICTED_SENTINEL, validation_diff_ack(diff)], }) } - fn reset_session(&self, _session_id: u64) -> Result<()> { + fn reset_session(&self, session_id: u64) -> Result<()> { + let mut validation = self + .validation + .lock() + .map_err(|_| anyhow!("TCP boundary validation lock poisoned"))?; + if validation + .as_ref() + .is_some_and(|(validated_session, _)| *validated_session == session_id) + { + *validation = None; + } Ok(()) } } @@ -425,16 +527,22 @@ fn reserve_loopback_addr() -> Result { } fn connect_ready(addr: SocketAddr) -> Result { - let deadline = Instant::now() + Duration::from_secs(5); + let deadline = Instant::now() + CONNECT_DEADLINE; loop { - let error = match TcpStream::connect(addr) { + let remaining = deadline.saturating_duration_since(Instant::now()); + ensure!( + !remaining.is_zero(), + "connect TCP boundary stage at {addr} timed out" + ); + let error = match TcpStream::connect_timeout(&addr, remaining.min(CONNECT_ATTEMPT_TIMEOUT)) + { Ok(mut stream) => { stream.set_nodelay(true).ok(); - stream.set_read_timeout(Some(Duration::from_millis(250)))?; - stream.set_write_timeout(Some(Duration::from_secs(5)))?; + stream.set_read_timeout(Some(remaining.min(READY_TIMEOUT)))?; + stream.set_write_timeout(Some(ROUNDTRIP_TIMEOUT))?; match recv_ready(&mut stream) { Ok(()) => { - stream.set_read_timeout(Some(Duration::from_secs(10)))?; + stream.set_read_timeout(Some(ROUNDTRIP_TIMEOUT))?; return Ok(stream); } Err(error) => anyhow!(error).context("receive TCP boundary ready handshake"), @@ -449,6 +557,22 @@ fn connect_ready(addr: SocketAddr) -> Result { } } +fn benchmark_session_id(metrics_run_id: &str) -> u64 { + let digest = Sha256::digest(metrics_run_id.as_bytes()); + let bytes: [u8; size_of::()] = digest[..size_of::()] + .try_into() + .expect("SHA-256 prefix has the requested length"); + u64::from_le_bytes(bytes) +} + +fn validation_diff_ack(diff: f32) -> i32 { + i32::from_le_bytes(diff.to_bits().to_le_bytes()) +} + +fn validation_diff_from_ack(ack: i32) -> f32 { + f32::from_bits(u32::from_le_bytes(ack.to_le_bytes())) +} + struct TcpTelemetryRuntime { telemetry: Telemetry, _runtime: tokio::runtime::Runtime, @@ -534,6 +658,7 @@ fn metrics_run_body(config: &MlxTcpBoundaryBenchConfig) -> Result { "width": config.width, "token_count": config.token_count, "wire_dtype": wire_dtype_label(config.wire_dtype)?, + "transport": transport_label(config), "warmup_iterations": config.warmup_iterations, "measured_iterations": config.measured_iterations, "stages": [{ @@ -623,6 +748,14 @@ fn activation_bytes(config: &MlxTcpBoundaryBenchConfig, element_bytes: usize) -> .context("TCP boundary byte count overflow") } +fn transport_label(config: &MlxTcpBoundaryBenchConfig) -> &'static str { + if config.connect_addr.is_some() { + "external_tcp" + } else { + "loopback" + } +} + #[cfg(test)] mod tests { use super::*; @@ -638,12 +771,14 @@ mod tests { metrics_otlp_grpc: "https://otlp.invalid/private".to_string(), metrics_run_id: "tcp-test-run".to_string(), metrics_report_path: PathBuf::from("/private/report.json"), + connect_addr: None, } } #[test] fn telemetry_contains_no_transport_targets_or_local_paths() { - let config = config(); + let mut config = config(); + config.connect_addr = Some("203.0.113.10:1234".parse().unwrap()); let exported = json!({ "stage_config": telemetry_stage_config(&config).unwrap(), "span_attributes": benchmark_attrs(&config).unwrap(), @@ -653,12 +788,13 @@ mod tests { assert!(!serialized.contains("collector.invalid")); assert!(!serialized.contains("otlp.invalid")); assert!(!serialized.contains("/private/")); + assert!(!serialized.contains("203.0.113.10")); } #[test] fn source_values_are_finite_and_in_the_codec_gate_range() { let config = config(); - let source = source_bytes(&config).unwrap(); + let source = source_bytes(config.width, config.token_count).unwrap(); assert_eq!(source.len(), 4096 * 32 * size_of::()); assert!( source @@ -677,8 +813,16 @@ mod tests { config.width = 8; config.token_count = 2; config.wire_dtype = wire_dtype; - let source = source_bytes(&config).unwrap(); - let engine = Arc::new(TcpBoundarySink::new(&config, Arc::clone(&source)).unwrap()); + let source = source_bytes(config.width, config.token_count).unwrap(); + let engine = Arc::new( + TcpBoundarySink::new( + config.width, + config.token_count, + config.wire_dtype, + Arc::clone(&source), + ) + .unwrap(), + ); let (server, mut client) = TcpStageServer::spawn_ready( engine.clone(), EngineStageServerOptions { @@ -689,13 +833,48 @@ mod tests { ) .unwrap(); - run_roundtrip(&config, &source, &mut client, 1).unwrap(); - assert!(engine.validation_diff().unwrap().unwrap() <= expected_diff); + let session_id = benchmark_session_id(&config.metrics_run_id); + let acknowledged_diff = + run_roundtrip(&config, &source, &mut client, session_id, 1).unwrap(); + assert!(acknowledged_diff <= expected_diff); + assert_eq!( + engine.validation_diff(session_id).unwrap(), + Some(acknowledged_diff) + ); drop(client); server.stop().unwrap(); } } + #[test] + fn sink_validates_each_session_and_failed_validation_does_not_poison_it() { + let config = config(); + let source = source_bytes(config.width, config.token_count).unwrap(); + let engine = TcpBoundarySink::new( + config.width, + config.token_count, + config.wire_dtype, + Arc::clone(&source), + ) + .unwrap(); + assert!(engine.validate_session(11, &source).unwrap() <= 0.001); + + let mut invalid = source.as_ref().clone(); + invalid[..size_of::()].copy_from_slice(&10.0_f32.to_le_bytes()); + assert!(engine.validate_session(12, &invalid).is_err()); + assert_eq!(engine.validation_diff(12).unwrap(), None); + + assert!(engine.validate_session(12, &source).unwrap() <= 0.001); + assert!(engine.validation_diff(12).unwrap().is_some()); + } + + #[test] + fn validation_ack_and_session_identity_round_trip() { + let diff = 0.000_452_160_84_f32; + assert_eq!(validation_diff_from_ack(validation_diff_ack(diff)), diff); + assert_ne!(benchmark_session_id("run-a"), benchmark_session_id("run-b")); + } + #[test] fn oversized_protocol_shapes_fail_before_source_allocation() { let mut config = config(); diff --git a/crates/skippy-engine-mlx/src/lib.rs b/crates/skippy-engine-mlx/src/lib.rs index 1e93b3f42a..1945997425 100644 --- a/crates/skippy-engine-mlx/src/lib.rs +++ b/crates/skippy-engine-mlx/src/lib.rs @@ -23,7 +23,8 @@ pub use backend::MlxBackend; #[cfg(all(feature = "mlx", target_os = "macos"))] pub use boundary_bench::{ MlxBoundaryBenchConfig, MlxBoundaryBenchReport, MlxTcpBoundaryBenchConfig, - MlxTcpBoundaryBenchReport, benchmark_mlx_boundary, benchmark_mlx_tcp_boundary, + MlxTcpBoundaryBenchReport, MlxTcpBoundarySinkConfig, benchmark_mlx_boundary, + benchmark_mlx_tcp_boundary, serve_mlx_tcp_boundary_sink, }; #[cfg(all(feature = "mlx", target_os = "macos"))] pub use derived::{ diff --git a/docs/design/MLX_STAGE_ENGINE_PLAN.md b/docs/design/MLX_STAGE_ENGINE_PLAN.md index 314a97b61c..9734a2c525 100644 --- a/docs/design/MLX_STAGE_ENGINE_PLAN.md +++ b/docs/design/MLX_STAGE_ENGINE_PLAN.md @@ -179,6 +179,17 @@ F32 16K p50 and broad p95 make this evidence a framing/host-copy validation, not a link-throughput estimate. A controlled remote TCP/QUIC sweep is still required before automatic wire-dtype selection. +**Update — the two-host TCP runner is now fail-closed.** TCP boundary schema v2 +can connect to a separately running production `engine_transport` sink. A +unique per-run wire session forces a fresh first-activation validation even +when the foreground sink is reused, and the sink returns the measured +exact-F32 / bounded-F16 error as an acknowledgement required by the sender. +Connect/READY and round-trip IO are bounded. Canonical telemetry records only +the neutral `external_tcp` mode, never the target address. The benchmark sink +is plaintext and unauthenticated, so it is restricted to trusted private +networks. Until sink revision is added to READY, controlled evidence must copy +and independently hash the identical release artifact on both hosts. + The derivation memory bound is the final packed routed bank, not one expert: six preallocated payload buffers total 718,405,632 bytes. Moving those buffers to a disk-backed random-write spool is the next step if preparation RSS must From a6802bf74ac512e5b5b42d6b4fb0903d8c45ad65 Mon Sep 17 00:00:00 2001 From: Michael Neale <14976+michaelneale@users.noreply.github.com> Date: Fri, 17 Jul 2026 23:42:32 +1000 Subject: [PATCH 29/37] docs(mlx): record two-host boundary evidence --- crates/skippy-engine-mlx/STAGED_EXECUTION.md | 43 ++++++++++++++++++++ docs/design/MLX_STAGE_ENGINE_PLAN.md | 13 ++++++ 2 files changed, 56 insertions(+) diff --git a/crates/skippy-engine-mlx/STAGED_EXECUTION.md b/crates/skippy-engine-mlx/STAGED_EXECUTION.md index d1b33188c5..30654a27e1 100644 --- a/crates/skippy-engine-mlx/STAGED_EXECUTION.md +++ b/crates/skippy-engine-mlx/STAGED_EXECUTION.md @@ -444,6 +444,49 @@ reply arrived, and `warmup_sink_acknowledged_max_abs_diff` is the value reported by that sink. The identical-artifact checksum procedure above is therefore part of the controlled evidence, not an optional provenance detail. +#### SSH-forwarded two-host V2 evidence + +Commit `27bd5880` was built once in release mode, signed once, and copied +unchanged to an M4 Max receiver from an M5 Max sender. Both hosts reported the +same executable SHA-256: +`ba5d1ea6f2613d0171d36eeaf9dfd86904d3ae19f6d335b3185bbb4ebe5a2222`. + +The receiver's application firewall allowed local sink traffic but suppressed +data after a direct-LAN TCP handshake for the ad-hoc research binary. The +controlled sweep therefore used one persistent SSH local forward to the +receiver's loopback-bound sink. Every timed sample still covers production +sender encoding/framing, cross-host transfer, production receiver +framing/reconstruction, sink acknowledgement, and the reply, but it also +includes SSH tunnelling and encryption. These are not raw-LAN or QUIC numbers. + +Each cell used three warmups and 20 sequential measured samples. The 16K pair +was repeated in reverse dtype order on a fresh tunnel after the initial F16 run +showed a severe transient. + +| Width | Tokens | F32 payload | F32 p50 / p95 | F16 payload | F16 p50 / p95 | +| ---: | ---: | ---: | ---: | ---: | ---: | +| 2,688 | 32 | 0.328 MiB | 16.447 / 21.106 ms | 0.164 MiB | 16.532 / 22.415 ms | +| 4,096 | 512 | 8 MiB | 150.690 / 173.837 ms | 4 MiB | 107.376 / 115.066 ms | +| 8,192 | 512 | 16 MiB | 288.479 / 874.937 ms | 8 MiB | 200.168 / 677.533 ms | +| 16,384 A | 512 | 32 MiB | 603.895 / 2346.909 ms | 16 MiB | 1441.627 / 6323.343 ms | +| 16,384 B | 512 | 32 MiB | 572.663 / 625.476 ms | 16 MiB | 345.498 / 923.549 ms | + +The actual 2,688×32 Nemotron boundary was effectively tied, consistent with +fixed SSH/tunnel overhead being large relative to its small payload. F16 +reduced p50 by about 29% at 4K and 31% at 8K. In the fresh-tunnel 16K repeat it +reduced p50 by about 40%, but the opposite result and multi-second tails in the +first 16K pair show that this setup cannot select a production wire dtype. It +is a functional two-host proof, with results consistent with payload reduction +mattering on this constrained SSH-forwarded path. Raw LAN/QUIC, repeated +interleaved trials, and pipeline overlap remain required for automatic policy. + +The 10 completed canonical runs contain exactly 200 +`stage.mlx_boundary_tcp_roundtrip` spans. All use schema +`mlx-tcp-boundary-v2`, revision `27bd588087b8186ccc902b000e79a90cc3b39d43`, +and transport `external_tcp`, with zero dropped spans or export errors. No span +falls outside its run lifecycle. F32 acknowledgements were exact and every F16 +acknowledgement reported maximum absolute error `0.00045216084`. + ## Reproduce Build once: diff --git a/docs/design/MLX_STAGE_ENGINE_PLAN.md b/docs/design/MLX_STAGE_ENGINE_PLAN.md index 9734a2c525..abf5c85ea6 100644 --- a/docs/design/MLX_STAGE_ENGINE_PLAN.md +++ b/docs/design/MLX_STAGE_ENGINE_PLAN.md @@ -190,6 +190,19 @@ is plaintext and unauthenticated, so it is restricted to trusted private networks. Until sink revision is added to READY, controlled evidence must copy and independently hash the identical release artifact on both hosts. +**Two-host result.** Commit `27bd5880` completed 10 canonical external-TCP +runs / 200 spans between an M5 Max sender and M4 Max sink with identical binary +SHA-256, explicit validation acknowledgements, and zero telemetry loss. The +receiver firewall required an SSH local forward, so the measurements include +SSH tunnelling and are not raw-LAN or QUIC evidence. F16 and F32 were effectively +tied at the real 2,688×32 Nemotron boundary; F16 reduced p50 by about 29% at +4K×512 and 31% at 8K×512. A fresh-tunnel 16K repeat favored F16 by about 40%, +but the first 16K F16 run suffered multi-second tunnel tails and lost badly. +This closes the cross-host production-framing functional proof and confirms +that payload reduction can matter on a constrained path. It does not close the +wire-dtype policy gate: direct LAN/QUIC, repeated interleaved cells, and +pipeline-overlap measurements remain. + The derivation memory bound is the final packed routed bank, not one expert: six preallocated payload buffers total 718,405,632 bytes. Moving those buffers to a disk-backed random-write spool is the next step if preparation RSS must From 58de032a6152bcb88a3f7f70379b692ca91e8b19 Mon Sep 17 00:00:00 2001 From: Michael Neale <14976+michaelneale@users.noreply.github.com> Date: Fri, 17 Jul 2026 23:58:39 +1000 Subject: [PATCH 30/37] feat(mlx): prove two-host stage execution --- Justfile | 5 +- crates/skippy-engine-mlx/STAGED_EXECUTION.md | 47 +++++++++++++++++-- crates/skippy-engine-mlx/src/bin/mlx-stage.rs | 2 +- docs/design/MLX_STAGE_ENGINE_PLAN.md | 12 +++++ 4 files changed, 59 insertions(+), 7 deletions(-) diff --git a/Justfile b/Justfile index 87e25d6bab..03f8d743ab 100644 --- a/Justfile +++ b/Justfile @@ -251,9 +251,12 @@ mlx-safetensors-stage-plan *ARGS: 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 production-shaped MLX stage server/client over Skippy's binary wire. +# 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 # Run `mlx-stage serve ...` or `mlx-stage prove ...` after `just mlx-stage-build`. mlx-stage *ARGS: diff --git a/crates/skippy-engine-mlx/STAGED_EXECUTION.md b/crates/skippy-engine-mlx/STAGED_EXECUTION.md index 30654a27e1..4a9a507d83 100644 --- a/crates/skippy-engine-mlx/STAGED_EXECUTION.md +++ b/crates/skippy-engine-mlx/STAGED_EXECUTION.md @@ -495,6 +495,13 @@ Build once: just mlx-stage-build ``` +This writes `target/release/mlx-stage` and the required sibling +`target/release/mlx.metallib`. Copy both files together when moving the CLI to +another Apple-Silicon host. For this pinned build the Metal library is about +157 MiB; its generated size can change with MLX or the Apple toolchain. It is a +runtime resource shared by every model on that host, not part of a stage +artifact. + Derive both quantized stage directories directly from immutable source ranges: ```bash @@ -525,21 +532,23 @@ Start the final stage: ```bash just mlx-stage serve \ - --model /tmp/mlx-split-smol/stage1 \ + --model /tmp/mlx-derived-smol-stage1 \ --model-id HuggingFaceTB/SmolLM2-135M-Instruct \ --stage-index 1 --layer-start 15 --layer-end 30 \ --bind 127.0.0.1:19091 --wire-dtype f16 --compute-dtype bf16 ``` -Add `--weight-quantization affine4` to both stage commands to reproduce the -JIT-quantized proof, then pass -`--expected 260,2240,314,253,1379,282,25801,28` to `mlx-stage prove`. +The directly derived directories are already affine-4. For a separate dense +`/tmp/mlx-split-smol` proof, use those paths and add +`--weight-quantization affine4` to both stage commands. The `prove` default is +the established affine-4 reference +`260,2240,314,253,1379,282,25801,28`. Start the first stage in another terminal: ```bash just mlx-stage serve \ - --model /tmp/mlx-split-smol/stage0 \ + --model /tmp/mlx-derived-smol-stage0 \ --model-id HuggingFaceTB/SmolLM2-135M-Instruct \ --stage-index 0 --layer-start 0 --layer-end 15 \ --bind 127.0.0.1:19090 --downstream 127.0.0.1:19091 \ @@ -552,6 +561,34 @@ Drive the chain: just mlx-stage prove --connect 127.0.0.1:19090 --wire-dtype f16 ``` +### Real two-host split proof + +The same M4 Max used for the boundary sweep independently derived only the +final `15..30` range from the immutable SmolLM2 checkpoint. In this recorded +pinned run it made 137 tensor payload requests for 162,827,136 source tensor +bytes and wrote a 45,889,726-byte affine-4 stage in three shards. All three +shard hashes exactly matched the earlier M5 Max derivation, demonstrating +deterministic stage-local materialization across the two machines for this +checkpoint, safemlx revision, and quantization recipe. + +The first real `MlxStageEngine` then loaded the pre-derived `0..15` stage on the +M5 Max and forwarded its width-576 residuals through an SSH local forward to +the M4 Max's real `15..30` final stage. Both used BF16 compute and the same +affine-4 artifacts. Two successive F16-wire runs, including session reset and +reuse, and one F32-wire run all produced the reference sequence: + +```text +[260, 2240, 314, 253, 1379, 282, 25801, 28] +``` + +This closes the intentionally unnecessary two-host small-Llama execution +proof: each layer server can hold only its own directly derived SafeTensors +stage, and the existing production binary stage protocol composes the two real +MLX engines. The transport was SSH-forwarded because of the receiver firewall, +and startup was manual through `mlx-stage`; this is not yet mesh coordinator, +placement, capability advertisement, OpenAI stage-0 orchestration, or raw +LAN/QUIC evidence. + ## Deliberate limitations of this checkpoint - `MlxStageEngine` supports dense Llama ranges and exactly one internal, diff --git a/crates/skippy-engine-mlx/src/bin/mlx-stage.rs b/crates/skippy-engine-mlx/src/bin/mlx-stage.rs index e137ed6ec3..ed7d1bfd8b 100644 --- a/crates/skippy-engine-mlx/src/bin/mlx-stage.rs +++ b/crates/skippy-engine-mlx/src/bin/mlx-stage.rs @@ -193,7 +193,7 @@ mod real { connect: SocketAddr, #[arg(long, default_value = "1,1531,314,260,3575,28")] tokens: String, - #[arg(long, default_value = "284,260,2240,314,1343,327,624,8685")] + #[arg(long, default_value = "260,2240,314,253,1379,282,25801,28")] expected: String, #[arg(long, value_enum, default_value_t = WireDtype::F16)] wire_dtype: WireDtype, diff --git a/docs/design/MLX_STAGE_ENGINE_PLAN.md b/docs/design/MLX_STAGE_ENGINE_PLAN.md index abf5c85ea6..3764ab7b41 100644 --- a/docs/design/MLX_STAGE_ENGINE_PLAN.md +++ b/docs/design/MLX_STAGE_ENGINE_PLAN.md @@ -203,6 +203,18 @@ that payload reduction can matter on a constrained path. It does not close the wire-dtype policy gate: direct LAN/QUIC, repeated interleaved cells, and pipeline-overlap measurements remain. +**Update — two real stages now execute across two hosts.** The M4 Max +independently fetched only SmolLM2 layers `15..30` (137 range requests, +162,827,136 tensor bytes) and produced a 45,889,726-byte affine-4 artifact whose +three shard hashes matched the M5 Max derivation. A real M5 Max `0..15` +`MlxStageEngine` chained through an SSH local forward to that real M4 Max final +stage. Two successive F16-wire runs and one F32-wire run all reproduced +`[260, 2240, 314, 253, 1379, 282, 25801, 28]`. A copied MLX executable also +requires the build-generated `mlx.metallib` beside it (about 157 MiB in this +pinned build); `just mlx-stage-build` now exports that sibling resource. This +closes manual two-host small-Llama execution and per-host range derivation, not +mesh-managed placement, OpenAI orchestration, or raw LAN/QUIC transport. + The derivation memory bound is the final packed routed bank, not one expert: six preallocated payload buffers total 718,405,632 bytes. Moving those buffers to a disk-backed random-write spool is the next step if preparation RSS must From 2690a1aeb72811dc4dcbf8f970d16dfb0c5d59e2 Mon Sep 17 00:00:00 2001 From: Michael Neale <14976+michaelneale@users.noreply.github.com> Date: Sat, 18 Jul 2026 02:46:02 +1000 Subject: [PATCH 31/37] feat(mlx): integrate safetensors mesh serving --- Cargo.lock | 1 + Justfile | 10 + crates/mesh-llm-host-runtime/src/api/tests.rs | 2 + .../src/inference/mlx.rs | 56 +- .../src/inference/skippy/stage/mod.rs | 1 + crates/mesh-llm-host-runtime/src/lib.rs | 26 +- .../mesh-llm-host-runtime/src/mesh/gossip.rs | 10 + crates/mesh-llm-host-runtime/src/mesh/mod.rs | 4 + .../mesh-llm-host-runtime/src/mesh/tests.rs | 11 + .../src/network/openai/transport.rs | 1 + .../src/protocol/convert.rs | 13 + .../mesh-llm-host-runtime/src/protocol/mod.rs | 13 + .../src/runtime/local.rs | 524 ++++++++++++++-- .../mesh-llm-host-runtime/src/runtime/mod.rs | 88 ++- .../src/runtime/split_planning.rs | 26 +- .../src/runtime_data/mod.rs | 2 + .../mesh-llm-host-runtime/src/system/mod.rs | 1 - .../src/system/native_runtime.rs | 9 +- .../model-hf/src/safetensors_stage/layout.rs | 228 ++++++- .../src/safetensors_stage/materialize.rs | 171 +++++- crates/model-hf/src/safetensors_stage/mod.rs | 7 +- .../model-hf/src/safetensors_stage/types.rs | 27 + crates/skippy-engine-mlx/Cargo.toml | 4 + .../SERVE_INTEGRATION_STATUS.md | 207 +++---- crates/skippy-engine-mlx/STAGED_EXECUTION.md | 14 +- crates/skippy-engine-mlx/src/backend.rs | 39 +- crates/skippy-engine-mlx/src/bin/mlx-serve.rs | 41 +- crates/skippy-engine-mlx/src/distributed.rs | 559 ++++++++++++++++++ crates/skippy-engine-mlx/src/engine.rs | 164 ++++- crates/skippy-engine-mlx/src/lib.rs | 8 +- crates/skippy-protocol/src/lib.rs | 2 + crates/skippy-server/src/engine_transport.rs | 146 ++++- docs/design/MLX_STAGE_ENGINE_PLAN.md | 15 + scripts/build-mac.sh | 29 +- 34 files changed, 2204 insertions(+), 255 deletions(-) create mode 100644 crates/skippy-engine-mlx/src/distributed.rs diff --git a/Cargo.lock b/Cargo.lock index 0742d5970e..0341f13e28 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -7784,6 +7784,7 @@ dependencies = [ "reqwest 0.12.28", "safemlx", "safemlx-lm", + "safemlx-lm-utils", "safetensors", "serde", "serde_json", diff --git a/Justfile b/Justfile index 03f8d743ab..ba1073bf7d 100644 --- a/Justfile +++ b/Justfile @@ -258,6 +258,16 @@ mlx-stage-build: 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 }} diff --git a/crates/mesh-llm-host-runtime/src/api/tests.rs b/crates/mesh-llm-host-runtime/src/api/tests.rs index b5f2acf1e9..811df20963 100644 --- a/crates/mesh-llm-host-runtime/src/api/tests.rs +++ b/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![], diff --git a/crates/mesh-llm-host-runtime/src/inference/mlx.rs b/crates/mesh-llm-host-runtime/src/inference/mlx.rs index 2afaf8ee8d..fe9d3d4758 100644 --- a/crates/mesh-llm-host-runtime/src/inference/mlx.rs +++ b/crates/mesh-llm-host-runtime/src/inference/mlx.rs @@ -15,7 +15,13 @@ use std::path::PathBuf; use std::sync::Arc; use anyhow::{Context, Result}; -use skippy_engine_mlx::{MlxBackend, MlxEngine, MlxEngineConfig}; +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 { @@ -26,11 +32,13 @@ 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 { + let weight_quantization = automatic_weight_quantization(&model_dir)?; let config = MlxEngineConfig { model_dir, model_id, - default_max_tokens: context_length.max(1) as usize, + 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, }; let engine = MlxEngine::spawn(config)?; Ok(Self { @@ -38,33 +46,57 @@ impl MlxModelHandle { }) } + /// 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 { + 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) fn start_http(&self, port: u16) -> MlxHttpHandle { + pub(crate) async fn start_http( + &self, + port: u16, + death_tx: tokio::sync::oneshot::Sender<()>, + ) -> Result { 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 listener = match tokio::net::TcpListener::bind(addr).await { - Ok(listener) => listener, - Err(error) => { - tracing::error!(%addr, %error, "MLX openai frontend failed to bind"); - return; - } - }; 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(()); }); - MlxHttpHandle { + Ok(MlxHttpHandle { port, shutdown_tx: Some(shutdown_tx), server: Some(server), - } + }) } } diff --git a/crates/mesh-llm-host-runtime/src/inference/skippy/stage/mod.rs b/crates/mesh-llm-host-runtime/src/inference/skippy/stage/mod.rs index 3b30704bd6..ac717ae17f 100644 --- a/crates/mesh-llm-host-runtime/src/inference/skippy/stage/mod.rs +++ b/crates/mesh-llm-host-runtime/src/inference/skippy/stage/mod.rs @@ -367,6 +367,7 @@ impl StageControlState { if load.backend == "mlx" { return mlx::load_stage(self, key, load, bind_addr).await; } + crate::system::native_runtime::ensure_native_runtime_loaded()?; let mut effective_load = load; effective_load.bind_addr = bind_addr.to_string(); super::configure_materialized_stage_cache(); diff --git a/crates/mesh-llm-host-runtime/src/lib.rs b/crates/mesh-llm-host-runtime/src/lib.rs index 403b959780..d61b64b589 100644 --- a/crates/mesh-llm-host-runtime/src/lib.rs +++ b/crates/mesh-llm-host-runtime/src/lib.rs @@ -93,14 +93,24 @@ pub async fn initialize_host_runtime_with_config(config_path: Option<&Path>) -> } None => system::native_runtime::NativeRuntimeStartupSelection::current(), }; - if let Some(runtime) = - system::native_runtime::try_load_installed_native_runtime(startup_selection).await? - { - tracing::info!( - native_runtime_id = %runtime.native_runtime_id, - libraries = ?runtime.libraries, - "Loaded MeshLLM native runtime" - ); + match system::native_runtime::try_load_installed_native_runtime(startup_selection).await { + Ok(Some(runtime)) => { + tracing::info!( + native_runtime_id = %runtime.native_runtime_id, + libraries = ?runtime.libraries, + "Loaded MeshLLM native runtime" + ); + } + Ok(None) => {} + #[cfg(all(feature = "mlx", target_os = "macos"))] + Err(error) => { + tracing::warn!( + error = %error, + "Continuing with MLX available and the Skippy/llama backend disabled" + ); + } + #[cfg(not(all(feature = "mlx", target_os = "macos")))] + Err(error) => return Err(error), } } #[cfg(not(feature = "dynamic-native-runtime"))] diff --git a/crates/mesh-llm-host-runtime/src/mesh/gossip.rs b/crates/mesh-llm-host-runtime/src/mesh/gossip.rs index 44cd659424..c2126c6688 100644 --- a/crates/mesh-llm-host-runtime/src/mesh/gossip.rs +++ b/crates/mesh-llm-host-runtime/src/mesh/gossip.rs @@ -252,6 +252,7 @@ pub(super) fn peer_meaningfully_changed(old: &PeerInfo, new: &PeerInfo) -> bool || old.artifact_transfer_supported != new.artifact_transfer_supported || old.stage_protocol_generation_supported != new.stage_protocol_generation_supported || old.stage_status_list_supported != new.stage_status_list_supported + || old.mlx_stage_supported != new.mlx_stage_supported || old.version != new.version || old.owner_summary != new.owner_summary || old.gpu_reserved_bytes != new.gpu_reserved_bytes @@ -331,6 +332,7 @@ pub(super) fn apply_transitive_ann( existing.artifact_transfer_supported = ann.artifact_transfer_supported; existing.stage_protocol_generation_supported = ann.stage_protocol_generation_supported; existing.stage_status_list_supported = ann.stage_status_list_supported; + existing.mlx_stage_supported = ann.mlx_stage_supported; existing.advertised_model_throughput = ann.advertised_model_throughput.clone(); if ann.experts_summary.is_some() { existing.experts_summary = ann.experts_summary.clone(); @@ -655,6 +657,7 @@ impl Node { existing.artifact_transfer_supported = ann.artifact_transfer_supported; existing.stage_protocol_generation_supported = ann.stage_protocol_generation_supported; existing.stage_status_list_supported = ann.stage_status_list_supported; + existing.mlx_stage_supported = ann.mlx_stage_supported; existing.advertised_model_throughput = ann.advertised_model_throughput.clone(); if ann.version.is_some() { existing.version = ann.version.clone(); @@ -976,6 +979,7 @@ impl Node { artifact_transfer_supported: peer.artifact_transfer_supported, stage_protocol_generation_supported: peer.stage_protocol_generation_supported, stage_status_list_supported: peer.stage_status_list_supported, + mlx_stage_supported: peer.mlx_stage_supported, advertised_model_throughput: peer.advertised_model_throughput.clone(), latency_ms: latency.latency_ms, latency_source: Some(match latency.source { @@ -1040,6 +1044,11 @@ impl Node { artifact_transfer_supported: data.artifact_transfer_supported, stage_protocol_generation_supported: true, stage_status_list_supported: true, + mlx_stage_supported: cfg!(all( + feature = "mlx", + target_os = "macos", + target_arch = "aarch64" + )), advertised_model_throughput: data.advertised_model_throughput, latency_ms: None, latency_source: None, @@ -1833,6 +1842,7 @@ mod tests { artifact_transfer_supported: true, stage_protocol_generation_supported: true, stage_status_list_supported: true, + mlx_stage_supported: false, advertised_model_throughput: vec![], latency_ms: None, latency_source: None, diff --git a/crates/mesh-llm-host-runtime/src/mesh/mod.rs b/crates/mesh-llm-host-runtime/src/mesh/mod.rs index 347c95722b..7a35353d25 100644 --- a/crates/mesh-llm-host-runtime/src/mesh/mod.rs +++ b/crates/mesh-llm-host-runtime/src/mesh/mod.rs @@ -1603,6 +1603,7 @@ pub(crate) struct PeerAnnouncement { pub(crate) artifact_transfer_supported: bool, pub(crate) stage_protocol_generation_supported: bool, pub(crate) stage_status_list_supported: bool, + pub(crate) mlx_stage_supported: bool, pub(crate) advertised_model_throughput: Vec, pub(crate) latency_ms: Option, pub(crate) latency_source: Option, @@ -1701,6 +1702,7 @@ pub struct PeerInfo { pub artifact_transfer_supported: bool, pub stage_protocol_generation_supported: bool, pub stage_status_list_supported: bool, + pub mlx_stage_supported: bool, pub(crate) advertised_model_throughput: Vec, /// Most recent direct RTT sample for display purposes (refreshed periodically). pub display_rtt: Option, @@ -1784,6 +1786,7 @@ impl PeerInfo { artifact_transfer_supported: ann.artifact_transfer_supported, stage_protocol_generation_supported: ann.stage_protocol_generation_supported, stage_status_list_supported: ann.stage_status_list_supported, + mlx_stage_supported: ann.mlx_stage_supported, advertised_model_throughput: ann.advertised_model_throughput.clone(), display_rtt: None, selected_path: None, @@ -7848,6 +7851,7 @@ impl Node { skippy_protocol::STAGE_SUBPROTOCOL_FEATURE_STATUS_LIST => { peer.stage_status_list_supported } + skippy_protocol::STAGE_SUBPROTOCOL_FEATURE_BACKEND_MLX => peer.mlx_stage_supported, _ => false, } } diff --git a/crates/mesh-llm-host-runtime/src/mesh/tests.rs b/crates/mesh-llm-host-runtime/src/mesh/tests.rs index 6f06d24abc..d47b6900ce 100644 --- a/crates/mesh-llm-host-runtime/src/mesh/tests.rs +++ b/crates/mesh-llm-host-runtime/src/mesh/tests.rs @@ -1251,6 +1251,7 @@ fn make_test_peer_info(peer_id: EndpointId) -> PeerInfo { artifact_transfer_supported: false, stage_protocol_generation_supported: false, stage_status_list_supported: false, + mlx_stage_supported: false, owner_summary: OwnershipSummary::default(), advertised_model_throughput: vec![], @@ -3801,6 +3802,7 @@ fn gossip_frame_roundtrip_preserves_scanned_model_metadata() { artifact_transfer_supported: false, stage_protocol_generation_supported: false, stage_status_list_supported: false, + mlx_stage_supported: false, advertised_model_throughput: vec![], latency_ms: None, latency_source: None, @@ -4089,6 +4091,7 @@ fn transitive_peer_update_refreshes_metadata_fields() { artifact_transfer_supported: true, stage_protocol_generation_supported: true, stage_status_list_supported: true, + mlx_stage_supported: false, advertised_model_throughput: vec![], latency_ms: None, latency_source: None, @@ -4179,6 +4182,7 @@ fn transitive_peer_merge_preserves_richer_direct_address() { artifact_transfer_supported: true, stage_protocol_generation_supported: true, stage_status_list_supported: true, + mlx_stage_supported: false, advertised_model_throughput: vec![], latency_ms: None, latency_source: None, @@ -4243,6 +4247,7 @@ fn transitive_peer_merge_preserves_richer_direct_address() { artifact_transfer_supported: true, stage_protocol_generation_supported: true, stage_status_list_supported: true, + mlx_stage_supported: false, advertised_model_throughput: vec![], latency_ms: None, latency_source: None, @@ -4830,6 +4835,7 @@ fn transitive_peer_update_refreshes_last_mentioned() { artifact_transfer_supported: true, stage_protocol_generation_supported: true, stage_status_list_supported: true, + mlx_stage_supported: false, advertised_model_throughput: vec![], latency_ms: None, latency_source: None, @@ -5582,6 +5588,7 @@ fn make_test_peer(id: EndpointId, rtt_ms: Option, vram_gb: u64) -> PeerInfo artifact_transfer_supported: false, stage_protocol_generation_supported: false, stage_status_list_supported: false, + mlx_stage_supported: false, owner_summary: OwnershipSummary::default(), advertised_model_throughput: vec![], @@ -6185,6 +6192,7 @@ fn requirement_peer_announcement( artifact_transfer_supported: true, stage_protocol_generation_supported: true, stage_status_list_supported: true, + mlx_stage_supported: false, advertised_model_throughput: vec![], latency_ms: None, latency_source: None, @@ -6815,6 +6823,7 @@ pub(crate) fn assert_mesh_requirements_add_peer_rejects_untrusted_release_signer artifact_transfer_supported: true, stage_protocol_generation_supported: true, stage_status_list_supported: true, + mlx_stage_supported: false, advertised_model_throughput: vec![], latency_ms: None, latency_source: None, @@ -6904,6 +6913,7 @@ pub(crate) fn assert_mesh_requirements_add_peer_rejects_invalid_release_attestat artifact_transfer_supported: true, stage_protocol_generation_supported: true, stage_status_list_supported: true, + mlx_stage_supported: false, advertised_model_throughput: vec![], latency_ms: None, latency_source: None, @@ -6990,6 +7000,7 @@ pub(crate) fn assert_mesh_requirements_add_peer_rejects_wrong_mesh_id() { artifact_transfer_supported: true, stage_protocol_generation_supported: true, stage_status_list_supported: true, + mlx_stage_supported: false, advertised_model_throughput: vec![], latency_ms: None, latency_source: None, diff --git a/crates/mesh-llm-host-runtime/src/network/openai/transport.rs b/crates/mesh-llm-host-runtime/src/network/openai/transport.rs index bf8f48a026..26fb10c466 100644 --- a/crates/mesh-llm-host-runtime/src/network/openai/transport.rs +++ b/crates/mesh-llm-host-runtime/src/network/openai/transport.rs @@ -5086,6 +5086,7 @@ mod tests { artifact_transfer_supported: false, stage_protocol_generation_supported: false, stage_status_list_supported: false, + mlx_stage_supported: false, advertised_model_throughput: vec![], display_rtt: None, selected_path: None, diff --git a/crates/mesh-llm-host-runtime/src/protocol/convert.rs b/crates/mesh-llm-host-runtime/src/protocol/convert.rs index fd1c2abcab..efdfcd5427 100644 --- a/crates/mesh-llm-host-runtime/src/protocol/convert.rs +++ b/crates/mesh-llm-host-runtime/src/protocol/convert.rs @@ -13,6 +13,7 @@ fn skippy_stage_subprotocols( artifact_transfer_supported: bool, stage_protocol_generation_supported: bool, status_list_supported: bool, + mlx_stage_supported: bool, ) -> Vec { let mut features = vec![skippy_protocol::STAGE_SUBPROTOCOL_FEATURE_STAGE_CONTROL.to_string()]; if stage_protocol_generation_supported { @@ -26,6 +27,9 @@ fn skippy_stage_subprotocols( if status_list_supported { features.push(skippy_protocol::STAGE_SUBPROTOCOL_FEATURE_STATUS_LIST.to_string()); } + if mlx_stage_supported { + features.push(skippy_protocol::STAGE_SUBPROTOCOL_FEATURE_BACKEND_MLX.to_string()); + } vec![crate::proto::node::MeshSubprotocol { name: skippy_protocol::STAGE_SUBPROTOCOL_NAME.to_string(), major: skippy_protocol::STAGE_SUBPROTOCOL_MAJOR, @@ -47,6 +51,13 @@ fn supports_skippy_status_list(subprotocols: &[crate::proto::node::MeshSubprotoc ) } +fn supports_skippy_mlx(subprotocols: &[crate::proto::node::MeshSubprotocol]) -> bool { + supports_skippy_stage_feature( + subprotocols, + skippy_protocol::STAGE_SUBPROTOCOL_FEATURE_BACKEND_MLX, + ) +} + fn supports_skippy_stage_generation(subprotocols: &[crate::proto::node::MeshSubprotocol]) -> bool { supports_skippy_stage_feature( subprotocols, @@ -697,6 +708,7 @@ pub(crate) fn local_ann_to_proto_ann( ann.artifact_transfer_supported, ann.stage_protocol_generation_supported, ann.stage_status_list_supported, + ann.mlx_stage_supported, ), } } @@ -884,6 +896,7 @@ pub(crate) fn proto_ann_to_local( artifact_transfer_supported: supports_skippy_artifact_transfer(&pa.subprotocols), stage_protocol_generation_supported: supports_skippy_stage_generation(&pa.subprotocols), stage_status_list_supported: supports_skippy_status_list(&pa.subprotocols), + mlx_stage_supported: supports_skippy_mlx(&pa.subprotocols), advertised_model_throughput: pa .advertised_model_throughput .iter() diff --git a/crates/mesh-llm-host-runtime/src/protocol/mod.rs b/crates/mesh-llm-host-runtime/src/protocol/mod.rs index ced38b9a41..431422bdd1 100644 --- a/crates/mesh-llm-host-runtime/src/protocol/mod.rs +++ b/crates/mesh-llm-host-runtime/src/protocol/mod.rs @@ -932,6 +932,7 @@ alias = "model-alias" artifact_transfer_supported: false, stage_protocol_generation_supported: false, stage_status_list_supported: false, + mlx_stage_supported: false, owner_summary: OwnershipSummary::default(), advertised_model_throughput: vec![], @@ -1373,6 +1374,7 @@ alias = "model-alias" artifact_transfer_supported: true, stage_protocol_generation_supported: true, stage_status_list_supported: true, + mlx_stage_supported: true, advertised_model_throughput: vec![], latency_ms: None, latency_source: None, @@ -1401,6 +1403,12 @@ alias = "model-alias" ); assert!(skippy.features.iter().any(|feature| feature == skippy_protocol::STAGE_SUBPROTOCOL_FEATURE_STAGE_PROTOCOL_GENERATION_V3)); + assert!( + skippy + .features + .iter() + .any(|feature| feature == skippy_protocol::STAGE_SUBPROTOCOL_FEATURE_BACKEND_MLX) + ); assert_eq!( proto_pa .owner_attestation @@ -1414,6 +1422,7 @@ alias = "model-alias" assert!(roundtripped.artifact_transfer_supported); assert!(roundtripped.stage_status_list_supported); assert!(roundtripped.stage_protocol_generation_supported); + assert!(roundtripped.mlx_stage_supported); let roundtripped = roundtripped .owner_attestation .expect("owner attestation must round-trip"); @@ -1497,6 +1506,7 @@ alias = "model-alias" artifact_transfer_supported: false, stage_protocol_generation_supported: false, stage_status_list_supported: false, + mlx_stage_supported: false, advertised_model_throughput: vec![ expected_hints[0].clone(), crate::network::metrics::ModelThroughputHint { @@ -1622,6 +1632,7 @@ alias = "model-alias" artifact_transfer_supported: true, stage_protocol_generation_supported: true, stage_status_list_supported: true, + mlx_stage_supported: false, advertised_model_throughput: vec![], latency_ms: None, latency_source: None, @@ -2479,6 +2490,7 @@ alias = "model-alias" artifact_transfer_supported: true, stage_protocol_generation_supported: true, stage_status_list_supported: true, + mlx_stage_supported: false, advertised_model_throughput: vec![], latency_ms: None, latency_source: None, @@ -2535,6 +2547,7 @@ alias = "model-alias" artifact_transfer_supported: false, stage_protocol_generation_supported: false, stage_status_list_supported: false, + mlx_stage_supported: false, advertised_model_throughput: vec![], latency_ms: None, latency_source: None, diff --git a/crates/mesh-llm-host-runtime/src/runtime/local.rs b/crates/mesh-llm-host-runtime/src/runtime/local.rs index 9c5fd47fcc..944b1f107e 100644 --- a/crates/mesh-llm-host-runtime/src/runtime/local.rs +++ b/crates/mesh-llm-host-runtime/src/runtime/local.rs @@ -85,7 +85,6 @@ pub(super) enum LocalRuntimeBackendHandle { // as the HTTP server is serving it; not read directly. _model: crate::inference::mlx::MlxModelHandle, http: crate::inference::mlx::MlxHttpHandle, - _death_tx: tokio::sync::oneshot::Sender<()>, }, } @@ -598,6 +597,8 @@ pub(super) async fn start_runtime_local_model( return start_runtime_mlx_model(spec, model_name).await; } + crate::system::native_runtime::ensure_native_runtime_loaded()?; + let package_ref = spec.model_path.to_string_lossy().to_string(); let layer_package = if skippy::is_layer_package_ref(&package_ref) { let package_ref_for_identity = package_ref.clone(); @@ -721,6 +722,12 @@ fn scan_layer_package_metadata( } pub(super) fn runtime_model_planning_bytes(model_path: &Path) -> Result { + #[cfg(all(feature = "mlx", target_os = "macos"))] + if let Ok(descriptor) = model_hf::safetensors_stage::read_checkpoint_descriptor(model_path) + && descriptor.estimated_affine4_weight_bytes > 0 + { + return Ok(descriptor.estimated_affine4_weight_bytes); + } let package_ref = model_path.to_string_lossy().to_string(); if skippy::is_layer_package_ref(&package_ref) { return Ok(skippy::identity_from_layer_package(&package_ref)?.source_model_bytes); @@ -759,8 +766,7 @@ pub(super) async fn start_runtime_split_model( let SplitRuntimeStartPreparation { package, participant_snapshot, - compact_meta, - kv_bytes_per_token, + planning_metadata, planned_topology, } = split_setup; let stages = planned_topology.stages; @@ -815,7 +821,8 @@ pub(super) async fn start_runtime_split_model( SPLIT_INITIAL_SHUTDOWN_GENERATION, planned_participants, stages, - ); + ) + .with_control_managed_stage0(package.package_ref.starts_with("hf-model://")); let mut loaded = load_split_runtime_generation(SplitGenerationLoadSpec { node: spec.node, mesh_config: spec.mesh_config, @@ -850,10 +857,10 @@ pub(super) async fn start_runtime_split_model( projector_path, ctx_size, topology_resources: SplitTopologyResourceInputs { - native_context_length: compact_meta.context_length, - kv_bytes_per_token, + native_context_length: planning_metadata.native_context_length, + kv_bytes_per_token: planning_metadata.kv_bytes_per_token, ctx_size_override: spec.ctx_size_override, - parallel_override: spec.parallel_override, + parallel_override: split_parallel_override(&package, spec.parallel_override), }, cache_type_k_override: spec.cache_type_k_override.map(str::to_string), cache_type_v_override: spec.cache_type_v_override.map(str::to_string), @@ -874,11 +881,16 @@ pub(super) async fn start_runtime_split_model( struct SplitRuntimeStartPreparation { package: skippy::SkippyPackageIdentity, participant_snapshot: SplitParticipantSnapshot, - compact_meta: models::gguf::GgufCompactMeta, - kv_bytes_per_token: u64, + planning_metadata: SplitRuntimePlanningMetadata, planned_topology: PlannedRuntimeSliceTopology, } +#[derive(Clone, Copy, Debug)] +struct SplitRuntimePlanningMetadata { + native_context_length: u32, + kv_bytes_per_token: u64, +} + async fn prepare_split_runtime_start( spec: &LocalRuntimeModelStartSpec<'_>, model_ref: &str, @@ -895,13 +907,12 @@ async fn prepare_split_runtime_start( timeout, ) .await?; - let compact_meta = split_runtime_compact_meta(&package).await?; - let kv_bytes_per_token = split_runtime_kv_bytes_per_token( + let planning_metadata = split_runtime_planning_metadata( &package, - &compact_meta, spec.cache_type_k_override, spec.cache_type_v_override, - )?; + ) + .await?; let planned_topology = plan_runtime_slice_topology_with_resources( topology_id, model_ref, @@ -909,21 +920,57 @@ async fn prepare_split_runtime_start( &participant_snapshot.participants, &participant_snapshot.excluded, SplitTopologyResourceInputs { - native_context_length: compact_meta.context_length, - kv_bytes_per_token, + native_context_length: planning_metadata.native_context_length, + kv_bytes_per_token: planning_metadata.kv_bytes_per_token, ctx_size_override: spec.ctx_size_override, - parallel_override: spec.parallel_override, + parallel_override: split_parallel_override(&package, spec.parallel_override), }, )?; Ok(SplitRuntimeStartPreparation { package, participant_snapshot, - compact_meta, - kv_bytes_per_token, + planning_metadata, planned_topology, }) } +fn split_parallel_override( + package: &skippy::SkippyPackageIdentity, + requested: Option, +) -> Option { + if package.package_ref.starts_with("hf-model://") { + Some(1) + } else { + requested + } +} + +async fn split_runtime_planning_metadata( + package: &skippy::SkippyPackageIdentity, + cache_type_k_override: Option<&str>, + cache_type_v_override: Option<&str>, +) -> Result { + #[cfg(all(feature = "mlx", target_os = "macos"))] + if is_mlx_split_package(package) { + let descriptor = describe_mlx_split_package(package).await?; + return Ok(SplitRuntimePlanningMetadata { + native_context_length: descriptor.native_context_length, + kv_bytes_per_token: descriptor.kv_bytes_per_token_bf16, + }); + } + + let compact_meta = split_runtime_compact_meta(package).await?; + Ok(SplitRuntimePlanningMetadata { + native_context_length: compact_meta.context_length, + kv_bytes_per_token: split_runtime_kv_bytes_per_token( + package, + &compact_meta, + cache_type_k_override, + cache_type_v_override, + )?, + }) +} + async fn split_runtime_compact_meta( package: &skippy::SkippyPackageIdentity, ) -> Result { @@ -979,6 +1026,10 @@ async fn resolve_split_runtime_package( model_path: &Path, model_ref: &str, ) -> Result { + #[cfg(all(feature = "mlx", target_os = "macos"))] + if is_safetensors_model_path(model_path) { + return resolve_mlx_split_package(model_path).await; + } let model_path_str = model_path.to_string_lossy().to_string(); if skippy::is_layer_package_ref(&model_path_str) { Ok(tokio::task::spawn_blocking(move || { @@ -993,6 +1044,119 @@ async fn resolve_split_runtime_package( } } +#[cfg(all(feature = "mlx", target_os = "macos"))] +fn is_mlx_split_package(package: &skippy::SkippyPackageIdentity) -> bool { + package.package_ref.starts_with("hf-model://") +} + +#[cfg(all(feature = "mlx", target_os = "macos"))] +async fn resolve_mlx_split_package(model_path: &Path) -> Result { + let prepared_descriptor = + model_hf::safetensors_stage::read_checkpoint_descriptor(model_path).ok(); + let (package_ref, descriptor) = if let Some(descriptor) = prepared_descriptor { + anyhow::ensure!( + is_immutable_huggingface_revision(&descriptor.revision), + "distributed MLX checkpoint descriptor revision must be an immutable 40-character commit SHA" + ); + ( + format!("hf-model://{}@{}", descriptor.repo, descriptor.revision), + descriptor, + ) + } else { + let identity = mlx_huggingface_identity(model_path).with_context(|| { + format!( + "distributed MLX requires a SafeTensors model in the Hugging Face cache with an immutable revision: {}", + model_path.display() + ) + })?; + let package_ref = format!("hf-model://{}@{}", identity.repo_id, identity.revision); + let descriptor = describe_mlx_checkpoint(&identity.repo_id, &identity.revision).await?; + (package_ref, descriptor) + }; + Ok(skippy::SkippyPackageIdentity { + package_ref, + manifest_sha256: descriptor.checkpoint_sha256.clone(), + source_model_path: mlx_model_dir(model_path), + source_model_sha256: descriptor.checkpoint_sha256, + source_model_bytes: descriptor.dense_tensor_bytes, + source_files: Vec::new(), + layer_weight_bytes: descriptor.estimated_affine4_layer_bytes, + layer_count: descriptor.layer_count, + activation_width: descriptor.hidden_size, + tensor_count: 0, + generation: None, + }) +} + +#[cfg(all(feature = "mlx", target_os = "macos"))] +fn is_immutable_huggingface_revision(revision: &str) -> bool { + revision.len() == 40 && revision.bytes().all(|byte| byte.is_ascii_hexdigit()) +} + +#[cfg(all(feature = "mlx", target_os = "macos"))] +fn mlx_huggingface_identity( + model_path: &Path, +) -> Option { + if model_path.is_file() { + return models::huggingface_identity_for_path(model_path); + } + [ + "model.safetensors.index.json", + "model.safetensors", + "model-00001-of-00001.safetensors", + ] + .into_iter() + .map(|name| model_path.join(name)) + .find_map(|path| models::huggingface_identity_for_path(&path)) + .or_else(|| { + std::fs::read_dir(model_path) + .ok()? + .filter_map(Result::ok) + .map(|entry| entry.path()) + .filter(|path| { + path.file_name() + .and_then(|name| name.to_str()) + .is_some_and(|name| name.ends_with(".safetensors")) + }) + .find_map(|path| models::huggingface_identity_for_path(&path)) + }) +} + +#[cfg(all(feature = "mlx", target_os = "macos"))] +async fn describe_mlx_split_package( + package: &skippy::SkippyPackageIdentity, +) -> Result { + if let Ok(descriptor) = + model_hf::safetensors_stage::read_checkpoint_descriptor(&package.source_model_path) + { + return Ok(descriptor); + } + let source = package + .package_ref + .strip_prefix("hf-model://") + .context("MLX split package is missing hf-model:// prefix")?; + let parsed = model_ref::parse_model_ref(source).context("parse MLX split package ref")?; + let revision = parsed + .revision + .context("MLX split package requires immutable revision")?; + describe_mlx_checkpoint(&parsed.repo, &revision).await +} + +#[cfg(all(feature = "mlx", target_os = "macos"))] +async fn describe_mlx_checkpoint( + repo: &str, + revision: &str, +) -> Result { + let repo = repo.to_string(); + let revision = revision.to_string(); + tokio::task::spawn_blocking(move || { + model_hf::safetensors_stage::SafetensorsStageMaterializer::from_environment()? + .describe_checkpoint(&repo, &revision) + }) + .await + .context("join MLX checkpoint descriptor task")? +} + fn split_kv_cache_quant( split_kv_policy: &skippy::KvCachePolicy, cache_type_k_override: Option<&str>, @@ -1108,6 +1272,7 @@ fn package_ref_has_independent_prepare_source(package_ref: &str) -> bool { // HF layer packages can be resolved by the selected worker during prepare; // peer artifact transfer is only an optional cache warm path. skippy_runtime::package::is_hf_package_ref(package_ref) + || package_ref.starts_with("hf-model://") } #[derive(Clone, Debug, Eq, PartialEq)] @@ -1128,6 +1293,7 @@ pub(super) enum SplitParticipantExclusionReason { MissingVram, MissingModelInterest, StageProtocolGeneration, + MlxBackendUnavailable, MissingStagePath, StagePathRelayOnly, StagePathTooSlow, @@ -1145,6 +1311,7 @@ impl SplitParticipantExclusionReason { Self::MissingVram => "missing_vram", Self::MissingModelInterest => "missing_model_interest", Self::StageProtocolGeneration => "stage_protocol_generation", + Self::MlxBackendUnavailable => "mlx_backend_unavailable", Self::MissingStagePath => "missing_stage_path", Self::StagePathRelayOnly => "stage_path_relay_only", Self::StagePathTooSlow => "stage_path_too_slow", @@ -1168,6 +1335,9 @@ impl SplitParticipantExclusionReason { Self::StageProtocolGeneration => { "Upgrade this peer so it advertises current stage protocol support." } + Self::MlxBackendUnavailable => { + "Use an Apple Silicon peer running an MLX-enabled mesh-llm build." + } Self::MissingStagePath => { "Wait for direct peer latency to be measured or fix direct QUIC connectivity." } @@ -1225,6 +1395,7 @@ struct SplitGenerationLoadSpec<'a> { struct SplitGenerationLoadSettings<'a> { stage0: &'a RuntimeSliceStagePlan, + backend: SplitRuntimeBackend, runtime_options: skippy_server::EmbeddedRuntimeOptions, embedded_openai: skippy::ResolvedEmbeddedOpenAiArgs, load_mode: LoadMode, @@ -1232,6 +1403,12 @@ struct SplitGenerationLoadSettings<'a> { activation_wire_dtype: skippy::StageWireDType, } +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +enum SplitRuntimeBackend { + Skippy, + Mlx, +} + async fn load_split_runtime_generation( spec: SplitGenerationLoadSpec<'_>, ) -> Result { @@ -1306,6 +1483,16 @@ async fn load_split_runtime_generation_inner( &stage0_return_endpoint, )) .await?; + if settings.backend == SplitRuntimeBackend::Mlx { + return Box::pin(load_mlx_split_runtime_generation( + spec, + &settings, + cleanup_on_error, + &mut ready_by_stage, + downstream, + )) + .await; + } let downstream_endpoint = if downstream.node_id == Some(spec.node.id()) { downstream.endpoint } else { @@ -1430,6 +1617,141 @@ async fn load_split_runtime_generation_inner( }) } +#[cfg(all(feature = "mlx", target_os = "macos"))] +async fn load_mlx_split_runtime_generation( + spec: &SplitGenerationLoadSpec<'_>, + settings: &SplitGenerationLoadSettings<'_>, + cleanup_on_error: &mut bool, + ready_by_stage: &mut HashMap, + downstream: skippy::StagePeerDescriptor, +) -> Result { + *cleanup_on_error = true; + let load = split_runtime_stage_load_request( + spec, + settings, + settings.stage0, + Some(downstream), + "127.0.0.1:0", + ); + prepare_split_stage(spec.node, settings.stage0.node_id, load.clone()).await?; + wait_for_split_stage_source( + spec.node, + settings.stage0.node_id, + &load, + Duration::from_secs(30 * 60), + ) + .await + .context("prepare local MLX stage 0")?; + let response = spec + .node + .send_local_stage_control(skippy::StageControlRequest::Load(load)) + .await + .context("load local MLX stage 0")?; + let skippy::StageControlResponse::Ready(ready) = response else { + anyhow::bail!("unexpected response while loading MLX stage 0"); + }; + anyhow::ensure!( + ready.accepted, + "MLX stage 0 rejected load: {}", + ready.error.unwrap_or_else(|| "unknown error".to_string()) + ); + let stage_addr = ready + .status + .bind_addr + .parse() + .with_context(|| format!("parse MLX stage 0 endpoint {}", ready.status.bind_addr))?; + ready_by_stage.insert(settings.stage0.stage_id.clone(), ready.status); + + let model_dir = mlx_model_dir(spec.model_path); + let model_ref = spec.model_ref.to_string(); + let context_length = spec.ctx_size; + let wire_dtype = skippy_wire_dtype(settings.activation_wire_dtype)?; + let _ = emit_event(OutputEvent::ModelLoading { + model: model_ref.clone(), + source: None, + }); + let load_model_ref = model_ref.clone(); + let mlx_model = tokio::task::spawn_blocking(move || { + crate::inference::mlx::MlxModelHandle::load_distributed( + model_dir, + load_model_ref, + context_length, + stage_addr, + wire_dtype, + ) + }) + .await + .context("join distributed MLX frontend load task")??; + let _ = emit_event(OutputEvent::ModelLoaded { + model: model_ref.clone(), + bytes: None, + }); + let (death_tx, death_rx) = tokio::sync::oneshot::channel(); + let http = mlx_model + .start_http(alloc_local_port().await?, death_tx) + .await?; + let capabilities = models::runtime_verified_model_capabilities( + spec.model_ref, + spec.model_path, + models::RuntimeMediaCapabilityEvidence { + vision_projector_loaded: false, + }, + ); + spec.node + .activate_stage_topology(split_stage_topology_instance( + &spec.generation.topology_id, + &spec.generation.run_id, + spec.model_ref, + spec.package, + &spec.generation.stages, + ready_by_stage, + )) + .await; + + Ok(SplitRuntimeGenerationHandle { + loaded_name: model_ref, + handle: LocalRuntimeModelHandle { + port: http.port(), + backend: "mlx".to_string(), + context_length, + slots: 1, + capabilities, + inner: LocalRuntimeBackendHandle::Mlx { + _model: mlx_model, + http, + }, + }, + death_rx, + cleanup: Some(SplitGenerationCleanup { + generation: spec.generation.clone(), + }), + coordinator_rx: None, + coordinator_task: None, + }) +} + +#[cfg(not(all(feature = "mlx", target_os = "macos")))] +async fn load_mlx_split_runtime_generation( + _spec: &SplitGenerationLoadSpec<'_>, + _settings: &SplitGenerationLoadSettings<'_>, + _cleanup_on_error: &mut bool, + _ready_by_stage: &mut HashMap, + _downstream: skippy::StagePeerDescriptor, +) -> Result { + anyhow::bail!("distributed MLX serving requires macOS and the mlx feature") +} + +#[cfg(all(feature = "mlx", target_os = "macos"))] +fn skippy_wire_dtype( + dtype: skippy::StageWireDType, +) -> Result { + match dtype { + skippy::StageWireDType::F32 => Ok(skippy_protocol::binary::WireActivationDType::F32), + skippy::StageWireDType::F16 => Ok(skippy_protocol::binary::WireActivationDType::F16), + skippy::StageWireDType::Q8 => anyhow::bail!("MLX stages do not support Q8 activations"), + } +} + async fn load_downstream_split_runtime_stages( spec: &SplitGenerationLoadSpec<'_>, settings: &SplitGenerationLoadSettings<'_>, @@ -1447,13 +1769,13 @@ async fn load_downstream_split_runtime_stages( downstream.clone(), stage0_return_endpoint, ); - prepare_split_stage(spec.node, stage.node_id, load.clone()).await?; - wait_for_split_stage_source( + Box::pin(prepare_split_stage(spec.node, stage.node_id, load.clone())).await?; + Box::pin(wait_for_split_stage_source( spec.node, stage.node_id, &load, Duration::from_secs(30 * 60), - ) + )) .await .with_context(|| { format!( @@ -1512,6 +1834,7 @@ fn split_runtime_stage_load_request( stage0_return_endpoint: &str, ) -> skippy::StageLoadRequest { let resolved_config = &settings.runtime_options.config; + let mlx = settings.backend == SplitRuntimeBackend::Mlx; let upstream = if downstream.is_none() { split_runtime_stage_upstream(spec, stage0_return_endpoint) } else { @@ -1521,7 +1844,7 @@ fn split_runtime_stage_load_request( topology_id: spec.generation.topology_id.clone(), run_id: spec.generation.run_id.clone(), model_id: spec.model_ref.to_string(), - backend: "skippy".to_string(), + backend: if mlx { "mlx" } else { "skippy" }.to_string(), package_ref: spec.package.package_ref.clone(), manifest_sha256: spec.package.manifest_sha256.clone(), stage_id: stage.stage_id.clone(), @@ -1534,23 +1857,27 @@ fn split_runtime_stage_load_request( spec.model_path, )), source_model_bytes: Some(spec.package.source_model_bytes), - projector_path: spec.projector_path.clone(), + projector_path: (!mlx).then(|| spec.projector_path.clone()).flatten(), selected_device: None, bind_addr: "127.0.0.1:0".to_string(), activation_width: settings.activation_width, wire_dtype: settings.activation_wire_dtype, ctx_size: spec.ctx_size, - lane_count: spec.slots as u32, - n_batch: resolved_config.n_batch, - n_ubatch: resolved_config.n_ubatch, + lane_count: if mlx { 1 } else { spec.slots as u32 }, + n_batch: (!mlx).then_some(resolved_config.n_batch).flatten(), + n_ubatch: (!mlx).then_some(resolved_config.n_ubatch).flatten(), n_gpu_layers: resolved_config.n_gpu_layers, - mmap: resolved_config.mmap, - mlock: resolved_config.mlock, + mmap: (!mlx).then_some(resolved_config.mmap).flatten(), + mlock: !mlx && resolved_config.mlock, weight_quantization: skippy::StageWeightQuantization::Auto, cache_type_k: resolved_config.cache_type_k.clone(), cache_type_v: resolved_config.cache_type_v.clone(), - flash_attn_type: resolved_config.flash_attn_type, - native_mtp_enabled: resolved_config.native_mtp_enabled, + flash_attn_type: if mlx { + FlashAttentionType::Auto + } else { + resolved_config.flash_attn_type + }, + native_mtp_enabled: !mlx && resolved_config.native_mtp_enabled, shutdown_generation: spec.generation.generation, coordinator_term: spec.generation.coordinator_term, coordinator_id: Some(spec.node.id()), @@ -1582,6 +1909,11 @@ fn split_generation_load_settings<'a>( .stages .first() .context("split topology did not produce stage 0")?; + let backend = if spec.package.package_ref.starts_with("hf-model://") { + SplitRuntimeBackend::Mlx + } else { + SplitRuntimeBackend::Skippy + }; let load_mode = split_generation_load_mode(spec.package); let activation_width = skippy_stage_activation_width(spec.package.activation_width, spec.model_ref)?; @@ -1631,16 +1963,23 @@ fn split_generation_load_settings<'a>( ); Ok(SplitGenerationLoadSettings { stage0, + backend, runtime_options, embedded_openai, load_mode, activation_width, - activation_wire_dtype: resolved.skippy.activation_wire_dtype, + activation_wire_dtype: if backend == SplitRuntimeBackend::Mlx { + skippy::StageWireDType::F16 + } else { + resolved.skippy.activation_wire_dtype + }, }) } fn split_generation_load_mode(package: &skippy::SkippyPackageIdentity) -> LoadMode { - if skippy::is_layer_package_ref(&package.package_ref) { + if package.package_ref.starts_with("hf-model://") { + LoadMode::ArtifactSlice + } else if skippy::is_layer_package_ref(&package.package_ref) { LoadMode::LayerPackage } else { LoadMode::RuntimeSlice @@ -1885,6 +2224,7 @@ struct SplitTopologyGeneration { lease_until_unix_ms: u64, participants: Vec, stages: Vec, + control_managed_stage0: bool, } impl SplitTopologyGeneration { @@ -1903,8 +2243,14 @@ impl SplitTopologyGeneration { lease_until_unix_ms: split_coordinator_lease_until_unix_ms(), participants, stages, + control_managed_stage0: false, } } + + fn with_control_managed_stage0(mut self, enabled: bool) -> Self { + self.control_managed_stage0 = enabled; + self + } } struct SplitTopologyCoordinator { @@ -2386,16 +2732,19 @@ impl SplitTopologyCoordinator { split_stages_meet_minimum(&stages), "split runtime needs at least two stage participants" ); - Ok(SplitTopologyGeneration::new( - topology_id, - run_id, - generation, - participants, - stages, - )) + Ok( + SplitTopologyGeneration::new(topology_id, run_id, generation, participants, stages) + .with_control_managed_stage0(self.active.control_managed_stage0), + ) } fn local_model_fits(&self) -> bool { + if self.package.package_ref.starts_with("hf-model://") { + // Metadata-only MLX startup deliberately has no whole checkpoint + // for a local fallback. Keep serving withdrawn until a replacement + // split can be elected. + return false; + } let local_capacity = self .pinned_gpu .as_ref() @@ -2750,7 +3099,8 @@ async fn stop_split_generation( ) .await; } - for stage in generation.stages.iter().skip(1) { + let skip = usize::from(!generation.control_managed_stage0); + for stage in generation.stages.iter().skip(skip) { let stop = skippy::StageStopRequest { topology_id: generation.topology_id.clone(), run_id: generation.run_id.clone(), @@ -2983,7 +3333,7 @@ fn split_participant_blocker( }) } -const fn split_participant_exclusion_reason_order() -> [SplitParticipantExclusionReason; 12] { +const fn split_participant_exclusion_reason_order() -> [SplitParticipantExclusionReason; 13] { [ SplitParticipantExclusionReason::StageControlUnreachable, SplitParticipantExclusionReason::PackageManifestMismatch, @@ -2994,6 +3344,7 @@ const fn split_participant_exclusion_reason_order() -> [SplitParticipantExclusio SplitParticipantExclusionReason::StagePathRelayOnly, SplitParticipantExclusionReason::StagePathTooSlow, SplitParticipantExclusionReason::StageProtocolGeneration, + SplitParticipantExclusionReason::MlxBackendUnavailable, SplitParticipantExclusionReason::MissingVram, SplitParticipantExclusionReason::MissingModelInterest, SplitParticipantExclusionReason::Client, @@ -3049,6 +3400,13 @@ async fn collect_split_participants( }); continue; } + if let Some(reason) = split_peer_backend_exclusion_reason(&peer, package) { + excluded.push(SplitParticipantExclusion { + node_id: peer.id, + reason, + }); + continue; + } if let Some(reason) = split_peer_stage_path_exclusion_reason(node.split_stage_path_snapshot(peer.id).await) { @@ -3097,6 +3455,14 @@ async fn collect_split_participants( } } +fn split_peer_backend_exclusion_reason( + peer: &mesh::PeerInfo, + package: &skippy::SkippyPackageIdentity, +) -> Option { + (package.package_ref.starts_with("hf-model://") && !peer.mlx_stage_supported) + .then_some(SplitParticipantExclusionReason::MlxBackendUnavailable) +} + fn split_peer_preflight_exclusion_reason( peer: &mesh::PeerInfo, model_name: &str, @@ -3195,7 +3561,9 @@ fn split_inventory_package_signal_result( if split_inventory_manifest_mismatch(inventory, package) { return Err(SplitParticipantExclusionReason::PackageManifestMismatch); } - if split_inventory_has_no_stage_surface(inventory) { + if split_inventory_has_no_stage_surface(inventory) + && !package_ref_has_independent_prepare_source(&package.package_ref) + { return Err(SplitParticipantExclusionReason::StageInventoryEmpty); } let signal = split_inventory_package_signal(inventory, package); @@ -3487,7 +3855,7 @@ async fn wait_for_split_stage_source( ) -> Result<()> { let deadline = tokio::time::Instant::now() + timeout; loop { - let inventory = query_stage_inventory(node, stage_node_id, load) + let inventory = Box::pin(query_stage_inventory(node, stage_node_id, load)) .await .with_context(|| stage_control_unreachable_message(&load.stage_id, stage_node_id))?; if split_stage_source_is_ready(&inventory, load) { @@ -3669,7 +4037,10 @@ fn is_safetensors_model_path(path: &Path) -> bool { } if path.is_dir() { return path.join("model.safetensors").exists() - || path.join("model.safetensors.index.json").exists(); + || path.join("model.safetensors.index.json").exists() + || path + .join(model_hf::safetensors_stage::CHECKPOINT_DESCRIPTOR_FILE) + .exists(); } false } @@ -3722,8 +4093,8 @@ async fn start_runtime_mlx_model( bytes: None, }); - let http = mlx_model.start_http(port); let (death_tx, death_rx) = tokio::sync::oneshot::channel(); + let http = mlx_model.start_http(port, death_tx).await?; Ok(( model_name, @@ -3736,7 +4107,6 @@ async fn start_runtime_mlx_model( inner: LocalRuntimeBackendHandle::Mlx { _model: mlx_model, http, - _death_tx: death_tx, }, }, death_rx, @@ -4129,6 +4499,7 @@ mod tests { artifact_transfer_supported: false, stage_protocol_generation_supported, stage_status_list_supported: false, + mlx_stage_supported: false, advertised_model_throughput: vec![], display_rtt: None, @@ -4872,6 +5243,33 @@ max_tokens = 222 assert!(signal.can_stage_with(&package, false)); } + #[test] + fn mlx_package_signal_allows_cold_independent_prepare() { + let mut package = package(10); + package.package_ref = format!("hf-model://org/model@{}", "a".repeat(40)); + package.manifest_sha256 = "b".repeat(64); + let inventory = skippy::StageLayerInventory { + model_id: "model-a".to_string(), + package_ref: package.package_ref.clone(), + manifest_sha256: package.manifest_sha256.clone(), + layer_count: 0, + ready_ranges: Vec::new(), + available_ranges: Vec::new(), + missing_ranges: Vec::new(), + preparing_ranges: Vec::new(), + source_model_path: None, + source_model_bytes: None, + source_model_kind: skippy::SourceModelKind::Unknown, + weight_quantization: skippy::StageWeightQuantization::Auto, + }; + + let signal = split_inventory_package_signal_result(&inventory, &package, false) + .expect("immutable MLX source should be independently preparable"); + + assert_eq!(signal.cached_slice_bytes, 0); + assert_eq!(signal.missing_artifact_bytes, package.source_model_bytes); + } + #[test] fn split_participant_timeout_error_reports_blocker_summary() { let participants = vec![SplitParticipant::new(make_id(1), 2_000_000_000, None)]; @@ -4929,6 +5327,33 @@ max_tokens = 222 ); } + #[test] + fn mlx_split_requires_explicit_peer_backend_capability() { + let mut peer = split_test_peer(0x71, "SmolLM2", true); + let mut package = package(30); + package.package_ref = format!("hf-model://org/model@{}", "a".repeat(40)); + + assert_eq!( + split_peer_backend_exclusion_reason(&peer, &package), + Some(SplitParticipantExclusionReason::MlxBackendUnavailable) + ); + + peer.mlx_stage_supported = true; + assert_eq!(split_peer_backend_exclusion_reason(&peer, &package), None); + } + + #[test] + fn mlx_split_uses_artifact_slice_and_one_lane() { + let mut package = package(30); + package.package_ref = format!("hf-model://org/model@{}", "a".repeat(40)); + + assert_eq!( + split_generation_load_mode(&package), + LoadMode::ArtifactSlice + ); + assert_eq!(split_parallel_override(&package, Some(8)), Some(1)); + } + #[test] fn split_peer_preflight_requires_measured_stage_path() { assert_eq!( @@ -5734,7 +6159,7 @@ max_tokens = 222 ); let mesh_config = plugin::MeshConfig::default(); - let error = match Box::pin(load_split_runtime_generation(SplitGenerationLoadSpec { + let load_future = load_split_runtime_generation(SplitGenerationLoadSpec { node: &node, mesh_config: &mesh_config, model_ref: "Qwen", @@ -5755,9 +6180,8 @@ max_tokens = 222 ), skippy_telemetry: skippy::SkippyTelemetryOptions::off(), survey_telemetry: survey::SurveyTelemetry::disabled(), - })) - .await - { + }); + let error = match Box::pin(load_future).await { Ok(_) => panic!("candidate split generation load unexpectedly succeeded"), Err(error) => error, }; diff --git a/crates/mesh-llm-host-runtime/src/runtime/mod.rs b/crates/mesh-llm-host-runtime/src/runtime/mod.rs index f905e2853d..494ccba454 100644 --- a/crates/mesh-llm-host-runtime/src/runtime/mod.rs +++ b/crates/mesh-llm-host-runtime/src/runtime/mod.rs @@ -67,6 +67,8 @@ use mesh_llm_events::{ emit_event, flush_output, output_sink, schedule_ready_prompt, sort_dashboard_endpoint_rows, }; use mesh_llm_node::serving::{UnloadOptions, UnloadTarget}; +#[cfg(all(feature = "mlx", target_os = "macos"))] +use model_artifact::ModelRepository; use skippy_protocol::FlashAttentionType; use std::cell::Cell; use std::collections::{BTreeMap, BTreeSet, HashMap}; @@ -4124,7 +4126,7 @@ fn build_startup_model_specs( async fn resolve_startup_models( specs: &[StartupModelSpec], - _split: bool, + split: bool, ) -> Result> { let mut plans = Vec::with_capacity(specs.len()); for spec in specs { @@ -4135,15 +4137,26 @@ async fn resolve_startup_models( // later, so layer-package discovery must not depend on `--split`. let requested_ref_for_catalog = requested_ref.to_string(); let model_ref_for_catalog = spec.model_ref.clone(); - let resolved_path = if let Some(package_ref) = tokio::task::spawn_blocking(move || { - resolve_split_layer_package(&requested_ref_for_catalog, &model_ref_for_catalog) - }) - .await - .context("join resolve layer package task")? - { - PathBuf::from(package_ref) + let mlx_checkpoint = if split { + resolve_split_mlx_checkpoint(&requested_ref).await? } else { - resolve_model(&spec.model_ref).await? + None + }; + let is_split_mlx_checkpoint = mlx_checkpoint.is_some(); + let resolved_path = match mlx_checkpoint { + Some(path) => path, + None => { + if let Some(package_ref) = tokio::task::spawn_blocking(move || { + resolve_split_layer_package(&requested_ref_for_catalog, &model_ref_for_catalog) + }) + .await + .context("join resolve layer package task")? + { + PathBuf::from(package_ref) + } else { + resolve_model(&spec.model_ref).await? + } + } }; let mmproj_path = match spec.mmproj_ref.as_ref() { @@ -4157,7 +4170,7 @@ async fn resolve_startup_models( // For hf:// layer package refs, use the requested ref as the model ref // rather than trying to parse the hf:// URL as a filesystem path. let path_str = resolved_path.to_string_lossy(); - if path_str.starts_with("hf://") { + if path_str.starts_with("hf://") || is_split_mlx_checkpoint { requested_ref.to_string() } else if resolved_path.join("model-package.json").is_file() { // Layer package directory: read the canonical model_id from the manifest @@ -4187,6 +4200,61 @@ async fn resolve_startup_models( Ok(plans) } +#[cfg(all(feature = "mlx", target_os = "macos"))] +async fn resolve_split_mlx_checkpoint(requested_ref: &str) -> Result> { + let Some((repo, revision)) = split_mlx_huggingface_source(requested_ref) else { + return Ok(None); + }; + let repository = model_hf::HfModelRepository::from_env()?; + let revision = repository + .resolve_revision(&repo, revision.as_deref()) + .await?; + let files = repository.list_files(&repo, &revision).await?; + let has_safetensors = files.iter().any(|file| { + file.path == "model.safetensors" + || file.path == "model.safetensors.index.json" + || model_resolver::is_split_mlx_first_shard(&file.path) + }); + if !has_safetensors { + return Ok(None); + } + anyhow::ensure!( + files.iter().any(|file| file.path == "tokenizer.json"), + "distributed MLX checkpoint {repo}@{revision} has no tokenizer.json" + ); + let prepared = tokio::task::spawn_blocking(move || { + model_hf::safetensors_stage::SafetensorsStageMaterializer::from_environment()? + .prepare_checkpoint(&repo, &revision) + }) + .await + .context("join distributed MLX checkpoint preparation")??; + tracing::info!( + checkpoint = %prepared.descriptor.checkpoint_sha256, + dense_tensor_bytes = prepared.descriptor.dense_tensor_bytes, + planned_affine4_bytes = prepared.descriptor.estimated_affine4_weight_bytes, + sidecar_path = %prepared.path.display(), + "prepared metadata-only distributed MLX checkpoint" + ); + Ok(Some(prepared.path)) +} + +#[cfg(all(feature = "mlx", target_os = "macos"))] +fn split_mlx_huggingface_source(input: &str) -> Option<(String, Option)> { + if let Some((repo, revision, file)) = model_resolver::parse_huggingface_file_ref(input) { + let safetensors = file == "model.safetensors" + || file == "model.safetensors.index.json" + || model_resolver::is_split_mlx_first_shard(&file); + return safetensors.then_some((repo, revision)); + } + let (repo, revision, selector) = model_resolver::parse_huggingface_repo_ref(input)?; + selector.is_none().then_some((repo, revision)) +} + +#[cfg(not(all(feature = "mlx", target_os = "macos")))] +async fn resolve_split_mlx_checkpoint(_requested_ref: &str) -> Result> { + Ok(None) +} + /// Read the `model_id` field from a layer package's `model-package.json`. fn read_layer_package_model_id(package_dir: &Path) -> Option { let manifest_path = package_dir.join("model-package.json"); diff --git a/crates/mesh-llm-host-runtime/src/runtime/split_planning.rs b/crates/mesh-llm-host-runtime/src/runtime/split_planning.rs index cc602d40f6..5a5d1ba1c9 100644 --- a/crates/mesh-llm-host-runtime/src/runtime/split_planning.rs +++ b/crates/mesh-llm-host-runtime/src/runtime/split_planning.rs @@ -421,9 +421,10 @@ pub(super) fn validate_split_capacity( .iter() .map(|participant| participant.vram_bytes) .sum::(); - // Use raw model weight for aggregate split check — the topology planner - // already performed detailed per-node budgeting with KV and headroom. - let required_total_bytes = package.source_model_bytes; + // Use exact runtime layer estimates when available. This is especially + // important for load-time-quantized SafeTensors checkpoints, whose dense + // source size can exceed the memory actually required by every stage. + let required_total_bytes = required_split_weight_bytes(package); anyhow::ensure!( total_vram_bytes >= required_total_bytes, "{}", @@ -460,6 +461,17 @@ pub(super) fn validate_split_capacity( Ok(()) } +fn required_split_weight_bytes(package: &skippy::SkippyPackageIdentity) -> u64 { + if package.layer_weight_bytes.len() != package.layer_count as usize { + return package.source_model_bytes; + } + package + .layer_weight_bytes + .iter() + .try_fold(0_u64, |total, bytes| total.checked_add(*bytes)) + .unwrap_or(u64::MAX) +} + pub(super) fn format_aggregate_split_capacity_error( model_ref: &str, required_bytes: u64, @@ -679,6 +691,14 @@ mod tests { assert_eq!(plan.stages[1].parameter_bytes, 8 * GIB); } + #[test] + fn aggregate_capacity_uses_exact_runtime_layer_weights() { + let mut package = package(2, 100); + package.layer_weight_bytes = vec![20, 30]; + + assert_eq!(required_split_weight_bytes(&package), 50); + } + #[test] fn resource_planner_prefers_lower_tpot_stage_count_from_participant_rtt() { let participants = vec![ diff --git a/crates/mesh-llm-host-runtime/src/runtime_data/mod.rs b/crates/mesh-llm-host-runtime/src/runtime_data/mod.rs index 6d5f5faaf0..d7dadefb33 100644 --- a/crates/mesh-llm-host-runtime/src/runtime_data/mod.rs +++ b/crates/mesh-llm-host-runtime/src/runtime_data/mod.rs @@ -533,6 +533,7 @@ pub(crate) mod tests { artifact_transfer_supported: false, stage_protocol_generation_supported: false, stage_status_list_supported: false, + mlx_stage_supported: false, advertised_model_throughput: vec![], display_rtt: None, selected_path: None, @@ -668,6 +669,7 @@ pub(crate) mod tests { artifact_transfer_supported: false, stage_protocol_generation_supported: false, stage_status_list_supported: false, + mlx_stage_supported: false, advertised_model_throughput: vec![crate::network::metrics::ModelThroughputHint { model_name: "Qwen/Qwen3-Coder".into(), avg_tokens_per_second_milli: 13_400, diff --git a/crates/mesh-llm-host-runtime/src/system/mod.rs b/crates/mesh-llm-host-runtime/src/system/mod.rs index a9a55ce94a..a1a87979da 100644 --- a/crates/mesh-llm-host-runtime/src/system/mod.rs +++ b/crates/mesh-llm-host-runtime/src/system/mod.rs @@ -1,4 +1,3 @@ -#[cfg(feature = "dynamic-native-runtime")] pub(crate) mod native_runtime; pub(crate) mod native_runtime_install; diff --git a/crates/mesh-llm-host-runtime/src/system/native_runtime.rs b/crates/mesh-llm-host-runtime/src/system/native_runtime.rs index 47b9aad7a8..76e49efd7f 100644 --- a/crates/mesh-llm-host-runtime/src/system/native_runtime.rs +++ b/crates/mesh-llm-host-runtime/src/system/native_runtime.rs @@ -739,7 +739,10 @@ mod dynamic { #[cfg(feature = "dynamic-native-runtime")] pub(crate) use dynamic::*; -#[cfg(not(feature = "dynamic-native-runtime"))] -pub(crate) fn try_load_installed_native_runtime() -> anyhow::Result> { - Ok(None) +pub(crate) fn ensure_native_runtime_loaded() -> anyhow::Result<()> { + anyhow::ensure!( + skippy_runtime::native_runtime_loaded(), + "the Skippy/llama backend requires a compatible MeshLLM native runtime; run `mesh-llm runtime install` or use a SafeTensors model with the MLX backend" + ); + Ok(()) } diff --git a/crates/model-hf/src/safetensors_stage/layout.rs b/crates/model-hf/src/safetensors_stage/layout.rs index 39f1ae21c5..3e97dc353a 100644 --- a/crates/model-hf/src/safetensors_stage/layout.rs +++ b/crates/model-hf/src/safetensors_stage/layout.rs @@ -10,8 +10,8 @@ use sha2::{Digest, Sha256}; use super::{ http::RemoteSource, types::{ - ByteRange, PreparedStage, SafetensorsShardPlan, SafetensorsSourceShard, - SafetensorsStagePlan, SafetensorsStageRequest, SelectedTensor, + ByteRange, PreparedStage, SafetensorsCheckpointDescriptor, SafetensorsShardPlan, + SafetensorsSourceShard, SafetensorsStagePlan, SafetensorsStageRequest, SelectedTensor, }, }; @@ -28,6 +28,15 @@ enum StageModelFamily { struct RawStageModelConfig { model_type: String, num_hidden_layers: u32, + hidden_size: u32, + #[serde(default)] + max_position_embeddings: u32, + #[serde(default)] + num_attention_heads: u32, + #[serde(default)] + num_key_value_heads: u32, + #[serde(default)] + head_dim: u32, } #[derive(Clone, Copy, Debug)] @@ -36,6 +45,160 @@ struct StageModelLayout { num_hidden_layers: u32, } +pub(crate) fn describe( + remote: &RemoteSource, + repo: &str, + revision: &str, +) -> Result { + let probe = SafetensorsStageRequest { + repo: repo.to_string(), + revision: revision.to_string(), + layer_start: 0, + layer_end: 1, + include_prefixes: Vec::new(), + } + .normalized()?; + let config_url = remote.url(repo, revision, "config.json")?; + let config_file = remote + .small_file(config_url, MAX_INDEX_BYTES) + .context("download SafeTensors model config")?; + let config = parse_raw_stage_model_config(&config_file.bytes)?; + ensure!( + config.model_type == "llama", + "automatic distributed MLX planning currently supports model_type=llama, got {:?}", + config.model_type + ); + validate_planning_config(&config)?; + + let full = prepare( + remote, + &SafetensorsStageRequest { + layer_end: config.num_hidden_layers, + ..probe + }, + )?; + let dense_tensor_bytes = full + .plan + .total_model_tensor_bytes + .context("SafeTensors checkpoint does not declare total tensor bytes")?; + let stage_layout = StageModelLayout { + family: StageModelFamily::Llama, + num_hidden_layers: config.num_hidden_layers, + }; + let estimated_affine4_weight_bytes = + checked_sum(full.tensors.iter().map(estimated_affine4_tensor_bytes))?; + let estimated_affine4_layer_bytes = estimated_affine4_layer_bytes(stage_layout, &full.tensors)?; + let kv_heads = if config.num_key_value_heads > 0 { + config.num_key_value_heads + } else { + config.num_attention_heads + }; + let head_dim = if config.head_dim > 0 { + config.head_dim + } else { + config.hidden_size / config.num_attention_heads + }; + let kv_bytes_per_token_bf16 = u64::from(config.num_hidden_layers) + .checked_mul(u64::from(kv_heads)) + .and_then(|bytes| bytes.checked_mul(u64::from(head_dim))) + .and_then(|bytes| bytes.checked_mul(2)) + .and_then(|bytes| bytes.checked_mul(2)) + .context("SafeTensors KV byte estimate overflow")?; + + Ok(SafetensorsCheckpointDescriptor { + checkpoint_sha256: full.checkpoint_sha256, + repo: repo.to_string(), + revision: revision.to_string(), + model_type: config.model_type, + layer_count: config.num_hidden_layers, + hidden_size: config.hidden_size, + native_context_length: if config.max_position_embeddings == 0 { + 2_048 + } else { + config.max_position_embeddings + }, + dense_tensor_bytes, + estimated_affine4_weight_bytes, + estimated_affine4_layer_bytes, + kv_bytes_per_token_bf16, + }) +} + +fn estimated_affine4_layer_bytes( + layout: StageModelLayout, + tensors: &[SelectedTensor], +) -> Result> { + let mut layer_bytes = vec![0_u64; layout.num_hidden_layers as usize]; + let final_layer = layer_bytes + .len() + .checked_sub(1) + .context("SafeTensors model has no layers")?; + for tensor in tensors { + let bytes = estimated_affine4_tensor_bytes(tensor); + if let Some(layer) = layout.layer_index(&tensor.name) { + add_estimated_bytes(&mut layer_bytes[layer as usize], bytes)?; + } else if tensor.name.starts_with(layout.embedding_prefix()) { + add_estimated_bytes(&mut layer_bytes[0], bytes)?; + if layout.family == StageModelFamily::Llama && final_layer != 0 { + // The current Llama final-stage selection reloads embeddings + // for tied readout compatibility, even when lm_head is also + // present. Charge the planner for what the stage really loads. + add_estimated_bytes(&mut layer_bytes[final_layer], bytes)?; + } + } else { + // Final norm/readout and any unclassified boundary tensors are + // selected only by the final stage. + add_estimated_bytes(&mut layer_bytes[final_layer], bytes)?; + } + } + Ok(layer_bytes) +} + +fn estimated_affine4_tensor_bytes(tensor: &SelectedTensor) -> u64 { + let shape = &tensor.header.shape; + let is_dense_float_weight = tensor.name.ends_with(".weight") + && shape.len() == 2 + && shape[1].is_multiple_of(64) + && matches!(tensor.header.dtype.as_str(), "F16" | "BF16" | "F32"); + if !is_dense_float_weight { + return tensor.source_range.len(); + } + let elements = shape.iter().copied().fold(1_u64, u64::saturating_mul); + // Packed 4-bit values plus a deliberately conservative allowance for + // group-64 scales, biases, alignment, and SafeTensors metadata. + elements.saturating_mul(3).div_ceil(4) +} + +fn add_estimated_bytes(total: &mut u64, bytes: u64) -> Result<()> { + *total = total + .checked_add(bytes) + .context("SafeTensors affine-4 layer byte estimate overflow")?; + Ok(()) +} + +fn validate_planning_config(config: &RawStageModelConfig) -> Result<()> { + ensure!( + config.hidden_size > 0, + "SafeTensors hidden_size must be positive" + ); + ensure!( + config.num_hidden_layers > 0, + "SafeTensors num_hidden_layers must be positive" + ); + ensure!( + config.num_attention_heads > 0, + "SafeTensors num_attention_heads must be positive" + ); + ensure!( + config + .hidden_size + .is_multiple_of(config.num_attention_heads) + || config.head_dim > 0, + "SafeTensors hidden_size is not divisible by num_attention_heads and head_dim is absent" + ); + Ok(()) +} + impl StageModelLayout { fn layer_index(self, name: &str) -> Option { self.layer_prefixes() @@ -174,16 +337,7 @@ pub(crate) fn prepare( } fn parse_stage_model_layout(bytes: &[u8]) -> Result { - let config: RawStageModelConfig = match serde_json::from_slice(bytes) { - Ok(config) => config, - Err(strict_error) => { - let text = - std::str::from_utf8(bytes).context("SafeTensors model config is not UTF-8")?; - json5::from_str(text).with_context(|| { - format!("parse SafeTensors model config as strict JSON ({strict_error}) or JSON5") - })? - } - }; + let config = parse_raw_stage_model_config(bytes)?; ensure!( config.num_hidden_layers > 0, "SafeTensors model num_hidden_layers must be non-zero" @@ -201,6 +355,20 @@ fn parse_stage_model_layout(bytes: &[u8]) -> Result { }) } +fn parse_raw_stage_model_config(bytes: &[u8]) -> Result { + let config: RawStageModelConfig = match serde_json::from_slice(bytes) { + Ok(config) => config, + Err(strict_error) => { + let text = + std::str::from_utf8(bytes).context("SafeTensors model config is not UTF-8")?; + json5::from_str(text).with_context(|| { + format!("parse SafeTensors model config as strict JSON ({strict_error}) or JSON5") + })? + } + }; + Ok(config) +} + fn validate_layer_range( request: &SafetensorsStageRequest, config: &StageModelLayout, @@ -595,6 +763,22 @@ fn checked_sum(values: impl IntoIterator) -> Result { mod tests { use super::*; + fn planning_tensor(name: &str, shape: Vec, source_bytes: u64) -> SelectedTensor { + SelectedTensor { + name: name.to_string(), + source_file: "model.safetensors".to_string(), + source_range: ByteRange { + start: 0, + end_exclusive: source_bytes, + }, + header: TensorHeader { + dtype: "BF16".to_string(), + shape, + data_offsets: [0, source_bytes], + }, + } + } + #[test] fn recognizes_family_specific_layer_paths() { let llama = parse_stage_model_layout( @@ -676,6 +860,26 @@ mod tests { ); } + #[test] + fn affine4_layer_estimates_charge_boundary_tensors_to_loading_stages() { + let layout = StageModelLayout { + family: StageModelFamily::Llama, + num_hidden_layers: 2, + }; + let tensors = vec![ + planning_tensor("model.embed_tokens.weight", vec![128, 64], 16_384), + planning_tensor("model.layers.0.mlp.weight", vec![64, 64], 8_192), + planning_tensor("model.layers.1.mlp.weight", vec![64, 64], 8_192), + planning_tensor("model.norm.weight", vec![64], 128), + planning_tensor("lm_head.weight", vec![128, 64], 16_384), + ]; + + let estimates = estimated_affine4_layer_bytes(layout, &tensors).unwrap(); + + assert_eq!(estimates, vec![9_216, 15_488]); + assert!(estimates[1] > estimates[0]); + } + #[test] fn assigns_embedding_and_readout_tensors_to_final_llama_stage() { let mut request = SafetensorsStageRequest { diff --git a/crates/model-hf/src/safetensors_stage/materialize.rs b/crates/model-hf/src/safetensors_stage/materialize.rs index 5aaa4b66ca..71b09e097d 100644 --- a/crates/model-hf/src/safetensors_stage/materialize.rs +++ b/crates/model-hf/src/safetensors_stage/materialize.rs @@ -16,7 +16,8 @@ use super::{ layout, locking::CacheKeyLock, types::{ - MANIFEST_SCHEMA_VERSION, PreparedStage, SafetensorsSourceShard, SafetensorsStageArtifact, + MANIFEST_SCHEMA_VERSION, PreparedSafetensorsCheckpoint, PreparedStage, + SafetensorsCheckpointDescriptor, SafetensorsSourceShard, SafetensorsStageArtifact, SafetensorsStageManifest, SafetensorsStagePlan, SafetensorsStageRequest, SelectedTensor, }, }; @@ -25,9 +26,19 @@ const MODEL_FILE: &str = "model.safetensors"; const CONFIG_FILE: &str = "config.json"; const PLAN_FILE: &str = "stage-plan.json"; const MANIFEST_FILE: &str = "stage-manifest.json"; +pub const CHECKPOINT_DESCRIPTOR_FILE: &str = "checkpoint-descriptor.json"; const MAX_LOCAL_HEADER_BYTES: u64 = 256 * 1024 * 1024; static TEMP_SEQUENCE: AtomicU64 = AtomicU64::new(0); +pub fn read_checkpoint_descriptor(path: &Path) -> Result { + read_json(&path.join(CHECKPOINT_DESCRIPTOR_FILE)).with_context(|| { + format!( + "read prepared SafeTensors checkpoint descriptor from {}", + path.display() + ) + }) +} + pub struct SafetensorsStageMaterializer { pub(crate) remote: RemoteSource, pub(crate) cache_root: PathBuf, @@ -55,6 +66,139 @@ impl SafetensorsStageMaterializer { Ok(layout::prepare(&self.remote, &request)?.plan) } + /// Fetches config/index/header metadata only. Tensor payload ranges are + /// planned but not downloaded. + pub fn describe_checkpoint( + &self, + repo: &str, + revision: &str, + ) -> Result { + layout::describe(&self.remote, repo, revision) + } + + /// Caches only the tokenizer/config/chat-template files needed by the + /// distributed stage-0 frontend. No model tensor payload is downloaded. + pub fn prepare_checkpoint( + &self, + repo: &str, + revision: &str, + ) -> Result { + let descriptor = self.describe_checkpoint(repo, revision)?; + let cache_key = format!( + "checkpoint-{:x}", + Sha256::digest(format!("{repo}@{revision}").as_bytes()) + ); + let sidecar_root = self.cache_root.join("checkpoint-sidecars"); + fs::create_dir_all(&sidecar_root).with_context(|| { + format!( + "create SafeTensors checkpoint sidecar cache {}", + sidecar_root.display() + ) + })?; + let _cache_lock = CacheKeyLock::acquire(&self.cache_root, &cache_key)?; + let destination = sidecar_root.join(&cache_key); + if checkpoint_sidecars_match(&destination, &descriptor) { + return Ok(PreparedSafetensorsCheckpoint { + path: destination, + descriptor, + }); + } + remove_stale_partials(&sidecar_root, &cache_key)?; + let temporary = temporary_path(&sidecar_root, &cache_key); + fs::create_dir(&temporary).with_context(|| { + format!( + "create temporary checkpoint sidecar cache {}", + temporary.display() + ) + })?; + let write_result = self.write_checkpoint_sidecars(repo, revision, &temporary, &descriptor); + if let Err(error) = write_result { + let _ = fs::remove_dir_all(&temporary); + return Err(error); + } + let quarantine = destination + .exists() + .then(|| temporary_path(&sidecar_root, &format!("{cache_key}.stale"))); + if let Some(quarantine) = quarantine.as_ref() { + fs::rename(&destination, quarantine).with_context(|| { + format!( + "quarantine incomplete checkpoint sidecar cache {}", + destination.display() + ) + })?; + } + if let Err(error) = fs::rename(&temporary, &destination) { + if let Some(quarantine) = quarantine.as_ref() { + let _ = fs::rename(quarantine, &destination); + } + return Err(error).context("publish checkpoint sidecar cache"); + } + if let Some(quarantine) = quarantine { + let _ = fs::remove_dir_all(quarantine); + } + sync_directory(&sidecar_root)?; + Ok(PreparedSafetensorsCheckpoint { + path: destination, + descriptor, + }) + } + + fn write_checkpoint_sidecars( + &self, + repo: &str, + revision: &str, + destination: &Path, + descriptor: &SafetensorsCheckpointDescriptor, + ) -> Result<()> { + self.cache_required_sidecar(repo, revision, "config.json", destination)?; + self.cache_required_sidecar(repo, revision, "tokenizer.json", destination)?; + for file in [ + "tokenizer_config.json", + "chat_template.jinja", + "chat_template.json", + "generation_config.json", + "special_tokens_map.json", + ] { + self.cache_optional_sidecar(repo, revision, file, destination)?; + } + write_json(destination.join(CHECKPOINT_DESCRIPTOR_FILE), descriptor)?; + sync_directory(destination) + } + + fn cache_required_sidecar( + &self, + repo: &str, + revision: &str, + file: &str, + destination: &Path, + ) -> Result<()> { + let remote = self + .remote + .small_file( + self.remote.url(repo, revision, file)?, + MAX_LOCAL_HEADER_BYTES, + ) + .with_context(|| format!("download required checkpoint sidecar {file}"))?; + write_synced(destination.join(file), &remote.bytes) + } + + fn cache_optional_sidecar( + &self, + repo: &str, + revision: &str, + file: &str, + destination: &Path, + ) -> Result<()> { + let Some(remote) = self.remote.optional_small_file( + self.remote.url(repo, revision, file)?, + MAX_LOCAL_HEADER_BYTES, + )? + else { + return Ok(()); + }; + write_synced(destination.join(file), &remote.bytes) + } + pub fn materialize( &self, request: SafetensorsStageRequest, @@ -289,15 +433,28 @@ impl SafetensorsStageMaterializer { } fn temporary_path(&self, cache_key: &str) -> PathBuf { - let sequence = TEMP_SEQUENCE.fetch_add(1, Ordering::Relaxed); - self.cache_root.join(format!( - ".{cache_key}.{}.{}.partial", - std::process::id(), - sequence - )) + temporary_path(&self.cache_root, cache_key) } } +fn checkpoint_sidecars_match( + destination: &Path, + expected: &SafetensorsCheckpointDescriptor, +) -> bool { + destination.join(CONFIG_FILE).is_file() + && destination.join("tokenizer.json").is_file() + && read_checkpoint_descriptor(destination).is_ok_and(|actual| actual == *expected) +} + +fn temporary_path(cache_root: &Path, cache_key: &str) -> PathBuf { + let sequence = TEMP_SEQUENCE.fetch_add(1, Ordering::Relaxed); + cache_root.join(format!( + ".{cache_key}.{}.{}.partial", + std::process::id(), + sequence + )) +} + struct MaterializationSpan { source_file: String, start: u64, diff --git a/crates/model-hf/src/safetensors_stage/mod.rs b/crates/model-hf/src/safetensors_stage/mod.rs index 687d36d39a..7ba254f73c 100644 --- a/crates/model-hf/src/safetensors_stage/mod.rs +++ b/crates/model-hf/src/safetensors_stage/mod.rs @@ -5,10 +5,13 @@ mod materialize; mod tensor_stream; mod types; -pub use materialize::SafetensorsStageMaterializer; +pub use materialize::{ + CHECKPOINT_DESCRIPTOR_FILE, SafetensorsStageMaterializer, read_checkpoint_descriptor, +}; pub use tensor_stream::SafetensorsStageTensorVisit; pub use types::{ - ByteRange, SafetensorsShardPlan, SafetensorsSourceShard, SafetensorsStageArtifact, + ByteRange, PreparedSafetensorsCheckpoint, SafetensorsCheckpointDescriptor, + SafetensorsShardPlan, SafetensorsSourceShard, SafetensorsStageArtifact, SafetensorsStageManifest, SafetensorsStagePlan, SafetensorsStageRequest, SafetensorsStageTensorFile, SafetensorsStageTensorVisitReport, }; diff --git a/crates/model-hf/src/safetensors_stage/types.rs b/crates/model-hf/src/safetensors_stage/types.rs index 5aab3f08d3..6ff6883dfb 100644 --- a/crates/model-hf/src/safetensors_stage/types.rs +++ b/crates/model-hf/src/safetensors_stage/types.rs @@ -73,6 +73,33 @@ pub struct SafetensorsStagePlan { pub shards: Vec, } +/// Metadata-only description used to plan a distributed MLX topology without +/// downloading checkpoint tensor payloads. +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub struct SafetensorsCheckpointDescriptor { + pub checkpoint_sha256: String, + pub repo: String, + pub revision: String, + pub model_type: String, + pub layer_count: u32, + pub hidden_size: u32, + pub native_context_length: u32, + pub dense_tensor_bytes: u64, + pub estimated_affine4_weight_bytes: u64, + /// Conservative affine-4 runtime bytes attributed to each transformer + /// layer. Boundary tensors are charged to the stages that load them. + #[serde(default)] + pub estimated_affine4_layer_bytes: Vec, + /// K + V cache bytes per token across all layers at BF16 precision. + pub kv_bytes_per_token_bf16: u64, +} + +#[derive(Clone, Debug)] +pub struct PreparedSafetensorsCheckpoint { + pub path: PathBuf, + pub descriptor: SafetensorsCheckpointDescriptor, +} + #[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] pub struct SafetensorsShardPlan { pub file: String, diff --git a/crates/skippy-engine-mlx/Cargo.toml b/crates/skippy-engine-mlx/Cargo.toml index bcc25dbff9..34c03b0b56 100644 --- a/crates/skippy-engine-mlx/Cargo.toml +++ b/crates/skippy-engine-mlx/Cargo.toml @@ -46,6 +46,7 @@ mlx = [ "dep:safetensors", "dep:safemlx", "dep:safemlx-lm", + "dep:safemlx-lm-utils", "dep:serde", "dep:sha2", "dep:skippy-metrics", @@ -92,6 +93,9 @@ safemlx = { git = "https://github.com/jbg/safemlx", rev = "4e53c5e", default-fea "safetensors", ], optional = true } safemlx-lm = { git = "https://github.com/jbg/safemlx", rev = "4e53c5e", optional = true } +safemlx-lm-utils = { git = "https://github.com/jbg/safemlx", rev = "4e53c5e", default-features = false, features = [ + "onig", +], optional = true } # Standalone tokenizer for incremental (streaming) detokenization without # borrowing the model. `onig` matches safemlx-lm so the build dedupes. tokenizers = { version = "0.23", default-features = false, features = [ diff --git a/crates/skippy-engine-mlx/SERVE_INTEGRATION_STATUS.md b/crates/skippy-engine-mlx/SERVE_INTEGRATION_STATUS.md index 26f0db0fdc..b4fda63c6b 100644 --- a/crates/skippy-engine-mlx/SERVE_INTEGRATION_STATUS.md +++ b/crates/skippy-engine-mlx/SERVE_INTEGRATION_STATUS.md @@ -1,113 +1,118 @@ -# MLX serve-integration — WIP status & blocker - -This branch (`micn/mlx-serve-wiring`) stacks the **`mesh-llm serve` integration** -on top of the standalone MLX engine crate from PR #1009 (`micn/mlx-redux`). - -Goal: on a Mac, `mesh-llm serve --model ` routes to the MLX -(Metal) engine and the model appears in `/v1/models` — with non-macOS / no-feature -builds byte-for-byte unaffected. - -**Status: functionally complete and verified locally, but blocked by one repo -publish invariant. Needs a maintainer decision before it can merge.** Captured -here so the work + reasoning aren't lost. - -## What works (verified on Apple Silicon this session) - -- `skippy-engine-mlx` is now a real workspace member (removed its private - `[workspace]`; added to root `members`, both CI crate-list scripts, and it - stays out of `default-members`). -- `mlx` feature added to `mesh-llm-host-runtime` and `mesh-llm`, wired as a - **macOS-target-gated optional dep** (`[target.'cfg(target_os = "macos")'.dependencies]`). -- `crates/mesh-llm-host-runtime/src/inference/mlx.rs`: `MlxModelHandle` + - `MlxHttpHandle` serving over the real `openai-frontend::router_for` + axum, with - graceful shutdown. -- `LocalRuntimeBackendHandle::Mlx` variant + all match arms (cfg-gated). -- `start_runtime_mlx_model` + `is_safetensors_model_path` routing: safetensors - models branch to MLX *before* the GGUF planning path. -- **Build gate PASSED:** `cargo build -p mesh-llm --features mlx` → exit 0. -- **Don't-break-default PASSED:** no-feature `cargo check`/clippy clean; - `cargo tree` confirms safemlx is absent unless `--features mlx`. -- fmt clean; clippy clean both with and without `--features mlx` - (boxed the `Skippy` enum variant to satisfy `large_enum_variant`). -- `xtask repo-consistency ci-crate-lists` PASSES. - -## Key finding: the two-native-stack link collision (RESOLVED) - -Linking MLX statically alongside the patched llama.cpp fails with: +# MLX serve integration status -``` -ld64.lld: error: duplicate symbol: gguf_get_key - >>> defined in .../gguflib-src/gguflib.c (MLX's vendored GGUF parser) - >>> defined in libskippy_ffi...(gguf.cpp.o) (patched llama.cpp) -``` +## Draft status -MLX statically links antirez's `gguflib` (via `MLX_BUILD_GGUF=ON`, and -`safemlx-sys` link-directs `gguflib` unconditionally); the patched llama.cpp -exports the same C symbols. Two independent GGUF parsers → duplicate symbol. +The `micn/mlx-redux` draft now has two connected serving proofs on Apple +Silicon. Both use mesh-llm's ordinary model/runtime and OpenAI surfaces rather +than a standalone demonstration server. -**Fix:** the `mlx` feature now implies `dynamic-native-runtime`. That loads the -llama.cpp runtime as a dylib (as release builds already do), so its GGUF symbols -live in a separate link unit and don't collide with MLX's static `gguflib`. -`cargo build -p mesh-llm --features mlx` alone links clean after this. +1. Whole-model serving resolves a Hugging Face SafeTensors repository through + the normal `mesh-llm serve --model ...` path, loads it with MLX, applies + automatic affine 4-bit/group-64 quantization while loading eligible dense + weights, and + serves `/v1/models` plus streaming and non-streaming chat completions. +2. Explicit `--split` serving resolves only metadata at the coordinator, + advertises an additive MLX stage capability, plans the ordinary stage + topology, and makes each stage fetch and derive only its assigned tensor + ranges. The coordinator's OpenAI frontend drives the stage chain over the + existing Skippy binary activation transport. -## BLOCKER: publish invariant (needs a maintainer call) +This is a substantial draft checkpoint, not production support for arbitrary +SafeTensors models. -``` -cargo run -p xtask -- repo-consistency release-targets - error: mesh-llm-host-runtime: publishable crate depends on - non-publishable workspace crate `skippy-engine-mlx` +## Verified whole-model proof + +`HuggingFaceTB/SmolLM2-135M-Instruct` was started through the shipped command +shape: + +```bash +mesh-llm serve --model HuggingFaceTB/SmolLM2-135M-Instruct ``` -Root-cause chain: -- `skippy-engine-mlx` git-pins `safemlx` / `safemlx-lm` to a public commit of - `github.com/jbg/safemlx` (crates.io's published safemlx produces gibberish for - dense models — verified; only the git commit serves correctly). -- crates.io forbids git dependencies → the crate must be `publish = false`. -- `mesh-llm-host-runtime` is a **published** SDK crate, and the repo invariant - (enforced by `xtask release-targets`, reflecting a real `cargo publish` - constraint) forbids a publishable crate from depending on a `publish = false` - one — even an optional, target-gated dep. - -This is **not solved elsewhere in the repo**: the other `publish = false` crates -(`skippy-quantize`, `mesh-llm-commands`) are only consumed by the non-published -binary crate `mesh-llm`, never by a published library crate. - -### Options - -- **A. `[patch.crates-io]` redirect + make the crate publishable.** Follows the - existing `hf-hub` precedent (a published-crate dep already redirected to a fork - via `[patch.crates-io]` in the root manifest). Add `skippy-engine-mlx` to - `scripts/publish-crates.sh`. Caveat that differs from hf-hub: safemlx's - *published* version is known-broken, so a hypothetically-published - `skippy-engine-mlx` with `mlx` on would reference broken upstream — acceptable - only because the feature is off by default and explicitly a stopgap until - safemlx cuts a working release. -- **B. Hold the host-runtime wiring.** Ship the standalone crate (PR #1009) as - is; keep this branch as the ready-to-go integration until safemlx releases a - working crates.io version, then flip git-pin → version-pin and merge. Most - honest; doesn't deliver "serve just works" yet. -- **C. Loosen the xtask invariant** to exempt optional/target deps. Not - recommended: it would allow a manifest that genuinely cannot `cargo publish`. - -**Recommendation: B for now** (this branch is the parked, working integration), -moving to **A** if/when we want it live before safemlx releases. The real -unlock is a working safemlx crates.io release, after which this is a trivial -git-pin → version-pin swap and the invariant is satisfied automatically. - -## How to reproduce / verify +The normal resolver downloaded the complete SafeTensors checkpoint, selected +backend `mlx`, and served model listing, non-streaming chat, and SSE chat. This +proves useful single-node MLX serving is part of the work. It intentionally +does not prove partial downloads: a whole-model server needs all model weights. + +The integrated loader now quantizes eligible, unquantized dense source tensors +to affine 4-bit as they load. Inkling, Nemotron-H, and checkpoints already +declaring a quantized representation retain their native representation. The +earlier solo-serving measurements showed that this +load-time representation has the same steady-state generation speed as loading +an equivalent pre-quantized artifact; see `../../spikes/mlx-solo/FINDINGS.md`. + +The 135M model is adequate as a serving and protocol oracle, but its weak agent +output is not evidence of Goose-quality model behavior. Larger single-node +models still need quality, memory-high-water, and agent-harness measurements. + +## Verified mesh split proof + +A two-host explicit split of the same 30-layer model completed through the +normal OpenAI frontend: + +| Stage | Layers | Derived affine-4 artifact | +| --- | ---: | ---: | +| coordinator stage | `0..29` | about 70 MB | +| remote final stage | `29..30` | about 17 MB | + +The remote host never stored the complete roughly 269 MB source checkpoint. +It fetched metadata plus the tensor ranges selected for its layer and final +boundary tensors. A non-streaming request traversed both real stages and +returned coherent text; SSE also completed with `[DONE]`. + +The apparently uneven artifacts are expected: embeddings and readout tensors +are much larger than a typical transformer block. The planner now derives +conservative per-layer affine-4 estimates from SafeTensors headers and charges +boundary tensors to every stage that loads them, rather than dividing total +bytes evenly by layer count. + +## Safety and lifecycle behavior + +- Only immutable 40-character Hugging Face revisions identify split tensor + ranges. +- Config, index, and shard headers are fetched before tensor payloads; exact + HTTP ranges are identity-checked with strong ETags. +- Coordinator sidecars and derived stage artifacts are lock-protected and + atomically published. +- Stage `Prepare` may derive a missing entry; `Load` validates and consumes an + already prepared entry. +- Connections track every touched session. EOF, transport error, timeout, or + normal stop resets local state and propagates `Stop` through downstream + stages. +- Stage and generation I/O have bounded timeouts. +- MLX frontend bind succeeds before the model is advertised as ready. +- The capability advertisement is additive and is emitted only by an + Apple-Silicon build containing the MLX feature. Older peers safely treat the + absent field as unsupported. + +## Current boundaries + +- Partial-download integration is currently entered by explicit `--split`; + automatic startup still follows the existing whole-model resolver. +- Integrated partial generation supports dense `model_type=llama` stages. +- Generation is serialized to one MLX lane and currently uses greedy sampling. +- Automatic affine-4 is the current default for eligible dense checkpoints, + not yet a general hardware/quality policy surface. +- Cache capacity and eviction are not yet owned by this integration. +- Frontier families require their own stage semantics. Nemotron-H has + metadata/range planning and a one-layer execution proof, but not a complete + hybrid-model topology. Inkling has whole-model support in the pinned safemlx + runtime and compelling exact-range storage evidence, but no integrated + partial-stage executor yet. + +## Build and focused verification + +Use the MLX recipes because they provide the full Xcode Metal toolchain and +copy the required `mlx.metallib` beside the executable: ```bash -export DEVELOPER_DIR=/Applications/Xcode.app/Contents/Developer # Metal toolchain -cargo build -p mesh-llm --features mlx # exit 0 (links clean) -cargo check -p mesh-llm-host-runtime # no-feature: clean, no safemlx -cargo run -p xtask -- repo-consistency ci-crate-lists # PASS -cargo run -p xtask -- repo-consistency release-targets # FAILS (the blocker) +just mlx-build +just mlx-release-build ``` -## Remaining once unblocked +Focused library checks require the same developer directory: -- End-to-end serve test: `mesh-llm serve --model ` → - confirm `/v1/models` + `/v1/chat/completions`. -- Fold this status into `WIRING.md` (note the `dynamic-native-runtime` - requirement and the publish resolution chosen). -- Rebase onto `main` and open/land the real PR. +```bash +DEVELOPER_DIR=/Applications/Xcode.app/Contents/Developer \ + cargo test -p skippy-engine-mlx --features mlx --lib +``` diff --git a/crates/skippy-engine-mlx/STAGED_EXECUTION.md b/crates/skippy-engine-mlx/STAGED_EXECUTION.md index 4a9a507d83..6bb332e668 100644 --- a/crates/skippy-engine-mlx/STAGED_EXECUTION.md +++ b/crates/skippy-engine-mlx/STAGED_EXECUTION.md @@ -6,8 +6,8 @@ Dense Llama-family MLX stages now run as separate OS processes from partial SafeTensors artifacts and communicate over Skippy's existing binary stage wire. The first proof uses `HuggingFaceTB/SmolLM2-135M-Instruct` split at layer 15. -This is a production-shaped bridge with an explicit host control path, not yet -the default automatic mesh launch path: +This is now integrated with the explicit `mesh-llm serve --split` launch path, +while automatic split selection remains future work: - `skippy-engine` owns the engine-neutral `StageEngine` contract and residual buffer descriptors. @@ -25,11 +25,21 @@ the default automatic mesh launch path: `hf-model://org/repo@` reference now derive or reuse a validated quantized stage and start the same engine through the normal host stage-control loop. +- The mesh advertises an additive `backend-mlx` capability, plans exact + SafeTensors ranges through its ordinary topology, and exposes the chain via + the normal OpenAI frontend. No process in the proof has access to the complete checkpoint. The tokenizer and config files are small shared metadata; tensor data comes only from that process's `model.safetensors`. +The same branch also includes a complementary single-node proof. Ordinary +`mesh-llm serve --model HuggingFaceTB/SmolLM2-135M-Instruct` resolves the full +SafeTensors checkpoint, automatically quantizes eligible unquantized dense +tensors to affine-4 during MLX load while preserving frontier/pre-quantized +representations, and serves normal and streaming OpenAI chat. See +`SERVE_INTEGRATION_STATUS.md` for the integrated status and limitations. + ## Verified result On Apple Silicon Metal, using two materialized 155.28 MiB partial files: diff --git a/crates/skippy-engine-mlx/src/backend.rs b/crates/skippy-engine-mlx/src/backend.rs index 16a1bad7fb..4f6dbfed58 100644 --- a/crates/skippy-engine-mlx/src/backend.rs +++ b/crates/skippy-engine-mlx/src/backend.rs @@ -19,16 +19,51 @@ use openai_frontend::common::{FinishReason, Usage, completion_id}; use openai_frontend::errors::OpenAiError; use openai_frontend::models::ModelObject; +use crate::distributed::MlxDistributedEngine; use crate::engine::{ChatTurn, FinishReason as EngineFinish, GenerateRequest, MlxEngine, TokenMsg}; +enum BackendEngine { + Local(MlxEngine), + Distributed(MlxDistributedEngine), +} + +impl BackendEngine { + fn model_id(&self) -> &str { + match self { + Self::Local(engine) => engine.model_id(), + Self::Distributed(engine) => engine.model_id(), + } + } + + fn clamp_max_tokens(&self, requested: Option) -> usize { + match self { + Self::Local(engine) => engine.clamp_max_tokens(requested), + Self::Distributed(engine) => engine.clamp_max_tokens(requested), + } + } + + fn submit(&self, request: GenerateRequest) -> tokio::sync::mpsc::UnboundedReceiver { + match self { + Self::Local(engine) => engine.submit(request), + Self::Distributed(engine) => engine.submit(request), + } + } +} + pub struct MlxBackend { - engine: Arc, + engine: Arc, } impl MlxBackend { pub fn new(engine: MlxEngine) -> Self { Self { - engine: Arc::new(engine), + engine: Arc::new(BackendEngine::Local(engine)), + } + } + + pub fn new_distributed(engine: MlxDistributedEngine) -> Self { + Self { + engine: Arc::new(BackendEngine::Distributed(engine)), } } diff --git a/crates/skippy-engine-mlx/src/bin/mlx-serve.rs b/crates/skippy-engine-mlx/src/bin/mlx-serve.rs index 7a95b69877..02e5c1a71c 100644 --- a/crates/skippy-engine-mlx/src/bin/mlx-serve.rs +++ b/crates/skippy-engine-mlx/src/bin/mlx-serve.rs @@ -11,8 +11,34 @@ mod real { use std::sync::Arc; use anyhow::{Context, Result}; - use clap::Parser; - use skippy_engine_mlx::{MlxBackend, MlxEngine, MlxEngineConfig}; + use clap::{Parser, ValueEnum}; + use skippy_engine_mlx::{MlxBackend, MlxEngine, MlxEngineConfig, MlxWeightQuantization}; + + #[derive(Clone, Copy, Debug, ValueEnum)] + enum WeightQuantization { + Auto, + None, + Affine4, + Affine8, + MxFp4, + } + + impl WeightQuantization { + fn engine(self) -> Option { + match self { + Self::Auto | Self::None => None, + Self::Affine4 => Some(MlxWeightQuantization::Affine { + group_size: 64, + bits: 4, + }), + Self::Affine8 => Some(MlxWeightQuantization::Affine { + group_size: 64, + bits: 8, + }), + Self::MxFp4 => Some(MlxWeightQuantization::MxFp4), + } + } + } #[derive(Parser, Debug)] #[command(about = "Serve an MLX safetensors model over the mesh-llm OpenAI frontend")] @@ -31,6 +57,10 @@ mod real { #[arg(long, default_value_t = 4096)] max_tokens_cap: usize, + /// Weight representation for dense checkpoints. + #[arg(long, value_enum, default_value = "auto")] + weight_quantization: WeightQuantization, + #[arg(long, default_value = "127.0.0.1:11434")] bind: String, } @@ -51,11 +81,18 @@ mod real { .unwrap_or_else(|| "mlx-model".to_string()) }); + let weight_quantization = match cli.weight_quantization { + WeightQuantization::Auto => { + skippy_engine_mlx::automatic_weight_quantization(&cli.model)? + } + explicit => explicit.engine(), + }; let config = MlxEngineConfig { model_dir: cli.model.clone(), model_id: model_id.clone(), default_max_tokens: cli.default_max_tokens, max_tokens_cap: cli.max_tokens_cap, + weight_quantization, }; tracing::info!("loading MLX model from {} ...", cli.model.display()); diff --git a/crates/skippy-engine-mlx/src/distributed.rs b/crates/skippy-engine-mlx/src/distributed.rs new file mode 100644 index 0000000000..2e12e8a8f2 --- /dev/null +++ b/crates/skippy-engine-mlx/src/distributed.rs @@ -0,0 +1,559 @@ +//! OpenAI generation driver for a mesh-managed chain of MLX layer stages. +//! +//! The coordinator owns tokenizer/chat-template sidecars only. Model weights +//! remain in the stage engines selected by the mesh topology. Generation uses +//! the existing Skippy binary stage wire: one final prefill followed by greedy +//! decode frames, with a stop frame resetting every stage in the chain. + +use std::fs; +use std::io::Write; +use std::net::{SocketAddr, TcpStream}; +use std::path::{Path, PathBuf}; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::thread; +use std::time::Duration; + +use anyhow::{Context, Result, anyhow, ensure}; +use safemlx_lm_utils::tokenizer::{Tokenizer as ChatTokenizer, load_model_chat_template_from_file}; +use serde_json::{Map, Value, json}; +use skippy_protocol::binary::{ + StageReply, StageStateHeader, StageWireMessage, WireActivationDType, WireMessageKind, + WireReplyKind, recv_ready, recv_reply, write_stage_message, +}; +use tokio::sync::mpsc; + +use crate::engine::{FinishReason, GenerateRequest, TokenMsg}; + +static NEXT_REQUEST_ID: AtomicU64 = AtomicU64::new(1); + +#[derive(Clone, Debug)] +pub struct MlxDistributedEngineConfig { + pub model_dir: PathBuf, + pub model_id: String, + pub stage_addr: SocketAddr, + pub wire_dtype: WireActivationDType, + pub default_max_tokens: usize, + pub max_tokens_cap: usize, + pub context_tokens: usize, +} + +struct Job { + request: GenerateRequest, + reply: mpsc::UnboundedSender, +} + +pub struct MlxDistributedEngine { + jobs: mpsc::UnboundedSender, + config: MlxDistributedEngineConfig, +} + +impl MlxDistributedEngine { + pub fn spawn(config: MlxDistributedEngineConfig) -> Result { + let (jobs, mut job_rx) = mpsc::unbounded_channel::(); + let worker_config = config.clone(); + let (ready_tx, ready_rx) = std::sync::mpsc::channel(); + thread::Builder::new() + .name("mlx-distributed-engine".to_string()) + .spawn(move || { + let tokenizer = DistributedTokenizer::load(&worker_config.model_dir); + match tokenizer { + Ok(mut tokenizer) => { + let _ = ready_tx.send(Ok(())); + while let Some(job) = job_rx.blocking_recv() { + if let Err(error) = generate_one(&worker_config, &mut tokenizer, &job) { + let _ = job.reply.send(TokenMsg::Error(format!("{error:#}"))); + } + } + } + Err(error) => { + let _ = ready_tx.send(Err(format!("{error:#}"))); + } + } + })?; + match ready_rx.recv() { + Ok(Ok(())) => Ok(Self { jobs, config }), + Ok(Err(error)) => Err(anyhow!("MLX distributed tokenizer load failed: {error}")), + Err(_) => Err(anyhow!("MLX distributed worker exited before readiness")), + } + } + + pub fn model_id(&self) -> &str { + &self.config.model_id + } + + pub fn clamp_max_tokens(&self, requested: Option) -> usize { + requested + .unwrap_or(self.config.default_max_tokens) + .clamp(1, self.config.max_tokens_cap) + } + + pub fn submit(&self, request: GenerateRequest) -> mpsc::UnboundedReceiver { + let (reply, rx) = mpsc::unbounded_channel(); + if self + .jobs + .send(Job { + request, + reply: reply.clone(), + }) + .is_err() + { + let _ = reply.send(TokenMsg::Error( + "MLX distributed worker is not running".to_string(), + )); + } + rx + } +} + +struct DistributedTokenizer { + tokenizer: ChatTokenizer, + chat_template: Option, + eos: Vec, +} + +impl DistributedTokenizer { + fn load(model_dir: &Path) -> Result { + let mut tokenizer = ChatTokenizer::from_file(model_dir.join("tokenizer.json")) + .map_err(|error| anyhow!("tokenizer.json: {error}"))?; + tokenizer.set_template_kwargs(load_tokenizer_template_kwargs(model_dir)?); + let chat_template = load_chat_template(model_dir)?; + let config: Value = serde_json::from_slice( + &fs::read(model_dir.join("config.json")).context("read MLX config.json")?, + ) + .context("parse MLX config.json")?; + Ok(Self { + tokenizer, + chat_template, + eos: eos_token_ids(&config), + }) + } + + fn prompt_tokens(&mut self, request: &GenerateRequest) -> Result> { + let (prompt, add_special_tokens) = self.render_prompt(request)?; + self.tokenizer + .encode(prompt, add_special_tokens) + .map_err(|error| anyhow!("encode distributed MLX prompt: {error}"))? + .get_ids() + .iter() + .copied() + .map(|token| i32::try_from(token).context("token id exceeds i32")) + .collect() + } + + fn render_prompt(&mut self, request: &GenerateRequest) -> Result<(String, bool)> { + if let Some(prompt) = request.raw_prompt.as_ref() { + return Ok((prompt.clone(), true)); + } + let messages = request + .messages + .iter() + .map(|turn| json!({"role": turn.role, "content": turn.content})) + .collect::>(); + let Some(template) = self.chat_template.clone() else { + return Ok((fallback_prompt(request), true)); + }; + let rendered = self + .tokenizer + .apply_chat_template_json(template, vec![messages], None, "mesh-llm-mlx", true, None) + .map_err(|error| anyhow!("render MLX chat template: {error}"))? + .into_iter() + .next() + .context("MLX chat template returned no prompt")?; + Ok((rendered, false)) + } +} + +fn load_chat_template(model_dir: &Path) -> Result> { + let tokenizer_config = model_dir.join("tokenizer_config.json"); + if tokenizer_config.is_file() + && let Some(template) = load_model_chat_template_from_file(&tokenizer_config)? + { + return Ok(Some(template)); + } + let standalone = model_dir.join("chat_template.jinja"); + if standalone.is_file() { + return fs::read_to_string(standalone) + .map(Some) + .context("read MLX chat_template.jinja"); + } + Ok(None) +} + +fn load_tokenizer_template_kwargs(model_dir: &Path) -> Result> { + let config_path = model_dir.join("tokenizer_config.json"); + if !config_path.is_file() { + return Ok(Map::new()); + } + let value: Value = + serde_json::from_slice(&fs::read(config_path).context("read MLX tokenizer_config.json")?) + .context("parse MLX tokenizer_config.json")?; + Ok(tokenizer_template_kwargs(&value)) +} + +fn tokenizer_template_kwargs(value: &Value) -> Map { + value + .as_object() + .into_iter() + .flatten() + .filter(|(key, value)| key.ends_with("_token") && (value.is_string() || value.is_null())) + .map(|(key, value)| (key.clone(), value.clone())) + .collect() +} + +fn eos_token_ids(config: &Value) -> Vec { + let value = config + .get("eos_token_id") + .or_else(|| config.get("text_config")?.get("eos_token_id")); + match value { + Some(Value::Number(value)) => value + .as_u64() + .and_then(|id| u32::try_from(id).ok()) + .into_iter() + .collect(), + Some(Value::Array(values)) => values + .iter() + .filter_map(Value::as_u64) + .filter_map(|id| u32::try_from(id).ok()) + .collect(), + _ => Vec::new(), + } +} + +fn fallback_prompt(request: &GenerateRequest) -> String { + request + .messages + .iter() + .map(|turn| format!("{}: {}", turn.role, turn.content)) + .chain(std::iter::once("assistant:".to_string())) + .collect::>() + .join("\n") +} + +fn generate_one( + config: &MlxDistributedEngineConfig, + tokenizer: &mut DistributedTokenizer, + job: &Job, +) -> Result<()> { + let prompt = tokenizer.prompt_tokens(&job.request)?; + ensure!(!prompt.is_empty(), "distributed MLX prompt has no tokens"); + let requested_max_tokens = job.request.max_tokens.clamp(1, config.max_tokens_cap); + let max_tokens = + context_bounded_max_tokens(prompt.len(), requested_max_tokens, config.context_tokens)?; + let request_id = NEXT_REQUEST_ID.fetch_add(1, Ordering::Relaxed).max(1); + let session_id = request_id; + let mut stream = connect_stage(config.stage_addr)?; + let result = generate_tokens( + &mut stream, + config.wire_dtype, + request_id, + session_id, + &prompt, + max_tokens, + &tokenizer.eos, + &tokenizer.tokenizer, + &job.reply, + ); + let stop_result = stop_session(&mut stream, config.wire_dtype, request_id, session_id); + result.and(stop_result) +} + +fn connect_stage(addr: SocketAddr) -> Result { + const CONNECT_TIMEOUT: Duration = Duration::from_secs(10); + const GENERATION_IO_TIMEOUT: Duration = Duration::from_secs(5 * 60); + let mut stream = TcpStream::connect_timeout(&addr, CONNECT_TIMEOUT) + .with_context(|| format!("connect MLX stage 0 at {addr}"))?; + stream.set_nodelay(true).ok(); + stream.set_read_timeout(Some(CONNECT_TIMEOUT))?; + stream.set_write_timeout(Some(CONNECT_TIMEOUT))?; + recv_ready(&mut stream).context("MLX stage 0 did not become ready")?; + stream.set_read_timeout(Some(GENERATION_IO_TIMEOUT))?; + stream.set_write_timeout(Some(GENERATION_IO_TIMEOUT))?; + Ok(stream) +} + +#[allow(clippy::too_many_arguments)] +fn generate_tokens( + stream: &mut TcpStream, + dtype: WireActivationDType, + request_id: u64, + session_id: u64, + prompt: &[i32], + max_tokens: usize, + eos: &[u32], + tokenizer: &ChatTokenizer, + reply: &mpsc::UnboundedSender, +) -> Result<()> { + send_message( + stream, + &prefill_message(dtype, request_id, session_id, prompt), + dtype, + )?; + let mut predicted = predicted_reply(stream)?; + let mut generated = Vec::with_capacity(max_tokens); + let mut decoder = tokenizer.decode_stream(true); + let mut emitted = String::new(); + let mut finish = FinishReason::Length; + while generated.len() < max_tokens { + let token = u32::try_from(predicted).context("negative predicted token")?; + if eos.contains(&token) { + finish = FinishReason::Stop; + break; + } + generated.push(token); + if let Some(delta) = decoder + .step(token) + .map_err(|error| anyhow!("decode distributed MLX token: {error}"))? + { + emitted.push_str(&delta); + let _ = reply.send(TokenMsg::Delta(delta)); + } + if reply.is_closed() { + return Ok(()); + } + if generated.len() == max_tokens { + break; + } + send_message( + stream, + &decode_message( + dtype, + request_id, + session_id, + prompt.len(), + generated.len() - 1, + predicted, + ), + dtype, + )?; + predicted = predicted_reply(stream)?; + } + emit_decode_suffix(tokenizer, &generated, &emitted, reply)?; + let _ = reply.send(TokenMsg::Done { + finish_reason: finish, + prompt_tokens: u32::try_from(prompt.len()).unwrap_or(u32::MAX), + completion_tokens: u32::try_from(generated.len()).unwrap_or(u32::MAX), + }); + Ok(()) +} + +fn emit_decode_suffix( + tokenizer: &ChatTokenizer, + token_ids: &[u32], + emitted: &str, + reply: &mpsc::UnboundedSender, +) -> Result<()> { + let decoded = tokenizer + .decode(token_ids, true) + .map_err(|error| anyhow!("finalize distributed MLX decode: {error}"))?; + if let Some(suffix) = decoded.strip_prefix(emitted) + && !suffix.is_empty() + { + let _ = reply.send(TokenMsg::Delta(suffix.to_string())); + } + Ok(()) +} + +fn context_bounded_max_tokens( + prompt_tokens: usize, + requested_max_tokens: usize, + context_tokens: usize, +) -> Result { + let available = context_tokens.saturating_sub(prompt_tokens); + ensure!( + available > 0, + "distributed MLX prompt uses {prompt_tokens} tokens, exceeding the {context_tokens}-token context" + ); + Ok(requested_max_tokens.min(available)) +} + +fn prefill_message( + dtype: WireActivationDType, + request_id: u64, + session_id: u64, + tokens: &[i32], +) -> StageWireMessage { + let token_count = i32::try_from(tokens.len()).unwrap_or(i32::MAX); + let mut state = StageStateHeader::new(WireMessageKind::PrefillFinalEmbd, dtype); + state.prompt_token_count = token_count; + StageWireMessage { + kind: WireMessageKind::PrefillFinalEmbd, + pos_start: 0, + token_count, + state, + request_id, + session_id, + sampling: None, + chat_sampling_metadata: None, + tokens: tokens.to_vec(), + positions: (0..token_count).collect(), + activation: Vec::new(), + raw_bytes: Vec::new(), + } +} + +fn decode_message( + dtype: WireActivationDType, + request_id: u64, + session_id: u64, + prompt_tokens: usize, + decode_step: usize, + current_token: i32, +) -> StageWireMessage { + let mut state = StageStateHeader::new(WireMessageKind::DecodeEmbd, dtype); + state.prompt_token_count = i32::try_from(prompt_tokens).unwrap_or(i32::MAX); + state.decode_step = i32::try_from(decode_step).unwrap_or(i32::MAX); + state.current_token = current_token; + let position = prompt_tokens.saturating_add(decode_step); + StageWireMessage { + kind: WireMessageKind::DecodeEmbd, + pos_start: i32::try_from(position).unwrap_or(i32::MAX), + token_count: 1, + state, + request_id, + session_id, + sampling: None, + chat_sampling_metadata: None, + tokens: Vec::new(), + positions: vec![i32::try_from(position).unwrap_or(i32::MAX)], + activation: Vec::new(), + raw_bytes: Vec::new(), + } +} + +fn send_message( + stream: &mut TcpStream, + message: &StageWireMessage, + dtype: WireActivationDType, +) -> Result<()> { + write_stage_message(&mut *stream, message, dtype)?; + stream.flush()?; + Ok(()) +} + +fn predicted_reply(stream: &mut TcpStream) -> Result { + let reply = recv_reply(stream)?; + ensure!( + matches!( + reply.kind, + WireReplyKind::PredictedToken | WireReplyKind::PredictedTokens + ), + "MLX stage chain returned {:?}, expected a prediction", + reply.kind + ); + prediction(&reply).context("MLX stage chain returned no predicted token") +} + +fn prediction(reply: &StageReply) -> Option { + reply + .predicted_tokens + .first() + .copied() + .or_else(|| (reply.kind == WireReplyKind::PredictedToken).then_some(reply.predicted)) +} + +fn stop_session( + stream: &mut TcpStream, + dtype: WireActivationDType, + request_id: u64, + session_id: u64, +) -> Result<()> { + send_message( + stream, + &StageWireMessage::stop_with_identity(dtype, request_id, session_id), + dtype, + )?; + let reply = recv_reply(stream)?; + ensure!( + reply.kind == WireReplyKind::Ack, + "MLX stage stop was not acknowledged" + ); + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn prefill_frame_requests_prediction_for_complete_prompt() { + let message = prefill_message(WireActivationDType::F16, 7, 9, &[1, 2, 3]); + + assert_eq!(message.kind, WireMessageKind::PrefillFinalEmbd); + assert_eq!(message.token_count, 3); + assert_eq!(message.tokens, vec![1, 2, 3]); + assert_eq!(message.positions, vec![0, 1, 2]); + assert_eq!(message.state.prompt_token_count, 3); + assert!(message.kind.requires_predicted_reply()); + } + + #[test] + fn decode_frame_carries_current_token_and_absolute_position() { + let message = decode_message(WireActivationDType::F16, 7, 9, 12, 2, 42); + + assert_eq!(message.kind, WireMessageKind::DecodeEmbd); + assert_eq!(message.pos_start, 14); + assert_eq!(message.positions, vec![14]); + assert_eq!(message.state.current_token, 42); + assert_eq!(message.state.decode_step, 2); + assert_eq!(message.state.prompt_token_count, 12); + } + + #[test] + fn prediction_prefers_multi_token_sideband() { + let reply = StageReply { + kind: WireReplyKind::PredictedTokens, + predicted: 3, + predicted_tokens: vec![4, 5], + stats: Default::default(), + }; + + assert_eq!(prediction(&reply), Some(4)); + } + + #[test] + fn scalar_prediction_falls_back_to_scalar_field() { + let reply = StageReply { + kind: WireReplyKind::PredictedToken, + predicted: 3, + predicted_tokens: Vec::new(), + stats: Default::default(), + }; + + assert_eq!(prediction(&reply), Some(3)); + } + + #[test] + fn empty_prediction_batch_is_rejected() { + let reply = StageReply { + kind: WireReplyKind::PredictedTokens, + predicted: 0, + predicted_tokens: Vec::new(), + stats: Default::default(), + }; + + assert_eq!(prediction(&reply), None); + } + + #[test] + fn generation_is_bounded_by_remaining_context() { + assert_eq!(context_bounded_max_tokens(12, 8, 16).unwrap(), 4); + assert!(context_bounded_max_tokens(16, 8, 16).is_err()); + } + + #[test] + fn tokenizer_template_kwargs_keep_only_special_token_values() { + let value = json!({ + "bos_token": "", + "eos_token": null, + "chat_template": "ignored", + "added_token": {"content": "ignored"}, + "padding_side": "left" + }); + + let kwargs = tokenizer_template_kwargs(&value); + + assert_eq!(kwargs.len(), 2); + assert_eq!(kwargs["bos_token"], ""); + assert!(kwargs["eos_token"].is_null()); + } +} diff --git a/crates/skippy-engine-mlx/src/engine.rs b/crates/skippy-engine-mlx/src/engine.rs index 9378feda7b..b18f6608dc 100644 --- a/crates/skippy-engine-mlx/src/engine.rs +++ b/crates/skippy-engine-mlx/src/engine.rs @@ -11,6 +11,8 @@ //! This also naturally serializes GPU access (one generation at a time), which //! matches how goose drives safemlx today. +use std::fs; +use std::path::Path; use std::path::PathBuf; use std::thread; use std::time::Instant; @@ -21,10 +23,12 @@ use tokio::sync::mpsc; use safemlx::transforms::async_eval; use safemlx::{Device, DeviceType, Stream}; -use safemlx_lm::models::LoadedModel; use safemlx_lm::models::input::{InputPart, ModelInput}; +use safemlx_lm::models::{LoadedModel, ModelLoadOptions}; use safemlx_lm::sampler::DefaultSampler; +use crate::stage::MlxWeightQuantization; + /// How the worker should load and run a model. #[derive(Clone, Debug)] pub struct MlxEngineConfig { @@ -32,6 +36,54 @@ pub struct MlxEngineConfig { pub model_id: String, pub default_max_tokens: usize, pub max_tokens_cap: usize, + /// Quantize eligible dense checkpoint tensors as they are loaded. Already + /// quantized checkpoints with matching metadata load directly. + pub weight_quantization: Option, +} + +/// Selects load-time affine-4 only for dense, unquantized model families that +/// the pinned safemlx runtime can quantize safely. Frontier grouped-expert +/// families and checkpoints already declaring a representation load natively. +pub fn automatic_weight_quantization(model_dir: &Path) -> Result> { + if model_dir.is_file() { + return Ok(None); + } + let config: Value = serde_json::from_slice( + &fs::read(model_dir.join("config.json")) + .map_err(|error| anyhow!("read MLX config.json: {error}"))?, + ) + .map_err(|error| anyhow!("parse MLX config.json: {error}"))?; + Ok(automatic_weight_quantization_for_config(&config)) +} + +fn automatic_weight_quantization_for_config(config: &Value) -> Option { + let model_type = config.get("model_type").and_then(Value::as_str); + let text_config = config.get("text_config"); + let text_model_type = text_config + .and_then(|value| value.get("model_type")) + .and_then(Value::as_str); + let is_grouped_expert = |model_type: Option<&str>| { + matches!( + model_type, + Some("inkling_mm_model" | "nemotron_h" | "nemotron_h_moe") + ) + }; + let unsupported_grouped_experts = + is_grouped_expert(model_type) || is_grouped_expert(text_model_type); + let declares_representation = declares_weight_representation(config) + || text_config.is_some_and(declares_weight_representation); + (!unsupported_grouped_experts && !declares_representation).then_some( + MlxWeightQuantization::Affine { + group_size: 64, + bits: 4, + }, + ) +} + +fn declares_weight_representation(config: &Value) -> bool { + ["quantization", "quantization_config", "compression_config"] + .iter() + .any(|key| config.get(*key).is_some_and(|value| !value.is_null())) } /// One chat turn, in `Send` form (no MLX types). @@ -138,17 +190,30 @@ fn load_engine(config: &MlxEngineConfig) -> Result { let weights_stream = Stream::new_with_device(&Device::new(DeviceType::Cpu, 0)); let started = Instant::now(); - let model = LoadedModel::load(&config.model_dir, &stream, &weights_stream) - .map_err(|e| anyhow!("load {}: {e}", config.model_dir.display()))?; + let options = config + .weight_quantization + .map(MlxWeightQuantization::safemlx) + .transpose()? + .map_or_else( + ModelLoadOptions::default, + ModelLoadOptions::with_quantization, + ); + let model = + LoadedModel::load_with_options(&config.model_dir, options, &stream, &weights_stream) + .map_err(|e| anyhow!("load {}: {e}", config.model_dir.display()))?; stream.synchronize().map_err(|e| anyhow!("sync: {e}"))?; let tokenizer = tokenizers::Tokenizer::from_file(config.model_dir.join("tokenizer.json")) .map_err(|e| anyhow!("tokenizer.json: {e}"))?; let eos = model.eos_token_ids().to_vec(); + let quantization_label = config + .weight_quantization + .map_or_else(|| "checkpoint".to_string(), MlxWeightQuantization::label); tracing::info!( model = %config.model_id, kind = model.model_type(), + weight_quantization = %quantization_label, load_secs = started.elapsed().as_secs_f64(), "MLX model loaded" ); @@ -178,7 +243,7 @@ fn run_worker( while let Some(job) = job_rx.blocking_recv() { let reply = job.reply.clone(); - if let Err(e) = generate_one(&mut engine, job) { + if let Err(e) = generate_one(&config, &mut engine, job) { let _ = reply.send(TokenMsg::Error(e.to_string())); } } @@ -209,7 +274,7 @@ fn build_prompt(model: &mut LoadedModel, req: &GenerateRequest) -> Result<(Strin } } -fn generate_one(engine: &mut LoadedEngine, job: Job) -> Result<()> { +fn generate_one(config: &MlxEngineConfig, engine: &mut LoadedEngine, job: Job) -> Result<()> { let LoadedEngine { model, stream, @@ -223,6 +288,13 @@ fn generate_one(engine: &mut LoadedEngine, job: Job) -> Result<()> { .encode_to_array(&prompt, add_special, stream) .map_err(|e| anyhow!("encode: {e}"))?; let prompt_tokens = tokens.shape()[1] as u32; + let available_tokens = config.max_tokens_cap.saturating_sub(prompt_tokens as usize); + anyhow::ensure!( + available_tokens > 0, + "MLX prompt uses {prompt_tokens} tokens, exceeding the {}-token context", + config.max_tokens_cap + ); + let max_tokens = job.req.max_tokens.min(available_tokens); let mut cache = model.new_cache(); let parts = [InputPart::text_token_ids(&tokens)]; @@ -236,12 +308,13 @@ fn generate_one(engine: &mut LoadedEngine, job: Job) -> Result<()> { DefaultSampler, ); - let mut ids: Vec = Vec::with_capacity(job.req.max_tokens); + let mut ids: Vec = Vec::with_capacity(max_tokens); + let mut decoder = tokenizer.decode_stream(true); let mut emitted = String::new(); let mut finish = FinishReason::Length; let mut current = generator.next().transpose().map_err(|e| anyhow!("{e}"))?; - for index in 0..job.req.max_tokens { + for index in 0..max_tokens { let Some(token) = current.take() else { finish = FinishReason::Stop; break; @@ -249,7 +322,7 @@ fn generate_one(engine: &mut LoadedEngine, job: Job) -> Result<()> { // Start the next decode before reading this token back (mlx-lm's // one-token async pipeline overlaps compute with host readback). - let next = if index + 1 < job.req.max_tokens { + let next = if index + 1 < max_tokens { let next = generator.next(); if let Some(Ok(next_token)) = next.as_ref() { async_eval([next_token]).map_err(|e| anyhow!("async_eval: {e}"))?; @@ -266,21 +339,14 @@ fn generate_one(engine: &mut LoadedEngine, job: Job) -> Result<()> { } ids.push(token_id); - // Incremental detokenization: decode the whole id sequence and emit only - // the newly appended suffix. Robust for byte-level BPE where one char can - // span multiple tokens. - let full = tokenizer - .decode(&ids, true) - .map_err(|e| anyhow!("decode: {e}"))?; - if let Some(delta) = full.strip_prefix(emitted.as_str()) { - if !delta.is_empty() { - if reply.send(TokenMsg::Delta(delta.to_string())).is_err() { - return Ok(()); // client hung up - } - emitted = full; + if let Some(delta) = decoder + .step(token_id) + .map_err(|error| anyhow!("decode: {error}"))? + { + emitted.push_str(&delta); + if reply.send(TokenMsg::Delta(delta)).is_err() { + return Ok(()); // client hung up } - } else { - emitted = full; // rare re-render; resync silently } if reply.is_closed() { @@ -289,6 +355,15 @@ fn generate_one(engine: &mut LoadedEngine, job: Job) -> Result<()> { current = next.transpose().map_err(|e| anyhow!("{e}"))?; } + let decoded = tokenizer + .decode(&ids, true) + .map_err(|error| anyhow!("finalize decode: {error}"))?; + if let Some(suffix) = decoded.strip_prefix(&emitted) + && !suffix.is_empty() + { + let _ = reply.send(TokenMsg::Delta(suffix.to_string())); + } + let _ = reply.send(TokenMsg::Done { finish_reason: finish, prompt_tokens, @@ -296,3 +371,48 @@ fn generate_one(engine: &mut LoadedEngine, job: Job) -> Result<()> { }); Ok(()) } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn automatic_quantization_preserves_frontier_and_prequantized_models() { + assert_eq!( + automatic_weight_quantization_for_config(&json!({"model_type": "llama"})), + Some(MlxWeightQuantization::Affine { + group_size: 64, + bits: 4 + }) + ); + assert_eq!( + automatic_weight_quantization_for_config(&json!({"model_type": "inkling_mm_model"})), + None + ); + assert_eq!( + automatic_weight_quantization_for_config(&json!({"model_type": "nemotron_h"})), + None + ); + assert_eq!( + automatic_weight_quantization_for_config(&json!({ + "model_type": "multimodal_wrapper", + "text_config": {"model_type": "nemotron_h"} + })), + None + ); + assert_eq!( + automatic_weight_quantization_for_config(&json!({ + "model_type": "qwen3", + "quantization_config": {"bits": 8} + })), + None + ); + assert_eq!( + automatic_weight_quantization_for_config(&json!({ + "model_type": "qwen3_vl", + "text_config": {"compression_config": {"format": "mxfp4"}} + })), + None + ); + } +} diff --git a/crates/skippy-engine-mlx/src/lib.rs b/crates/skippy-engine-mlx/src/lib.rs index 1945997425..4a544da616 100644 --- a/crates/skippy-engine-mlx/src/lib.rs +++ b/crates/skippy-engine-mlx/src/lib.rs @@ -14,6 +14,8 @@ mod boundary_bench; #[cfg(all(feature = "mlx", target_os = "macos"))] mod derived; #[cfg(all(feature = "mlx", target_os = "macos"))] +mod distributed; +#[cfg(all(feature = "mlx", target_os = "macos"))] mod engine; #[cfg(all(feature = "mlx", target_os = "macos"))] mod stage; @@ -34,7 +36,11 @@ pub use derived::{ load_prepared_quantized_stage, mlx_derived_stage_cache_root, validate_nemotron_h_moe_stage, }; #[cfg(all(feature = "mlx", target_os = "macos"))] -pub use engine::{ChatTurn, GenerateRequest, MlxEngine, MlxEngineConfig}; +pub use distributed::{MlxDistributedEngine, MlxDistributedEngineConfig}; +#[cfg(all(feature = "mlx", target_os = "macos"))] +pub use engine::{ + ChatTurn, GenerateRequest, MlxEngine, MlxEngineConfig, automatic_weight_quantization, +}; #[cfg(all(feature = "mlx", target_os = "macos"))] pub use stage::{ MlxComputeDtype, MlxNemotronHStageValidationReport, MlxNemotronHWireValidationReport, diff --git a/crates/skippy-protocol/src/lib.rs b/crates/skippy-protocol/src/lib.rs index 3053ea0718..4270095afa 100644 --- a/crates/skippy-protocol/src/lib.rs +++ b/crates/skippy-protocol/src/lib.rs @@ -21,6 +21,8 @@ pub const STAGE_SUBPROTOCOL_FEATURE_STAGE_GENERATION: &str = STAGE_SUBPROTOCOL_FEATURE_STAGE_PROTOCOL_GENERATION_V3; pub const STAGE_SUBPROTOCOL_FEATURE_ARTIFACT_TRANSFER: &str = "artifact-transfer"; pub const STAGE_SUBPROTOCOL_FEATURE_STATUS_LIST: &str = "status-list"; +/// The peer can derive and execute MLX SafeTensors artifact slices. +pub const STAGE_SUBPROTOCOL_FEATURE_BACKEND_MLX: &str = "backend-mlx"; pub const STAGE_STREAM_CONTROL: u8 = 0x01; pub const STAGE_STREAM_TRANSPORT: u8 = 0x02; pub const STAGE_STREAM_ARTIFACT_TRANSFER: u8 = 0x03; diff --git a/crates/skippy-server/src/engine_transport.rs b/crates/skippy-server/src/engine_transport.rs index 69ceff7a13..7eda810df3 100644 --- a/crates/skippy-server/src/engine_transport.rs +++ b/crates/skippy-server/src/engine_transport.rs @@ -7,6 +7,7 @@ //! second wire protocol. Advanced operations stay capability-gated. use std::{ + collections::HashSet, io::{self, Write}, net::{Shutdown, SocketAddr, TcpListener, TcpStream}, sync::{ @@ -17,6 +18,8 @@ use std::{ time::Duration, }; +const ENGINE_STAGE_IO_TIMEOUT: Duration = Duration::from_secs(5 * 60); + use anyhow::{Context, Result, bail, ensure}; use skippy_engine::{ StageActivation, StageEngine, StageExecutionKind, StageExecutionOutput, StageExecutionRequest, @@ -170,6 +173,8 @@ fn handle_connection( mut upstream: TcpStream, downstream_control: Arc>>, ) -> Result<()> { + upstream.set_read_timeout(Some(ENGINE_STAGE_IO_TIMEOUT))?; + upstream.set_write_timeout(Some(ENGINE_STAGE_IO_TIMEOUT))?; let mut downstream = options .downstream_addr .map(connect_downstream) @@ -183,9 +188,41 @@ fn handle_connection( upstream.flush().ok(); let activation_width = i32::try_from(engine.info().activation_width).context("activation width exceeds i32")?; + let mut active_sessions = HashSet::new(); + let result = handle_connection_messages( + engine.as_ref(), + &options, + &mut upstream, + downstream.as_mut(), + activation_width, + &mut active_sessions, + ); + let cleanup = cleanup_connection_sessions( + engine.as_ref(), + downstream.as_mut(), + options.wire_dtype, + &active_sessions, + ); + match (result, cleanup) { + (Err(error), Err(cleanup_error)) => { + eprintln!("engine stage session cleanup also failed: {cleanup_error:#}"); + Err(error) + } + (Err(error), _) => Err(error), + (Ok(()), cleanup) => cleanup, + } +} +fn handle_connection_messages( + engine: &dyn StageEngine, + options: &EngineStageServerOptions, + upstream: &mut TcpStream, + mut downstream: Option<&mut TcpStream>, + activation_width: i32, + active_sessions: &mut HashSet, +) -> Result<()> { loop { - let message = match read_stage_message(&mut upstream, activation_width) { + let message = match read_stage_message(&mut *upstream, activation_width) { Ok(message) => message, Err(error) if error.kind() == io::ErrorKind::UnexpectedEof => return Ok(()), Err(error) => return Err(error).context("read engine stage message"), @@ -193,35 +230,69 @@ fn handle_connection( if message.kind == WireMessageKind::Stop { engine.reset_session(message.session_id)?; let downstream_reply = - forward_control(downstream.as_mut(), &message, options.wire_dtype)?; - send_ack(&mut upstream, downstream_reply)?; + forward_control(downstream.as_deref_mut(), &message, options.wire_dtype)?; + send_ack(upstream, downstream_reply)?; + active_sessions.remove(&message.session_id); continue; } if message.kind.is_session_control() { - execute_session_control(engine.as_ref(), &message)?; + active_sessions.insert(message.session_id); + execute_session_control(engine, &message)?; let downstream_reply = - forward_control(downstream.as_mut(), &message, options.wire_dtype)?; - send_ack(&mut upstream, downstream_reply)?; + forward_control(downstream.as_deref_mut(), &message, options.wire_dtype)?; + send_ack(upstream, downstream_reply)?; continue; } + active_sessions.insert(message.session_id); let request = execution_request(&message, activation_width)?; let output = engine.execute(request)?; - match downstream.as_mut() { + match downstream.as_deref_mut() { Some(downstream) => { - let forwarded = - forwarded_message(engine.as_ref(), &message, output, options.wire_dtype)?; + let forwarded = forwarded_message(engine, &message, output, options.wire_dtype)?; write_stage_message(&mut *downstream, &forwarded, options.wire_dtype) .context("forward engine stage message")?; downstream.flush().ok(); let reply = recv_reply(&mut *downstream).context("receive downstream reply")?; - send_reply(&mut upstream, reply)?; + send_reply(upstream, reply)?; } - None => send_final_reply(&mut upstream, &message, output)?, + None => send_final_reply(upstream, &message, output)?, } } } +fn cleanup_connection_sessions( + engine: &dyn StageEngine, + mut downstream: Option<&mut TcpStream>, + wire_dtype: WireActivationDType, + active_sessions: &HashSet, +) -> Result<()> { + let mut first_error = None; + for session_id in active_sessions { + if let Err(error) = engine.reset_session(*session_id) + && first_error.is_none() + { + first_error = Some(error.context(format!("reset abandoned session {session_id}"))); + } + let stop = StageWireMessage::stop_with_identity(wire_dtype, *session_id, *session_id); + let downstream_result = forward_control(downstream.as_deref_mut(), &stop, wire_dtype) + .and_then(|reply| { + ensure!( + reply.is_none_or(|reply| reply.kind == WireReplyKind::Ack), + "abandoned session stop expected downstream ACK" + ); + Ok(()) + }); + if let Err(error) = downstream_result + && first_error.is_none() + { + first_error = + Some(error.context(format!("propagate abandoned session {session_id} stop"))); + } + } + first_error.map_or(Ok(()), Err) +} + fn connect_downstream(addr: SocketAddr) -> Result { const CONNECT_TIMEOUT: Duration = Duration::from_secs(5); let mut stream = TcpStream::connect_timeout(&addr, CONNECT_TIMEOUT) @@ -230,8 +301,8 @@ fn connect_downstream(addr: SocketAddr) -> Result { stream.set_read_timeout(Some(CONNECT_TIMEOUT))?; stream.set_write_timeout(Some(CONNECT_TIMEOUT))?; recv_ready(&mut stream).context("downstream engine stage did not become ready")?; - stream.set_read_timeout(None)?; - stream.set_write_timeout(None)?; + stream.set_read_timeout(Some(ENGINE_STAGE_IO_TIMEOUT))?; + stream.set_write_timeout(Some(ENGINE_STAGE_IO_TIMEOUT))?; Ok(stream) } @@ -413,6 +484,43 @@ mod tests { use super::*; use skippy_protocol::binary::StageStateHeader; + struct RecordingEngine { + info: skippy_engine::StageEngineInfo, + resets: Mutex>, + } + + impl RecordingEngine { + fn new() -> Self { + Self { + info: skippy_engine::StageEngineInfo { + engine: "test".to_string(), + model_id: "test/model".to_string(), + stage_index: 0, + layer_start: 0, + layer_end: 1, + total_layers: 1, + activation_width: 4, + }, + resets: Mutex::new(Vec::new()), + } + } + } + + impl StageEngine for RecordingEngine { + fn info(&self) -> &skippy_engine::StageEngineInfo { + &self.info + } + + fn execute(&self, _request: StageExecutionRequest) -> Result { + Ok(StageExecutionOutput::default()) + } + + fn reset_session(&self, session_id: u64) -> Result<()> { + self.resets.lock().unwrap().push(session_id); + Ok(()) + } + } + fn decode_message(tokens: Vec, current_token: i32) -> StageWireMessage { let kind = WireMessageKind::DecodeEmbd; let mut state = StageStateHeader::new(kind, WireActivationDType::F16); @@ -439,4 +547,16 @@ mod tests { assert_eq!(request.token_ids, vec![7]); assert_eq!(request.kind, StageExecutionKind::Decode); } + + #[test] + fn connection_cleanup_resets_every_abandoned_session() { + let engine = RecordingEngine::new(); + let sessions = HashSet::from([7, 9]); + + cleanup_connection_sessions(&engine, None, WireActivationDType::F16, &sessions).unwrap(); + + let mut resets = engine.resets.lock().unwrap().clone(); + resets.sort_unstable(); + assert_eq!(resets, vec![7, 9]); + } } diff --git a/docs/design/MLX_STAGE_ENGINE_PLAN.md b/docs/design/MLX_STAGE_ENGINE_PLAN.md index 3764ab7b41..8f956d08a4 100644 --- a/docs/design/MLX_STAGE_ENGINE_PLAN.md +++ b/docs/design/MLX_STAGE_ENGINE_PLAN.md @@ -50,6 +50,21 @@ cache/session operations, additional staged families, and bounded range-to-derived-cache quantization remain. See `crates/skippy-engine-mlx/STAGED_EXECUTION.md`. +**Update — whole-model and explicit mesh serving are now integrated.** The +ordinary `mesh-llm serve --model` path can select MLX for a complete Hugging +Face SafeTensors model, automatically quantize eligible unquantized dense +weights to affine-4 at load time while preserving frontier/pre-quantized +representations, and serve streaming or non-streaming OpenAI chat. The explicit `--split` +path avoids the coordinator's full checkpoint download, advertises an additive +MLX stage capability, derives per-host range-only affine artifacts, and drives +the resulting chain from the same OpenAI surface. A real two-host 29/1-layer +SmolLM2 run completed with roughly 70 MB and 17 MB stage artifacts while the +remote never stored the roughly 269 MB source checkpoint. Planning now uses +header-derived per-layer estimates so large embedding/readout tensors are not +hidden by an equal-layer average. Automatic split selection and partial +executors beyond dense Llama remain open. See +`crates/skippy-engine-mlx/SERVE_INTEGRATION_STATUS.md`. + **Update — host `StagePrepare` / `StageLoad` now consumes range-only MLX stages.** An immutable `hf-model://org/repo@` request is validated before network work, checked against a topology-wide checkpoint identity, diff --git a/scripts/build-mac.sh b/scripts/build-mac.sh index 319a0147e0..25c9181e1f 100755 --- a/scripts/build-mac.sh +++ b/scripts/build-mac.sh @@ -14,6 +14,7 @@ build_profile="${MESH_LLM_BUILD_PROFILE:-debug}" MESH_LLM_LOCAL_CODESIGN_IDENTITY="${MESH_LLM_LOCAL_CODESIGN_IDENTITY:-Mesh-LLM Local Codesign}" MESH_LLM_AUTO_GENERATE_CODESIGN="${MESH_LLM_AUTO_GENERATE_CODESIGN:-1}" rustc_wrapper="" +MESH_LLM_CARGO_FEATURES="${MESH_LLM_CARGO_FEATURES:-}" build_profile="${build_profile:l}" append_rustflag() { @@ -497,14 +498,18 @@ if [[ -d "$MESH_DIR" ]]; then configure_rust_cache mesh_binary="" + cargo_features=() + if [[ -n "$MESH_LLM_CARGO_FEATURES" ]]; then + cargo_features=(--features "$MESH_LLM_CARGO_FEATURES") + fi case "$build_profile" in dev|debug) echo "Building mesh-llm (profile: dev, bin only)..." stamp_build_version if [[ -n "$rustc_wrapper" ]]; then - (cd "$REPO_ROOT" && RUSTC_WRAPPER="$rustc_wrapper" cargo build -p mesh-llm --bin mesh-llm) + (cd "$REPO_ROOT" && RUSTC_WRAPPER="$rustc_wrapper" cargo build -p mesh-llm --bin mesh-llm "${cargo_features[@]}") else - (cd "$REPO_ROOT" && cargo build -p mesh-llm --bin mesh-llm) + (cd "$REPO_ROOT" && cargo build -p mesh-llm --bin mesh-llm "${cargo_features[@]}") fi mesh_binary="target/debug/mesh-llm" echo "Mesh binary: $mesh_binary" @@ -513,9 +518,9 @@ if [[ -d "$MESH_DIR" ]]; then echo "Building mesh-llm (profile: release)..." stamp_build_version if [[ -n "$rustc_wrapper" ]]; then - (cd "$REPO_ROOT" && RUSTC_WRAPPER="$rustc_wrapper" cargo build --release -p mesh-llm) + (cd "$REPO_ROOT" && RUSTC_WRAPPER="$rustc_wrapper" cargo build --release -p mesh-llm "${cargo_features[@]}") else - (cd "$REPO_ROOT" && cargo build --release -p mesh-llm) + (cd "$REPO_ROOT" && cargo build --release -p mesh-llm "${cargo_features[@]}") fi mesh_binary="target/release/mesh-llm" echo "Mesh binary: $mesh_binary" @@ -526,5 +531,21 @@ if [[ -d "$MESH_DIR" ]]; then ;; esac + if [[ "$MESH_LLM_CARGO_FEATURES" =~ '(^|[[:space:],])mlx($|[[:space:],])' ]]; then + resource_dir="${mesh_binary:h}/safemlx-resources" + metallib="${resource_dir}/mlx.metallib" + sibling="${mesh_binary:h}/mlx.metallib" + [[ -s "$REPO_ROOT/$metallib" ]] || { + echo "MLX build did not produce $metallib" >&2 + exit 1 + } + cp "$REPO_ROOT/$metallib" "$REPO_ROOT/$sibling" + cmp -s "$REPO_ROOT/$metallib" "$REPO_ROOT/$sibling" || { + echo "MLX metallib copy verification failed" >&2 + exit 1 + } + echo "MLX Metal library: $sibling" + fi + sign_with_keychain_identity_if_available "$REPO_ROOT/$mesh_binary" fi From 9b1cdcd09cc3f297e6f54e26f048874c2c724a02 Mon Sep 17 00:00:00 2001 From: Michael Neale <14976+michaelneale@users.noreply.github.com> Date: Sat, 18 Jul 2026 02:46:48 +1000 Subject: [PATCH 32/37] docs(mlx): reconcile integrated serving status --- crates/skippy-engine-mlx/STAGED_EXECUTION.md | 14 +++++++------- docs/design/MLX_STAGE_ENGINE_PLAN.md | 17 +++++++++-------- 2 files changed, 16 insertions(+), 15 deletions(-) diff --git a/crates/skippy-engine-mlx/STAGED_EXECUTION.md b/crates/skippy-engine-mlx/STAGED_EXECUTION.md index 6bb332e668..19f8e9ff2f 100644 --- a/crates/skippy-engine-mlx/STAGED_EXECUTION.md +++ b/crates/skippy-engine-mlx/STAGED_EXECUTION.md @@ -625,13 +625,13 @@ LAN/QUIC evidence. - `engine_transport` is the reduced compatibility lane. The mature llama.cpp binary server remains unchanged and still owns telemetry, exact-prefix cache, batching, and OpenAI orchestration. -- Mesh topology planning does not yet produce MLX stage assignments. The host - can consume explicit `backend=mlx` Prepare/Load requests, but automatic - placement, capability advertisement, coordinator model planning, and an - OpenAI stage-0 frontend remain. Explicit host requests now derive and reuse - quantized artifacts, but cache eviction and an optional local - request-to-recipe locator remain. The quantization field is an additive mesh - protocol change; old peers omit it and therefore mean `auto`, while unknown +- Explicit `mesh-llm serve --split` now capability-gates MLX participants, + produces stage assignments, and drives them from an OpenAI stage-0 frontend. + Automatic split selection and additional family adapters remain. Explicit + host requests derive and reuse quantized artifacts, but cache eviction and + an optional local request-to-recipe locator remain. The quantization field + is an additive mesh protocol change; old peers omit it and therefore mean + `auto`, while unknown values fail closed on new peers. Automatic placement must capability-gate explicit non-default profiles before mixed-version deployment. No Skippy ABI changed. diff --git a/docs/design/MLX_STAGE_ENGINE_PLAN.md b/docs/design/MLX_STAGE_ENGINE_PLAN.md index 8f956d08a4..8aef61dacf 100644 --- a/docs/design/MLX_STAGE_ENGINE_PLAN.md +++ b/docs/design/MLX_STAGE_ENGINE_PLAN.md @@ -74,10 +74,10 @@ Skippy binary wire. `StageLoad` requires that prepared entry and cannot derive or download tensor payloads on a miss. A clean-cache SmolLM2 proof ran both 15-layer ranges through `spawn_stage_control_loop`, reproduced the affine-4 eight-token reference, and -stopped both stages without retaining dense stage artifacts. Automatic MLX -topology production, capability advertisement, remote two-node proof, and -cache eviction remain; this checkpoint proves the host consumer path, not -automatic placement. +stopped both stages without retaining dense stage artifacts. At that +checkpoint, topology production, capability advertisement, remote two-node +proof, and cache eviction remained; the newer integrated update above +supersedes the first three items. Cache eviction still remains. **Update — partial MLX stages now JIT-quantize one tensor at a time.** The pinned safemlx strict loader already contains the required bounded lazy-graph @@ -275,8 +275,9 @@ unknown values fail closed. Apple Metal currently maps `auto` to affine host lifecycle fetched 162,825,984 and 162,827,136 source tensor bytes and wrote 45,859,713 and 45,861,308-byte derived artifacts. It passed in 120.74 seconds; the immediate validated-cache run passed in 8.87 seconds with 258,162,688 B max -RSS. `MESH_MLX_DERIVED_CACHE_DIR` overrides the host cache root. Eviction, -automatic topology selection, and the remote two-node proof remain. +RSS. `MESH_MLX_DERIVED_CACHE_DIR` overrides the host cache root. This was the +pre-integration checkpoint; eviction remains, while the newer update above +records explicit topology selection and the mesh-managed two-node proof. Prepare is the only lifecycle operation allowed to build a missing entry; Load validates and consumes the prepared artifact or fails closed. @@ -1018,8 +1019,8 @@ Spikes 1 and 2 are more decisive than any standalone token/s benchmark. 1. Add capacity and eviction ownership to the host derived-stage cache, and decide whether a local request-to-recipe locator should eliminate warm-path metadata probes. -2. Teach topology planning and capability advertisement to select MLX stages; - explicit host Prepare/Load and the engine-neutral server lane are proven. +2. Extend the explicit, capability-gated dense-Llama topology into automatic + split selection and additional certified model-family stage adapters. 3. Extend the initial metrics-backed synthetic **Spike 2 (boundary fence)** matrix to real model outputs and TCP/QUIC links. Preserve the separate eval/synchronize, host copy, codec, and network phase evidence; do not assume From 96cec9f53ca6ecd2fcf39785dcd5d92a4a706d91 Mon Sep 17 00:00:00 2001 From: Michael Neale <14976+michaelneale@users.noreply.github.com> Date: Sat, 18 Jul 2026 02:52:42 +1000 Subject: [PATCH 33/37] build(mlx): make engine publish-chain compatible --- Cargo.toml | 3 + crates/skippy-engine-mlx/Cargo.toml | 43 ++-- .../SERVE_INTEGRATION_STATUS.md | 2 + crates/skippy-engine-mlx/WIRING.md | 227 +++++------------- docs/design/MLX_STAGE_ENGINE_PLAN.md | 13 +- scripts/publish-crates.sh | 11 + 6 files changed, 108 insertions(+), 191 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index 8f5fedd0a4..56cad99f74 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -83,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/jbg/safemlx", rev = "4e53c5e" } +safemlx-lm = { git = "https://github.com/jbg/safemlx", rev = "4e53c5e" } +safemlx-lm-utils = { git = "https://github.com/jbg/safemlx", rev = "4e53c5e" } [workspace.lints.clippy] cognitive_complexity = "warn" diff --git a/crates/skippy-engine-mlx/Cargo.toml b/crates/skippy-engine-mlx/Cargo.toml index 34c03b0b56..c02166637d 100644 --- a/crates/skippy-engine-mlx/Cargo.toml +++ b/crates/skippy-engine-mlx/Cargo.toml @@ -3,22 +3,18 @@ # Serves HF safetensors tensor models over mesh-llm's REAL OpenAI-compatible # frontend (`openai-frontend::OpenAiBackend`), goose-style, on Apple Silicon. # -# It depends on the MLX-in-Rust stack (safemlx / safemlx-lm) pinned to a specific -# PUBLIC commit of github.com/jbg/safemlx. This is a deliberate git-rev pin, not -# a private fork: safemlx's crates.io releases currently produce garbage output -# for dense models (Qwen3, Llama) — verified this to be a library bug, not a -# prompting/template issue — while a recent upstream commit serves them -# correctly. We pin that commit and will swap to a normal version pin once -# safemlx cuts a working release. See WIRING.md. -# -# The crate keeps its OWN `[workspace]` table for now so the heavy MLX native -# build (CMake + full MLX C++ compile) never runs in unrelated builds/CI. -# Promotion to a real workspace member is documented in WIRING.md. +# The published manifest uses crates.io versions. This repository patches them +# to a proven public safemlx commit until an equivalent release is available. +# See WIRING.md. [package] name = "skippy-engine-mlx" version.workspace = true edition.workspace = true -publish = false +license.workspace = true +description = "Apple MLX stage and OpenAI serving engine for Skippy" +repository = "https://github.com/Mesh-LLM/mesh-llm" +homepage = "https://github.com/Mesh-LLM/mesh-llm" +readme = "STAGED_EXECUTION.md" [lib] name = "skippy_engine_mlx" @@ -59,12 +55,12 @@ mlx = [ [dependencies] # The real mesh-llm OpenAI frontend — this is what proves we serve over the # same surface the shipped binary uses, not a toy. -openai-frontend = { path = "../openai-frontend" } -model-hf = { path = "../model-hf", optional = true } -skippy-engine = { path = "../skippy-engine" } -skippy-metrics = { path = "../skippy-metrics", optional = true } -skippy-protocol = { path = "../skippy-protocol" } -skippy-server = { path = "../skippy-server" } +openai-frontend = { path = "../openai-frontend", version = "0.72.1" } +model-hf = { path = "../model-hf", version = "0.72.1", optional = true } +skippy-engine = { path = "../skippy-engine", version = "0.72.1" } +skippy-metrics = { path = "../skippy-metrics", version = "0.72.1", optional = true } +skippy-protocol = { path = "../skippy-protocol", version = "0.72.1" } +skippy-server = { path = "../skippy-server", version = "0.72.1" } async-trait = "0.1" async-stream = "0.3" anyhow = "1" @@ -84,16 +80,15 @@ tokio = { version = "1", features = ["full"] } tracing = "0.1" tracing-subscriber = { version = "0.3", features = ["env-filter"] } -# MLX stack — pinned to a specific PUBLIC commit of github.com/jbg/safemlx. -# See the header comment for why this is a git-rev pin rather than a crates.io -# version pin. Same feature set goose uses: accelerate + metal + safetensors. -safemlx = { git = "https://github.com/jbg/safemlx", rev = "4e53c5e", default-features = false, features = [ +# MLX stack. The workspace root patches these versions to the proven public +# commit; published consumers see ordinary registry dependencies. +safemlx = { version = "0.1.3", default-features = false, features = [ "accelerate", "metal", "safetensors", ], optional = true } -safemlx-lm = { git = "https://github.com/jbg/safemlx", rev = "4e53c5e", optional = true } -safemlx-lm-utils = { git = "https://github.com/jbg/safemlx", rev = "4e53c5e", default-features = false, features = [ +safemlx-lm = { version = "0.4.1", optional = true } +safemlx-lm-utils = { version = "0.1.4", default-features = false, features = [ "onig", ], optional = true } # Standalone tokenizer for incremental (streaming) detokenization without diff --git a/crates/skippy-engine-mlx/SERVE_INTEGRATION_STATUS.md b/crates/skippy-engine-mlx/SERVE_INTEGRATION_STATUS.md index b4fda63c6b..b8563a7866 100644 --- a/crates/skippy-engine-mlx/SERVE_INTEGRATION_STATUS.md +++ b/crates/skippy-engine-mlx/SERVE_INTEGRATION_STATUS.md @@ -94,6 +94,8 @@ bytes evenly by layer count. - Automatic affine-4 is the current default for eligible dense checkpoints, not yet a general hardware/quality policy surface. - Cache capacity and eviction are not yet owned by this integration. +- The publishable crate uses registry requirements while the workspace root + patches them to the certified public safemlx revision. - Frontier families require their own stage semantics. Nemotron-H has metadata/range planning and a one-layer execution proof, but not a complete hybrid-model topology. Inkling has whole-model support in the pinned safemlx diff --git a/crates/skippy-engine-mlx/WIRING.md b/crates/skippy-engine-mlx/WIRING.md index 96d1599d1e..3b15040b2a 100644 --- a/crates/skippy-engine-mlx/WIRING.md +++ b/crates/skippy-engine-mlx/WIRING.md @@ -1,170 +1,75 @@ -# Wiring `skippy-engine-mlx` into mesh-llm - -This crate is a **working, self-contained MLX (Metal) serving engine** that -already serves HF safetensors models over mesh-llm's real OpenAI frontend -(`openai_frontend::router_for`). It is intentionally standalone right now — its -own cargo workspace — so the heavy MLX native build never runs in unrelated -builds/CI. This document is the concrete plan to promote it into the shipped -binary so that **on a Mac, `mesh-llm serve` can run an MLX tensor model and -users can pick one from `/v1/models`.** - -## What already works (this crate, today) - -- `MlxEngine` — a dedicated OS worker thread owns the non-`Send` MLX objects - (model, streams, arrays); the outside world talks to it with `Send` channels. -- `MlxBackend: openai_frontend::OpenAiBackend` — `models`, `chat_completion`, - `chat_completion_stream` (SSE), with usage accounting and incremental - detokenization. -- `mlx-serve` bin — `router_for(Arc)` + `axum::serve`. -- Verified on Apple Silicon (Metal), source precision, over the real frontend: - Qwen3-0.6B (non-stream + streaming) and SmolLM2 (Llama arch) both generate - **coherent** output. Source precision ~321 tok/s (see - `../../spikes/mlx-solo/FINDINGS.md`). - -## Dependency: why a git-rev pin (not crates.io) - -`Cargo.toml` pins `safemlx` / `safemlx-lm` to a specific **public** commit of -`github.com/jbg/safemlx` (`rev = "4e53c5e"`), not a crates.io version. This is a -deliberate, reproducible git pin — **not a private fork and no local patches**: - -- safemlx's **published** crates (both `0.1.5` and `0.4.1`) produce **garbage - output for dense models** (Qwen3 *and* Llama both emit repeated-token gibberish - with this exact same crate code). This was verified to be a **library bug**, - not a prompting/template problem — the Qwen3 chat template is confirmed applied - correctly, and greedy sampling (`temp=0` → argmax) is used. -- The pinned upstream commit serves those families correctly. Swapping the same - crate between published and the git pin flips the output coherent↔gibberish, - which isolates the cause to the safemlx version. -- **Action item:** swap the git-rev pin for a normal version pin once safemlx - cuts a crates.io release that serves dense models correctly. -- No JIT quantization is used (plain `LoadedModel::load`), so none of the - quant-path loader quirks apply; this path needs zero source patches. - -## Promotion plan (the actual PR) - -### 1. Make it a real workspace member - -- Add `crates/skippy-engine-mlx` to root `Cargo.toml` `members` and remove its - local `[workspace]` table. The safemlx deps are already git-rev pinned to a - public commit (see "Dependency" above), so nothing else changes about them — - a workspace member with a git dependency is fine. -- Keep the crate's `mlx` feature; gate all MLX code with - `#[cfg(all(feature = "mlx", target_os = "macos"))]` (already done). -- Update `scripts/affected-crates.sh`, `scripts/plan-clippy-batches.sh`, - `scripts/publish-crates.sh` `WORKSPACE_MEMBERS`, and run - `cargo run -p xtask -- repo-consistency ci-crate-lists`. Because MLX is a heavy - native lane, add it to the CI backend-gating like the other native features - (only build/test the `mlx` feature on macOS runners). - -### 2. Depend on it from host-runtime, macOS-gated - -In `crates/mesh-llm-host-runtime/Cargo.toml`: - -```toml -[target.'cfg(target_os = "macos")'.dependencies] -skippy-engine-mlx = { path = "../skippy-engine-mlx", features = ["mlx"], optional = true } - -[features] -mlx = ["dep:skippy-engine-mlx"] +# MLX wiring in mesh-llm + +`skippy-engine-mlx` is a workspace member and the Apple-Silicon serving engine +used by the host runtime's optional `mlx` feature. It supports both whole-model +OpenAI serving and explicit mesh-managed dense-Llama stage chains. Current +status and evidence are in `SERVE_INTEGRATION_STATUS.md` and +`STAGED_EXECUTION.md`. + +## Dependency and publication shape + +The crate manifest uses ordinary crates.io requirements: + +- `safemlx = 0.1.3` +- `safemlx-lm = 0.4.1` +- `safemlx-lm-utils = 0.1.4` + +The repository root patches those packages to public safemlx commit +`4e53c5e`. That commit contains the loader/model behavior certified by this +draft; the published releases do not yet contain every required correction. +Using `[patch.crates-io]` keeps the workspace reproducibly pinned without +putting illegal git dependencies in the published crate manifest. + +`skippy-engine-mlx` is therefore part of `scripts/publish-crates.sh` after its +workspace dependencies and before `mesh-llm-host-runtime`. Both +`repo-consistency publish-crates` and `repo-consistency release-targets` must +pass when this dependency shape changes. + +## Feature and platform gating + +- MLX code and native dependencies are behind the crate's `mlx` feature. +- Host integration is behind `mesh-llm-host-runtime/mlx` and macOS target + dependencies. +- Capability advertisement additionally requires `target_arch = "aarch64"`. +- Enabling MLX also enables the dynamic llama native runtime. MLX and static + llama.cpp both expose GGUF C symbols, so keeping llama.cpp in a separate + dynamic link unit avoids duplicate-symbol failures. +- An MLX-only start may proceed without an installed llama native runtime. + Actual GGUF/Skippy loading still checks and reports that requirement. + +Use the dedicated recipes so the complete Xcode Metal toolchain is selected +and the generated library is copied beside the executable: + +```bash +just mlx-build +just mlx-release-build ``` -(The `path` here is the in-repo crate path once it is a workspace member; its own -safemlx deps stay git-rev pinned.) Propagate a `mlx` feature up through -`crates/mesh-llm/Cargo.toml`, and enable it by default only on macOS builds in -the release packaging. +Both produce `mesh-llm` and a sibling `mlx.metallib`; the build verifies that +the copied resource is byte-identical. -### 3. Add an `Mlx` variant to the launch enum +## Whole-model selection -`crates/mesh-llm-host-runtime/src/runtime/local.rs`: +On an MLX build, a resolved SafeTensors model routes to `MlxModelHandle` before +the GGUF path. The automatic weight policy is intentionally conservative: -- `LocalRuntimeBackendHandle` currently has one variant (`Skippy { .. }`). Add - (macOS+feature gated): +- eligible unquantized dense families request affine-4/group-64 at load time; +- Inkling and Nemotron-H load their native representation because their routed + rank-3 experts do not support that transform; +- checkpoints already declaring quantization/compression load natively; +- explicit `mlx-serve` users may choose auto, none, affine4, affine8, or mxfp4. - ```rust - #[cfg(all(feature = "mlx", target_os = "macos"))] - Mlx { backend: Arc, http: MlxHttpHandle, _death_tx: ... }, - ``` +The model is exposed through the shared `openai-frontend` router. The HTTP +listener binds before runtime readiness is reported. -- Every `match &self.inner { LocalRuntimeBackendHandle::Skippy { .. } => ... }` - in `local.rs` (pid, ctx_used_tokens, openai_guardrails, llama_slots_snapshot, - set_openai_guardrail_mode, shutdown/http accessor) needs an `Mlx` arm. Most map - to simple/None-ish values since MLX has no llama slots or GGUF guardrail state. +## Split selection -### 4. Route safetensors models to the MLX branch at launch +Explicit `mesh-llm serve --split` uses metadata-only checkpoint description, +additive MLX peer capability, ordinary resource-aware topology planning, exact +per-stage HTTP tensor ranges, recipe-keyed affine artifacts, and the existing +Skippy binary stage wire. The coordinator retains tokenizer/config sidecars but +does not download model tensor payloads. -`start_runtime_local_model` currently branches: - -``` -if is_layer_package_ref(..) { layer_package } else { skippy (direct GGUF) } -``` - -Add a first branch: if the resolved model is `ModelFormat::Safetensors` (the -model layer already classifies this — `crates/model-artifact` `ModelFormat`, and -`models/resolve` already detects `is_primary_mlx_weight_file`), and we're on -macOS with the `mlx` feature, start an `MlxEngine` instead: - -```rust -#[cfg(all(feature = "mlx", target_os = "macos"))] -if resolved_format == ModelFormat::Safetensors { - return start_runtime_mlx_model(spec, model_name, plan).await; -} -``` - -`start_runtime_mlx_model` mirrors `start_runtime_skippy_model`: build an -`MlxEngineConfig` from the resolved model dir + planned ctx/limits, `spawn` the -engine on a blocking task, wrap `MlxBackend` in the embedded HTTP handle -(`openai_frontend::router_for`), and return a `LocalRuntimeModelHandle` whose -`backend` string is `"mlx"`. - -### 5. Model discovery / listing already works - -The model layer already discovers, downloads, catalogs, and lists MLX -safetensors repos (`crates/mesh-llm-host-runtime/src/models/catalog.rs`, -`.../models/resolve`, `crates/model-resolver`). No change needed for a user to -*see* MLX models; the missing piece was purely the serving engine, which this -crate provides. Auto-behavior: on a Mac, a resolved safetensors model simply -routes to the MLX engine. - -### 6. Auto on Mac + user selection - -- With the `mlx` feature enabled by default on macOS builds, serving a - safetensors model "just works" with no extra flags. -- Users pick a model the same way as today: `mesh-llm serve --model ` - or from `~/.mesh-llm/config.toml`; safetensors → MLX, GGUF → llama.cpp. -- Optionally add `--serving-backend mlx|llama` to force the engine when a model - is available in both formats (parallels the existing skippy backend selector - noted in `docs/SKIPPY.md`). - -## Out of scope for the first PR - -- **Splits / staged execution.** This is single-stage, whole-model serving. The - staged `StageEngine` trait, partial-load, and activation-frame work (plan §6–§8) - are separate and remain gated on the go/no-go spikes. -- **Tool calling / reasoning parsing.** goose's `mlx.rs` has native + emulated - tool parsing and thinking-output filtering worth porting later; this crate - streams raw model text (including `` blocks) for now. -- **Draft/speculative decoding** (goose's `gemma4_mtp`). -- **JIT quantization on load.** safemlx can affine-quantize dense weights at load - time (~604 tok/s 4-bit in earlier spikes), but it is deliberately excluded here - to keep the first PR to the goose-style plain-load path. - -## Future: Linux + NVIDIA (CUDA) - -MLX is **not Apple-only** — `jbg/safemlx` supports **Linux + NVIDIA GPUs via -CUDA** (a `cuda` cargo feature, plus `nccl` for multi-GPU), gated in -`safemlx-sys/build.rs` with a hard `panic!` to Linux targets, for both -`x86_64-linux` and `sbsa-linux` (ARM). The generation code -(`LoadedModel::load` + `generate_with_cache`) is backend-independent — only the -native build backend differs (Metal vs CUDA). So a later change could serve MLX -on Linux/NVIDIA mesh nodes by widening the gate from `target_os = "macos"` to -also allow `cfg(all(target_os = "linux", feature = "cuda"))` and enabling -`safemlx/cuda`. This is **out of scope for this PR** (Mac/Metal first) and is -tracked as future research; ROCm/Vulkan/Windows are not supported upstream. - -## Testing the promoted path - -- `just build` on macOS with the `mlx` feature. -- `mesh-llm serve --model ` → confirm the model appears in - `/v1/models` and `/v1/chat/completions` returns a generation. -- Confirm non-macOS / no-feature builds are byte-for-byte unaffected (the crate - and its deps compile out entirely). +Automatic split selection and partial adapters beyond dense Llama remain +future work. Inkling and Nemotron-H whole-model support must not be confused +with certified partial-stage support. diff --git a/docs/design/MLX_STAGE_ENGINE_PLAN.md b/docs/design/MLX_STAGE_ENGINE_PLAN.md index 8aef61dacf..03d10f7e9b 100644 --- a/docs/design/MLX_STAGE_ENGINE_PLAN.md +++ b/docs/design/MLX_STAGE_ENGINE_PLAN.md @@ -992,18 +992,19 @@ Spikes 1 and 2 are more decisive than any standalone token/s benchmark. them out of early phases. - **Two artifact pipelines** add storage + certification cost; mitigate with a single canonical BF16 source and reproducible derivation. -- **safemlx supply chain — pin to a git rev, not a crates.io version (confirmed - this session).** The published crates collide version strings with the fork +- **safemlx supply chain — certify a git revision behind registry requirements + (confirmed this session).** The published crates collide version strings with the fork HEAD: crates.io `safemlx-lm 0.4.1` is a *different, older* codebase than the fork's `0.4.1` (851 vs 2221 lines in `qwen3.rs`), because the fork develops on a fixed version without bumping. A fork-free build against published crates **compiled and ran but produced gibberish for Qwen3 source precision and crashed on a pre-quantized repo** (`rms_norm` size mismatch) — the working dense-Qwen3/Llama + JIT-quant code exists only in unpublished fork HEAD. So - MLX-for-Skippy must **pin a specific git commit** of `jbg/safemlx` (and carry - the small loader fixes until upstreamed), or coordinate a real published - release. This makes "track upstream + pin + possibly patch" a **standing cost**, - not a one-off. + The published `skippy-engine-mlx` manifest therefore uses normal registry + requirements while this workspace patches them to a **specific public git + commit** of `jbg/safemlx`. A future safemlx release can remove that root + patch. This makes "track upstream + certify + patch" a **standing cost**, not + a one-off. - **Hardware coverage is a moving target with two gates.** New backends must land in upstream `ml-explore/mlx` *then* be wired through safemlx (which authors no backends itself). ROCm is an active-but-unmerged upstream experiment (#2300); diff --git a/scripts/publish-crates.sh b/scripts/publish-crates.sh index e5097cd116..5ca812fa0e 100755 --- a/scripts/publish-crates.sh +++ b/scripts/publish-crates.sh @@ -391,6 +391,15 @@ unpublished_registry_deps() { skippy-protocol \ skippy-runtime ;; + skippy-engine-mlx) + printf '%s\n' \ + model-hf \ + openai-frontend \ + skippy-engine \ + skippy-metrics \ + skippy-protocol \ + skippy-server + ;; mesh-llm-host-runtime) printf '%s\n' \ mesh-llm-api-server \ @@ -418,6 +427,7 @@ unpublished_registry_deps() { model-resolver \ openai-frontend \ skippy-coordinator \ + skippy-engine-mlx \ skippy-protocol \ skippy-runtime \ skippy-server \ @@ -500,6 +510,7 @@ publish_crates=( skippy-runtime openai-frontend skippy-server + skippy-engine-mlx mesh-llm-plugin-manager mesh-mixture-of-agents mesh-llm-system From 226d00459abd6426be771210f7b5a9be62f87e10 Mon Sep 17 00:00:00 2001 From: Michael Neale <14976+michaelneale@users.noreply.github.com> Date: Sat, 18 Jul 2026 03:00:46 +1000 Subject: [PATCH 34/37] fix(mlx): preserve native load on auto quant mismatch --- .../src/inference/mlx.rs | 1 + .../SERVE_INTEGRATION_STATUS.md | 13 +++- crates/skippy-engine-mlx/WIRING.md | 2 + crates/skippy-engine-mlx/src/bin/mlx-serve.rs | 3 + crates/skippy-engine-mlx/src/engine.rs | 69 +++++++++++++++++-- 5 files changed, 81 insertions(+), 7 deletions(-) diff --git a/crates/mesh-llm-host-runtime/src/inference/mlx.rs b/crates/mesh-llm-host-runtime/src/inference/mlx.rs index fe9d3d4758..3353511d19 100644 --- a/crates/mesh-llm-host-runtime/src/inference/mlx.rs +++ b/crates/mesh-llm-host-runtime/src/inference/mlx.rs @@ -39,6 +39,7 @@ impl MlxModelHandle { 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 { diff --git a/crates/skippy-engine-mlx/SERVE_INTEGRATION_STATUS.md b/crates/skippy-engine-mlx/SERVE_INTEGRATION_STATUS.md index b8563a7866..eb87c4eeea 100644 --- a/crates/skippy-engine-mlx/SERVE_INTEGRATION_STATUS.md +++ b/crates/skippy-engine-mlx/SERVE_INTEGRATION_STATUS.md @@ -37,7 +37,9 @@ does not prove partial downloads: a whole-model server needs all model weights. The integrated loader now quantizes eligible, unquantized dense source tensors to affine 4-bit as they load. Inkling, Nemotron-H, and checkpoints already declaring a quantized representation retain their native representation. The -earlier solo-serving measurements showed that this +automatic policy also retries the native representation when a family-specific +strict loader rejects the optional transform. Explicit affine modes remain +fail-closed. The earlier solo-serving measurements showed that this load-time representation has the same steady-state generation speed as loading an equivalent pre-quantized artifact; see `../../spikes/mlx-solo/FINDINGS.md`. @@ -45,6 +47,15 @@ The 135M model is adequate as a serving and protocol oracle, but its weak agent output is not evidence of Goose-quality model behavior. Larger single-node models still need quality, memory-high-water, and agent-harness measurements. +The integrated path was therefore also exercised with `Qwen/Qwen3-0.6B` at a +16K context. Its redundant tied `lm_head.weight` is incompatible with the +pinned strict affine loader, so auto mode reported that incompatibility and +retried the native checkpoint representation. The model then answered a basic +arithmetic prompt correctly and completed a Goose OpenAI-provider run with the +requested exact response. This proves the adaptive fallback preserves useful +single-node serving rather than turning optional quantization into a startup +requirement. + ## Verified mesh split proof A two-host explicit split of the same 30-layer model completed through the diff --git a/crates/skippy-engine-mlx/WIRING.md b/crates/skippy-engine-mlx/WIRING.md index 3b15040b2a..f5077a7de9 100644 --- a/crates/skippy-engine-mlx/WIRING.md +++ b/crates/skippy-engine-mlx/WIRING.md @@ -57,6 +57,8 @@ the GGUF path. The automatic weight policy is intentionally conservative: - Inkling and Nemotron-H load their native representation because their routed rank-3 experts do not support that transform; - checkpoints already declaring quantization/compression load natively; +- auto retries native loading when a family-specific strict/quantization check + rejects the optional transform; - explicit `mlx-serve` users may choose auto, none, affine4, affine8, or mxfp4. The model is exposed through the shared `openai-frontend` router. The HTTP diff --git a/crates/skippy-engine-mlx/src/bin/mlx-serve.rs b/crates/skippy-engine-mlx/src/bin/mlx-serve.rs index 02e5c1a71c..59c083bcb7 100644 --- a/crates/skippy-engine-mlx/src/bin/mlx-serve.rs +++ b/crates/skippy-engine-mlx/src/bin/mlx-serve.rs @@ -87,12 +87,15 @@ mod real { } explicit => explicit.engine(), }; + let allow_native_quantization_fallback = + matches!(cli.weight_quantization, WeightQuantization::Auto); let config = MlxEngineConfig { model_dir: cli.model.clone(), model_id: model_id.clone(), default_max_tokens: cli.default_max_tokens, max_tokens_cap: cli.max_tokens_cap, weight_quantization, + allow_native_quantization_fallback, }; tracing::info!("loading MLX model from {} ...", cli.model.display()); diff --git a/crates/skippy-engine-mlx/src/engine.rs b/crates/skippy-engine-mlx/src/engine.rs index b18f6608dc..54a13fc9f3 100644 --- a/crates/skippy-engine-mlx/src/engine.rs +++ b/crates/skippy-engine-mlx/src/engine.rs @@ -39,6 +39,9 @@ pub struct MlxEngineConfig { /// Quantize eligible dense checkpoint tensors as they are loaded. Already /// quantized checkpoints with matching metadata load directly. pub weight_quantization: Option, + /// Auto policy may retry native loading when a model-specific strict or + /// quantization check rejects the optional transform. + pub allow_native_quantization_fallback: bool, } /// Selects load-time affine-4 only for dense, unquantized model families that @@ -198,17 +201,47 @@ fn load_engine(config: &MlxEngineConfig) -> Result { ModelLoadOptions::default, ModelLoadOptions::with_quantization, ); - let model = - LoadedModel::load_with_options(&config.model_dir, options, &stream, &weights_stream) - .map_err(|e| anyhow!("load {}: {e}", config.model_dir.display()))?; + let mut used_native_fallback = false; + let model = match LoadedModel::load_with_options( + &config.model_dir, + options, + &stream, + &weights_stream, + ) { + Ok(model) => model, + Err(error) + if config.allow_native_quantization_fallback + && config.weight_quantization.is_some() + && optional_quantization_incompatible(&error) => + { + used_native_fallback = true; + tracing::warn!( + model = %config.model_id, + %error, + "MLX automatic quantization is incompatible; retrying native checkpoint representation" + ); + let _ = stream.synchronize(); + LoadedModel::load(&config.model_dir, &stream, &weights_stream).map_err(|native_error| { + anyhow!( + "load {} with automatic quantization failed ({error}); native fallback also failed: {native_error}", + config.model_dir.display() + ) + })? + } + Err(error) => return Err(anyhow!("load {}: {error}", config.model_dir.display())), + }; stream.synchronize().map_err(|e| anyhow!("sync: {e}"))?; let tokenizer = tokenizers::Tokenizer::from_file(config.model_dir.join("tokenizer.json")) .map_err(|e| anyhow!("tokenizer.json: {e}"))?; let eos = model.eos_token_ids().to_vec(); - let quantization_label = config - .weight_quantization - .map_or_else(|| "checkpoint".to_string(), MlxWeightQuantization::label); + let quantization_label = if used_native_fallback { + "checkpoint-fallback".to_string() + } else { + config + .weight_quantization + .map_or_else(|| "checkpoint".to_string(), MlxWeightQuantization::label) + }; tracing::info!( model = %config.model_id, @@ -225,6 +258,14 @@ fn load_engine(config: &MlxEngineConfig) -> Result { }) } +fn optional_quantization_incompatible(error: &safemlx_lm::error::Error) -> bool { + matches!( + error, + safemlx_lm::error::Error::Quantization(_) + | safemlx_lm::error::Error::StrictLoadValidation { .. } + ) +} + fn run_worker( config: MlxEngineConfig, mut job_rx: mpsc::UnboundedReceiver, @@ -415,4 +456,20 @@ mod tests { None ); } + + #[test] + fn automatic_quantization_retries_only_compatibility_failures() { + assert!(optional_quantization_incompatible( + &safemlx_lm::error::Error::Quantization("unsupported".to_string()) + )); + assert!(optional_quantization_incompatible( + &safemlx_lm::error::Error::StrictLoadValidation { + missing: Vec::new(), + unused: vec!["lm_head.weight".to_string()], + } + )); + assert!(!optional_quantization_incompatible( + &safemlx_lm::error::Error::UnsupportedArchitecture("unsupported".to_string()) + )); + } } From 41c4fbf2a9b39d7a5261e994c4bf9f1aa2f85aec Mon Sep 17 00:00:00 2001 From: Michael Neale <14976+michaelneale@users.noreply.github.com> Date: Sat, 18 Jul 2026 03:09:01 +1000 Subject: [PATCH 35/37] fix(mlx): fail closed on strict load errors --- .../SERVE_INTEGRATION_STATUS.md | 7 +- crates/skippy-engine-mlx/WIRING.md | 5 +- crates/skippy-engine-mlx/src/engine.rs | 100 +++++++++++++++--- 3 files changed, 91 insertions(+), 21 deletions(-) diff --git a/crates/skippy-engine-mlx/SERVE_INTEGRATION_STATUS.md b/crates/skippy-engine-mlx/SERVE_INTEGRATION_STATUS.md index eb87c4eeea..126fa46b1a 100644 --- a/crates/skippy-engine-mlx/SERVE_INTEGRATION_STATUS.md +++ b/crates/skippy-engine-mlx/SERVE_INTEGRATION_STATUS.md @@ -37,9 +37,10 @@ does not prove partial downloads: a whole-model server needs all model weights. The integrated loader now quantizes eligible, unquantized dense source tensors to affine 4-bit as they load. Inkling, Nemotron-H, and checkpoints already declaring a quantized representation retain their native representation. The -automatic policy also retries the native representation when a family-specific -strict loader rejects the optional transform. Explicit affine modes remain -fail-closed. The earlier solo-serving measurements showed that this +automatic policy also retries the native representation for quantization +incompatibility or the known benign tied-Qwen `lm_head.weight` rejection. Other +strict-loader failures remain fail-closed, as do explicit affine modes. The +earlier solo-serving measurements showed that this load-time representation has the same steady-state generation speed as loading an equivalent pre-quantized artifact; see `../../spikes/mlx-solo/FINDINGS.md`. diff --git a/crates/skippy-engine-mlx/WIRING.md b/crates/skippy-engine-mlx/WIRING.md index f5077a7de9..93626e8798 100644 --- a/crates/skippy-engine-mlx/WIRING.md +++ b/crates/skippy-engine-mlx/WIRING.md @@ -57,8 +57,9 @@ the GGUF path. The automatic weight policy is intentionally conservative: - Inkling and Nemotron-H load their native representation because their routed rank-3 experts do not support that transform; - checkpoints already declaring quantization/compression load natively; -- auto retries native loading when a family-specific strict/quantization check - rejects the optional transform; +- auto retries native loading for quantization incompatibility or the known + benign tied-Qwen `lm_head.weight` strict-loader rejection; every other strict + validation failure remains fail-closed; - explicit `mlx-serve` users may choose auto, none, affine4, affine8, or mxfp4. The model is exposed through the shared `openai-frontend` router. The HTTP diff --git a/crates/skippy-engine-mlx/src/engine.rs b/crates/skippy-engine-mlx/src/engine.rs index 54a13fc9f3..f6c632e9c4 100644 --- a/crates/skippy-engine-mlx/src/engine.rs +++ b/crates/skippy-engine-mlx/src/engine.rs @@ -39,8 +39,8 @@ pub struct MlxEngineConfig { /// Quantize eligible dense checkpoint tensors as they are loaded. Already /// quantized checkpoints with matching metadata load directly. pub weight_quantization: Option, - /// Auto policy may retry native loading when a model-specific strict or - /// quantization check rejects the optional transform. + /// Auto policy may retry native loading for a quantization incompatibility + /// or a narrowly recognized benign strict-loader rejection. pub allow_native_quantization_fallback: bool, } @@ -212,7 +212,7 @@ fn load_engine(config: &MlxEngineConfig) -> Result { Err(error) if config.allow_native_quantization_fallback && config.weight_quantization.is_some() - && optional_quantization_incompatible(&error) => + && optional_quantization_incompatible(&config.model_dir, &error) => { used_native_fallback = true; tracing::warn!( @@ -258,12 +258,32 @@ fn load_engine(config: &MlxEngineConfig) -> Result { }) } -fn optional_quantization_incompatible(error: &safemlx_lm::error::Error) -> bool { - matches!( - error, - safemlx_lm::error::Error::Quantization(_) - | safemlx_lm::error::Error::StrictLoadValidation { .. } - ) +fn optional_quantization_incompatible(model_dir: &Path, error: &safemlx_lm::error::Error) -> bool { + match error { + safemlx_lm::error::Error::Quantization(_) => true, + safemlx_lm::error::Error::StrictLoadValidation { missing, unused } => { + known_tied_qwen_lm_head_failure(model_dir, missing, unused) + } + _ => false, + } +} + +fn known_tied_qwen_lm_head_failure( + model_dir: &Path, + missing: &[String], + unused: &[String], +) -> bool { + if !missing.is_empty() || unused != ["lm_head.weight"] { + return false; + } + let Ok(bytes) = fs::read(model_dir.join("config.json")) else { + return false; + }; + let Ok(config) = serde_json::from_slice::(&bytes) else { + return false; + }; + config.get("model_type").and_then(Value::as_str) == Some("qwen3") + && config.get("tie_word_embeddings").and_then(Value::as_bool) == Some(true) } fn run_worker( @@ -458,18 +478,66 @@ mod tests { } #[test] - fn automatic_quantization_retries_only_compatibility_failures() { + fn automatic_quantization_retries_quantization_failures() { assert!(optional_quantization_incompatible( + Path::new("unused"), &safemlx_lm::error::Error::Quantization("unsupported".to_string()) )); - assert!(optional_quantization_incompatible( - &safemlx_lm::error::Error::StrictLoadValidation { - missing: Vec::new(), - unused: vec!["lm_head.weight".to_string()], - } - )); assert!(!optional_quantization_incompatible( + Path::new("unused"), &safemlx_lm::error::Error::UnsupportedArchitecture("unsupported".to_string()) )); } + + #[test] + fn automatic_quantization_retries_only_known_tied_qwen_strict_failure() { + let temp = tempfile::tempdir().unwrap(); + fs::write( + temp.path().join("config.json"), + br#"{"model_type":"qwen3","tie_word_embeddings":true}"#, + ) + .unwrap(); + let error = safemlx_lm::error::Error::StrictLoadValidation { + missing: Vec::new(), + unused: vec!["lm_head.weight".to_string()], + }; + assert!(optional_quantization_incompatible(temp.path(), &error)); + + let missing_core = safemlx_lm::error::Error::StrictLoadValidation { + missing: vec!["model.layers.0.self_attn.q_proj.weight".to_string()], + unused: vec!["lm_head.weight".to_string()], + }; + assert!(!optional_quantization_incompatible( + temp.path(), + &missing_core + )); + let unexpected_unused = safemlx_lm::error::Error::StrictLoadValidation { + missing: Vec::new(), + unused: vec![ + "lm_head.weight".to_string(), + "unexpected.weight".to_string(), + ], + }; + assert!(!optional_quantization_incompatible( + temp.path(), + &unexpected_unused + )); + let shape_mismatch = safemlx_lm::error::Error::StrictLoadValidation { + missing: vec!["model.layers.0.mlp.down_proj.weight".to_string()], + unused: vec![ + "model.layers.0.mlp.down_proj.weight -> model.layers.0.mlp.down_proj.weight: expected [1024, 3072], got [1024, 2048]".to_string(), + ], + }; + assert!(!optional_quantization_incompatible( + temp.path(), + &shape_mismatch + )); + + fs::write( + temp.path().join("config.json"), + br#"{"model_type":"llama","tie_word_embeddings":true}"#, + ) + .unwrap(); + assert!(!optional_quantization_incompatible(temp.path(), &error)); + } } From c4ad58fa0c3c8a841f727534ad8331358d3291c8 Mon Sep 17 00:00:00 2001 From: Michael Neale <14976+michaelneale@users.noreply.github.com> Date: Sat, 18 Jul 2026 09:54:38 +1000 Subject: [PATCH 36/37] fix(mlx): serve standard prequantized checkpoints --- Cargo.lock | 12 +++++----- Cargo.toml | 6 ++--- crates/skippy-engine-mlx/Cargo.toml | 6 +++-- .../SERVE_INTEGRATION_STATUS.md | 24 +++++++++++++++++-- crates/skippy-engine-mlx/WIRING.md | 18 ++++++++++---- crates/skippy-engine-mlx/src/derived.rs | 2 +- crates/skippy-engine-mlx/src/engine.rs | 14 +++++++++++ docs/design/MLX_STAGE_ENGINE_PLAN.md | 19 +++++++++------ spikes/mlx-solo/FINDINGS.md | 3 ++- 9 files changed, 78 insertions(+), 26 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 0341f13e28..37a8ac5423 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -7094,7 +7094,7 @@ checksum = "3944826ff8fa8093089aba3acb4ef44b9446a99a16f3bf4e74af3f77d340ab7d" [[package]] name = "safemlx" version = "0.1.3" -source = "git+https://github.com/jbg/safemlx?rev=4e53c5e#4e53c5ecd7cbd91c0dfd0992a3c731ca2c36e9c7" +source = "git+https://github.com/michaelneale/safemlx?rev=c6b4741#c6b47418f3ea0e7b304464a80d8bc8f63f3bbc22" dependencies = [ "bytemuck", "dyn-clone", @@ -7118,7 +7118,7 @@ dependencies = [ [[package]] name = "safemlx-internal-macros" version = "0.1.1" -source = "git+https://github.com/jbg/safemlx?rev=4e53c5e#4e53c5ecd7cbd91c0dfd0992a3c731ca2c36e9c7" +source = "git+https://github.com/michaelneale/safemlx?rev=c6b4741#c6b47418f3ea0e7b304464a80d8bc8f63f3bbc22" dependencies = [ "darling 0.23.0", "itertools 0.15.0", @@ -7130,7 +7130,7 @@ dependencies = [ [[package]] name = "safemlx-lm" version = "0.4.1" -source = "git+https://github.com/jbg/safemlx?rev=4e53c5e#4e53c5ecd7cbd91c0dfd0992a3c731ca2c36e9c7" +source = "git+https://github.com/michaelneale/safemlx?rev=c6b4741#c6b47418f3ea0e7b304464a80d8bc8f63f3bbc22" dependencies = [ "anyhow", "clap", @@ -7149,7 +7149,7 @@ dependencies = [ [[package]] name = "safemlx-lm-utils" version = "0.1.4" -source = "git+https://github.com/jbg/safemlx?rev=4e53c5e#4e53c5ecd7cbd91c0dfd0992a3c731ca2c36e9c7" +source = "git+https://github.com/michaelneale/safemlx?rev=c6b4741#c6b47418f3ea0e7b304464a80d8bc8f63f3bbc22" dependencies = [ "minijinja", "minijinja-contrib", @@ -7162,7 +7162,7 @@ dependencies = [ [[package]] name = "safemlx-macros" version = "0.1.1" -source = "git+https://github.com/jbg/safemlx?rev=4e53c5e#4e53c5ecd7cbd91c0dfd0992a3c731ca2c36e9c7" +source = "git+https://github.com/michaelneale/safemlx?rev=c6b4741#c6b47418f3ea0e7b304464a80d8bc8f63f3bbc22" dependencies = [ "darling 0.23.0", "proc-macro2", @@ -7173,7 +7173,7 @@ dependencies = [ [[package]] name = "safemlx-sys" version = "0.1.3" -source = "git+https://github.com/jbg/safemlx?rev=4e53c5e#4e53c5ecd7cbd91c0dfd0992a3c731ca2c36e9c7" +source = "git+https://github.com/michaelneale/safemlx?rev=c6b4741#c6b47418f3ea0e7b304464a80d8bc8f63f3bbc22" dependencies = [ "cc", "cmake", diff --git a/Cargo.toml b/Cargo.toml index 56cad99f74..a9ba500d7e 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -83,9 +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/jbg/safemlx", rev = "4e53c5e" } -safemlx-lm = { git = "https://github.com/jbg/safemlx", rev = "4e53c5e" } -safemlx-lm-utils = { git = "https://github.com/jbg/safemlx", rev = "4e53c5e" } +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" diff --git a/crates/skippy-engine-mlx/Cargo.toml b/crates/skippy-engine-mlx/Cargo.toml index c02166637d..b6b6f7b976 100644 --- a/crates/skippy-engine-mlx/Cargo.toml +++ b/crates/skippy-engine-mlx/Cargo.toml @@ -5,7 +5,8 @@ # # The published manifest uses crates.io versions. This repository patches them # to a proven public safemlx commit until an equivalent release is available. -# See WIRING.md. +# Root patches do not propagate, so standalone published consumers cannot use +# the `mlx` feature until those safemlx APIs are released. See WIRING.md. [package] name = "skippy-engine-mlx" version.workspace = true @@ -81,7 +82,8 @@ tracing = "0.1" tracing-subscriber = { version = "0.3", features = ["env-filter"] } # MLX stack. The workspace root patches these versions to the proven public -# commit; published consumers see ordinary registry dependencies. +# commit. The registry versions do not yet provide all APIs used by this +# feature, so this is currently a workspace/source-build integration. safemlx = { version = "0.1.3", default-features = false, features = [ "accelerate", "metal", diff --git a/crates/skippy-engine-mlx/SERVE_INTEGRATION_STATUS.md b/crates/skippy-engine-mlx/SERVE_INTEGRATION_STATUS.md index 126fa46b1a..c792a7576d 100644 --- a/crates/skippy-engine-mlx/SERVE_INTEGRATION_STATUS.md +++ b/crates/skippy-engine-mlx/SERVE_INTEGRATION_STATUS.md @@ -57,6 +57,24 @@ requested exact response. This proves the adaptive fallback preserves useful single-node serving rather than turning optional quantization into a startup requirement. +The ordinary command also now accepts an unchanged published MLX-LM artifact: + +```bash +mesh-llm serve --model mlx-community/Qwen3-0.6B-4bit --ctx-size 16384 +``` + +The recorded run resolved Hugging Face revision +`73e3e38d981303bc594367cd910ea6eb48349da8`; the unpinned command above follows +the repository's current default revision. + +The repository's omitted `quantization.mode` is interpreted as MLX-LM's +standard `affine` default by the pinned safemlx correction proposed upstream in +`jbg/safemlx#2`. The integrated server loaded the cached 4-bit checkpoint in +about 1.2 seconds, advertised it through `/v1/models`, returned `391` for +`23 * 17`, completed SSE with `[DONE]`, and served Goose through the same +OpenAI endpoint. A 4K-context Goose attempt was correctly rejected because its +assembled request required about 11.2K tokens; the 16K run completed. + ## Verified mesh split proof A two-host explicit split of the same 30-layer model completed through the @@ -106,8 +124,10 @@ bytes evenly by layer count. - Automatic affine-4 is the current default for eligible dense checkpoints, not yet a general hardware/quality policy surface. - Cache capacity and eviction are not yet owned by this integration. -- The publishable crate uses registry requirements while the workspace root - patches them to the certified public safemlx revision. +- The workspace root must patch the registry requirements to the certified + public safemlx revision. Root patches do not propagate to crates.io + consumers, so standalone published `skippy-engine-mlx --features mlx` support + is not usable until compatible safemlx releases contain the required APIs. - Frontier families require their own stage semantics. Nemotron-H has metadata/range planning and a one-layer execution proof, but not a complete hybrid-model topology. Inkling has whole-model support in the pinned safemlx diff --git a/crates/skippy-engine-mlx/WIRING.md b/crates/skippy-engine-mlx/WIRING.md index 93626e8798..bb417ddf45 100644 --- a/crates/skippy-engine-mlx/WIRING.md +++ b/crates/skippy-engine-mlx/WIRING.md @@ -15,10 +15,20 @@ The crate manifest uses ordinary crates.io requirements: - `safemlx-lm-utils = 0.1.4` The repository root patches those packages to public safemlx commit -`4e53c5e`. That commit contains the loader/model behavior certified by this -draft; the published releases do not yet contain every required correction. -Using `[patch.crates-io]` keeps the workspace reproducibly pinned without -putting illegal git dependencies in the published crate manifest. +`c6b4741`, based on upstream `4e53c5e`. It adds MLX-LM-compatible handling for +published checkpoints that omit `quantization.mode` and therefore imply +`affine`; the correction is proposed upstream in `jbg/safemlx#2`. The published +releases do not yet contain every required correction. Using +`[patch.crates-io]` keeps the workspace reproducibly pinned without putting +illegal git dependencies in the published crate manifest. + +This patch is required, not optional downstream polish: Cargo root patches do +not propagate through published crates, and the current registry +`safemlx-lm = 0.4.1` lacks APIs used by the MLX engine. Consequently the +workspace and source-built `mesh-llm` MLX feature are usable, while standalone +published `skippy-engine-mlx --features mlx` consumers must wait for compatible +safemlx releases. Keeping the crate in the publish graph does not certify that +feature shape yet. `skippy-engine-mlx` is therefore part of `scripts/publish-crates.sh` after its workspace dependencies and before `mesh-llm-host-runtime`. Both diff --git a/crates/skippy-engine-mlx/src/derived.rs b/crates/skippy-engine-mlx/src/derived.rs index 3754d22584..73e43e1796 100644 --- a/crates/skippy-engine-mlx/src/derived.rs +++ b/crates/skippy-engine-mlx/src/derived.rs @@ -46,7 +46,7 @@ pub use nemotron_h::{MlxNemotronHValidationReport, validate_nemotron_h_moe_stage pub(super) const DERIVED_STAGE_SCHEMA_VERSION: u32 = 1; const DERIVED_STAGE_IMPLEMENTATION: &str = "mesh-mlx-range-derived-v1"; -const SAFEMLX_REVISION: &str = "4e53c5ecd7cbd91c0dfd0992a3c731ca2c36e9c7"; +const SAFEMLX_REVISION: &str = "c6b47418f3ea0e7b304464a80d8bc8f63f3bbc22"; const PLAN_FILE: &str = "stage-plan.json"; pub(super) const REPORT_FILE: &str = "derived-stage.json"; static DERIVED_SEQUENCE: AtomicU64 = AtomicU64::new(0); diff --git a/crates/skippy-engine-mlx/src/engine.rs b/crates/skippy-engine-mlx/src/engine.rs index f6c632e9c4..43ba7b28b5 100644 --- a/crates/skippy-engine-mlx/src/engine.rs +++ b/crates/skippy-engine-mlx/src/engine.rs @@ -477,6 +477,20 @@ mod tests { ); } + #[test] + fn published_mlx_lm_quantization_without_mode_is_supported() { + let quantization = serde_json::from_value::( + json!({"group_size": 64, "bits": 4}), + ) + .unwrap(); + assert_eq!( + quantization, + safemlx_lm::quantization::WeightQuantization::Affine( + safemlx_lm::quantization::AffineQuantization::new(64, 4).unwrap() + ) + ); + } + #[test] fn automatic_quantization_retries_quantization_failures() { assert!(optional_quantization_incompatible( diff --git a/docs/design/MLX_STAGE_ENGINE_PLAN.md b/docs/design/MLX_STAGE_ENGINE_PLAN.md index 03d10f7e9b..04c5d8f54b 100644 --- a/docs/design/MLX_STAGE_ENGINE_PLAN.md +++ b/docs/design/MLX_STAGE_ENGINE_PLAN.md @@ -987,7 +987,9 @@ Spikes 1 and 2 are more decisive than any standalone token/s benchmark. Qwen3-0.6B alone: (1) the published crate hard-enables the `metal` feature, so a Metal-less/CI build needs a **workspace-level** `default-features = false`; (2) tied-embedding `lm_head.weight` fails the *quantized* strict loader - (dense load tolerates it). Both fixed in the fork; expect more per-family. + (dense load tolerates it). The pinned fork fixes the omitted-mode loader gap; + mesh handles the exact tied-head rejection with a narrow native-load fallback. + Expect more per-family. - **Recurrent/hybrid + MoE** splitting is materially harder than dense; scope them out of early phases. - **Two artifact pipelines** add storage + certification cost; mitigate with a @@ -999,12 +1001,15 @@ Spikes 1 and 2 are more decisive than any standalone token/s benchmark. a fixed version without bumping. A fork-free build against published crates **compiled and ran but produced gibberish for Qwen3 source precision and crashed on a pre-quantized repo** (`rms_norm` size mismatch) — the working - dense-Qwen3/Llama + JIT-quant code exists only in unpublished fork HEAD. So - The published `skippy-engine-mlx` manifest therefore uses normal registry - requirements while this workspace patches them to a **specific public git - commit** of `jbg/safemlx`. A future safemlx release can remove that root - patch. This makes "track upstream + certify + patch" a **standing cost**, not - a one-off. + dense-Qwen3/Llama + JIT-quant code exists only in unpublished fork HEAD. The + `skippy-engine-mlx` manifest uses normal registry requirements while this + workspace patches them to a **specific public git commit** based on + `jbg/safemlx`. Root patches do not propagate to crates.io consumers, and the + registry release lacks APIs used by the engine, so standalone published MLX + consumers are not usable yet. The current compatibility fix is proposed + upstream in `jbg/safemlx#2`; compatible future safemlx releases can remove + the root patch and unblock that feature shape. This makes "track upstream + + certify + patch" a **standing cost**, not a one-off. - **Hardware coverage is a moving target with two gates.** New backends must land in upstream `ml-explore/mlx` *then* be wired through safemlx (which authors no backends itself). ROCm is an active-but-unmerged upstream experiment (#2300); diff --git a/spikes/mlx-solo/FINDINGS.md b/spikes/mlx-solo/FINDINGS.md index 379b63b949..9fc55e8832 100644 --- a/spikes/mlx-solo/FINDINGS.md +++ b/spikes/mlx-solo/FINDINGS.md @@ -149,4 +149,5 @@ Still open: Skippy's F16/F32 binary activation codec; see `../mlx-safetensors-stages/FINDINGS.md`. Two `skippy-server` processes, the boundary-fence benchmark, and bounded-memory load-time quantization remain. -- Upstreaming the two fixes to `jbg/safemlx` (or carrying a thin fork). +- The omitted-mode compatibility fix is proposed in `jbg/safemlx#2`; the + tied-`lm_head.weight` quantized-loader fix still needs upstream resolution. From 75fb8b2411238c9cb3c7a37db98b6a5368670f9d Mon Sep 17 00:00:00 2001 From: Michael Neale <14976+michaelneale@users.noreply.github.com> Date: Sat, 18 Jul 2026 09:54:41 +1000 Subject: [PATCH 37/37] fix(models): ignore bit width in parameter labels --- .../src/models/profile.rs | 22 ++++++++++++++----- 1 file changed, 16 insertions(+), 6 deletions(-) diff --git a/crates/mesh-llm-host-runtime/src/models/profile.rs b/crates/mesh-llm-host-runtime/src/models/profile.rs index 8730f7e8f0..bc6d3f1b06 100644 --- a/crates/mesh-llm-host-runtime/src/models/profile.rs +++ b/crates/mesh-llm-host-runtime/src/models/profile.rs @@ -77,10 +77,11 @@ fn quant_from_text(value: &str) -> Option { } fn parameter_size_from_text(text: &str) -> Option { - static MULTIPLIED_RE: LazyLock = - LazyLock::new(|| Regex::new(r"(?i)(\d+(?:\.\d+)?)x(\d+(?:\.\d+)?)([bm])").unwrap()); + static MULTIPLIED_RE: LazyLock = LazyLock::new(|| { + Regex::new(r"(?i)(\d+(?:\.\d+)?)x(\d+(?:\.\d+)?)([bm])(?:$|[^a-z])").unwrap() + }); static SIMPLE_RE: LazyLock = - LazyLock::new(|| Regex::new(r"(?i)(\d+(?:\.\d+)?)([bm])").unwrap()); + LazyLock::new(|| Regex::new(r"(?i)(\d+(?:\.\d+)?)([bm])(?:$|[^a-z])").unwrap()); MULTIPLIED_RE .captures(text) @@ -100,10 +101,11 @@ fn parameter_size_from_text(text: &str) -> Option { } fn parameter_count_b_from_text(text: &str) -> Option { - static MULTIPLIED_RE: LazyLock = - LazyLock::new(|| Regex::new(r"(?i)(\d+(?:\.\d+)?)x(\d+(?:\.\d+)?)([bm])").unwrap()); + static MULTIPLIED_RE: LazyLock = LazyLock::new(|| { + Regex::new(r"(?i)(\d+(?:\.\d+)?)x(\d+(?:\.\d+)?)([bm])(?:$|[^a-z])").unwrap() + }); static SIMPLE_RE: LazyLock = - LazyLock::new(|| Regex::new(r"(?i)(\d+(?:\.\d+)?)([bm])").unwrap()); + LazyLock::new(|| Regex::new(r"(?i)(\d+(?:\.\d+)?)([bm])(?:$|[^a-z])").unwrap()); let mut best: Option = None; for captures in MULTIPLIED_RE.captures_iter(text) { @@ -154,6 +156,10 @@ mod tests { parameter_size_from_text("mixtral-8x7b").as_deref(), Some("8x7B") ); + assert_eq!( + parameter_size_from_text("mlx-community/Qwen3-0.6B-4bit").as_deref(), + Some("0.6B") + ); } #[test] @@ -161,5 +167,9 @@ mod tests { assert_eq!(parameter_count_b_from_text("Qwen3-32B-Q4_K_M"), Some(32.0)); assert_eq!(parameter_count_b_from_text("mixtral-8x7b"), Some(56.0)); assert_eq!(parameter_count_b_from_text("235B-A22B"), Some(235.0)); + assert_eq!( + parameter_count_b_from_text("mlx-community/Qwen3-0.6B-4bit"), + Some(0.6) + ); } }