From 00925b2fa10d57682f116721877ab5f5b4235460 Mon Sep 17 00:00:00 2001 From: xiaoweis <39303645+Shixiaowei02@users.noreply.github.com> Date: Mon, 13 Jul 2026 08:08:56 -0700 Subject: [PATCH] [None][fix] Invalidate cached special-id set in slow convert_ids_to_tokens The slow-tokenizer optimization in convert_ids_to_tokens snapshotted all_special_ids into a set at __init__ and reused it forever. If special tokens were (re)registered after the wrapper was constructed (e.g. add_special_tokens during chat-template / adapter setup), the snapshot went stale: the affected ids stopped being skipped and leaked into detokenized output when skip_special_tokens=True -- an accuracy regression. Recomputing all_special_ids on every call is correct but a severe perf regression: convert_tokens_to_ids(all_special_tokens) is expensive on slow tokenizers and runs once per streamed token (measured ~10^2-10^4x slower, superlinear in the number of special tokens). Instead keep the memoized set and rebuild it only when a cheap O(1) fingerprint of the tokenizer's special tokens changes: the count of registered added tokens (covers newly added specials, including additional_special_tokens) plus the ids of the named special tokens (covers reassigning/clearing bos/eos/pad/unk/sep/cls/mask). Steady-state streaming stays O(1) per call while the cache can never be stale. Add regression tests on a genuine slow tokenizer for a late-added special token and a reassigned named special token. Signed-off-by: xiaoweis <39303645+Shixiaowei02@users.noreply.github.com> --- tensorrt_llm/tokenizer/tokenizer.py | 58 +++++++---- .../unittest/llmapi/test_tokenizer_decode.py | 97 +++++++++++++++++++ 2 files changed, 137 insertions(+), 18 deletions(-) diff --git a/tensorrt_llm/tokenizer/tokenizer.py b/tensorrt_llm/tokenizer/tokenizer.py index 1422cc3804b5..5bafcf01a823 100644 --- a/tensorrt_llm/tokenizer/tokenizer.py +++ b/tensorrt_llm/tokenizer/tokenizer.py @@ -202,14 +202,8 @@ def __init__(self, tokenizer): self.tokenizer.all_special_tokens) else: self._all_special_tokens_set = set() - # Cache of special-token ids used by convert_ids_to_tokens to work - # around an O(N*K) bug in transformers 5.x's slow tokenizer path - # (see convert_ids_to_tokens below). Computed once here so we don't - # rebuild it on every per-token streaming call. - try: - self._all_special_ids_set = set(self.tokenizer.all_special_ids) - except (AttributeError, NotImplementedError): - self._all_special_ids_set = set() + self._special_ids_cache = set() + self._special_ids_fingerprint = None def __reduce__(self): # In multi-node scenarios, AutoTokenizer.from_pretrained with @@ -339,19 +333,47 @@ def convert_ids_to_tokens( # the buggy branch is short-circuited. return inner.convert_ids_to_tokens( ids, skip_special_tokens=skip_special_tokens) - # Slow PreTrainedTokenizer.convert_ids_to_tokens in transformers 5.x - # tests `index in self.all_special_ids` inside the per-id loop, and - # `all_special_ids` is an @property that rebuilds the list on every - # access via convert_tokens_to_ids(self.all_special_tokens). The fast - # subclass already hoists this out of the loop (see - # tokenization_utils_tokenizers.py), but the slow base class wasn't - # updated. Mirror that fix using the set we cached in __init__, - # which keeps streaming detokenization at O(1) all-special-ids cost - # per call regardless of batch size or stream_interval. + # transformers 5.x's slow convert_ids_to_tokens checks + # `index in self.all_special_ids` per id, and all_special_ids rebuilds + # itself on every access. Look ids up in a cached set instead: O(N*K) + # -> O(N) per call. tokens = inner.convert_ids_to_tokens(ids, skip_special_tokens=False) - special_ids = self._all_special_ids_set + special_ids = self._special_token_ids() return [t for idx, t in zip(ids, tokens) if int(idx) not in special_ids] + def _special_token_ids(self) -> set: + """Set of special-token ids, memoized across streaming calls. + + A snapshot taken once in __init__ goes stale if special tokens are + (re)registered later, leaking them into detokenized output. Recomputing + all_special_ids every call is correct but slow on slow tokenizers. + Rebuild the cache only when a cheap O(1) fingerprint changes: the + added-token count (new specials) plus the ids of the named special + tokens (bos/eos/... reassigned or cleared). + """ + inner = self.tokenizer + try: + # additional_special_tokens are covered by the added-token count. + named_ids = tuple( + getattr(inner, attr + "_id") + for attr in inner.SPECIAL_TOKENS_ATTRIBUTES + if attr != "additional_special_tokens") + fingerprint = (len(inner._added_tokens_decoder), named_ids) + except (AttributeError, NotImplementedError): + fingerprint = None + if fingerprint is None: + try: + return set(inner.all_special_ids) + except (AttributeError, NotImplementedError): + return set() + if fingerprint != self._special_ids_fingerprint: + try: + self._special_ids_cache = set(inner.all_special_ids) + except (AttributeError, NotImplementedError): + self._special_ids_cache = set() + self._special_ids_fingerprint = fingerprint + return self._special_ids_cache + def convert_tokens_to_string( self, tokens: List[str], diff --git a/tests/unittest/llmapi/test_tokenizer_decode.py b/tests/unittest/llmapi/test_tokenizer_decode.py index b1df835b29b6..1fb2a364db07 100644 --- a/tests/unittest/llmapi/test_tokenizer_decode.py +++ b/tests/unittest/llmapi/test_tokenizer_decode.py @@ -17,10 +17,107 @@ from tokenizers.decoders import ByteFallback, Fuse, Sequence from tokenizers.models import WordLevel from transformers import PreTrainedTokenizerFast +from transformers.tokenization_utils import PreTrainedTokenizer from tensorrt_llm.tokenizer import TransformersTokenizer +class _ToySlowTokenizer(PreTrainedTokenizer): + """A minimal genuine *slow* tokenizer (``is_fast`` is False). + + The slow ``convert_ids_to_tokens`` path is the only one that filters + special tokens in TransformersTokenizer's wrapper; fast tokenizers delegate + straight to the backend. This toy tokenizer lets us exercise that path + deterministically without downloading weights. + """ + + def __init__(self, **kwargs): + self._vocab = {f"tok{i}": i for i in range(64)} + self._ids = {v: k for k, v in self._vocab.items()} + super().__init__(**kwargs) + + @property + def vocab_size(self) -> int: + return len(self._vocab) + + def get_vocab(self): + return dict(self._vocab) + + def _convert_token_to_id(self, token): + return self._vocab.get(token, 0) + + def _convert_id_to_token(self, index): + return self._ids.get(int(index), "tok0") + + def _tokenize(self, text): + return text.split() + + +def test_convert_ids_to_tokens_reflects_late_registered_special_tokens() -> None: + # Regression test: convert_ids_to_tokens(skip_special_tokens=True) on a slow + # tokenizer must skip special tokens that were registered *after* the + # TransformersTokenizer wrapper was constructed. A previous optimization + # snapshotted all_special_ids once in __init__, so any special token added + # later (e.g. via add_special_tokens during chat-template / adapter setup) + # went unskipped and leaked into the detokenized output. + inner = _ToySlowTokenizer() + inner.add_special_tokens( + { + "bos_token": "tok1", + "eos_token": "tok2", + "additional_special_tokens": ["tok5"], + } + ) + tokenizer = TransformersTokenizer(inner) + assert inner.is_fast is False + + # Register a brand-new special token after the wrapper already exists. + inner.add_special_tokens({"additional_special_tokens": ["tok5", "tok42"]}) + + ids = [10, 42, 11, 5, 12] + # Prime the cache first (with the pre-mutation special ids), then mutate, + # to make sure the cache is invalidated rather than serving a stale set. + tokenizer.convert_ids_to_tokens(ids, skip_special_tokens=True) + inner.add_special_tokens({"additional_special_tokens": ["tok5", "tok42"]}) + got = tokenizer.convert_ids_to_tokens(ids, skip_special_tokens=True) + + # Ground truth is the underlying tokenizer's own skip behavior. + assert got == inner.convert_ids_to_tokens(ids, skip_special_tokens=True) + # The late-added special token (id 42) and the earlier one (id 5) must be + # dropped; only the ordinary tokens survive. + assert got == ["tok10", "tok11", "tok12"] + + +def test_convert_ids_to_tokens_reflects_reassigned_special_token() -> None: + # The cached special-id set is keyed on a fingerprint that includes the + # named special-token ids, so re-assigning eos to a token that is *already* + # registered (which leaves the added-tokens count unchanged) still + # invalidates the cache. Without the id component this would serve a stale + # set and drop/keep the wrong tokens. + inner = _ToySlowTokenizer() + inner.add_special_tokens( + { + "bos_token": "tok1", + "eos_token": "tok2", + "additional_special_tokens": ["tok5"], + } + ) + tokenizer = TransformersTokenizer(inner) + assert inner.is_fast is False + + ids = [10, 2, 11, 5, 12] + # Prime the cache: at this point id 2 (eos) is special and must be dropped. + primed = tokenizer.convert_ids_to_tokens(ids, skip_special_tokens=True) + assert primed == ["tok10", "tok11", "tok12"] + + # Re-point eos onto tok5 (id 5), which is already a registered special + # token. Now id 2 is no longer special and must survive; id 5 must not. + inner.add_special_tokens({"eos_token": "tok5"}) + got = tokenizer.convert_ids_to_tokens(ids, skip_special_tokens=True) + assert got == inner.convert_ids_to_tokens(ids, skip_special_tokens=True) + assert got == ["tok10", "tok2", "tok11", "tok12"] + + def test_hf_decode_incrementally_recovers_from_invalid_prefix() -> None: backend = Tokenizer( WordLevel(