Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions CONTRIBUTING.md
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@ Thank you for your interest in contributing to mlxcel! This document covers the
| Build and test locally | [`docs/installation.md`](docs/installation.md) |
| Understand the architecture | [`docs/architecture.md`](docs/architecture.md) |
| Add a new model family | [`docs/adding-models.md`](docs/adding-models.md) |
| Add or validate an embedding family (`/v1/embeddings`, `mlxcel embed`) | [`docs/embeddings.md`](docs/embeddings.md) |
| Understand the code guidelines (shared functions, file size and module splits) | [`docs/code-guidelines.md`](docs/code-guidelines.md) |

## How to contribute
Expand Down
1 change: 1 addition & 0 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -396,6 +396,7 @@ Install with `pip install ./python`. See [Python client](docs/python-client.md)
- [Speculative-decoding acceptance](docs/speculative-acceptance.md)
- [Adaptive MTP policy API](docs/mtp-policy-api.md)
- [OpenAI Responses API](docs/responses-api.md)
- [Embeddings API and `mlxcel embed`](docs/embeddings.md)
- [Audio input preprocessing](docs/audio-preprocessing.md)
- [Python client](docs/python-client.md)
- [Adding a new model](docs/adding-models.md)
Expand Down
1 change: 1 addition & 0 deletions docs/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,7 @@ Current GitHub-facing docs:
18. `sparse-paged-decode.md` - sparse attention reduced to page indirection over the fused v2 decode kernel: the addressing argument, the MiniMax-M3 routing, why DeepSeek Sparse Attention is not routed, where the dispatch floor sits, the kill switch, and the benchmark harness.
19. `mtp-policy-api.md` - the supported read interface for the adaptive B=1 MTP verdict (`GET /v1/internal/mtp-policy`): the response body, the four states, the unavailable reasons, and the schema versioning and compatibility policy.
20. `code-guidelines.md` - the file-size and module-split thresholds, including when to extract a `<name>_helpers.rs` and when inline tests move to a sibling `_tests.rs`, the `// Used by:` annotation convention for shared functions, the JIT kernel rule that every varying input dtype must appear in `template_args` because CUDA keys the compiled-module cache on it, and the `HashMap` iteration-order rule covering what counts as an order-sensitive consumer, why a stable sort on a non-total key does not pin the result, the fresh-map-per-iteration testing requirement, and why no static check enforces it.
21. `embeddings.md` - the `/v1/embeddings` endpoint and `mlxcel embed`: embedding-checkpoint detection rules, pooling modes and the `1_Pooling/config.json` reader, L2 normalization and `dimensions`, `max_length` derivation, the padding-aware attention-mask builders, the request and response schema with error codes, the server flags, and the checklist for registering a new embedding family.

## Architecture Decision Records

Expand Down
183 changes: 183 additions & 0 deletions docs/embeddings.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,183 @@
# Embeddings API (`/v1/embeddings`) and `mlxcel embed`

mlxcel serves embedding checkpoints through the OpenAI-compatible `POST /v1/embeddings` endpoint and the offline `mlxcel embed` command. This page covers the shared foundation every embedding family builds on: detection, pooling, normalization, length limits, the request and response schema, the server flags, and the CLI. Per-family details (prompt prefixes, instruction formats, image support) are filled in by each family as it lands; the family table lives in [supported-models.md](supported-models.md#embedding-models).

**Current status.** The route, the worker thread, batching, pooling and the CLI are in place. Detection recognizes every family listed below, and the loader dispatches on it, but no family forward pass has landed yet: loading a recognized checkpoint reports `<family> is detected as an embedding checkpoint, but this embedding family is not yet supported by /v1/embeddings` until its port merges (epic #1348).

## Implemented endpoints

| Method | Path | Description |
|--------|------|-------------|
| POST | `/v1/embeddings` | Embed a string, a list of strings, a token-id array, a list of token-id arrays, or a list of typed parts. |
| POST | `/embeddings` | Alias without the `/v1` prefix. |
| GET | `/v1/models` | Lists the embedding model next to the chat model when `--embedding-model` loads a second checkpoint. |

## Implementation source map

| Module | Responsibility |
|--------|----------------|
| `src/embeddings/model.rs` | `EmbeddingModel` trait, `EmbeddingBatch`, `EmbeddingOutput`, `ImageInput`. |
| `src/embeddings/pooling.rs` | `PoolingMode`, `1_Pooling/config.json` reader, `pool`, `normalize_l2`, `truncate_dimensions`, the `MLXCEL_EMBEDDING_POOLING` override. |
| `src/embeddings/limits.rs` | `max_length` derivation, pad token and vocabulary size. |
| `src/embeddings/tokenize.rs` | Right-padded batch tokenization, trailing-special truncation, pair encoding. |
| `src/embeddings/loader.rs` | `load_embedding_model`: family dispatch, subfolder weights, bf16 rule. |
| `src/embeddings/engine.rs` | Length-sorted micro-batching, normalization, `dimensions`, readback. |
| `src/lib/mlxcel-core/src/utils.rs` | `create_bidirectional_padding_mask`, `create_causal_padding_mask`, `create_bidirectional_window_mask`. |
| `src/lib/mlxcel-core/src/weights.rs` | `load_weights_from_dir_with_subfolders` (`2_Dense/...` tensors prefixed `2_Dense.`). |
| `src/models/detection.rs` | `is_embedding_checkpoint`: the detection rules below. |
| `src/server/embedding_model.rs` | `EmbeddingModelProvider` trait and `EmbeddingError`. |
| `src/server/embedding_worker.rs` | `EmbeddingWorker`: the dedicated MLX-owning thread, bounded queue, timeout, panic boundary. |
| `src/server/routes/embeddings.rs` | HTTP handler, validation order, error mapping. |
| `src/server/types/embeddings.rs` | Request and response types, base64 encoding. |
| `src/commands/embed.rs` | `mlxcel embed`. |

## Detection

A checkpoint is served as an embedding model, before the ordinary `model_type` dispatch runs, when any of these hold:

- `config.json` `model_type` is an encoder-only family: `bert`, `xlm-roberta`, `modernbert`, `siglip`. These never generate text; a `BertForMaskedLM` or `ModernBertForMaskedLM` export loads as an embedder with its MLM head dropped.
- `config.json` `architectures[0]` is an embedding architecture: `BertModel`, `XLMRobertaModel`, `ModernBertModel`, `SiglipModel`, `SiglipTextModel`, `Gemma3TextModel` with `use_bidirectional_attention: true`, `LlamaBidirectionalModel`, `LlamaNemotronVLModel`, `Lfm2BidirectionalModel`, `Ministral3Model` with `is_causal: false`, `ColIdefics3`, `ColQwen2_5`, `ColQwen2ForRetrieval`.
- `modules.json` exists and lists a module whose `type` ends with `.Pooling` (the sentence-transformers layout). A `modules.json` whose only extra module is `1_LogitScore` (Qwen3-VL-Reranker) does not qualify.
- `1_Pooling/config.json` exists.

A checkpoint whose `architectures[0]` ends with `ForSequenceClassification` is a reranker, never an embedder, whatever else its layout says.

The resolved family is keyed on `model_type`:

| `model_type` | `ModelType` |
|--------------|-------------|
| `bert` | `Bert` |
| `xlm-roberta` | `XlmRoberta` |
| `modernbert` | `ModernBert` |
| `siglip` | `SiglipText` |
| `gemma3_text`, `gemma3` | `Gemma3Embedding` |
| `qwen3` | `Qwen3Embedding` |
| `qwen3_vl` | `Qwen3VLEmbedding` |
| `lfm2` | `Lfm2Embedding` |
| `ministral3` | `Ministral3Embedding` |
| `llama`, `llama_bidirec` | `LlamaBidirec` |
| `llama_nemotron_vl` | `LlamaNemotronVLEmbedding` |
| `idefics3` | `ColIdefics3` |
| `qwen2_5_vl`, `colqwen2` | `ColQwen25` |

A pooling layout on any other `model_type` is reported as an error naming the `model_type`, rather than misrouted to a causal generator. `Qwen3ForCausalLM` without a pooling layout still detects as the `Qwen3` generator.

`mlxcel arch` lists the embedding variants under the `Embedding` family. The generation loader (`load_model`, used by `mlxcel generate`, `mlxcel run` and the chat worker) rejects every embedding variant with a message pointing at `/v1/embeddings` and `mlxcel embed`, so `mlxcel-server -m <embedding checkpoint>` serves embeddings and leaves chat unloaded (the same way a Whisper checkpoint leaves chat unloaded), and `/v1/chat/completions` returns the existing "model is not loaded" error.

## Pooling

Given hidden states `[B, L, D]` (cast to f32) and the attention mask `[B, L]` (int32, `1` = real token, `0` = padding):

| Mode | Definition |
|------|------------|
| `cls` | `hidden[b, first_real_index[b], :]` with `first_real_index = argmax(mask, axis=1)`; index 0 for right padding, the first `1` for left padding. |
| `mean` | `sum(hidden * mask[..., None], axis=1) / max(sum(mask, axis=1), 1e-9)`. |
| `max` | `max(where(mask == 0, -inf, hidden), axis=1)`. |
| `lasttoken` | `hidden[b, last_real_index[b], :]`; an all-padding row uses `L - 1`. Correct for left and right padding. |

Unsupported modes (`weightedmean`, `mean_sqrt_len_tokens`, `include_prompt: false`, more than one legacy flag set) are a load error naming the mode, not a silent fallback.

Resolution order:

1. `<model_dir>/1_Pooling/config.json` when present. Both key styles are accepted: the new-style `"pooling_mode": "cls" | "mean" | "max" | "lasttoken"` and the legacy booleans `pooling_mode_cls_token`, `pooling_mode_mean_tokens`, `pooling_mode_max_tokens`, `pooling_mode_lasttoken`, `pooling_mode_weightedmean_tokens`, `pooling_mode_mean_sqrt_len_tokens`. Exactly one legacy flag may be true; none true means `mean`.
2. Otherwise the family default (`EmbeddingModel::default_pooling()`).
3. `MLXCEL_EMBEDDING_POOLING=cls|mean|max|lasttoken` overrides both, for debugging, and is logged at startup.

Normalization is L2 per vector, `v / max(||v||_2, 1e-9)`, applied when `EmbeddingModel::normalize()` is true (the default; `config.json` `normalize: false` turns it off for families that declare it). The request field `dimensions` keeps the first `dimensions` components and re-normalizes when normalization is on; `1 <= dimensions <= D` is required, any other value is a `400`.

## Length limits and padding

`max_length` is the smallest of: `sentence_bert_config.json` `max_seq_length` when present; `tokenizer_config.json` `model_max_length` when `0 < value < 1_000_000`; `config.json` `max_position_embeddings` for absolute-position encoders (BERT, XLM-RoBERTa, SigLIP); the hard cap `8192`; and `--embedding-max-length` when set. Inputs longer than `max_length` are truncated from the right after special tokens are added, keeping the trailing special token when the tokenizer appends one, so a BERT input keeps its `[SEP]` and a Qwen3-Embedding input keeps its `<|endoftext|>`. Token-id inputs are truncated from the right verbatim.

For `sentence-transformers/all-MiniLM-L6-v2` this resolves to `256` (from `sentence_bert_config.json`); for `Qwen/Qwen3-Embedding-0.6B` to the cap of `8192` (no `sentence_bert_config.json`, `model_max_length` 131072, rotary positions so `max_position_embeddings` is not consulted).

Micro-batches are right-padded to their longest member (or to a fixed width when the family requires one, as SigLIP text does) with the pad token id: `tokenizer_config.json` `pad_token`, falling back to `eos_token`, then `0`. Any padding or truncation baked into `tokenizer.json` (MiniLM pads to a fixed 128) is stripped at load, because the engine pads per micro-batch and truncates per checkpoint limit.

## Attention masks

Families build additive f32 masks (`0.0` = attend, `-inf` = blocked) from the `[B, L]` attention mask with three builders in `mlxcel_core::utils`, next to `create_causal_mask`:

- `create_bidirectional_padding_mask(mask) -> [B, 1, 1, L]`: key column `k` is blocked iff `mask[b, k] == 0`. Padding query rows still see every real key, so no row is fully blocked.
- `create_causal_padding_mask(mask, offset) -> [B, 1, L, L + offset]`: `k <= q + offset` and `mask[b, k] == 1`. A fully padded query row keeps its diagonal column, the same rescue `create_causal_mask_with_left_padding` applies.
- `create_bidirectional_window_mask(mask, window) -> [B, 1, L, L]`: blocked iff `mask[b, k] == 0` or `|q - k| >= window`.

A mask is never built as `(1 - m) * C` with a finite `C` in the activation dtype: in f16, `-1e9` overflows to `-inf` and `0 * -inf` is NaN. The builders produce `0 / -inf` in f32, which survives the cast `fast_scaled_dot_product_attention` applies for f16 and bf16 activations.

## Request and response

```json
{
"model": "optional; must equal the served embedding model id when given",
"input": "string | [string] | [int] | [[int]] | [{\"type\": \"text\", \"text\": ...} | {\"type\": \"image_url\", \"image_url\": {\"url\": ...}}]",
"encoding_format": "float | base64",
"dimensions": 256,
"instruction": "optional; forwarded to EmbeddingModel::format_text",
"user": "ignored"
}
```

```json
{
"object": "list",
"data": [{"object": "embedding", "index": 0, "embedding": [0.01, 0.02]}],
"model": "all-MiniLM-L6-v2",
"usage": {"prompt_tokens": 17, "total_tokens": 17}
}
```

- Token-id inputs are used verbatim (no special tokens added) and every id must be `< vocab_size`.
- `encoding_format: base64` encodes each vector as little-endian f32 bytes, standard base64 with padding. For multi-vector (late-interaction) models `embedding` is a list of vectors, `[num_real_tokens][D]`, in float mode; in base64 mode it is the row-major bytes plus a sibling `"shape": [num_real_tokens, D]`.
- `usage.prompt_tokens` counts the real (non-padding) tokens across all inputs, special tokens included; `total_tokens` equals it.
- `image_url` items accept a `data:image/...;base64,` URI, a `file://` path, an `http(s)://` URL or a local path, bounded by the shared image limits. Images are embedded one at a time; text items are sorted by token length, cut into micro-batches of `--embedding-batch-size`, and written back in request order.

### Errors

| Status | Type | When |
|--------|------|------|
| 400 | `invalid_request_error` | Malformed body, empty `input`, an empty string item, an empty token list, a token id `>= vocab_size`, an image item for a model with `supports_images() == false`, `dimensions` out of range, an unsupported `encoding_format`, or `model` not matching the served id. |
| 501 | `not_implemented` | No embedding model loaded: `No embedding model loaded; start with -m <embedding checkpoint> or --embedding-model <path>`. |
| 503 | `server_busy` | The bounded worker queue is full (`--embedding-queue-depth`). |
| 504 | `server_timeout` | The worker did not reply within `--embedding-request-timeout-secs`. |
| 500 | `server_error` | The forward pass failed. |

## Server flags

| Flag | Env | Default | Meaning |
|------|-----|---------|---------|
| `-m <embedding checkpoint>` | `LLAMA_ARG_MODEL` | | Serves the checkpoint on `/v1/embeddings`; chat stays unloaded. |
| `--embedding-model <path or repo-id>` | `LLAMA_ARG_EMBEDDING_MODEL`, `MLXCEL_EMBEDDING_MODEL` | unset | A second checkpoint served on `/v1/embeddings` next to the chat model in `-m`. Resolved and auto-downloaded like `-m`. Combining it with an embedding checkpoint in `-m` is a startup error ("two embedding models"). A load failure of an explicit `--embedding-model` is a startup error; a load failure of an `-m` embedding checkpoint is logged and the route answers `501`. |
| `--embedding-batch-size N` | `MLXCEL_EMBEDDING_BATCH_SIZE` | `16` | Texts per forward pass. |
| `--embedding-max-length N` | `MLXCEL_EMBEDDING_MAX_LENGTH` | derived | Lowers the derived `max_length`. |
| `--embedding-queue-depth N` | `MLXCEL_EMBEDDING_QUEUE_DEPTH` | `8` | Bound on the worker command queue; a full queue returns `503`. |
| `--embedding-request-timeout-secs N` | `MLXCEL_EMBEDDING_REQUEST_TIMEOUT_SECS` | `120` | Per-request reply timeout; `0` falls back to the default. |

Both `mlxcel serve` and `mlxcel-server` accept every flag. The embedding model runs on its own dedicated MLX thread (the model, tokenizer and every array live there), so chat and embeddings never share a stream.

```sh
# Embeddings only
mlxcel-server -m sentence-transformers/all-MiniLM-L6-v2 --port 8080

# Chat plus embeddings
mlxcel-server -m mlx-community/Qwen3-4B-4bit --embedding-model Qwen/Qwen3-Embedding-0.6B

curl -s localhost:8080/v1/embeddings -H 'Content-Type: application/json' \
-d '{"input": ["The weather is lovely today.", "It is so sunny outside!", "He drove to the stadium."]}'
```

## `mlxcel embed`

```text
mlxcel embed -m <path or repo-id> -p "text" [-p "text2" ...] [--image <file> ...]
[--instruction "..."] [--dimensions N] [--max-length N] [--batch-size N] [--json]
```

Prints one vector per input (`[v1, v2, ...]`, one line each; a list of rows for multi-vector models) and, with two or more inputs, the cosine-similarity matrix (for multi-vector models the MaxSim score averaged over the query rows). `--json` prints one object with `embeddings`, `shapes`, `prompt_tokens` and `similarity` instead. This is the offline validation tool for every family: the same loader, pooling, normalization and batching as the server, without a listener.

## Adding a family

1. Add the family module under `src/models/` with an `EmbeddingModel` implementation. Read weights with `crate::embeddings::loader::load_embedding_weights` (module subfolders included, text bf16 rule applied) and resolve the pooling mode with `crate::embeddings::resolve_pooling_mode(model_dir, family_default)`. Build attention masks from `EmbeddingBatch::attention_mask` with the builders above. Return pooled `[B, D]` vectors (`[B, L, D]` with padding rows zeroed for multi-vector families); the engine normalizes and truncates.
2. Replace the family's `not yet supported` arm in `build_family_model` (`src/embeddings/loader.rs`) with the constructor.
3. Quantized checkpoints (`config.quantization = {group_size, bits}`) go through `UnifiedLinear::from_weights` / `UnifiedEmbedding::from_weights`, which accept a tensor with or without `.scales`; `quantization_params(config)` reads the block.
4. Validate with `mlxcel embed` against a real checkpoint and add the family row to [supported-models.md](supported-models.md#embedding-models).

Detection (`src/models/detection.rs`), the `ModelType` variant, the `mlxcel arch` entry and the `ModelKind::Embedding` registration in `src/model_metadata.rs` already exist for every family in the table above.
Loading