Skip to content

[https://nvbugs/6227203][fix] Remove redundant TikTokenTokenizer shim from KimiK25InputProcessor #14741

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
lancelly merged 1 commit into from
Jun 9, 2026
Merged

[https://nvbugs/6227203][fix] Remove redundant TikTokenTokenizer shim from KimiK25InputProcessor #14741

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
84 changes: 0 additions & 84 deletions tensorrt_llm/_torch/models/modeling_kimi_k25.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down
1 change: 0 additions & 1 deletion tests/integration/test_lists/waives.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
Loading