[None][fix] Invalidate cached special-id set in slow convert_ids_to_tokens

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],

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(