feat(zmq): multimodal inputs over the direct ZMQ backend (3/7) - #2056
Conversation
|
Warning Review limit reached
Next review available in: 8 minutes You've used all free OSS reviews for now. Wait for the free limit to reset to keep reviewing this public repository. How can I continue?After more reviews become available, a review can be triggered using the To avoid repeated limits, reduce automatic review volume by pausing incremental auto-reviews earlier, using label-based review opt-in, excluding WIP or generated PR titles, or requesting reviews manually when the PR is ready. If your team needs uninterrupted high-volume reviews, an organization admin can enable usage-based reviews. How do review limits work?CodeRabbit enforces per-developer PR review limits for each organization. Most developers receive the normal plan review availability. For paid Pro and Pro+ PR reviews, CodeRabbit uses adaptive limits for sustained high-volume activity. When a developer's recent PR review activity reaches the 95th percentile or higher among CodeRabbit users, additional reviews become available more gradually as earlier reviews age out of the rolling window. Please refer docs for additional details. Review details⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (9)
📝 WalkthroughSummary by CodeRabbit
WalkthroughThis change adds typed vLLM multimodal protocol structures, direct-ZMQ tensor preprocessing, model-dtype casting, placeholder masking, and gateway integration. vLLM requests now carry assembled multimodal features, while TokenSpeed requests reject multimodal inputs. ChangesZMQ multimodal support
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant GrpcRequest
participant BackendClient
participant ZmqMultimodal
participant EngineCoreRequest
GrpcRequest->>BackendClient: Provide multimodal request data
BackendClient->>ZmqMultimodal: Convert inline tensors and metadata
ZmqMultimodal-->>BackendClient: Return typed MmFeatures
BackendClient->>EngineCoreRequest: Forward mm_features to vLLM
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
grpc
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
| MmFieldElem { | ||
| data: Some(MmKwargValue::Tensor(decoded.whole())), | ||
| field: MmField::Shared(MmSharedField { | ||
| batch_size: num_items, |
There was a problem hiding this comment.
🟡 Nit: keep_on_cpu is hardcoded to false for shared fields, ignoring the on_cpu variable that the batched and flat branches use correctly. If a shared tensor's key appears in keep_on_cpu_keys, the engine will incorrectly move it to the accelerator.
| MmFieldElem { | |
| data: Some(MmKwargValue::Tensor(decoded.whole())), | |
| field: MmField::Shared(MmSharedField { | |
| batch_size: num_items, | |
| MmFieldElem { | |
| data: Some(MmKwargValue::Tensor(decoded.whole())), | |
| field: MmField::Shared(MmSharedField { | |
| batch_size: num_items, | |
| keep_on_cpu: on_cpu, | |
| }), |
There was a problem hiding this comment.
Clean, well-structured PR. The multimodal type definitions mirror the Python side faithfully, the per-item split logic (batched/flat/shared) is correct, and the dtype casting is properly handled. Good defensive error handling at multiple layers (assembly, backend dispatch, translation). One minor nit posted about keep_on_cpu in the shared-field branch.
Summary: 0 🔴 Important · 1 🟡 Nit · 0 🟣 Pre-existing
fc7ed51 to
da17295
Compare
d7077ef to
a1f3dfb
Compare
da17295 to
d8562c1
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: 5
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
model_gateway/src/routers/grpc/zmq_client.rs (1)
194-208: 🚀 Performance & Scalability | 🟠 Major | 🏗️ Heavy liftShare multimodal features across
n > 1fan-out subs.
fan_out_requests(*req)clones fullvllm::GenerateRequests, including inlinemm_inputs;translate_requestthen callszmq_multimodal::build_mm_featuresseparately per subsample, decoding and casting float32 tensors again and creatingWireTensorbytes in everyEngineCoreRequest. BuildMmFeaturesonce before the loop and assign a refcounted slice for each sub; this removes repeated multimedia preprocessing and cuts peak gateway memory from scaling withn. Also establish a bounded inline multimodal payload limit, because the current comment explicitly validates arbitrary sizes.🤖 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 194 - 208, Update the fan-out flow around fan_out_requests, translate_request, and zmq_multimodal::build_mm_features to preprocess inline multimodal inputs once before the loop, then assign each generated EngineCoreRequest a refcounted shared MmFeatures slice. Add validation enforcing a bounded inline multimodal payload size before preprocessing, while preserving existing request translation and stream behavior.Source: Coding guidelines
🧹 Nitpick comments (5)
crates/engine_zmq_client/src/protocol/vllm/multimodal.rs (1)
248-260: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win🟡 Nit: Add
MmField::Flatto the round-trip test.
field_round_trips_python_factory_tuplecovers onlyBatchedandShared.Flatis the variant whoseMmFieldWireInneruntagged resolution depends ondeny_unknown_fieldsdistinguishing three field sets, and it carries the nestedMmSlice/SliceSpectuple encoding. That is the most fragile decode path in this module and it is only exercised on the encode side.♻️ Add the Flat case
for field in [ MmField::Batched(MmBatchedField { keep_on_cpu: true }), MmField::Shared(MmSharedField { batch_size: 4, keep_on_cpu: false, }), + MmField::Flat(MmFlatField { + slices: vec![MmSlice::Slice(SliceSpec { + start: Some(0), + stop: Some(1200), + step: None, + })], + dim: 0, + keep_on_cpu: true, + }), ] {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 `@crates/engine_zmq_client/src/protocol/vllm/multimodal.rs` around lines 248 - 260, Add the MmField::Flat variant, including representative MmSlice/SliceSpec data, to the field_round_trips_python_factory_tuple test so its nested tuple encoding and untagged decoding are verified alongside Batched and Shared.Source: Coding guidelines
model_gateway/src/routers/grpc/backend_client.rs (1)
368-379: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value🟡 Nit:
zmq_tokenspeed_mmis unreachable today, and theinto_proto(true)flag is unexplained.Two other layers already reject TokenSpeed multimodal input:
assemble_multimodal_data_implbails for a non-vLLM ZMQ runtime, andtranslate_request_tokenspeedrejects any request withmm_inputsset. So this converter can only ever producemm_inputsthat the translator immediately rejects with a generic message. Add a short comment stating that the converter exists for the future TokenSpeed wire, so a reader does not assume the path works.Also document what
trueselects ininto_proto(true), or use a named binding. The literal carries no meaning at the call site.🤖 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/backend_client.rs` around lines 368 - 379, Add a brief comment above zmq_tokenspeed_mm explaining that it is retained for a future TokenSpeed wire path despite being currently unreachable. Replace the unexplained true argument in MultimodalData::TokenSpeed conversion with a named binding or concise explanation identifying the serialization mode it selects, without changing behavior.model_gateway/src/routers/grpc/zmq_multimodal.rs (1)
384-405: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win🟡 Nit: Two coverage gaps, and a fixture that hides the bounds check.
base_inputs()uses placeholders(1, 3)and(6, 3), but several tests pass&[]asprompt_token_ids. Those ranges exceed an empty prompt. The tests pass only becauseis_embed_maskreturns early whenim_token_idisNone. If you apply the bounds-check fix, give these tests a prompt of at least 9 tokens.- No test sets
keep_on_cpu_keys, so the droppedon_cpuin the shared branch is not caught.As per coding guidelines: "Run the pr-test-analyzer agent to verify that tests adequately cover new or changed functionality."
Also applies to: 447-469
🤖 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_multimodal.rs` around lines 384 - 405, Update tests using base_inputs and empty prompt_token_ids to provide a prompt with at least nine tokens, so placeholder ranges are valid once bounds checks are enforced. Add coverage for keep_on_cpu_keys through the shared feature-building path and assert the resulting on-CPU behavior is preserved. Run the pr-test-analyzer agent to verify the new and changed functionality is adequately covered.Source: Coding guidelines
model_gateway/src/routers/grpc/multimodal/assemble.rs (1)
114-127: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value🟡 Nit: The ZMQ runtime dispatch shape differs from
backend_client.rs.This arm treats
Vllm | Unspecifiedas vLLM and bails for anything else.BackendClient::build_chat_requestandbuild_messages_requestinmodel_gateway/src/routers/grpc/backend_client.rs(lines 209-236, 254-281) use the opposite shape:TokenSpeedis explicit and_falls through to vLLM. Both functions switch on the sameclient.runtime(). TodayZmqEngineClient::runtime()returns onlyVllmorTokenSpeed, so no live path diverges. If a third ZMQ runtime is added, assembly will reject the request while request building silently treats it as vLLM. Pick one shape for both.Separately,
assemble_vllmcomputesshm_enabledandshm_min_bytes, and this arm then overwritesshm_enabledandrdma_enabled.shm_min_bytesis left at its resolved value even though it no longer applies. Passing the inline-transport intent intoassemble_vllmwould keep the three transport fields consistent in one place.As per coding guidelines: "Configuration changes must be checked across CLI arguments, types.rs, both conversion paths in main.rs, Python bindings, and the Go SDK."
🤖 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/multimodal/assemble.rs` around lines 114 - 127, Align the ZMQ runtime dispatch in the BackendClient::Zmq arm with BackendClient::build_chat_request and build_messages_request, keeping TokenSpeed explicit and routing other runtimes to vLLM. Update assemble_vllm to receive the inline-transport intent so it sets shm_enabled, rdma_enabled, and shm_min_bytes consistently, removing the post-assembly field overrides.Source: Coding guidelines
crates/engine_zmq_client/src/codec/tensor.rs (1)
158-159: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick win🟡 Nit: Reuse the existing float32 payload instead of copying it.
from_f32_bytes_castalready owns the allocation afterdata.len().is_multiple_of(4)andas_chunks::<4>incrates/engine_zmq_client/src/codec/tensor.rs; the Float32 branch still callsdata.to_vec(), sodecode_tensorimmediately discards the originalVec<u8>and allocates a second copy. Allowfrom_raw_bytes/from_rawto accept or takeBytesso the no-cast path can reuse the inline payload for large pixel tensors.🤖 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 `@crates/engine_zmq_client/src/codec/tensor.rs` around lines 158 - 159, Update the Float32 branch in decode_tensor and the related from_raw_bytes/from_raw constructors to accept or take ownership of Bytes, reusing the existing from_f32_bytes_cast allocation instead of calling data.to_vec(). Preserve the current validation and decoding behavior while eliminating the redundant payload copy.
🤖 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/zmq_client.rs`:
- Around line 873-885: Update the mm_inputs handling in translate_request to
validate that multimodal inputs containing tensors or hashes are not paired with
empty mm_placeholders before applying the empty-feature filter; return an error
instead of silently producing a text-only request. Preserve the existing
behavior for genuinely empty multimodal inputs, and add translation-level test
coverage in zmq_client.rs.
In `@model_gateway/src/routers/grpc/zmq_multimodal.rs`:
- Around line 257-271: Update the shared branch in the item-building logic to
set MmSharedField.keep_on_cpu from the existing on_cpu value, matching the
batched and flat branches. Preserve the shared tensor replication behavior while
ensuring keys listed in mm.keep_on_cpu_keys retain the requested CPU placement.
- Around line 217-231: Update the bounds accumulation loop in the flat-size
handling to use checked addition for total and size, returning a descriptive
error if the sum overflows instead of panicking or wrapping. Preserve the
existing negative-size validation and leading-dimension comparison after
successful accumulation.
- Around line 304-321: Update is_embed_mask to validate
offset.checked_add(length) against prompt_token_ids.len() before unwrapping or
checking im_token_id. Preserve the existing error for invalid ranges, then
return Ok(None) when im_token_id is absent only after validation succeeds.
- Around line 103-120: Update the tensor decoding flow around the float32 branch
to validate every payload length against the shape element count and the
declared dtype’s element size before returning. Continue casting float32 tensors
through WireTensor::from_f32_bytes_cast, but reject float16, bfloat16, and
float64 inputs instead of forwarding them; preserve passthrough only for
supported non-floating dtypes after validation.
---
Outside diff comments:
In `@model_gateway/src/routers/grpc/zmq_client.rs`:
- Around line 194-208: Update the fan-out flow around fan_out_requests,
translate_request, and zmq_multimodal::build_mm_features to preprocess inline
multimodal inputs once before the loop, then assign each generated
EngineCoreRequest a refcounted shared MmFeatures slice. Add validation enforcing
a bounded inline multimodal payload size before preprocessing, while preserving
existing request translation and stream behavior.
---
Nitpick comments:
In `@crates/engine_zmq_client/src/codec/tensor.rs`:
- Around line 158-159: Update the Float32 branch in decode_tensor and the
related from_raw_bytes/from_raw constructors to accept or take ownership of
Bytes, reusing the existing from_f32_bytes_cast allocation instead of calling
data.to_vec(). Preserve the current validation and decoding behavior while
eliminating the redundant payload copy.
In `@crates/engine_zmq_client/src/protocol/vllm/multimodal.rs`:
- Around line 248-260: Add the MmField::Flat variant, including representative
MmSlice/SliceSpec data, to the field_round_trips_python_factory_tuple test so
its nested tuple encoding and untagged decoding are verified alongside Batched
and Shared.
In `@model_gateway/src/routers/grpc/backend_client.rs`:
- Around line 368-379: Add a brief comment above zmq_tokenspeed_mm explaining
that it is retained for a future TokenSpeed wire path despite being currently
unreachable. Replace the unexplained true argument in MultimodalData::TokenSpeed
conversion with a named binding or concise explanation identifying the
serialization mode it selects, without changing behavior.
In `@model_gateway/src/routers/grpc/multimodal/assemble.rs`:
- Around line 114-127: Align the ZMQ runtime dispatch in the BackendClient::Zmq
arm with BackendClient::build_chat_request and build_messages_request, keeping
TokenSpeed explicit and routing other runtimes to vLLM. Update assemble_vllm to
receive the inline-transport intent so it sets shm_enabled, rdma_enabled, and
shm_min_bytes consistently, removing the post-assembly field overrides.
In `@model_gateway/src/routers/grpc/zmq_multimodal.rs`:
- Around line 384-405: Update tests using base_inputs and empty prompt_token_ids
to provide a prompt with at least nine tokens, so placeholder ranges are valid
once bounds checks are enforced. Add coverage for keep_on_cpu_keys through the
shared feature-building path and assert the resulting on-CPU behavior is
preserved. Run the pr-test-analyzer agent to verify the new and changed
functionality is adequately covered.
🪄 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: 88f0b447-fe17-4d59-aaaf-9d995717fd74
📒 Files selected for processing (9)
crates/engine_zmq_client/src/codec/tensor.rscrates/engine_zmq_client/src/protocol/vllm/mod.rscrates/engine_zmq_client/src/protocol/vllm/multimodal.rscrates/engine_zmq_client/src/protocol/vllm/request.rsmodel_gateway/src/routers/grpc/backend_client.rsmodel_gateway/src/routers/grpc/mod.rsmodel_gateway/src/routers/grpc/multimodal/assemble.rsmodel_gateway/src/routers/grpc/zmq_client.rsmodel_gateway/src/routers/grpc/zmq_multimodal.rs
d8562c1 to
447543f
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. |
There was a problem hiding this comment.
🧹 Nitpick comments (1)
crates/engine_zmq_client/src/protocol/vllm/multimodal.rs (1)
247-259: 🗄️ Data Integrity & Integration | 🔵 Trivial | ⚡ Quick win🟡 Nit: Add round-trip coverage for
MmField::Flat.The test only round-trips
BatchedandShared.Flatuses the distinctslicestuple encoding and factory conversion path. AddFlatcases for bothMmSlice::SliceandMmSlice::Slices.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 `@crates/engine_zmq_client/src/protocol/vllm/multimodal.rs` around lines 247 - 259, Extend field_round_trips_python_factory_tuple to include MmField::Flat cases containing both MmSlice::Slice and MmSlice::Slices, then verify each encodes and decodes back identically through the existing round-trip assertions. Run the pr-test-analyzer agent to confirm 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 `@crates/engine_zmq_client/src/protocol/vllm/multimodal.rs`:
- Around line 247-259: Extend field_round_trips_python_factory_tuple to include
MmField::Flat cases containing both MmSlice::Slice and MmSlice::Slices, then
verify each encodes and decodes back identically through the existing round-trip
assertions. Run the pr-test-analyzer agent to confirm coverage.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 0513b6d9-6ae8-4201-9187-1457d0920615
📒 Files selected for processing (9)
crates/engine_zmq_client/src/codec/tensor.rscrates/engine_zmq_client/src/protocol/vllm/mod.rscrates/engine_zmq_client/src/protocol/vllm/multimodal.rscrates/engine_zmq_client/src/protocol/vllm/request.rsmodel_gateway/src/routers/grpc/backend_client.rsmodel_gateway/src/routers/grpc/mod.rsmodel_gateway/src/routers/grpc/multimodal/assemble.rsmodel_gateway/src/routers/grpc/zmq_client.rsmodel_gateway/src/routers/grpc/zmq_multimodal.rs
🚧 Files skipped from review as they are similar to previous changes (8)
- crates/engine_zmq_client/src/codec/tensor.rs
- model_gateway/src/routers/grpc/mod.rs
- crates/engine_zmq_client/src/protocol/vllm/request.rs
- crates/engine_zmq_client/src/protocol/vllm/mod.rs
- model_gateway/src/routers/grpc/multimodal/assemble.rs
- model_gateway/src/routers/grpc/backend_client.rs
- model_gateway/src/routers/grpc/zmq_client.rs
- model_gateway/src/routers/grpc/zmq_multimodal.rs
Carry preprocessed multimodal features to a vLLM EngineCore over the direct ZMQ wire, replacing the earlier "not supported" rejection. - Add a typed MmFeatures to the vLLM EngineCore protocol and carry it on EngineCoreRequest in place of the untyped placeholder; tensors ride inline as raw views (no /dev/shm or RDMA pull on this wire). - Add zmq_multimodal, which performs the per-item mm-feature split the Python servicer would otherwise do, casting float tensors to the engine's model dtype (from_f32_bytes_cast). - Assemble vLLM multimodal data for a ZMQ backend with shm/rdma disabled; TokenSpeed multimodal over ZMQ stays rejected (no wire slot yet). - BackendClient converts assembled mm data to the runtime's proto, surfacing a backend/variant mismatch as a build error rather than a panic. Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com>
447543f to
20a273d
Compare
1 participant
Description
Problem
Multimodal requests (image/audio/video) cannot be served over a direct-ZMQ
engine: the adapter has no way to carry processed multimodal features and their
placeholder ranges across the vLLM EngineCore wire.
Solution
Type vLLM's multimodal feature payload (
mm_features) — feature specs,placeholder ranges, processed kwargs, and the tensor field factories
(
batched/flat/shared) — and assemble the gateway's multimodal inputsinto that wire shape, including the aux tensor frames.
Changes
crates/engine_zmq_client/src/protocol/vllm/multimodal.rs(new) —MmFeatureSpec,PlaceholderRange,MmKwargsItem, and theMmFieldfactory-tuple ser/de.crates/engine_zmq_client/src/codec/tensor.rs—WireTensorsupport used bymultimodal payloads.
crates/engine_zmq_client/src/protocol/vllm/{mod,request}.rs— carrymm_featureson the EngineCore request and emit aux tensor frames.model_gateway/src/routers/grpc/{mod,zmq_client,zmq_multimodal}.rs,.../multimodal/assemble.rs— assemble gateway multimodal inputs into the wirepayload (
zmq_multimodal.rsnew).model_gateway/src/routers/grpc/backend_client.rs— return a typed build erroron backend/variant mismatch for the ZMQ multimodal path (see note below).
Test Plan
cargo +nightly fmt -- --checkandcargo clippy --all-featuresclean.cargo testunit coverage inmultimodal.rs: field→Python factory-tupleencoding, factory-tuple round-trip, and feature-spec map encoding with the
tensor ext-3 payload.
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 (this PR) → 04-eos-forwarding →
05-harmony-stop-matcher → 06-e2e-infra → 07-e2e-tests. Review/merge after its
parent
zmq/02-structured-outputs.