From 2eb4ac9317ead0c10c3a78da8a4034ea88382a92 Mon Sep 17 00:00:00 2001 From: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com> Date: Thu, 29 Oct 2026 19:33:08 -0700 Subject: [PATCH 1/8] [None][doc] Recommend a single serving frontend with the prefix-tokenization cache The cache is per frontend process, and multiple frontends spread a conversation's connections across processes, so most turns miss. Signed-off-by: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com> --- .../features/prefix-tokenization-cache.md | 24 +++++++++++++++++++ 1 file changed, 24 insertions(+) diff --git a/docs/source/features/prefix-tokenization-cache.md b/docs/source/features/prefix-tokenization-cache.md index 7a3065b52213..d9074a378603 100644 --- a/docs/source/features/prefix-tokenization-cache.md +++ b/docs/source/features/prefix-tokenization-cache.md @@ -77,6 +77,30 @@ feature you turned on. fallbacks, and its occupancy, so a cache that is enabled but never hitting is visible in the server log. +## Multiple serving frontends + +Use a single serving frontend (`num_serve_frontends: 1`, the default) when you +enable the cache for multi-turn traffic. + +The cache lives in the memory of one frontend process and is not shared +between processes. With `num_serve_frontends: K`, the K frontend processes +share the serving port through `SO_REUSEPORT`, and the kernel assigns each new +TCP connection to one of them without looking at the request. A client that +opens a new connection per request, as many benchmark clients do, therefore +sends each turn of a conversation to an arbitrary frontend. That frontend +usually has not seen the conversation's previous turn, so the prompt is +tokenized in full. + +On an agentic multi-turn trace with prompts of about 160k tokens on average, +at one session, the p90 time spent in the frontend was 453 ms with 8 +frontends and 166 ms with one. Decode throughput per user was the same up to +32 concurrent sessions, because at that concurrency the frontend is not the +bottleneck. + +Multiple frontends help when a single Python frontend cannot keep up, typically +at high request rates with short prompts. That is not the workload this cache +targets. + ## Measured effect An A/B on GLM-5.2 (GB300, disaggregated, matched pair, 3600 s, ~29.6k requests From 6aa531a8ba8d5012e4f4f7127660c44de30063f2 Mon Sep 17 00:00:00 2001 From: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com> Date: Thu, 29 Oct 2026 07:05:45 -0700 Subject: [PATCH 2/8] [None][perf] Splice diverging prompts in the prefix-tokenization cache Agentic clients rewrite the end of the history between turns (a trimmed tool result, a replaced last message), so the new prompt no longer extends the cached one and was tokenized in full. Record a checkpoint every 256 tokens of a cached prompt and splice a diverging prompt at the last checkpoint inside the longest common prefix, keeping the cached entry for a branch that may continue from it. Document the splice. Signed-off-by: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com> --- .../features/prefix-tokenization-cache.md | 24 +++- tensorrt_llm/inputs/prefix_token_cache.py | 133 +++++++++++++++--- .../inputs/test_prefix_token_cache.py | 23 ++- 3 files changed, 152 insertions(+), 28 deletions(-) diff --git a/docs/source/features/prefix-tokenization-cache.md b/docs/source/features/prefix-tokenization-cache.md index d9074a378603..0cae5f84805a 100644 --- a/docs/source/features/prefix-tokenization-cache.md +++ b/docs/source/features/prefix-tokenization-cache.md @@ -14,9 +14,18 @@ whole prompt. The new prompt may also rewrite the last tokens of the cached one (those after the seam described below), as a chat template does when it replaces the previous turn's generation prefix with the response. +Agentic clients also rewrite the end of the history between turns, for example +by trimming a tool result or replacing the last message, so the new prompt is +no longer an extension of the cached one. The cache records a checkpoint every +256 tokens of a cached prompt. When a new prompt diverges from a cached one, the +cache reuses the cached ids up to the last checkpoint that still lies a full +checkpoint interval inside their longest common prefix, and tokenizes only the +rest. + Tokenizing a tail in isolation is not generally identical to tokenizing the whole prompt, because BPE merges can straddle the split. The cache therefore -backs off a fixed number of tokens from the end of the cached prefix, +backs off a fixed number of tokens from the end of the cached prefix (for a +diverging prompt, at least one checkpoint interval before the divergence), re-tokenizes from there, and accepts the splice only if the re-tokenized ids equal the cached ids over that seam. Otherwise the prompt is tokenized in full. The output is always identical to tokenizing the whole prompt. @@ -62,7 +71,10 @@ feature you turned on. `add_special_tokens=False`, so they benefit. `/v1/completions` defaults to `add_special_tokens=True` and is not accelerated. - A prompt that reuses a cached entry replaces that entry, so a conversation - costs one entry regardless of how many turns it has. Lookup is bucketed by a + costs one entry regardless of how many turns it has. A prompt that diverges + from a cached entry before its seam is added next to it instead: the + divergence may be a branch, such as a subagent started from a shared history, + and the original conversation can still continue from the old entry. Lookup is bucketed by a hash of the first `TLLM_PREFIX_TOKEN_CACHE_MIN_CHARS` characters, and each entry keeps a short probe of the characters before its seam so that conversations sharing a long system preamble are told apart without @@ -73,7 +85,8 @@ feature you turned on. - An unexpected exception inside the cache disables it, with a warning, and the request is tokenized normally. A cache problem never fails a request. - The cache logs its configuration at startup and, every 1000 cache-eligible - requests, an `INFO` line with the hit rate, the number of re-synchronization + requests, an `INFO` line with the hit rate (and how many hits were on diverging + prompts), the number of re-synchronization fallbacks, and its occupancy, so a cache that is enabled but never hitting is visible in the server log. @@ -121,3 +134,8 @@ The splice was also checked against real `gpt2`, Llama-3, and Qwen3 tokenizers on 300 randomly grown multi-turn prompts each, with deltas deliberately starting mid-word: 300/300 identical ids, 280 cache hits, and 0 re-synchronization fallbacks per tokenizer. + +Splicing diverging prompts was checked on 300 multi-turn prompts whose history +is rewritten between turns: all 300 matched tokenizing the whole prompt, 99.3% +were cache hits, and tokenization took about a tenth of the time of +tokenizing every prompt in full. diff --git a/tensorrt_llm/inputs/prefix_token_cache.py b/tensorrt_llm/inputs/prefix_token_cache.py index fe5ebf128dfe..58b7fdcab6b9 100644 --- a/tensorrt_llm/inputs/prefix_token_cache.py +++ b/tensorrt_llm/inputs/prefix_token_cache.py @@ -10,11 +10,16 @@ of the cached prompt, so the new prompt may also rewrite that end, as a chat template does when the previous turn's generation prefix is replaced by the response. +Agentic clients also rewrite the end of the history between turns (a trimmed +tool result, a replaced last message), so a prompt that diverges from a cached +one before its re-tokenization point is spliced at the last checkpoint inside +their longest common prefix. Tokenizing a tail in isolation is not generally identical to tokenizing the whole prompt, because BPE merges can straddle the split. The cache therefore -backs off ``overlap`` tokens from the end of the cached prefix, re-tokenizes -from there, and accepts the splice only if the first ``resync`` re-tokenized ids +backs off ``overlap`` tokens from the end of the cached prefix (or, for a +diverging prompt, at least ``checkpoint_tokens`` tokens from the divergence), +re-tokenizes from there, and accepts the splice only if the first ``resync`` re-tokenized ids equal the cached ids over the same span. Otherwise the prompt is tokenized in full. @@ -26,12 +31,13 @@ from __future__ import annotations +import bisect import os import threading from array import array from collections import OrderedDict from dataclasses import dataclass -from typing import ClassVar, Mapping, Protocol, Sequence +from typing import ClassVar, Mapping, NamedTuple, Protocol, Sequence from ..logger import logger @@ -81,6 +87,10 @@ class PrefixTokenCacheConfig: overlap: int = 64 # Re-tokenized ids that must equal the cached ids for the splice to be used. resync: int = 32 + # Spacing, in tokens, of the split points a diverging prompt can be spliced + # at. A split is used only when the next checkpoint still lies inside the + # common prefix, so it must exceed ``resync``. + checkpoint_tokens: int = 256 ENV_VARS: ClassVar[dict[str, str]] = { "max_entries": "TLLM_PREFIX_TOKEN_CACHE_ENTRIES", @@ -115,6 +125,17 @@ class _Entry: split_token: int split_char: int probe: str # the PROBE_CHARS characters of ``text`` before ``split_char`` + # Start character of every ``checkpoint_tokens``-th token (0, C, 2C, ...), + # where a prompt that diverges from ``text`` can be re-tokenized from. + checkpoints: array # int32 + + +class _Match(NamedTuple): + eid: int + entry: _Entry + split_token: int + split_char: int + extends: bool # the prompt keeps ``entry.text[: entry.split_char]`` class PrefixTokenCache: @@ -143,6 +164,8 @@ def __init__(self, config: PrefixTokenCacheConfig) -> None: self.hits = 0 self.misses = 0 self.resync_failures = 0 + # Hits on a cached prompt the new one diverges from (also in ``hits``). + self.divergent_hits = 0 def encode(self, tokenizer: OffsetTokenizer, text: str) -> list[int]: """Return the token ids of ``text`` without special tokens.""" @@ -162,7 +185,7 @@ def summary(self) -> str: total = self.hits + self.misses rate = self.hits / total if total else 0.0 return ( - f"{self.hits}/{total} hits ({rate:.1%}), " + f"{self.hits}/{total} hits ({rate:.1%}, {self.divergent_hits} divergent), " f"{self.resync_failures} resync fallbacks, " f"{len(self._entries)} entries, {self._total_chars} cached chars" ) @@ -171,11 +194,11 @@ def _encode(self, tokenizer: OffsetTokenizer, text: str) -> list[int]: found = self._lookup(text) ids = None if found is not None: - eid, entry = found - tail_ids, tail_offsets = self._tokenize(tokenizer, text[entry.split_char :]) - if self._resynced(entry, tail_ids): - ids = entry.ids[: entry.split_token].tolist() + list(tail_ids) - base_token, base_char = entry.split_token, entry.split_char + entry = found.entry + tail_ids, tail_offsets = self._tokenize(tokenizer, text[found.split_char :]) + if self._resynced(entry, found.split_token, tail_ids): + ids = entry.ids[: found.split_token].tolist() + list(tail_ids) + base_token, base_char = found.split_token, found.split_char if ids is None: tail_ids, tail_offsets = self._tokenize(tokenizer, text) ids, base_token, base_char = tail_ids, 0, 0 @@ -192,10 +215,19 @@ def _encode(self, tokenizer: OffsetTokenizer, text: str) -> list[int]: self.resync_failures += 1 else: self.hits += 1 - if found is not None: - self._remove(found[0]) + self.divergent_hits += not found.extends + # A prompt that reuses an entry supersedes it. A diverging prompt + # may be a branch (a subagent off a shared history), so the entry + # it was spliced from stays for the other branch. + if found is not None and found.extends: + self._remove(found.eid) if split_token > 0 and tail_index >= 0: - self._insert(text, ids, split_token, base_char + tail_offsets[tail_index][0]) + checkpoints = self._checkpoints( + found, base_token, base_char, tail_offsets, len(ids) + ) + self._insert( + text, ids, split_token, base_char + tail_offsets[tail_index][0], checkpoints + ) report = (self.hits + self.misses) % LOG_EVERY_REQUESTS == 0 summary = self.summary() if report else None if summary is not None: @@ -209,41 +241,86 @@ def _tokenize( enc = tokenizer(text, add_special_tokens=False, return_offsets_mapping=True) return enc["input_ids"], enc["offset_mapping"] - def _resynced(self, entry: _Entry, tail_ids: Sequence[int]) -> bool: + def _checkpoints( + self, + found: _Match | None, + base_token: int, + base_char: int, + tail_offsets: Sequence[tuple[int, int]], + num_ids: int, + ) -> array: + """Checkpoints of a prompt whose ids are ``base_token`` cached ids plus a tail.""" + step = self._config.checkpoint_tokens + # Spliced ids before base_token are the cached entry's, at the same chars. + num_kept = -(-base_token // step) + out = array("i", found.entry.checkpoints[:num_kept] if base_token else ()) + for index in range(num_kept * step, num_ids, step): + out.append(base_char + tail_offsets[index - base_token][0]) + return out + + def _resynced(self, entry: _Entry, split_token: int, tail_ids: Sequence[int]) -> bool: """Whether the tokenizer re-synchronized with the cached ids at the seam.""" - cached_tail = entry.ids[entry.split_token :] + cached_tail = entry.ids[split_token:] span = min(self._config.resync, len(tail_ids), len(cached_tail)) return span > 0 and list(tail_ids[:span]) == cached_tail[:span].tolist() def _bucket_key(self, text: str) -> int: return hash(text[: self._config.min_chars]) - def _lookup(self, text: str) -> tuple[int, _Entry] | None: - """Cached entry with the longest re-tokenization prefix ``text`` keeps, promoted to MRU.""" + def _lookup(self, text: str) -> _Match | None: + """Best cached entry to splice ``text`` from, promoted to MRU. + + The entry with the longest re-tokenization prefix that ``text`` keeps + wins. Without one, the + entry sharing the longest common prefix with ``text`` is used, split at + its last checkpoint that leaves a full checkpoint interval of common + text after the split. + """ with self._lock: candidates = list(self._buckets.get(self._bucket_key(text), {}).items()) # Entries are immutable, so the comparisons run outside the lock. The # probe check rejects other conversations in O(PROBE_CHARS); only a true # prefix pays the full compare over the shared preamble. - best: tuple[int, _Entry] | None = None + best: _Match | None = None for eid, entry in candidates: n = entry.split_char if ( n <= len(text) - and (best is None or n > best[1].split_char) + and (best is None or n > best.split_char) and text[n - len(entry.probe) : n] == entry.probe and text.startswith(entry.text[:n]) ): - best = (eid, entry) + best = _Match(eid, entry, entry.split_token, entry.split_char, True) + if best is None: + for eid, entry in candidates: + match = self._divergent_match(eid, entry, text) + if match is not None and (best is None or match.split_token > best.split_token): + best = match if best is not None: with self._lock: - if best[0] in self._entries: - self._entries.move_to_end(best[0]) + if best.eid in self._entries: + self._entries.move_to_end(best.eid) # If it was evicted meanwhile the entry is still a valid prefix to # splice from; _remove() on it is a no-op. return best - def _insert(self, text: str, ids: Sequence[int], split_token: int, split_char: int) -> None: + def _divergent_match(self, eid: int, entry: _Entry, text: str) -> _Match | None: + common = _common_prefix_len(entry.text, text) + # The checkpoint after the split must start inside the common prefix. + index = bisect.bisect_right(entry.checkpoints, common) - 2 + if index < 1: + return None + split_token = index * self._config.checkpoint_tokens + return _Match(eid, entry, split_token, entry.checkpoints[index], False) + + def _insert( + self, + text: str, + ids: Sequence[int], + split_token: int, + split_char: int, + checkpoints: array, + ) -> None: """Caller holds the lock.""" key = self._bucket_key(text) bucket = self._buckets.setdefault(key, {}) @@ -251,7 +328,7 @@ def _insert(self, text: str, ids: Sequence[int], split_token: int, split_char: i return # concurrent misses on the same prompt eid, self._next_id = self._next_id, self._next_id + 1 probe = text[max(split_char - PROBE_CHARS, 0) : split_char] - entry = _Entry(text, array("i", ids), split_token, split_char, probe) + entry = _Entry(text, array("i", ids), split_token, split_char, probe, checkpoints) self._entries[eid] = entry bucket[eid] = entry self._total_chars += len(text) @@ -274,6 +351,16 @@ def _remove(self, eid: int) -> None: del self._buckets[key] +def _common_prefix_len(a: str, b: str) -> int: + """Length of the longest common prefix, compared in blocks.""" + n, i, block = min(len(a), len(b)), 0, 4096 + while i < n and a[i : i + block] == b[i : i + block]: + i += block + while i < n and a[i] == b[i]: + i += 1 + return min(i, n) + + def create_prefix_token_cache(tokenizer: object) -> PrefixTokenCache | None: """Return a cache for ``tokenizer`` if the tokenizer can support one. diff --git a/tests/unittest/inputs/test_prefix_token_cache.py b/tests/unittest/inputs/test_prefix_token_cache.py index d92cd4783458..c22490df39de 100644 --- a/tests/unittest/inputs/test_prefix_token_cache.py +++ b/tests/unittest/inputs/test_prefix_token_cache.py @@ -82,7 +82,7 @@ def __call__(self, text: str, **kwargs: Any) -> dict[str, Any]: def _cache(**overrides: int) -> PrefixTokenCache: - config = dict(max_entries=64, overlap=4, resync=2, min_chars=_MIN_CHARS) + config = dict(max_entries=64, overlap=4, resync=2, min_chars=_MIN_CHARS, checkpoint_tokens=8) return PrefixTokenCache(PrefixTokenCacheConfig(**{**config, **overrides})) @@ -129,6 +129,24 @@ def test_resync_failure_falls_back_to_full_tokenization() -> None: assert cache.resync_failures > 0 +@pytest.mark.parametrize("tokenizer", [_MergeTokenizer(), _UnsyncableTokenizer()]) +def test_rewritten_history_is_identical_to_full_tokenization(tokenizer: Any) -> None: + """A turn that rewrites the end of the history splices from the common prefix.""" + cache = _cache() + turns = _turns(8) + for i, prompt in enumerate(turns): + assert cache.encode(tokenizer, prompt) == tokenizer(prompt)["input_ids"] + # Trim the last message and append a different one, as agents do. + rewritten = prompt[: len(prompt) - 20] + f"rewritten tail {i} there. " + assert cache.encode(tokenizer, rewritten) == tokenizer(rewritten)["input_ids"] + if isinstance(tokenizer, _MergeTokenizer): + assert cache.divergent_hits >= len(turns) + # The rewritten branch does not evict the history it diverged from. + assert any(entry.text == turns[-1] for entry in cache._entries.values()) + else: + assert cache.divergent_hits == 0 and cache.resync_failures > 0 + + def test_short_prompts_are_tokenized_but_not_cached() -> None: tokenizer, cache = _MergeTokenizer(), _cache() prompt = "the cab" @@ -160,13 +178,14 @@ def test_conversations_sharing_a_long_preamble_find_their_own_entry() -> None: for prompt in turns.values(): cache.encode(tokenizer, prompt) assert len(cache._buckets) == 1 + hits_before = cache.hits for i, prompt in turns.items(): # The next turn keeps the previous prompt, generation suffix included. extended = prompt + f"reply to {i}.<|user|>\nmore from {i}.{suffix}" found = cache._lookup(extended) assert found is not None and found[1].text == prompt assert cache.encode(tokenizer, extended) == tokenizer(extended)["input_ids"] - assert cache.hits == 16 and cache.resync_failures == 0 + assert cache.hits - hits_before == 16 and cache.resync_failures == 0 def test_rewritten_end_of_cached_prompt_reuses_its_prefix() -> None: From 2965b503feac9677d9193af2ffcdc8cc48371190 Mon Sep 17 00:00:00 2001 From: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com> Date: Tue, 3 Nov 2026 05:07:59 -0800 Subject: [PATCH 3/8] [None][perf] Optionally tokenize each conversation in one serving frontend With num_serve_frontends > 1, SO_REUSEPORT hands each connection to any frontend, so consecutive turns of a conversation usually reach a frontend whose prefix-tokenization cache has not seen the previous turn. With TLLM_PREFIX_TOKEN_CACHE_ROUTE_BY_CONVERSATION=1, every frontend serves its cache on a Unix socket in the multi-frontend ipc dir, and a chat prompt whose conversation id hashes to another frontend is tokenized there. Only the prompt text and the token ids cross processes; a prompt without a conversation id, or whose owner does not answer, is tokenized locally. The switch is off by default, and without it the chat path is unchanged. Signed-off-by: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com> --- .../features/prefix-tokenization-cache.md | 14 ++ tensorrt_llm/commands/serve.py | 8 +- tensorrt_llm/inputs/prefix_token_cache.py | 37 +++- tensorrt_llm/inputs/registry.py | 10 + tensorrt_llm/serve/multi_frontend.py | 2 + tensorrt_llm/serve/openai_server.py | 39 +++- tensorrt_llm/serve/tokenization_owners.py | 200 ++++++++++++++++++ .../llmapi/test_serve_tokenization_owners.py | 108 ++++++++++ 8 files changed, 409 insertions(+), 9 deletions(-) create mode 100644 tensorrt_llm/serve/tokenization_owners.py create mode 100644 tests/unittest/llmapi/test_serve_tokenization_owners.py diff --git a/docs/source/features/prefix-tokenization-cache.md b/docs/source/features/prefix-tokenization-cache.md index 0cae5f84805a..d14f8217ed26 100644 --- a/docs/source/features/prefix-tokenization-cache.md +++ b/docs/source/features/prefix-tokenization-cache.md @@ -114,6 +114,20 @@ Multiple frontends help when a single Python frontend cannot keep up, typically at high request rates with short prompts. That is not the workload this cache targets. +If you need several frontends, set +`TLLM_PREFIX_TOKEN_CACHE_ROUTE_BY_CONVERSATION=1` to tokenize each +conversation in one frontend. A chat request's conversation id (the same id +attention DP routes by, from `conversation_params.conversation_id` or headers +such as `X-Correlation-ID`) is hashed to an owner frontend, and the frontend the +request arrived at sends the prompt text to the owner over a local Unix socket +and gets back the token ids; the request itself stays where it arrived. A +request without a conversation id, or whose owner does not answer, is +tokenized locally. The switch is off by default and has no effect with a +single frontend or with the cache disabled. On the trace above with 8 +frontends, it brought the p90 frontend time at one session from about 165 ms +to about 58 ms, against 44 ms with a single frontend (both measured with the +divergent splice). + ## Measured effect An A/B on GLM-5.2 (GB300, disaggregated, matched pair, 3600 s, ~29.6k requests diff --git a/tensorrt_llm/commands/serve.py b/tensorrt_llm/commands/serve.py index 2fcd075d71df..ca0095fef847 100644 --- a/tensorrt_llm/commands/serve.py +++ b/tensorrt_llm/commands/serve.py @@ -805,7 +805,9 @@ def launch_server( multi_frontend_serving = MultiFrontendServing( launcher_uds=attached_launcher_uds, is_launcher=False, - launcher_pid=os.getppid()) + launcher_pid=os.getppid(), + frontend_id=int(os.environ["TLLM_EXECUTOR_FRONTEND_ID"]), + num_frontends=multi_frontend.num_frontends) if enable_runtime_control_endpoints and multi_frontend.num_frontends > 1: raise click.BadParameter( "Runtime control endpoints require num_serve_frontends=1 because " @@ -906,7 +908,9 @@ def launch_server( stop_watching_frontends = _watch_attached_frontends( frontend_children) multi_frontend_serving = MultiFrontendServing( - launcher_uds=launcher_uds_path, is_launcher=True) + launcher_uds=launcher_uds_path, + is_launcher=True, + num_frontends=multi_frontend.num_frontends) sockets.append(launcher_uds_sock) server = OpenAIServer( diff --git a/tensorrt_llm/inputs/prefix_token_cache.py b/tensorrt_llm/inputs/prefix_token_cache.py index 58b7fdcab6b9..875e40d36757 100644 --- a/tensorrt_llm/inputs/prefix_token_cache.py +++ b/tensorrt_llm/inputs/prefix_token_cache.py @@ -32,12 +32,14 @@ from __future__ import annotations import bisect +import contextlib import os import threading from array import array from collections import OrderedDict +from contextvars import ContextVar from dataclasses import dataclass -from typing import ClassVar, Mapping, NamedTuple, Protocol, Sequence +from typing import Callable, ClassVar, Iterator, Mapping, NamedTuple, Protocol, Sequence from ..logger import logger @@ -45,6 +47,7 @@ "PrefixTokenCache", "PrefixTokenCacheConfig", "create_prefix_token_cache", + "tokenization_routing_key", ] # Log the hit rate every this many cache-eligible requests, so an operator can @@ -57,6 +60,23 @@ # same position. PROBE_CHARS = 256 +# The conversation the prompt being tokenized belongs to, if the caller knows it. +_ROUTING_KEY: ContextVar[str | None] = ContextVar("tokenization_routing_key", default=None) + + +@contextlib.contextmanager +def tokenization_routing_key(key: str | None) -> Iterator[None]: + """Tag the prompts tokenized in this block with conversation ``key``. + + A cache with a ``route`` hands such prompts to the process that caches that + conversation. + """ + token = _ROUTING_KEY.set(key) + try: + yield + finally: + _ROUTING_KEY.reset(token) + class OffsetTokenizer(Protocol): """The subset of the HF fast-tokenizer call interface the cache relies on.""" @@ -166,11 +186,24 @@ def __init__(self, config: PrefixTokenCacheConfig) -> None: self.resync_failures = 0 # Hits on a cached prompt the new one diverges from (also in ``hits``). self.divergent_hits = 0 + # Tokenizes a prompt tagged with a routing key in the process that + # owns that conversation, or returns None to tokenize it here. Set by + # multi-frontend serving. + self.route: Callable[[str, str], list[int] | None] | None = None + # Prompts another process tokenized through ``route``. + self.routed = 0 def encode(self, tokenizer: OffsetTokenizer, text: str) -> list[int]: """Return the token ids of ``text`` without special tokens.""" if self.disabled or len(text) < self._config.min_chars: return self._tokenize(tokenizer, text)[0] + key = _ROUTING_KEY.get() + if self.route is not None and key is not None: + ids = self.route(key, text) + if ids is not None: + with self._lock: + self.routed += 1 + return ids try: return self._encode(tokenizer, text) except Exception as e: @@ -186,7 +219,7 @@ def summary(self) -> str: rate = self.hits / total if total else 0.0 return ( f"{self.hits}/{total} hits ({rate:.1%}, {self.divergent_hits} divergent), " - f"{self.resync_failures} resync fallbacks, " + f"{self.resync_failures} resync fallbacks, {self.routed} routed, " f"{len(self._entries)} entries, {self._total_chars} cached chars" ) diff --git a/tensorrt_llm/inputs/registry.py b/tensorrt_llm/inputs/registry.py index 71cfa6cefd23..4c88f2006335 100644 --- a/tensorrt_llm/inputs/registry.py +++ b/tensorrt_llm/inputs/registry.py @@ -136,6 +136,7 @@ def __init__(self, self.config = config self.model_path = model_path self.multimodal_hashing_supported = None + self._prefix_token_cache_tokenizer = tokenizer # None when disabled or the tokenizer cannot support the cache. self._prefix_token_cache = (create_prefix_token_cache(tokenizer) if enable_tokenization_cache else None) @@ -153,6 +154,15 @@ def __call__( return token_ids, None +def prefix_token_cache_of( + input_processor: Any) -> Optional[Tuple[PrefixTokenCache, Any]]: + """The prefix-tokenization cache of ``input_processor`` and its tokenizer, if enabled.""" + cache = getattr(input_processor, "_prefix_token_cache", None) + if cache is None: + return None + return cache, input_processor._prefix_token_cache_tokenizer + + class BaseMultimodalInputProcessor(ABC): """ Base class for multimodal input processors with default implementations diff --git a/tensorrt_llm/serve/multi_frontend.py b/tensorrt_llm/serve/multi_frontend.py index ea7aa6bd4cea..bf63266c744c 100644 --- a/tensorrt_llm/serve/multi_frontend.py +++ b/tensorrt_llm/serve/multi_frontend.py @@ -85,6 +85,8 @@ class MultiFrontendServing: launcher_uds: str is_launcher: bool launcher_pid: Optional[int] = None # attached frontends: their parent + frontend_id: int = 0 # the launcher is frontend 0 + num_frontends: int = 1 @property def is_attached(self) -> bool: diff --git a/tensorrt_llm/serve/openai_server.py b/tensorrt_llm/serve/openai_server.py index 233512d9788b..c40f43f9bade 100644 --- a/tensorrt_llm/serve/openai_server.py +++ b/tensorrt_llm/serve/openai_server.py @@ -51,7 +51,9 @@ from tensorrt_llm.inputs.media_io import (BaseMediaIO, log_allowed_local_media_path) from tensorrt_llm.inputs.multimodal import MultimodalServerConfig -from tensorrt_llm.inputs.registry import BaseMultimodalInputProcessor +from tensorrt_llm.inputs.prefix_token_cache import tokenization_routing_key +from tensorrt_llm.inputs.registry import (BaseMultimodalInputProcessor, + prefix_token_cache_of) from tensorrt_llm.inputs.utils import ConversationMessage from tensorrt_llm.llmapi import DisaggregatedParams as LlmDisaggregatedParams from tensorrt_llm.llmapi import MultimodalEncoder, SchedulingParams, tracing @@ -86,7 +88,8 @@ from tensorrt_llm.serve.chat_utils import (load_chat_template, parse_chat_messages_coroutines) from tensorrt_llm.serve.cluster_storage import create_cluster_storage_client -from tensorrt_llm.serve.conversation_id import resolve_request_conversation_id +from tensorrt_llm.serve.conversation_id import (get_request_routing_id, + resolve_request_conversation_id) from tensorrt_llm.serve.disagg_auth import ( request_requires_internal_disagg_auth, validate_internal_disagg_request, validate_subagent_affinity) @@ -145,6 +148,8 @@ from tensorrt_llm.serve.runtime_control_auth import ( RuntimeControlReplayCache, validate_runtime_control_request) from tensorrt_llm.serve.serving_extensions import load_builtin_extensions +from tensorrt_llm.serve.tokenization_owners import (TokenizationOwners, + routing_enabled) from tensorrt_llm.serve.tool_parser.tool_parser_factory import ToolParserFactory from tensorrt_llm.serve.visual_gen_metrics import ( build_visual_gen_server_timings, build_visual_gen_timing_headers) @@ -164,6 +169,13 @@ from tensorrt_llm.visual_gen import VisualGen +def _preprocess_in_conversation(routing_id: Optional[str], preprocess_fn, + *args): + """Run ``preprocess_fn`` with its prompt tagged as conversation ``routing_id``.""" + with tokenization_routing_key(routing_id): + return preprocess_fn(*args) + + def _is_visual_gen_instance(obj) -> bool: """isinstance(obj, VisualGen) without importing the visual_gen tree. @@ -754,6 +766,16 @@ def __init__( if multi_frontend_serving is not None and multi_frontend_serving.is_attached: self._launcher_forwarder = LauncherForwarder( multi_frontend_serving.launcher_uds) + self._tokenization_owners: Optional[TokenizationOwners] = None + prefix_cache = prefix_token_cache_of( + getattr(generator, "input_processor", None)) + if (routing_enabled() and multi_frontend_serving is not None + and multi_frontend_serving.num_frontends > 1 + and prefix_cache is not None): + self._tokenization_owners = TokenizationOwners( + os.path.dirname(multi_frontend_serving.launcher_uds), + multi_frontend_serving.frontend_id, + multi_frontend_serving.num_frontends, *prefix_cache) self._shared_clock_offset: Optional[SharedClockOffset] = ( SharedClockOffset( os.path.dirname(multi_frontend_serving.launcher_uds)) @@ -866,6 +888,8 @@ async def lifespan(app: FastAPI): pass if self._launcher_forwarder is not None: await self._launcher_forwarder.close() + if self._tokenization_owners is not None: + self._tokenization_owners.close() if self.embedding_batcher is not None: await self.embedding_batcher.shutdown() logger.info("Stopped encode dynamic batcher") @@ -2481,11 +2505,16 @@ async def chat_stream_generator( generate_inputs = prompt preprocess_fn = getattr(self.generator, "preprocess", None) if preprocess_fn is not None: + preprocess = functools.partial(preprocess_fn, prompt, + sampling_params, + disaggregated_params) + if self._tokenization_owners is not None: + preprocess = functools.partial( + _preprocess_in_conversation, + get_request_routing_id(request), preprocess) loop = asyncio.get_event_loop() generate_inputs = await loop.run_in_executor( - self._input_proc_executor, - functools.partial(preprocess_fn, prompt, sampling_params, - disaggregated_params)) + self._input_proc_executor, preprocess) promise = self.generator.generate_async( inputs=generate_inputs, diff --git a/tensorrt_llm/serve/tokenization_owners.py b/tensorrt_llm/serve/tokenization_owners.py new file mode 100644 index 000000000000..6012529b313d --- /dev/null +++ b/tensorrt_llm/serve/tokenization_owners.py @@ -0,0 +1,200 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Tokenize each conversation in the frontend that caches it. + +With ``num_serve_frontends > 1`` the kernel hands each connection to any +frontend, so consecutive turns of a conversation land in different processes, +each with its own prefix-tokenization cache. Every frontend therefore serves +its cache on a Unix socket in the multi-frontend ipc dir, and a prompt whose +conversation id hashes to another frontend is tokenized there, by the cache +that has seen the conversation's earlier turns. Only the prompt text and the +token ids cross the socket; the request itself stays where it arrived. +""" + +import contextlib +import os +import socket +import socketserver +import struct +import threading +import zlib +from array import array +from typing import Any, Optional + +from tensorrt_llm.inputs.prefix_token_cache import PrefixTokenCache +from tensorrt_llm.logger import logger + +SOCKET_NAME = "tokenize-{}.sock" +# Set to 1 to route tokenization by conversation id across serving frontends. +ENABLE_ENV_VAR = "TLLM_PREFIX_TOKEN_CACHE_ROUTE_BY_CONVERSATION" + +_LENGTH = struct.Struct("!q") + + +def routing_enabled() -> bool: + """Whether tokenization is routed by conversation id (off unless enabled).""" + return os.environ.get(ENABLE_ENV_VAR) == "1" + + +def owner_of(key: str, num_frontends: int) -> int: + """The frontend that tokenizes conversation ``key``, the same in every process.""" + return zlib.crc32(key.encode()) % num_frontends + + +class TokenizationOwners: + """Route ``cache``'s prompts to the frontend that owns their conversation. + + Serves this frontend's ``cache`` to the other frontends and installs + ``cache.route``. A prompt the owner cannot tokenize, for example because it + has not started or has died, is tokenized locally. + """ + + def __init__( + self, + ipc_dir: str, + frontend_id: int, + num_frontends: int, + cache: PrefixTokenCache, + tokenizer: Any, + timeout_s: float = 30.0, + ): + self._ipc_dir = ipc_dir + self._frontend_id = frontend_id + self._num_frontends = num_frontends + self._timeout_s = timeout_s + # One connection per owner per tokenizing thread. + self._connections = threading.local() + self._all_connections: set[socket.socket] = set() + self._lock = threading.Lock() + + path = os.path.join(ipc_dir, SOCKET_NAME.format(frontend_id)) + self._server = _Server(path, cache, tokenizer) + threading.Thread( + target=self._server.serve_forever, name="trtllm_tokenize_owner", daemon=True + ).start() + cache.route = self._route + + def close(self) -> None: + """Stop serving and close every connection, in both directions.""" + self._server.shutdown() + self._server.server_close() + self._server.close_connections() + with self._lock: + for conn in self._all_connections: + conn.close() + self._all_connections.clear() + + def _route(self, key: str, text: str) -> Optional[list[int]]: + owner = owner_of(key, self._num_frontends) + if owner == self._frontend_id: + return None + try: + return self._request(owner, text) + except OSError as e: + self._drop(owner) + logger.debug(f"Tokenizing locally, frontend {owner} did not answer: {e!r}") + return None + + def _request(self, owner: int, text: str) -> list[int]: + conn = self._connection(owner) + data = text.encode() + conn.sendall(_LENGTH.pack(len(data)) + data) + (count,) = _LENGTH.unpack(_recv_exactly(conn, _LENGTH.size)) + ids = array("i") + ids.frombytes(_recv_exactly(conn, count * ids.itemsize)) + return ids.tolist() + + def _connection(self, owner: int) -> socket.socket: + conns = self._connections.__dict__ + conn = conns.get(owner) + if conn is None: + conn = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) + conn.settimeout(self._timeout_s) + try: + conn.connect(os.path.join(self._ipc_dir, SOCKET_NAME.format(owner))) + except OSError: + conn.close() + raise + conns[owner] = conn + with self._lock: + self._all_connections.add(conn) + return conn + + def _drop(self, owner: int) -> None: + conn = self._connections.__dict__.pop(owner, None) + if conn is not None: + with self._lock: + self._all_connections.discard(conn) + conn.close() + + +class _Server(socketserver.ThreadingUnixStreamServer): + daemon_threads = True + + def __init__(self, path: str, cache: PrefixTokenCache, tokenizer: Any): + self.cache = cache + self.tokenizer = tokenizer + # Accepted connections and the threads serving them. + self.handlers: dict[socket.socket, threading.Thread] = {} + self.handlers_lock = threading.Lock() + if os.path.exists(path): + os.unlink(path) + super().__init__(path, _Handler) + + def close_connections(self, timeout_s: float = 5.0) -> None: + """Wake every handler blocked on its peer and wait for it to exit.""" + with self.handlers_lock: + handlers = list(self.handlers.items()) + for conn, _ in handlers: + with contextlib.suppress(OSError): + conn.shutdown(socket.SHUT_RDWR) + for _, thread in handlers: + thread.join(timeout_s) + + +class _Handler(socketserver.BaseRequestHandler): + server: _Server + + def setup(self) -> None: + with self.server.handlers_lock: + self.server.handlers[self.request] = threading.current_thread() + + def finish(self) -> None: + with self.server.handlers_lock: + self.server.handlers.pop(self.request, None) + + def handle(self) -> None: + # Ends when the peer closes the connection or close_connections() runs. + with contextlib.suppress(OSError): + while True: + self._answer_one() + + def _answer_one(self) -> None: + conn = self.request + (size,) = _LENGTH.unpack(_recv_exactly(conn, _LENGTH.size)) + text = _recv_exactly(conn, size).decode() + ids = array("i", self.server.cache.encode(self.server.tokenizer, text)) + conn.sendall(_LENGTH.pack(len(ids)) + ids.tobytes()) + + +def _recv_exactly(conn: socket.socket, size: int) -> bytearray: + buf = bytearray(size) + view = memoryview(buf) + while view: + n = conn.recv_into(view) + if n == 0: + raise ConnectionError("peer closed the connection") + view = view[n:] + return buf diff --git a/tests/unittest/llmapi/test_serve_tokenization_owners.py b/tests/unittest/llmapi/test_serve_tokenization_owners.py new file mode 100644 index 000000000000..07e9878653ac --- /dev/null +++ b/tests/unittest/llmapi/test_serve_tokenization_owners.py @@ -0,0 +1,108 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Each conversation is tokenized by the frontend that owns it. + +Two frontends are simulated in one process: each has its own cache and serves it +on its own socket, as two trtllm-serve frontend processes do. +""" + +from typing import Any + +import pytest + +from tensorrt_llm.inputs.prefix_token_cache import ( + PrefixTokenCache, + PrefixTokenCacheConfig, + tokenization_routing_key, +) +from tensorrt_llm.serve.tokenization_owners import TokenizationOwners, owner_of + +pytestmark = pytest.mark.cpu_only + + +class _CharTokenizer: + is_fast = True + + def __call__(self, text: str, **kwargs: Any) -> dict[str, Any]: + return { + "input_ids": [ord(c) for c in text], + "offset_mapping": [(i, i + 1) for i in range(len(text))], + } + + +def _cache() -> PrefixTokenCache: + return PrefixTokenCache( + PrefixTokenCacheConfig( + max_entries=64, overlap=4, resync=2, min_chars=64, checkpoint_tokens=8 + ) + ) + + +def _turns(n: int = 6) -> list[str]: + text, out = "system preamble " * 8, [] + for i in range(n): + text += f"turn {i} says something new. " + out.append(text) + return out + + +def _key_owned_by(owner: int, num_frontends: int) -> str: + return next( + key + for key in (f"conversation-{i}" for i in range(100)) + if owner_of(key, num_frontends) == owner + ) + + +def test_a_conversation_is_tokenized_by_its_owner_wherever_its_turns_arrive(tmp_path) -> None: + tokenizer = _CharTokenizer() + caches = [_cache(), _cache()] + owners = [ + TokenizationOwners(str(tmp_path), frontend_id, 2, cache, tokenizer) + for frontend_id, cache in enumerate(caches) + ] + try: + key = _key_owned_by(1, 2) + turns = _turns() + # Turns alternate between the frontends, as SO_REUSEPORT spreads them. + for i, prompt in enumerate(turns): + with tokenization_routing_key(key): + assert caches[i % 2].encode(tokenizer, prompt) == [ord(c) for c in prompt] + + assert caches[0].routed == len(turns) // 2 + assert caches[0].hits + caches[0].misses == 0 + # The owner saw every turn, so every turn after the first hit. + assert caches[1].hits == len(turns) - 1 + + # A prompt without a conversation stays where it arrived. + caches[0].encode(tokenizer, turns[0]) + assert caches[0].misses == 1 + finally: + for owner in owners: + owner.close() + + +def test_a_prompt_is_tokenized_locally_when_its_owner_is_unreachable(tmp_path) -> None: + tokenizer, cache = _CharTokenizer(), _cache() + # Frontend 1 never starts. + owners = TokenizationOwners(str(tmp_path), 0, 2, cache, tokenizer) + try: + prompt = _turns(1)[0] + with tokenization_routing_key(_key_owned_by(1, 2)): + assert cache.encode(tokenizer, prompt) == [ord(c) for c in prompt] + assert cache.routed == 0 + assert cache.misses == 1 + finally: + owners.close() From ade688d2aad20fc7d42da10a37e7a8a51e0fb049 Mon Sep 17 00:00:00 2001 From: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com> Date: Tue, 3 Nov 2026 18:40:08 -0800 Subject: [PATCH 4/8] [None][perf] Count routing failures and accept every frontend's connections - The owner's listen backlog is SOMAXCONN: a Unix-socket connect to a full backlog fails at once, so with the default of 5 the first turns of conversations starting together were tokenized locally. - An owner that cannot tokenize a prompt replies with a failure instead of closing the connection, and the asking frontend tokenizes it itself. - Prompts tokenized locally because routing failed are counted in the periodic cache log line. The owner timeout is 5 s. - The doc's single-frontend numbers are measured with the divergent splice. Signed-off-by: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com> --- .../features/prefix-tokenization-cache.md | 33 +++++++++--------- tensorrt_llm/inputs/prefix_token_cache.py | 26 ++++++++++---- tensorrt_llm/serve/tokenization_owners.py | 34 ++++++++++++++----- .../llmapi/test_serve_tokenization_owners.py | 24 +++++++++---- 4 files changed, 80 insertions(+), 37 deletions(-) diff --git a/docs/source/features/prefix-tokenization-cache.md b/docs/source/features/prefix-tokenization-cache.md index d14f8217ed26..b709a39431a5 100644 --- a/docs/source/features/prefix-tokenization-cache.md +++ b/docs/source/features/prefix-tokenization-cache.md @@ -85,10 +85,11 @@ feature you turned on. - An unexpected exception inside the cache disables it, with a warning, and the request is tokenized normally. A cache problem never fails a request. - The cache logs its configuration at startup and, every 1000 cache-eligible - requests, an `INFO` line with the hit rate (and how many hits were on diverging - prompts), the number of re-synchronization - fallbacks, and its occupancy, so a cache that is enabled but never hitting is - visible in the server log. + requests, an `INFO` line with the hit rate (and how many hits were on + diverging prompts), the number of re-synchronization fallbacks, how many + prompts were routed to another frontend and how many were tokenized locally + because routing failed (see below), and its occupancy, so a cache that is + enabled but never hitting is visible in the server log. ## Multiple serving frontends @@ -105,8 +106,8 @@ usually has not seen the conversation's previous turn, so the prompt is tokenized in full. On an agentic multi-turn trace with prompts of about 160k tokens on average, -at one session, the p90 time spent in the frontend was 453 ms with 8 -frontends and 166 ms with one. Decode throughput per user was the same up to +at one session, the p90 time spent in the frontend was about 165 ms with 8 +frontends and 44 ms with one. Decode throughput per user was the same up to 32 concurrent sessions, because at that concurrency the frontend is not the bottleneck. @@ -115,18 +116,18 @@ at high request rates with short prompts. That is not the workload this cache targets. If you need several frontends, set -`TLLM_PREFIX_TOKEN_CACHE_ROUTE_BY_CONVERSATION=1` to tokenize each -conversation in one frontend. A chat request's conversation id (the same id -attention DP routes by, from `conversation_params.conversation_id` or headers -such as `X-Correlation-ID`) is hashed to an owner frontend, and the frontend the +`TLLM_PREFIX_TOKEN_CACHE_ROUTE_BY_CONVERSATION=1` to tokenize each conversation +in one frontend. A chat request's conversation id (the same id attention DP +routes by, from `conversation_params.conversation_id` or headers such as +`X-Correlation-ID`) is hashed to an owner frontend, and the frontend the request arrived at sends the prompt text to the owner over a local Unix socket and gets back the token ids; the request itself stays where it arrived. A -request without a conversation id, or whose owner does not answer, is -tokenized locally. The switch is off by default and has no effect with a -single frontend or with the cache disabled. On the trace above with 8 -frontends, it brought the p90 frontend time at one session from about 165 ms -to about 58 ms, against 44 ms with a single frontend (both measured with the -divergent splice). +request without a conversation id is tokenized locally, and so is one whose +owner does not answer within 5 seconds or cannot tokenize the prompt; the +latter are counted as route failures in the periodic log line. The switch is +off by default and has no effect with a single frontend or with the cache +disabled. On the trace above with 8 frontends, it brought the p90 frontend time +at one session from about 165 ms to about 58 ms. ## Measured effect diff --git a/tensorrt_llm/inputs/prefix_token_cache.py b/tensorrt_llm/inputs/prefix_token_cache.py index 875e40d36757..30b46257ba32 100644 --- a/tensorrt_llm/inputs/prefix_token_cache.py +++ b/tensorrt_llm/inputs/prefix_token_cache.py @@ -46,6 +46,7 @@ __all__ = [ "PrefixTokenCache", "PrefixTokenCacheConfig", + "RouteFailed", "create_prefix_token_cache", "tokenization_routing_key", ] @@ -78,6 +79,10 @@ def tokenization_routing_key(key: str | None) -> Iterator[None]: _ROUTING_KEY.reset(token) +class RouteFailed(Exception): + """``PrefixTokenCache.route`` could not get the ids from the owning process.""" + + class OffsetTokenizer(Protocol): """The subset of the HF fast-tokenizer call interface the cache relies on.""" @@ -187,11 +192,13 @@ def __init__(self, config: PrefixTokenCacheConfig) -> None: # Hits on a cached prompt the new one diverges from (also in ``hits``). self.divergent_hits = 0 # Tokenizes a prompt tagged with a routing key in the process that - # owns that conversation, or returns None to tokenize it here. Set by - # multi-frontend serving. + # owns that conversation, or returns None to tokenize it here. Raises + # RouteFailed if the owner did not answer. Set by multi-frontend serving. self.route: Callable[[str, str], list[int] | None] | None = None # Prompts another process tokenized through ``route``. self.routed = 0 + # Prompts tokenized here because ``route`` failed. + self.route_failures = 0 def encode(self, tokenizer: OffsetTokenizer, text: str) -> list[int]: """Return the token ids of ``text`` without special tokens.""" @@ -199,11 +206,17 @@ def encode(self, tokenizer: OffsetTokenizer, text: str) -> list[int]: return self._tokenize(tokenizer, text)[0] key = _ROUTING_KEY.get() if self.route is not None and key is not None: - ids = self.route(key, text) - if ids is not None: + try: + ids = self.route(key, text) + except RouteFailed as e: + logger.debug(f"Tokenizing a prompt locally: {e}") with self._lock: - self.routed += 1 - return ids + self.route_failures += 1 + else: + if ids is not None: + with self._lock: + self.routed += 1 + return ids try: return self._encode(tokenizer, text) except Exception as e: @@ -220,6 +233,7 @@ def summary(self) -> str: return ( f"{self.hits}/{total} hits ({rate:.1%}, {self.divergent_hits} divergent), " f"{self.resync_failures} resync fallbacks, {self.routed} routed, " + f"{self.route_failures} route failures, " f"{len(self._entries)} entries, {self._total_chars} cached chars" ) diff --git a/tensorrt_llm/serve/tokenization_owners.py b/tensorrt_llm/serve/tokenization_owners.py index 6012529b313d..b15c41c5dcd0 100644 --- a/tensorrt_llm/serve/tokenization_owners.py +++ b/tensorrt_llm/serve/tokenization_owners.py @@ -33,7 +33,7 @@ from array import array from typing import Any, Optional -from tensorrt_llm.inputs.prefix_token_cache import PrefixTokenCache +from tensorrt_llm.inputs.prefix_token_cache import PrefixTokenCache, RouteFailed from tensorrt_llm.logger import logger SOCKET_NAME = "tokenize-{}.sock" @@ -41,6 +41,11 @@ ENABLE_ENV_VAR = "TLLM_PREFIX_TOKEN_CACHE_ROUTE_BY_CONVERSATION" _LENGTH = struct.Struct("!q") +# Sent in place of a token count when the owner could not tokenize the prompt. +_FAILED = -1 +# Lone surrogates, which JSON can carry, survive the trip and fail in the +# tokenizer as they would locally. +_TEXT_ERRORS = "surrogatepass" def routing_enabled() -> bool: @@ -57,8 +62,8 @@ class TokenizationOwners: """Route ``cache``'s prompts to the frontend that owns their conversation. Serves this frontend's ``cache`` to the other frontends and installs - ``cache.route``. A prompt the owner cannot tokenize, for example because it - has not started or has died, is tokenized locally. + ``cache.route``. A prompt the owner does not tokenize within ``timeout_s``, + for example because it has not started or has died, is tokenized locally. """ def __init__( @@ -68,7 +73,7 @@ def __init__( num_frontends: int, cache: PrefixTokenCache, tokenizer: Any, - timeout_s: float = 30.0, + timeout_s: float = 5.0, ): self._ipc_dir = ipc_dir self._frontend_id = frontend_id @@ -104,14 +109,15 @@ def _route(self, key: str, text: str) -> Optional[list[int]]: return self._request(owner, text) except OSError as e: self._drop(owner) - logger.debug(f"Tokenizing locally, frontend {owner} did not answer: {e!r}") - return None + raise RouteFailed(f"frontend {owner} did not answer: {e!r}") from e def _request(self, owner: int, text: str) -> list[int]: conn = self._connection(owner) - data = text.encode() + data = text.encode(errors=_TEXT_ERRORS) conn.sendall(_LENGTH.pack(len(data)) + data) (count,) = _LENGTH.unpack(_recv_exactly(conn, _LENGTH.size)) + if count == _FAILED: + raise RouteFailed(f"frontend {owner} could not tokenize the prompt") ids = array("i") ids.frombytes(_recv_exactly(conn, count * ids.itemsize)) return ids.tolist() @@ -142,6 +148,9 @@ def _drop(self, owner: int) -> None: class _Server(socketserver.ThreadingUnixStreamServer): daemon_threads = True + # Every tokenizing thread of every other frontend may connect at once, and a + # Unix-socket connect to a full backlog fails at once instead of waiting. + request_queue_size = socket.SOMAXCONN def __init__(self, path: str, cache: PrefixTokenCache, tokenizer: Any): self.cache = cache @@ -184,8 +193,15 @@ def handle(self) -> None: def _answer_one(self) -> None: conn = self.request (size,) = _LENGTH.unpack(_recv_exactly(conn, _LENGTH.size)) - text = _recv_exactly(conn, size).decode() - ids = array("i", self.server.cache.encode(self.server.tokenizer, text)) + text = _recv_exactly(conn, size).decode(errors=_TEXT_ERRORS) + try: + ids = array("i", self.server.cache.encode(self.server.tokenizer, text)) + except Exception as e: + # The asking frontend tokenizes the prompt itself and reports the + # error with its request. + logger.debug(f"Could not tokenize a prompt for another frontend: {e!r}") + conn.sendall(_LENGTH.pack(_FAILED)) + return conn.sendall(_LENGTH.pack(len(ids)) + ids.tobytes()) diff --git a/tests/unittest/llmapi/test_serve_tokenization_owners.py b/tests/unittest/llmapi/test_serve_tokenization_owners.py index 07e9878653ac..83f3469189e2 100644 --- a/tests/unittest/llmapi/test_serve_tokenization_owners.py +++ b/tests/unittest/llmapi/test_serve_tokenization_owners.py @@ -94,15 +94,27 @@ def test_a_conversation_is_tokenized_by_its_owner_wherever_its_turns_arrive(tmp_ owner.close() -def test_a_prompt_is_tokenized_locally_when_its_owner_is_unreachable(tmp_path) -> None: +class _FailingTokenizer(_CharTokenizer): + def __call__(self, text: str, **kwargs: Any) -> dict[str, Any]: + raise ValueError("cannot tokenize") + + +def test_a_prompt_is_tokenized_locally_when_its_owner_does_not_tokenize_it(tmp_path) -> None: tokenizer, cache = _CharTokenizer(), _cache() - # Frontend 1 never starts. - owners = TokenizationOwners(str(tmp_path), 0, 2, cache, tokenizer) + owners = [TokenizationOwners(str(tmp_path), 0, 2, cache, tokenizer)] try: prompt = _turns(1)[0] with tokenization_routing_key(_key_owned_by(1, 2)): + # Frontend 1 has not started. assert cache.encode(tokenizer, prompt) == [ord(c) for c in prompt] - assert cache.routed == 0 - assert cache.misses == 1 + assert (cache.routed, cache.route_failures, cache.misses) == (0, 1, 1) + + # Frontend 1 runs but its tokenizer fails; the connection survives. + owners.append(TokenizationOwners(str(tmp_path), 1, 2, _cache(), _FailingTokenizer())) + for _ in range(2): + assert cache.encode(tokenizer, prompt) == [ord(c) for c in prompt] + assert (cache.routed, cache.route_failures) == (0, 3) + assert "3 route failures" in cache.summary() finally: - owners.close() + for owner in owners: + owner.close() From d11ca967028b7ee0f2bfcfd851a90a49b803fdbe Mon Sep 17 00:00:00 2001 From: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com> Date: Wed, 4 Nov 2026 01:36:16 -0800 Subject: [PATCH 5/8] [None][fix] Address review of the prefix-tokenization cache changes - Reject a splice whose re-tokenized tail opens with several tokens at one character offset (byte-fallback characters, expanded ligatures), where a run of equal ids could hide a duplicated token. - Hash a conversation id with a lone surrogate instead of failing the request. - Assert the hit or resync fallback of every rewritten prompt. - Note that multiple serving frontends need the classic IPC executor. Signed-off-by: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com> --- .../features/prefix-tokenization-cache.md | 3 +- tensorrt_llm/inputs/prefix_token_cache.py | 17 ++++++++- tensorrt_llm/serve/tokenization_owners.py | 4 +- .../inputs/test_prefix_token_cache.py | 38 ++++++++++++++++++- .../llmapi/test_serve_tokenization_owners.py | 5 +++ 5 files changed, 61 insertions(+), 6 deletions(-) diff --git a/docs/source/features/prefix-tokenization-cache.md b/docs/source/features/prefix-tokenization-cache.md index b709a39431a5..c24ed6544940 100644 --- a/docs/source/features/prefix-tokenization-cache.md +++ b/docs/source/features/prefix-tokenization-cache.md @@ -94,7 +94,8 @@ feature you turned on. ## Multiple serving frontends Use a single serving frontend (`num_serve_frontends: 1`, the default) when you -enable the cache for multi-turn traffic. +enable the cache for multi-turn traffic. Multiple frontends are available only +with the classic IPC executor, so this section applies only there. The cache lives in the memory of one frontend process and is not shared between processes. With `num_serve_frontends: K`, the K frontend processes diff --git a/tensorrt_llm/inputs/prefix_token_cache.py b/tensorrt_llm/inputs/prefix_token_cache.py index 30b46257ba32..432b1b2112df 100644 --- a/tensorrt_llm/inputs/prefix_token_cache.py +++ b/tensorrt_llm/inputs/prefix_token_cache.py @@ -243,7 +243,7 @@ def _encode(self, tokenizer: OffsetTokenizer, text: str) -> list[int]: if found is not None: entry = found.entry tail_ids, tail_offsets = self._tokenize(tokenizer, text[found.split_char :]) - if self._resynced(entry, found.split_token, tail_ids): + if self._resynced(entry, found.split_token, tail_ids, tail_offsets): ids = entry.ids[: found.split_token].tolist() + list(tail_ids) base_token, base_char = found.split_token, found.split_char if ids is None: @@ -305,8 +305,21 @@ def _checkpoints( out.append(base_char + tail_offsets[index - base_token][0]) return out - def _resynced(self, entry: _Entry, split_token: int, tail_ids: Sequence[int]) -> bool: + def _resynced( + self, + entry: _Entry, + split_token: int, + tail_ids: Sequence[int], + tail_offsets: Sequence[tuple[int, int]], + ) -> bool: """Whether the tokenizer re-synchronized with the cached ids at the seam.""" + # A character tokenized into several tokens (a byte-fallback character, + # an expanded ligature) reports the same start for each. If the tail + # opens with one, the split token may not be the first of them and the + # tail would repeat the cached ids before it, which a run of equal ids + # can hide from the comparison below. + if len(tail_offsets) > 1 and tail_offsets[1][0] == tail_offsets[0][0]: + return False cached_tail = entry.ids[split_token:] span = min(self._config.resync, len(tail_ids), len(cached_tail)) return span > 0 and list(tail_ids[:span]) == cached_tail[:span].tolist() diff --git a/tensorrt_llm/serve/tokenization_owners.py b/tensorrt_llm/serve/tokenization_owners.py index b15c41c5dcd0..e10c1b20981d 100644 --- a/tensorrt_llm/serve/tokenization_owners.py +++ b/tensorrt_llm/serve/tokenization_owners.py @@ -44,7 +44,7 @@ # Sent in place of a token count when the owner could not tokenize the prompt. _FAILED = -1 # Lone surrogates, which JSON can carry, survive the trip and fail in the -# tokenizer as they would locally. +# tokenizer as they would locally, and do not fail a conversation id's hash. _TEXT_ERRORS = "surrogatepass" @@ -55,7 +55,7 @@ def routing_enabled() -> bool: def owner_of(key: str, num_frontends: int) -> int: """The frontend that tokenizes conversation ``key``, the same in every process.""" - return zlib.crc32(key.encode()) % num_frontends + return zlib.crc32(key.encode(errors=_TEXT_ERRORS)) % num_frontends class TokenizationOwners: diff --git a/tests/unittest/inputs/test_prefix_token_cache.py b/tests/unittest/inputs/test_prefix_token_cache.py index c22490df39de..1ba639f5c862 100644 --- a/tests/unittest/inputs/test_prefix_token_cache.py +++ b/tests/unittest/inputs/test_prefix_token_cache.py @@ -75,6 +75,25 @@ def __call__(self, text: str, **kwargs: Any) -> dict[str, Any]: return {"input_ids": ids, "offset_mapping": offsets} +class _ExpandingTokenizer: + """Char-level, except that "ff" becomes two equal tokens with the same offsets. + + As with a byte-fallback character or an NFKC ligature, a split can land on + the second token of the character, and a long run of them makes the tail + agree with the cached ids however it is misaligned. + """ + + is_fast = True + + def __call__(self, text: str, **kwargs: Any) -> dict[str, Any]: + ids, offsets = [], [] + for i, char in enumerate(text): + for _ in range(2 if char == "ff" else 1): + ids.append(_tok_id(char)) + offsets.append((i, i + 1)) + return {"input_ids": ids, "offset_mapping": offsets} + + _MIN_CHARS = 16 # A shared opening longer than min_chars, so growing prompts land in the same # bucket and the prefix lookup can find them. @@ -138,13 +157,30 @@ def test_rewritten_history_is_identical_to_full_tokenization(tokenizer: Any) -> assert cache.encode(tokenizer, prompt) == tokenizer(prompt)["input_ids"] # Trim the last message and append a different one, as agents do. rewritten = prompt[: len(prompt) - 20] + f"rewritten tail {i} there. " + hits, failures = cache.hits, cache.resync_failures assert cache.encode(tokenizer, rewritten) == tokenizer(rewritten)["input_ids"] + if isinstance(tokenizer, _MergeTokenizer): + assert cache.hits == hits + 1 + else: + # Spliced from the history and rejected at the seam. + assert cache.resync_failures == failures + 1 if isinstance(tokenizer, _MergeTokenizer): assert cache.divergent_hits >= len(turns) # The rewritten branch does not evict the history it diverged from. assert any(entry.text == turns[-1] for entry in cache._entries.values()) else: - assert cache.divergent_hits == 0 and cache.resync_failures > 0 + assert cache.divergent_hits == 0 + + +def test_a_split_inside_a_multi_token_character_is_not_used() -> None: + tokenizer, cache = _ExpandingTokenizer(), _cache() + text = _PREAMBLE + for i in range(12): + # The odd character shifts every later split onto the other token of a pair. + text += "ff" * 40 + "." * (i % 2) + assert cache.encode(tokenizer, text) == tokenizer(text)["input_ids"] + rewritten = text[:-10] + "rewritten. " + assert cache.encode(tokenizer, rewritten) == tokenizer(rewritten)["input_ids"] def test_short_prompts_are_tokenized_but_not_cached() -> None: diff --git a/tests/unittest/llmapi/test_serve_tokenization_owners.py b/tests/unittest/llmapi/test_serve_tokenization_owners.py index 83f3469189e2..ed0ac5c8d706 100644 --- a/tests/unittest/llmapi/test_serve_tokenization_owners.py +++ b/tests/unittest/llmapi/test_serve_tokenization_owners.py @@ -118,3 +118,8 @@ def test_a_prompt_is_tokenized_locally_when_its_owner_does_not_tokenize_it(tmp_p finally: for owner in owners: owner.close() + + +def test_a_conversation_id_with_a_lone_surrogate_has_an_owner() -> None: + # JSON can carry one; it must not fail the request it arrives with. + assert owner_of("conversation-\ud800", 2) in (0, 1) From ba3ded1c8f4a766be3b7600e8fe555625ce10a31 Mon Sep 17 00:00:00 2001 From: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com> Date: Fri, 6 Nov 2026 00:43:20 -0800 Subject: [PATCH 6/8] [None][fix] Address the second review of the prefix-tokenization cache changes - Split at the first token of a multi-token character instead of rejecting such seams, so byte-fallback-heavy text keeps its hits. - Find the end of a common prefix by halving the mismatching block. - Do not promote the entry a diverging prompt was spliced from. - Skip a tokenization owner for a while after it does not answer. - Type the tokenizer and the preprocess callable; keep per-thread owner connections in a threading.local subclass. - Run test_serve_tokenization_owners.py in l0_cpu. Signed-off-by: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com> --- .../features/prefix-tokenization-cache.md | 3 +- tensorrt_llm/inputs/prefix_token_cache.py | 91 ++++++++++++------- tensorrt_llm/serve/openai_server.py | 11 ++- tensorrt_llm/serve/tokenization_owners.py | 32 +++++-- .../integration/test_lists/test-db/l0_cpu.yml | 1 + .../inputs/test_prefix_token_cache.py | 5 +- .../llmapi/test_serve_tokenization_owners.py | 29 +++++- 7 files changed, 121 insertions(+), 51 deletions(-) diff --git a/docs/source/features/prefix-tokenization-cache.md b/docs/source/features/prefix-tokenization-cache.md index c24ed6544940..ed734bf2a197 100644 --- a/docs/source/features/prefix-tokenization-cache.md +++ b/docs/source/features/prefix-tokenization-cache.md @@ -125,7 +125,8 @@ request arrived at sends the prompt text to the owner over a local Unix socket and gets back the token ids; the request itself stays where it arrived. A request without a conversation id is tokenized locally, and so is one whose owner does not answer within 5 seconds or cannot tokenize the prompt; the -latter are counted as route failures in the periodic log line. The switch is +latter are counted as route failures in the periodic log line. An owner that +did not answer is not asked again for 30 seconds. The switch is off by default and has no effect with a single frontend or with the cache disabled. On the trace above with 8 frontends, it brought the p90 frontend time at one session from about 165 ms to about 58 ms. diff --git a/tensorrt_llm/inputs/prefix_token_cache.py b/tensorrt_llm/inputs/prefix_token_cache.py index 432b1b2112df..e443a45158b5 100644 --- a/tensorrt_llm/inputs/prefix_token_cache.py +++ b/tensorrt_llm/inputs/prefix_token_cache.py @@ -44,6 +44,7 @@ from ..logger import logger __all__ = [ + "OffsetTokenizer", "PrefixTokenCache", "PrefixTokenCacheConfig", "RouteFailed", @@ -146,12 +147,14 @@ class _Entry: text: str ids: array # int32 # Every prompt that keeps ``text[:split_char]`` is re-tokenized from the same - # place: token ``split_token``, which starts at character ``split_char``. + # place: token ``split_token``, the first token of character ``split_char``. split_token: int split_char: int probe: str # the PROBE_CHARS characters of ``text`` before ``split_char`` - # Start character of every ``checkpoint_tokens``-th token (0, C, 2C, ...), - # where a prompt that diverges from ``text`` can be re-tokenized from. + # Where a prompt that diverges from ``text`` can be re-tokenized from: the + # first token of the character holding every ``checkpoint_tokens``-th token + # (0, C, 2C, ...), and the start character of that token. + checkpoint_tokens: array # int32 checkpoints: array # int32 @@ -243,7 +246,7 @@ def _encode(self, tokenizer: OffsetTokenizer, text: str) -> list[int]: if found is not None: entry = found.entry tail_ids, tail_offsets = self._tokenize(tokenizer, text[found.split_char :]) - if self._resynced(entry, found.split_token, tail_ids, tail_offsets): + if self._resynced(entry, found.split_token, tail_ids): ids = entry.ids[: found.split_token].tolist() + list(tail_ids) base_token, base_char = found.split_token, found.split_char if ids is None: @@ -253,6 +256,9 @@ def _encode(self, tokenizer: OffsetTokenizer, text: str) -> list[int]: # Where a prompt extending this one will be re-tokenized from. split_token = len(ids) - self._config.overlap tail_index = split_token - base_token + if tail_index >= 0: + tail_index = _first_of_char(tail_offsets, tail_index) + split_token = base_token + tail_index with self._lock: if found is None: @@ -269,11 +275,16 @@ def _encode(self, tokenizer: OffsetTokenizer, text: str) -> list[int]: if found is not None and found.extends: self._remove(found.eid) if split_token > 0 and tail_index >= 0: - checkpoints = self._checkpoints( + checkpoint_tokens, checkpoints = self._checkpoints( found, base_token, base_char, tail_offsets, len(ids) ) self._insert( - text, ids, split_token, base_char + tail_offsets[tail_index][0], checkpoints + text, + ids, + split_token, + base_char + tail_offsets[tail_index][0], + checkpoint_tokens, + checkpoints, ) report = (self.hits + self.misses) % LOG_EVERY_REQUESTS == 0 summary = self.summary() if report else None @@ -295,31 +306,21 @@ def _checkpoints( base_char: int, tail_offsets: Sequence[tuple[int, int]], num_ids: int, - ) -> array: - """Checkpoints of a prompt whose ids are ``base_token`` cached ids plus a tail.""" + ) -> tuple[array, array]: + """Checkpoint tokens and chars of ``base_token`` cached ids plus a tail.""" step = self._config.checkpoint_tokens # Spliced ids before base_token are the cached entry's, at the same chars. - num_kept = -(-base_token // step) - out = array("i", found.entry.checkpoints[:num_kept] if base_token else ()) + num_kept = -(-base_token // step) if base_token else 0 + tokens = array("i", found.entry.checkpoint_tokens[:num_kept] if num_kept else ()) + chars = array("i", found.entry.checkpoints[:num_kept] if num_kept else ()) for index in range(num_kept * step, num_ids, step): - out.append(base_char + tail_offsets[index - base_token][0]) - return out + tail_index = _first_of_char(tail_offsets, index - base_token) + tokens.append(base_token + tail_index) + chars.append(base_char + tail_offsets[tail_index][0]) + return tokens, chars - def _resynced( - self, - entry: _Entry, - split_token: int, - tail_ids: Sequence[int], - tail_offsets: Sequence[tuple[int, int]], - ) -> bool: + def _resynced(self, entry: _Entry, split_token: int, tail_ids: Sequence[int]) -> bool: """Whether the tokenizer re-synchronized with the cached ids at the seam.""" - # A character tokenized into several tokens (a byte-fallback character, - # an expanded ligature) reports the same start for each. If the tail - # opens with one, the split token may not be the first of them and the - # tail would repeat the cached ids before it, which a run of equal ids - # can hide from the comparison below. - if len(tail_offsets) > 1 and tail_offsets[1][0] == tail_offsets[0][0]: - return False cached_tail = entry.ids[split_token:] span = min(self._config.resync, len(tail_ids), len(cached_tail)) return span > 0 and list(tail_ids[:span]) == cached_tail[:span].tolist() @@ -328,7 +329,7 @@ def _bucket_key(self, text: str) -> int: return hash(text[: self._config.min_chars]) def _lookup(self, text: str) -> _Match | None: - """Best cached entry to splice ``text`` from, promoted to MRU. + """Best cached entry to splice ``text`` from, promoted to MRU if ``text`` extends it. The entry with the longest re-tokenization prefix that ``text`` keeps wins. Without one, the @@ -356,7 +357,9 @@ def _lookup(self, text: str) -> _Match | None: match = self._divergent_match(eid, entry, text) if match is not None and (best is None or match.split_token > best.split_token): best = match - if best is not None: + # A diverging prompt is not promoted: when the history was rewritten + # rather than branched, the entry it diverged from is dead. + if best is not None and best.extends: with self._lock: if best.eid in self._entries: self._entries.move_to_end(best.eid) @@ -370,8 +373,7 @@ def _divergent_match(self, eid: int, entry: _Entry, text: str) -> _Match | None: index = bisect.bisect_right(entry.checkpoints, common) - 2 if index < 1: return None - split_token = index * self._config.checkpoint_tokens - return _Match(eid, entry, split_token, entry.checkpoints[index], False) + return _Match(eid, entry, entry.checkpoint_tokens[index], entry.checkpoints[index], False) def _insert( self, @@ -379,6 +381,7 @@ def _insert( ids: Sequence[int], split_token: int, split_char: int, + checkpoint_tokens: array, checkpoints: array, ) -> None: """Caller holds the lock.""" @@ -388,7 +391,9 @@ def _insert( return # concurrent misses on the same prompt eid, self._next_id = self._next_id, self._next_id + 1 probe = text[max(split_char - PROBE_CHARS, 0) : split_char] - entry = _Entry(text, array("i", ids), split_token, split_char, probe, checkpoints) + entry = _Entry( + text, array("i", ids), split_token, split_char, probe, checkpoint_tokens, checkpoints + ) self._entries[eid] = entry bucket[eid] = entry self._total_chars += len(text) @@ -416,9 +421,27 @@ def _common_prefix_len(a: str, b: str) -> int: n, i, block = min(len(a), len(b)), 0, 4096 while i < n and a[i : i + block] == b[i : i + block]: i += block - while i < n and a[i] == b[i]: - i += 1 - return min(i, n) + # The first mismatch, if any, lies in the block at i; halve it to find it. + lo, hi = min(i, n), min(i + block, n) + while lo < hi: + mid = (lo + hi + 1) // 2 + if a[lo:mid] == b[lo:mid]: + lo = mid + else: + hi = mid - 1 + return lo + + +def _first_of_char(offsets: Sequence[tuple[int, int]], index: int) -> int: + """First token of the character token ``index`` starts in. + + A character tokenized into several tokens (a byte-fallback character, an + expanded ligature) reports the same start for each; a split must not + separate them, or the tail would repeat the tokens of that character. + """ + while index > 0 and offsets[index - 1][0] == offsets[index][0]: + index -= 1 + return index def create_prefix_token_cache(tokenizer: object) -> PrefixTokenCache | None: diff --git a/tensorrt_llm/serve/openai_server.py b/tensorrt_llm/serve/openai_server.py index c40f43f9bade..b6c30daa4f19 100644 --- a/tensorrt_llm/serve/openai_server.py +++ b/tensorrt_llm/serve/openai_server.py @@ -23,7 +23,7 @@ from http import HTTPStatus from pathlib import Path from typing import (TYPE_CHECKING, Annotated, Any, AsyncGenerator, - AsyncIterator, Dict, List, Optional, Tuple, Union) + AsyncIterator, Callable, Dict, List, Optional, Tuple, Union) import msgspec import uvicorn @@ -59,7 +59,7 @@ from tensorrt_llm.llmapi import MultimodalEncoder, SchedulingParams, tracing from tensorrt_llm.llmapi.disagg_utils import (DisaggClusterConfig, MetadataServerConfig, ServerRole) -from tensorrt_llm.llmapi.llm import LLM, RequestOutput +from tensorrt_llm.llmapi.llm import LLM, PreprocessedInputs, RequestOutput from tensorrt_llm.llmapi.reasoning_parser import ReasoningParserFactory from tensorrt_llm.llmapi.thinking_budget import \ add_thinking_budget_logits_processor @@ -169,11 +169,12 @@ from tensorrt_llm.visual_gen import VisualGen -def _preprocess_in_conversation(routing_id: Optional[str], preprocess_fn, - *args): +def _preprocess_in_conversation( + routing_id: Optional[str], + preprocess_fn: Callable[[], PreprocessedInputs]) -> PreprocessedInputs: """Run ``preprocess_fn`` with its prompt tagged as conversation ``routing_id``.""" with tokenization_routing_key(routing_id): - return preprocess_fn(*args) + return preprocess_fn() def _is_visual_gen_instance(obj) -> bool: diff --git a/tensorrt_llm/serve/tokenization_owners.py b/tensorrt_llm/serve/tokenization_owners.py index e10c1b20981d..3f07be0bed84 100644 --- a/tensorrt_llm/serve/tokenization_owners.py +++ b/tensorrt_llm/serve/tokenization_owners.py @@ -29,11 +29,12 @@ import socketserver import struct import threading +import time import zlib from array import array -from typing import Any, Optional +from typing import Optional -from tensorrt_llm.inputs.prefix_token_cache import PrefixTokenCache, RouteFailed +from tensorrt_llm.inputs.prefix_token_cache import OffsetTokenizer, PrefixTokenCache, RouteFailed from tensorrt_llm.logger import logger SOCKET_NAME = "tokenize-{}.sock" @@ -58,12 +59,20 @@ def owner_of(key: str, num_frontends: int) -> int: return zlib.crc32(key.encode(errors=_TEXT_ERRORS)) % num_frontends +class _Connections(threading.local): + """This thread's connection to each owner.""" + + def __init__(self) -> None: + self.by_owner: dict[int, socket.socket] = {} + + class TokenizationOwners: """Route ``cache``'s prompts to the frontend that owns their conversation. Serves this frontend's ``cache`` to the other frontends and installs ``cache.route``. A prompt the owner does not tokenize within ``timeout_s``, - for example because it has not started or has died, is tokenized locally. + for example because it has not started, has died or has hung, is tokenized + locally, and so is every prompt it owns for the next ``retry_after_s``. """ def __init__( @@ -72,15 +81,19 @@ def __init__( frontend_id: int, num_frontends: int, cache: PrefixTokenCache, - tokenizer: Any, + tokenizer: OffsetTokenizer, timeout_s: float = 5.0, + retry_after_s: float = 30.0, ): self._ipc_dir = ipc_dir self._frontend_id = frontend_id self._num_frontends = num_frontends self._timeout_s = timeout_s + self._retry_after_s = retry_after_s + # Monotonic time before which an owner that did not answer is not asked. + self._skip_until: dict[int, float] = {} # One connection per owner per tokenizing thread. - self._connections = threading.local() + self._connections = _Connections() self._all_connections: set[socket.socket] = set() self._lock = threading.Lock() @@ -105,9 +118,12 @@ def _route(self, key: str, text: str) -> Optional[list[int]]: owner = owner_of(key, self._num_frontends) if owner == self._frontend_id: return None + if time.monotonic() < self._skip_until.get(owner, 0.0): + raise RouteFailed(f"frontend {owner} did not answer recently") try: return self._request(owner, text) except OSError as e: + self._skip_until[owner] = time.monotonic() + self._retry_after_s self._drop(owner) raise RouteFailed(f"frontend {owner} did not answer: {e!r}") from e @@ -123,7 +139,7 @@ def _request(self, owner: int, text: str) -> list[int]: return ids.tolist() def _connection(self, owner: int) -> socket.socket: - conns = self._connections.__dict__ + conns = self._connections.by_owner conn = conns.get(owner) if conn is None: conn = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) @@ -139,7 +155,7 @@ def _connection(self, owner: int) -> socket.socket: return conn def _drop(self, owner: int) -> None: - conn = self._connections.__dict__.pop(owner, None) + conn = self._connections.by_owner.pop(owner, None) if conn is not None: with self._lock: self._all_connections.discard(conn) @@ -152,7 +168,7 @@ class _Server(socketserver.ThreadingUnixStreamServer): # Unix-socket connect to a full backlog fails at once instead of waiting. request_queue_size = socket.SOMAXCONN - def __init__(self, path: str, cache: PrefixTokenCache, tokenizer: Any): + def __init__(self, path: str, cache: PrefixTokenCache, tokenizer: OffsetTokenizer): self.cache = cache self.tokenizer = tokenizer # Accepted connections and the threads serving them. diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml index ff302874d5d6..8c796023fd3a 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -406,6 +406,7 @@ l0_cpu: - unittest/llmapi/test_serve_context_only_settings.py - unittest/llmapi/test_serve_multi_frontend_helpers.py - unittest/llmapi/test_serve_multi_frontend_route_gate.py + - unittest/llmapi/test_serve_tokenization_owners.py - unittest/llmapi/test_serve_num_frontends_default.py - unittest/llmapi/test_serve_report_addr.py - unittest/llmapi/test_serve_sleep_mode.py diff --git a/tests/unittest/inputs/test_prefix_token_cache.py b/tests/unittest/inputs/test_prefix_token_cache.py index 1ba639f5c862..0621aed9dacc 100644 --- a/tests/unittest/inputs/test_prefix_token_cache.py +++ b/tests/unittest/inputs/test_prefix_token_cache.py @@ -172,7 +172,7 @@ def test_rewritten_history_is_identical_to_full_tokenization(tokenizer: Any) -> assert cache.divergent_hits == 0 -def test_a_split_inside_a_multi_token_character_is_not_used() -> None: +def test_a_split_is_moved_to_the_first_token_of_a_multi_token_character() -> None: tokenizer, cache = _ExpandingTokenizer(), _cache() text = _PREAMBLE for i in range(12): @@ -181,6 +181,9 @@ def test_a_split_inside_a_multi_token_character_is_not_used() -> None: assert cache.encode(tokenizer, text) == tokenizer(text)["input_ids"] rewritten = text[:-10] + "rewritten. " assert cache.encode(tokenizer, rewritten) == tokenizer(rewritten)["input_ids"] + # The splits are usable, not rejected at the seam. + assert cache.resync_failures == 0 + assert cache.divergent_hits > 0 def test_short_prompts_are_tokenized_but_not_cached() -> None: diff --git a/tests/unittest/llmapi/test_serve_tokenization_owners.py b/tests/unittest/llmapi/test_serve_tokenization_owners.py index ed0ac5c8d706..edfb86ab4ac6 100644 --- a/tests/unittest/llmapi/test_serve_tokenization_owners.py +++ b/tests/unittest/llmapi/test_serve_tokenization_owners.py @@ -18,6 +18,9 @@ on its own socket, as two trtllm-serve frontend processes do. """ +import os +import socket +import time from typing import Any import pytest @@ -27,7 +30,7 @@ PrefixTokenCacheConfig, tokenization_routing_key, ) -from tensorrt_llm.serve.tokenization_owners import TokenizationOwners, owner_of +from tensorrt_llm.serve.tokenization_owners import SOCKET_NAME, TokenizationOwners, owner_of pytestmark = pytest.mark.cpu_only @@ -101,7 +104,7 @@ def __call__(self, text: str, **kwargs: Any) -> dict[str, Any]: def test_a_prompt_is_tokenized_locally_when_its_owner_does_not_tokenize_it(tmp_path) -> None: tokenizer, cache = _CharTokenizer(), _cache() - owners = [TokenizationOwners(str(tmp_path), 0, 2, cache, tokenizer)] + owners = [TokenizationOwners(str(tmp_path), 0, 2, cache, tokenizer, retry_after_s=0.0)] try: prompt = _turns(1)[0] with tokenization_routing_key(_key_owned_by(1, 2)): @@ -120,6 +123,28 @@ def test_a_prompt_is_tokenized_locally_when_its_owner_does_not_tokenize_it(tmp_p owner.close() +def test_a_hung_owner_is_skipped_after_it_times_out(tmp_path) -> None: + tokenizer, cache = _CharTokenizer(), _cache() + # Frontend 1 accepts connections but never answers. + hung = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) + hung.bind(os.path.join(str(tmp_path), SOCKET_NAME.format(1))) + hung.listen() + owners = TokenizationOwners(str(tmp_path), 0, 2, cache, tokenizer, timeout_s=0.5) + try: + prompt = _turns(1)[0] + with tokenization_routing_key(_key_owned_by(1, 2)): + for _ in range(3): + start = time.monotonic() + assert cache.encode(tokenizer, prompt) == [ord(c) for c in prompt] + elapsed = time.monotonic() - start + # Only the first prompt waited for the owner. + assert elapsed < 0.25 + assert (cache.routed, cache.route_failures) == (0, 3) + finally: + owners.close() + hung.close() + + def test_a_conversation_id_with_a_lone_surrogate_has_an_owner() -> None: # JSON can carry one; it must not fail the request it arrives with. assert owner_of("conversation-\ud800", 2) in (0, 1) From f5c283a433f71c0bce7de17a49884088b96864c2 Mon Sep 17 00:00:00 2001 From: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com> Date: Mon, 9 Nov 2026 00:35:50 -0800 Subject: [PATCH 7/8] [None][fix] Address sunnyqgg's review of the prefix-tokenization cache changes - Reject divergent splices whose tail is shifted by a character: the tail must match the cached ids over the whole interval up to the next checkpoint, which now starts strictly inside the common prefix. Probe the tokenizer once at cache creation and disable the cache when a tail can be shifted (Prepend normalizers, e.g. Mixtral-8x7B-v0.1). - A prompt the tokenizer rejects no longer disables the cache. - Count routed prompts toward the periodic log line. - Try the divergent scan when an extension match leaves a long tail. - Share one connection per owner per frontend, bounding the routing mesh to 2*(K-1) sockets and K-1 threads per frontend; skip an owner only after a timeout, with a warning; turn any reply error into a route failure. - Close tokenization owners off the event loop with one shared deadline. - Key routing on the signed sub-agent affinity id on disagg context workers. - Validate TLLM_PREFIX_TOKEN_CACHE_ROUTE_BY_CONVERSATION (0/1), log when routing starts or cannot engage, and document it. - Share the conversation-to-frontend mapping and socket path helpers in multi_frontend.py; add get_prefix_token_cache() accessors and PrefixTokenCache.set_route(). - Docs: trtllm-serve starts 8 frontends by default; stable conversation ids; affinity ports; measurement context. - Tests: shifted-tail tokenizer, probe, deep edit, _common_prefix_len table, checkpoint validity, connection sharing, close deadline, switch parsing, start-up gating, serve.py wiring. Signed-off-by: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com> --- .../features/prefix-tokenization-cache.md | 86 +++-- tensorrt_llm/inputs/prefix_token_cache.py | 245 ++++++++++--- tensorrt_llm/inputs/registry.py | 33 +- tensorrt_llm/serve/multi_frontend.py | 20 ++ tensorrt_llm/serve/openai_client.py | 4 +- tensorrt_llm/serve/openai_server.py | 28 +- tensorrt_llm/serve/tokenization_owners.py | 212 ++++++++---- .../inputs/test_prefix_token_cache.py | 142 ++++++++ .../llmapi/test_serve_frontend_affinity.py | 8 + .../llmapi/test_serve_tokenization_owners.py | 322 +++++++++++++++--- 10 files changed, 869 insertions(+), 231 deletions(-) diff --git a/docs/source/features/prefix-tokenization-cache.md b/docs/source/features/prefix-tokenization-cache.md index ed734bf2a197..0bb65cd4272a 100644 --- a/docs/source/features/prefix-tokenization-cache.md +++ b/docs/source/features/prefix-tokenization-cache.md @@ -18,17 +18,26 @@ Agentic clients also rewrite the end of the history between turns, for example by trimming a tool result or replacing the last message, so the new prompt is no longer an extension of the cached one. The cache records a checkpoint every 256 tokens of a cached prompt. When a new prompt diverges from a cached one, the -cache reuses the cached ids up to the last checkpoint that still lies a full -checkpoint interval inside their longest common prefix, and tokenizes only the -rest. +cache reuses the cached ids up to the checkpoint one interval before the last +checkpoint inside their longest common prefix, and tokenizes only the rest. The +same is tried when a new prompt extends a much shorter cached prompt, such as +another conversation with the same opening, and the later split is used. Tokenizing a tail in isolation is not generally identical to tokenizing the whole prompt, because BPE merges can straddle the split. The cache therefore backs off a fixed number of tokens from the end of the cached prefix (for a diverging prompt, at least one checkpoint interval before the divergence), re-tokenizes from there, and accepts the splice only if the re-tokenized ids -equal the cached ids over that seam. Otherwise the prompt is tokenized in full. -The output is always identical to tokenizing the whole prompt. +equal the cached ids over that seam (for a diverging prompt, over the whole +interval up to the next checkpoint). Otherwise the prompt is tokenized in full. Every split is moved to +the first token of a character, so a character tokenized into several tokens is +never split. The output is always identical to tokenizing the whole prompt. + +Some tokenizers alter every text they tokenize; for example, a `tokenizer.json` +whose normalizer prepends `▁` (as for Phi-3 and Mixtral-8x7B-v0.1) adds a space +to every tail. Such a shift can pass the check at the seam, so the cache probes +the tokenizer when it is created and stays disabled, with a warning, if a tail +can be shifted. ## Enabling the cache @@ -54,10 +63,11 @@ The cache size can be tuned with environment variables: | `TLLM_PREFIX_TOKEN_CACHE_ENTRIES` | 512 | Maximum number of cached prompts. Eviction is least-recently-used. | | `TLLM_PREFIX_TOKEN_CACHE_MAX_CHARS` | 67108864 | Maximum total characters of cached prompt text. Cached ids are stored as int32, about one byte per character of English text, so this bounds host memory to roughly twice this many bytes. | | `TLLM_PREFIX_TOKEN_CACHE_MIN_CHARS` | 4096 | Prompts shorter than this are tokenized normally and never cached. | +| `TLLM_PREFIX_TOKEN_CACHE_ROUTE_BY_CONVERSATION` | 0 | Set to `1` to tokenize each conversation in one serving frontend; see [Multiple serving frontends](#multiple-serving-frontends). | -Each value must be a positive integer. A malformed value raises at server -startup so that misconfiguration fails loudly rather than silently disabling a -feature you turned on. +The sizing values must be positive integers, and the routing switch must be `0` +or `1`. A malformed value raises at server startup so that misconfiguration +fails loudly rather than silently disabling a feature you turned on. ## Scope and safety @@ -85,7 +95,7 @@ feature you turned on. - An unexpected exception inside the cache disables it, with a warning, and the request is tokenized normally. A cache problem never fails a request. - The cache logs its configuration at startup and, every 1000 cache-eligible - requests, an `INFO` line with the hit rate (and how many hits were on + requests (routed ones included), an `INFO` line with the hit rate (and how many hits were on diverging prompts), the number of re-synchronization fallbacks, how many prompts were routed to another frontend and how many were tokenized locally because routing failed (see below), and its occupancy, so a cache that is @@ -93,9 +103,11 @@ feature you turned on. ## Multiple serving frontends -Use a single serving frontend (`num_serve_frontends: 1`, the default) when you -enable the cache for multi-turn traffic. Multiple frontends are available only -with the classic IPC executor, so this section applies only there. +`trtllm-serve` starts 8 frontend processes by default. When you enable the +cache for multi-turn traffic, use a single serving frontend by passing +`--num_serve_frontends 1` (or setting `num_serve_frontends: 1` in the +`--config` YAML file). Multiple frontends are available only with the classic +IPC executor, so this section applies only there. The cache lives in the memory of one frontend process and is not shared between processes. With `num_serve_frontends: K`, the K frontend processes @@ -106,30 +118,46 @@ sends each turn of a conversation to an arbitrary frontend. That frontend usually has not seen the conversation's previous turn, so the prompt is tokenized in full. -On an agentic multi-turn trace with prompts of about 160k tokens on average, -at one session, the p90 time spent in the frontend was about 165 ms with 8 -frontends and 44 ms with one. Decode throughput per user was the same up to -32 concurrent sessions, because at that concurrency the frontend is not the -bottleneck. +On an agentic multi-turn trace with prompts of about 160k tokens on average +(GLM-5.3 and its tokenizer on GB300), at one session, the p90 time spent in the +frontend was about 165 ms with 8 frontends and 44 ms with one. Decode throughput +per user was the same from 1 to 32 concurrent sessions, the range measured, +because at that concurrency the frontend is not the bottleneck. Multiple frontends help when a single Python frontend cannot keep up, typically at high request rates with short prompts. That is not the workload this cache targets. -If you need several frontends, set +A client that can pick the frontend itself avoids the problem: with +`serve_frontend_affinity_port_stride` set, each frontend also listens on its own +port, and a disaggregated server sends every context request that carries a +conversation id to the port of the frontend that conversation maps to. + +Otherwise, if you need several frontends, set `TLLM_PREFIX_TOKEN_CACHE_ROUTE_BY_CONVERSATION=1` to tokenize each conversation in one frontend. A chat request's conversation id (the same id attention DP -routes by, from `conversation_params.conversation_id` or headers such as -`X-Correlation-ID`) is hashed to an owner frontend, and the frontend the -request arrived at sends the prompt text to the owner over a local Unix socket -and gets back the token ids; the request itself stays where it arrived. A -request without a conversation id is tokenized locally, and so is one whose -owner does not answer within 5 seconds or cannot tokenize the prompt; the -latter are counted as route failures in the periodic log line. An owner that -did not answer is not asked again for 30 seconds. The switch is -off by default and has no effect with a single frontend or with the cache -disabled. On the trace above with 8 frontends, it brought the p90 frontend time -at one session from about 165 ms to about 58 ms. +routes by: from `conversation_params.conversation_id` or headers such as +`X-Correlation-ID`, or, on a disaggregated context server, the sub-agent parent +id forwarded by the disaggregated server) is mapped to an owner frontend, the +same one the affinity ports use, and the frontend the request arrived at sends +the prompt text to the owner over a local Unix socket and gets back the token +ids; the request itself stays where it arrived. Routing helps only when the id +is stable across a conversation's turns: with an id that is fresh on every +request, as many gateways set `X-Correlation-ID`, it only adds a hop per turn. + +A request without a conversation id is tokenized locally. So is one whose owner +is not serving, does not answer within 5 seconds, or cannot tokenize the +prompt; these are counted as route failures in the periodic log line. An owner +that times out is logged with a warning and not asked again for 30 seconds. +Each frontend keeps one connection to every other frontend, shared by its +tokenizing threads, so with K frontends each one holds 2 × (K − 1) extra Unix +sockets and K − 1 extra threads. + +The switch is off by default. When it is set but cannot take effect, with a +single frontend or with the cache disabled, the server logs a warning; when it +takes effect, each frontend logs an `INFO` line. On the trace above with 8 +frontends, it brought the p90 frontend time at one session from about 165 ms to +about 58 ms. ## Measured effect diff --git a/tensorrt_llm/inputs/prefix_token_cache.py b/tensorrt_llm/inputs/prefix_token_cache.py index e443a45158b5..dc2ac841db98 100644 --- a/tensorrt_llm/inputs/prefix_token_cache.py +++ b/tensorrt_llm/inputs/prefix_token_cache.py @@ -12,16 +12,20 @@ response. Agentic clients also rewrite the end of the history between turns (a trimmed tool result, a replaced last message), so a prompt that diverges from a cached -one before its re-tokenization point is spliced at the last checkpoint inside -their longest common prefix. +one before its re-tokenization point is spliced at a checkpoint of the cached +prompt: the last one that has another checkpoint after it inside their longest +common prefix. Tokenizing a tail in isolation is not generally identical to tokenizing the whole prompt, because BPE merges can straddle the split. The cache therefore backs off ``overlap`` tokens from the end of the cached prefix (or, for a diverging prompt, at least ``checkpoint_tokens`` tokens from the divergence), -re-tokenizes from there, and accepts the splice only if the first ``resync`` re-tokenized ids -equal the cached ids over the same span. Otherwise the prompt is tokenized in -full. +re-tokenizes from there, and accepts the splice only if the first ``resync`` +re-tokenized ids equal the cached ids over the same span (for a diverging +prompt, over the whole span up to the next checkpoint, which must also start a +re-tokenized token). Otherwise the prompt is tokenized in full. Tokenizers that alter +every text they are called on, such as those whose normalizer prepends a +``▁``, are probed for once when the cache is created and get no cache. Off by default; enable with the ``enable_tokenization_cache`` LLM argument. Each ``DefaultInputProcessor`` owns its own cache, so cached ids are never @@ -61,6 +65,9 @@ # candidate is only compared in full when the incoming text has that probe at the # same position. PROBE_CHARS = 256 +# Block size, in characters, of the coarse scan for the first differing +# character of two prompts. +_COMPARE_BLOCK_CHARS = 4096 # The conversation the prompt being tokenized belongs to, if the caller knows it. _ROUTING_KEY: ContextVar[str | None] = ContextVar("tokenization_routing_key", default=None) @@ -96,8 +103,8 @@ def __call__( class PrefixTokenCacheConfig: """Tunables for :class:`PrefixTokenCache`. - The sizing fields have env overrides (see ``ENV_VARS``); ``overlap`` and - ``resync`` are correctness internals and do not. + The sizing fields have env overrides (see ``ENV_VARS``); ``overlap``, + ``resync`` and ``checkpoint_tokens`` are correctness internals and do not. """ max_entries: int = 512 @@ -124,6 +131,12 @@ class PrefixTokenCacheConfig: "min_chars": "TLLM_PREFIX_TOKEN_CACHE_MIN_CHARS", } + def __post_init__(self) -> None: + if self.checkpoint_tokens <= self.resync: + raise ValueError( + f"checkpoint_tokens ({self.checkpoint_tokens}) must exceed resync ({self.resync})" + ) + @classmethod def from_env(cls) -> PrefixTokenCacheConfig: """Build a config from the environment. @@ -154,8 +167,8 @@ class _Entry: # Where a prompt that diverges from ``text`` can be re-tokenized from: the # first token of the character holding every ``checkpoint_tokens``-th token # (0, C, 2C, ...), and the start character of that token. - checkpoint_tokens: array # int32 - checkpoints: array # int32 + checkpoint_token_indices: array # int32 + checkpoint_chars: array # int32 class _Match(NamedTuple): @@ -164,18 +177,25 @@ class _Match(NamedTuple): split_token: int split_char: int extends: bool # the prompt keeps ``entry.text[: entry.split_char]`` + # For a diverging prompt, the token and character of the checkpoint after + # the split: the re-tokenized tail must match the cached ids up to there. + next_checkpoint: tuple[int, int] | None = None class PrefixTokenCache: """Splice a cached prefix's token ids with a freshly tokenized tail. Thread-safe. Eviction is least-recently-used, bounded by both entry count - and total cached characters. When a prompt reuses a cached entry, the new - prompt replaces that entry rather than sitting alongside it, so a - conversation costs one entry however many turns it has. + and total cached characters. A prompt that extends a cached entry replaces + that entry rather than sitting alongside it, so a conversation costs one + entry however many turns it has. A prompt that diverges from a cached entry + is added next to it, since the divergence may be a branch of a history that + continues. An unexpected error disables the cache, with a warning, and every prompt is - tokenized in full from then on: a cache bug must never fail a request. + tokenized in full from then on: a cache bug must never fail a request. A + prompt the tokenizer itself rejects fails as it would without the cache and + leaves the cache enabled. """ def __init__(self, config: PrefixTokenCacheConfig) -> None: @@ -194,40 +214,76 @@ def __init__(self, config: PrefixTokenCacheConfig) -> None: self.resync_failures = 0 # Hits on a cached prompt the new one diverges from (also in ``hits``). self.divergent_hits = 0 - # Tokenizes a prompt tagged with a routing key in the process that - # owns that conversation, or returns None to tokenize it here. Raises - # RouteFailed if the owner did not answer. Set by multi-frontend serving. - self.route: Callable[[str, str], list[int] | None] | None = None - # Prompts another process tokenized through ``route``. + # Cache-eligible prompts, routed or not; paces the periodic log line. + self.requests = 0 + # Prompts another process tokenized through the route. self.routed = 0 - # Prompts tokenized here because ``route`` failed. + # Prompts tokenized here because the route failed. self.route_failures = 0 + self._route: Callable[[str, str], list[int] | None] | None = None + + def set_route(self, route: Callable[[str, str], list[int] | None] | None) -> None: + """Hand prompts tagged with a routing key to ``route``, or stop if None. + + ``route(key, text)`` returns the ids of ``text`` as tokenized by the + process that caches conversation ``key``, or None to tokenize it here, + and raises :class:`RouteFailed` if that process did not answer; the + prompt is then tokenized here too. + """ + self._route = route def encode(self, tokenizer: OffsetTokenizer, text: str) -> list[int]: - """Return the token ids of ``text`` without special tokens.""" + """Return the token ids of ``text`` without special tokens. + + Raises whatever the tokenizer raises for a prompt it cannot tokenize. + """ if self.disabled or len(text) < self._config.min_chars: return self._tokenize(tokenizer, text)[0] - key = _ROUTING_KEY.get() - if self.route is not None and key is not None: - try: - ids = self.route(key, text) - except RouteFailed as e: - logger.debug(f"Tokenizing a prompt locally: {e}") - with self._lock: - self.route_failures += 1 - else: - if ids is not None: - with self._lock: - self.routed += 1 - return ids + try: + ids = self._routed(text) + if ids is None: + ids = self._encode_or_disable(tokenizer, text) + return ids + finally: + self._count_request() + + def _routed(self, text: str) -> list[int] | None: + """Ids from the process that owns ``text``'s conversation, if any.""" + route, key = self._route, _ROUTING_KEY.get() + if route is None or key is None: + return None + try: + ids = route(key, text) + except RouteFailed as e: + logger.debug(f"Tokenizing a prompt locally: {e}") + with self._lock: + self.route_failures += 1 + return None + if ids is not None: + with self._lock: + self.routed += 1 + return ids + + def _encode_or_disable(self, tokenizer: OffsetTokenizer, text: str) -> list[int]: try: return self._encode(tokenizer, text) except Exception as e: + # Raises, leaving the cache enabled, if the tokenizer rejects the + # prompt: that is the request's error, not the cache's. + ids = self._tokenize(tokenizer, text)[0] self.disabled = True logger.warning( f"Disabling the prefix token cache after an error: {e!r} ({self.summary()})" ) - return self._tokenize(tokenizer, text)[0] + return ids + + def _count_request(self) -> None: + with self._lock: + self.requests += 1 + report = self.requests % LOG_EVERY_REQUESTS == 0 + summary = self.summary() if report else None + if summary is not None: + logger.info(f"Prefix token cache: {summary}") def summary(self) -> str: """One-line hit-rate and occupancy summary. Caller need not hold the lock.""" @@ -246,7 +302,9 @@ def _encode(self, tokenizer: OffsetTokenizer, text: str) -> list[int]: if found is not None: entry = found.entry tail_ids, tail_offsets = self._tokenize(tokenizer, text[found.split_char :]) - if self._resynced(entry, found.split_token, tail_ids): + if self._resynced(entry, found.split_token, tail_ids) and self._aligned( + found, tail_ids, tail_offsets + ): ids = entry.ids[: found.split_token].tolist() + list(tail_ids) base_token, base_char = found.split_token, found.split_char if ids is None: @@ -275,7 +333,7 @@ def _encode(self, tokenizer: OffsetTokenizer, text: str) -> list[int]: if found is not None and found.extends: self._remove(found.eid) if split_token > 0 and tail_index >= 0: - checkpoint_tokens, checkpoints = self._checkpoints( + checkpoint_token_indices, checkpoint_chars = self._checkpoints( found, base_token, base_char, tail_offsets, len(ids) ) self._insert( @@ -283,13 +341,9 @@ def _encode(self, tokenizer: OffsetTokenizer, text: str) -> list[int]: ids, split_token, base_char + tail_offsets[tail_index][0], - checkpoint_tokens, - checkpoints, + checkpoint_token_indices, + checkpoint_chars, ) - report = (self.hits + self.misses) % LOG_EVERY_REQUESTS == 0 - summary = self.summary() if report else None - if summary is not None: - logger.info(f"Prefix token cache: {summary}") return ids @staticmethod @@ -311,8 +365,8 @@ def _checkpoints( step = self._config.checkpoint_tokens # Spliced ids before base_token are the cached entry's, at the same chars. num_kept = -(-base_token // step) if base_token else 0 - tokens = array("i", found.entry.checkpoint_tokens[:num_kept] if num_kept else ()) - chars = array("i", found.entry.checkpoints[:num_kept] if num_kept else ()) + tokens = array("i", found.entry.checkpoint_token_indices[:num_kept] if num_kept else ()) + chars = array("i", found.entry.checkpoint_chars[:num_kept] if num_kept else ()) for index in range(num_kept * step, num_ids, step): tail_index = _first_of_char(tail_offsets, index - base_token) tokens.append(base_token + tail_index) @@ -325,6 +379,29 @@ def _resynced(self, entry: _Entry, split_token: int, tail_ids: Sequence[int]) -> span = min(self._config.resync, len(tail_ids), len(cached_tail)) return span > 0 and list(tail_ids[:span]) == cached_tail[:span].tolist() + @staticmethod + def _aligned( + found: _Match, tail_ids: Sequence[int], tail_offsets: Sequence[tuple[int, int]] + ) -> bool: + """Whether a diverging prompt's tail matches the cached ids up to the next checkpoint. + + The resync span can lie inside a run of equal tokens, such as a long + run of spaces, where a tail shifted by a character still agrees with + the cached ids; the shift shows where the run ends, or as tokens that + start one character off. The next checkpoint starts inside the common + prefix, so an unshifted tail has the cached ids up to it and starts a + token there. + """ + if found.next_checkpoint is None: + return True + token, char = found.next_checkpoint + index = token - found.split_token + return ( + index < len(tail_ids) + and tail_offsets[index][0] == char - found.split_char + and list(tail_ids[:index]) == found.entry.ids[found.split_token : token].tolist() + ) + def _bucket_key(self, text: str) -> int: return hash(text[: self._config.min_chars]) @@ -332,10 +409,10 @@ def _lookup(self, text: str) -> _Match | None: """Best cached entry to splice ``text`` from, promoted to MRU if ``text`` extends it. The entry with the longest re-tokenization prefix that ``text`` keeps - wins. Without one, the - entry sharing the longest common prefix with ``text`` is used, split at - its last checkpoint that leaves a full checkpoint interval of common - text after the split. + wins. Without one, or when it leaves more than two checkpoint intervals + to re-tokenize, the entries ``text`` diverges from are also tried, each + split at the checkpoint one interval before its last checkpoint inside + the common prefix, and the latest split wins. """ with self._lock: candidates = list(self._buckets.get(self._bucket_key(text), {}).items()) @@ -352,7 +429,7 @@ def _lookup(self, text: str) -> _Match | None: and text.startswith(entry.text[:n]) ): best = _Match(eid, entry, entry.split_token, entry.split_char, True) - if best is None: + if best is None or self._leaves_long_tail(best, text): for eid, entry in candidates: match = self._divergent_match(eid, entry, text) if match is not None and (best is None or match.split_token > best.split_token): @@ -367,13 +444,27 @@ def _lookup(self, text: str) -> _Match | None: # splice from; _remove() on it is a no-op. return best + def _leaves_long_tail(self, match: _Match, text: str) -> bool: + """Whether ``text`` extends ``match`` by more than two checkpoint intervals.""" + entry = match.entry + chars_per_token = entry.split_char / max(entry.split_token, 1) + tail_tokens = (len(text) - match.split_char) / max(chars_per_token, 1e-6) + return tail_tokens > 2 * self._config.checkpoint_tokens + def _divergent_match(self, eid: int, entry: _Entry, text: str) -> _Match | None: common = _common_prefix_len(entry.text, text) # The checkpoint after the split must start inside the common prefix. - index = bisect.bisect_right(entry.checkpoints, common) - 2 + index = bisect.bisect_left(entry.checkpoint_chars, common) - 2 if index < 1: return None - return _Match(eid, entry, entry.checkpoint_tokens[index], entry.checkpoints[index], False) + return _Match( + eid, + entry, + entry.checkpoint_token_indices[index], + entry.checkpoint_chars[index], + False, + (entry.checkpoint_token_indices[index + 1], entry.checkpoint_chars[index + 1]), + ) def _insert( self, @@ -381,8 +472,8 @@ def _insert( ids: Sequence[int], split_token: int, split_char: int, - checkpoint_tokens: array, - checkpoints: array, + checkpoint_token_indices: array, + checkpoint_chars: array, ) -> None: """Caller holds the lock.""" key = self._bucket_key(text) @@ -392,7 +483,13 @@ def _insert( eid, self._next_id = self._next_id, self._next_id + 1 probe = text[max(split_char - PROBE_CHARS, 0) : split_char] entry = _Entry( - text, array("i", ids), split_token, split_char, probe, checkpoint_tokens, checkpoints + text, + array("i", ids), + split_token, + split_char, + probe, + checkpoint_token_indices, + checkpoint_chars, ) self._entries[eid] = entry bucket[eid] = entry @@ -418,7 +515,7 @@ def _remove(self, eid: int) -> None: def _common_prefix_len(a: str, b: str) -> int: """Length of the longest common prefix, compared in blocks.""" - n, i, block = min(len(a), len(b)), 0, 4096 + n, i, block = min(len(a), len(b)), 0, _COMPARE_BLOCK_CHARS while i < n and a[i : i + block] == b[i : i + block]: i += block # The first mismatch, if any, lies in the block at i; halve it to find it. @@ -444,10 +541,31 @@ def _first_of_char(offsets: Sequence[tuple[int, int]], index: int) -> int: return index +def _splices_wrongly(tokenizer: OffsetTokenizer, config: PrefixTokenCacheConfig) -> bool: + """Whether a tail the cache would accept at the seam can differ from full tokenization. + + Tokenizers that alter every text they are called on, such as one whose + normalizer prepends ``▁``, shift a tail by a character. Inside a run of + equal tokens the shift passes the resync check and only shows at the end + of the run, so the probe splits inside a long run of spaces. + """ + text = "x" + " " * 4096 + "y" + ids, offsets = PrefixTokenCache._tokenize(tokenizer, text) + split = next((i for i, (start, _) in enumerate(offsets) if start >= 1024), None) + if split is None: + return False + split = _first_of_char(offsets, split) + tail_ids = PrefixTokenCache._tokenize(tokenizer, text[offsets[split][0] :])[0] + span = min(config.resync, len(tail_ids), len(ids) - split) + resynced = list(tail_ids[:span]) == list(ids[split : split + span]) + return resynced and list(ids[:split]) + list(tail_ids) != list(ids) + + def create_prefix_token_cache(tokenizer: object) -> PrefixTokenCache | None: """Return a cache for ``tokenizer`` if the tokenizer can support one. - Returns None when there is no tokenizer or it cannot report offsets. + Returns None when there is no tokenizer, it cannot report offsets, or a + spliced tail can differ from tokenizing the whole prompt. Raises ``ValueError`` for a malformed env override. """ if tokenizer is None: @@ -460,5 +578,20 @@ def create_prefix_token_cache(tokenizer: object) -> PrefixTokenCache | None: ) return None config = PrefixTokenCacheConfig.from_env() + try: + unsafe = _splices_wrongly(tokenizer, config) + except Exception as e: + logger.warning( + f"Could not probe the tokenizer for the prefix token cache ({e!r}); disabling it." + ) + return None + if unsafe: + logger.warning( + "enable_tokenization_cache is set but the tokenizer changes the start " + "of every text it tokenizes (for example, its normalizer prepends a " + "space marker), so spliced token ids could differ from tokenizing " + "the whole prompt; disabling the prefix token cache." + ) + return None logger.info(f"Prefix token cache enabled: {config}") return PrefixTokenCache(config) diff --git a/tensorrt_llm/inputs/registry.py b/tensorrt_llm/inputs/registry.py index 4c88f2006335..83965750fc91 100644 --- a/tensorrt_llm/inputs/registry.py +++ b/tensorrt_llm/inputs/registry.py @@ -52,7 +52,8 @@ validate_mm_inputs) # yapf: enable from .multimodal_data import MEDIA_MODALITIES, serialize_item -from .prefix_token_cache import PrefixTokenCache, create_prefix_token_cache +from .prefix_token_cache import (OffsetTokenizer, PrefixTokenCache, + create_prefix_token_cache) from .tokenization import tokenize_prompt N = TypeVar("N", bound=Type[nn.Module]) @@ -136,11 +137,17 @@ def __init__(self, self.config = config self.model_path = model_path self.multimodal_hashing_supported = None - self._prefix_token_cache_tokenizer = tokenizer # None when disabled or the tokenizer cannot support the cache. self._prefix_token_cache = (create_prefix_token_cache(tokenizer) if enable_tokenization_cache else None) + def get_prefix_token_cache( + self) -> Optional[Tuple[PrefixTokenCache, OffsetTokenizer]]: + """This processor's prefix-tokenization cache and the tokenizer it caches, if enabled.""" + if self._prefix_token_cache is None: + return None + return self._prefix_token_cache, self.tokenizer + def __call__( self, inputs: TextPrompt, sampling_params: SamplingParams ) -> Tuple[List[int], Optional[ExtraProcessedInputs]]: @@ -155,12 +162,15 @@ def __call__( def prefix_token_cache_of( - input_processor: Any) -> Optional[Tuple[PrefixTokenCache, Any]]: - """The prefix-tokenization cache of ``input_processor`` and its tokenizer, if enabled.""" - cache = getattr(input_processor, "_prefix_token_cache", None) - if cache is None: - return None - return cache, input_processor._prefix_token_cache_tokenizer + input_processor: object +) -> Optional[Tuple[PrefixTokenCache, OffsetTokenizer]]: + """The prefix-tokenization cache of ``input_processor`` and its tokenizer. + + None when the processor has no cache: it is disabled, or the processor + does not support one. + """ + get_cache = getattr(input_processor, "get_prefix_token_cache", None) + return None if get_cache is None else get_cache() class BaseMultimodalInputProcessor(ABC): @@ -282,6 +292,13 @@ def _init_tokenization_cache(self, enable_tokenization_cache: bool, self._prefix_token_cache = (create_prefix_token_cache(tokenizer) if enable_tokenization_cache else None) + def get_prefix_token_cache( + self) -> Optional[Tuple[PrefixTokenCache, OffsetTokenizer]]: + """This processor's prefix-tokenization cache and the tokenizer it caches, if enabled.""" + if self._prefix_token_cache is None: + return None + return self._prefix_token_cache, self._prefix_token_cache_tokenizer + def _encode_with_tokenization_cache( self, prompt: str, diff --git a/tensorrt_llm/serve/multi_frontend.py b/tensorrt_llm/serve/multi_frontend.py index bf63266c744c..49766455a570 100644 --- a/tensorrt_llm/serve/multi_frontend.py +++ b/tensorrt_llm/serve/multi_frontend.py @@ -26,6 +26,7 @@ import math import os import tempfile +import zlib from dataclasses import dataclass from typing import Callable, Optional @@ -38,6 +39,9 @@ LAUNCHER_UDS_NAME = "launcher.sock" LAUNCHER_UDS_ENV = "TLLM_SERVE_LAUNCHER_UDS" CLOCK_OFFSET_NAME = "clock_offset" +# Every frontend serves its prefix-tokenization cache on this Unix socket when +# tokenization is routed by conversation. +_TOKENIZATION_UDS_NAME = "tokenize-{}.sock" # Launcher-owned routes. Attached frontends register them ahead of the local # handlers, which they shadow since FastAPI matches in registration order. @@ -93,6 +97,22 @@ def is_attached(self) -> bool: return not self.is_launcher +def frontend_of_conversation(conversation_id: str, num_frontends: int) -> int: + """The frontend that serves conversation ``conversation_id``, the same in every process. + + A disaggregated server sends the conversation's context requests to this + frontend's affinity port, and with tokenization routing this frontend + tokenizes its prompts, so both must use this mapping. + """ + # Lone surrogates, which JSON can carry, must not fail the request. + return zlib.crc32(conversation_id.encode(errors="surrogatepass")) % num_frontends + + +def tokenization_uds_path(ipc_dir: str, frontend_id: int) -> str: + """The Unix socket frontend ``frontend_id`` serves its tokenization cache on.""" + return os.path.join(ipc_dir, _TOKENIZATION_UDS_NAME.format(frontend_id)) + + class LauncherForwarder: """Forward HTTP requests from an attached frontend to the launcher.""" diff --git a/tensorrt_llm/serve/openai_client.py b/tensorrt_llm/serve/openai_client.py index 1bf03ddebbef..542550e186a0 100644 --- a/tensorrt_llm/serve/openai_client.py +++ b/tensorrt_llm/serve/openai_client.py @@ -18,7 +18,6 @@ import json import os import traceback -import zlib from abc import ABC, abstractmethod from contextlib import aclosing from typing import ( @@ -49,6 +48,7 @@ build_subagent_affinity_headers, request_requires_internal_disagg_auth, ) +from tensorrt_llm.serve.multi_frontend import frontend_of_conversation from tensorrt_llm.serve.openai_protocol import ( ChatCompletionRequest, ChatCompletionResponse, @@ -309,7 +309,7 @@ def _url_server(self, server: str, request: UCompletionRequest) -> str: return server num_frontends, stride = self._frontend_affinity host, _, port = server.rpartition(":") - frontend = zlib.crc32(conversation_id.encode()) % num_frontends + frontend = frontend_of_conversation(conversation_id, num_frontends) affinity_server = f"{host}:{int(port) + stride * (frontend + 1)}" refused_at = self._refused_affinity_ports.get(affinity_server) if refused_at is not None: diff --git a/tensorrt_llm/serve/openai_server.py b/tensorrt_llm/serve/openai_server.py index b6c30daa4f19..d83c8865aad9 100644 --- a/tensorrt_llm/serve/openai_server.py +++ b/tensorrt_llm/serve/openai_server.py @@ -52,8 +52,7 @@ log_allowed_local_media_path) from tensorrt_llm.inputs.multimodal import MultimodalServerConfig from tensorrt_llm.inputs.prefix_token_cache import tokenization_routing_key -from tensorrt_llm.inputs.registry import (BaseMultimodalInputProcessor, - prefix_token_cache_of) +from tensorrt_llm.inputs.registry import BaseMultimodalInputProcessor from tensorrt_llm.inputs.utils import ConversationMessage from tensorrt_llm.llmapi import DisaggregatedParams as LlmDisaggregatedParams from tensorrt_llm.llmapi import MultimodalEncoder, SchedulingParams, tracing @@ -88,7 +87,7 @@ from tensorrt_llm.serve.chat_utils import (load_chat_template, parse_chat_messages_coroutines) from tensorrt_llm.serve.cluster_storage import create_cluster_storage_client -from tensorrt_llm.serve.conversation_id import (get_request_routing_id, +from tensorrt_llm.serve.conversation_id import (get_request_conversation_id, resolve_request_conversation_id) from tensorrt_llm.serve.disagg_auth import ( request_requires_internal_disagg_auth, validate_internal_disagg_request, @@ -149,7 +148,7 @@ RuntimeControlReplayCache, validate_runtime_control_request) from tensorrt_llm.serve.serving_extensions import load_builtin_extensions from tensorrt_llm.serve.tokenization_owners import (TokenizationOwners, - routing_enabled) + start_tokenization_owners) from tensorrt_llm.serve.tool_parser.tool_parser_factory import ToolParserFactory from tensorrt_llm.serve.visual_gen_metrics import ( build_visual_gen_server_timings, build_visual_gen_timing_headers) @@ -767,16 +766,10 @@ def __init__( if multi_frontend_serving is not None and multi_frontend_serving.is_attached: self._launcher_forwarder = LauncherForwarder( multi_frontend_serving.launcher_uds) - self._tokenization_owners: Optional[TokenizationOwners] = None - prefix_cache = prefix_token_cache_of( - getattr(generator, "input_processor", None)) - if (routing_enabled() and multi_frontend_serving is not None - and multi_frontend_serving.num_frontends > 1 - and prefix_cache is not None): - self._tokenization_owners = TokenizationOwners( - os.path.dirname(multi_frontend_serving.launcher_uds), - multi_frontend_serving.frontend_id, - multi_frontend_serving.num_frontends, *prefix_cache) + self._tokenization_owners: Optional[ + TokenizationOwners] = start_tokenization_owners( + multi_frontend_serving, + getattr(generator, "input_processor", None)) self._shared_clock_offset: Optional[SharedClockOffset] = ( SharedClockOffset( os.path.dirname(multi_frontend_serving.launcher_uds)) @@ -890,7 +883,7 @@ async def lifespan(app: FastAPI): if self._launcher_forwarder is not None: await self._launcher_forwarder.close() if self._tokenization_owners is not None: - self._tokenization_owners.close() + await asyncio.to_thread(self._tokenization_owners.close) if self.embedding_batcher is not None: await self.embedding_batcher.shutdown() logger.info("Stopped encode dynamic batcher") @@ -2510,9 +2503,12 @@ async def chat_stream_generator( sampling_params, disaggregated_params) if self._tokenization_owners is not None: + # A disaggregated context worker gets the sub-agent's + # parent id only from the signed header. preprocess = functools.partial( _preprocess_in_conversation, - get_request_routing_id(request), preprocess) + scheduling_params.subagent_affinity_id + or get_request_conversation_id(request), preprocess) loop = asyncio.get_event_loop() generate_inputs = await loop.run_in_executor( self._input_proc_executor, preprocess) diff --git a/tensorrt_llm/serve/tokenization_owners.py b/tensorrt_llm/serve/tokenization_owners.py index 3f07be0bed84..e2c1c4870d03 100644 --- a/tensorrt_llm/serve/tokenization_owners.py +++ b/tensorrt_llm/serve/tokenization_owners.py @@ -18,9 +18,13 @@ frontend, so consecutive turns of a conversation land in different processes, each with its own prefix-tokenization cache. Every frontend therefore serves its cache on a Unix socket in the multi-frontend ipc dir, and a prompt whose -conversation id hashes to another frontend is tokenized there, by the cache -that has seen the conversation's earlier turns. Only the prompt text and the -token ids cross the socket; the request itself stays where it arrived. +conversation id maps to another frontend is tokenized there, by the cache that +has seen the conversation's earlier turns. Only the prompt text and the token +ids cross the socket; the request itself stays where it arrived. + +Each frontend keeps one connection to every other frontend, shared by its +tokenizing threads, so a server with K frontends holds 2 * (K - 1) routing +sockets and K - 1 serving threads per frontend. """ import contextlib @@ -30,14 +34,18 @@ import struct import threading import time -import zlib from array import array from typing import Optional from tensorrt_llm.inputs.prefix_token_cache import OffsetTokenizer, PrefixTokenCache, RouteFailed +from tensorrt_llm.inputs.registry import prefix_token_cache_of from tensorrt_llm.logger import logger +from tensorrt_llm.serve.multi_frontend import ( + MultiFrontendServing, + frontend_of_conversation, + tokenization_uds_path, +) -SOCKET_NAME = "tokenize-{}.sock" # Set to 1 to route tokenization by conversation id across serving frontends. ENABLE_ENV_VAR = "TLLM_PREFIX_TOKEN_CACHE_ROUTE_BY_CONVERSATION" @@ -45,34 +53,74 @@ # Sent in place of a token count when the owner could not tokenize the prompt. _FAILED = -1 # Lone surrogates, which JSON can carry, survive the trip and fail in the -# tokenizer as they would locally, and do not fail a conversation id's hash. +# tokenizer as they would locally. _TEXT_ERRORS = "surrogatepass" def routing_enabled() -> bool: - """Whether tokenization is routed by conversation id (off unless enabled).""" - return os.environ.get(ENABLE_ENV_VAR) == "1" + """Whether tokenization is routed by conversation id (off unless enabled). + Raises: + ValueError: if the switch is set to anything but ``0`` or ``1``. + """ + raw = os.environ.get(ENABLE_ENV_VAR) + if raw is None or raw == "0": + return False + if raw == "1": + return True + raise ValueError(f"{ENABLE_ENV_VAR} must be 0 or 1, got {raw!r}") -def owner_of(key: str, num_frontends: int) -> int: - """The frontend that tokenizes conversation ``key``, the same in every process.""" - return zlib.crc32(key.encode(errors=_TEXT_ERRORS)) % num_frontends +def start_tokenization_owners( + multi_frontend_serving: Optional[MultiFrontendServing], input_processor: object +) -> Optional["TokenizationOwners"]: + """Route this frontend's prompts by conversation if the switch is on and can engage. -class _Connections(threading.local): - """This thread's connection to each owner.""" + Raises: + ValueError: if the switch is malformed. + """ + if not routing_enabled(): + return None + num_frontends = 1 if multi_frontend_serving is None else multi_frontend_serving.num_frontends + prefix_cache = prefix_token_cache_of(input_processor) + if num_frontends <= 1 or prefix_cache is None: + logger.warning( + f"{ENABLE_ENV_VAR}=1 has no effect: it needs num_serve_frontends > 1 " + f"(got {num_frontends}) and an enabled prefix token cache " + f"(enable_tokenization_cache: {prefix_cache is not None})." + ) + return None + logger.info( + f"Routing tokenization by conversation id: frontend " + f"{multi_frontend_serving.frontend_id} of {num_frontends}." + ) + return TokenizationOwners( + os.path.dirname(multi_frontend_serving.launcher_uds), + multi_frontend_serving.frontend_id, + num_frontends, + *prefix_cache, + ) + + +class _Peer: + """This frontend's connection to one owner, shared by its tokenizing threads.""" def __init__(self) -> None: - self.by_owner: dict[int, socket.socket] = {} + # Held for a whole exchange, which keeps requests and replies paired. + self.lock = threading.Lock() + self.conn: Optional[socket.socket] = None + # Monotonic time before which the owner, which timed out, is not asked. + self.skip_until = 0.0 class TokenizationOwners: """Route ``cache``'s prompts to the frontend that owns their conversation. - Serves this frontend's ``cache`` to the other frontends and installs - ``cache.route``. A prompt the owner does not tokenize within ``timeout_s``, - for example because it has not started, has died or has hung, is tokenized - locally, and so is every prompt it owns for the next ``retry_after_s``. + Serves this frontend's ``cache`` to the other frontends and installs its + route. A prompt the owner does not tokenize within ``timeout_s``, for + example because it has hung, is tokenized locally, and so is every prompt + it owns for the next ``retry_after_s``. A prompt whose owner is not serving + (not started yet, or dead) or cannot tokenize it is tokenized locally too. """ def __init__( @@ -84,91 +132,107 @@ def __init__( tokenizer: OffsetTokenizer, timeout_s: float = 5.0, retry_after_s: float = 30.0, - ): + ) -> None: self._ipc_dir = ipc_dir self._frontend_id = frontend_id self._num_frontends = num_frontends + self._cache = cache self._timeout_s = timeout_s self._retry_after_s = retry_after_s - # Monotonic time before which an owner that did not answer is not asked. - self._skip_until: dict[int, float] = {} - # One connection per owner per tokenizing thread. - self._connections = _Connections() - self._all_connections: set[socket.socket] = set() - self._lock = threading.Lock() - - path = os.path.join(ipc_dir, SOCKET_NAME.format(frontend_id)) - self._server = _Server(path, cache, tokenizer) + self._peers = {owner: _Peer() for owner in range(num_frontends) if owner != frontend_id} + + self._server = _Server(tokenization_uds_path(ipc_dir, frontend_id), cache, tokenizer) threading.Thread( target=self._server.serve_forever, name="trtllm_tokenize_owner", daemon=True ).start() - cache.route = self._route + cache.set_route(self._route) + + def close(self, timeout_s: float = 5.0) -> None: + """Stop routing and serving, and close every connection in both directions. - def close(self) -> None: - """Stop serving and close every connection, in both directions.""" + Blocks for up to about ``timeout_s`` while handlers finish the prompt + they are tokenizing; those still running then are left to exit with + the process. + """ + deadline = time.monotonic() + timeout_s + self._cache.set_route(None) self._server.shutdown() self._server.server_close() - self._server.close_connections() - with self._lock: - for conn in self._all_connections: + self._server.close_connections(deadline) + for peer in self._peers.values(): + conn = peer.conn + if conn is not None: + # Wakes a thread waiting on the owner, which then tokenizes + # its prompt locally. + with contextlib.suppress(OSError): + conn.shutdown(socket.SHUT_RDWR) conn.close() - self._all_connections.clear() def _route(self, key: str, text: str) -> Optional[list[int]]: - owner = owner_of(key, self._num_frontends) + owner = frontend_of_conversation(key, self._num_frontends) if owner == self._frontend_id: return None - if time.monotonic() < self._skip_until.get(owner, 0.0): - raise RouteFailed(f"frontend {owner} did not answer recently") + peer = self._peers[owner] + if time.monotonic() < peer.skip_until: + raise RouteFailed(f"frontend {owner} timed out recently") + if not peer.lock.acquire(timeout=self._timeout_s): + raise RouteFailed(f"frontend {owner} is busy") try: - return self._request(owner, text) - except OSError as e: - self._skip_until[owner] = time.monotonic() + self._retry_after_s - self._drop(owner) + return self._request(owner, peer, text) + except RouteFailed: + raise + except TimeoutError as e: + _drop(peer) + peer.skip_until = time.monotonic() + self._retry_after_s + logger.warning( + f"Frontend {owner} did not tokenize a prompt within {self._timeout_s} s; " + f"tokenizing its conversations locally for {self._retry_after_s} s." + ) + raise RouteFailed(f"frontend {owner} timed out") from e + except Exception as e: + # Not serving (connect failed), died mid-exchange, or sent a + # malformed reply: the connection cannot be reused. + _drop(peer) raise RouteFailed(f"frontend {owner} did not answer: {e!r}") from e + finally: + peer.lock.release() - def _request(self, owner: int, text: str) -> list[int]: - conn = self._connection(owner) - data = text.encode(errors=_TEXT_ERRORS) - conn.sendall(_LENGTH.pack(len(data)) + data) - (count,) = _LENGTH.unpack(_recv_exactly(conn, _LENGTH.size)) - if count == _FAILED: - raise RouteFailed(f"frontend {owner} could not tokenize the prompt") - ids = array("i") - ids.frombytes(_recv_exactly(conn, count * ids.itemsize)) - return ids.tolist() - - def _connection(self, owner: int) -> socket.socket: - conns = self._connections.by_owner - conn = conns.get(owner) - if conn is None: + def _request(self, owner: int, peer: _Peer, text: str) -> list[int]: + if peer.conn is None: conn = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) conn.settimeout(self._timeout_s) try: - conn.connect(os.path.join(self._ipc_dir, SOCKET_NAME.format(owner))) + conn.connect(tokenization_uds_path(self._ipc_dir, owner)) except OSError: conn.close() raise - conns[owner] = conn - with self._lock: - self._all_connections.add(conn) - return conn + peer.conn = conn + data = text.encode(errors=_TEXT_ERRORS) + peer.conn.sendall(_LENGTH.pack(len(data)) + data) + (count,) = _LENGTH.unpack(_recv_exactly(peer.conn, _LENGTH.size)) + if count == _FAILED: + raise RouteFailed(f"frontend {owner} could not tokenize the prompt") + if count < 0: + raise ValueError(f"invalid token count {count}") + ids = array("i") + ids.frombytes(_recv_exactly(peer.conn, count * ids.itemsize)) + return ids.tolist() + - def _drop(self, owner: int) -> None: - conn = self._connections.by_owner.pop(owner, None) - if conn is not None: - with self._lock: - self._all_connections.discard(conn) - conn.close() +def _drop(peer: _Peer) -> None: + """Caller holds ``peer.lock``.""" + if peer.conn is not None: + peer.conn.close() + peer.conn = None class _Server(socketserver.ThreadingUnixStreamServer): daemon_threads = True - # Every tokenizing thread of every other frontend may connect at once, and a - # Unix-socket connect to a full backlog fails at once instead of waiting. + # Every other frontend may connect at once, and a Unix-socket connect to a + # full backlog fails at once instead of waiting. request_queue_size = socket.SOMAXCONN - def __init__(self, path: str, cache: PrefixTokenCache, tokenizer: OffsetTokenizer): + def __init__(self, path: str, cache: PrefixTokenCache, tokenizer: OffsetTokenizer) -> None: self.cache = cache self.tokenizer = tokenizer # Accepted connections and the threads serving them. @@ -178,15 +242,15 @@ def __init__(self, path: str, cache: PrefixTokenCache, tokenizer: OffsetTokenize os.unlink(path) super().__init__(path, _Handler) - def close_connections(self, timeout_s: float = 5.0) -> None: - """Wake every handler blocked on its peer and wait for it to exit.""" + def close_connections(self, deadline: float) -> None: + """Wake every handler blocked on its peer and wait for it to exit until ``deadline``.""" with self.handlers_lock: handlers = list(self.handlers.items()) for conn, _ in handlers: with contextlib.suppress(OSError): conn.shutdown(socket.SHUT_RDWR) for _, thread in handlers: - thread.join(timeout_s) + thread.join(max(0.0, deadline - time.monotonic())) class _Handler(socketserver.BaseRequestHandler): diff --git a/tests/unittest/inputs/test_prefix_token_cache.py b/tests/unittest/inputs/test_prefix_token_cache.py index 0621aed9dacc..fc8550d71e43 100644 --- a/tests/unittest/inputs/test_prefix_token_cache.py +++ b/tests/unittest/inputs/test_prefix_token_cache.py @@ -20,6 +20,7 @@ from tensorrt_llm.inputs.prefix_token_cache import ( PrefixTokenCache, PrefixTokenCacheConfig, + _common_prefix_len, create_prefix_token_cache, ) from tensorrt_llm.inputs.registry import DefaultInputProcessor @@ -94,6 +95,40 @@ def __call__(self, text: str, **kwargs: Any) -> dict[str, Any]: return {"input_ids": ids, "offset_mapping": offsets} +class _PrependingTokenizer: + """Prepends a space to every text, as a ``Prepend("▁")`` normalizer does. + + Runs of spaces become tokens of up to four spaces, counted from the start + of the run, so a tail split inside a run is shifted by the prepended space + and still agrees with the cached ids until the end of the run. + """ + + is_fast = True + + def __call__(self, text: str, **kwargs: Any) -> dict[str, Any]: + normalized = " " + text + ids, offsets, i = [], [], 0 + while i < len(normalized): + j = i + 1 + if normalized[i] == " ": + while j < len(normalized) and j - i < 4 and normalized[j] == " ": + j += 1 + ids.append(_tok_id(normalized[i:j])) + # The prepended space maps to the start of the text. + offsets.append((max(i - 1, 0), j - 1)) + i = j + return {"input_ids": ids, "offset_mapping": offsets} + + +class _RejectingTokenizer(_MergeTokenizer): + """Rejects a lone surrogate, as HF tokenizers do.""" + + def __call__(self, text: str, **kwargs: Any) -> dict[str, Any]: + if "\ud800" in text: + raise TypeError("cannot tokenize a lone surrogate") + return super().__call__(text, **kwargs) + + _MIN_CHARS = 16 # A shared opening longer than min_chars, so growing prompts land in the same # bucket and the prefix lookup can find them. @@ -116,6 +151,23 @@ def _turns(n: int = 12, preamble: str = _PREAMBLE) -> list[str]: return out +def _assert_checkpoints_are_valid(cache: PrefixTokenCache, tokenizer: Any) -> None: + """Checkpoint i is the first token of the character holding token i * step.""" + step = cache._config.checkpoint_tokens + for entry in cache._entries.values(): + enc = tokenizer(entry.text) + assert entry.ids.tolist() == enc["input_ids"] + starts = [start for start, _ in enc["offset_mapping"]] + tokens = entry.checkpoint_token_indices.tolist() + chars = entry.checkpoint_chars.tolist() + assert len(tokens) == len(chars) == -(-len(starts) // step) + assert tokens == sorted(tokens) and chars == sorted(chars) + for i, (token, char) in enumerate(zip(tokens, chars)): + assert token <= i * step + assert starts[token] == starts[i * step] == char + assert token == 0 or starts[token - 1] != char + + def _conversations(n: int) -> list[str]: """First turns of ``n`` unrelated conversations.""" return [_turns(1, preamble=f"conversation {i:03d}: {_PREAMBLE}")[0] for i in range(n)] @@ -170,6 +222,78 @@ def test_rewritten_history_is_identical_to_full_tokenization(tokenizer: Any) -> assert any(entry.text == turns[-1] for entry in cache._entries.values()) else: assert cache.divergent_hits == 0 + assert not cache.disabled + _assert_checkpoints_are_valid(cache, tokenizer) + + +def test_a_deep_edit_splices_from_a_late_checkpoint() -> None: + """A long history edited deep inside is spliced across several compare blocks.""" + tokenizer, cache = _MergeTokenizer(), _cache() + history = _turns(400)[-1][:10_000] + assert cache.encode(tokenizer, history) == tokenizer(history)["input_ids"] + edited = history[:7000] + "#" + history[7001:] + found = cache._lookup(edited) + assert found is not None and not found.extends + # One checkpoint interval (8 tokens of at most 2 chars) before the edit, at most. + assert 7000 - 3 * 2 * 8 <= found.split_char < 7000 + assert cache.encode(tokenizer, edited) == tokenizer(edited)["input_ids"] + assert cache.divergent_hits == 1 and cache.resync_failures == 0 + assert not cache.disabled + _assert_checkpoints_are_valid(cache, tokenizer) + + +def test_a_rewritten_turn_splices_from_its_history_not_a_shorter_shared_opening() -> None: + tokenizer, cache = _MergeTokenizer(), _cache() + opening = _PREAMBLE + history = _turns(20)[-1] + # Another conversation that so far holds only the shared opening. + for prompt in (history, opening): + assert cache.encode(tokenizer, prompt) == tokenizer(prompt)["input_ids"] + rewritten = history[:-20] + "a rewritten ending. " + # The rewritten turn extends the opening's entry too. + assert any(entry.text == opening for entry in cache._entries.values()) + found = cache._lookup(rewritten) + assert found is not None and found.entry.text == history and not found.extends + divergent_hits = cache.divergent_hits + assert cache.encode(tokenizer, rewritten) == tokenizer(rewritten)["input_ids"] + assert cache.divergent_hits == divergent_hits + 1 + + +def test_a_tail_shifted_inside_a_run_of_spaces_is_rejected() -> None: + """A shift the resync span cannot see is caught at the next checkpoint.""" + tokenizer, cache = _PrependingTokenizer(), _cache() + run_end = len(_PREAMBLE) + 400 + history = _PREAMBLE + " " * 400 + "end of the padded tool output. " * 4 + for kept in range(12): + cache = _cache() + assert cache.encode(tokenizer, history) == tokenizer(history)["input_ids"] + # Diverge just after the run, so the split lands inside it. + rewritten = history[: run_end + kept] + "#rewritten" + found = cache._lookup(rewritten) + assert found is not None and run_end - 100 < found.split_char < run_end + assert cache.encode(tokenizer, rewritten) == tokenizer(rewritten)["input_ids"] + assert cache.resync_failures == 1 + assert not cache.disabled + + +def test_a_tokenizer_that_shifts_every_tail_gets_no_cache() -> None: + assert create_prefix_token_cache(_PrependingTokenizer()) is None + for tokenizer in (_MergeTokenizer(), _ExpandingTokenizer(), _UnsyncableTokenizer()): + assert create_prefix_token_cache(tokenizer) is not None + + +@pytest.mark.parametrize("mismatch", [0, 1, 4095, 4096, 4097, 6000, 8191, None]) +def test_common_prefix_len_finds_the_first_mismatch(mismatch: int | None) -> None: + a = "".join(chr(97 + (i * 7) % 26) for i in range(8192)) + b = a if mismatch is None else a[:mismatch] + "#" + a[mismatch + 1 :] + expected = len(a) if mismatch is None else mismatch + assert _common_prefix_len(a, b) == _common_prefix_len(b, a) == expected + + +@pytest.mark.parametrize("n", [0, 4095, 4096, 4097, 8192]) +def test_common_prefix_len_of_a_prefix_is_its_length(n: int) -> None: + a = "".join(chr(97 + (i * 7) % 26) for i in range(8192)) + assert _common_prefix_len(a, a[:n]) == _common_prefix_len(a[:n], a) == n def test_a_split_is_moved_to_the_first_token_of_a_multi_token_character() -> None: @@ -184,6 +308,8 @@ def test_a_split_is_moved_to_the_first_token_of_a_multi_token_character() -> Non # The splits are usable, not rejected at the seam. assert cache.resync_failures == 0 assert cache.divergent_hits > 0 + assert not cache.disabled + _assert_checkpoints_are_valid(cache, tokenizer) def test_short_prompts_are_tokenized_but_not_cached() -> None: @@ -315,6 +441,22 @@ def test_error_disables_cache_and_still_returns_ids() -> None: assert not cache._entries +def test_a_prompt_the_tokenizer_rejects_fails_without_disabling_the_cache() -> None: + tokenizer, cache = _RejectingTokenizer(), _cache() + turns = _turns(3) + cache.encode(tokenizer, turns[0]) + with pytest.raises(TypeError): + cache.encode(tokenizer, turns[1] + "\ud800") + assert not cache.disabled + assert cache.encode(tokenizer, turns[2]) == tokenizer(turns[2])["input_ids"] + assert cache.hits == 1 + + +def test_checkpoints_must_be_wider_than_the_resync_span() -> None: + with pytest.raises(ValueError, match="checkpoint_tokens"): + PrefixTokenCacheConfig(resync=32, checkpoint_tokens=32) + + def test_config_from_env_reads_overrides(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setenv("TLLM_PREFIX_TOKEN_CACHE_ENTRIES", "7") monkeypatch.setenv("TLLM_PREFIX_TOKEN_CACHE_MIN_CHARS", "100") diff --git a/tests/unittest/llmapi/test_serve_frontend_affinity.py b/tests/unittest/llmapi/test_serve_frontend_affinity.py index 90569b6b8750..78a3c963ee29 100644 --- a/tests/unittest/llmapi/test_serve_frontend_affinity.py +++ b/tests/unittest/llmapi/test_serve_frontend_affinity.py @@ -106,6 +106,7 @@ def fake_spawn(llm, num_frontends, *, extra_env, report_failure, affinity_socket class FakeServer: def __init__(self, **kwargs): + served["multi_frontend_serving"] = kwargs["multi_frontend_serving"] self.app = Mock() def __call__(self, host, port, sockets): @@ -128,6 +129,13 @@ def __call__(self, host, port, sockets): ) assert served["ports"] == [base, base + stride] + # Tokenization routing needs the launcher's frontend id and the frontend count. + serving = served["multi_frontend_serving"] + assert (serving.is_launcher, serving.frontend_id, serving.num_frontends) == ( + True, + 0, + num_frontends, + ) assert served["uds"] == [str(tmp_path / serve.LAUNCHER_UDS_NAME)] assert handed_off == {1: base + 2 * stride, 2: base + 3 * stride} # Everything the launcher bound is released once it stops serving. diff --git a/tests/unittest/llmapi/test_serve_tokenization_owners.py b/tests/unittest/llmapi/test_serve_tokenization_owners.py index edfb86ab4ac6..f5c3946d0b9a 100644 --- a/tests/unittest/llmapi/test_serve_tokenization_owners.py +++ b/tests/unittest/llmapi/test_serve_tokenization_owners.py @@ -18,19 +18,36 @@ on its own socket, as two trtllm-serve frontend processes do. """ +import contextlib import os import socket +import tempfile +import threading import time -from typing import Any +from types import SimpleNamespace +from typing import Any, Iterator +from unittest.mock import patch import pytest +from tensorrt_llm.inputs import prefix_token_cache from tensorrt_llm.inputs.prefix_token_cache import ( PrefixTokenCache, PrefixTokenCacheConfig, tokenization_routing_key, ) -from tensorrt_llm.serve.tokenization_owners import SOCKET_NAME, TokenizationOwners, owner_of +from tensorrt_llm.serve import tokenization_owners +from tensorrt_llm.serve.multi_frontend import ( + MultiFrontendServing, + frontend_of_conversation, + tokenization_uds_path, +) +from tensorrt_llm.serve.tokenization_owners import ( + ENABLE_ENV_VAR, + TokenizationOwners, + routing_enabled, + start_tokenization_owners, +) pytestmark = pytest.mark.cpu_only @@ -38,13 +55,41 @@ class _CharTokenizer: is_fast = True + def __init__(self) -> None: + self.calls = 0 + def __call__(self, text: str, **kwargs: Any) -> dict[str, Any]: + self.calls += 1 return { "input_ids": [ord(c) for c in text], "offset_mapping": [(i, i + 1) for i in range(len(text))], } +class _RejectingTokenizer(_CharTokenizer): + """Rejects prompts containing ``MARKER``, as HF tokenizers reject a lone surrogate.""" + + MARKER = "" + + def __call__(self, text: str, **kwargs: Any) -> dict[str, Any]: + if self.MARKER in text: + self.calls += 1 + raise TypeError("cannot tokenize") + return super().__call__(text, **kwargs) + + +class _SlowTokenizer(_CharTokenizer): + def __init__(self, delay_s: float) -> None: + super().__init__() + self.delay_s = delay_s + self.entered = threading.Event() + + def __call__(self, text: str, **kwargs: Any) -> dict[str, Any]: + self.entered.set() + time.sleep(self.delay_s) + return super().__call__(text, **kwargs) + + def _cache() -> PrefixTokenCache: return PrefixTokenCache( PrefixTokenCacheConfig( @@ -61,22 +106,40 @@ def _turns(n: int = 6) -> list[str]: return out -def _key_owned_by(owner: int, num_frontends: int) -> str: - return next( +def _key_owned_by(owner: int, num_frontends: int, skip: int = 0) -> str: + keys = ( key for key in (f"conversation-{i}" for i in range(100)) - if owner_of(key, num_frontends) == owner + if frontend_of_conversation(key, num_frontends) == owner ) + for _ in range(skip): + next(keys) + return next(keys) + + +@pytest.fixture +def ipc_dir() -> Iterator[str]: + # Short, so socket paths stay under the AF_UNIX path limit. + with tempfile.TemporaryDirectory(dir="/tmp") as path: + yield path + + +@contextlib.contextmanager +def _owners(ipc_dir: str, *specs: tuple[int, PrefixTokenCache, Any], **kwargs: Any): + started: list[TokenizationOwners] = [] + try: + for frontend_id, cache, tokenizer in specs: + started.append(TokenizationOwners(ipc_dir, frontend_id, 2, cache, tokenizer, **kwargs)) + yield started + finally: + for owner in started: + owner.close() -def test_a_conversation_is_tokenized_by_its_owner_wherever_its_turns_arrive(tmp_path) -> None: +def test_a_conversation_is_tokenized_by_its_owner_wherever_its_turns_arrive(ipc_dir) -> None: tokenizer = _CharTokenizer() caches = [_cache(), _cache()] - owners = [ - TokenizationOwners(str(tmp_path), frontend_id, 2, cache, tokenizer) - for frontend_id, cache in enumerate(caches) - ] - try: + with _owners(ipc_dir, (0, caches[0], tokenizer), (1, caches[1], tokenizer)): key = _key_owned_by(1, 2) turns = _turns() # Turns alternate between the frontends, as SO_REUSEPORT spreads them. @@ -92,59 +155,226 @@ def test_a_conversation_is_tokenized_by_its_owner_wherever_its_turns_arrive(tmp_ # A prompt without a conversation stays where it arrived. caches[0].encode(tokenizer, turns[0]) assert caches[0].misses == 1 - finally: - for owner in owners: - owner.close() -class _FailingTokenizer(_CharTokenizer): - def __call__(self, text: str, **kwargs: Any) -> dict[str, Any]: - raise ValueError("cannot tokenize") +def test_frontends_share_one_connection_per_owner(ipc_dir) -> None: + tokenizer = _CharTokenizer() + caches = [_cache(), _cache()] + with _owners(ipc_dir, (0, caches[0], tokenizer), (1, caches[1], tokenizer)) as owners: + prompts = _turns(8) + errors: list[BaseException] = [] + + def worker(thread: int) -> None: + try: + with tokenization_routing_key(_key_owned_by(1, 2, skip=thread)): + for prompt in prompts: + assert caches[0].encode(tokenizer, prompt) == [ord(c) for c in prompt] + except BaseException as exc: # surface in the main thread + errors.append(exc) + + threads = [threading.Thread(target=worker, args=(i,)) for i in range(4)] + for t in threads: + t.start() + for t in threads: + t.join() + assert not errors + assert caches[0].routed == 4 * len(prompts) + assert len(owners[1]._server.handlers) == 1 -def test_a_prompt_is_tokenized_locally_when_its_owner_does_not_tokenize_it(tmp_path) -> None: +def test_a_prompt_is_tokenized_locally_while_its_owner_is_not_serving(ipc_dir) -> None: tokenizer, cache = _CharTokenizer(), _cache() - owners = [TokenizationOwners(str(tmp_path), 0, 2, cache, tokenizer, retry_after_s=0.0)] - try: - prompt = _turns(1)[0] + prompt = _turns(1)[0] + with _owners(ipc_dir, (0, cache, tokenizer)) as owners: with tokenization_routing_key(_key_owned_by(1, 2)): # Frontend 1 has not started. assert cache.encode(tokenizer, prompt) == [ord(c) for c in prompt] assert (cache.routed, cache.route_failures, cache.misses) == (0, 1, 1) - # Frontend 1 runs but its tokenizer fails; the connection survives. - owners.append(TokenizationOwners(str(tmp_path), 1, 2, _cache(), _FailingTokenizer())) - for _ in range(2): + # Once it serves, it is asked at once: a refused connect is not a timeout. + owners.append(TokenizationOwners(ipc_dir, 1, 2, _cache(), tokenizer)) + assert cache.encode(tokenizer, prompt) == [ord(c) for c in prompt] + assert (cache.routed, cache.route_failures) == (1, 1) + assert "1 route failures" in cache.summary() + + +def test_a_prompt_the_owner_cannot_tokenize_disables_no_cache(ipc_dir) -> None: + tokenizer, owner_tokenizer = _RejectingTokenizer(), _RejectingTokenizer() + cache, owner_cache = _cache(), _cache() + with _owners(ipc_dir, (0, cache, tokenizer), (1, owner_cache, owner_tokenizer)) as owners: + turns = _turns(3) + rejected = turns[0] + _RejectingTokenizer.MARKER + with tokenization_routing_key(_key_owned_by(1, 2)): + assert cache.encode(tokenizer, turns[0]) == [ord(c) for c in turns[0]] + connection = owners[0]._peers[1].conn + calls = owner_tokenizer.calls + + with pytest.raises(TypeError): + cache.encode(tokenizer, rejected) + # The owner tried it, with and without the cache, and so did the + # asking frontend. + assert owner_tokenizer.calls == calls + 2 + assert tokenizer.calls == 2 + assert not cache.disabled and not owner_cache.disabled + assert cache.route_failures == 1 + + # The connection survives, and the conversation still hits at its owner. + for prompt in turns[1:]: assert cache.encode(tokenizer, prompt) == [ord(c) for c in prompt] - assert (cache.routed, cache.route_failures) == (0, 3) - assert "3 route failures" in cache.summary() - finally: - for owner in owners: - owner.close() + assert owners[0]._peers[1].conn is connection + assert cache.routed == 3 + assert owner_cache.hits == 2 -def test_a_hung_owner_is_skipped_after_it_times_out(tmp_path) -> None: +def test_a_hung_owner_is_skipped_after_it_times_out(ipc_dir) -> None: tokenizer, cache = _CharTokenizer(), _cache() + timeout_s = 0.5 # Frontend 1 accepts connections but never answers. - hung = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) - hung.bind(os.path.join(str(tmp_path), SOCKET_NAME.format(1))) - hung.listen() - owners = TokenizationOwners(str(tmp_path), 0, 2, cache, tokenizer, timeout_s=0.5) - try: + with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as hung: + hung.bind(tokenization_uds_path(ipc_dir, 1)) + hung.listen() + with ( + _owners(ipc_dir, (0, cache, tokenizer), timeout_s=timeout_s), + patch.object(tokenization_owners.logger, "warning") as warning, + ): + prompt = _turns(1)[0] + with tokenization_routing_key(_key_owned_by(1, 2)): + elapsed = [] + for _ in range(3): + start = time.monotonic() + assert cache.encode(tokenizer, prompt) == [ord(c) for c in prompt] + elapsed.append(time.monotonic() - start) + # Only the first prompt waited for the owner. + assert elapsed[0] >= timeout_s + assert max(elapsed[1:]) < timeout_s / 2 + assert (cache.routed, cache.route_failures) == (0, 3) + # Logged once, when the owner started being skipped. + assert warning.call_count == 1 + assert "Frontend 1" in warning.call_args.args[0] + + +def test_close_does_not_wait_long_for_a_handler_in_the_middle_of_a_prompt(ipc_dir) -> None: + tokenizer, slow = _CharTokenizer(), _SlowTokenizer(delay_s=3.0) + cache = _cache() + owner = TokenizationOwners(ipc_dir, 1, 2, _cache(), slow) + with _owners(ipc_dir, (0, cache, tokenizer)): prompt = _turns(1)[0] + + def ask() -> None: + with tokenization_routing_key(_key_owned_by(1, 2)): + cache.encode(tokenizer, prompt) + + asking = threading.Thread(target=ask, daemon=True) + asking.start() + assert slow.entered.wait(5.0) + start = time.monotonic() + owner.close(timeout_s=0.2) + # Up to the serve_forever poll interval (0.5 s) plus the shared deadline. + assert time.monotonic() - start < 1.5 + + +def test_routed_prompts_count_toward_the_periodic_log_line(ipc_dir) -> None: + tokenizer = _CharTokenizer() + caches = [_cache(), _cache()] + with ( + _owners(ipc_dir, (0, caches[0], tokenizer), (1, caches[1], tokenizer)), + patch.object(prefix_token_cache, "LOG_EVERY_REQUESTS", 2), + patch.object(prefix_token_cache.logger, "info") as info, + ): with tokenization_routing_key(_key_owned_by(1, 2)): - for _ in range(3): - start = time.monotonic() - assert cache.encode(tokenizer, prompt) == [ord(c) for c in prompt] - elapsed = time.monotonic() - start - # Only the first prompt waited for the owner. - assert elapsed < 0.25 - assert (cache.routed, cache.route_failures) == (0, 3) - finally: - owners.close() - hung.close() + for prompt in _turns(4): + caches[0].encode(tokenizer, prompt) + assert caches[0].requests == 4 + # Two lines from frontend 0, whose prompts were all routed, and two from the owner. + lines = [call.args[0] for call in info.call_args_list] + assert sum("4 routed" in line or "2 routed" in line for line in lines) == 2 + assert len(lines) == 4 + + +def test_close_removes_the_route(ipc_dir) -> None: + tokenizer, cache = _CharTokenizer(), _cache() + with _owners(ipc_dir, (0, cache, tokenizer)): + pass + prompt = _turns(1)[0] + with tokenization_routing_key(_key_owned_by(1, 2)): + assert cache.encode(tokenizer, prompt) == [ord(c) for c in prompt] + assert (cache.routed, cache.route_failures, cache.misses) == (0, 0, 1) + + +def test_the_owner_of_a_conversation_is_fixed() -> None: + # Disaggregated servers pick the affinity port with the same mapping, in + # other processes, so it must not change (for example to the seeded hash()). + assert [frontend_of_conversation(f"conversation-{i}", 8) for i in range(4)] == [2, 4, 6, 0] + assert frontend_of_conversation("session-a", 3) == 1 def test_a_conversation_id_with_a_lone_surrogate_has_an_owner() -> None: # JSON can carry one; it must not fail the request it arrives with. - assert owner_of("conversation-\ud800", 2) in (0, 1) + assert frontend_of_conversation("conversation-\ud800", 2) in (0, 1) + + +@pytest.mark.parametrize("value, enabled", [(None, False), ("0", False), ("1", True)]) +def test_routing_switch_values(monkeypatch: pytest.MonkeyPatch, value: Any, enabled: bool) -> None: + if value is None: + monkeypatch.delenv(ENABLE_ENV_VAR, raising=False) + else: + monkeypatch.setenv(ENABLE_ENV_VAR, value) + assert routing_enabled() is enabled + + +@pytest.mark.parametrize("value", ["true", "on", "1 ", "", "2"]) +def test_a_malformed_routing_switch_raises(monkeypatch: pytest.MonkeyPatch, value: str) -> None: + monkeypatch.setenv(ENABLE_ENV_VAR, value) + with pytest.raises(ValueError, match=ENABLE_ENV_VAR): + routing_enabled() + + +def _processor(cache: Any) -> SimpleNamespace: + tokenizer = _CharTokenizer() + return SimpleNamespace( + get_prefix_token_cache=lambda: None if cache is None else (cache, tokenizer) + ) + + +def _serving(ipc_dir: str, num_frontends: int, frontend_id: int = 0) -> MultiFrontendServing: + return MultiFrontendServing( + launcher_uds=os.path.join(ipc_dir, "launcher.sock"), + is_launcher=frontend_id == 0, + frontend_id=frontend_id, + num_frontends=num_frontends, + ) + + +def test_routing_starts_when_switched_on_with_frontends_and_a_cache( + monkeypatch: pytest.MonkeyPatch, ipc_dir: str +) -> None: + monkeypatch.setenv(ENABLE_ENV_VAR, "1") + cache = _cache() + with patch.object(tokenization_owners.logger, "info") as info: + owners = start_tokenization_owners(_serving(ipc_dir, 4, frontend_id=2), _processor(cache)) + try: + assert owners is not None + assert os.path.exists(tokenization_uds_path(ipc_dir, 2)) + assert "frontend 2 of 4" in info.call_args.args[0] + finally: + owners.close() + + +@pytest.mark.parametrize("num_frontends, has_cache", [(1, True), (4, False)]) +def test_routing_warns_when_switched_on_but_it_cannot_engage( + monkeypatch: pytest.MonkeyPatch, ipc_dir: str, num_frontends: int, has_cache: bool +) -> None: + monkeypatch.setenv(ENABLE_ENV_VAR, "1") + processor = _processor(_cache() if has_cache else None) + with patch.object(tokenization_owners.logger, "warning") as warning: + assert start_tokenization_owners(_serving(ipc_dir, num_frontends), processor) is None + assert start_tokenization_owners(None, processor) is None + assert warning.call_count == 2 + assert ENABLE_ENV_VAR in warning.call_args.args[0] + assert not os.listdir(ipc_dir) + + +def test_routing_is_off_by_default(monkeypatch: pytest.MonkeyPatch, ipc_dir: str) -> None: + monkeypatch.delenv(ENABLE_ENV_VAR, raising=False) + assert start_tokenization_owners(_serving(ipc_dir, 4), _processor(_cache())) is None + assert not os.listdir(ipc_dir) From bb9d467e783216faab1cddcacff3b030d70002bf Mon Sep 17 00:00:00 2001 From: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com> Date: Mon, 9 Nov 2026 22:23:55 -0800 Subject: [PATCH 8/8] [None][fix] Let the stalled tokenization handler finish before the close test ends close() returns at its deadline and leaves a handler that is still tokenizing running, so the test left the owner's handler thread alive and pytest-threadleak failed it. Release the slow tokenizer after the timing check and join the asking and handler threads. Signed-off-by: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com> --- tests/unittest/llmapi/test_serve_tokenization_owners.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/tests/unittest/llmapi/test_serve_tokenization_owners.py b/tests/unittest/llmapi/test_serve_tokenization_owners.py index f5c3946d0b9a..b2abe7491719 100644 --- a/tests/unittest/llmapi/test_serve_tokenization_owners.py +++ b/tests/unittest/llmapi/test_serve_tokenization_owners.py @@ -83,10 +83,11 @@ def __init__(self, delay_s: float) -> None: super().__init__() self.delay_s = delay_s self.entered = threading.Event() + self.release = threading.Event() def __call__(self, text: str, **kwargs: Any) -> dict[str, Any]: self.entered.set() - time.sleep(self.delay_s) + self.release.wait(self.delay_s) return super().__call__(text, **kwargs) @@ -267,10 +268,16 @@ def ask() -> None: asking = threading.Thread(target=ask, daemon=True) asking.start() assert slow.entered.wait(5.0) + handlers = list(owner._server.handlers.values()) start = time.monotonic() owner.close(timeout_s=0.2) # Up to the serve_forever poll interval (0.5 s) plus the shared deadline. assert time.monotonic() - start < 1.5 + # close() left the handler tokenizing; finish it so no thread outlives the test. + slow.release.set() + for thread in [asking, *handlers]: + thread.join(5.0) + assert not thread.is_alive() def test_routed_prompts_count_toward_the_periodic_log_line(ipc_dir) -> None: