feat(router): forward media references to vLLM workers that process them - #2400
feat(router): forward media references to vLLM workers that process them#2400CatherineSue wants to merge 1 commit into
Conversation
Make router-side multimodal preprocessing optional for vLLM gRPC workers.
When a request resolves to worker-side processing, preparation keeps the
rendered prompt with one unexpanded placeholder anchor per media item and
stashes the media plan; request building attaches it as
GenerateRequest.media_refs instead of pixel tensors, and the worker's own
processor fetches, expands and hashes.
Resolution runs in preparation, before worker selection, from
SMG_MM_PROCESSING (auto|router|worker, default auto): `auto` forwards only
when every registered worker of the model advertises `mm_processor` and
the model's spec opts in via ModelProcessorSpec::worker_expandable (Qwen3-VL
and Llama 4; Phi-3.5 and Qwen-VL render anchors vLLM does not target);
`router` is the kill switch; `worker` is strict and rejects hints and
un-opted models with named 400s. Worker selection filters candidates to
advertising workers on both PD legs, keeps the pin across retries through
WireConstraint.requires_media_refs, and sheds 503
no_media_ref_capable_worker when none exist; a post-selection check also
rejects EPD and unadvertised URL schemes. media_refs is a sibling of
mm_inputs, so the PD decode clone carries it and both legs process the
same references. ZMQ workers never qualify. Resolutions are counted in
smg_mm_processing_total{model,mode,reason} and logged once per change.
Signed-off-by: Chang Su <8605658+CatherineSue@users.noreply.github.com>
|
Important Draft PR not reviewedDraft PRs are not automatically reviewed by default.
To automatically review draft PRs, update your CodeRabbit configuration: reviews:
auto_review:
drafts: trueWarning Your free Security trial is over. An organization admin can upgrade to Advanced for continuous pull request security review or dismiss this notice. Comment |
github-actions
Bot
added
grpc
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
| wire: Option<WireConstraint>, | ||
| media_refs: bool, | ||
| ) -> Response { | ||
| if media_refs { |
There was a problem hiding this comment.
🔴 Important: This short-circuit runs before any diagnosis, so in worker mode every selection failure is reported as no_media_ref_capable_worker, including the ones that have nothing to do with capability.
Concretely, under auto the mode only resolves to Worker when the whole registered fleet already advertises mm_processor, so the realistic way to reach selection_failure(.., media_refs = true) is that the capable workers were filtered out for another reason — all overloaded, unhealthy, or circuit-broken. In that case the request now gets:
- error code
no_media_ref_capable_workerwith a message telling the operator toset SMG_VLLM_MM_PROCESSOR on the workers(the labels are already set — wrong instruction on-call will chase), - no
Retry-Afterheader and nomark_non_retryable, unlikeoverload::shed_if_all_overloaded, so the load-shed contract documented inoverload.rs("a terminal shed is what keeps the counter per-request rather than per-attempt") is lost for multimodal requests, - no
record_worker_overload_shedcounter increment, so the shed is invisible in metrics, - a 503 where a genuinely absent model would previously have been a 404 (
SMG_MM_PROCESSING=workerwith zero registered workers).
Suggest only taking this branch when the capability filter is actually what emptied the pool — e.g. run the existing leg_candidates loop first and fall back to this 503 only if no candidate in any leg satisfies multimodal::worker_accepts_media_refs, and mark it non-retryable (or route it through overload::shed-style construction) so it behaves like the other 503 the selection stage emits.
| let Some(primary) = legs.first() else { | ||
| return Err(MmRefsError::WorkerNotCapable); | ||
| }; | ||
| let accepted = worker_media_ref_schemes(*primary); |
There was a problem hiding this comment.
🟡 Nit: The scheme gate only consults legs.first() — the prefill leg in PD — but the whole point of the sibling media_refs field is that both legs fetch and process the same references (clone_without_mm_pixels carries media_refs unchanged, and select_pd_pair filters both pools). If prefill advertises mm_media_ref_schemes=http,https,file (started with --allowed-local-media-path) and decode advertises only http,https, a file:// reference passes this check and then fails on the decode leg at dispatch time, which is a much worse failure mode than the 400 this function exists to produce.
Intersecting the schemes across all legs (and treating "advertised none" as unconstrained, as today) would keep the check aligned with the both-legs rule the rest of the PR enforces.
| MmProcessingMode::Router => (MmProcessing::Router, "config"), | ||
| MmProcessingMode::Worker => { | ||
| if !plan.is_forwardable() { | ||
| return Err(MmRefsError::HintUnsupported); |
There was a problem hiding this comment.
🟡 Nit: In strict worker mode any non-forwardable plan is reported as HintUnsupported, but MediaPlan::is_forwardable is false for far more than per-item hints: inline ImageData/VideoData, audio parts, and ImageEmbeds all fail it. A request with a base64 image_data part then gets code multimodal_hint_unsupported_in_worker_mode and the message "per-item media hints (max_long_side_pixel, fps) cannot be forwarded to a worker", which names a field the request never set.
assemble_media_refs already distinguishes these cases correctly (UnsupportedPart("inline image bytes") etc.). Reusing that classification here — e.g. a helper that returns the specific MmRefsError for the first non-forwardable part — would make the 400 actionable.
| headers, | ||
| rid_key, | ||
| wire, | ||
| ctx.wire.requires_media_refs, |
There was a problem hiding this comment.
🟡 Nit: reselect re-applies the capability filter (via requires_media_refs) but never re-runs ensure_selection_supports_media_refs, so the URL-scheme half of the gate only ever runs on the first selection. A retry can land on a capable worker that advertises a different mm_media_ref_schemes set than the one validated at ingress (e.g. only the original worker was started with --allowed-local-media-path, so only it accepts file://), and the request is then dispatched with a reference that worker will refuse.
The plan itself is gone by then (request_building does ctx.state.multimodal_refs.take()), so fixing this needs the schemes carried into DispatchContext — e.g. stash the set of schemes the plan actually uses alongside WireConstraint::requires_media_refs and filter candidates on it, which would also subsume the separate post-selection pass in execute.
| media_items = media_plan.parts().len(), | ||
| "Forwarding media references for worker-side processing" | ||
| ); | ||
| multimodal_refs = Some(media_plan); |
There was a problem hiding this comment.
🟡 Nit: In the worker branch token_ids stays unexpanded, and those are exactly the ids that become RoutingSnapshot::token_ids and feed worker selection (prep.token_ids() → tokens in WorkerSelectionStage::execute). So for a multimodal request the prefix/cache-aware policies (cache_aware.rs, prefix_hash.rs, bucket.rs) now see one anchor token per image instead of the full pad run — a prompt the router models as ~30 tokens while the worker actually prefills a few thousand.
The routing stays correct (vLLM re-expands), but for a fleet flipped to worker mode the load/prefix accounting for image traffic is off by an order of magnitude, which is the kind of thing that only shows up as unexplained imbalance later. Worth either documenting explicitly here (and in the mode's docs) or feeding the policies an estimated expanded length. Same applies to the mirrored branch in messages/preparation.rs:246.
| let legs: Vec<&dyn Worker> = match workers { | ||
| WorkerSelection::Single { worker } => vec![worker.as_ref()], | ||
| WorkerSelection::Disaggregated { | ||
| encode_assignments: Some(_), |
There was a problem hiding this comment.
🟡 Nit: This guard is unreachable from the pipeline, and what happens instead is a silent EPD bypass.
encode_assignments is Some(_) only when assign_encode_workers got a non-empty hash list, and encode_item_hashes returns an empty vec whenever multimodal_intermediate is None (worker_selection.rs:858) — which is exactly the worker-mode state. So an EPD router in worker mode always lands in the Disaggregated { encode_assignments: None, .. } arm and the request is dispatched P/D-only with the refs, never hitting EncodeNotSupported.
That means an EPD deployment that sets SMG_VLLM_MM_PROCESSOR fleet-wide (encode workers are vLLM gRPC workers under the same model, so get_by_model sees them in the auto uniformity check) silently stops using its encode fleet for images: the dedicated encode GPUs go idle and the prefill workers do the vision work, with nothing in the logs or the smg_mm_processing_total reason to say so — auto reports auto_uniform.
If EPD is meant to be out of scope for worker mode, the cleanest fix is to resolve to Router in resolve_mm_processing when the router runs in EPD mode (reason e.g. epd_unsupported), rather than relying on a post-selection check that cannot fire.
1 participant
Description
Problem
The gRPC router fetches, preprocesses and placeholder-expands every multimodal request itself, before the backend is even known, and ships pixel tensors to vLLM. With the wire contract from the previous PR in place, the router still has no way to skip that work and hand a capable vLLM worker the media references instead.
Solution
Part 2 of the media-refs stack (stacked on the proto + servicer PR). Router-side preprocessing becomes optional per model:
SMG_MM_PROCESSINGselectsauto(default),router(kill switch) orworker(strict). Underautoa request is forwarded only when the plan is forwardable (plain image/video URLs, no MiniMax per-item hints), the model's spec opts in, and every registered worker of the model is a vLLM gRPC worker advertisingmm_processor. Otherwise today's path runs, byte-identical, and the reason is counted insmg_mm_processing_total{model,mode,reason}and logged once per change.ModelProcessorSpec::worker_expandabledefaults to false. Qwen3-VL (image, video) and Llama 4 (image) opt in because their anchors are exactly the single token vLLM's prompt updates target. Phi-3.5 (<|image|>vs vLLM's<|image_{i}|>) and Qwen-VL v1 (<image>vs[img_start, img_end]) stay on the router path.MediaPlan; worker selection filters candidates to advertising workers on both PD legs, pins retries throughWireConstraint.requires_media_refs, and sheds 503no_media_ref_capable_workerwhen none exist; a post-selection check rejects EPD and URL schemes the worker did not advertise; request building attachesmedia_refsvia a post-build setter (no newMultimodalDatavariant), caps inlinedata:payloads with the existingSMG_*_MAX_INPUT_BYTES, and rejects ZMQ clients.media_refsis a sibling ofmm_inputs, soclone_without_mm_pixelscarries it unchanged: both legs process the same references, and the decode leg gets vLLM-computed hashes and M-RoPE grids natively (nocache_salt).Changes
crates/protocols/src/worker.rs:MmProcessingMode { Auto, Router, Worker }.crates/multimodal:ModelProcessorSpec::worker_expandable(+ Qwen3-VL, Llama 4 opt-ins);media::{image,video}_max_input_bytesmade public.model_gateway/src/routers/grpc/multimodal/refs.rs(new): resolution, capability predicates, selection check,assemble_media_refs, named error codes.multimodal/{plan,config,mod}.rs:MediaPlan::{parts,is_forwardable},PlaceholderTokens::worker_expandable,MultimodalComponents::{processing, mm_mode_log}fromSMG_MM_PROCESSING.context.rs:ProcessingState::{multimodal_refs, media_refs_forwarded},WireConstraint::requires_media_refs.regular/stages/{chat,messages}/preparation.rs: worker-mode branch after anchor validation.common/stages/worker_selection.rs: candidate filter, retry pin, 503, post-selection gate.regular/stages/{chat,messages}/request_building.rs: attach refs, ZMQ rejection.proto_wrapper.rs:set_vllm_media_refs/has_vllm_media_refs;metrics.rs:smg_mm_processing_total.The diff is above the usual size guideline because roughly 40% of it is tests; the halves (resolution + payload vs pipeline wiring) do not ship independently, so they are kept together.
Test Plan
Not run here: GPU e2e. The
E2E_MM_PROCESSING=workerlane (Qwen3-VL chat + PD suites withSMG_VLLM_MM_PROCESSOR=inprocesson the workers) is the next PR in the stack.Checklist
cargo +nightly fmtpassescargo clippy --all-targets --all-features -- -D warningspasses (run without--all-featureslocally; that needs system OpenCV)