[https://nvbugs/6248987][fix] Made the slow-tokenizer swap lazy and idempotent. __init__ now just sets `_slo - #14846
Conversation
…_` now just sets `_slo Signed-off-by: Chenfei Zhang <chenfeiz@nvidia.com>
📝 WalkthroughWalkthroughThis PR introduces lazy slow-tokenizer loading for K2.5 special-token handling in the Kimi K25 model. Rather than eagerly swapping tokenizers, the processor now detects K2.5-specific markers in prompt text and conditionally activates the slow tokenizer only when needed, improving performance on marker-free prompts. ChangesK2.5 Lazy Slow-Tokenizer Loading
🎯 3 (Moderate) | ⏱️ ~20 minutes 🚥 Pre-merge checks | ✅ 3 | ❌ 2❌ Failed checks (2 warnings)
✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
tensorrt_llm/_torch/models/modeling_kimi_k25.py (1)
1220-1241:⚠️ Potential issue | 🟠 Major | ⚡ Quick winPreserve the fast tokenizer after the first slow-tokenizer request.
This still mutates the shared tokenizer permanently. Once a multimodal request or
get_prompt_token_ids()hits_ensure_k25_slow_tokenizer(), later pure-text requests at Lines 1274-1279 also run through the slow Python tokenizer becauseself._tokenizernow points at it. That brings the perf regression back for the lifetime of thisKimiK25InputProcessor, which contradicts the PR goal of keeping pure-text prompts on the fast path. Keep separate fast/slow tokenizer references and select per call instead of overwritingself._tokenizer.Also applies to: 1274-1279, 1524-1526
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/_torch/models/modeling_kimi_k25.py` around lines 1220 - 1241, The current _ensure_k25_slow_tokenizer implementation permanently overwrites self._tokenizer with the slow Python tokenizer, causing later pure-text calls (e.g., get_prompt_token_ids and the pure-text path) to use the slow path; instead, keep the original fast tokenizer and store the slow tokenizer separately (e.g., self._slow_tokenizer) and a flag self._slow_tokenizer_active, do not assign slow_tok to self._tokenizer; update only self._processor.tokenizer if the image-only code truly requires the slow instance, and change call sites (where you currently reference self._tokenizer for tokenization in multimodal vs pure-text paths) to select slow vs fast explicitly (use self._slow_tokenizer when self._slow_tokenizer_active and multimodal/image path, otherwise use the preserved fast tokenizer).
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Outside diff comments:
In `@tensorrt_llm/_torch/models/modeling_kimi_k25.py`:
- Around line 1220-1241: The current _ensure_k25_slow_tokenizer implementation
permanently overwrites self._tokenizer with the slow Python tokenizer, causing
later pure-text calls (e.g., get_prompt_token_ids and the pure-text path) to use
the slow path; instead, keep the original fast tokenizer and store the slow
tokenizer separately (e.g., self._slow_tokenizer) and a flag
self._slow_tokenizer_active, do not assign slow_tok to self._tokenizer; update
only self._processor.tokenizer if the image-only code truly requires the slow
instance, and change call sites (where you currently reference self._tokenizer
for tokenization in multimodal vs pure-text paths) to select slow vs fast
explicitly (use self._slow_tokenizer when self._slow_tokenizer_active and
multimodal/image path, otherwise use the preserved fast tokenizer).
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 00aa6cee-0714-48cc-8011-ba3ec6f5645a
📒 Files selected for processing (1)
tensorrt_llm/_torch/models/modeling_kimi_k25.py
|
/bot run --disable-fail-fast --stage-list "DGX_B200-8_GPUs-PyTorch-PerfSanity-Post-Merge-1,GB200-4_GPUs-PyTorch-PerfSanity-Post-Merge-7" |
|
PR_Github #51500 [ run ] triggered by Bot. Commit: |
|
PR_Github #51500 [ run ] completed with state
|
|
/bot run --disable-fail-fast --stage-list "GB200-4_GPUs-PyTorch-PerfSanity-Post-Merge-1,GB200-4_GPUs-PyTorch-PerfSanity-Post-Merge-2,GB200-4_GPUs-PyTorch-PerfSanity-Post-Merge-3,GB200-4_GPUs-PyTorch-PerfSanity-Post-Merge-4,GB200-4_GPUs-PyTorch-PerfSanity-Post-Merge-5,GB200-4_GPUs-PyTorch-PerfSanity-Post-Merge-6" |
|
PR_Github #51618 [ run ] triggered by Bot. Commit: |
|
PR_Github #51618 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #51769 [ run ] triggered by Bot. Commit: |
|
PR_Github #51769 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #51817 [ run ] triggered by Bot. Commit: |
|
PR_Github #51817 [ run ] completed with state |
5 participants
Summary
_ensure_k25_slow_tokenizer()inKimiK25InputProcessor.__init__, swapping in the pure-PythonTikTokenTokenizer. On the text-only k25_thinking_fp4 perf path (ISL=8K, concurrency=2) every request is tokenized on the orchestrator's GIL, inflating_fetch_new_requests/broadcast_requests~34% and TPOT 9.8× (output_token_throughput drops ~92%). GPU kernel inventory and runtimes were unchanged — pure host overhead.__init__now just sets_slow_tokenizer_active=False._ensure_k25_slow_tokenizer()early-returns once active and is invoked only when an input actually needs correct mapping of K2.5 special tokens — i.e. multimodal data is present, the disaggget_prompt_token_idspath runs, or the prompt itself contains markers like<|media_pad|>/<|im_user|>(checked via_input_needs_slow_tokenizer). Pure-text prompts keep the fast Rust tokenizer; the multimodal accuracy path NVBug 6182617 still gets the slowtokens_trie.Test plan
Links
Summary by CodeRabbit
Bug Fixes