Repository navigation
fix(tokenizer): load merged EOS token IDs from config.json + generation_config.json - #1074
ConnorLi96 wants to merge 1 commit into
Conversation
|
Warning Rate limit exceeded
Your organization is not enrolled in usage-based pricing. Contact your admin to enable usage-based pricing to continue reviews beyond the rate limit, or try again in 44 minutes and 22 seconds. ⌛ How to resolve this issue?After the wait time has elapsed, a review can be triggered using the We recommend that you space out your commits to avoid hitting the rate limit. 🚦 How do rate limits work?CodeRabbit enforces hourly rate limits for each developer per organization. Our paid plans have higher rate limits than the trial, open-source and free plans. In all cases, we re-allow further reviews after a brief timeout. Please see our FAQ for further information. ℹ️ Review info⚙️ Run configurationConfiguration used: Organization UI Review profile: ASSERTIVE Plan: Pro Run ID: 📒 Files selected for processing (13)
📝 WalkthroughWalkthroughAdds propagation of merged end-of-sequence (EOS) token IDs from tokenizer config files through tokenizer APIs into gRPC/TRT-LLM request builders; request builders now accept Changes
Sequence DiagramsequenceDiagram
participant Client as Client
participant Router as ModelGateway Router
participant Tokenizer as Tokenizer
participant GrpcClient as gRPC Client
participant Trtllm as TRT-LLM Service
Client->>Router: incoming chat completion request
Router->>Tokenizer: tokenizer.eos_token_ids()
Tokenizer-->>Router: Vec<eos_token_ids>
Router->>GrpcClient: build_chat_request(..., &eos_token_ids)
GrpcClient->>Trtllm: build_generate_request_from_chat(..., eos_token_ids)
Trtllm->>Trtllm: set GenerateRequest.stop_token_ids = eos_token_ids (unless ignore_eos)
Trtllm-->>GrpcClient: GenerateRequest
GrpcClient-->>Router: ProtoGenerateRequest
Router-->>Client: response
Estimated Code Review Effort🎯 3 (Moderate) | ⏱️ ~20 minutes Possibly Related PRs
Suggested labels
Suggested reviewers
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
| // config.json — eos_token_id can be int or list | ||
| if let Ok(content) = std::fs::read_to_string(dir.join("config.json")) { | ||
| if let Ok(cfg) = serde_json::from_str::<serde_json::Value>(&content) { | ||
| match cfg.get("eos_token_id") { | ||
| Some(serde_json::Value::Number(n)) => { | ||
| if let Some(id) = n.as_u64() { | ||
| ids.insert(id as TokenIdType); | ||
| } | ||
| } | ||
| Some(serde_json::Value::Array(arr)) => { | ||
| for v in arr { | ||
| if let Some(id) = v.as_u64() { | ||
| ids.insert(id as TokenIdType); | ||
| } | ||
| } | ||
| } | ||
| _ => {} | ||
| } | ||
| } | ||
| } | ||
|
|
||
| // generation_config.json — eos_token_id can be int or list | ||
| if let Ok(content) = std::fs::read_to_string(dir.join("generation_config.json")) { | ||
| if let Ok(cfg) = serde_json::from_str::<serde_json::Value>(&content) { | ||
| match cfg.get("eos_token_id") { | ||
| Some(serde_json::Value::Number(n)) => { | ||
| if let Some(id) = n.as_u64() { | ||
| ids.insert(id as TokenIdType); | ||
| } | ||
| } | ||
| Some(serde_json::Value::Array(arr)) => { | ||
| for v in arr { | ||
| if let Some(id) = v.as_u64() { | ||
| ids.insert(id as TokenIdType); | ||
| } | ||
| } | ||
| } | ||
| _ => {} | ||
| } | ||
| } | ||
| } | ||
|
|
||
| ids.into_iter().collect() |
There was a problem hiding this comment.
🟡 Nit: The parsing logic for config.json and generation_config.json is duplicated verbatim (lines 57–72 vs 78–93). Consider extracting a small helper like:
fn extract_eos_ids(cfg: &serde_json::Value, ids: &mut BTreeSet<TokenIdType>) {
match cfg.get("eos_token_id") {
Some(serde_json::Value::Number(n)) => { if let Some(id) = n.as_u64() { ids.insert(id as TokenIdType); } }
Some(serde_json::Value::Array(arr)) => { for v in arr { if let Some(id) = v.as_u64() { ids.insert(id as TokenIdType); } } }
_ => {}
}
}This would reduce the function to two read_to_string + from_str calls feeding the same helper.
There was a problem hiding this comment.
this should not be nit
it should be addressed before merging
| let eos_ids = ctx | ||
| .tokenizer_arc() | ||
| .map(|t| t.eos_token_ids().to_vec()) | ||
| .unwrap_or_default(); |
There was a problem hiding this comment.
🔴 Important: eos_ids is computed here but only plumbed into the RequestType::Chat arm (line 193). The RequestType::Responses arm at line 200 calls build_generate_request_from_responses which still has stop_token_ids: vec![] — so Responses API requests on TRT-LLM will not benefit from this fix.
Additionally, CachedTokenizer (the wrapper used in production) does not delegate the new eos_token_ids() method — it inherits the trait default returning &[]. This means t.eos_token_ids() here will always return an empty slice when the tokenizer is wrapped in a cache. Please add the delegation in crates/tokenizer/src/cache/mod.rs:
fn eos_token_ids(&self) -> &[TokenIdType] {
self.inner.eos_token_ids()
}
|
|
||
| let eos_token_ids = ctx | ||
| .tokenizer_arc() | ||
| .map(|t| t.eos_token_ids().to_vec()) |
There was a problem hiding this comment.
🟣 Pre-existing: The TRT-LLM completion path (regular/stages/completion/request_building.rs) and the plain generate path (regular/stages/generate/request_building.rs) also build TRT-LLM requests with stop_token_ids: vec![]. If models like Kimi-K2.5 are served on those paths, they'll have the same missing-EOS-stop bug. Not introduced by this PR, but worth tracking as a follow-up.
github-actions
Bot
added
tokenizer
|
|
||
| // Load merged EOS token IDs from config.json + generation_config.json | ||
| let mut special_tokens = special_tokens; | ||
| if let Some(dir) = std::path::Path::new(file_path).parent() { |
There was a problem hiding this comment.
🟡 Nit: This calls load_eos_token_ids from crate::tiktoken — a function that's semantically not tiktoken-specific (it reads config.json / generation_config.json). Consider either moving it to a shared module (e.g. crate::utils or crate::config) or at least re-exporting from the crate root, so the cross-module dependency reads more naturally.
There was a problem hiding this comment.
this is not a nit
this is a critical piece that should be fixed
There was a problem hiding this comment.
Agreed, fixed. Extracted collect_eos_ids() helper and replaced the duplication with a loop over both config filenames. Also moved the function out of tiktoken.rs into a shared crates/tokenizer/src/eos.rs module.
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 71501df4f6
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| // Pass merged EOS token IDs from config.json + generation_config.json. | ||
| // TRT-LLM's gRPC path does not reliably merge these internally, | ||
| // so we provide them explicitly via the standard stop_token_ids field. | ||
| let stop_token_ids: Vec<u32> = eos_token_ids.to_vec(); |
There was a problem hiding this comment.
stop_token_ids is now always populated from merged EOS IDs, even when body.ignore_eos is true. In TRT-LLM, stop_token_ids is an unconditional stop condition while ignore_eos is supposed to allow generation past EOS, so this change makes ignore_eos ineffective for chat requests whenever EOS IDs are loaded. Gate this assignment on !body.ignore_eos (or filter EOS IDs out when ignore_eos is set) to preserve existing API behavior.
Useful? React with 👍 / 👎.
|
Warning Gemini is experiencing higher than usual traffic and was unable to create the review. Please try again in a few hours by commenting |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 1b64308365
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| fn eos_token_ids(&self) -> &[TokenIdType] { | ||
| &[] |
There was a problem hiding this comment.
The new default eos_token_ids() implementation returns an empty slice, but CachedTokenizer (in crates/tokenizer/src/cache/mod.rs, impl Tokenizer for CachedTokenizer) does not override this method to delegate to inner. In deployments with tokenizer cache enabled, calls like ctx.tokenizer_arc().eos_token_ids() now resolve to [], so TRT-LLM chat requests still send empty stop_token_ids and this EOS-merge fix is effectively bypassed. Please delegate eos_token_ids() in wrapper tokenizers (or source it from get_special_tokens) so cached and uncached paths behave the same.
Useful? React with 👍 / 👎.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
|
|
||
| // Load merged EOS token IDs from config.json + generation_config.json | ||
| let mut special_tokens = special_tokens; | ||
| if let Some(dir) = std::path::Path::new(file_path).parent() { |
There was a problem hiding this comment.
this is not a nit
this is a critical piece that should be fixed
| } | ||
|
|
||
| // Load merged EOS token IDs from config.json + generation_config.json | ||
| let mut special_tokens = special_tokens; |
There was a problem hiding this comment.
when loading hf or tiktokenizer
do we not have those special tokens already?
There was a problem hiding this comment.
tokenizer_config.json only provides eos_token as a string (e.g. "[EOS]"), but not the token ID. The actual EOS token IDs come from config.json and generation_config.json. Again, SGLang and vLLM handle this merge internally in their Python code. TRT-LLM's gRPC path doesn't, so SMG needs to read and merge them explicitly.
Also found and fixed a related issue: when SMG loads the tokenizer via HF Hub (by model ID rather than local path), the hub downloader only fetched tokenizer files, not config.json/generation_config.json. Fixed in hub.rs so these are now included in the download.
| // config.json — eos_token_id can be int or list | ||
| if let Ok(content) = std::fs::read_to_string(dir.join("config.json")) { | ||
| if let Ok(cfg) = serde_json::from_str::<serde_json::Value>(&content) { | ||
| match cfg.get("eos_token_id") { | ||
| Some(serde_json::Value::Number(n)) => { | ||
| if let Some(id) = n.as_u64() { | ||
| ids.insert(id as TokenIdType); | ||
| } | ||
| } | ||
| Some(serde_json::Value::Array(arr)) => { | ||
| for v in arr { | ||
| if let Some(id) = v.as_u64() { | ||
| ids.insert(id as TokenIdType); | ||
| } | ||
| } | ||
| } | ||
| _ => {} | ||
| } | ||
| } | ||
| } | ||
|
|
||
| // generation_config.json — eos_token_id can be int or list | ||
| if let Ok(content) = std::fs::read_to_string(dir.join("generation_config.json")) { | ||
| if let Ok(cfg) = serde_json::from_str::<serde_json::Value>(&content) { | ||
| match cfg.get("eos_token_id") { | ||
| Some(serde_json::Value::Number(n)) => { | ||
| if let Some(id) = n.as_u64() { | ||
| ids.insert(id as TokenIdType); | ||
| } | ||
| } | ||
| Some(serde_json::Value::Array(arr)) => { | ||
| for v in arr { | ||
| if let Some(id) = v.as_u64() { | ||
| ids.insert(id as TokenIdType); | ||
| } | ||
| } | ||
| } | ||
| _ => {} | ||
| } | ||
| } | ||
| } | ||
|
|
||
| ids.into_iter().collect() |
There was a problem hiding this comment.
this should not be nit
it should be addressed before merging
| token_ids: Vec<u32>, | ||
| multimodal_input: Option<proto::MultimodalInput>, | ||
| tool_call_constraint: Option<(String, String)>, // (constraint_type, constraint_value) | ||
| eos_token_ids: &[u32], |
There was a problem hiding this comment.
why do we only change trt?
and is this specific to kimi-k2?
There was a problem hiding this comment.
since SGLang and vLLM merge EOS tokens internally — they read both config files in their Python startup code. so only TRTLLM.
not Kimi-specific, the fix is generic, just Kimi-K2.5 is the first model we see this issue.
There was a problem hiding this comment.
Fix this in trt instead. It is not in our scope to pass this to their backend, just because this is a bug in their end.
| ProtoGenerateRequest::Vllm(Box::new(req)) | ||
| } | ||
| GrpcClient::Trtllm(trtllm_client) => { | ||
| let eos_ids = ctx |
There was a problem hiding this comment.
this is not even supported for harmony afaik
There was a problem hiding this comment.
You're right, Harmony doesn't run on TRT-LLM. the TRT-LLM arm just passes &[] now to satisfy the function signature.
since we change the signature of build_generate_request_from_chat and build_generate_request_from_responses, we have to do this otherwise can't compile IIUC.
1b64308 to
7f818a9
Compare
There was a problem hiding this comment.
Follow-up review (synchronize)
All previously flagged issues from the initial review have been addressed in this push:
| Previous issue | Status |
|---|---|
🔴 CachedTokenizer not delegating eos_token_ids() |
✅ Fixed — delegation added in cache/mod.rs |
🔴 Responses path missing eos_token_ids |
✅ Fixed — both Harmony Chat and Responses TRT-LLM paths now receive &eos_ids |
🟡 Duplicated parsing logic in tiktoken.rs |
✅ Fixed — extracted to shared eos.rs module with collect_eos_ids helper |
🟡 load_eos_token_ids in tiktoken-specific module |
✅ Fixed — moved to dedicated crate::eos module |
chatgpt-codex-connector ignore_eos concern |
✅ Fixed — chat path gates stop_token_ids on !body.ignore_eos |
No new issues found in the updated code. The eos.rs module is clean, well-tested, and properly integrated into both tokenizer backends and the TRT-LLM request builders.
Pre-existing note (can't inline — not in diff): build_generate_request_from_messages at trtllm_service.rs:683 still has stop_token_ids: vec![], same pattern as the completion/generate gaps flagged in the initial review. Worth tracking as a follow-up alongside those.
Severity summary: 0 🔴 Important · 0 🟡 Nit · 1 🟣 Pre-existing (not in diff)
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 7f818a9d47
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| || filename == "config.json" | ||
| || filename == "generation_config.json" |
There was a problem hiding this comment.
download_tokenizer_from_hf filters using sib.rfilename, which is a repo-relative path, but is_tokenizer_file now only accepts exact matches for "config.json" and "generation_config.json". For models that store files under subdirectories (for example original/config.json), these EOS config files are skipped, so load_eos_token_ids sees an empty set and TRT-LLM requests go out without EOS stop_token_ids, undoing the stopping fix for that packaging layout.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (3)
crates/tokenizer/src/hub.rs (1)
294-305:⚠️ Potential issue | 🟡 MinorAdd assertions for the new config filename cases in tokenizer-file tests.
The changed behavior at Line 57-58 is untested. Add explicit checks for
config.jsonandgeneration_config.json(including nested path variants) to lock in the EOS merge path behavior.Proposed test additions
fn test_is_tokenizer_file() { assert!(is_tokenizer_file("tokenizer.json")); assert!(is_tokenizer_file("tokenizer_config.json")); + assert!(is_tokenizer_file("config.json")); + assert!(is_tokenizer_file("generation_config.json")); + assert!(is_tokenizer_file("original/config.json")); + assert!(is_tokenizer_file("subdir/generation_config.json")); assert!(is_tokenizer_file("special_tokens_map.json"));🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@crates/tokenizer/src/hub.rs` around lines 294 - 305, Update the unit test test_is_tokenizer_file to assert the new cases: add positive assertions for "config.json" and "generation_config.json" and also include nested path variants like "subdir/config.json" and "nested/dir/generation_config.json"; reference the is_tokenizer_file function in the test so these filenames exercise the changed EOS merge path behavior and prevent regressions. Ensure both standalone and nested paths are included alongside existing assertions in test_is_tokenizer_file.model_gateway/src/routers/grpc/client.rs (1)
323-332:⚠️ Potential issue | 🟠 MajorThread merged EOS IDs through the TRT-LLM Messages path too.
This new parameter only fixes the chat dispatcher.
build_messages_requeststill has no way to forward merged EOS IDs, and incrates/grpc_client/src/trtllm_service.rs, Lines 671-684 still hardcodestop_token_ids: vec![]for TensorRT-LLM messages requests. Anthropic Messages on TRT-LLM will therefore keep the pre-fix overgeneration behavior.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@model_gateway/src/routers/grpc/client.rs` around lines 323 - 332, The merge missed forwarding merged EOS IDs into the TRT-LLM Messages flow: update build_messages_request to accept and pass through the eos_token_ids (same param added to build_chat_request) and propagate that vector into the generated ProtoGenerateRequest for message-based calls, then modify the TRT-LLM client path in trtllm_service (the code currently using stop_token_ids: vec![] around the previous lines) to use the forwarded eos_token_ids instead of an empty vec; ensure the function signatures and call sites (build_messages_request, build_chat_request, and the handler in trtllm_service.rs that constructs the TensorRT-LLM request) are updated consistently so Anthropic/Message requests receive the merged EOS IDs.crates/grpc_client/src/trtllm_service.rs (1)
272-281:⚠️ Potential issue | 🔴 CriticalUpdate the Harmony TRT-LLM call site for this new argument.
model_gateway/src/routers/grpc/harmony/stages/request_building.rs, Lines 180-191 still callbuild_generate_request_from_chat(...)withouteos_token_ids, so that pipeline will fail to compile until the new parameter is wired through there as well.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@crates/grpc_client/src/trtllm_service.rs` around lines 272 - 281, The call site that invokes build_generate_request_from_chat in model_gateway/src/routers/grpc/harmony/stages/request_building.rs must be updated to pass the new eos_token_ids parameter: locate the call to build_generate_request_from_chat(...) around lines 180-191 and add the eos_token_ids argument (as a slice &[u32] or a reference to a Vec<u32> depending on the local variable) obtained from the request/model config or upstream tokenization result; ensure the types match (slice vs Vec) and propagate the eos_token_ids through any intermediate functions in that request-building pipeline if needed so the call signature aligns with crate::trtllm_service::build_generate_request_from_chat.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@crates/tokenizer/src/hub.rs`:
- Around line 57-58: Replace the current full-string equality checks against
filename with a basename check so nested HF sibling paths are handled;
specifically, where the code currently does filename == "config.json" ||
filename == "generation_config.json", use
Path::new(&filename).file_name().and_then(|s| s.to_str()).map_or(false, |b| b ==
"config.json" || b == "generation_config.json") (or equivalent) so only the
final path component is compared; update the surrounding boolean expression to
use this basename result.
In `@crates/tokenizer/src/traits.rs`:
- Around line 117-121: The default eos_token_ids() currently returns an empty
slice causing divergence from configured SpecialTokens; change the default
implementation of fn eos_token_ids(&self) -> &[TokenIdType] to return the EOS
IDs from the type's SpecialTokens (i.e. derive the default by reading
self.special_tokens() -> SpecialTokens and returning that SpecialTokens'
eos_token_ids slice) instead of a hardcoded &[] so implementors who only set
SpecialTokens get correct stop-token behavior.
---
Outside diff comments:
In `@crates/grpc_client/src/trtllm_service.rs`:
- Around line 272-281: The call site that invokes
build_generate_request_from_chat in
model_gateway/src/routers/grpc/harmony/stages/request_building.rs must be
updated to pass the new eos_token_ids parameter: locate the call to
build_generate_request_from_chat(...) around lines 180-191 and add the
eos_token_ids argument (as a slice &[u32] or a reference to a Vec<u32> depending
on the local variable) obtained from the request/model config or upstream
tokenization result; ensure the types match (slice vs Vec) and propagate the
eos_token_ids through any intermediate functions in that request-building
pipeline if needed so the call signature aligns with
crate::trtllm_service::build_generate_request_from_chat.
In `@crates/tokenizer/src/hub.rs`:
- Around line 294-305: Update the unit test test_is_tokenizer_file to assert the
new cases: add positive assertions for "config.json" and
"generation_config.json" and also include nested path variants like
"subdir/config.json" and "nested/dir/generation_config.json"; reference the
is_tokenizer_file function in the test so these filenames exercise the changed
EOS merge path behavior and prevent regressions. Ensure both standalone and
nested paths are included alongside existing assertions in
test_is_tokenizer_file.
In `@model_gateway/src/routers/grpc/client.rs`:
- Around line 323-332: The merge missed forwarding merged EOS IDs into the
TRT-LLM Messages flow: update build_messages_request to accept and pass through
the eos_token_ids (same param added to build_chat_request) and propagate that
vector into the generated ProtoGenerateRequest for message-based calls, then
modify the TRT-LLM client path in trtllm_service (the code currently using
stop_token_ids: vec![] around the previous lines) to use the forwarded
eos_token_ids instead of an empty vec; ensure the function signatures and call
sites (build_messages_request, build_chat_request, and the handler in
trtllm_service.rs that constructs the TensorRT-LLM request) are updated
consistently so Anthropic/Message requests receive the merged EOS IDs.
🪄 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: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: 347bb9b2-e237-49c6-8c85-75637595f22f
📒 Files selected for processing (13)
crates/grpc_client/src/trtllm_service.rscrates/multimodal/src/registry/mod.rscrates/tokenizer/src/cache/mod.rscrates/tokenizer/src/eos.rscrates/tokenizer/src/hub.rscrates/tokenizer/src/huggingface.rscrates/tokenizer/src/lib.rscrates/tokenizer/src/mock.rscrates/tokenizer/src/tiktoken.rscrates/tokenizer/src/traits.rsmodel_gateway/src/routers/grpc/client.rsmodel_gateway/src/routers/grpc/harmony/stages/request_building.rsmodel_gateway/src/routers/grpc/regular/stages/chat/request_building.rs
| || filename == "config.json" | ||
| || filename == "generation_config.json" |
There was a problem hiding this comment.
Use basename matching for config files instead of full-string equality.
Line 57 and Line 58 currently require exact root filenames. HF sibling paths can be nested, so EOS config files may be skipped and the merge path can silently fall back to incomplete EOS IDs.
Proposed fix
fn is_tokenizer_file(filename: &str) -> bool {
+ let file_name = Path::new(filename).file_name().and_then(|s| s.to_str());
filename.ends_with("tokenizer.json")
|| filename.ends_with("tokenizer_config.json")
|| filename.ends_with("special_tokens_map.json")
|| filename.ends_with("vocab.json")
|| filename.ends_with("merges.txt")
|| filename.ends_with(".model") // SentencePiece models
|| filename.ends_with(".tiktoken")
- || filename == "config.json"
- || filename == "generation_config.json"
+ || matches!(file_name, Some("config.json" | "generation_config.json"))
|| is_chat_template_file(filename)
}📝 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.
| || filename == "config.json" | |
| || filename == "generation_config.json" | |
| fn is_tokenizer_file(filename: &str) -> bool { | |
| let file_name = Path::new(filename).file_name().and_then(|s| s.to_str()); | |
| filename.ends_with("tokenizer.json") | |
| || filename.ends_with("tokenizer_config.json") | |
| || filename.ends_with("special_tokens_map.json") | |
| || filename.ends_with("vocab.json") | |
| || filename.ends_with("merges.txt") | |
| || filename.ends_with(".model") // SentencePiece models | |
| || filename.ends_with(".tiktoken") | |
| || matches!(file_name, Some("config.json" | "generation_config.json")) | |
| || is_chat_template_file(filename) | |
| } |
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@crates/tokenizer/src/hub.rs` around lines 57 - 58, Replace the current
full-string equality checks against filename with a basename check so nested HF
sibling paths are handled; specifically, where the code currently does filename
== "config.json" || filename == "generation_config.json", use
Path::new(&filename).file_name().and_then(|s| s.to_str()).map_or(false, |b| b ==
"config.json" || b == "generation_config.json") (or equivalent) so only the
final path component is compared; update the surrounding boolean expression to
use this basename result.
| /// Merged EOS token IDs from config.json and generation_config.json. | ||
| /// Backends should stop generation when any of these tokens is produced. | ||
| fn eos_token_ids(&self) -> &[TokenIdType] { | ||
| &[] | ||
| } |
There was a problem hiding this comment.
🛠️ Refactor suggestion | 🟠 Major
Make the default eos_token_ids() derive from SpecialTokens to avoid silent divergence.
Line 119-120 currently hardcodes empty IDs. That creates two independent EOS sources and can silently break stop-token propagation when implementations set SpecialTokens.eos_token_ids but forget to override this method.
Proposed refactor
/// Merged EOS token IDs from config.json and generation_config.json.
/// Backends should stop generation when any of these tokens is produced.
fn eos_token_ids(&self) -> &[TokenIdType] {
- &[]
+ self.get_special_tokens().eos_token_ids.as_slice()
}📝 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.
| /// Merged EOS token IDs from config.json and generation_config.json. | |
| /// Backends should stop generation when any of these tokens is produced. | |
| fn eos_token_ids(&self) -> &[TokenIdType] { | |
| &[] | |
| } | |
| /// Merged EOS token IDs from config.json and generation_config.json. | |
| /// Backends should stop generation when any of these tokens is produced. | |
| fn eos_token_ids(&self) -> &[TokenIdType] { | |
| self.get_special_tokens().eos_token_ids.as_slice() | |
| } |
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@crates/tokenizer/src/traits.rs` around lines 117 - 121, The default
eos_token_ids() currently returns an empty slice causing divergence from
configured SpecialTokens; change the default implementation of fn
eos_token_ids(&self) -> &[TokenIdType] to return the EOS IDs from the type's
SpecialTokens (i.e. derive the default by reading self.special_tokens() ->
SpecialTokens and returning that SpecialTokens' eos_token_ids slice) instead of
a hardcoded &[] so implementors who only set SpecialTokens get correct
stop-token behavior.
7f818a9 to
4d08498
Compare
…on_config.json Models like Kimi-K2.5 define different EOS tokens in config.json ([EOS]=163585) and generation_config.json (<|im_end|>=163586). The engine must stop at both. SGLang and vLLM handle EOS merging internally; only the TRT-LLM gRPC path needs explicit stop_token_ids because it does not merge them. Changes: - Add crates/tokenizer/src/eos.rs with load_eos_token_ids() in a shared module (not tiktoken-specific) with extracted helper to eliminate code duplication - Expose eos_token_ids() on Tokenizer trait + CachedTokenizer delegation - Pass merged IDs as stop_token_ids in TRT-LLM gRPC requests (both Chat and Responses paths) - Respect ignore_eos: skip stop_token_ids when ignore_eos=true - Unit tests for EOS ID parsing (single int, array, missing, merge) Signed-off-by: ConnorLi96 <ConnorLi96@users.noreply.github.com> Made-with: Cursor Signed-off-by: ConnorLi96 <ConnorLi96@users.noreply.github.com> Made-with: Cursor Signed-off-by: ConnorLi96 <ConnorLi96@users.noreply.github.com> Made-with: Cursor
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 4d08498fc4
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| || filename == "config.json" | ||
| || filename == "generation_config.json" |
There was a problem hiding this comment.
Including config.json/generation_config.json in is_tokenizer_file() can make download_tokenizer_from_hf() set cache_dir from a root config file before any tokenizer file is seen. For models that keep tokenizer files in a subdirectory (for example tokenizer/tokenizer.json), create_tokenizer_async_with_chat_template() then probes only the returned root directory and fails to find a tokenizer, so loading regresses from working to failing for that layout. Please ensure cache-dir selection is based on tokenizer artifacts (or add recursive tokenizer discovery) rather than config-file matches.
Useful? React with 👍 / 👎.
|
This pull request has been automatically marked as stale because it has not had any activity within 14 days. It will be automatically closed if no further activity occurs within 16 days. Leave a comment if you feel this pull request should remain open. Thank you! |
|
Hi @ConnorLi96, this PR has merge conflicts that must be resolved before it can be merged. Please rebase your branch: git fetch origin main
git rebase origin/main
# resolve any conflicts, then:
git push --force-with-lease |
mergify
Bot
added
the
needs-rebase
|
This pull request has been automatically marked as stale because it has not had any activity within 14 days. It will be automatically closed if no further activity occurs within 16 days. Leave a comment if you feel this pull request should remain open. Thank you! |
|
@slin1237 could you take a look on the latest update here? this is quite huge issue that prevents seamlessly using Kimi models served via vllm with smg. thanks! |
|
This pull request has been automatically marked as stale because it has not had any activity within 14 days. It will be automatically closed if no further activity occurs within 16 days. Leave a comment if you feel this pull request should remain open. Thank you! |
|
This pull request has been automatically closed due to inactivity. Please feel free to reopen if you intend to continue working on it. Thank you! |
4 participants
Problem
Models like Kimi-K2.5 define different EOS tokens in
config.json(EOS=163585) andgeneration_config.json(<|im_end|>=163586). TRT-LLM's gRPC path only uses one source, causing infinite generation past turn boundaries (18/60 → 55/60 on fc-dash).https://huggingface.co/nvidia/Kimi-K2.5-NVFP4/blob/main/config.json#L12
https://huggingface.co/nvidia/Kimi-K2.5-NVFP4/blob/main/generation_config.json#L3
Solution
load_eos_token_ids()to merge EOS IDs from both config fileseos_token_ids()on theTokenizertraitstop_token_idsin TRT-LLM gRPC requests<|im_end|>stop string workaroundChanges
crates/tokenizer/src/traits.rs— addeos_token_ids()to traitcrates/tokenizer/src/tiktoken.rs—load_eos_token_ids()implcrates/tokenizer/src/huggingface.rs— impl for HF tokenizercrates/tokenizer/src/mock.rs— mock implcrates/grpc_client/src/trtllm_service.rs— pass merged IDs, remove hardcoded stop stringmodel_gateway/src/routers/grpc/client.rs— plumb new fieldmodel_gateway/.../harmony/stages/request_building.rs— pass EOS IDsmodel_gateway/.../regular/stages/chat/request_building.rs— pass EOS IDsTest Plan
Checklist:
cargo +nightly fmtpassescargo clippy --all-targets --all-features -- -D warningspassesMade with Cursor
Summary by CodeRabbit
New Features
Tests
Closes #1112