Skip to content

[None][fix] Invalidate cached special-id set in slow convert_ids_to_tokens#16331

Draft
Shixiaowei02 wants to merge 1 commit into
NVIDIA:mainfrom
Shixiaowei02:user/xiaoweis/fix-tokenizer-special-ids-cache-staleness
Draft

[None][fix] Invalidate cached special-id set in slow convert_ids_to_tokens#16331
Shixiaowei02 wants to merge 1 commit into
NVIDIA:mainfrom
Shixiaowei02:user/xiaoweis/fix-tokenizer-special-ids-cache-staleness

Conversation

@Shixiaowei02

@Shixiaowei02 Shixiaowei02 commented Jul 13, 2026

Copy link
Copy Markdown
Collaborator

Description

This pull request improves the handling of special tokens in the TransformersTokenizer wrapper, ensuring that late-registered or reassigned special tokens are correctly excluded during token detokenization, particularly for slow tokenizers. The changes address a bug where the set of special token IDs could become stale, causing incorrect output, and introduce tests to verify the fix.

Special Token Handling Improvements

  • Replaced the static snapshot of special token IDs in TransformersTokenizer.__init__ with a dynamic, memoized cache that updates only when the set of special tokens changes, ensuring correct handling of late-registered or reassigned special tokens. ([1] [2])
  • Added the _special_token_ids method to efficiently and correctly compute the set of special token IDs, using a fingerprint to detect changes and avoid stale caches. (tensorrt_llm/tokenizer/tokenizer.pyL342-R376)

Testing

  • Added a minimal slow tokenizer implementation (_ToySlowTokenizer) and two regression tests to verify that convert_ids_to_tokens correctly reflects late-registered and reassigned special tokens, preventing stale caches from leaking special tokens into detokenized output. (tests/unittest/llmapi/test_tokenizer_decode.pyR20-R120)

Test Coverage

PR Checklist

Please review the following before submitting your PR:

  • PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.

  • PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.

  • Test cases are provided for new code paths (see test instructions)

  • If PR introduces API changes, an appropriate PR label is added - either api-compatible or api-breaking. For api-breaking, include BREAKING in the PR title.

  • Any new dependencies have been scanned for license and vulnerabilities

  • CODEOWNERS updated if ownership changes

  • Documentation updated as needed

  • Update tava architecture diagram if there is a significant design change in PR.

  • The reviewers assigned automatically/manually are appropriate for the PR.

  • Please check this after reviewing the above items as appropriate for this PR.

GitHub Bot Help

To see a list of available CI bot commands, please comment /bot help.

@Shixiaowei02

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #58990 [ run ] triggered by Bot. Commit: 4068de7 Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #58990 [ run ] completed with state SUCCESS. Commit: 4068de7
/LLM/main/L0_MergeRequest_PR pipeline #47520 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

  • Please check the failed tests and fix your PR
  • If you cannot view the failures, ask the CI triggerer to share details
  • Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

@Shixiaowei02
Shixiaowei02 force-pushed the user/xiaoweis/fix-tokenizer-special-ids-cache-staleness branch from 4068de7 to fcfa02f Compare July 16, 2026 06:18
@Shixiaowei02

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #59639 [ run ] triggered by Bot. Commit: fcfa02f Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #59639 [ run ] completed with state SUCCESS. Commit: fcfa02f
/LLM/main/L0_MergeRequest_PR pipeline #48076 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

  • Please check the failed tests and fix your PR
  • If you cannot view the failures, ask the CI triggerer to share details
  • Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

…okens

The slow-tokenizer optimization in convert_ids_to_tokens snapshotted
all_special_ids into a set at __init__ and reused it forever. If special
tokens were (re)registered after the wrapper was constructed (e.g.
add_special_tokens during chat-template / adapter setup), the snapshot went
stale: the affected ids stopped being skipped and leaked into detokenized
output when skip_special_tokens=True -- an accuracy regression.

Recomputing all_special_ids on every call is correct but a severe perf
regression: convert_tokens_to_ids(all_special_tokens) is expensive on slow
tokenizers and runs once per streamed token (measured ~10^2-10^4x slower,
superlinear in the number of special tokens). Instead keep the memoized set
and rebuild it only when a cheap O(1) fingerprint of the tokenizer's special
tokens changes: the count of registered added tokens (covers newly added
specials, including additional_special_tokens) plus the ids of the named
special tokens (covers reassigning/clearing bos/eos/pad/unk/sep/cls/mask).
Steady-state streaming stays O(1) per call while the cache can never be stale.

Add regression tests on a genuine slow tokenizer for a late-added special
token and a reassigned named special token.

Signed-off-by: xiaoweis <39303645+Shixiaowei02@users.noreply.github.com>
@Shixiaowei02
Shixiaowei02 force-pushed the user/xiaoweis/fix-tokenizer-special-ids-cache-staleness branch from fcfa02f to 00925b2 Compare July 17, 2026 06:50
@Shixiaowei02

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #59912 [ run ] triggered by Bot. Commit: 00925b2 Link to invocation

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants