Add typed encoder backend and dense XLM-RoBERTa embed - #605
Closed
Aryan4443 wants to merge 5 commits into
Closed
Conversation
Unblock encoder embedding checkpoints on Metal by loading through the optional mlx-embeddings extra and serving CLS dense embed. Sparse token_classify remains a follow-up for vllm-project#589. Signed-off-by: Aryan <aryannlakhani21@gmail.com> Co-authored-by: Cursor <cursoragent@cursor.com>
Encoder embedding models own bidirectional attention and have no mlx-lm-style layer list to wrap. Skipping patch setup unblocks the offline LLM.embed path for BGE-M3 / XLM-RoBERTa. Signed-off-by: Aryan <aryannlakhani21@gmail.com> Co-authored-by: Cursor <cursoragent@cursor.com>
Move load, bidirectional segment forward, CLS pooling defaults, and paged-attention patch skipping behind EncoderEmbeddingAdapter so the encoder path is reviewable as a single ownership boundary. Signed-off-by: Aryan <aryannlakhani21@gmail.com> Co-authored-by: Cursor <cursoragent@cursor.com>
Drop the incompatible mlx-embeddings dependency and load BGE-M3 / XLM-RoBERTa through an in-tree MLX encoder that keeps the adapter-owned load, forward, CLS defaults, and patch-skip boundary. Signed-off-by: Aryan <aryannlakhani21@gmail.com> Co-authored-by: Cursor <cursoragent@cursor.com>
Keep pooling, lifecycle, cache setup, and the runner family-agnostic behind EncoderEmbeddingBackend, and register XLM-RoBERTa as the first dense family with real load plus CLS pooling parity. Signed-off-by: Aryan <aryannlakhani21@gmail.com> Co-authored-by: Cursor <cursoragent@cursor.com>
LxYuan0420
reviewed
Aug 13, 2026
Collaborator
There was a problem hiding this comment.
Thanks for working on this, but I don’t think this PR is the right base.
It mixes the pooling seam, encoder lifecycle, cache behavior, and native XLM-R/BGE-M3 model code in one stack. I’m submitting a smaller PR first to make the current decoder pooling path clean and extensible:
vLLM pooling API
-> runner scheduling/output
-> pooling backend
-> decoder now
-> encoder later
-> XLM-R/BGE-M3 later
After that lands, encoder support should restart on top of this seam in smaller PRs.
Author
|
Thanks for the direction, that split makes sense. I'll wait for #606 to land, then restart encoder support on that seam in smaller PRs. Happy to take PR 2 (encoder pooling backend + no-KV lifecycle) or PR 3 (native XLM-R dense embed) once the decoder backend is in. |
LxYuan0420
reviewed
Aug 15, 2026
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
2 participants
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.
Summary
EncoderEmbeddingBackend+ family registry. XLM-R / BGE names stay in the family module.mlx-embeddings), with CLS pooling on the existing pooling path.Fixes #589 (dense embed only; sparse
token_classifyremains follow-up).Test plan
python -m pytest tests/test_encoder_embeddings.py tests/test_v1_pooling.pyLLM.embedonmlx-community/bge-m3-mlx-8bitwith tight memory (VLLM_METAL_MEMORY_FRACTION=0.35,max_model_len=128)