[Bugfix] Content-aware routing should work with token ids - #237
[Bugfix] Content-aware routing should work with token ids#237simondanielsson wants to merge 8 commits into
Conversation
Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com>
Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com>
Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com>
simondanielsson
changed the title
![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: 171ddf4532
ℹ️ 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".
Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com>
Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com>
Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com>
simondanielsson
changed the title
simondanielsson
changed the title
…with-tokenids Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com>
|
@hsliuustc0106 @wuhang2014 could you please take a look? Thanks 🙏 |
There was a problem hiding this comment.
🟡 Changes recommended
Token-key detection collides with valid text inputs, and one token-ID request path still uses character-based matching.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Pull request overview
Updates cache-aware routing to use actual token IDs and count only complete-token prefix matches.
Changes:
- Encodes token IDs with unit separators.
- Adds token-aware cache match-rate calculation.
- Updates routing-key and match-rate tests.
File summaries
| File | Description |
|---|---|
src/protocols/spec.rs |
Encodes token-based prompts for routing. |
src/policies/cache_aware.rs |
Calculates complete-token cache matches. |
tests/test_prompt_input.rs |
Updates expected token routing keys. |
Review details
- Files reviewed: 3/3 changed files
- Comments generated: 3
- Review effort level: Balanced
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
| if input_char_count == 0 { | ||
| return 0.0; | ||
| } | ||
| if text.contains(TOKEN_ID_SEPARATOR) { |
There was a problem hiding this comment.
Good idea, added a prefix token to discern between token ids and regular text
Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com>
Purpose
For purposes of routing, previously token ids were (strangely) encoded as
token_ids:{num_tokens}. For cache aware routing, that meant prompts of equal length were routed to the same vLLM instance. The actual content of the prompt was ignored.This PR fixes this to by encoding the token ids and not just length of the prompt.
Extra care is taken so we only match on full tokens during cache aware routing, to avoid token
1234is a partial match to token12345(because they are in fact distinct tokens and are unlikely prefixes when detokenized).Test Plan
Tests pass
Test Result
OK
Related work
#242 unifies routing logic for PD and DP, so it will also get this benefit of correctly routing when token ids are passed.
Essential Elements of an Effective PR Description Checklist