Tiny speed up kimi detokenizer by 10x (#16427)

python/sglang/srt/environ.py

@@ -416,6 +416,9 @@ class Envs:
     SGLANG_ENABLE_METRICS_DEVICE_TIMER = EnvBool(False)
     SGLANG_ENABLE_METRICS_DP_ATTENTION = EnvBool(False)
 
+    # Tokenizer
+    SGLANG_PATCH_TOKENIZER = EnvBool(False)  # TODO enable by default
+
     # fmt: on
 
 

python/sglang/srt/utils/hf_transformers_utils.py

@@ -68,6 +68,7 @@ from sglang.srt.configs.internvl import InternVLChatConfig
 from sglang.srt.connector import create_remote_connector
 from sglang.srt.multimodal.customized_mm_processor_utils import _CUSTOMIZED_MM_PROCESSOR
 from sglang.srt.utils import is_remote_url, logger, lru_cache_frozenset, mistral_utils
+from sglang.srt.utils.patch_tokenizer import patch_tokenizer
 
 _CONFIG_REGISTRY: List[Type[PretrainedConfig]] = [
     ChatGLMConfig,
@@ -501,6 +502,7 @@ def get_tokenizer(
         )
 
     attach_additional_stop_token_ids(tokenizer)
+    tokenizer = patch_tokenizer(tokenizer)
     return tokenizer
 
 

python/sglang/srt/utils/patch_tokenizer.py

@@ -0,0 +1,116 @@
+import logging
+
+from sglang.srt.environ import envs
+
+logger = logging.getLogger(__name__)
+
+
+def patch_tokenizer(tokenizer):
+    if not envs.SGLANG_PATCH_TOKENIZER.get():
+        return tokenizer
+
+    if _is_kimi_tiktoken_tokenizer(tokenizer):
+        logger.info(
+            f"Applying special tokens cache patch for Kimi tokenizer: {type(tokenizer)}"
+        )
+        return _SpecialTokensCachePatcher.patch(tokenizer)
+
+    return tokenizer
+
+
+def unpatch_tokenizer(tokenizer):
+    return _SpecialTokensCachePatcher.unpatch(tokenizer)
+
+
+def _is_kimi_tiktoken_tokenizer(tokenizer):
+    cls = type(tokenizer)
+    class_name = cls.__name__
+    module_name = cls.__module__ or ""
+    return class_name == "TikTokenTokenizer" and "tokenization_kimi" in module_name
+
+
+class _SpecialTokensCachePatcher:
+    _PATCHED_FLAG = "_sglang_special_tokens_patched"
+    _CACHED_TOKENS_ATTR = "_sglang_cached_special_tokens"
+    _CACHED_IDS_ATTR = "_sglang_cached_special_ids"
+
+    @classmethod
+    def patch(cls, tokenizer):
+        tokenizer_cls = type(tokenizer)
+
+        if getattr(tokenizer_cls, cls._PATCHED_FLAG, False):
+            return tokenizer
+
+        tokenizer_cls._original_all_special_tokens = (
+            tokenizer_cls.all_special_tokens.fget
+        )
+        tokenizer_cls._original_all_special_ids = tokenizer_cls.all_special_ids.fget
+        tokenizer_cls._original_add_special_tokens = tokenizer_cls.add_special_tokens
+        tokenizer_cls._original_add_tokens = tokenizer_cls.add_tokens
+
+        patched_all_special_tokens = _make_cached_property(
+            cls._CACHED_TOKENS_ATTR, tokenizer_cls._original_all_special_tokens
+        )
+        patched_all_special_ids = _make_cached_property(
+            cls._CACHED_IDS_ATTR, tokenizer_cls._original_all_special_ids
+        )
+
+        def patched_add_special_tokens(self, *args, **kwargs):
+            assert (
+                False
+            ), "Cannot modify special tokens after patch. Call unpatch_tokenizer first."
+
+        def patched_add_tokens(self, new_tokens, special_tokens=False):
+            assert (
+                not special_tokens
+            ), "Cannot add special tokens after patch. Call unpatch_tokenizer first."
+            return tokenizer_cls._original_add_tokens(
+                self, new_tokens, special_tokens=False
+            )
+
+        tokenizer_cls.all_special_tokens = patched_all_special_tokens
+        tokenizer_cls.all_special_ids = patched_all_special_ids
+        tokenizer_cls.add_special_tokens = patched_add_special_tokens
+        tokenizer_cls.add_tokens = patched_add_tokens
+        setattr(tokenizer_cls, cls._PATCHED_FLAG, True)
+
+        return tokenizer
+
+    @classmethod
+    def unpatch(cls, tokenizer):
+        tokenizer_cls = type(tokenizer)
+
+        if not getattr(tokenizer_cls, cls._PATCHED_FLAG, False):
+            return tokenizer
+
+        tokenizer_cls.all_special_tokens = property(
+            tokenizer_cls._original_all_special_tokens
+        )
+        tokenizer_cls.all_special_ids = property(
+            tokenizer_cls._original_all_special_ids
+        )
+        tokenizer_cls.add_special_tokens = tokenizer_cls._original_add_special_tokens
+        tokenizer_cls.add_tokens = tokenizer_cls._original_add_tokens
+
+        del tokenizer_cls._original_all_special_tokens
+        del tokenizer_cls._original_all_special_ids
+        del tokenizer_cls._original_add_special_tokens
+        del tokenizer_cls._original_add_tokens
+        delattr(tokenizer_cls, cls._PATCHED_FLAG)
+
+        for attr in [cls._CACHED_TOKENS_ATTR, cls._CACHED_IDS_ATTR]:
+            if hasattr(tokenizer, attr):
+                delattr(tokenizer, attr)
+
+        logger.info(f"Unpatched special tokens cache for {tokenizer_cls.__name__}")
+        return tokenizer
+
+
+def _make_cached_property(cache_attr, original_fn):
+    @property
+    def cached_prop(self):
+        if getattr(self, cache_attr, None) is None:
+            setattr(self, cache_attr, original_fn(self))
+        return getattr(self, cache_attr)
+
+    return cached_prop