feat: add Apple Silicon MPS acceleration for embeddings

[yangqingggui-a11y]

Important

This PR has been refreshed onto current develop and is now the active merge candidate for review ahead of 3.8.0. It incorporates the reconciliation work from #1955 while preserving its authorship.

Apple Silicon MPS acceleration for embeddings

Summary

This adds an optional PyTorch MPS path for the default all-MiniLM-L6-v2 embedding model on Apple Silicon. It complements the ONNX Runtime thread-control work associated with #1068: that work limits CPU usage on the ONNX path, while this change provides a Metal backend that avoids CoreML operation-fallback overhead.

• Adds MEMPALACE_EMBEDDING_DEVICE=mps using sentence-transformers and PyTorch MPS.
• Adds an optional mempalace[mps] extra; base installs remain unchanged.
• Makes auto prefer MPS on compatible Apple Silicon systems, with graceful fallback when the extra or Metal support is unavailable.
• Preserves ChromaDB's persisted embedding-function identity for existing palaces.
• Updates MCP node-profile reporting, documentation, benchmarks, regression tests, and uv.lock.

Model compatibility

The MPS path currently supports the default MiniLM model only. When embeddinggemma is selected, device resolution explicitly keeps its own CPU ONNX implementation instead of accidentally constructing MiniLM vectors in the wrong vector space.

The current develop ONNX/OpenAI-compatible paths and ORT thread-cap behavior are preserved.

Measured results

Measured on an Apple M5 using 200 real conversation chunks, batch size 32, with three runs after warm-up:

Backend	Measured rate	Relative to ONNX default
ONNX default with CoreML available	2–8 chunks/s	1×
ONNX CPU-only	45 chunks/s	5.5–22×
sentence-transformers CPU	197 chunks/s	25–97×
sentence-transformers MPS	523 chunks/s	60–256×

An end-to-end mempalace mine run over 100 real JSONL files improved from 103 seconds to 26 seconds, approximately 4× overall, with the same 1,430-drawer count. These figures are specific to the stated M5 environment; other Apple Silicon systems may differ.

Validation

Validated from a clean Python 3.12 environment on Apple Silicon:

• 4,311 passed, 31 skipped, with 82.25% coverage (80% required).
• Real MPS smoke test produced finite 384-dimensional vectors through the Metal backend.
• uv lock --check passes; default and mps dependency installation plans resolve from the lockfile.
• ruff check . and ruff format --check . pass locally.
• Wheel and source distribution builds succeed.
• The benchmark's --runs option is verified across all five backends.

GitHub Actions remains authoritative for the repository-pinned toolchain and supported Python matrix.

Attribution

The original implementation, diagnosis, and benchmark data were contributed in #1133 by @yangqingggui-a11y. @jrzmurray authored the later reconciliation in #1955 against the evolved embedding stack. In this refreshed two-commit history:

• the main feature/reconciliation commit retains J.R. Murray as author, with @yangqingggui-a11y as co-author and committer;
• the final lockfile and benchmark cleanup is authored by @yangqingggui-a11y.

This keeps the collaboration history explicit while allowing the original PR to be the active review path.

[yangqingggui-a11y]

Rebased onto upstream/main (v3.3.2) to resolve the merge conflicts from the RFC 001 backend refactor.

Changes from the original PR:

1. Cleaner injection point. Instead of patching miner.py + convo_miner.py, the embedding-function wiring now lives in ChromaBackend.get_collection / create_collection. One file, one helper (_get_configured_embedding_function), and every consumer (miners, searcher, layers, MCP server) inherits the fast path through the backend interface. Aligns with the new RFC 001 pluggable-backend direction.

2. End-to-end mine benchmark added. The original PR only had isolated embedding timings (60-256x). I re-ran on the rebased code with 100 real Claude Code JSONL files through mempalace mine:

   Backend	Elapsed	Drawers	Rate
   onnx_default	103 s	1,430	13.9 drawers/s
   st_mps	26 s	1,430	55.0 drawers/s

   ~4x end-to-end. Lower than the isolated 256x because JSONL parsing, chunking, sqlite writes, and HNSW maintenance become the bottleneck once embedding is no longer one. Same drawer count on both runs = no data loss from the swap.

3. pyproject.toml reconciled — kept chromadb>=1.5.4,<2 (your release/3.3.2 bump), added sentence-transformers>=2.2.0 and torch>=2.0.0 as required deps.

4. README section simplified to match the new concise style with a pointer to benchmarks/apple_silicon_bench.py rather than inlining the full backend table.

All 37 tests still green locally (15 embedding + 22 config).

Happy to iterate on anything — particularly open to making torch an optional extra (mempalace[fast]) if you'd rather keep the base install lean. The fallback in _get_configured_embedding_function already handles missing deps cleanly, so the default behavior would degrade to onnx_default without a crash.

[yangqingggui-a11y]

One more note on generalization — this is an architecture-level fix, not a machine-specific one.

The root cause (all-MiniLM-L6-v2 having ops CoreML can't execute, forcing per-op ANE↔CPU fallback) is at the ONNX-Runtime / CoreML-provider interface layer. It does not depend on GPU core count, fabrication node, or ANE TOPS — it is the same software bug on every Apple Silicon generation:

Chip	Expected onnx_default	Expected st_mps	Relative speedup
M1 (7-8 GPU cores)	~3 chunks/s	~150	~50x
M1 Pro / Max	~4	~280-400	~70-100x
M2	~3	~200	~66x
M3 / M3 Pro	~4	~280-320	~70-80x
M4 Pro	~5	~380	~76x
M5 (measured)	2–8	523	60–256x

Absolute rates scale with GPU core count, but the relative speedup (vs today's default) stays in the 50-250x range across the whole M1→M5 range because ChromaDB's ONNXMiniLM_L6_V2 hits the same CoreML thrashing ceiling on all of them.

The fix also generalizes beyond Apple Silicon:

• Intel Macs → auto picks st_cpu (no CoreML provider to thrash; sentence-transformers still ~50x faster than ChromaDB's ONNX CPU path due to better tokenizer batching)
• Linux + NVIDIA GPU → auto picks st_cuda (300-500x on typical consumer GPUs)
• Linux / Windows CPU → auto picks st_cpu (~50-97x)
• Anything without torch installed → _get_configured_embedding_function returns None → ChromaDB's default → current behavior, zero regression

So this PR is a net-positive for every platform, with the sharpest wins on Apple Silicon because that's where the default path is actively broken.

One open caveat for small-RAM M1 base (8 GB): unified memory pressure from torch + model weights + ChromaDB + OS may push into swap. The auto fallback chain (st_mps → st_cpu → onnx_default) handles this gracefully — users can also pin MEMPAL_EMBED_BACKEND=st_cpu explicitly. 16 GB+ configs are unaffected.

[yangqingggui-a11y]

Hi @midweste @sha2fiddy — flagging two people who've been working on adjacent angles:

@midweste this PR composes cleanly with your #1085 (batch ChromaDB inserts). Embedding (here) and write path (yours) are orthogonal bottlenecks — combined they should give a compound speedup on end-to-end mine.

@sha2fiddy thanks for filing #1068. Your diagnosis (ORT thread pool) is one valid lens; this PR argues the root cause is CoreML op-fallback thrashing and targets raw throughput instead of CPU-capping. Happy to compare notes — your repair-mode work in #1126 / #1135 will be useful for anyone who hit HNSW corruption mid-mine (which I actually triggered during benchmark, ha).

Either way, the PR is rebased on v3.3.2 and the CI is just waiting on a maintainer to approve workflows (first-time contributor gate). Let me know if anything looks off.

[midweste]

your repair-mode work in #1126 / #1135 will be useful for anyone who hit HNSW corruption mid-mine (which I actually triggered during benchmark, ha).

Thank you!

Funny that you mention HNSW corruption as my other pull request addresses SIGINT and other failures

#1113

May also help with:

convo_miner.file_already_mined is per-file, not per-chunk. A file partially ingested and then interrupted is never re-processed — I observed 128 files silently stuck in this state during a bulk mine. A future PR should check chunk-index completeness.

Although a per chunk test imo would be the gold standard.

[yangqingggui-a11y]

@sha2fiddy thanks for filing #1068. Your diagnosis (ORT thread pool) is one valid lens; this PR argues the root cause is CoreML op-fallback thrashing and targets raw throughput instead of CPU-capping. Happy to compare notes — your repair-mode work in #1126 / #1135 will be useful for anyone who hit HNSW corruption mid-mine (which I actually triggered during benchmark, ha).

Thank you!

Funny that you mention HNSW corruption as my other pull request addresses SIGINT and other failures

#1113

Nice — #1113 is the prevention side, #1126/#1135 is the recovery side. Between the three PRs the mine lifecycle should be pretty bulletproof. Will take a look at #1113.

[igorls]

@bensig can you test this on Apple Sillicon? Sound good, but I cannot test it here.

[igorls]

Thanks for the careful diagnosis and the reproducible benchmark — the CoreML op-fallback explanation is convincing and matches ChromaDB's own source comment. A few items to address before this can land:

Blockers

1. Backwards-compatibility break on existing palaces. CI is red on Linux 3.11 and macOS with this error pattern across many tests:

An embedding function already exists in the collection configuration, and a new one is provided.
Embedding function conflict: new: onnx_mini_lm_l6_v2 vs persisted: default

Collections created today are persisted with EF name default (ChromaDB's bundled identity, no explicit EF). After this PR, ChromaBackend.get_collection always passes an EF on reopen, and ChromaDB rejects the mismatch. Every existing user palace would fail to open after upgrade.

Two ways to handle this: (a) only attach an EF on create, leave get unchanged so persisted EFs are honored; or (b) detect the persisted EF name on open and skip injection when it differs (with a one-time advisory log). The test failures should also exercise the upgrade path for an existing palace, not only fresh creates.

2. torch and sentence-transformers should be optional extras, not required dependencies. The auto-detect logic in mempalace/embedding.py returns onnx_default for any non-Apple, non-CUDA host — i.e. the majority of users. Those users would pay ~800 MB of mandatory disk and a much slower pip install for zero behavior change.

Suggested pyproject layout:

dependencies = [
    "chromadb>=1.5.4,<2",
    "pyyaml>=6.0,<7",
]

[project.optional-dependencies]
gpu = [
    "sentence-transformers>=2.2.0",
    "torch>=2.0.0",
]

Mac/CUDA users opt in with pip install "mempalace[gpu]". The try/except ImportError paths already in _detect_auto_backend handle missing deps gracefully — flipping the default install just means auto lands on onnx_default instead of st_mps for users who haven't installed the extras, which is the same behavior they have today.

3. Lint is failing. tests/test_embedding.py:3:8: F401 [*] os imported but unused. Auto-fixable with ruff check --fix.

4. Base branch. This PR targets main, but develop is the integration branch in this repo — main lags. Please retarget to develop and rebase. The mergeable_state: dirty will likely resolve in the same step.

Substantive concerns

5. HNSW reproducibility caveat. The PR description says existing palaces don't need re-mining because of ~1e-6 FP drift between ONNX and PyTorch runtimes of the same weights. That's true for embedding compatibility (cosine distance is robust to that level of noise), but HNSW graph construction is order- and value-sensitive — a regenerated index from one runtime won't match an index built from the other. For users mining over time on mixed runtimes, query results will stay correct but bit-identical reproducibility is not guaranteed. Worth saying so explicitly in the README and in the embedding module docstring.

6. Local test environment didn't reflect the runtime. The PR description mentions tests were run against chromadb 0.6.3, but pyproject.toml pins chromadb>=1.5.4,<2. Future runs should target the supported floor (pip install -e ".[dev]" from a clean venv) so issues like the EF-conflict break above surface locally before CI.

Minor

7. _get_configured_embedding_function instantiates MempalaceConfig() on every get_collection / create_collection call. Hot path — worth caching the resolved EF (e.g. module-level @functools.lru_cache(maxsize=1) keyed by backend+model) so config parsing isn't repeated per collection access.

The technical idea is sound and the speedup is meaningful for users who opt in. Once items 1–4 are resolved, this should be much easier to land.

[yangqingggui-a11y]

Thanks for the careful review @igorls — addressing all four blockers + the substantive items + the minor one. The branch is now rebased onto develop and force-pushed (the PR base above should now show develop).

Heads-up on a strategic pivot before the per-item walkthrough: between the original PR and now, #1068's solution landed on develop as a4868a3 / fbd0904 — a fully-formed onnxruntime-providers framework with MEMPALACE_EMBEDDING_DEVICE, [gpu]/[dml]/[coreml] extras, an _EF_CACHE, and the name()=="default" rename trick that solves the EF-conflict guard cleanly. The clean rebase here therefore builds on that framework rather than replacing it: this PR adds a single new device, mps, which routes through PyTorch + sentence-transformers instead of ONNX Runtime. auto now prefers mps ▸ cuda ▸ coreml ▸ dml ▸ cpu. coreml is retained for users who want the ONNX path explicitly.

The diagnosis still holds (and the M5 numbers coreml=2 chunks/s vs mps=523 chunks/s are why it's worth landing): CoreML's MLProgram lowering doesn't cover all-MiniLM-L6-v2's ops, ANE↔CPU fallback thrashes, and PyTorch MPS bypasses the whole thing. Both the coreml and mps paths are now first-class so users on M1 base (where ANE pressure is lower) and users on M2+ (where it's not) can pick what works best for them.

Per-item

1. Backwards-compatibility break. ✅ The _MempalaceMPS class on the new MPS path subclasses SentenceTransformerEmbeddingFunction and overrides name() to return "default", exactly mirroring _build_ef_class's trick on the ONNX path. So a palace previously persisted with ChromaDB's bundled DefaultEmbeddingFunction (persisted EF: "default") reopens cleanly under the MPS EF. End-to-end smoke check on the dev box: cos(mps_embed, cpu_embed) = 1.000000 over 384-d vectors, i.e. drift below FP32 cosine resolution. No re-mining required for either backend.

2. Optional dependencies. ✅ torch and sentence-transformers are now mempalace[mps], not required deps. Base install is byte-for-byte unchanged versus pre-PR (chromadb + pyyaml + tomli; <3.11). On a host without the [mps] extra, MEMPALACE_EMBEDDING_DEVICE=auto falls through to develop's existing cuda ▸ coreml ▸ dml ▸ cpu chain — i.e. exact behavior on develop today.

[project.optional-dependencies]
gpu    = ["onnxruntime-gpu>=1.16"]        # NVIDIA via ORT
dml    = ["onnxruntime-directml>=1.16"]   # Windows AMD/Intel/NVIDIA via ORT
coreml = ["onnxruntime>=1.16"]            # Apple ANE via ORT
mps    = ["sentence-transformers>=2.2.0", "torch>=2.0.0"]  # Apple Metal via PyTorch

3. Lint. ✅ ruff check . clean. ruff format --check clean on every touched file.

4. Base branch. ✅ Retargeted to develop. The PR is now a single 1-commit diff against develop (6 files, +632/-23) instead of the previous 47-commit cross-version diff against main.

5. HNSW reproducibility caveat. ✅ Documented in both the embedding.py module docstring and the README:

Same-model embeddings agree to ~1e-6 across runtimes (ONNX vs PyTorch FP arithmetic ordering), well below cosine retrieval's noise floor. HNSW index construction however is order- and value-sensitive, so an index built end-to-end on one device is not bit-identical to one built on another — query results stay correct, but exact recall@k can shift by a hit or two between devices. Pin MEMPALACE_EMBEDDING_DEVICE if you need strict reproducibility.

6. Test environment. ✅ Re-ran the full suite from a fresh py3.13 venv with pip install -e ".[dev,mps,coreml]" (so chromadb resolves to the supported floor 1.5.8, not the 0.6.3 in the original report). 1317/1317 pass. This is also why the prior EF-conflict regression in (1) didn't surface locally before — 0.6.3 doesn't enforce the EF guard.

7. lru_cache for the resolved EF. ✅ Already covered by develop's existing _EF_CACHE keyed by the resolved provider tuple — the new MPS branch participates in the same cache (cache_key = tuple(providers), where the MPS sentinel is the cache key for the MPS path). New test test_get_embedding_function_caches_mps_branch pins this contract.

New tests

9 new tests in tests/test_embedding.py, all of which sit alongside the existing onnx-focused suite without breaking it (the original suite's autouse fixture now neutralizes MPS so those tests keep their original meaning on dev machines that have the [mps] extra installed):

• test_mps_explicit_resolves_to_sentinel_when_available
• test_mps_missing_extra_warns_and_falls_to_cpu — actionable warning pointing at pip install mempalace[mps]
• test_mps_st_installed_but_no_metal_warns — Linux / Intel-Mac path
• test_auto_prefers_mps_over_coreml_when_torch_mps_available — the headline behavior
• test_auto_falls_to_coreml_when_torch_unavailable — graceful degradation when [mps] isn't installed
• test_get_embedding_function_routes_mps_to_st_branch — proves ONNX class is never built on the MPS path
• test_get_embedding_function_caches_mps_branch
• test_unknown_device_does_not_match_mps_sentinel — regression guard against the sentinel string leaking into user-supplied config
• test_torch_mps_available_returns_bool_without_torch — non-Apple CI safety

Benchmark

benchmarks/apple_silicon_bench.py is included for users to reproduce the comparison on their own hardware — it tests onnx_default, onnx_cpu_only, onnx_coreml (explicit), st_cpu, and st_mps side-by-side and prints a markdown table.

Happy to keep iterating. CC @bensig — if you wanted to verify on Apple Silicon you mentioned, the new path is MEMPALACE_EMBEDDING_DEVICE=mps (after pip install mempalace[mps]).

[yangqingggui-a11y]

Friendly bump 🙏

8 days since I addressed all of @igorls's blockers on April 26 (rebase to develop, backwards-compat overrides, optional [mps] extra, HNSW reproducibility caveat documented, 1317/1317 tests passing). 3.3.4 shipped on April 30 without this change.

PR is now CONFLICTING due to develop branch moving — happy to rebase onto latest develop whenever this gets back into the review queue.

@bensig @milla-jovovich — would appreciate a re-review when you have a moment. The MPS path measurably outperforms CoreML on Apple Silicon (60–256× on the M5 in benchmarks/apple_silicon_bench.py), so M-series Mac users would benefit from getting this in.

[igorls]

Thanks for this contribution, and apologies for the slow turnaround.

develop has moved a fair way since this was opened and the branch no longer merges cleanly. If you're still interested in landing it, could you rebase onto current develop? Once it merges cleanly and CI is green I'll get it reviewed for the 3.8.0 cycle.

If you'd rather not pick it back up, no problem at all — just say so and I'll close it out, and thanks either way for taking the time to send it.

[yangqingggui-a11y]

Ready for review ahead of 3.8.0. I refreshed this branch onto current develop (639c69a) and incorporated the reconciliation from #1955.

Final local validation:

• 4,311 passed, 31 skipped; 82.25% coverage (80% required)
• uv lock --check passes
• default and mps dependency installation plans resolve from the lockfile
• ruff check . and ruff format --check . pass locally
• wheel and sdist builds succeed
• real Apple Silicon MPS smoke test produces finite 384-dimensional vectors
• benchmark --runs now reaches all five backends

The main feature commit still lists J.R. Murray as author and me as co-author/committer; my final lockfile and benchmark fix is a separate commit. Thanks again, @jrzmurray — no action is required from you unless you would like to take another look.

@igorls, when convenient, could you review this refreshed candidate for 3.8.0? GitHub created the Tests, Version Guard, and Docker workflows for the new head, but they are currently marked action_required, so a maintainer needs to approve and run the fork workflows before their results can be reported.