[None][perf] Add boundary-aware incremental router tokenization - #17462
Draft
lishicheng1996-nv wants to merge 5 commits into
Draft
[None][perf] Add boundary-aware incremental router tokenization#17462lishicheng1996-nv wants to merge 5 commits into
lishicheng1996-nv wants to merge 5 commits into
Conversation
Signed-off-by: Shicheng Li <shicli@nvidia.com>
Contributor
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (2)
🚧 Files skipped from review as they are similar to previous changes (2)
Walkthrough
ChangesIncremental tokenization
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant ChatRequest
participant OpenAIServer
participant InputProcessingExecutor
participant IncrementalTokenizationCache
participant OffsetTokenizer
ChatRequest->>OpenAIServer: provide conversation and messages
OpenAIServer->>OpenAIServer: resolve conversation ID and render prompt
OpenAIServer->>InputProcessingExecutor: submit eligible text prompt
InputProcessingExecutor->>IncrementalTokenizationCache: encode rendered prompt
IncrementalTokenizationCache->>OffsetTokenizer: encode changed suffix with offsets
OffsetTokenizer-->>IncrementalTokenizationCache: return token IDs and offsets
IncrementalTokenizationCache-->>OpenAIServer: return prompt token IDs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 4
🤖 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.
Inline comments:
In `@tensorrt_llm/serve/router_utils.py`:
- Around line 320-322: Define a structural tokenizer Protocol near
_encode_with_offsets with encode() and offset-aware __call__() signatures
matching the operations used by the function, then annotate the tokenizer
parameter with that Protocol. Preserve the existing return type and encoding
behavior.
- Around line 225-228: Update the logger calls in
tensorrt_llm/serve/router_utils.py at lines 225-228 and 305-309 to preformat
their messages as single f-strings, embedding
self._incremental_tokenize_rollback in the first and key in the second; do not
pass printf-style format arguments separately.
- Around line 279-281: Update the cut-point calculation in the incremental
tokenization flow so `cut_char` is explicitly set to zero when `cut_token == 0`;
otherwise continue using `previous_offsets[cut_token][0]`. Ensure suffix
encoding includes the entire rendered string when no cached token is reusable.
In `@tests/unittest/disaggregated/test_router.py`:
- Around line 2387-2403: Expand coverage around
KvCacheAwareRouter._encode_with_prefix_cache beyond
test_prefix_cache_rolls_back_boundary_token: add cases for unavailable offsets,
invalid offsets, periodic verification, and an initial offset greater than zero,
asserting correct token results and tokenizer calls. Add
unittest/disaggregated/test_router.py to the appropriate CI test list under
tests/integration/test_lists/test-db/ or qa/, and run pytest tests/unittest/.
🪄 Autofix
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: Enterprise
Run ID: bc3aed6d-5ce2-4ba5-886c-940782266563
📒 Files selected for processing (2)
tensorrt_llm/serve/router_utils.pytests/unittest/disaggregated/test_router.py
Signed-off-by: Shicheng Li <shicli@nvidia.com>
Signed-off-by: Shicheng Li <shicli@nvidia.com>
Contributor
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)
tests/unittest/disaggregated/test_router.py (1)
2456-2482: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winCover a non-default rollback value.
Line 2463 enables incremental tokenization but does not set
TRTLLM_INCREMENTAL_TOKENIZE_ROLLBACK. This test validates only the default one-token rollback. If configuration parsing or application fails, the current suite can pass.Add a case with a non-default rollback such as
2. Assert canonical token IDs and the earlier suffix passed to the tokenizer.🤖 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 `@tests/unittest/disaggregated/test_router.py` around lines 2456 - 2482, Extend test_prefix_cache_rolls_back_boundary_token to configure a non-default TRTLLM_INCREMENTAL_TOKENIZE_ROLLBACK value such as 2, then assert tokenizer.encode(current) remains the canonical result and tokenizer.calls contains the suffix beginning at the expected two-token rollback boundary. Keep the existing prefix-cache setup and verify the earlier suffix passed to _encode_with_prefix_cache, ensuring the configured rollback is actually applied rather than only relying on the default.
🤖 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 `@tests/unittest/disaggregated/test_router.py`:
- Around line 2456-2482: Extend test_prefix_cache_rolls_back_boundary_token to
configure a non-default TRTLLM_INCREMENTAL_TOKENIZE_ROLLBACK value such as 2,
then assert tokenizer.encode(current) remains the canonical result and
tokenizer.calls contains the suffix beginning at the expected two-token rollback
boundary. Keep the existing prefix-cache setup and verify the earlier suffix
passed to _encode_with_prefix_cache, ensuring the configured rollback is
actually applied rather than only relying on the default.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: da4f1e47-10fa-4117-bd36-6a1b46706f86
📒 Files selected for processing (2)
tensorrt_llm/serve/router_utils.pytests/unittest/disaggregated/test_router.py
🚧 Files skipped from review as they are similar to previous changes (1)
- tensorrt_llm/serve/router_utils.py
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (1)
tensorrt_llm/serve/chat_tokenization.py (1)
72-75: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueUse Python 3.10 annotation syntax in new code.
tensorrt_llm/serve/chat_tokenization.py#L72-L75: replaceOptional[list[tuple[int, int]]]withlist[tuple[int, int]] | None.tensorrt_llm/serve/chat_tokenization.py#L177-L177: replaceOptional[...]in the return type with... | None.tensorrt_llm/serve/openai_server.py#L1464-L1467: replaceUnion[str, List[int]]withstr | list[int].As per coding guidelines, use Python 3.10+ and prefer built-in generics and
|. Retrieved learnings confirm this repository supports Python 3.10+ syntax.🤖 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/serve/chat_tokenization.py` around lines 72 - 75, Update the type annotations in tensorrt_llm/serve/chat_tokenization.py lines 72-75 and 177 to use Python 3.10 union syntax, replacing Optional[...] with the equivalent | None form; update tensorrt_llm/serve/openai_server.py lines 1464-1467 to replace Union[str, List[int]] with str | list[int].Sources: Coding guidelines, Learnings
🤖 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.
Inline comments:
In `@tensorrt_llm/serve/chat_tokenization.py`:
- Around line 146-151: Preformat all three affected logger messages as single
f-strings because tensorrt_llm.logger does not interpolate printf-style
arguments: update the mismatch log in tensorrt_llm/serve/chat_tokenization.py
lines 146-151, the hit-count log in tensorrt_llm/serve/chat_tokenization.py
lines 166-171, and the server log in tensorrt_llm/serve/openai_server.py lines
338-340. In the hit-count condition, require a nonzero hit count before applying
the periodic logging check so zero-hit requests do not flood the logs.
- Around line 106-114: Replace the binary-search prefix computation in the
surrounding tokenization flow with a single-pass comparison that advances
through matching characters in previous_text and rendered, avoiding repeated
growing slices. Preserve the resulting common_prefix_chars value and existing
suffix tokenization behavior.
---
Nitpick comments:
In `@tensorrt_llm/serve/chat_tokenization.py`:
- Around line 72-75: Update the type annotations in
tensorrt_llm/serve/chat_tokenization.py lines 72-75 and 177 to use Python 3.10
union syntax, replacing Optional[...] with the equivalent | None form; update
tensorrt_llm/serve/openai_server.py lines 1464-1467 to replace Union[str,
List[int]] with str | list[int].
🪄 Autofix
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: Enterprise
Run ID: f078bb57-a1b2-49ef-908e-2773833adeba
📒 Files selected for processing (4)
tensorrt_llm/serve/chat_tokenization.pytensorrt_llm/serve/openai_server.pytensorrt_llm/serve/router_utils.pytests/unittest/disaggregated/test_router.py
🚧 Files skipped from review as they are similar to previous changes (2)
- tensorrt_llm/serve/router_utils.py
- tests/unittest/disaggregated/test_router.py
Comment on lines
+106
to
+114
| prefix_limit = min(len(previous_text), len(rendered)) | ||
| low, high = 0, prefix_limit | ||
| while low < high: | ||
| middle = (low + high + 1) // 2 | ||
| if previous_text[:middle] == rendered[:middle]: | ||
| low = middle | ||
| else: | ||
| high = middle - 1 | ||
| common_prefix_chars = low |
Contributor
There was a problem hiding this comment.
🚀 Performance & Scalability | 🟠 Major | ⚡ Quick win
Compute the common prefix in one pass.
Lines 106-114 allocate and compare increasingly large string slices. A long shared prompt performs Θ(n log n) copied-character work before suffix tokenization. This can reduce the cache speedup for long conversations.
Proposed fix
- prefix_limit = min(len(previous_text), len(rendered))
- low, high = 0, prefix_limit
- while low < high:
- middle = (low + high + 1) // 2
- if previous_text[:middle] == rendered[:middle]:
- low = middle
- else:
- high = middle - 1
- common_prefix_chars = low
+ common_prefix_chars = 0
+ for previous_char, rendered_char in zip(previous_text, rendered):
+ if previous_char != rendered_char:
+ break
+ common_prefix_chars += 1📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| prefix_limit = min(len(previous_text), len(rendered)) | |
| low, high = 0, prefix_limit | |
| while low < high: | |
| middle = (low + high + 1) // 2 | |
| if previous_text[:middle] == rendered[:middle]: | |
| low = middle | |
| else: | |
| high = middle - 1 | |
| common_prefix_chars = low | |
| common_prefix_chars = 0 | |
| for previous_char, rendered_char in zip(previous_text, rendered): | |
| if previous_char != rendered_char: | |
| break | |
| common_prefix_chars += 1 |
🤖 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/serve/chat_tokenization.py` around lines 106 - 114, Replace the
binary-search prefix computation in the surrounding tokenization flow with a
single-pass comparison that advances through matching characters in
previous_text and rendered, avoiding repeated growing slices. Preserve the
resulting common_prefix_chars value and existing suffix tokenization behavior.
lishicheng1996-nv
marked this pull request as draft
August 10, 2026 12:05
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
Real AgentX case
Qwen3.5's chat template does not preserve the previous rendered prompt as a literal prefix across AgentX turns. Before generation, a prompt ends with a synthetic generation prefix:
When the next turn is rendered, the actual assistant response replaces that final
<think>\n, followed by the next user turn and a fresh generation prefix:The character LCP therefore stops immediately before the old
<think>\n. In 12 inspected same-key, non-literal-prefix transitions, the complete unmatched old tail was exactly the 8-character<think>\n; across all 1,087 consecutive transitions, the median changed old tail was also 8 characters.This defeats the exact-prefix method from #15040. In the AgentX v41 trace:
Why rollback is required
Starting suffix tokenization exactly at the character LCP is not generally composable: a BPE token can merge characters from both sides of that boundary. The implementation maps the LCP to the cached offset boundary and rolls back one complete token before retokenizing.
On the full AgentX transition set, cutting without rollback matched canonical tokenization for 1,085/1,087 transitions. One-token rollback fixed both mismatches and matched 1,087/1,087.
Validation
Qwen3.5 AgentX v41 trace on an 18-CPU frontend allocation:
Additional validation:
<think>\nrewrite, a cross-boundary merge, first-token offsets starting after character zero, unavailable/invalid offsets, and periodic verification fallback.tests/integration/test_lists/test-db/l0_cpu.ymlalready runs the entireunittest/disaggregateddirectory.Configuration
Enable with
TRTLLM_INCREMENTAL_TOKENIZE=1.Optional safeguards:
TRTLLM_INCREMENTAL_TOKENIZE_ROLLBACK_TOKENScontrols the number of cached tokens rolled back; the minimum is one.TRTLLM_INCREMENTAL_TOKENIZE_VERIFY_EVERYperiodically compares incremental output with canonical full tokenization and falls back on mismatch.Dev Engineer Review
OffsetTokenizerandIncrementalTokenizationCacheAPIs.CODING_GUIDELINES.md.QA Engineer Review
tests/unittest/disaggregated/test_router.py.test-db/orqa/coverage was identified.