Repository navigation
fix(grpc): add matched_stop support for vLLM and TensorRT-LLM - #602
Conversation
|
Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits. |
Summary of ChangesHello, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request enhances the gRPC responses for vLLM and TensorRT-LLM by introducing comprehensive support for Highlights
Changelog
Activity
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
📝 WalkthroughWalkthroughProtos changed to represent matched stop data as a oneof (token id or string) in GenerateComplete for multiple backends; the gRPC wrapper removed direct matched_stop access and added a backend-aware matched_stop_json() serializer. Changes
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 3✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches
🧪 Generate unit tests (beta)
Comment |
github-actions
Bot
added
grpc
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request adds support for matched_stop in vLLM and TensorRT-LLM gRPC backends, aligning them with the existing SGLang implementation. The changes involve updating the vLLM protobuf definition and modifying the Rust wrapper to handle the new fields. The implementation is clean and correct. The suggestion to refactor a small piece of duplicated code in proto_wrapper.rs to improve maintainability is valid and aligns with best practices.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
vLLM gRPC was missing stop_reason entirely (proto had no field), and
TensorRT-LLM had optional string stop_reason which lost the int type
for token IDs. Add oneof matched_stop { matched_token_id, matched_stop_str }
to both protos (wire-compatible for TRT-LLM) and wire through the Rust
wrapper using a local macro to deduplicate the three backend arms.
Signed-off-by: Chang Su <chang.s.su@oracle.com>
d3c536f to
b82bb23
Compare
…ixes The previous pin (fd080fc7) is the commit immediately before lightseekorg/tokenspeed#578, which adds defensive Finished-state handlers to the scheduler FSM. Without #578 the engine crashes under retract pressure with: RuntimeError: FSM transition invalid: event=tokenspeed::fsm::ExtendResultEvent; state=tokenspeed::fsm::Finished Reproduced on the nightly Qwen3-30B-A3B bench: when the host KV cache fills up and a retract fails, AbortEvent terminalizes the request → Finished, but overlap scheduling has already dispatched a forward batch including it, and the late ExtendResultEvent commit hits a strict FSM handler that throws and kills the scheduler event loop. Bump to current lightseekorg/tokenspeed main (eabeb106) so we also pick up #602 (release scheduler slot + cancel non-stream handlers on client disconnect), which removes the long pre-crash stream of ``Received output for rid=... but the state was deleted in AsyncLLM`` warnings caused by aborted requests still occupying engine slots. Signed-off-by: Chang Su <chang.s.su@oracle.com>
…ixes The previous pin (fd080fc7) is the commit immediately before lightseekorg/tokenspeed#578, which adds defensive Finished-state handlers to the scheduler FSM. Without #578 the engine crashes under retract pressure with: RuntimeError: FSM transition invalid: event=tokenspeed::fsm::ExtendResultEvent; state=tokenspeed::fsm::Finished Reproduced on the nightly Qwen3-30B-A3B bench: when the host KV cache fills up and a retract fails, AbortEvent terminalizes the request → Finished, but overlap scheduling has already dispatched a forward batch including it, and the late ExtendResultEvent commit hits a strict FSM handler that throws and kills the scheduler event loop. Bump to current lightseekorg/tokenspeed main (eabeb106) so we also pick up #602 (release scheduler slot + cancel non-stream handlers on client disconnect), which removes the long pre-crash stream of ``Received output for rid=... but the state was deleted in AsyncLLM`` warnings caused by aborted requests still occupying engine slots. Signed-off-by: key4ng <rukeyang@gmail.com>
1 participant
Description
Problem
vLLM gRPC responses were missing
stop_reason(which specific stop string/token triggered a stop). vLLM HTTP returns it (e.g.,stop_reason=','), but the gRPC proto never defined the field. TensorRT-LLM hadoptional string stop_reasonin its proto but the Rust wrapper ignored it. Only SGLang'soneof matched_stopwas wired through.Solution
oneof matched_stop { matched_token_id, matched_stop_str }to vLLM'sGenerateCompleteproto, matching SGLang's pattern to preserveint | strtype distinction from vLLM'sCompletionOutput.stop_reasonmatched_stop_json()in the proto wrappermatched_stop()helperChanges
grpc_client/proto/vllm_engine.proto: Addoneof matched_stop(fields 10-11) toGenerateCompletemodel_gateway/src/routers/grpc/proto_wrapper.rs: Rewritematched_stop_json()to handle all three backends (SGLang oneof, vLLM oneof, TensorRT-LLM string), remove deadmatched_stop()methodNote: The corresponding vLLM
grpc_server.pychange to populate the new field is in the vLLM fork.Test Plan
cargo build -p smg(proto re-generation + compilation)cargo clippy -p smg -- -D warnings— zero warningscargo test -p smg— all tests passpre-commit run --all-files— all hooks passstop_reasonappears in vLLM gRPC chat completion responses with stop sequencesChecklist
cargo +nightly fmtpassescargo clippy --all-targets --all-features -- -D warningspassesSummary by CodeRabbit