refactor(grpc): split MultimodalData into backend-specific variants - #588
Conversation
Replace the unified MultimodalData struct with an enum of three backend-specific structs (SglangMultimodalData, VllmMultimodalData, TrtllmMultimodalData). Each variant carries only the fields its backend needs: - SGLang: pixel_values + model_specific + patch-only placeholders - vLLM: pixel_values + model_specific + structural placeholders + hashes + field keys - TRT-LLM: raw image bytes only (no tensor serialization) Introduce MultimodalIntermediate as a lightweight preparation output that holds preprocessing results without serializing tensors. The assembly into backend-specific data is deferred to request_building, where the target backend is known after worker selection. This avoids wasting work (e.g. TRT-LLM no longer serializes megabytes of pixel_values it never sends). Signed-off-by: Chang Su <chang.s.su@oracle.com>
|
Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits. |
github-actions
Bot
added
grpc
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review infoConfiguration used: Organization UI Review profile: ASSERTIVE Plan: Pro 📒 Files selected for processing (1)
📝 WalkthroughWalkthroughRefactors multimodal handling to produce a MultimodalIntermediate in preparation, adds assemble_multimodal_data(...) to build backend-specific MultimodalData variants (Sglang/Vllm/Trtllm) at request-build time, converts MultimodalData into an enum with per-backend types, and updates call sites with variant-aware matching and unreachable guards. Changes
Sequence DiagramsequenceDiagram
participant Client as Client/Request
participant Prep as PreparationStage
participant Intermediate as MultimodalIntermediate
participant Assembler as AssemblyLayer
participant Variant as MultimodalDataVariant
participant Proto as ProtoConversion
Client->>Prep: submit message (may include images)
Prep->>Intermediate: produce MultimodalIntermediate (preprocessed, images, placeholders, layouts)
Prep-->>Client: return ProcessedMessages with intermediate
Client->>Assembler: request_building invokes assemble_multimodal_data(intermediate, GrpcClient)
Assembler->>Assembler: select backend assembler (Sglang/Vllm/Trtllm)
Assembler->>Variant: assemble backend-specific MultimodalData
Variant->>Proto: call into_proto()
Proto-->>Client: return backend-specific proto payload
Estimated code review effort🎯 4 (Complex) | ⏱️ ~55 minutes Possibly related issues
Possibly related PRs
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 |
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 significantly refactors the handling of multimodal data within the system's gRPC layer. By replacing a monolithic 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
|
![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 is a well-executed refactoring that improves the handling of multimodal data. By splitting the monolithic MultimodalData struct into backend-specific variants within an enum, the code is now cleaner and more efficient, as each backend only deals with the data it requires. The introduction of MultimodalIntermediate to defer data assembly until after worker selection is a smart design choice that enhances modularity.
I have a couple of suggestions for further improvement:
- An optimization to reduce memory allocations and copies when collecting raw image data, aligning with the principle of avoiding unnecessary intermediate allocations.
- A small refactoring to make a data serialization function more idiomatic and concise.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
…pecific Replace mutable HashMap loop with idiomatic filter_map().collect() pattern. Signed-off-by: Chang Su <chang.s.su@oracle.com>
There was a problem hiding this comment.
Actionable comments posted: 1
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/multimodal.rs (1)
453-551: 🧹 Nitpick | 🔵 TrivialAdd unit tests for backend assembly invariants.
This PR adds critical backend-specific assembly paths (
assemble_sglang,assemble_vllm,assemble_trtllm) but current tests in this file do not cover them. Please add focused tests for placeholder selection, hash/image mapping, and serialization shapes/dtypes to prevent silent payload regressions.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@model_gateway/src/routers/grpc/multimodal.rs` around lines 453 - 551, Add unit tests that exercise assemble_sglang, assemble_vllm, and assemble_trtllm to lock-in assembly invariants: (1) For assemble_sglang, test that when MultimodalIntermediate.patch_offsets is Some the returned SglangMultimodalData.mm_placeholders uses those patch-only offsets and when None it falls back to MultimodalIntermediate.placeholders (check offset/length types and ordering); (2) For assemble_vllm, assert mm_hashes matches intermediate.images[*].hash and image-derived fields (pixel_values, pixel_values_shape, model_specific_tensors) are populated, and validate batched_keys and flat_keys come from PreprocessedImages::batched_keys and ::flat_keys respectively; (3) For assemble_trtllm, verify image_data is a Vec of the images' raw_bytes; and (4) directly test serialize_pixel_values and serialize_model_specific produce correct byte-layout (little-endian f32 bytes and shape Vec<u32>) and that model_specific_to_tensor_bytes entries become keys in the HashMap of TensorBytes. Use small deterministic PreprocessedImages and MultimodalIntermediate fixtures to assert types, lengths, and exact bytes where practical.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@model_gateway/src/routers/grpc/multimodal.rs`:
- Around line 461-468: The fallback currently only triggers when
intermediate.patch_offsets is None, but when patch_offsets is Some(vec![]) we
must treat it as unavailable; update the mm_placeholders selection to use
patch_offsets only if it is Some and non-empty (e.g., check .as_ref().filter(|v|
!v.is_empty()) or match and treat empty vec as None) and otherwise build
mm_placeholders from intermediate.placeholders (referencing mm_placeholders,
intermediate.patch_offsets, and intermediate.placeholders) so empty
patch_offsets will not produce an empty placeholder list.
---
Outside diff comments:
In `@model_gateway/src/routers/grpc/multimodal.rs`:
- Around line 453-551: Add unit tests that exercise assemble_sglang,
assemble_vllm, and assemble_trtllm to lock-in assembly invariants: (1) For
assemble_sglang, test that when MultimodalIntermediate.patch_offsets is Some the
returned SglangMultimodalData.mm_placeholders uses those patch-only offsets and
when None it falls back to MultimodalIntermediate.placeholders (check
offset/length types and ordering); (2) For assemble_vllm, assert mm_hashes
matches intermediate.images[*].hash and image-derived fields (pixel_values,
pixel_values_shape, model_specific_tensors) are populated, and validate
batched_keys and flat_keys come from PreprocessedImages::batched_keys and
::flat_keys respectively; (3) For assemble_trtllm, verify image_data is a Vec of
the images' raw_bytes; and (4) directly test serialize_pixel_values and
serialize_model_specific produce correct byte-layout (little-endian f32 bytes
and shape Vec<u32>) and that model_specific_to_tensor_bytes entries become keys
in the HashMap of TensorBytes. Use small deterministic PreprocessedImages and
MultimodalIntermediate fixtures to assert types, lengths, and exact bytes where
practical.
…ssembly Fall back to full structural placeholder ranges when patch_offsets is Some(vec![]) rather than only on None, preventing empty mm_placeholders from being sent to SGLang. Signed-off-by: Chang Su <chang.s.su@oracle.com>
Signed-off-by: Chang Su <chang.s.su@oracle.com>
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@model_gateway/src/routers/grpc/multimodal.rs`:
- Around line 453-514: Add unit tests that exercise assemble_sglang,
assemble_vllm, and assemble_trtllm using minimal MultimodalIntermediate
fixtures: verify assemble_sglang returns SglangMultimodalData with correct
image_data, pixel_values, pixel_values_shape, model_specific_tensors,
im_token_id, and that mm_placeholders uses patch_offsets when present and
non-empty and falls back to placeholders mapping (offset,length) when
patch_offsets is None or empty; verify assemble_vllm returns VllmMultimodalData
with expected pixel_values, pixel_values_shape, model_specific_tensors,
im_token_id, mm_placeholders (from placeholders), mm_hashes (from images.hash),
and batched_keys/flat_keys (use PreprocessedImages::batched_keys/flat_keys
deterministically); verify assemble_trtllm returns TrtllmMultimodalData with
image_data from images.raw_bytes; create small helper builders for
MultimodalIntermediate/PreprocessedImages to set images,
preprocessed.model_specific, im_token_id, placeholders and patch_offsets to
assert both code paths.
- Around line 545-553: The serializer currently drops unsupported
ModelSpecificValue entries silently in serialize_model_specific by using
filter_map with model_specific_to_tensor_bytes; change this so dropped keys are
surfaced: either (preferred) make serialize_model_specific return
Result<HashMap<String,TensorBytes>, SerializeError> and return an Err listing
the unsupported keys (using model_specific_to_tensor_bytes failures), or (if
non-breaking) collect the dropped keys and emit a warning (e.g., tracing::warn!
or your crate logger) that includes the key names and their value types before
returning the partial map; apply the same fix to the analogous map-to-tensor
conversion at the other site referenced (lines ~590-592) so unsupported entries
are consistently reported.
Signed-off-by: Chang Su <chang.s.su@oracle.com>
1 participant
Description
Problem
The unified
MultimodalDatastruct carries fields for all three backends (vLLM, SGLang, TRT-LLM), but each backend only uses a subset:pixel_valuesit never sendspixel_values,model_specific_tensors,patch_offsets,im_token_id— never usesmm_hashes,batched_keys,flat_keyspixel_values,model_specific_tensors,mm_hashes,batched_keys,flat_keys— never usessglang_patch_offsetsorimage_dataAs backends diverge (e.g. vLLM's
keep_on_cpufield configs, TRT-LLM's upcoming preprocessed tensor support), the god struct keeps growing with fields that most backends ignore.Solution
Replace the unified
MultimodalDatastruct with an enum of three backend-specific structs. IntroduceMultimodalIntermediateas a lightweight preparation output that holds preprocessing results without serializing tensors. The assembly into backend-specific data is deferred to request_building, where the target backend is known after worker selection.Changes
proto_wrapper.rs: Replace unifiedMultimodalDatastruct withMultimodalDataenum +SglangMultimodalData,VllmMultimodalData,TrtllmMultimodalDatastructs. Moveinto_*_proto()to per-structinto_proto().multimodal.rs: AddMultimodalIntermediatestruct. Extractserialize_pixel_values()andserialize_model_specific()helpers. Addassemble_sglang(),assemble_vllm(),assemble_trtllm(),assemble_multimodal_data(). Updateprocess_multimodal()to return intermediate. Removebuild_multimodal_data().mod.rs:ProcessedMessages.multimodal_data→multimodal_intermediate.preparation.rs: Store intermediate instead ofMultimodalData.request_building.rs: Callassemble_multimodal_data(intermediate, builder_client)after worker selection.client.rs: Updatebuild_chat_request()to match on enum variant.Test Plan
Checklist
cargo +nightly fmtpassescargo clippy --all-targets --all-features -- -D warningspassesSummary by CodeRabbit
Refactor
New Features