lattice-embed 0.12.0

SIMD-accelerated vector operations and embedding generation
Documentation
# Design: lattice-embed

## Current Design

**Multi-model embedding service.** 8 model variants in `EmbeddingModel` enum (`model.rs:102-135`), `#[non_exhaustive]` for future extension:

| Variant                        | Architecture  | Native Dims | Max Tokens | MRL | Notes                                   |
| ------------------------------ | ------------- | ----------- | ---------- | --- | --------------------------------------- |
| `BgeSmallEnV15` (enum default) | BERT encoder  | 384         | 512        | no  | English-only                            |
| `BgeBaseEnV15`                 | BERT encoder  | 768         | 512        | no  | English-only                            |
| `BgeLargeEnV15`                | BERT encoder  | 1024        | 512        | no  | English-only                            |
| `MultilingualE5Small`          | BERT encoder  | 384         | 512        | no  | Recommended production default          |
| `MultilingualE5Base`           | BERT encoder  | 768         | 512        | no  | Multilingual                            |
| `Qwen3Embedding0_6B`           | Qwen3 decoder | 1024        | 2048       | yes | Last-token pooling, instruction prefix  |
| `Qwen3Embedding4B`             | Qwen3 decoder | 2560        | 2048       | yes | Sharded safetensors, instruction prefix |
| `TextEmbedding3Small`          | OpenAI API    | 1536        | 8191       | no  | Remote only, no local weights           |

`ModelConfig` (`model.rs:381`) pairs a model variant with optional `output_dim` for MRL truncation. `validate()` enforces bounds. `EmbeddingKey` (`types.rs:119`) identifies an embedding space, but `ModelConfig` builds no `EmbeddingKey` itself -- callers construct one directly (`EmbeddingKey::new`, `types.rs:136`). Dimension separation is enforced downstream instead: different `output_dim` values change `dimensions()` (`model.rs:447-456`), which feeds `EmbeddingCache::compute_key()`'s hashed input, keeping differently-dimensioned embeddings in separate cache entries. Production default is `multilingual-e5-small` via application config, not the enum default (`BgeSmallEnV15`). The enum default serves code paths without application config (tests, standalone usage). Qwen3 models require instruction prefix for queries (`model.rs:227-234`): `"Instruct: Given a web search query, retrieve relevant passages...\nQuery: "`. Documents get no prefix (`document_instruction()` returns `None`).

**MRL truncation.** Matryoshka Representation Learning allows Qwen3-Embedding models to produce native-dimension vectors (1024d for 0.6B, 2560d for 4B) that are truncated to smaller dimensions while preserving quality. Only Qwen3 models support this (`supports_output_dim()` at `model.rs:263-268`); non-MRL models reject any `output_dim` value (`model.rs:420-425`). Truncation to 1024d or 1536d retains >95% quality (per Qwen3 paper) at 40-60% storage reduction. Implementation: `NativeEmbeddingService` passes `output_dim` to `QwenModel::set_output_dim()` at `native.rs:230`. Truncation happens inside the inference crate after forward pass + pooling, before L2 normalization -- the embed crate only passes the config. Min dimension: `MIN_MRL_OUTPUT_DIM = 32` (`model.rs:375`). Max: model's `native_dimensions()`. `LATTICE_EMBED_DIM` env var overrides `output_dim` for Qwen models (`native.rs:74-93`). Different `output_dim` values produce different `dimensions()` (`model.rs:447-456`) results, which enter `EmbeddingCache::compute_key()`'s hashed input (`cache.rs:149-167`), keeping differently-dimensioned embeddings in separate cache entries -- mixing dimensionalities together would produce garbage similarity scores.

**Two-layer cache.** Layer 1 (embed crate): in-memory sharded LRU. `EmbeddingCache` at `cache.rs:97` -- 16 shards, each `RwLock<LruCache<CacheKey, Arc<[f32]>>>` with per-shard `AtomicU64` hit/miss counters. Shard selection: `key[0] & 0xF`. Default capacity 4000 entries (250/shard). SHA-256-based cache key (`compute_key`, `cache.rs:149-167`): hashes `text_bytes` followed by `"{model}:{model.key_version()}:{dimensions}:{role.cache_tag()}"` -- different MRL dims and different roles (e.g. query vs. document) enter the hashed input as distinct bytes, so the same text embedded for different purposes keeps those namespaces separate in the digest input rather than sharing an entry. Eviction: LRU per-shard. Layer 2 (inference crate): `QwenModel` persistent disk cache at `~/.lattice/cache/embed_{model}_{dim}d.bin` (`native.rs:219-248`). Loaded at model init, saved explicitly. Uses `HashMap<u64, Vec<f32>>` with flush-all eviction at 10K cap (see `crates/inference/DESIGN.md` KC-3). BERT models have no persistent cache. The two layers use different key formats: Layer 1 uses SHA-256 of text + model/role metadata (embed crate doesn't have token IDs), Layer 2 uses hash of token ID sequence (inference crate operates post-tokenization). This avoids coupling between crates. `CachedEmbeddingService` (`cached.rs:42`) wraps any `EmbeddingService`: compute SHA-256 keys -> batch `get_many()` -> embed misses only -> `put_many()` -> merge preserving order. FP-035 fix at `cached.rs:153-159` validates embedding count matches input count.

**Batch processing.** Entry: `NativeEmbeddingService::embed()` at `native.rs:326-357`. Input validation: model mismatch -> error, empty input -> error, batch <= 1000 (`DEFAULT_MAX_BATCH_SIZE`), per-text <= 32768 UTF-8 bytes (`MAX_TEXT_BYTES`). Model loads exactly once via `OnceLock` on a blocking thread (`native.rs:161-189`) -- cancellation-safe, survives async drops. BERT models: native batch processing via `BertModel::encode_batch()`. Qwen models: sequential per-item `QwenModel::encode()` (checks internal cache per item) -- sequential because `QwenModel` holds a `Mutex<ForwardBuffers>` preventing parallel forward passes. Batch parallelism would require per-item buffer allocation, adding memory pressure for marginal gain (GPU forward is already fast). Output: `Vec<Vec<f32>>` with order preserved. Core trait: `EmbeddingService` (`service/mod.rs:55-86`) with `embed()`, `embed_one()`, `model_config()`, `supports_model()`, `name()`.

## Alternatives Rejected

| Alternative                                             | Rejected Because                                                                                                                                                                                                        | Date |
| ------------------------------------------------------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ---- |
| Single global LRU cache (no sharding)                   | Write contention under concurrent requests. 16-shard design gives ~16x write throughput at cost of slightly less optimal global LRU ordering. SHA-256 key distribution ensures balanced sharding.                       | --   |
| Runtime dimension selection (no ModelConfig validation) | Would allow invalid dimensions (0, negative, above native) to reach inference. Validation at config construction catches errors early with clear messages.                                                              | --   |
| Unified cache key format across both layers             | Layer 1 needs text-level keys (embed crate doesn't have token IDs). Layer 2 needs token-level keys (inference crate operates post-tokenization). Forcing one format adds unnecessary coupling between crates.           | --   |
| Async model loading (non-blocking)                      | Model load involves mmap + GPU init -- CPU-bound, not I/O-bound. `spawn_blocking` + `OnceLock` is simpler and correct. Async model load would require `async fn init()` on the trait, complicating all implementations. | --   |

## Invariants

**INV-1.** Different `output_dim` values produce different `dimensions()` results (`model.rs:447-456`), which feed distinct `EmbeddingCache::compute_key()` hashes (`cache.rs:149-167`) -- `ModelConfig` builds no `EmbeddingKey` itself; that type is caller-constructed (`EmbeddingKey::new`, `types.rs:136`). If broken: embeddings of different dimensionality share a cache entry, returning a wrong-dimension embedding.
**INV-2.** MRL only for models where `supports_output_dim()` returns true (Qwen3 only, `model.rs:263-268`). If broken: truncation of non-MRL model produces semantically invalid embeddings (BERT embeddings are not trained for truncation).
**INV-3.** MRL dimension in `[MIN_MRL_OUTPUT_DIM(32), native_dimensions()]` (`model.rs:426-437`). If broken: dimension below training minimum degrades quality severely; above native is padding with zeros.
**INV-4.** Model loads exactly once per `NativeEmbeddingService` lifetime (`OnceLock` at `native.rs:63-64`). If broken: double model load wastes memory, or concurrent load attempts race.
**INV-5.** Cache key hashes the text bytes followed by a formatted string of the model, its key version, its dimensions, and the role tag (`compute_key`, `cache.rs:149-167`). If broken: same text with a different model, dimension, or role shares a cache entry, returning a wrong-dimension or wrong-role embedding.
**INV-6.** `CachedEmbeddingService` validates `embeddings.len() == texts_to_embed.len()` (FP-035, `cached.rs:153-159`). If broken: cache merge produces misaligned results (embedding N assigned to text M).
**INV-7.** Batch size <= 1000, text length <= `MAX_TEXT_BYTES` (32768 UTF-8 bytes, `str::len()`). If broken: OOM on large batches or excessive tokenization time. Validated at both `native.rs` and `cached.rs` entry points.
**INV-8.** Requested model must equal loaded model (`native.rs:327-332`). If broken: model produces embeddings in wrong dimension/space.

## Known Concerns Acknowledged

**KC-1.** Enum default (`BgeSmallEnV15`) differs from production default (`multilingual-e5-small`). Intentional: enum default serves code paths without application config (tests, standalone usage). Production routing goes through application config. Documenting both prevents confusion.

**KC-2.** Qwen models process batch items sequentially, not in parallel (`native.rs:29-32`). Intentional: `QwenModel`'s internal cache check is per-item, and the model holds a `Mutex<ForwardBuffers>` preventing parallel forward passes. Batch parallelism would require per-item buffer allocation, adding memory pressure for marginal gain (GPU forward is already fast). (ref: `native.rs:29-32`)

**KC-3.** `EmbeddingCache` marked Unstable in public API (`lib.rs:87-94`). Intentional: cache internals (shard count, key format, LRU parameters) may change. The stable surface is `EmbeddingService` trait + `ModelConfig`.

**KC-4.** 36 unsafe blocks in the `simd/` module (`distance.rs`, `normalize.rs`, `cosine.rs`, `dot_product.rs`, `quantized.rs`, `binary.rs`, `int4.rs`). Intentional: NEON/AVX2 SIMD intrinsics for vectorized distance computation and normalization in performance-critical embedding operations. TCB-justified -- safe abstractions would impose measurable overhead on hot-path distance calculations.

## Baseline Metrics

| Dimension | Metric                                       | Value                   | Measured   | Threshold | Command                                                           |
| --------- | -------------------------------------------- | ----------------------- | ---------- | --------- | ----------------------------------------------------------------- |
| perf      | Embed latency (mE5-small, 100 texts, cached) | pending                 | --         | +20%      | `cargo bench -p lattice-embed --bench embed`                      |
| perf      | Embed latency (Qwen3-0.6B, 100 texts, cold)  | pending                 | --         | +20%      | `cargo bench -p lattice-embed --bench embed`                      |
| perf      | Cache hit rate (steady state)                | pending                 | --         | -10%      | (logged in cache stats)                                           |
| quality   | MRL 1024d vs native 2560d cosine similarity  | >0.95 (per Qwen3 paper) | 2026-04-29 | <0.93     | pending eval harness                                              |
| security  | unsafe block count                           | 36                      | 2026-04-29 | +5        | `grep -rw "unsafe {" crates/embed/src/ --include="*.rs" \| wc -l` |
| quality   | Model variant count                          | 8                       | 2026-04-29 | -1        | count `EmbeddingModel` enum variants                              |

## Change Protocol

1. Read this DESIGN.md before modifying any file in `crates/embed/`.
2. Check Baseline Metrics -- run the measurement commands and compare against thresholds.
3. Check Known Concerns -- ensure your change doesn't re-introduce a concern already acknowledged.
4. Re-measure baselines after your change and update the table if values shift.