Micro-batch local embeddings into one batched MLX forward#1917
Open
mimeding wants to merge 3 commits into
Open
Micro-batch local embeddings into one batched MLX forward#1917mimeding wants to merge 3 commits into
mimeding wants to merge 3 commits into
Conversation
Memory search and distillation fan out many concurrent single-text embeds (VecturaKit embeds each search query via embed(text:)), and each one serialized through MetalGate as its own MLX forward. Two changes make those requests genuinely batch: - EmbeddingBatcher: an actor that coalesces single-text requests arriving within a 25 ms window (first request opens the window; later arrivals never extend it, so a lone caller's added latency is bounded by one window) into one backend call, capped at 16 per batch with an immediate flush when full. Errors propagate to every waiter; cancelling one waiter resumes only that caller and leaves the batch running. The window duration and clock are injectable so tests never depend on wall-clock time. - VMLXModel2VecEmbedder now runs a real batched forward: vmlx-swift's Model2VecStaticEmbeddingPipeline.embed(texts:) just loops embedOne per text, so coalescing through it would have changed nothing. The embedder loads the same bundle (embeddings tensor, normalize flag, tokenizer) and computes a padded [N, L] gather + masked mean + row L2-normalize with a single eval for the whole batch. The batching layer sits above MetalGate: the batcher's backend is MetalSafeEmbedder.embed(texts:), so one coalesced batch acquires the gate exactly once, and gate semantics are unchanged. EmbeddingService.sharedEmbedder is now a CoalescingEmbedder that routes single-text embeds through the batcher and passes caller-assembled batches straight through (they are already one forward); directEmbedder keeps the old non-coalescing path callable. Batch parity against the real potion-base-4M bundle (padding and masking must not change any vector, max delta < 1e-4, unit norms) is covered by an opt-in test: MLX aborts the test process when its metallib is not locatable, so it only runs when OSAURUS_EMBEDDING_MODEL_DIR is set explicitly. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The parity test previously compared the batched forward against its own embed(text:) path, which routes through embed(texts:) — both sides ran the new implementation, so only padding/masking invariance was proven. A systematic semantic deviation from vmlx's Model2Vec pipeline (tokenizer handling, unknown-token filtering, normalize behavior) would have passed unnoticed. Compare each text's batched vector against vmlx-swift's original sequential Model2VecStaticEmbeddingPipeline loaded from the same bundle (observed max delta on potion-base-4M: 1.5e-08). The one intentional convention difference — an empty input array returns [] here but throws emptyBatch in vmlx — is asserted on both sides so a silent change in either implementation fails the test. MLXEmbedders joins the test target's dependencies for the reference import. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Adversarial review finding on the embedding micro-batch path: the CoalescingEmbedder hands caller-assembled arrays straight to VMLXModel2VecEmbedder.embed(texts:), whose padded gather materializes [N, L_max, D] floats with N client-controlled via /v1/embeddings — an unbounded request meant an unbounded GPU allocation. embed(texts:) now processes the input in chunks of at most EmbeddingBatcher.defaultMaxBatchSize (16 — the same bound the micro-batcher keeps: ~4 MB per chunk at 512 tokens x 128 dims), concatenating results in order. The chunks run sequentially inside the CALLER's single MetalGate acquisition — MetalSafeEmbedder.embed(texts:) wraps the whole method, so chunking never releases/re-acquires the gate mid-batch. The per-chunk forward is unchanged (extracted to embedChunk); chunk boundary math lives in the pure static forwardChunkRanges so it is unit-testable without loading MLX. Tests: chunked == unchunked equality for 40 texts at the batcher level (every backend batch capped at 16, vectors match the direct reference), a real-model 40-text chunk-parity case in the env-gated suite (chunk boundaries invisible in output), forwardChunkRanges boundary/edge cases, and the previously unexercised mismatchedResultCount path via a mock backend returning the wrong vector count. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Contributor
Author
|
Adversarial-review fix pushed (06f84bc): the direct embed(texts:) pass-through now chunks client-controlled arrays to ≤16 per forward inside the single MetalGate acquisition — previously /v1/embeddings could materialize an unbounded [N, L_max, D] padded tensor (~8.6 GB for a hostile batch). Chunk-parity and mismatched-count tests added. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
Micro-batches the local embedding path into genuinely batched MLX forwards:
EmbeddingBatcher(actor): the first request in an idle period opens a 25 ms window (later arrivals never extend it, so a lone caller's added latency is bounded by one window); up to 16 requests coalesce, with immediate flush when full; per-batch errors propagate to every waiter; cancelling one waiter removes only that caller. Clock and window are injectable for tests.[N, L]batched forward inVMLXModel2VecEmbedder: padded gather → masked mean → row L2-normalize → oneevalper batch. This was necessary because vmlx-swift'sembed(texts:)is secretly a sequential loop (texts.map(embedOne)— one gather+eval per text), so a coalescing front alone would have been a fake batch./v1/embeddingsand the pluginembed(texts:)paths are pass-through.Why
Embeddings are the highest-frequency small inference in osaurus (memory search + distillation), and they serialize today — both at the gate and in vmlx's per-text loop. Batched execution amortizes the gate acquisition, kernel launches, and eval/
asArrayround-trips; the biggest practical win is under generation load, where each embedding previously re-queued behind the generation gate individually.Measurement
MLXEmbedders) is test-target-only.[]here vs vmlx throwingemptyBatch.OSAURUS_EMBEDDING_MODEL_DIR) because MLX hard-aborts the test process when its metallib isn't locatable under plainswift test— the same reason no existing test runs MLX compute.Risk & rollback
Worst case for a lone embedding call is +25 ms (documented; negligible against gate-contention waits it replaces). The batched math is validated to 1e-8 against the prior implementation. A follow-up worth doing upstream: vectorize
Model2VecStaticEmbeddingPipeline.embed(texts:)in vmlx-swift so the ~40 mirrored loader lines here can shrink back to a coalescing front. Revert = two commits.🤖 Generated with Claude Code