[None][fix] Support custom_tokenizer in KvCacheAwareRouter for disagg serving - #12990
Conversation
📝 WalkthroughWalkthroughThe changes add support for custom tokenizers to the routing layer by introducing an optional Changes
Sequence Diagram(s)sequenceDiagram
participant Router as KvCacheAwareRouter
participant Aliases as TOKENIZER_ALIASES
participant Importer as Dynamic Importer
participant Tokenizer as Custom Tokenizer Class
Router->>Router: __init__(custom_tokenizer)
Router->>Router: store custom_tokenizer
Router->>Router: _get_tokenizer()
alt custom_tokenizer is set
Router->>Aliases: lookup custom_tokenizer
Aliases-->>Router: resolved class path
Router->>Importer: dynamically import class
Importer-->>Router: tokenizer class
Router->>Tokenizer: from_pretrained(model)
Tokenizer-->>Router: tokenizer instance
else
Router->>Router: use AutoTokenizer.from_pretrained()
end
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes 🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 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.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@tensorrt_llm/serve/router.py`:
- Around line 659-668: The dynamic custom-tokenizer import block (using
self._custom_tokenizer, TOKENIZER_ALIASES, tokenizer_path, module_path,
class_name, tokenizer_class, and assigning into self._tokenizers[model]) must be
wrapped in a try/except that catches errors from rsplit, import_module, and
getattr and raises a clear, deterministic configuration error (instead of
bubbling raw exceptions); mirror the pattern used in
tensorrt_llm/llmapi/llm_args.py and
tensorrt_llm/serve/scripts/backend_request_func.py: attempt to resolve
tokenizer_path and instantiate via tokenizer_class.from_pretrained(model) in the
try, and on exception raise a concise ValueError or custom config error with a
message that includes the invalid alias/import/class name and the original
exception message to aid diagnostics, while leaving the
AutoTokenizer.from_pretrained fallback branch unchanged.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: 04cc2cf7-9a1c-4168-8b52-9ef56b13bb20
📒 Files selected for processing (2)
tensorrt_llm/llmapi/llm_args.pytensorrt_llm/serve/router.py
cf35b05 to
8d8bd50
Compare
svc-trtllm-gh-bot
added
the
Community want to contribute
8d8bd50 to
ee2a484
Compare
|
/bot run |
ee2a484 to
97cc478
Compare
|
/bot run |
|
PR_Github #43316 [ run ] triggered by Bot. Commit: |
|
PR_Github #43316 [ run ] completed with state
|
97cc478 to
fcf570a
Compare
|
/bot run |
|
PR_Github #43371 [ run ] triggered by Bot. Commit: |
|
PR_Github #43371 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #43397 [ run ] triggered by Bot. Commit: |
|
PR_Github #43397 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #43630 [ run ] triggered by Bot. Commit: |
fcf570a to
e0eec5b
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #43649 [ run ] triggered by Bot. Commit: |
|
PR_Github #43649 [ run ] completed with state
|
… serving The KvCacheAwareRouter needs to tokenize incoming ChatCompletionRequests to compute block hashes for cache-aware routing. Models like DeepSeek-V3.2 use a custom tokenizer whose HF tokenizer lacks a chat_template, causing apply_chat_template to fail. Add custom_tokenizer support to KvCacheAwareRouter: - Accept custom_tokenizer parameter (e.g. 'deepseek_v32') passed via router_args in the disagg server config - Load custom tokenizer via load_custom_tokenizer from tokenizer module - Handle apply_chat_template returning a string instead of token IDs (some custom tokenizers ignore the tokenize=True flag) Signed-off-by: Shicheng Li <shicli@nvidia.com>
e0eec5b to
4ab90e5
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #43762 [ run ] triggered by Bot. Commit: |
|
PR_Github #43762 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #44007 [ run ] triggered by Bot. Commit: |
|
/bot run --disable-fail-fast |
|
PR_Github #44101 [ run ] triggered by Bot. Commit: |
|
PR_Github #44101 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #44139 [ run ] triggered by Bot. Commit: |
|
PR_Github #44139 [ run ] completed with state |
Mirror the change made on `main` in PR NVIDIA#12990 ("Support custom_tokenizer in KvCacheAwareRouter for disagg serving"), which moved TOKENIZER_ALIASES from a method-local variable inside `validate_and_init_tokenizer` to a top-level module constant. This makes the dict importable by external code that needs to resolve the same short aliases — for example, the Dynamo TRT-LLM worker (`dynamo.trtllm.workers.llm_worker`, github.com/ai-dynamo/dynamo PR NVIDIA#8079) does: from tensorrt_llm.llmapi.llm_args import TOKENIZER_ALIASES so it can construct the V4 tokenizer ahead of LLM() init (Dynamo needs the tokenizer object early to populate SamplingParams defaults and build its logits processor). Without the lift, that import fails at module load time with `ImportError: cannot import name 'TOKENIZER_ALIASES'` and the worker crash-loops. Difference vs PR NVIDIA#12990 on main: the dict here also includes 'deepseek_v4' → DeepseekV4Tokenizer, since DeepSeek-V4 lives on this branch and isn't on main yet. The function body stays behavior-equivalent — `validate_and_init_tokenizer` now references the module-level constant instead of redefining its own copy. No semantic change for existing callers. Signed-off-by: Yuewei Na <nv-yna@users.noreply.github.com>
6 participants
Summary
The KvCacheAwareRouter needs to tokenize incoming ChatCompletionRequests to compute block hashes for cache-aware routing. Models like DeepSeek-V3.2 use a custom tokenizer whose HF tokenizer lacks
chat_template, causingapply_chat_templateto fail.Changes
router.py: Addcustom_tokenizerparameter toKvCacheAwareRouter(passed viarouter_argsin disagg server config). Resolve aliases usingTOKENIZER_ALIASESfromllm_args.py. Load custom tokenizer viafrom_pretrained. Handleapply_chat_templatereturning string instead of token IDs (some custom tokenizers ignoretokenize=True).llm_args.py: Use existing module-levelTOKENIZER_ALIASESin_setup_tokenizer(remove local duplicate).Usage
In disagg server config:
Test plan
ValueError: chat_template is not set→ 100% request failures🤖 Generated with Claude Code
Summary by CodeRabbit
New Features
deepseek_v32,glm_moe_dsa).Bug Fixes