feat(zmq): honor vLLM prompt logprobs over the ZMQ backend - #2077
Conversation
The ZMQ path rejected any request carrying `prompt_logprobs` because the adapter had no mapping for the engine's prompt tensors, so a feature the gRPC servicer supports 400'd purely on transport. Forward `prompt_logprobs` on the wire and map `new_prompt_logprobs_tensors` back the way the servicer does: accumulate across chunked-prefill ticks, seed the first prompt token with a null logprob (nothing precedes it to condition on), and attach the set as `input_logprobs` to the first token-bearing chunk (the proto carries it in the first chunk only) and to the terminal `Complete`. Shaping the ranked candidates per position is now shared between the output and prompt sides, which also gives `-1` its documented meaning of "every candidate the engine returned" instead of silently dropping them. Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com>
github-actions
Bot
added
grpc
📝 WalkthroughSummary by CodeRabbit
WalkthroughZMQ vLLM generation now forwards prompt logprob requests, supports ranked candidates including ChangesvLLM prompt logprob streaming
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant Generate
participant EngineCore
participant VllmGenerateStream
Generate->>EngineCore: Forward prompt logprob and candidate settings
EngineCore-->>VllmGenerateStream: Return prompt logprob responses
VllmGenerateStream->>VllmGenerateStream: Accumulate prompt tokens and ranked candidates
VllmGenerateStream-->>Generate: Emit prompt logprobs on first chunk and terminal completion
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 |
|
👋 The PR description doesn't fully follow
Please update the PR description so reviewers have the context they need. |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (2)
model_gateway/src/routers/grpc/zmq_client.rs (2)
957-980: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick win🟡 Nit Construct
InputLogProbsonly for responses that carry it.
attach_input_logprobsclones all accumulated prompt logprobs and candidates on every later decode tick. Later chunks discard that payload. Long prompts andprompt_logprobs = -1cause repeated large allocations during streaming.Create the payload only for a parked
Complete, the first token-bearingChunk, or a directComplete.Proposed hot-path reduction
- let input_logprobs = vllm::InputLogProbs { - token_logprobs: self.state.prompt_logprobs.clone(), - token_ids: self.state.prompt_token_ids.clone(), - top_logprobs: self.state.prompt_top_logprobs.clone(), - }; if let Some(vllm::generate_response::Response::Complete(parked)) = self .pending .as_mut() .and_then(|pending| pending.response.as_mut()) { - parked.input_logprobs = Some(input_logprobs.clone()); + parked.input_logprobs = Some(Self::input_logprobs(&self.state)); } match response.response.as_mut() { Some(vllm::generate_response::Response::Chunk(chunk)) if !self.input_logprobs_emitted && !chunk.token_ids.is_empty() => { - chunk.input_logprobs = Some(input_logprobs); + chunk.input_logprobs = Some(Self::input_logprobs(&self.state)); self.input_logprobs_emitted = true; } Some(vllm::generate_response::Response::Complete(complete)) => { - complete.input_logprobs = Some(input_logprobs); + complete.input_logprobs = Some(Self::input_logprobs(&self.state)); } _ => {} } } + + fn input_logprobs(state: &StreamState) -> vllm::InputLogProbs { + vllm::InputLogProbs { + token_logprobs: state.prompt_logprobs.clone(), + token_ids: state.prompt_token_ids.clone(), + top_logprobs: state.prompt_top_logprobs.clone(), + } + }As per coding guidelines, “Avoid unnecessary clone() calls in gRPC streaming hot paths, especially during per-token response processing.”
🤖 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 957 - 980, Update attach_input_logprobs so InputLogProbs is constructed only inside the parked Complete, first token-bearing Chunk, or direct Complete branches that retain it. Avoid cloning prompt_logprobs, prompt_token_ids, or prompt_top_logprobs on later token-less or already-emitted chunks, while preserving the existing assignment behavior for responses that carry input logprobs.Source: Coding guidelines
1965-2081: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win🟡 Nit Add a multi-tick chunked-prefill test.
This test sends one
new_prompt_logprobs_tensorspayload. It does not verify ordered accumulation across multiple prefill ticks.Send two prompt-tensor outputs before the first token-bearing response. Assert that the first chunk and
Completecontain the null first-token entry and all scored prompt positions in order.As per coding guidelines, “Run the pr-test-analyzer agent to verify that tests adequately cover new or changed functionality.”
🤖 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 1965 - 2081, Extend the test around the engine_task and expect_input_logprobs to emit two prefill RequestBatch outputs with separate new_prompt_logprobs_tensors before the token-bearing decode output. Update the expected first chunk and Complete input logprobs to verify the null first-token entry plus every scored prompt position in original order, while preserving the assertion that only the first chunk carries prompt logprobs. Run the pr-test-analyzer agent to verify coverage.Source: Coding guidelines
🤖 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.
Nitpick comments:
In `@model_gateway/src/routers/grpc/zmq_client.rs`:
- Around line 957-980: Update attach_input_logprobs so InputLogProbs is
constructed only inside the parked Complete, first token-bearing Chunk, or
direct Complete branches that retain it. Avoid cloning prompt_logprobs,
prompt_token_ids, or prompt_top_logprobs on later token-less or already-emitted
chunks, while preserving the existing assignment behavior for responses that
carry input logprobs.
- Around line 1965-2081: Extend the test around the engine_task and
expect_input_logprobs to emit two prefill RequestBatch outputs with separate
new_prompt_logprobs_tensors before the token-bearing decode output. Update the
expected first chunk and Complete input logprobs to verify the null first-token
entry plus every scored prompt position in original order, while preserving the
assertion that only the first chunk carries prompt logprobs. Run the
pr-test-analyzer agent to verify coverage.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 26ec90f3-a662-49a0-9e0f-54f9a54f17f1
📒 Files selected for processing (1)
model_gateway/src/routers/grpc/zmq_client.rs
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
| if self.state.prompt_logprobs.is_empty() { | ||
| return; | ||
| } | ||
| let input_logprobs = vllm::InputLogProbs { | ||
| token_logprobs: self.state.prompt_logprobs.clone(), | ||
| token_ids: self.state.prompt_token_ids.clone(), | ||
| top_logprobs: self.state.prompt_top_logprobs.clone(), |
There was a problem hiding this comment.
🟡 Nit: clone() on every decode tick in a streaming hot path. After input_logprobs_emitted is set and there's no pending Complete, subsequent chunk ticks clone all three vecs just to enter the match and hit the no-op Chunk(_) arm. For a long generation against a large prompt this is O(prompt_len) allocation per output token for data that's immediately dropped.
Not a bug — prompt logprobs aren't wired to any frontend yet — but once they are, an early return here would avoid the per-tick waste:
if self.state.prompt_logprobs.is_empty() {
return;
}
// Decode-only ticks after the first token-bearing chunk: the data
// was already emitted and won't be used again until the Complete.
if self.input_logprobs_emitted
&& self.pending.is_none()
&& matches!(response.response, Some(vllm::generate_response::Response::Chunk(_)))
{
return;
}
There was a problem hiding this comment.
Clean implementation of prompt logprobs parity for the ZMQ backend. The accumulation logic mirrors the existing output-logprobs pattern well, the null-seeded first token follows the API contract, and the test coverage is thorough (end-to-end over IPC, unit tests for the sentinel mapping and forwarding).
One minor nit flagged inline: per-tick cloning in attach_input_logprobs can be short-circuited after the first chunk is emitted. No blocking issues.
Summary: 0 🔴 Important · 1 🟡 Nit · 0 🟣 Pre-existing
1 participant
What
The ZMQ path rejected any vLLM request carrying
prompt_logprobs:The gRPC servicer supports the field, so the same request 400'd or
succeeded purely on which transport the worker happened to use. This
closes that asymmetry.
prompt_logprobsverbatim on the EngineCore wire.new_prompt_logprobs_tensorsback the way the Python servicer's_build_input_logprobsdoes: accumulate across chunked-prefill ticks,seed the first prompt token with a
nulllogprob (nothing precedes itto condition on), and emit the accumulated set as
input_logprobsonthe first token-bearing chunk (the proto carries it in the first chunk
only) and on the terminal
Complete.prompt sides (
ranked_candidate_count). As a side effect-1nowmeans what the proto says it means — every candidate the engine
returned — instead of collapsing to zero candidates.
Scope note
This is transport parity, not a new user-visible feature: no frontend
populates
SamplingParams.prompt_logprobstoday (/v1/completionsecho+logprobsrequests output logprobs only), so neither backendreturns
input_logprobsin practice yet. Wiring the frontend is aseparate, backend-independent change.
Testing
generate_streams_prompt_logprobsdrives the full adapter overipc://against the mock engine: assertsprompt_logprobsreachesthe engine, then that a two-tick stream puts the null-seeded prompt
set on the first chunk, not on the second, and on the
Complete.vllm_forwards_prompt_logprobsreplaces the old rejection test.ranked_candidate_count_maps_the_sentinelspins the0/n/-1semantics.
cargo test -p smg --lib(one pre-existing unrelated failure,middleware::metrics::tests::distinct_ids_on_matched_route_do_not_grow_interner,which also fails on
main),cargo +nightly fmt --check, andcargo +1.95.0 clippy --all-targets -- -D warningsare clean.