From fcde4bda49f83b23b9dbaa79e59816860caff151 Mon Sep 17 00:00:00 2001
From: Daigo Kimura <84138151+imD-5@users.noreply.github.com>
Date: Fri, 23 Oct 2026 22:02:56 +0900
Subject: [PATCH] fix(server): a re-sent chat prompt encodes only the pieces it
 has not sent before

A chat client re-sends its whole conversation every turn, and the CUDA server
tokenized all of it each time. For a long agent conversation that is millions
of small native allocations a request, made on the request's own thread; glibc
keeps much of that memory in the thread's arena after it is freed, so rank 0
grows by GiBs over a session until the node runs out of memory.

The tokenizer splits its input at the added tokens it matches before
normalization and encodes each part on its own. PromptPieces cuts the rendered
prompt after each such token, encodes only the pieces it has not seen and
reuses the ids of the rest, within a token budget. The ids are the tokenizer's
own: a tokenizer whose added tokens could move a split (rstrip, single_word,
special tokens encoded as text) keeps the full encode.
---
 src/tensorfold/cuda/server.py          |  13 ++-
 src/tensorfold/server/prompt_pieces.py |  72 +++++++++++++++
 tests/test_prompt_pieces.py            | 121 +++++++++++++++++++++++++
 3 files changed, 205 insertions(+), 1 deletion(-)
 create mode 100644 src/tensorfold/server/prompt_pieces.py
 create mode 100644 tests/test_prompt_pieces.py

diff --git a/src/tensorfold/cuda/server.py b/src/tensorfold/cuda/server.py
index 43cbc02a69..b66c853fb5 100644
--- a/src/tensorfold/cuda/server.py
+++ b/src/tensorfold/cuda/server.py
@@ -18,6 +18,7 @@
 from tensorfold.server.messages import validate_modalities
 from tensorfold.server.plain_text import PlainText
 from tensorfold.server.probabilities import TokenBytes, probability_options
+from tensorfold.server.prompt_pieces import PromptPieces
 from tensorfold.server.request_options import heard_effort, parse_numbers, thinking_fields
 from tensorfold.server.stopping import matched_stop, stop_options
 from tensorfold.server.thinking_notes import unanswered
@@ -119,6 +120,15 @@ def _plain_text(self) -> PlainText | None:
                     self.tok, assistant_keep=self.plain_assistant_keep)
             return made
 
+    def _prompt_pieces(self) -> PromptPieces | None:
+        """The chat prompts' encoded pieces, made once; None for a tokenizer they cannot follow."""
+
+        with _MADE:
+            made = self.__dict__.get("_prompt_pieces_made")
+            if made is None:
+                made = self.__dict__["_prompt_pieces_made"] = PromptPieces.of(self.tok)
+            return made
+
     @property
     def model_ids(self) -> list[str]:
         """The ids this endpoint answers to, as the MLX server lists them: ``--name`` first, then each ``--alias``."""
@@ -322,7 +332,8 @@ def render(messages: list[dict[str, Any]], **images: bool) -> str:   # text rend
                 prompt = plain.encode(text, shield)
             else:
                 text = render(body["messages"])
-                prompt = self._encode(text, False)
+                pieces = self._prompt_pieces()
+                prompt = self._encode(text, False) if pieces is None else pieces.ids(text)
         elif isinstance(body.get("prompt"), list):       # token ids (vLLM's and OpenAI's form): served as given
             prompt = self.token_ids(body["prompt"])
         else:
diff --git a/src/tensorfold/server/prompt_pieces.py b/src/tensorfold/server/prompt_pieces.py
new file mode 100644
index 0000000000..79f00a611d
--- /dev/null
+++ b/src/tensorfold/server/prompt_pieces.py
@@ -0,0 +1,72 @@
+"""Token ids of re-sent chat prompts: each piece between added tokens is encoded once, then reused.
+
+A tokenizer splits its input at the added tokens it matches before normalization and encodes each part on its own, so
+a piece that ends in one encodes alone to the ids it gets inside the whole text. A chat client re-sends its whole
+conversation every turn; only the pieces it has not sent before need the tokenizer.
+"""
+
+from __future__ import annotations
+
+import re
+import threading
+from array import array
+from collections import OrderedDict
+from typing import Any
+
+BUDGET = 1 << 21        # tokens kept across all pieces (4 bytes each, plus the pieces' text)
+
+
+class PromptPieces:
+    """``ids(text)`` equals ``tok.encode(text, add_special_tokens=False).ids``."""
+
+    def __init__(self, tok: Any, cut: re.Pattern[str], budget: int = BUDGET) -> None:
+        self._tok, self._cut, self._budget = tok, cut, budget
+        self._kept: OrderedDict[str, array] = OrderedDict()      # least recently used first
+        self._held = 0
+        self._lock = threading.Lock()
+
+    @classmethod
+    def of(cls, tok: Any, budget: int = BUDGET) -> PromptPieces | None:
+        """Pieces for ``tok``, or None where a cut could fall where the tokenizer does not split."""
+
+        if not all(callable(getattr(tok, name, None)) for name in ("encode_batch", "get_added_tokens_decoder")):
+            return None
+        if getattr(tok, "encode_special_tokens", False):        # special tokens are encoded as text, not split out
+            return None
+        first = [token for token in tok.get_added_tokens_decoder().values() if not token.normalized]
+        # rstrip takes the spaces after a token into its match, and single_word skips a match inside a word
+        if not first or any(token.rstrip or token.single_word for token in first):
+            return None
+        spellings = sorted({token.content for token in first if token.content}, key=len, reverse=True)
+        return cls(tok, re.compile("|".join(map(re.escape, spellings))), budget)     # the longest at the leftmost
+
+    def ids(self, text: str) -> list[int]:
+        pieces, at = [], 0
+        for match in self._cut.finditer(text):
+            pieces.append(text[at:match.end()])
+            at = match.end()
+        if at < len(text):
+            pieces.append(text[at:])
+        with self._lock:
+            found = {piece: self._kept[piece] for piece in pieces if piece in self._kept}
+            for piece in found:
+                self._kept.move_to_end(piece)
+        missing = [piece for piece in dict.fromkeys(pieces) if piece not in found]
+        if missing:                       # encode_batch releases the GIL while it encodes
+            encoded = self._tok.encode_batch(missing, add_special_tokens=False)
+            new = {piece: array("I", encoding.ids) for piece, encoding in zip(missing, encoded)}
+            found.update(new)
+            self._keep(new)
+        ids: list[int] = []
+        for piece in pieces:
+            ids.extend(found[piece])
+        return ids
+
+    def _keep(self, new: dict[str, array]) -> None:
+        with self._lock:
+            for piece, ids in new.items():
+                if piece not in self._kept and len(ids) <= self._budget:
+                    self._kept[piece] = ids
+                    self._held += len(ids)
+            while self._held > self._budget:
+                self._held -= len(self._kept.popitem(last=False)[1])
diff --git a/tests/test_prompt_pieces.py b/tests/test_prompt_pieces.py
new file mode 100644
index 0000000000..819bd6f761
--- /dev/null
+++ b/tests/test_prompt_pieces.py
@@ -0,0 +1,121 @@
+"""PromptPieces: a chat prompt's ids equal the tokenizer's, and a re-sent conversation encodes only its new pieces."""
+
+import json
+import random
+from types import SimpleNamespace
+
+import pytest
+
+tokenizers = pytest.importorskip("tokenizers")
+from tokenizers import AddedToken, Regex, Tokenizer, models, normalizers, pre_tokenizers, trainers
+
+from tensorfold.server.prompt_pieces import PromptPieces
+
+TURN, END, TOOL, CALL, MARK = "<｜User｜>", "<｜end▁of▁sentence｜>", "<tool>", "<tool><call>", "<x>"
+WORDS = ["a", "b", " ", "  ", "\n", "\n\n", "\t", "1", "12", "1234", "日本", "語", "🙂", ".", "'s", "def", " x",
+         "(", ")", "<", ">", "｜", "<｜", "Call", "call", "tool", "end"]
+
+
+def tokenizer(**mark) -> Tokenizer:
+    """A byte-level BPE like DeepSeek's, with a normalizer and added tokens; ``mark`` sets the flags of ``MARK``."""
+
+    tok = Tokenizer(models.BPE())
+    tok.normalizer = normalizers.Lowercase()
+    tok.pre_tokenizer = pre_tokenizers.Sequence([pre_tokenizers.Split(Regex(r"\p{N}{1,3}"), "isolated"),
+                                                 pre_tokenizers.ByteLevel(add_prefix_space=False)])
+    trainer = trainers.BpeTrainer(vocab_size=300, initial_alphabet=pre_tokenizers.ByteLevel.alphabet(),
+                                  show_progress=False)
+    tok.train_from_iterator(["def call(x): return x  # 日本語 🙂 tool end 1234\n"] * 50, trainer)
+    tok.add_special_tokens([TURN, END])
+    tok.add_tokens([AddedToken(TOOL, normalized=False), AddedToken(CALL, normalized=False),
+                    AddedToken("call", normalized=True), AddedToken(MARK, normalized=False, **mark)])
+    return tok
+
+
+class Counting:
+    """The tokenizer, recording each text it is asked to encode."""
+
+    def __init__(self, tok):
+        self.tok, self.encoded = tok, []
+
+    def get_added_tokens_decoder(self):
+        return self.tok.get_added_tokens_decoder()
+
+    def encode_batch(self, texts, add_special_tokens=False):
+        self.encoded += texts
+        return self.tok.encode_batch(texts, add_special_tokens=add_special_tokens)
+
+
+def full(tok, text):
+    return tok.encode(text, add_special_tokens=False).ids
+
+
+def pieces_of(tok, **kwargs) -> PromptPieces:
+    pieces = PromptPieces.of(tok, **kwargs)
+    assert pieces is not None
+    return pieces
+
+
+@pytest.mark.parametrize("mark", [{}, {"lstrip": True}])
+def test_ids_equal_the_tokenizers_on_texts_with_added_tokens_anywhere(mark):
+    tok, rng = tokenizer(**mark), random.Random(0)
+    pieces = pieces_of(tok, budget=200)       # small: pieces are dropped and encoded again along the way
+    marks = [TURN, END, TOOL, CALL, MARK, " " + MARK, TOOL + "<call", TOOL + CALL]
+    texts = ["", TURN, END + END, " " + TURN + " ", "a" * 3000 + TURN + "b"]
+    texts += ["".join(rng.choice(marks) if rng.random() < 0.2 else rng.choice(WORDS)
+                      for _ in range(rng.randint(1, 50))) for _ in range(1500)]
+    for text in texts:
+        assert pieces.ids(text) == full(tok, text), text
+
+
+def test_a_resent_conversation_encodes_only_its_new_pieces():
+    tok = Counting(tokenizer())
+    pieces = pieces_of(tok)
+    turn1 = TURN + "def call(x):\n  return 12 * x\n" * 200 + END
+    turn2 = TURN + "and 日本語? 🙂" + END
+    assert pieces.ids(turn1) == full(tok.tok, turn1)
+    tok.encoded.clear()
+    assert pieces.ids(turn1 + turn2 + "tail") == full(tok.tok, turn1 + turn2 + "tail")
+    assert tok.encoded == ["and 日本語? 🙂" + END, "tail"]       # TURN alone is a piece seen before
+
+
+@pytest.mark.parametrize("mark", [{"rstrip": True}, {"single_word": True}])
+def test_a_tokenizer_whose_added_tokens_move_a_split_gets_no_pieces(mark):
+    assert PromptPieces.of(tokenizer(**mark)) is None
+
+
+def test_special_tokens_encoded_as_text_get_no_pieces():
+    tok = tokenizer()
+    tok.encode_special_tokens = True
+    assert PromptPieces.of(tok) is None
+
+
+def test_a_tokenizer_without_added_tokens_or_batches_gets_no_pieces():
+    assert PromptPieces.of(Tokenizer(models.BPE())) is None
+    assert PromptPieces.of(object()) is None
+
+
+def test_the_pieces_kept_stay_within_the_budget():
+    tok = tokenizer()
+    pieces = pieces_of(tok, budget=50)
+    for i in range(40):
+        text = f"{TURN}turn {i} " + "x" * i + END
+        assert pieces.ids(text) == full(tok, text)
+        assert pieces._held <= 50 and pieces._held == sum(len(ids) for ids in pieces._kept.values())
+
+
+def test_the_cuda_server_encodes_chat_prompts_through_its_pieces(tmp_path):
+    pytest.importorskip("jinja2")
+    from tensorfold.cuda.server import App
+
+    tokenizer().save(str(tmp_path / "tokenizer.json"))
+    template = "{% for m in messages %}" + TURN + "{{ m.content }}" + END + "{% endfor %}"
+    (tmp_path / "tokenizer_config.json").write_text(json.dumps({"chat_template": template}))
+    app = App(SimpleNamespace(eos=(0,)), tmp_path, "test")
+    messages = [{"role": "user", "content": "def call(x): return 12 * x"}]
+    for turn in range(3):
+        prompt = app.prepare({"messages": messages, "max_tokens": 1}, True).prompt
+        assert prompt == full(app.tok, app.template.render(messages, tools=None, enable_thinking=False))
+        messages += [{"role": "assistant", "content": f"answer {turn}"}, {"role": "user", "content": f"more {turn}"}]
+    pieces = app._prompt_pieces()
+    assert pieces is not None and "def call(x): return 12 * x" + END in pieces._kept