Skip to content

fix: Bound embedding input length and attention memory per batch - #837

Draft
buildingjoshbetter wants to merge 2 commits into
maintenance/reranker-token-budget-v1from
maintenance/embed-token-budget-v1
Draft

buildingjoshbetter wants to merge 2 commits into
maintenance/reranker-token-budget-v1from
maintenance/embed-token-budget-v1

Conversation

@buildingjoshbetter

Copy link
Copy Markdown
Owner

Embedding models were run with their full input limit (32,768 tokens for Qwen3-Embedding-0.6B) and a fixed 32-text batch. Attention memory grows with batch × length², so one batch of long texts could need tens of GB. On a Mac it fails on MPS first and is then retried on CPU. This applies the same dynamic bound as the reranker change (#836) to every embedding path and tier.

  • Caps embedder input at 2,048 tokens. That is above the longest stored memory, so existing vectors do not change; only oversized texts such as pasted documents are truncated.
  • Embeds texts longest-first in batches sized so batch × tokens² stays within the cost of 32 texts of 1,024 tokens. The budget is halved below 6 GiB of available memory and quartered below 2 GiB.
  • Keeps the original single encode call, with no measurement, for single texts, tensor or list outputs, and batches that fit even at the smallest budget.
  • Applies to every SentenceTransformer embedder: base/pro Qwen3, legacy MiniLM/bge-small (already below the cap) and the custom tier. The edge static model2vec embedder is untouched.
  • Covers every construction site (local, prepared target, model server, custom), the local encode wrappers and the model server embed loop. The single-query fast lane is bounded by the cap.

Real Qwen3-Embedding-0.6B on CPU, 32 synthetic texts per call, 16 GiB container, peak cgroup memory before → after:

Text length Before After Vectors
50 words 1.50 GB, 0.83 s 0.66 GB, 0.59 s identical
800 words 1.32 GB, 23.1 s 1.32 GB, 23.5 s identical
3,000 words 3.82 GB, 129 s 1.03 GB, 61 s change (2,048-token cap)
12,000 words killed at 16 GiB 1.05 GB, 61 s before never completed

Single-query latency: 30.77 ms → 30.87 ms median over 200 calls (within noise).

Validation at aa37923, stacked on #836:

  • Combined suite: 4,529 tests / 5,048 subtests passed (14 new tests), 19 skips, 3 network cases deselected.
  • Focused suites: 782 / 2,336 and 1,327 / 1,469 passed.
  • Installed wheel smoke: 95 Python source files verified.
  • Separate original gates: 739 native source/inert cases and 276 memory cases passed.
  • Independent review accepted.

tests/test-model-daemon-admission-297.py pins the exact bytes of mps_utils.py to guard against unintended edits. The pin is updated to the reviewed bytes; the mps_utils.py change is the three guarded budgeted_encode call sites.

Base is maintenance/reranker-token-budget-v1 (#836). Draft; no merges or releases.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant