Add BGE-M3 encoder pooling (#612)
This PR is:
- To add BGE-M3 dense embedding and sparse `token_classify` support
through the encoder pooling backend.
- To load canonical `BAAI/bge-m3` on the MLX runtime, including PyTorch
backbone weights and `sparse_linear.pt`.
- To keep BGE-specific pooling behavior in the model-owned encoder
backend instead of the runner.
- To share encoder checkpoint loading while avoiding unused repository
downloads.
Note:
This supports canonical, unquantized `BAAI/bge-m3`. MLX Community BGE-M3
checkpoints are not supported in this PR because they do not include
`sparse_linear.pt`, and this change intentionally ships dense + sparse
BGE-M3 together.
For sparse BGE-M3, use:
```bash
vllm serve BAAI/bge-m3 \
--runner pooling \
--max-model-len 8192 \
--pooler-config '{"task":"token_classify"}' \
--hf-overrides '{"architectures":["BgeM3EmbeddingModel"]}'
```
Proof it works:
Model revision: `5617a9f61b028005a4858fdac845db406aefb181`.
Dense command:
```bash
vllm serve BAAI/bge-m3 --runner pooling --max-model-len 8192 --port 8012
```
Dense `/v1/embeddings` returned shape `(2, 1024)` with normalized
embeddings and matched the upstream similarity reference:
```text
norms [1.0, 1.0]
similarity [[0.625909, 0.347491], [0.33097, 0.673372]]
reference [[0.6259, 0.3474], [0.3309, 0.6734]]
```
Sparse `/pooling` with `task=token_classify` returned one sparse
lexical-score vector with length `8` and matched the upstream
lexical-score reference:
```text
lexical_score 0.19554413560964878
reference 0.19554901123046875
abs_diff 4.8756208199662865e-06
```
---------
Signed-off-by: Yuan Lik Xun <lxyuan0420@gmail.com> L
Lik Xun Yuan (Lx) committed
ef3ec8c9c3547f1402f86a50f3a2e5ac73c0d96c
Parent: 083f581
Committed by GitHub <noreply@github.com>
on 8/16/2026, 8:46:41 AM