From 9482784662325c0d0fcb0f3e2056d9399bd111dd Mon Sep 17 00:00:00 2001
From: buildingjoshbetter
 <236675935+buildingjoshbetter@users.noreply.github.com>
Date: Tue, 13 Oct 2026 20:45:30 -0500
Subject: [PATCH 1/6] fix: Bound embedding input length and attention memory
 per batch

---
 tests/test-embed-token-budget.py         | 127 +++++++++++++++++++++++
 tests/test-model-daemon-admission-297.py |   2 +-
 truememory/_embed_budget.py              |  56 ++++++++++
 truememory/_rerank_budget.py             |  17 +--
 truememory/embedding_target.py           |  17 +--
 truememory/model_server.py               |  41 ++++++--
 truememory/mps_utils.py                  |  10 ++
 truememory/vector_search.py              |  17 ++-
 8 files changed, 263 insertions(+), 24 deletions(-)
 create mode 100644 tests/test-embed-token-budget.py
 create mode 100644 truememory/_embed_budget.py

diff --git a/tests/test-embed-token-budget.py b/tests/test-embed-token-budget.py
new file mode 100644
index 00000000..8530cbe8
--- /dev/null
+++ b/tests/test-embed-token-budget.py
@@ -0,0 +1,127 @@
+"""Embedding input cap and token-budgeted batches bound attention memory."""
+from __future__ import annotations
+
+import numpy as np
+import pytest
+
+from truememory import _embed_budget as embed_budget
+from truememory._embed_budget import (
+    EMBED_MAX_TOKENS, EMBED_TOKEN_BUDGET, budgeted_encode, cap_embed_length, embed_token_budget,
+    text_token_lengths,
+)
+
+
+class FakeEmbedder:
+    """Word count stands in for tokens; each vector depends only on its text."""
+
+    def __init__(self, max_seq_length: int | None = 32768) -> None:
+        self.max_seq_length = max_seq_length
+        self.calls: list[tuple[int, int]] = []
+
+    def tokenizer(self, texts, truncation, max_length):
+        assert truncation is True
+        return {"input_ids": [[0] * min(len(text.split()) + 2, max_length) for text in texts]}
+
+    def encode(self, texts, batch_size=32, show_progress_bar=False, **kwargs):
+        self.calls.append((len(texts), batch_size))
+        return np.array([[len(text.split()), len(text) % 7] for text in texts], dtype=np.float32)
+
+
+def _texts(*word_counts: int) -> list[str]:
+    return [" ".join(["word"] * count) for count in word_counts]
+
+
+def test_cap_lowers_long_limits_only():
+    for current, expected in ((32768, EMBED_MAX_TOKENS), (512, 512), (None, None)):
+        model = FakeEmbedder(current)
+        cap_embed_length(model)
+        assert model.max_seq_length == expected
+    cap_embed_length(object())
+
+
+def test_budget_shrinks_under_memory_pressure():
+    assert embed_token_budget(16 * 1024 ** 3) == EMBED_TOKEN_BUDGET
+    assert embed_token_budget(4 * 1024 ** 3) == EMBED_TOKEN_BUDGET // 2
+    assert embed_token_budget(1024 ** 3) == EMBED_TOKEN_BUDGET // 4
+
+
+def test_texts_up_to_1024_tokens_keep_one_unchanged_call(monkeypatch):
+    monkeypatch.setattr(embed_budget, "embed_token_budget", lambda: EMBED_TOKEN_BUDGET)
+    model = FakeEmbedder(EMBED_MAX_TOKENS)
+    texts = _texts(*([1000] * 40 + [5] * 40))
+    values = budgeted_encode(model, texts, batch_size=32, show_progress_bar=False)
+    assert model.calls == [(80, 32)]
+    assert values.shape == (80, 2)
+
+
+def test_long_texts_split_longest_first_and_keep_order(monkeypatch):
+    monkeypatch.setattr(embed_budget, "embed_token_budget", lambda: EMBED_TOKEN_BUDGET)
+    model = FakeEmbedder(EMBED_MAX_TOKENS)
+    texts = _texts(*([20] * 60 + [5000] * 12))
+    values = budgeted_encode(model, texts, batch_size=32, show_progress_bar=False)
+    assert np.array_equal(values, FakeEmbedder().encode(texts))
+    assert sum(size for size, _ in model.calls) == 72
+    assert model.calls[0] == (EMBED_TOKEN_BUDGET // EMBED_MAX_TOKENS ** 2,) * 2
+    assert max(text_token_lengths(model, texts)) == EMBED_MAX_TOKENS
+    assert all(batch == size <= 32 for size, batch in model.calls)
+
+
+def test_tensor_output_and_static_models_keep_single_call():
+    model = FakeEmbedder(EMBED_MAX_TOKENS)
+    budgeted_encode(model, _texts(5000, 5000), batch_size=32, convert_to_tensor=True)
+    assert model.calls == [(2, 32)]
+
+    class Static:
+        calls = []
+
+        def encode(self, texts, **kwargs):
+            self.calls.append(len(texts))
+            return np.zeros((len(texts), 2), dtype=np.float32)
+
+    static = Static()
+    from truememory.mps_utils import encode_with_model_ownership
+    encode_with_model_ownership(static, _texts(5000, 5000), batch_size=32)
+    assert static.calls == [2]
+
+
+def test_local_encode_wrappers_use_the_budget(monkeypatch):
+    monkeypatch.setattr(embed_budget, "embed_token_budget", lambda: EMBED_TOKEN_BUDGET)
+    from truememory.mps_utils import encode_with_model_ownership, encode_with_mps_fallback
+    for wrapper in (encode_with_model_ownership, encode_with_mps_fallback):
+        model = FakeEmbedder(EMBED_MAX_TOKENS)
+        texts = _texts(*([10] * 30 + [4000] * 9))
+        values = wrapper(model, texts, batch_size=32, show_progress_bar=False)
+        assert np.array_equal(values, FakeEmbedder().encode(texts))
+        assert model.calls[0][0] == EMBED_TOKEN_BUDGET // EMBED_MAX_TOKENS ** 2
+
+
+def test_model_server_embed_uses_budgeted_slices(monkeypatch):
+    from truememory.model_server import ModelServer, _EmbedState
+
+    monkeypatch.setattr(embed_budget, "embed_token_budget", lambda: EMBED_TOKEN_BUDGET)
+    server = ModelServer()
+    model = FakeEmbedder(EMBED_MAX_TOKENS)
+    server._embed_state = _EmbedState(model=model, tier="", model_id="model2vec")
+    texts = _texts(*([30] * 50 + [3000] * 10))
+    response = server.handle_request({"op": "embed", "texts": texts, "tier": ""})
+    assert response["ok"] is True
+    assert np.array_equal(np.asarray(response["vectors"]), FakeEmbedder().encode(texts))
+    assert model.calls[0] == (EMBED_TOKEN_BUDGET // EMBED_MAX_TOKENS ** 2,) * 2
+    assert all(size <= 32 for size, _ in model.calls)
+
+
+def test_model_server_short_embed_unchanged():
+    from truememory.model_server import ModelServer, _EmbedState
+
+    server = ModelServer()
+    model = FakeEmbedder(EMBED_MAX_TOKENS)
+    server._embed_state = _EmbedState(model=model, tier="", model_id="model2vec")
+    response = server.handle_request({"op": "embed", "texts": _texts(*range(5, 45)), "tier": ""})
+    assert response["ok"] is True
+    assert model.calls == [(32, 32), (8, 32)]
+
+
+@pytest.mark.parametrize("length", [1, 1024, 1025, 2048])
+def test_batch_size_bounds(length):
+    from truememory._rerank_budget import budget_batch_size
+    assert 1 <= budget_batch_size(length, 32, EMBED_TOKEN_BUDGET) <= 32
diff --git a/tests/test-model-daemon-admission-297.py b/tests/test-model-daemon-admission-297.py
index abd47447..3fdbf6e4 100644
--- a/tests/test-model-daemon-admission-297.py
+++ b/tests/test-model-daemon-admission-297.py
@@ -201,7 +201,7 @@ def test_original_module_and_console_launch_contract_is_preserved(self) -> None:
         self.assertIn('[executable, "-m", "truememory.model_server"]', client)
         self.assertIn('truememory-model-server = "truememory.model_server:main"', (ROOT / "pyproject.toml").read_text())
         self.assertEqual(hashlib.sha256((ROOT / "truememory/mps_utils.py").read_bytes()).hexdigest(),
-                         "7057b1e5cdbd23d4de625e6e0c999d2a71455df5fec421842dd9841dabe1b69b")
+                         "adbe1e27ca5dff795f98b2ea3937b9174205ddd8ec327e8316dd9a56b5e3159a")
 
 
 if __name__ == "__main__":
diff --git a/truememory/_embed_budget.py b/truememory/_embed_budget.py
new file mode 100644
index 00000000..e80b0ccf
--- /dev/null
+++ b/truememory/_embed_budget.py
@@ -0,0 +1,56 @@
+"""Embedding input cap and token-budgeted batches that bound attention memory."""
+from __future__ import annotations
+
+from truememory._rerank_budget import budgeted_slices, memory_pressure_divisor
+
+# Above the longest stored memory, so existing vectors are unchanged; only
+# oversized texts such as pasted documents are truncated.
+EMBED_MAX_TOKENS = 2048
+# Keeps the default 32-text batch for every text up to 1,024 tokens.
+EMBED_TOKEN_BUDGET = 32 * 1024 * 1024
+
+
+def cap_embed_length(model: object) -> None:
+    """Lower an embedder's integer input limit to the cap; static models are untouched."""
+    current = getattr(model, "max_seq_length", None)
+    if type(current) is int and current > EMBED_MAX_TOKENS:
+        model.max_seq_length = EMBED_MAX_TOKENS
+
+
+def embed_token_budget(available: int | None = None) -> int:
+    return EMBED_TOKEN_BUDGET // memory_pressure_divisor(available)
+
+
+def measurable(model: object) -> bool:
+    return getattr(model, "tokenizer", None) is not None and type(getattr(model, "max_seq_length", None)) is int
+
+
+def text_token_lengths(model: object, texts: object) -> list[int] | None:
+    """Truncated token length of each text, or None when it cannot be measured."""
+    if not measurable(model) or not isinstance(texts, list) or not all(isinstance(text, str) for text in texts):
+        return None
+    if not texts:
+        return []
+    encoded = model.tokenizer(texts, truncation=True, max_length=model.max_seq_length)
+    return [len(ids) for ids in encoded["input_ids"]]
+
+
+def budgeted_encode(model: object, texts: object, **kwargs: object):
+    """SentenceTransformer.encode with long texts embedded in smaller batches.
+
+    Inputs whose default batches already fit keep the original single call.
+    """
+    batch_size = kwargs.get("batch_size", 32)
+    plain = (type(batch_size) is int and batch_size > 0 and not kwargs.get("convert_to_tensor")
+             and kwargs.get("output_value", "sentence_embedding") == "sentence_embedding")
+    lengths = text_token_lengths(model, texts) if plain else None
+    slices = None if lengths is None else budgeted_slices(lengths, batch_size, embed_token_budget())
+    if slices is None:
+        return model.encode(texts, **kwargs)
+    import numpy as np
+    results: list = [None] * len(texts)
+    for indices in slices:
+        values = model.encode([texts[index] for index in indices], **{**kwargs, "batch_size": len(indices)})
+        for position, index in enumerate(indices):
+            results[index] = values[position]
+    return np.stack(results)
diff --git a/truememory/_rerank_budget.py b/truememory/_rerank_budget.py
index c2688921..b21b90fb 100644
--- a/truememory/_rerank_budget.py
+++ b/truememory/_rerank_budget.py
@@ -17,19 +17,24 @@ def cap_reranker_length(model: object) -> None:
         model.max_seq_length = RERANK_MAX_TOKENS
 
 
-def rerank_token_budget(available: int | None = None) -> int:
-    """Per-batch attention budget, reduced while the host is short of memory."""
+def memory_pressure_divisor(available: int | None = None) -> int:
+    """1 normally, 2 below 6 GiB available and 4 below 2 GiB."""
     if available is None:
         try:
             import psutil
             available = psutil.virtual_memory().available
         except (ImportError, OSError):
-            return RERANK_TOKEN_BUDGET
+            return 1
     if available < 2 * 1024 ** 3:
-        return RERANK_TOKEN_BUDGET // 4
+        return 4
     if available < 6 * 1024 ** 3:
-        return RERANK_TOKEN_BUDGET // 2
-    return RERANK_TOKEN_BUDGET
+        return 2
+    return 1
+
+
+def rerank_token_budget(available: int | None = None) -> int:
+    """Per-batch attention budget, reduced while the host is short of memory."""
+    return RERANK_TOKEN_BUDGET // memory_pressure_divisor(available)
 
 
 def pair_token_lengths(model: object, pairs: list) -> list[int] | None:
diff --git a/truememory/embedding_target.py b/truememory/embedding_target.py
index 5604881b..511068b1 100644
--- a/truememory/embedding_target.py
+++ b/truememory/embedding_target.py
@@ -113,13 +113,18 @@ def admit() -> None:
     if target.model_id == "qwen3_256":
         model_kwargs = {"attn_implementation": "eager"} if sys.platform == "darwin" else None
         admit()
-        return SentenceTransformer(
+        model = SentenceTransformer(
             "Qwen/Qwen3-Embedding-0.6B", truncate_dim=256,
             config_kwargs={"use_cache": False},
             model_kwargs=model_kwargs, device=device,
         )
-    admit()
-    return SentenceTransformer(
-        target.model_id, truncate_dim=target.dimension,
-        trust_remote_code=False, device=device,
-    )
+    else:
+        admit()
+        model = SentenceTransformer(
+            target.model_id, truncate_dim=target.dimension,
+            trust_remote_code=False, device=device,
+        )
+    if type(getattr(model, "max_seq_length", None)) is int:
+        from truememory._embed_budget import cap_embed_length
+        cap_embed_length(model)
+    return model
diff --git a/truememory/model_server.py b/truememory/model_server.py
index bed0fbdb..e51f1d12 100644
--- a/truememory/model_server.py
+++ b/truememory/model_server.py
@@ -722,13 +722,17 @@ def _build_embed_model(model_id: str, device: str | None):
             mkwargs = {}
             if sys.platform == "darwin":
                 mkwargs["attn_implementation"] = "eager"
-            return SentenceTransformer(
+            model = SentenceTransformer(
                 "Qwen/Qwen3-Embedding-0.6B",
                 truncate_dim=256,
                 config_kwargs={"use_cache": False},
                 model_kwargs=mkwargs or None,
                 device=device,
             )
+            if type(getattr(model, "max_seq_length", None)) is int:
+                from truememory._embed_budget import cap_embed_length
+                cap_embed_length(model)
+            return model
         if model_id not in ("model2vec", "minilm", "bge-small", "qwen3_256"):
             # Custom model: require explicit opt-in for arbitrary downloads
             if os.environ.get("TRUEMEMORY_CUSTOM_ALLOW_DOWNLOAD", "").strip() != "1":
@@ -748,11 +752,15 @@ def _build_embed_model(model_id: str, device: str | None):
             from truememory.mps_utils import ensure_mps_memory_budget
             ensure_mps_memory_budget(device)
             from sentence_transformers import SentenceTransformer
-            return SentenceTransformer(
+            model = SentenceTransformer(
                 model_id, truncate_dim=custom_dim,
                 trust_remote_code=False,
                 device=device,
             )
+            if type(getattr(model, "max_seq_length", None)) is int:
+                from truememory._embed_budget import cap_embed_length
+                cap_embed_length(model)
+            return model
         from model2vec import StaticModel
         return StaticModel.from_pretrained(
             "minishlab/potion-base-8M", force_download=False
@@ -1336,6 +1344,7 @@ def _handle_request_inner(self, request: dict) -> dict:
                     vectors = None
                     offset = 0
                     order = None
+                    steps = None
                     while offset < len(texts) or vectors is None:
                         retry = None
                         with deadline.locked(self._lock):
@@ -1352,18 +1361,36 @@ def _handle_request_inner(self, request: dict) -> dict:
                                          else self._get_prepared_embed_model(target, deadline))
                                 self._check_serving_embed_model(model, expected_target)
                                 order = self._embed_global_order(model, texts, limit, deadline)
+                                if (getattr(model, "tokenizer", None) is not None
+                                        and type(getattr(model, "max_seq_length", None)) is int):
+                                    from truememory._embed_budget import embed_token_budget, text_token_lengths
+                                    from truememory._rerank_budget import budgeted_slices
+                                    lengths = text_token_lengths(model, texts)
+                                    slices = (None if lengths is None else
+                                              budgeted_slices(lengths, limit, embed_token_budget()))
+                                    if slices is not None:
+                                        # Long texts: smaller longest-first batches bound attention memory.
+                                        order = np.array([index for group in slices for index in group], dtype=np.intp)
+                                        steps, start = {}, 0
+                                        for group in slices:
+                                            steps[start] = len(group)
+                                            start += len(group)
                             while offset < len(texts) or vectors is None:
                                 deadline.check()
-                                indices = None if order is None else self._embed_slice_indices(
-                                    model, texts, order, offset, limit, deadline,
-                                )
+                                step = limit if steps is None else steps[offset]
+                                if steps is not None:
+                                    indices = order[offset:offset + step]
+                                else:
+                                    indices = None if order is None else self._embed_slice_indices(
+                                        model, texts, order, offset, limit, deadline,
+                                    )
                                 batch = (texts[offset:offset + limit] if indices is None else
                                          [texts[index] for index in indices])
                                 self._check_process_memory("embed slice", deadline)
                                 deadline.check()
                                 recover = False
                                 try:
-                                    values = model.encode(batch, batch_size=limit, show_progress_bar=False)
+                                    values = model.encode(batch, batch_size=step, show_progress_bar=False)
                                 except RuntimeError as exc:
                                     from truememory.mps_utils import is_mps_oom
                                     if not is_mps_oom(exc):
@@ -1387,7 +1414,7 @@ def _handle_request_inner(self, request: dict) -> dict:
                             self._check_process_memory("embed CPU retry", deadline)
                             deadline.check()
                             log.warning("MPS OOM during encoding; retrying microbatch on CPU")
-                            values = model.encode(retry, batch_size=limit, show_progress_bar=False)
+                            values = model.encode(retry, batch_size=len(retry) or limit, show_progress_bar=False)
                             if target is not None or expected_target is not None:
                                 check_target_vectors(target or expected_target, values, len(retry))
                             vectors = _store_batch_result(
diff --git a/truememory/mps_utils.py b/truememory/mps_utils.py
index e8a9e30e..1ae19cb4 100644
--- a/truememory/mps_utils.py
+++ b/truememory/mps_utils.py
@@ -311,6 +311,9 @@ def encode_with_model_ownership(model: object, texts: object, **kwargs: object)
     if isinstance(model, EmbeddingProxy):
         return model.encode(texts, **kwargs)
     with _own_model(model):
+        if getattr(model, "tokenizer", None) is not None and type(getattr(model, "max_seq_length", None)) is int:
+            from truememory._embed_budget import budgeted_encode
+            return budgeted_encode(model, texts, **kwargs)
         return model.encode(texts, **kwargs)
 
 
@@ -323,8 +326,12 @@ def encode_with_mps_fallback(model: object, texts: object, **kwargs: object) ->
         # prevent recall from using the daemon's independent fast lane.
         return model.encode(texts, **kwargs)
 
+    bounded = getattr(model, "tokenizer", None) is not None and type(getattr(model, "max_seq_length", None)) is int
     with _own_model(model):
         try:
+            if bounded:
+                from truememory._embed_budget import budgeted_encode
+                return budgeted_encode(model, texts, **kwargs)
             return model.encode(texts, **kwargs)
         except RuntimeError as exc:
             if not is_mps_oom(exc):
@@ -335,4 +342,7 @@ def encode_with_mps_fallback(model: object, texts: object, **kwargs: object) ->
         if hasattr(model, "to"):
             model.to("cpu")
         flush_mps_cache()
+        if bounded:
+            from truememory._embed_budget import budgeted_encode
+            return budgeted_encode(model, texts, **kwargs)
         return model.encode(texts, **kwargs)
diff --git a/truememory/vector_search.py b/truememory/vector_search.py
index 9510b5f9..5590715e 100644
--- a/truememory/vector_search.py
+++ b/truememory/vector_search.py
@@ -374,6 +374,9 @@ def get_model():
                         from model2vec import StaticModel
                         _model = StaticModel.from_pretrained("minishlab/potion-base-8M", force_download=False)
                         _embedding_dim = 256
+                    if type(getattr(_model, "max_seq_length", None)) is int:
+                        from truememory._embed_budget import cap_embed_length
+                        cap_embed_length(_model)
                     if budget_owner is not None:
                         bind_model_owner(_model, budget_owner)
             except BaseException:
@@ -615,10 +618,16 @@ def encode(
                                                  native_tensor_shape=(count, self.target.dimension)),
                                       known_bytes=payload_bytes, retained_bytes=payload_bytes,
                                       deadline=lambda: _target_remaining(expires), group="result assembly")
-                    values = encoder.encode(
-                        [TARGET_PROBE_TEXT] if probe else texts,
-                        batch_size=batch_size, show_progress_bar=False,
-                    )
+                    if (getattr(encoder, "tokenizer", None) is not None
+                            and type(getattr(encoder, "max_seq_length", None)) is int):
+                        from truememory._embed_budget import budgeted_encode
+                        values = budgeted_encode(encoder, [TARGET_PROBE_TEXT] if probe else texts,
+                                                 batch_size=batch_size, show_progress_bar=False)
+                    else:
+                        values = encoder.encode(
+                            [TARGET_PROBE_TEXT] if probe else texts,
+                            batch_size=batch_size, show_progress_bar=False,
+                        )
                     _target_remaining(expires)
                     values = np.asarray(values, dtype=np.float32)
                     check_target_vectors(self.target, values, count)

From aa379232425d7380d63e9cc34e0e027a247c5f09 Mon Sep 17 00:00:00 2001
From: buildingjoshbetter
 <236675935+buildingjoshbetter@users.noreply.github.com>
Date: Tue, 13 Oct 2026 21:01:07 -0500
Subject: [PATCH 2/6] perf: Skip embedding budget measurement when a batch
 always fits

---
 tests/test-embed-token-budget.py | 19 +++++++++++++++++++
 truememory/_embed_budget.py      | 11 ++++++++---
 2 files changed, 27 insertions(+), 3 deletions(-)

diff --git a/tests/test-embed-token-budget.py b/tests/test-embed-token-budget.py
index 8530cbe8..2e714fa6 100644
--- a/tests/test-embed-token-budget.py
+++ b/tests/test-embed-token-budget.py
@@ -20,6 +20,7 @@ def __init__(self, max_seq_length: int | None = 32768) -> None:
 
     def tokenizer(self, texts, truncation, max_length):
         assert truncation is True
+        self.tokenized = getattr(self, "tokenized", 0) + 1
         return {"input_ids": [[0] * min(len(text.split()) + 2, max_length) for text in texts]}
 
     def encode(self, texts, batch_size=32, show_progress_bar=False, **kwargs):
@@ -125,3 +126,21 @@ def test_model_server_short_embed_unchanged():
 def test_batch_size_bounds(length):
     from truememory._rerank_budget import budget_batch_size
     assert 1 <= budget_batch_size(length, 32, EMBED_TOKEN_BUDGET) <= 32
+
+
+def test_single_text_and_list_outputs_skip_measurement():
+    model = FakeEmbedder(EMBED_MAX_TOKENS)
+    budgeted_encode(model, _texts(5000), batch_size=32)
+    budgeted_encode(model, _texts(5000, 5000), batch_size=32, convert_to_numpy=False)
+    assert getattr(model, "tokenized", 0) == 0
+    assert model.calls == [(1, 32), (2, 32)]
+
+
+def test_short_batches_skip_the_memory_read(monkeypatch):
+    def unexpected():
+        raise AssertionError("memory read for a batch that always fits")
+
+    monkeypatch.setattr(embed_budget, "embed_token_budget", unexpected)
+    model = FakeEmbedder(EMBED_MAX_TOKENS)
+    budgeted_encode(model, _texts(*([500] * 32)), batch_size=32)
+    assert model.calls == [(32, 32)]
diff --git a/truememory/_embed_budget.py b/truememory/_embed_budget.py
index e80b0ccf..dba6dcd1 100644
--- a/truememory/_embed_budget.py
+++ b/truememory/_embed_budget.py
@@ -1,7 +1,7 @@
 """Embedding input cap and token-budgeted batches that bound attention memory."""
 from __future__ import annotations
 
-from truememory._rerank_budget import budgeted_slices, memory_pressure_divisor
+from truememory._rerank_budget import budget_batch_size, budgeted_slices, memory_pressure_divisor
 
 # Above the longest stored memory, so existing vectors are unchanged; only
 # oversized texts such as pasted documents are truncated.
@@ -41,10 +41,15 @@ def budgeted_encode(model: object, texts: object, **kwargs: object):
     Inputs whose default batches already fit keep the original single call.
     """
     batch_size = kwargs.get("batch_size", 32)
-    plain = (type(batch_size) is int and batch_size > 0 and not kwargs.get("convert_to_tensor")
+    # One text is always one batch, so single queries skip measurement entirely.
+    plain = (isinstance(texts, list) and len(texts) > 1 and type(batch_size) is int and batch_size > 0
+             and not kwargs.get("convert_to_tensor") and kwargs.get("convert_to_numpy", True) is not False
              and kwargs.get("output_value", "sentence_embedding") == "sentence_embedding")
     lengths = text_token_lengths(model, texts) if plain else None
-    slices = None if lengths is None else budgeted_slices(lengths, batch_size, embed_token_budget())
+    # Batches that fit even the smallest budget skip the memory read.
+    if lengths is None or budget_batch_size(max(lengths), batch_size, EMBED_TOKEN_BUDGET // 4) >= batch_size:
+        return model.encode(texts, **kwargs)
+    slices = budgeted_slices(lengths, batch_size, embed_token_budget())
     if slices is None:
         return model.encode(texts, **kwargs)
     import numpy as np

From 01da5a8f574e8d274398d1b0c671f05fac719b85 Mon Sep 17 00:00:00 2001
From: buildingjoshbetter
 <236675935+buildingjoshbetter@users.noreply.github.com>
Date: Wed, 14 Oct 2026 03:41:04 -0500
Subject: [PATCH 3/6] perf: Size embedding and rerank batches to the free MPS
 allocation

The token budgets bounded attention by system memory only. On MPS the
process cap is far smaller, so a batch with one long memory went MPS OOM
and moved the model to CPU for the server lifetime. Batches on MPS now
also fit the free share of the cap, using the measured peak footprint
(48 bytes per batch x tokens^2 x head).

Measured on Apple silicon at the 3.25 GiB cap with both models resident:
31 short texts plus one 860-token text and 32 texts of about 540 tokens
both completed on MPS (4.0 s and 15.4 s; previously MPS OOM), cosine
similarity >= 0.9996 against float32 CPU vectors. Reranking 100
candidates including ten at the 1,024-token cap completed on MPS in
6.7 s (previously MPS OOM).
---
 tests/test-embed-token-budget.py | 43 ++++++++++++++++++++++++++++
 truememory/_embed_budget.py      | 11 ++++++--
 truememory/_rerank_budget.py     | 48 +++++++++++++++++++++++++++++++-
 truememory/model_server.py       | 13 +++++----
 4 files changed, 105 insertions(+), 10 deletions(-)

diff --git a/tests/test-embed-token-budget.py b/tests/test-embed-token-budget.py
index 2e714fa6..02857921 100644
--- a/tests/test-embed-token-budget.py
+++ b/tests/test-embed-token-budget.py
@@ -144,3 +144,46 @@ def unexpected():
     model = FakeEmbedder(EMBED_MAX_TOKENS)
     budgeted_encode(model, _texts(*([500] * 32)), batch_size=32)
     assert model.calls == [(32, 32)]
+
+
+def _on_mps(monkeypatch, model, allocated_mib, cap_mib=1536, heads=16):
+    import sys
+    import types
+
+    from truememory import mps_utils
+    model.device = types.SimpleNamespace(type="mps")
+    model.config = types.SimpleNamespace(num_attention_heads=heads)
+    monkeypatch.setattr(mps_utils, "get_mps_memory_budget",
+                        lambda: types.SimpleNamespace(effective_bytes=cap_mib * 1024 ** 2))
+    monkeypatch.setitem(sys.modules, "torch", types.SimpleNamespace(
+        mps=types.SimpleNamespace(current_allocated_memory=lambda: allocated_mib * 1024 ** 2)))
+
+
+def test_mps_batches_fit_the_free_allocation(monkeypatch):
+    monkeypatch.setattr(embed_budget, "embed_token_budget", lambda: EMBED_TOKEN_BUDGET)
+    model = FakeEmbedder(EMBED_MAX_TOKENS)
+    _on_mps(monkeypatch, model, allocated_mib=1190)
+    texts = _texts(*([45] * 31 + [1000]))
+    values = budgeted_encode(model, texts, batch_size=32)
+    # 80% of the free 346 MiB over 16 heads x 48 bytes per score: 377,924 tokens^2.
+    fit = int(346 * 1024 ** 2 * 0.8) // (16 * 48)
+    # Below one long text's 1,004,004 tokens^2, so it is embedded alone.
+    assert fit < 1002 ** 2
+    assert model.calls[0] == (1, 1)
+    assert sum(size for size, _ in model.calls) == 32
+    assert np.array_equal(values, FakeEmbedder().encode(texts))
+
+
+def test_mps_short_batches_keep_one_call(monkeypatch):
+    model = FakeEmbedder(EMBED_MAX_TOKENS)
+    _on_mps(monkeypatch, model, allocated_mib=1190)
+    budgeted_encode(model, _texts(*([45] * 32)), batch_size=32)
+    assert model.calls == [(32, 32)]
+
+
+def test_cpu_models_ignore_the_mps_allocation(monkeypatch):
+    from truememory._rerank_budget import accelerator_token_budget
+    model = FakeEmbedder(EMBED_MAX_TOKENS)
+    _on_mps(monkeypatch, model, allocated_mib=1500)
+    model.device = "cpu"
+    assert accelerator_token_budget(model, EMBED_TOKEN_BUDGET) == EMBED_TOKEN_BUDGET
diff --git a/truememory/_embed_budget.py b/truememory/_embed_budget.py
index dba6dcd1..00e91030 100644
--- a/truememory/_embed_budget.py
+++ b/truememory/_embed_budget.py
@@ -1,7 +1,9 @@
 """Embedding input cap and token-budgeted batches that bound attention memory."""
 from __future__ import annotations
 
-from truememory._rerank_budget import budget_batch_size, budgeted_slices, memory_pressure_divisor
+from truememory._rerank_budget import (
+    accelerator_token_budget, budget_batch_size, budgeted_slices, memory_pressure_divisor,
+)
 
 # Above the longest stored memory, so existing vectors are unchanged; only
 # oversized texts such as pasted documents are truncated.
@@ -46,10 +48,13 @@ def budgeted_encode(model: object, texts: object, **kwargs: object):
              and not kwargs.get("convert_to_tensor") and kwargs.get("convert_to_numpy", True) is not False
              and kwargs.get("output_value", "sentence_embedding") == "sentence_embedding")
     lengths = text_token_lengths(model, texts) if plain else None
+    if lengths is None:
+        return model.encode(texts, **kwargs)
     # Batches that fit even the smallest budget skip the memory read.
-    if lengths is None or budget_batch_size(max(lengths), batch_size, EMBED_TOKEN_BUDGET // 4) >= batch_size:
+    if budget_batch_size(max(lengths), batch_size,
+                         accelerator_token_budget(model, EMBED_TOKEN_BUDGET // 4)) >= batch_size:
         return model.encode(texts, **kwargs)
-    slices = budgeted_slices(lengths, batch_size, embed_token_budget())
+    slices = budgeted_slices(lengths, batch_size, accelerator_token_budget(model, embed_token_budget()))
     if slices is None:
         return model.encode(texts, **kwargs)
     import numpy as np
diff --git a/truememory/_rerank_budget.py b/truememory/_rerank_budget.py
index c42a7418..30574007 100644
--- a/truememory/_rerank_budget.py
+++ b/truememory/_rerank_budget.py
@@ -52,6 +52,51 @@ def memory_pressure_divisor(available: int | None = None) -> int:
     return 1
 
 
+# Peak MPS footprint per batch x tokens^2 x attention head, measured on Apple
+# silicon for Qwen3-Embedding and the ModernBERT reranker with both resident
+# (scores, masks, softmax and the allocator's heaps), with a fifth in reserve.
+_MPS_BYTES_PER_SCORE = 48
+_MPS_HEADROOM_SHARE = 0.8
+
+
+def _attention_heads(model: object) -> int | None:
+    configs = [getattr(model, "config", None), getattr(getattr(model, "model", None), "config", None)]
+    first = getattr(model, "_first_module", None)
+    if callable(first):
+        try:
+            configs.append(getattr(getattr(first(), "auto_model", None), "config", None))
+        except (AttributeError, IndexError, KeyError, StopIteration, TypeError):
+            pass
+    for config in configs:
+        heads = getattr(config, "num_attention_heads", None)
+        if type(heads) is int and heads > 0:
+            return heads
+    return None
+
+
+def accelerator_token_budget(model: object, budget: int) -> int:
+    """Lower *budget* so one batch's attention fits the free MPS allocation.
+
+    MPS allocations are capped per process. A batch above the remaining cap
+    fails with an OOM that degrades the model to CPU for the server lifetime.
+    """
+    device = getattr(model, "device", None)
+    if getattr(device, "type", device) != "mps":
+        return budget
+    from truememory.mps_utils import get_mps_memory_budget
+    policy = get_mps_memory_budget()
+    heads = _attention_heads(model)
+    if policy is None or not policy.effective_bytes or heads is None:
+        return budget
+    try:
+        import torch
+        free = policy.effective_bytes - torch.mps.current_allocated_memory()
+    except (ImportError, AttributeError, RuntimeError):
+        return budget
+    fit = int(max(0, free) * _MPS_HEADROOM_SHARE) // (heads * _MPS_BYTES_PER_SCORE)
+    return max(1, min(budget, fit))
+
+
 def rerank_token_budget(available: int | None = None) -> int:
     """Per-batch attention budget, reduced while the host is short of memory."""
     return RERANK_TOKEN_BUDGET // memory_pressure_divisor(available)
@@ -97,7 +142,8 @@ def budgeted_predict(model: object, pairs: list, batch_size: int):
     """CrossEncoder.predict with long pairs scored in smaller batches."""
     lengths = (pair_token_lengths(model, pairs)
                if type(batch_size) is int and batch_size > 0 else None)
-    slices = None if lengths is None else budgeted_slices(lengths, batch_size, rerank_token_budget())
+    slices = (None if lengths is None else
+              budgeted_slices(lengths, batch_size, accelerator_token_budget(model, rerank_token_budget())))
     if slices is None:
         return model.predict(pairs, batch_size=batch_size, show_progress_bar=False)
     import numpy as np
diff --git a/truememory/model_server.py b/truememory/model_server.py
index 6f3c43f2..ebd3a3f3 100644
--- a/truememory/model_server.py
+++ b/truememory/model_server.py
@@ -1447,10 +1447,10 @@ def _handle_request_inner(self, request: dict) -> dict:
                                 if (getattr(model, "tokenizer", None) is not None
                                         and type(getattr(model, "max_seq_length", None)) is int):
                                     from truememory._embed_budget import embed_token_budget, text_token_lengths
-                                    from truememory._rerank_budget import budgeted_slices
+                                    from truememory._rerank_budget import accelerator_token_budget, budgeted_slices
                                     lengths = text_token_lengths(model, texts)
-                                    slices = (None if lengths is None else
-                                              budgeted_slices(lengths, limit, embed_token_budget()))
+                                    slices = (None if lengths is None else budgeted_slices(
+                                        lengths, limit, accelerator_token_budget(model, embed_token_budget())))
                                     if slices is not None:
                                         # Long texts: smaller longest-first batches bound attention memory.
                                         order = np.array([index for group in slices for index in group], dtype=np.intp)
@@ -1560,11 +1560,12 @@ def _handle_request_inner(self, request: dict) -> dict:
                                         type(getattr(reranker, "max_seq_length", None)) is int
                                         or hasattr(reranker, "max_length")):
                                     from truememory._rerank_budget import (
-                                        budgeted_slices, pair_token_lengths, rerank_token_budget,
+                                        accelerator_token_budget, budgeted_slices, pair_token_lengths,
+                                        rerank_token_budget,
                                     )
                                     lengths = pair_token_lengths(reranker, pairs)
-                                    slices = (None if lengths is None else
-                                              budgeted_slices(lengths, limit, rerank_token_budget()))
+                                    slices = (None if lengths is None else budgeted_slices(
+                                        lengths, limit, accelerator_token_budget(reranker, rerank_token_budget())))
                                 if slices is not None:
                                     # Long pairs: smaller longest-first batches bound attention scratch.
                                     order = np.array([index for group in slices for index in group], dtype=np.intp)

From 7cb8c51a5e9180e201ea9c6c6287e4f7087edd42 Mon Sep 17 00:00:00 2001
From: buildingjoshbetter
 <236675935+buildingjoshbetter@users.noreply.github.com>
Date: Wed, 14 Oct 2026 06:33:44 -0500
Subject: [PATCH 4/6] test: Cover MPS batch sizing in the server loops and
 bounded fast-lane waits

The MPS allocation sizing was only exercised through the local encode
helper. Pin it in the model server's embed and rerank loops, the head
lookup through sentence-transformers modules, and the bookkeeping wait
under a deadline shorter than its 0.25 s bound.
---
 tests/test-embed-token-budget.py           | 27 ++++++++++++++++++++++
 tests/test-model-server-memory-pressure.py | 24 +++++++++++++++++++
 tests/test-reranker-token-budget.py        | 25 ++++++++++++++++++++
 3 files changed, 76 insertions(+)

diff --git a/tests/test-embed-token-budget.py b/tests/test-embed-token-budget.py
index 02857921..c4f34e3e 100644
--- a/tests/test-embed-token-budget.py
+++ b/tests/test-embed-token-budget.py
@@ -187,3 +187,30 @@ def test_cpu_models_ignore_the_mps_allocation(monkeypatch):
     _on_mps(monkeypatch, model, allocated_mib=1500)
     model.device = "cpu"
     assert accelerator_token_budget(model, EMBED_TOKEN_BUDGET) == EMBED_TOKEN_BUDGET
+
+
+def test_server_embed_loop_fits_the_free_mps_allocation(monkeypatch):
+    from truememory.model_server import ModelServer, _EmbedState
+
+    monkeypatch.setattr(embed_budget, "embed_token_budget", lambda: EMBED_TOKEN_BUDGET)
+    model = FakeEmbedder(EMBED_MAX_TOKENS)
+    _on_mps(monkeypatch, model, allocated_mib=1190, cap_mib=4096)
+    server = ModelServer()
+    server._embed_state = _EmbedState(model=model, tier="", model_id="model2vec")
+    texts = _texts(*([30] * 40 + [1000] * 8))
+    response = server.handle_request({"op": "embed", "texts": texts, "tier": ""})
+    assert response["ok"] is True
+    assert np.array_equal(np.asarray(response["vectors"]), FakeEmbedder().encode(texts))
+    # 80% of the free 2,906 MiB over 16 heads x 48 bytes: 3,174,345 tokens^2, so three 1,002-token texts.
+    fit = int(2906 * 1024 ** 2 * 0.8) // (16 * 48)
+    assert model.calls[0][0] == fit // 1002 ** 2 == 3
+
+
+def test_attention_heads_are_read_through_sentence_transformer_modules():
+    import types
+
+    from truememory._rerank_budget import _attention_heads
+    module = types.SimpleNamespace(auto_model=types.SimpleNamespace(config=types.SimpleNamespace(num_attention_heads=16)))
+    model = types.SimpleNamespace(_first_module=lambda: module)
+    assert _attention_heads(model) == 16
+    assert _attention_heads(types.SimpleNamespace()) is None
diff --git a/tests/test-model-server-memory-pressure.py b/tests/test-model-server-memory-pressure.py
index 79db1a12..a9a3338f 100644
--- a/tests/test-model-server-memory-pressure.py
+++ b/tests/test-model-server-memory-pressure.py
@@ -339,3 +339,27 @@ def test_session_release_returns_one_stage(budgeted):
         session.release("synthetic stage")
         assert lease.reserved_known_bytes == 0
 
+
+
+def test_short_deadline_bounds_the_bookkeeping_wait(plain, monkeypatch):
+    plain._publish_embed_state(ms._EmbedState(model=Model("cpu"), tier="", model_id="model2vec"))
+    monkeypatch.setattr(plain, "_build_embed_model", lambda model_id, device: Model(device))
+    held, done = threading.Event(), threading.Event()
+
+    def bookkeeping():
+        with plain._lock:
+            held.set()
+            done.wait(5)
+
+    worker = threading.Thread(target=bookkeeping)
+    worker.start()
+    held.wait()
+    try:
+        started = time.monotonic()
+        plain._handle_fast_embed(["synthetic"], "", ms._RequestDeadline(time.monotonic() + 0.05))
+        assert time.monotonic() - started < 0.25
+    except ms._RequestDeadlineExceeded:
+        assert time.monotonic() - started < 0.25
+    finally:
+        done.set()
+        worker.join()
diff --git a/tests/test-reranker-token-budget.py b/tests/test-reranker-token-budget.py
index 68db22f3..54669a90 100644
--- a/tests/test-reranker-token-budget.py
+++ b/tests/test-reranker-token-budget.py
@@ -198,3 +198,28 @@ def test_model_server_caps_and_slices_legacy_crossencoder(monkeypatch):
     assert response["ok"] is True
     assert server._reranker.max_length == RERANK_MAX_TOKENS
     assert server._reranker.calls[0][0] == RERANK_TOKEN_BUDGET // RERANK_MAX_TOKENS ** 2
+
+
+def test_server_rerank_loop_fits_the_free_mps_allocation(monkeypatch):
+    import sys
+    import types
+
+    from truememory import mps_utils
+    from truememory.model_server import ModelServer
+
+    fake = FakeCrossEncoder()
+    fake.device = types.SimpleNamespace(type="mps")
+    fake.config = types.SimpleNamespace(num_attention_heads=12)
+    monkeypatch.setattr(mps_utils, "get_mps_memory_budget",
+                        lambda: types.SimpleNamespace(effective_bytes=4096 * 1024 ** 2))
+    monkeypatch.setitem(sys.modules, "torch", types.SimpleNamespace(
+        mps=types.SimpleNamespace(current_allocated_memory=lambda: 1790 * 1024 ** 2)))
+    monkeypatch.setattr(ModelServer, "_get_reranker", lambda self, model_name=None, deadline=None: fake)
+    monkeypatch.setattr(rerank_budget, "rerank_token_budget", lambda: RERANK_TOKEN_BUDGET)
+    pairs = _pairs(*([30] * 90 + [2000] * 10))
+    response = ModelServer().handle_request({"op": "rerank", "pairs": pairs, "model_name": "m"})
+    assert response["ok"] is True
+    assert np.allclose(response["scores"], FakeCrossEncoder().predict(pairs, batch_size=len(pairs)), rtol=0, atol=0)
+    # 80% of the free 2,306 MiB over 12 heads x 48 bytes: 3,358,665 tokens^2, three 1,024-token pairs.
+    fit = int(2306 * 1024 ** 2 * 0.8) // (12 * 48)
+    assert fake.calls[0][0] == fit // RERANK_MAX_TOKENS ** 2 == 3

From 18a94e3c559b7276f180c4c7fe61f88b27d95853 Mon Sep 17 00:00:00 2001
From: buildingjoshbetter
 <236675935+buildingjoshbetter@users.noreply.github.com>
Date: Fri, 16 Oct 2026 15:52:19 -0500
Subject: [PATCH 5/6] fix: Size CPU fallback pieces against the RSS budget, not
 free memory

After an MPS OOM the server retried the failed piece, and every piece
after it, at the size chosen for MPS. In the soak the first CPU piece was
11 texts of 1,718 tokens: about 4.16 GB of eager attention per layer, on
top of 3.56 GB resident under a 4 GiB budget. CPU batches were otherwise
sized by free system memory, which the server's budget does not track.

accelerator_token_budget now sizes CPU batches so one batch's attention
(8 bytes per score at the softmax peak, per head) fits 80% of the process
RSS budget minus resident memory: the model server's owner limit, or the
automatic budget in local processes. After an embed or rerank MPS OOM the
server retries the failed piece and re-slices every remaining piece under
that budget, keeping result order. Models whose head count is unknown keep
their previous budget.
---
 tests/test-cpu-fallback-token-budget.py | 122 ++++++++++++++++++++++++
 truememory/_rerank_budget.py            |  37 ++++++-
 truememory/memory_budget.py             |  11 +++
 truememory/model_server.py              |  68 ++++++++++++-
 4 files changed, 232 insertions(+), 6 deletions(-)
 create mode 100644 tests/test-cpu-fallback-token-budget.py

diff --git a/tests/test-cpu-fallback-token-budget.py b/tests/test-cpu-fallback-token-budget.py
new file mode 100644
index 00000000..8a265c56
--- /dev/null
+++ b/tests/test-cpu-fallback-token-budget.py
@@ -0,0 +1,122 @@
+"""After an MPS OOM, CPU pieces are sized against the RSS budget, not free system memory."""
+from __future__ import annotations
+
+import sys
+import types
+
+import numpy as np
+import pytest
+
+from truememory import memory_budget
+from truememory import model_server as ms
+from truememory import mps_utils
+from truememory._embed_budget import EMBED_TOKEN_BUDGET
+from truememory._rerank_budget import accelerator_token_budget, budget_batch_size
+
+GIB = 1024 ** 3
+RESIDENT = 3_560_000_000  # the soak's resident set beside a 4 GiB budget (16 GiB Mac)
+
+
+class FakeRSS:
+    def __init__(self, rss: int) -> None:
+        self.rss = rss
+
+    def Process(self) -> "FakeRSS":
+        return self
+
+    def virtual_memory(self) -> types.SimpleNamespace:
+        # Plenty of free system memory: the old sizing saw no pressure at all.
+        return types.SimpleNamespace(available=16 * GIB, total=16 * GIB)
+
+    def memory_info(self) -> types.SimpleNamespace:
+        return types.SimpleNamespace(rss=self.rss)
+
+
+class Embedder:
+    """Qwen-shaped: 16 heads, word count stands in for tokens, MPS can OOM."""
+
+    config = types.SimpleNamespace(num_attention_heads=16)
+
+    def __init__(self, device: str) -> None:
+        self.device = device
+        self.max_seq_length = 2048
+        self.calls: list[tuple[str, int, int]] = []
+        self.oom_on_mps = True
+
+    def tokenizer(self, texts, truncation, max_length):
+        return {"input_ids": [[0] * min(len(text.split()), max_length) for text in texts]}
+
+    def encode(self, texts, batch_size=32, show_progress_bar=False, **kwargs):
+        longest = max(len(text.split()) for text in texts)
+        self.calls.append((self.device, len(texts), longest))
+        if self.device == "mps" and self.oom_on_mps:
+            raise RuntimeError("MPS backend out of memory (MPS allocated: 3.9 GiB, max allowed: 4.0 GiB)")
+        return np.array([[len(text.split()), index] for index, text in enumerate(texts)], dtype=np.float32)
+
+    def to(self, device: str) -> "Embedder":
+        self.device = device
+        return self
+
+
+@pytest.fixture
+def rss_budget(monkeypatch):
+    sensor = FakeRSS(RESIDENT)
+    monkeypatch.setitem(sys.modules, "psutil", sensor)
+    monkeypatch.setattr(memory_budget, "process_rss_limit", lambda: 4 * GIB, raising=False)
+    return sensor
+
+
+def _cpu_fit() -> int:
+    return int((4 * GIB - RESIDENT) * 0.8) // (16 * 8)
+
+
+def test_cpu_budget_is_the_rss_room_not_free_system_memory(rss_budget):
+    model = Embedder("cpu")
+    fit = accelerator_token_budget(model, EMBED_TOKEN_BUDGET)
+    assert fit == _cpu_fit() == 4_593_545
+    # The soak's first CPU piece was 11 texts of 1,718 tokens: 4.16 GB of attention.
+    assert budget_batch_size(1718, 32, EMBED_TOKEN_BUDGET) == 11
+    assert budget_batch_size(1718, 32, fit) == 1
+    assert 1 * 1718 ** 2 * 16 * 8 < (4 * GIB - RESIDENT)
+
+
+def test_cpu_budget_floors_at_one_text_and_keeps_room_when_plentiful(rss_budget, monkeypatch):
+    model = Embedder("cpu")
+    rss_budget.rss = 5 * GIB
+    assert accelerator_token_budget(model, EMBED_TOKEN_BUDGET) == 1
+    rss_budget.rss = 1 * GIB
+    monkeypatch.setattr(memory_budget, "process_rss_limit", lambda: 8 * GIB)
+    assert accelerator_token_budget(model, EMBED_TOKEN_BUDGET) == EMBED_TOKEN_BUDGET
+
+
+def test_models_without_known_heads_keep_their_budget(rss_budget):
+    model = types.SimpleNamespace(device="cpu")
+    assert accelerator_token_budget(model, EMBED_TOKEN_BUDGET) == EMBED_TOKEN_BUDGET
+
+
+def test_server_retries_the_oom_piece_and_the_rest_in_rss_sized_pieces(rss_budget, monkeypatch):
+    monkeypatch.delenv("TRUEMEMORY_DEVICE", raising=False)
+    monkeypatch.setattr(mps_utils, "auto_detect_device", lambda: "mps")
+    monkeypatch.setattr(mps_utils, "ensure_mps_memory_budget", lambda device: None)
+    # No configured MPS cap: the MPS pieces keep the plain 32 x 1,024^2 budget.
+    monkeypatch.setattr(mps_utils, "get_mps_memory_budget", lambda: None)
+    server = ms.ModelServer()
+    server._write_status_file = lambda: None
+    server._embed_device = lambda: "mps"
+    model = Embedder("mps")
+    monkeypatch.setattr(server, "_build_embed_model", lambda model_id, device: model)
+    texts = [" ".join(["memory"] * 1718) for _ in range(22)] + [" ".join(["short"] * 20) for _ in range(10)]
+    try:
+        response = server.handle_request({"op": "embed", "texts": texts, "tier": "qwen3_256"})
+    finally:
+        server._workers.shutdown(wait=False)
+    assert response["ok"] is True
+    # Row order is preserved: each vector carries its own text's length.
+    assert [int(row[0]) for row in response["vectors"]] == [1718] * 22 + [20] * 10
+    mps_calls = [call for call in model.calls if call[0] == "mps"]
+    cpu_calls = [call for call in model.calls if call[0] == "cpu"]
+    assert mps_calls == [("mps", 11, 1718)], "The MPS piece is sized by the MPS budget"
+    fit = _cpu_fit()
+    assert cpu_calls and all(size * longest ** 2 <= fit or size == 1 for _, size, longest in cpu_calls)
+    assert all(size == 1 for _, size, longest in cpu_calls if longest == 1718)
+    assert sum(size for _, size, _ in cpu_calls) == 32
diff --git a/truememory/_rerank_budget.py b/truememory/_rerank_budget.py
index 30574007..21b1a5de 100644
--- a/truememory/_rerank_budget.py
+++ b/truememory/_rerank_budget.py
@@ -74,14 +74,43 @@ def _attention_heads(model: object) -> int | None:
     return None
 
 
+# Peak CPU bytes per batch x tokens^2 x attention head: eager attention keeps
+# two float32 score tensors alive at its softmax (11 x 1,718 tokens x 16 heads
+# is 4.16 GB per layer).
+_CPU_BYTES_PER_SCORE = 8
+
+
+def _cpu_token_budget(model: object, budget: int) -> int:
+    """Lower *budget* so one CPU batch's attention fits the process RSS budget.
+
+    CPU attention lands in resident memory. Pieces are sized against the RSS
+    budget minus what is already resident, not against free system memory.
+    """
+    heads = _attention_heads(model)
+    if heads is None:
+        return budget
+    try:
+        import psutil
+        from truememory.memory_budget import process_rss_limit
+        room = process_rss_limit() - psutil.Process().memory_info().rss
+    except Exception:
+        return budget
+    fit = int(max(0, room) * _MPS_HEADROOM_SHARE) // (heads * _CPU_BYTES_PER_SCORE)
+    return max(1, min(budget, fit))
+
+
 def accelerator_token_budget(model: object, budget: int) -> int:
-    """Lower *budget* so one batch's attention fits the free MPS allocation.
+    """Lower *budget* so one batch's attention fits where the model runs.
 
-    MPS allocations are capped per process. A batch above the remaining cap
-    fails with an OOM that degrades the model to CPU for the server lifetime.
+    MPS allocations are capped per process: a batch above the remaining cap
+    fails with an OOM that degrades the model to CPU. On CPU the attention is
+    resident memory, so a batch must fit the process RSS budget.
     """
     device = getattr(model, "device", None)
-    if getattr(device, "type", device) != "mps":
+    kind = getattr(device, "type", device)
+    if kind == "cpu":
+        return _cpu_token_budget(model, budget)
+    if kind != "mps":
         return budget
     from truememory.mps_utils import get_mps_memory_budget
     policy = get_mps_memory_budget()
diff --git a/truememory/memory_budget.py b/truememory/memory_budget.py
index afca47d3..05a60d6e 100644
--- a/truememory/memory_budget.py
+++ b/truememory/memory_budget.py
@@ -414,6 +414,11 @@ def _record(self, token: int) -> _OwnerRecord:
             raise StaleBudgetOwner("Owner has retired")
         return record
 
+    def ceiling(self) -> int:
+        """Tightest enrolled owner limit in bytes, or 0 when none is enrolled."""
+        with self._lock:
+            return self._ceiling()
+
     def _ceiling(self) -> int:
         ceiling = 0
         for record in self._owners.values():
@@ -708,6 +713,12 @@ def resolve_allocation_plan(key: ProfileKey) -> AllocationPlan:
     return (registry or get_process_registry()).resolve_profile(key)
 
 
+def process_rss_limit() -> int:
+    """The process RSS ceiling: the tightest enrolled owner, else the automatic budget."""
+    ceiling = get_process_registry().ceiling()
+    return ceiling if ceiling else automatic_budget_bytes()
+
+
 def automatic_budget_bytes() -> int:
     try:
         import psutil
diff --git a/truememory/model_server.py b/truememory/model_server.py
index 6eb78a33..9d87dc53 100644
--- a/truememory/model_server.py
+++ b/truememory/model_server.py
@@ -117,6 +117,40 @@ class _FastLaneNotReady(_ProcessMemoryRefused):
     fast_lane_reloading = True
 
 
+def _run_in_pieces(run: Callable[[list, int], object], items: list, lengths: list[int] | None,
+                   budget: int) -> object:
+    """Run *items* in longest-first pieces whose batch x tokens^2 fits *budget*.
+
+    Results come back in item order. Without lengths, one call runs them all.
+    """
+    from truememory._rerank_budget import budgeted_slices
+    slices = None if lengths is None else budgeted_slices(lengths, max(1, len(items)), budget)
+    if slices is None:
+        return run(items, max(1, len(items)))
+    results: list = [None] * len(items)
+    for group in slices:
+        values = run([items[index] for index in group], len(group))
+        for position, index in enumerate(group):
+            results[index] = values[position]
+    return np.stack(results)
+
+
+def _reslice_remaining(order: np.ndarray | None, lengths: list[int], start: int, total: int, limit: int,
+                       budget: int) -> tuple[np.ndarray, dict[int, int]]:
+    """Order and piece sizes for positions from *start* on, under a new budget."""
+    from truememory._rerank_budget import budgeted_slices
+    done = [int(index) for index in order[:start]] if order is not None else list(range(start))
+    remaining = [int(index) for index in order[start:]] if order is not None else list(range(start, total))
+    slices = budgeted_slices([lengths[index] for index in remaining], limit, budget)
+    groups = ([[remaining[position] for position in group] for group in slices] if slices is not None
+              else [remaining[offset:offset + limit] for offset in range(0, len(remaining), limit)])
+    steps, position = {}, start
+    for group in groups:
+        steps[position] = len(group)
+        position += len(group)
+    return np.array(done + [index for group in groups for index in group], dtype=np.intp), steps
+
+
 def _array_metadata(shape: tuple, encoded: str = "") -> dict:
     return {"__ndarray__": encoded, "dtype": "float32", "shape": list(shape)}
 
@@ -1908,6 +1942,7 @@ def _handle_request_inner(self, request: dict) -> dict:
                     offset = 0
                     order = None
                     steps = None
+                    lengths = None
                     while offset < len(texts) or vectors is None:
                         retry = None
                         with deadline.locked(self._lock):
@@ -1977,7 +2012,20 @@ def _handle_request_inner(self, request: dict) -> dict:
                             self._check_process_memory("embed CPU retry", deadline)
                             deadline.check()
                             log.warning("MPS OOM during encoding; retrying microbatch on CPU")
-                            values = model.encode(retry, batch_size=len(retry) or limit, show_progress_bar=False)
+                            if lengths is None:
+                                values = model.encode(retry, batch_size=len(retry) or limit, show_progress_bar=False)
+                            else:
+                                # MPS-sized pieces do not fit CPU memory: size the retry and
+                                # every remaining piece against the RSS budget instead.
+                                from truememory._embed_budget import embed_token_budget
+                                from truememory._rerank_budget import accelerator_token_budget
+                                budget = accelerator_token_budget(model, embed_token_budget())
+                                retry_lengths = ([lengths[int(index)] for index in indices] if indices is not None
+                                                 else lengths[offset:offset + len(retry)])
+                                values = _run_in_pieces(
+                                    lambda items, size: model.encode(items, batch_size=size, show_progress_bar=False),
+                                    retry, retry_lengths, budget,
+                                )
                             if target is not None or expected_target is not None:
                                 check_target_vectors(target or expected_target, values, len(retry))
                             vectors = _store_batch_result(
@@ -1985,6 +2033,8 @@ def _handle_request_inner(self, request: dict) -> dict:
                                 **result_options,
                             )
                             offset += len(retry)
+                            if lengths is not None and offset < len(texts):
+                                order, steps = _reslice_remaining(order, lengths, offset, len(texts), limit, budget)
                     self._after_request_batches(throttler, len(texts), encode_start, deadline)
 
                     with deadline.locked(self._lock):
@@ -2022,6 +2072,7 @@ def _handle_request_inner(self, request: dict) -> dict:
                     offset = 0
                     order = None
                     steps = None
+                    lengths = None
                     recovery_name = model_name
                     while offset < len(pairs) or scores is None:
                         retry = None
@@ -2108,7 +2159,18 @@ def _handle_request_inner(self, request: dict) -> dict:
                         if retry is not None:
                             self._check_process_memory("reranker CPU retry", deadline)
                             deadline.check()
-                            values = reranker.predict(retry, batch_size=len(retry) or limit, show_progress_bar=False)
+                            if lengths is None:
+                                values = reranker.predict(retry, batch_size=len(retry) or limit, show_progress_bar=False)
+                            else:
+                                # Size the retry and remaining pieces for the CPU replacement.
+                                from truememory._rerank_budget import accelerator_token_budget, rerank_token_budget
+                                budget = accelerator_token_budget(reranker, rerank_token_budget())
+                                retry_lengths = ([lengths[int(index)] for index in indices] if indices is not None
+                                                 else lengths[offset:offset + len(retry)])
+                                values = _run_in_pieces(
+                                    lambda items, size: reranker.predict(items, batch_size=size, show_progress_bar=False),
+                                    retry, retry_lengths, budget,
+                                )
                             if indices is None:
                                 scores = _store_batch_result(scores, values, offset, len(retry), len(pairs),
                                                              "scores", **result_options)
@@ -2118,6 +2180,8 @@ def _handle_request_inner(self, request: dict) -> dict:
                                     **result_options,
                                 )
                             offset += len(retry)
+                            if lengths is not None and offset < len(pairs):
+                                order, steps = _reslice_remaining(order, lengths, offset, len(pairs), limit, budget)
                     self._after_request_batches(throttler, len(pairs), predict_start, deadline)
                     return ({"ok": True, "scores": scores} if expected_name is None else
                             {"ok": True, "scores": scores, "protocol": PROTOCOL_VERSION,

From 8e1f248898de4598f2784b61963f4ebff441ea32 Mon Sep 17 00:00:00 2001
From: buildingjoshbetter
 <236675935+buildingjoshbetter@users.noreply.github.com>
Date: Mon, 26 Oct 2026 01:51:39 -0500
Subject: [PATCH 6/6] test: Give the CPU token budget test ample RSS room

test_cpu_models_ignore_the_mps_allocation checks that a CPU model's attention budget ignores the MPS allocation. The CPU budget is sized against the process RSS limit minus resident memory, so on a hosted runner after thousands of tests it came out lower (20,661,350 vs 33,554,432). The test now pins process_rss_limit high; the assertion is unchanged.
---
 tests/test-embed-token-budget.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/tests/test-embed-token-budget.py b/tests/test-embed-token-budget.py
index c4f34e3e..974c74bb 100644
--- a/tests/test-embed-token-budget.py
+++ b/tests/test-embed-token-budget.py
@@ -186,6 +186,10 @@ def test_cpu_models_ignore_the_mps_allocation(monkeypatch):
     model = FakeEmbedder(EMBED_MAX_TOKENS)
     _on_mps(monkeypatch, model, allocated_mib=1500)
     model.device = "cpu"
+    # CPU pieces are sized against the process RSS budget; give it ample room
+    # so the result does not depend on how much the test process holds.
+    from truememory import memory_budget
+    monkeypatch.setattr(memory_budget, "process_rss_limit", lambda: 1 << 40)
     assert accelerator_token_budget(model, EMBED_TOKEN_BUDGET) == EMBED_TOKEN_BUDGET