From b8f574699ed54d8650b5acd27100ec3aac76a9b0 Mon Sep 17 00:00:00 2001 From: Tianyu Xiong <117647511+tianyuxbear@users.noreply.github.com> Date: Fri, 29 May 2026 00:32:30 -0700 Subject: [PATCH] [https://nvbugs/6227203][fix] Remove redundant TikTokenTokenizer shim from KimiK25InputProcessor Signed-off-by: Tianyu Xiong <117647511+tianyuxbear@users.noreply.github.com> --- .../_torch/models/modeling_kimi_k25.py | 84 ------------------- tests/integration/test_lists/waives.txt | 1 - 2 files changed, 85 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_kimi_k25.py b/tensorrt_llm/_torch/models/modeling_kimi_k25.py index d130210376ba..a00eb05615b8 100644 --- a/tensorrt_llm/_torch/models/modeling_kimi_k25.py +++ b/tensorrt_llm/_torch/models/modeling_kimi_k25.py @@ -291,23 +291,6 @@ def _frames_to_chunks( # Default vocabulary size for K2.5 _VOCAB_SIZE = 163840 -# K2.5 special token markers that the transformers 5.5.x Rust fast tokenizer -# BPE-splits instead of mapping to canonical IDs. When any of these appear in -# a prompt, we must route tokenization through the slow ``TikTokenTokenizer``. -# Pure text (no markers and no multimodal data) keeps the fast tokenizer. -# See NVBug 6182617 (correctness) and NVBug 6248987 (perf). -_K25_SPECIAL_TOKEN_MARKERS = ( - "<|media_begin|>", - "<|media_content|>", - "<|media_pad|>", - "<|media_end|>", - "<|im_user|>", - "<|im_assistant|>", - "<|im_system|>", - "<|im_end|>", - "<|im_middle|>", -) - # --------------------------------------------------------------------------- # Native MoonViT3d Vision Encoder Components @@ -1069,16 +1052,6 @@ def __init__( config, "media_placeholder_token_id", _MEDIA_PLACEHOLDER_TOKEN_ID ) - # transformers 5.5.x ``AutoTokenizer`` may route K2.5 to the Rust - # fast backend, which BPE-splits ``<|media_pad|>`` / ``<|im_user|>`` - # / etc. instead of mapping them to their canonical IDs. The slow - # ``TikTokenTokenizer`` preserves them. Swap is deferred until we - # actually see an input that needs it (multimodal data or a K2.5 - # special token marker in the prompt) — the text-only thinking - # path keeps the fast tokenizer to avoid a GIL-bound 9x TPOT - # regression. See NVBug 6182617 (correctness) / 6248987 (perf). - self._slow_tokenizer_active = False - @property def config(self) -> PretrainedConfig: return self._config @@ -1195,51 +1168,6 @@ def get_num_tokens_per_video(self, *, video: List, **kwargs) -> int: total_tokens += self.get_num_tokens_per_image(image=chunk[0]) return total_tokens - @staticmethod - def _input_needs_slow_tokenizer(text: Optional[str]) -> bool: - """Return True iff ``text`` contains any K2.5 special token marker - that the Rust fast tokenizer would BPE-split incorrectly.""" - if not text: - return False - return any(marker in text for marker in _K25_SPECIAL_TOKEN_MARKERS) - - def _ensure_k25_slow_tokenizer(self) -> None: - """Override ``self._tokenizer`` and ``self._processor.tokenizer`` - with the model's slow ``TikTokenTokenizer``. - - Idempotent: callers invoke this lazily, on the first request that - actually requires correct mapping of K2.5 special tokens. Done this - way (instead of unconditionally in ``__init__``) so text-only - prompts keep the fast Rust tokenizer — running the slow Python - ``TikTokenTokenizer`` on the orchestrator GIL adds ~100 ms per - ``_fetch_new_requests`` / ``broadcast_requests`` step at 8 K-token - prompts, an order-of-magnitude TPOT regression. The slow class' - ``tokens_trie`` always splits the special tokens correctly. - See NVBug 6182617 (correctness) and NVBug 6248987 (perf). - """ - if self._slow_tokenizer_active: - return - from transformers.dynamic_module_utils import get_class_from_dynamic_module - - slow_cls = get_class_from_dynamic_module( - "tokenization_kimi.TikTokenTokenizer", - self._model_path, - ) - slow_tok = slow_cls.from_pretrained(self._model_path, trust_remote_code=True) - - logger.info( - "K2.5 InputProcessor swapping in slow TikTokenTokenizer " - "(originally %s). See NVBug 6182617.", - type(self._tokenizer).__name__, - ) - - self._tokenizer = slow_tok - # Image-only path uses ``self._processor.tokenizer`` (an - # independent instance from ``AutoProcessor``); swap it too. - if getattr(self._processor, "tokenizer", None) is not None: - self._processor.tokenizer = slow_tok - self._slow_tokenizer_active = True - @torch.inference_mode() def call_with_text_prompt( self, @@ -1271,18 +1199,9 @@ def call_with_text_prompt( # Text-only path if not images and not videos: - # Fast tokenizer is fine unless the prompt itself carries K2.5 - # special tokens (rare on the thinking perf path); only fall - # back to the slow ``TikTokenTokenizer`` then. See NVBug 6248987. - if self._input_needs_slow_tokenizer(text_prompt): - self._ensure_k25_slow_tokenizer() token_ids = self._tokenizer(text_prompt, return_tensors="pt").input_ids[0] return token_ids.to(torch.int32).tolist(), {} - # Multimodal path: prompt is rewritten with media placeholders that - # the fast tokenizer would BPE-split, so we always need the slow one. - self._ensure_k25_slow_tokenizer() - # Build the ``medias`` list expected by KimiK25Processor. # The HF processor accepts either ``messages`` (chat format) or # both ``medias`` and ``text``. Since we already have the @@ -1521,9 +1440,6 @@ def get_prompt_token_ids( f"must match model hidden size {expected_hidden_size}" ) - # Disagg-serving multimodal path: prompt has media placeholders that - # must map to canonical IDs, so the slow tokenizer is required. - self._ensure_k25_slow_tokenizer() input_ids = self._tokenizer(text_prompt, return_tensors="pt").input_ids[0] placeholder_id = self._media_placeholder_token_id diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index 3866d63e2df4..93aa6e7a0ad1 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -295,7 +295,6 @@ test_e2e.py::test_multi_nodes_eval[Qwen3/saved_models_Qwen3-235B-A22B_nvfp4_hf-t test_e2e.py::test_openai_chat_example[trt] SKIP (https://nvbugs/5477444) test_e2e.py::test_openai_completions_example[trt] SKIP (https://nvbugs/5701450) test_e2e.py::test_openai_disagg_multi_nodes_completion[ctx_tp1pp2-gen_tp1pp2] SKIP (https://nvbugs/6190759) -test_e2e.py::test_openai_kv_cache_contamination SKIP (https://nvbugs/6227203) test_e2e.py::test_ptp_quickstart_advanced_deepseek_r1_w4afp8_8gpus[DeepSeek-R1-W4AFP8-DeepSeek-R1/DeepSeek-R1-W4AFP8] SKIP (https://nvbugs/5836830) test_e2e.py::test_trtllm_bench_iteration_log[TRT-streaming-meta-llama/Llama-3.1-8B-llama-3.1-model/Meta-Llama-3.1-8B] SKIP (https://nvbugs/5448523) test_e2e.py::test_trtllm_multimodal_benchmark_serving SKIP (https://nvbugs/5864769)