feat(server+app): serve embeddings and reranking from the MTPLX daemon - #212
feat(server+app): serve embeddings and reranking from the MTPLX daemon#212Cyb3rb1ade wants to merge 12 commits into
Conversation
MTPLX serves chat only, so any retrieval-backed setup — RAG, agent
memory — has to keep a second inference server alive next to it just to
answer /v1/embeddings. This adds both retrieval endpoints to the daemon
that is already running.
mtplx serve \
--embedding-model mlx-community/Qwen3-Embedding-8B-4bit-DWQ \
--reranker-model vserifsaglam/Qwen3-Reranker-4B-4bit-MLX
Both flags repeat, take a Hugging Face id or a local path with an
optional REF=served-id alias, and are picked per request via "model".
/v1/embeddings follows the OpenAI shape, /v1/rerank the Cohere/Jina one.
Retrieval models deliberately bypass the MTP path: multi-token
prediction makes next-token decoding cheaper, which means nothing for a
model that returns a vector instead of a token stream. They run the
transformer stack directly — last-token pooling plus L2 normalisation
for embeddings, a softmax over the yes/no logits for reranking. Batches
are padded on the right, because with causal attention a real token
never attends to a later pad token, so every real position stays
bit-identical to an unpadded run; left padding would corrupt it.
Backends are cached by resolved path, so listing one reference as both
an embedder and a reranker loads a single copy of the weights and serves
both roles from it. --retrieval-max-resident caps how many stay in
memory and unloads the least recently used beyond it. Models load on
first request, so an unused endpoint costs nothing.
Additive throughout: with no flag the endpoints answer 404 and the chat
runtime is untouched. /v1/models gains a "capability" field (chat,
embedding, rerank) alongside the existing entry.
The flags are wired through three separate paths, because the server
argv is rebuilt rather than inherited: the CLI parsers for serve and
quickstart, the child argv builder in commands/public.py, and the
server module's own parser. Missing any one of them leaves the endpoints
silently unconfigured on every path but a bare `mtplx serve`.
The macOS app exposes the same settings under Settings → Retrieval
endpoints and persists them to settings.json; ~/.mtplx/config.toml gains
embedding_models, reranker_models and retrieval_max_resident.
Evidence, on an M-series Mac with 48 GB:
- 28 Python tests and 8 Swift tests covering flag parsing, served-id
resolution, the shared-backend guarantee, eviction, settings round
trip, and the HTTP contract including the chat-only 404 path.
- Live against Qwen3-Embedding-8B-4bit-DWQ and Qwen3-Reranker-4B-4bit-MLX
through the real routes: 4096-dim vectors with norm 1.000000, and a
reranker separating a matching from a non-matching document 0.9994 to
0.0000, while /v1/chat/completions kept answering from the same daemon.
- Vectors were compared against the same models served by another local
runtime: worst-case cosine similarity 0.9998 with identical reranker
ordering, so existing indexes do not need re-embedding.
No performance claim is made about the chat path, which is unchanged.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01FqHNxHa5fUW7BRDtq5nghu
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: d9a41c9018
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| others = [ | ||
| ref | ||
| for ref in self._resident | ||
| if ref != spec.model_ref and self._backends[ref].loaded | ||
| ] |
There was a problem hiding this comment.
When two first-use requests select different retrieval models concurrently, each _backend call sees the other backend as not yet loaded and therefore excludes it from others; both requests can then load their models after releasing _lock. With retrieval_max_resident=1, this leaves two large models resident until another request happens to trigger eviction, defeating the memory cap and potentially causing an out-of-memory failure. Treat in-progress backends as occupying a resident slot or serialize slot reservation and loading.
Useful? React with 👍 / 👎.
| registry = RetrievalRegistry( | ||
| max_resident=int(getattr(args, "retrieval_max_resident", DEFAULT_MAX_RESIDENT) or DEFAULT_MAX_RESIDENT), | ||
| cache_dir=getattr(args, "model_dir", None), | ||
| ) |
There was a problem hiding this comment.
For the normal mtplx serve/quickstart subprocess path, the server namespace has no model_dir attribute: the public CLI uses cache_dir to resolve the chat model, then forwards retrieval references without forwarding that directory. Consequently, an embedding or reranker cached under --cache-dir or the configured model_dir cannot be found unless it also exists in the default cache, so its first request fails despite the model being locally available.
Useful? React with 👍 / 👎.
| "data": [ | ||
| { | ||
| "id": state.model_id, | ||
| "object": "model", | ||
| "created": now, | ||
| "owned_by": "mtplx", | ||
| "context_length": state.context_window, | ||
| "max_context_length": state.context_window, | ||
| "max_model_len": state.context_window, | ||
| } | ||
| {"object": "embedding", "index": index, "embedding": vector} | ||
| for index, vector in enumerate(vectors) |
There was a problem hiding this comment.
When an OpenAI-compatible client requests encoding_format: "base64", the handler still returns each embedding as a JSON float array because request.encoding_format is never consulted. Clients that explicitly select base64 will receive a response in the wrong representation and may fail to decode it; either encode the vectors as base64 or reject unsupported formats instead of silently returning floats.
Useful? React with 👍 / 👎.
Three defects from the Codex review on youssofal#212, each with a regression test. Reserve a resident slot on acquisition, not on completed load. Loading happens outside the registry lock, so two concurrent first-use requests for different models each saw the other as not yet loaded, skipped eviction, and both stayed resident — with --retrieval-max-resident 1 that silently doubles peak memory exactly when it is tightest. Every entry in the residency map now holds a slot regardless of load state. Forward the cache directory to retrieval resolution. The chat model is resolved to an absolute path before the server subprocess starts, so that process carries no cache directory at all; retrieval references stay symbolic and were resolved against the default cache. A model pulled into a custom --cache-dir was therefore invisible despite being on disk. The directory now travels with the references via a new --retrieval-cache-dir, and registry_from_args prefers it over the in-process cache_dir and model_dir fallbacks. Honour encoding_format on /v1/embeddings. The field was accepted and ignored, so a client asking for base64 received float arrays and could fail to decode them. base64 now returns the little-endian float32 buffer OpenAI clients expect, float stays the default, and any other value is rejected with 400 rather than silently misinterpreted. Verified live against Qwen3-Embedding-8B-4bit-DWQ through the running daemon: base64 decodes to 4096 dimensions with norm 1.000000, float still returns a JSON array, and float16 is refused with 400. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01FqHNxHa5fUW7BRDtq5nghu
|
Thanks — all three findings were reproducible and are fixed in P1 — resident slots ( P2 — cache directory. Confirmed, and the cause is slightly different from the description, which matters for the fix: the server subprocess has no cache-directory argument at all. The chat model is resolved to an absolute path before launch, so nothing downstream needs one; retrieval references stay symbolic and were resolved against the default cache. Added P2 — Verified live against Python retrieval tests are now 34, all passing, alongside the 8 Swift tests. |
Configuring a retrieval model was unobservable: the settings took, the endpoints answered, and nothing in the app said whether a model had ever loaded, whether it was being used, or what it cost. That is a form to fill in, not an integration. The registry now tracks per-model requests, items, compute time, load time and last use, derives throughput and average latency, and reports them — together with load and residency state — under a "retrieval" key in the dashboard snapshot the app already polls. The key is always present, so a client can tell "no retrieval configured" apart from "this daemon predates retrieval". Load time is measured but deliberately excluded from request latency. Counting it would have made a cold first request read as 10 s average latency on a model that actually serves in 450 ms, which is exactly the number a user would act on. The app decodes the status and renders it in the Activity tab: per model the role, whether it is loaded and resident, throughput in items per second, average latency, cumulative work, load duration, and any last error. Retrieval models load on first request, so "configured but never loaded" is a normal state and is labelled as such rather than being shown as broken or as working. Measured on the running daemon: 13.4 texts/s at 447 ms per request with the 6.0 s weight load reported separately, and an unused reranker correctly shown as not loaded. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01FqHNxHa5fUW7BRDtq5nghu
|
@codex review Two commits landed since the first pass: the three findings from that review are fixed in |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 3c6c2dbf9f
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| while len(self._resident) > self.max_resident: | ||
| oldest = next(iter(self._resident)) | ||
| if oldest == spec.model_ref: | ||
| break | ||
| self._backends[oldest].unload() |
There was a problem hiding this comment.
When two first-use requests target different models with max_resident=1, the first request can leave _backend() before entering ensure_loaded() or the inference lock; the second request then unloads and removes that backend, after which the first still loads it through its retained reference. Similarly, eviction can clear _model after ensure_loaded() returns while the request retains model_obj. Both models can therefore remain live and run concurrently even though only one appears in _resident, defeating the memory cap precisely under concurrent cold traffic. Keep an acquisition pinned until the request finishes, or serialize eviction against the entire load-and-inference lifetime.
Useful? React with 👍 / 👎.
|
|
||
| def _backend(self, spec: RetrievalSpec) -> _Backend: | ||
| with self._lock: | ||
| backend = self._backends.get(spec.model_ref) |
There was a problem hiding this comment.
When two configured references resolve to the same filesystem directory—for example a Hugging Face id and its local path or two aliases—the lookup uses the raw model_ref, so separate _Backend objects are created and the weights are loaded and counted as resident twice. This contradicts the resolved-path sharing guarantee and can waste several GiB; resolve first and use a canonical path as the backend and residency key.
Useful? React with 👍 / 👎.
| self.items += items | ||
| self.compute_seconds += seconds | ||
| self.last_used_s = time.time() | ||
| self.last_error = None |
There was a problem hiding this comment.
If model resolution/loading, token validation, or a forward pass fails, embed() and rerank() exit before record() and no path ever assigns last_error; the only assignment clears it after a successful request. Consequently the new dashboard reports the failed model as unused or idle and its “Last error” row can never appear. Record the exception before re-raising it, while retaining this clear-on-success behavior.
Useful? React with 👍 / 👎.
| for index, vector in enumerate(vectors) | ||
| ], | ||
| "model": spec.served_id, | ||
| "usage": {"prompt_tokens": 0, "total_tokens": 0}, |
There was a problem hiding this comment.
Every successful non-empty embedding request reports zero prompt and total tokens, and the rerank response likewise always reports zero. Clients commonly use these OpenAI/Cohere-compatible usage fields for accounting and limits, so this silently produces incorrect totals for all retrieval traffic; return the token counts already produced during input encoding rather than fixed zeros.
Useful? React with 👍 / 👎.
… tokens Four findings from the second Codex review on youssofal#212, each with a regression test. Pin an acquired backend for its whole load-and-inference lifetime. Reserving a residency slot was not enough: a request that had left the acquisition path but not finished still held its weights, so evicting them freed nothing and forced an immediate reload — leaving two models live while only one was counted. Eviction now skips pinned backends and briefly exceeds the cap instead of unloading a model in use. Key backends by resolved filesystem path. The lookup used the raw reference, so a Hugging Face id and the local path it resolves to produced two backends and loaded the same weights twice — contradicting the sharing guarantee this module documents. Resolution is cached per reference, and descriptors resolve through the same key, which also fixes load state being reported against a backend shared with an alias. Record failures in the statistics. Nothing ever assigned last_error; the only assignment cleared it after success, so a model that failed to resolve or load was displayed as merely idle and the dashboard's error row could never appear. Failures are now recorded before re-raising, and still cleared by the next success. Report real token counts. Embedding and rerank responses returned a fixed zero usage, which silently corrupts the totals clients keep for accounting and limits. The counts are taken from the encoding already performed for each batch. Verified live: embeddings report 14 prompt tokens for two short inputs, rerank 160 for two documents, and the snapshot shows per-model token and error counters alongside residency keyed by resolved path. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01FqHNxHa5fUW7BRDtq5nghu
|
All four findings from the second pass were reproducible and are fixed in P1 — pin acquired backends. Correct, and my previous fix was insufficient: reserving a slot still left the window where a request had exited the acquisition path but not finished, so eviction unloaded weights it was about to use and they were immediately reloaded — two models live, one counted. A backend is now pinned for its whole load-and-inference lifetime, and eviction skips pinned backends. When every resident backend is busy the cap is briefly exceeded, which is the lesser evil versus pulling weights out from under an in-flight request. P2 — key by resolved path. Correct, and it contradicted this module's own docstring. Lookup used the raw reference, so P2 — record failures. Correct. P2 — real token usage. Correct. Counts now come from the encoding already performed per batch. Verified live: 14 prompt tokens for two short embedding inputs, 160 total for a two-document rerank. Suites after the change: 48 Python retrieval tests, 533 Swift, all green. |
|
@codex review
|
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: fd2a238605
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| finally: | ||
| with self._lock: | ||
| backend.users = max(0, backend.users - 1) |
There was a problem hiding this comment.
When requests for more models than max_resident overlap, acquisition correctly skips every pinned victim, but this finally block only decrements users and never retries eviction. For example, with a cap of 1, overlapping requests for A and B finish with both backends idle, loaded, and resident indefinitely until another retrieval request happens, defeating the memory cap and potentially retaining multiple multi-GB models. Run the eviction pass when a backend becomes unpinned.
Useful? React with 👍 / 👎.
| ) | ||
| } content: { | ||
| VStack(alignment: .leading, spacing: 12) { | ||
| ForEach(retrieval.models) { model in |
There was a problem hiding this comment.
When the same unaliased reference is configured for embedding and reranking—the supported shared-backend case—the registry emits two status entries with the same served id but different roles. RetrievalModelStatus uses that served ID as its Identifiable identity, so this ForEach receives duplicate IDs and may reuse or omit one row, displaying the wrong role or metrics. Key rows by a composite of role and served ID instead.
Useful? React with 👍 / 👎.
…tities Two findings from the third Codex review on youssofal#212. Enforce the cap when the last pin is released. Overlapping requests for more models than the cap legitimately exceed it while every backend is pinned, but the release path only decremented the pin count. Both models then stayed loaded and resident indefinitely until some later request happened to trigger eviction — multiple multi-GB models retained for no reason. Eviction now runs again as backends become unpinned. Writing the regression test surfaced a second defect in the same pass: eviction could drop the *most recently used* backend whenever an older one was pinned, which inverts LRU. The pass now never considers the most recently used entry — the one just acquired or just released — so a surplus is resolved by dropping the oldest unpinned model instead. Give role-shared models unique SwiftUI identities. One reference registered as both embedder and reranker is the shared-backend case this feature advertises, and it yields two status entries under the same served id. Using that id as `Identifiable` identity made a ForEach see duplicates and drop or reuse a row, showing the wrong role or metrics. Identity is now a composite of role and served id; the served id remains available for display and for the `"model"` field clients send. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01FqHNxHa5fUW7BRDtq5nghu
|
Both findings were reproducible and are fixed in P1 — enforce the cap on release. Correct. Acquisition skipped pinned victims as intended, but the release path only decremented the pin count, so a legitimate overshoot from overlapping requests stayed loaded until the next request happened to trigger eviction. Eviction now runs again as backends become unpinned. Writing the regression test surfaced a second defect in the same code: eviction could drop the most recently used backend whenever an older one was pinned, which inverts LRU — with a cap of 1, finishing the inner of two overlapping requests evicted the model that had just been used and kept the older one. The pass now never considers the most recently used entry, so a surplus resolves by dropping the oldest unpinned model. Same rule at both call sites, which also removed the P2 — role-shared identities. Correct, and it hit the case this feature advertises: one reference in both roles yields two entries under the same served id, so Suites: 49 Python retrieval tests, 16 Swift retrieval tests, 366 Python and 533 Swift overall — green. Verified live on the running daemon afterwards. |
MLX holds model weights in unified memory that never appears in process RSS, so a daemon serving 24 GB of models shows 0.5 GB in Activity Monitor. The app's Memory Detail card already attributes the chat model's weights, the session cache and the working set on that basis — but retrieval models were missing from it, leaving several GB unaccounted for and looking like a leak somewhere else on the machine. The registry now records each backend's shard size on load and reports it per model, plus a `resident_bytes` total summed per backend rather than per served id, so one reference serving both roles is counted once. The app adds a "Retrieval models" row to the same breakdown, showing "none loaded" while the endpoints are configured but idle. Measured on the running daemon: 0 GB before use, 3.96 GB after the first embedding request, 6.07 GB with the reranker also resident, the increase matching the reranker's shard size exactly. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01FqHNxHa5fUW7BRDtq5nghu
…on bank Records the scope agreed with the user: idle unload for retrieval models plus SSD archival of the session bank, with the chat model explicitly out of scope because no unload path exists for it in the core. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01FqHNxHa5fUW7BRDtq5nghu
…on bank Retrieval models loaded on first request and then stayed resident until the resident cap was exceeded. With two configured models and a cap of two the cap can never be exceeded, so nothing was ever released: several GB were held for the lifetime of the daemon even with no retrieval traffic for hours. Idle release simply did not exist anywhere. A watcher runs beside the memory-pressure guard, every 30 s, and only when --retrieval-idle-timeout is configured — an unconfigured daemon behaves exactly as before. It unloads loaded, unpinned backends idle for longer than the threshold, then archives the session bank to the existing SSD cold tier once nothing is resident, so a resumed conversation restores its prefix from disk instead of re-prefilling. Backends inside an in-flight request are never candidates: the pin count is checked under the same lock that hands them out. The watcher swallows and logs its own errors — one that can take the daemon down is worse than one that misses a cycle — and a failed archive does not block the unload. The snapshot reports the configured timeout and per-model idle seconds so the app can show how long a model has left rather than only that it is loaded. Configurable from the CLI, config.toml and the app. The chat model is deliberately untouched: no unload path exists for it in the core, and adding one is a separate change with its own design. Measured end to end through the app-launched daemon: 17.87 GB idle, 21.83 GB with the embedder resident, back to 17.87 GB after the timeout with nothing loaded. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01FqHNxHa5fUW7BRDtq5nghu
The idle watcher covers the quiet case, but the existing memory-pressure guard only shed the MLX buffer pool and session-bank weight — retrieval models were never candidates, so several GB of the cheapest-to-reload memory survived a CRITICAL reading untouched. The guard now releases unpinned retrieval backends on the rising edge into CRITICAL, before clearing the buffer pool. Weights that reload in seconds are the right thing to give back first when the system is out of memory. This exposed a conflict in unload_idle's contract: 0 means "disabled" as a configured timeout but "every unpinned model" as an explicit request from the pressure path. An omitted threshold now falls back to the configured timeout (0 disables), while an explicitly passed threshold is always honoured, including 0. Pinning is unaffected either way — a model inside an in-flight request is never released, under pressure or not. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01FqHNxHa5fUW7BRDtq5nghu
The card sat sixth in the Activity tab, below in-flight requests, two handoff notices, recent requests and Speed Truth — so "which models are loaded" required scrolling past four cards to answer. It now follows in-flight requests directly, still conditional so a chat-only setup is unchanged. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01FqHNxHa5fUW7BRDtq5nghu
Embedding requests took ~900 ms warm while the forward pass itself took 31 ms. The gap was serialisation: returning a dict lets FastAPI run jsonable_encoder over the response, which walks all 4096 floats of every vector individually in Python. For an 8B embedder that is roughly 30x the cost of the inference it wraps. The handler now serialises with json.dumps and returns a Response directly. Measured on the running daemon: ~900 ms before, 41 ms after, byte-identical payload. A hypothesis worth recording as refuted: mx.clear_cache() in the batch loop looked like the culprit, since MTPLX's own pressure guard documents that clearing mid-decode taxes subsequent steps. Measured back to back it costs 2 ms per request (31 vs 29), so it stays — the cost was never in the MLX path at all. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01FqHNxHa5fUW7BRDtq5nghu
…seven short ones A batch is padded out to its longest row, and the model pays rows * width token positions regardless of how short the other rows are. Both embed() and rerank() chunked their input in arrival order, so the mixed request that Hermes recalls produce all the time — one long passage beside a handful of short ones — cost as much as a batch of nothing but long passages. Sorting alone would not have fixed it. At the default batch size of 8 a mixed request is a single batch, and reordering inside one batch leaves max(lengths) exactly where it was. What moves the boundary is a second limit in token positions: sort by length, then close a batch once either the configured row count or BATCH_TOKEN_BUDGET would be exceeded. The long text falls into a batch of its own and the short ones stay packed. The budget is set from measurement rather than taste. On a 4-bit Qwen3-Embedding-8B, long sequences gain nothing from being batched (~950 ms each at 853 tokens, alone or eight at a time) while short ones gain about threefold, so a row count is the wrong unit in both directions. 2048 positions keeps any one batch under a second. Measured A/B in one process, alternating variants so thermal drift hits both sides equally: embedding, 1 long + 7 short 8239 ms -> 1249 ms 6.60x embedding, 8 short 187 ms -> 186 ms 1.00x embedding, 8 long 8921 ms -> 9119 ms 0.98x rerank, 1 long + 3 short 1558 ms -> 702 ms 2.22x rerank, 4 short 286 ms -> 293 ms 0.97x Splitting the all-long case into batches of two costs 2% — inside the noise — and cuts peak tensor memory to a quarter, which is welcome on a daemon already holding chat weights and a session bank. Planning reorders, so results are scattered back by index instead of appended. Callers see their own order; rerank() depends on this because it ranks by index into the caller's document list. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01CAYAMGYxYz6d1MbZv3rLHc
|
this is a serious contribution and the self-review discipline shows. it's on the bench for a heavy pass this weekend: embedding/rerank correctness against reference implementations, residency behavior under the memory pillar, concurrent load next to a decode session. not merging before that, but expect movement. |
2 participants
What
Adds
POST /v1/embeddings(OpenAI shape) andPOST /v1/rerank(Cohere/Jina shape) to the MTPLX daemon, so a retrieval-backed setup — RAG, agent memory — no longer needs a second inference server running beside MTPLX purely to answer embedding calls. Configurable from the CLI and the macOS app, and observable in the dashboard.Both flags repeat, accept a Hugging Face id or a local path with an optional
REF=served-idalias, and are selected per request via"model". Also onquickstart, in the app under Settings → Retrieval endpoints, and in~/.mtplx/config.tomlasembedding_models/reranker_models/retrieval_max_resident.Why not through MTP
Retrieval models deliberately bypass the MTP path. Multi-token prediction makes next-token decoding cheaper, which means nothing for a model that returns a vector instead of a token stream — there is no draft to accept or reject. They run the transformer stack directly:
<|endoftext|>and pools that position, whichpooling="last"reproduces.yes/nologits at the last real position.Batches are padded on the right. With causal attention a real token never attends to a later pad token, so right padding leaves every real position bit-identical to an unpadded run; left padding would corrupt it.
Memory behaviour
--retrieval-max-resident(default 2) unloads the least recently used backend beyond the cap, but a backend stays pinned for its whole load-and-inference lifetime and is never evicted while in use. When every resident backend is busy the cap is briefly exceeded — the alternative is unloading weights an in-flight request is about to use, which frees nothing and forces an immediate reload.Additive by construction
With no flag configured the endpoints answer 404 and the chat runtime is untouched.
/v1/modelsgains acapabilityfield (chat/embedding/rerank) alongside the existing entry.Three wiring points, worth knowing if this area is refactored: the server argv is rebuilt rather than inherited, so the flags are threaded through the
serve/quickstartparsers, the child argv builder incommands/public.py, and the server module's own parser. Missing any one leaves the endpoints silently unconfigured on every path but a baremtplx serve. The same applies to the cache directory: the chat model reaches the server pre-resolved, so retrieval references need--retrieval-cache-dirto find models pulled into a custom--cache-dir.Observability
Configuration that cannot be observed is a guess, so the dashboard snapshot carries a
retrievalsection: per model the role, load and residency state, request/item/token/error counts, cumulative compute, load duration, derived throughput and average latency, and the last error. The key is always present, letting a client tell "no retrieval configured" apart from "this daemon predates retrieval".Load time is measured but excluded from request latency. Counting it made a cold first request read as ~10 s average latency on a model that serves in ~450 ms — precisely the number a user would act on.
The app renders this in the Activity tab and distinguishes the states that matter: configured but never loaded (normal, because loading is lazy), loaded but unused, working, and failing.
Evidence
Hardware: Apple Silicon, 48 GB unified memory. Models:
mlx-community/Qwen3-Embedding-8B-4bit-DWQ(4-bit DWQ) andvserifsaglam/Qwen3-Reranker-4B-4bit-MLX(4-bit). Profilesustained, fan mode default (no fan-controlled runs). Base commit 97fb388.tests/test_retrieval.py) and 15 Swift tests, covering flag parsing, served-id resolution, resolved-path backend sharing, pinning against eviction, slot accounting, metric derivation, error recording, token usage, settings round trip, backward-compatible decoding, and the HTTP contract including the chat-only 404 path. Full suites green: 365 Python in the touched areas, 533 Swift.1.000000; reranker separating a matching from a non-matching document0.9994vs0.0000; base64 decoding to the same vectors;13.4 texts/sat447 msper request with the6.0 sweight load reported separately; real token usage (14for two short embedding inputs,160for a two-document rerank);/v1/chat/completionsanswering from the same daemon throughout.No performance claim is made about the chat path, which is unchanged.
Commits
feat(server+app)— the endpoints, CLI/app/config wiring, tests, docs.fix(retrieval)— round one of review: resident-slot reservation, forwarding the cache directory, honouringencoding_format(base64, unknown values rejected rather than silently misread).feat(retrieval)— live load state and throughput in the snapshot and the app.fix(retrieval)— round two of review: pinning backends through inference, keying by resolved path, recording failures, and reporting real token counts.Contribution gate
tests/test_no_mlx_imports.pyandtests/test_runtime_kpis.pypass;python -m buildproduces both artifacts withmtplx/retrieval.pyin the wheel.tests/test_public_cli.pyhas one failure,test_start_hermes_live_path_writes_profile_and_handoff, which assertsport == 18085. It fails identically on a pristine checkout of the base commit in this environment because port 18085 was already bound — verified in a separategit worktreeatHEAD, so it is not a regression from this change.