feat(zmq): forward EOS and resolve string stops for the direct backend (4/7) - #2057
Merged
Conversation
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (12)
🚧 Files skipped from review as they are similar to previous changes (12)
📝 WalkthroughSummary by CodeRabbit
WalkthroughThe PR adds matched stop reporting to the tokenizer, resolves string stops for applicable requests, propagates model EOS IDs into ZMQ vLLM sampling parameters, and preserves local stop results across regular and streaming responses. ChangesStop and EOS processing
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant Client
participant RequestBuilder
participant ZmqEngineClient
participant EngineCore
participant StopSequenceDecoder
Client->>RequestBuilder: submit generation request
RequestBuilder->>ZmqEngineClient: send normalized stops
ZmqEngineClient->>EngineCore: send EOS and stop-token parameters
EngineCore-->>StopSequenceDecoder: stream generated chunks
StopSequenceDecoder-->>Client: emit text and local stop metadata
Possibly related PRs
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 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 |
github-actions
Bot
added
tokenizer
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
There was a problem hiding this comment.
Clean, well-structured PR. The stop-sequence resolution and EOS forwarding logic is correct across all paths (streaming, non-streaming, multi-choice). Key observations:
StopSequenceDecoder.matched_stop: Correctly tracks the matched string and resets onreset(). Tests cover both string and token-level stops.resolve_string_stops: Properly drains string stops from the proto while forwarding single-token ones asstop_token_ids. SGLang always resolves; vLLM only over ZMQ. Theis_vllm()→is_zmq()rename onBackendClientcorrectly narrows the scope (gRPC-vLLM handles stops server-side).- Dual EOS injection: Tokenizer EOS ids (in
resolve_string_stops) and model-dir EOS ids (intranslate_sampling) are complementary sources with dedup — defensive and correct. - Streaming finish-reason pinning:
stopped_indices/stoppedflags correctly prevent the engine's Complete message from overwriting a local stop-decoder match. - Test coverage: Comprehensive — covers single/multi-token stops, dedup, ignore_eos, no-tokenizer fallback, model-dir resolution, and matched-stop reporting.
No issues found. 0 important · 0 nit · 0 pre-existing
slin1237
force-pushed
the
zmq/03-multimodal
branch
from
August 5, 2026 20:20
fc7ed51 to
da17295
Compare
slin1237
force-pushed
the
zmq/04-eos-forwarding
branch
from
August 5, 2026 20:20
1bddbe3 to
ad35b4b
Compare
Comment on lines
118
to
124
| @@ -71,6 +124,9 @@ pub struct ZmqEngineClient { | |||
| /// Model id advertised for metadata (the engine does not report it on the | |||
There was a problem hiding this comment.
🟡 Nit: eos_ids_from_file silently swallows JSON parse failures — if config.json exists but contains corrupt JSON, this returns empty EOS ids with no log. Since this runs once at connect time (not per-request), a warn! on parse failure would help diagnose "generation runs to max_tokens" issues without adding noise.
The tokenizer-based EOS injection in resolve_string_stops serves as a backstop, but if both paths fail silently, diagnosing the root cause becomes harder.
Suggested change
| fn eos_ids_from_file(path: &Path) -> Vec<u32> { | |
| let text = match std::fs::read_to_string(path) { | |
| Ok(t) => t, | |
| Err(_) => return Vec::new(), | |
| }; | |
| match serde_json::from_str::<serde_json::Value>(&text) { | |
| Ok(config) => eos_ids_from_value(config.get("eos_token_id")), | |
| Err(e) => { | |
| tracing::warn!(path = %path.display(), error = %e, "failed to parse model config for EOS ids"); | |
| Vec::new() | |
| } | |
| } | |
| } |
slin1237
force-pushed
the
zmq/03-multimodal
branch
3 times, most recently
from
August 6, 2026 00:30
447543f to
20a273d
Compare
slin1237
force-pushed
the
zmq/04-eos-forwarding
branch
from
August 6, 2026 02:17
ad35b4b to
ca6a4b1
Compare
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 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 `@model_gateway/src/routers/grpc/common/stages/helpers.rs`:
- Around line 372-403: Extend the request match handling to cover
ProtoGenerateRequest::TokenSpeed when is_zmq is true, applying the same stop
extraction, single-token encoding, and deduplication used for the Vllm ZMQ
branch before translate_request_tokenspeed runs. Add a TokenSpeed ZMQ test that
verifies a one-token string stop is serialized on the wire in stop_token_ids.
In `@model_gateway/src/routers/grpc/regular/streaming.rs`:
- Around line 418-437: Add regression tests covering local string-stop handling
in the chat and Messages streaming paths, including PD decode streams. Verify
pre-stop text is emitted, subsequent chunks are ignored, final responses report
finish reason "stop" with the matched sequence, and chat requests with more than
one choice apply the behavior independently to each choice.
In `@model_gateway/src/routers/grpc/zmq_client.rs`:
- Around line 101-107: Update eos_ids_from_file to distinguish a missing
optional file from a present file that fails reading or JSON parsing: preserve
the empty result for absent files, but return or propagate an error for read or
parse failures so connect_zmq_backend cannot connect with invalid EOS
configuration. Adjust callers accordingly and add a test covering malformed
JSON.
🪄 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: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: a4253928-13e9-414d-853c-ffeaf098e7ed
📒 Files selected for processing (12)
crates/tokenizer/src/mock.rscrates/tokenizer/src/stop.rsmodel_gateway/src/routers/grpc/backend_client.rsmodel_gateway/src/routers/grpc/common/stages/helpers.rsmodel_gateway/src/routers/grpc/regular/processor.rsmodel_gateway/src/routers/grpc/regular/stages/chat/request_building.rsmodel_gateway/src/routers/grpc/regular/stages/completion/request_building.rsmodel_gateway/src/routers/grpc/regular/stages/generate/request_building.rsmodel_gateway/src/routers/grpc/regular/stages/messages/request_building.rsmodel_gateway/src/routers/grpc/regular/streaming.rsmodel_gateway/src/routers/grpc/zmq_client.rsmodel_gateway/src/worker/worker.rs
Comment on lines
+418
to
+437
| let (chunk_text, should_stop) = | ||
| Self::process_chunk_tokens(stop_decoder, chunk.token_ids()); | ||
|
|
||
| if should_stop { | ||
| // Stop-decoder match takes precedence: pin "stop" even if | ||
| // the backend's eventual Complete carries "length" (the | ||
| // local stop sequence fired first). Any pre-stop text in | ||
| // `chunk_text` is still emitted below before the finish | ||
| // reason is flushed in Phase 4. | ||
| finish_reasons | ||
| .entry(index) | ||
| .or_insert_with(|| "stop".to_string()); | ||
| matched_stops.entry(index).or_insert_with(|| { | ||
| stop_decoder | ||
| .matched_stop() | ||
| .map(|s| Value::String(s.to_string())) | ||
| }); | ||
| stopped_indices.insert(index); | ||
| } | ||
|
|
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Locate existing Rust tests for local stop handling and streaming metadata.
rg -n -C 3 --type rust \
'matched_stop|StopSequence|stopped_indices|process_streaming_chunks|process_messages_streaming_chunks' \
model_gateway cratesRepository: smg-project/smg
Length of output: 50371
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "Relevant test file sizes and focused search"
wc -l model_gateway/tests/messages_streaming_test.rs model_gateway/tests/chat_streaming_test.rs 2>/dev/null || true
rg -n -C 4 --type rust \
'local.*stop|string stop|stop string|matched_stop_json|local-stop|length|MaxTokens|stop_sequence|n>1' \
model_gateway/tests crates/protocols/tests 2>/dev/null || true
echo
echo "Streaming source relevant functions outline"
ast-grep outline model_gateway/src/routers/grpc/regular/streaming.rs \
--view expanded --match process_streaming_chunks --match process_messages_streaming_chunks \
2>/dev/null || true
echo
echo "Read streaming source relevant ranges"
sed -n '360,470p' model_gateway/src/routers/grpc/regular/streaming.rs
echo
sed -n '560,625p' model_gateway/src/routers/grpc/regular/streaming.rsRepository: smg-project/smg
Length of output: 43703
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "chat streaming tests focused search"
if [ -f model_gateway/tests/chat_streaming_test.rs ]; then
wc -l model_gateway/tests/chat_streaming_test.rs
rg -n -C 4 --type rust \
'matched_stop|stop.*length|length.*stop|n>1|multi.*choice|stop_sequence|local.*stop|string.*stop|stop_decoder|process_chunk_tokens|process_messages_streaming_chunks|process_streaming_chunks' \
model_gateway/tests/chat_streaming_test.rs
else
echo "missing model_gateway/tests/chat_streaming_test.rs"
fi
echo
echo "Messages streaming tests focused search"
wc -l model_gateway/tests/messages_streaming_test.rs
rg -n -C 4 --type rust \
'matched_stop|stop.*length|length.*stop|n>1|multi.*choice|stop_sequence|local.*stop|string.*stop|stop_decoder|process_messages_streaming_chunks' \
model_gateway/tests/messages_streaming_test.rs
echo
echo "streaming functions call sites/routes"
rg -n -C 3 \
'process_streaming_chunks|process_messages_streaming_chunks|process_decode_streaming_chunk|DecodeStream|Decode' \
model_gateway/src/routers/grpc \
crates/tokenizer/src \
| head -n 240Repository: smg-project/smg
Length of output: 27098
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "Find test files containing process_streaming_chunks or stream function names"
rg -n --type rust \
'fn (.*stream|.*chat.*stream|.*regular.*stream|test.*stream)|process_streaming_chunks|process_messages_streaming_chunks|LocalStop|matched_stop|stopped_indices' \
model_gateway/tests model_gateway/src/routers/grpc crates/protocols/src/builders/chat 2>/dev/null | head -n 240
echo
echo "Locate exact function definitions"
rg -n --type rust \
'pub\(crate\)? fn process_streaming_chunks|pub\(crate\)? fn process_messages_streaming_chunks|fn process_streaming_chunks|fn process_messages_streaming_chunks' \
model_gateway/src/routers/grpc crates/protocols/src/builders/chat 2>/dev/null | head -n 80
echo
echo "Inspect surrounding implementation if locations known"
if true; then
for f in \
model_gateway/src/routers/grpc/regular/streaming.rs \
model_gateway/src/routers/grpc/pd/streaming.rs \
model_gateway/src/routers/grpc/messages/streaming.rs \
crates/protocols/src/builders/chat/stream_response.rs
do
if [ -f "$f" ]; then
echo "--- $f line count ---"
wc -l "$f"
fi
done
fiRepository: smg-project/smg
Length of output: 18534
🟡 Add local-stop streaming regression coverage.
Add tests where a local string stop fires before a backend "length" completion for chat and Messages streams. Check that pre-stop text is emitted, later chunks are ignored, the final response uses "stop", the matched sequence is present, n > 1 chat choices are handled, and PD decode streams use the same path.
🤖 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 `@model_gateway/src/routers/grpc/regular/streaming.rs` around lines 418 - 437,
Add regression tests covering local string-stop handling in the chat and
Messages streaming paths, including PD decode streams. Verify pre-stop text is
emitted, subsequent chunks are ignored, final responses report finish reason
"stop" with the matched sequence, and chat requests with more than one choice
apply the behavior independently to each choice.
Source: Coding guidelines
Comment on lines
+101
to
+107
| /// Read a config file's `eos_token_id`, which is a single id or a list. | ||
| fn eos_ids_from_file(path: &Path) -> Vec<u32> { | ||
| std::fs::read_to_string(path) | ||
| .ok() | ||
| .and_then(|text| serde_json::from_str::<serde_json::Value>(&text).ok()) | ||
| .map(|config| eos_ids_from_value(config.get("eos_token_id"))) | ||
| .unwrap_or_default() |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
🔴 Important Fail when a present EOS configuration cannot be read or parsed.
These calls convert an unreadable or malformed local configuration into an empty EOS set. connect_zmq_backend then connects successfully, and translate_sampling can send no EOS ID. If no request tokenizer is available for the later fallback, generation stops only at explicit stops or max_tokens.
Keep missing optional files supported if required. Return an error for a present file that cannot be read or parsed. Add a malformed-JSON test.
As per coding guidelines, “Do not silently fall back to None or a default when configuration validation should fail loudly.”
🤖 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 `@model_gateway/src/routers/grpc/zmq_client.rs` around lines 101 - 107, Update
eos_ids_from_file to distinguish a missing optional file from a present file
that fails reading or JSON parsing: preserve the empty result for absent files,
but return or propagate an error for read or parse failures so
connect_zmq_backend cannot connect with invalid EOS configuration. Adjust
callers accordingly and add a test covering malformed JSON.
Source: Coding guidelines
A direct-ZMQ engine has no tokenizer or model config, so it can neither stop at EOS on its own nor match string stop sequences — the router must do both. - Resolve EOS ids from the worker's local model dir at connect time and attach them to every vLLM request: the primary rides _eos_token_id, extra ids merge into stop_token_ids, and the union feeds _all_stop_token_ids (min_tokens masking). ignore_eos drops the wire stops but keeps the bookkeeping set, mirroring the reference frontend. - Add resolve_string_stops, the single point shared by SGLang gRPC (skip_tokenizer_init) and every ZMQ backend: drop the string stop list, forward single-token stops as stop_token_ids, and leave multi-token stops to the router-side StopSequenceDecoder. Wire it into the chat, generate, messages, and completion request-building stages via a new BackendClient::is_zmq (renamed from is_vllm). - Teach StopSequenceDecoder to report the matched stop string; the regular processor and streaming paths let a local decoder match take precedence over the engine's finish reason (which is "length" when stops are enforced gateway-side) and surface the matched sequence. Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com>
slin1237
force-pushed
the
zmq/04-eos-forwarding
branch
from
August 6, 2026 04:02
ca6a4b1 to
252ee46
Compare
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
This was referenced Aug 6, 2026
2 tasks
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.
Description
Problem
A direct-ZMQ engine receives token ids only and cannot match string stops itself,
and it is not told the model's EOS token ids. Without those, generation over ZMQ
can run to
max_tokens(endless generation) and stringstopsequences arenever honored.
Solution
Forward the model's EOS token ids to the engine at connect time, and resolve
string stop sequences on the gateway side for the ZMQ path: single-token stops
are encoded and sent as stop-token ids, while multi-token string stops are matched
locally by the stop decoder, whose match takes precedence over the engine's finish
reason.
Changes
crates/engine_zmq_client/.../zmq_client.rs—EosTokenIds(primary + extra,from_model_dir), threaded throughconnect()and request/sampling translation.model_gateway/src/worker/worker.rs— load EOS ids and pass them intoconnect_zmq_backend.crates/tokenizer/src/stop.rs—matched_stopfield + accessor (+ tests);crates/tokenizer/src/mock.rs— mockeos_token_ids.model_gateway/.../common/stages/helpers.rs—encode_single_token_stops/resolve_string_stops(+ tests).model_gateway/.../regular/stages/{chat,generate,messages,completion}/request_building.rs— resolve string stops using the
is_zmqbackend flag.model_gateway/.../regular/{processor,streaming}.rs— local stop-decoder matchtakes precedence over the engine finish reason.
model_gateway/src/routers/grpc/backend_client.rs— renameis_vllm→is_zmqto name the actual capability gap (token-only backend).
Test Plan
cargo +nightly fmt -- --checkandcargo clippy --all-featuresclean.cargo test:stop.rsmatched-stop accessor tests andhelpers.rssingle-token / string-stop resolution tests.
Checklist
cargo +nightly fmtpassescargo clippy --all-targets --all-features -- -D warningspassesStack (split of the oversized #2041, merge bottom-up): 01-core-dispatch →
02-structured-outputs → 03-multimodal → 04-eos-forwarding (this PR) →
05-harmony-stop-matcher → 06-e2e-infra → 07-e2e-tests. Review/merge after its
parent
zmq/03-multimodal.