feat(mm-routing): skip MM-aware routing work when KV router is load-balancing only - #10515
feat(mm-routing): skip MM-aware routing work when KV router is load-balancing only#10515krishung5 wants to merge 4 commits into
Conversation
…alancing only When the KV router is enabled but not using prefix overlap (--load-aware, overlap_score_credit=0, use_kv_events=false, or a non-KV mode), the frontend still computed per-image mm_hash + dim fetches for MM-aware routing. That's wasted work: the router ignores prefix overlap, so there's no KV-reuse benefit (e.g. when prefix caching is disabled on the engine) but the routing cost is still paid. Gate the routing-only MM work on whether the router actually uses overlap: - Signal: KvRouterConfig::should_subscribe_to_kv_events() (use_kv_events && overlap_score_credit > 0). The model watcher computes router_mode == KV && that, and passes mm_routing_enabled into OpenAIPreprocessor. - Rust frontend (covers vLLM + SGLang via --dyn-chat-processor dynamo): when disabled, skips the URL-passthrough fetch_image_dims, the mm_image_entries build, and gather_mm_exact_routing_info. Media transfer (frontend decode) is unaffected. - vLLM chat-processor (--dyn-chat-processor vllm): skips building mm_routing_info; process_inputs and the mm_kwargs transfer (#8065) are kept — they serve the backend, not routing. Non-KV-mode / no-router-context callers (C API, static pipelines, tests) default to enabled, so behavior is unchanged there. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…-when-load-balancing
github-actions
Bot
added
feat
frontend
WalkthroughThis PR introduces an ChangesMultimodal routing feature flag gating
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
components/src/dynamo/frontend/vllm_processor.py (1)
878-898:⚠️ Potential issue | 🟠 Major | 🏗️ Heavy liftDerive
mm_routing_enabledfrom the effective worker-set router config, not the global frontend config.
lib/llm/src/discovery/watcher.rsresolvesrouter_config = card.router_config.as_ref().unwrap_or(&self.router_config)before it builds the routed engine, but Lines 883-887 recompute the flag fromFrontendConfig. If an MDC overridesuse_kv_eventsoroverlap_score_credit, the Python preprocessor and the Rust router can disagree: you either keep paying the MM-routing cost unnecessarily or disable MM overlap routing while the router still expects it.Please source this from the MDC override as well, or thread the already-resolved bool from Rust into
chat_engine_factory, so both paths share one gate.🤖 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 `@components/src/dynamo/frontend/vllm_processor.py` around lines 878 - 898, mm_routing_enabled is being computed from self.config (FrontendConfig) but must reflect the effective worker/router config that may be overridden by the MDC; update the code that sets mm_routing_enabled so it reads the resolved router config used to build the routed engine (the same boolean resolved in lib/llm/src/discovery/watcher.rs where router_config = card.router_config.as_ref().unwrap_or(&self.router_config)), or instead accept and thread a pre-resolved bool from chat_engine_factory into the VllmProcessor call; specifically, ensure the mm_routing_enabled passed into VllmProcessor comes from the worker-set/router_config override (or the threaded resolved flag) rather than directly from getattr(self.config, ...) so both Python preprocessor and the Rust router share the same gate.lib/llm/src/preprocessor.rs (1)
544-554:⚠️ Potential issue | 🟠 Major | ⚡ Quick winGate the dim-fetch client warmup on
mm_routing_enabled.Lines 551-554 still force
DIM_FETCH_MEDIA_FETCHER/DIM_FETCH_HTTP_CLIENTwhenever the model is MM-routable, even whenmm_routing_enabledisfalse. In load-balancing-only or non-KV setups, that means this path still does the TLS/env/media-fetcher initialization the new flag is supposed to skip, and a bad dim-fetch configuration can still fail worker-set registration even though MM overlap routing is disabled.Suggested fix
- #[cfg(feature = "mm-routing")] - if image_token_counter.is_some() || routing_image_token_id.is_some() { + #[cfg(feature = "mm-routing")] + if mm_routing_enabled + && (image_token_counter.is_some() || routing_image_token_id.is_some()) + { std::sync::LazyLock::force(&DIM_FETCH_MEDIA_FETCHER); std::sync::LazyLock::force(&DIM_FETCH_HTTP_CLIENT); }🤖 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 `@lib/llm/src/preprocessor.rs` around lines 544 - 554, The warmup logic currently forces DIM_FETCH_MEDIA_FETCHER and DIM_FETCH_HTTP_CLIENT when a model is MM-routable (checks image_token_counter or routing_image_token_id) even if mm_routing_enabled is false; update the condition that surrounds std::sync::LazyLock::force(&DIM_FETCH_MEDIA_FETCHER) and ::force(&DIM_FETCH_HTTP_CLIENT) to also require mm_routing_enabled to be true (i.e., only trigger the dim-fetch client warmup when mm_routing_enabled && (image_token_counter.is_some() || routing_image_token_id.is_some())); adjust the same guarded block that currently uses #[cfg(feature = "mm-routing")] so the mm_routing_enabled runtime flag prevents initialization in load-balancing/non-KV setups.
🤖 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.
Outside diff comments:
In `@components/src/dynamo/frontend/vllm_processor.py`:
- Around line 878-898: mm_routing_enabled is being computed from self.config
(FrontendConfig) but must reflect the effective worker/router config that may be
overridden by the MDC; update the code that sets mm_routing_enabled so it reads
the resolved router config used to build the routed engine (the same boolean
resolved in lib/llm/src/discovery/watcher.rs where router_config =
card.router_config.as_ref().unwrap_or(&self.router_config)), or instead accept
and thread a pre-resolved bool from chat_engine_factory into the VllmProcessor
call; specifically, ensure the mm_routing_enabled passed into VllmProcessor
comes from the worker-set/router_config override (or the threaded resolved flag)
rather than directly from getattr(self.config, ...) so both Python preprocessor
and the Rust router share the same gate.
In `@lib/llm/src/preprocessor.rs`:
- Around line 544-554: The warmup logic currently forces DIM_FETCH_MEDIA_FETCHER
and DIM_FETCH_HTTP_CLIENT when a model is MM-routable (checks
image_token_counter or routing_image_token_id) even if mm_routing_enabled is
false; update the condition that surrounds
std::sync::LazyLock::force(&DIM_FETCH_MEDIA_FETCHER) and
::force(&DIM_FETCH_HTTP_CLIENT) to also require mm_routing_enabled to be true
(i.e., only trigger the dim-fetch client warmup when mm_routing_enabled &&
(image_token_counter.is_some() || routing_image_token_id.is_some())); adjust the
same guarded block that currently uses #[cfg(feature = "mm-routing")] so the
mm_routing_enabled runtime flag prevents initialization in load-balancing/non-KV
setups.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 275b4272-6080-4977-a49e-aacfc905eaf3
📒 Files selected for processing (5)
components/src/dynamo/frontend/vllm_processor.pylib/kv-router/src/scheduling/config.rslib/llm/src/discovery/watcher.rslib/llm/src/entrypoint/input/common.rslib/llm/src/preprocessor.rs
| let mm_routing_enabled = router_config.router_mode == RouterMode::KV | ||
| && router_config | ||
| .kv_router_config | ||
| .should_subscribe_to_kv_events(); |
There was a problem hiding this comment.
should_subscribe_to_kv_events() is false for --no-router-kv-events, but approximate KV mode still does prefix-overlap routing from routing decisions, so this disables MM routing while the router still scores prefix overlap. Fix: derive mm_routing_enabled from KV mode plus overlap scoring being active, not from KV event subscription.
🤖 AI Fix
In lib/llm/src/discovery/watcher.rs inside ModelWatcher::do_worker_set_registration, replace the mm_routing_enabled expression that calls router_config.kv_router_config.should_subscribe_to_kv_events() with router_config.router_mode == RouterMode::KV && router_config.kv_router_config.overlap_score_credit > 0.0 or an equivalent KvRouterConfig::uses_prefix_overlap() helper so --no-router-kv-events approximate routing keeps MM routing enabled.
| # or any non-KV mode → skip the wasted mm_routing_info build. | ||
| mm_routing_enabled = bool( | ||
| getattr(self.config, "router_mode", None) == "kv" | ||
| and getattr(self.config, "use_kv_events", False) |
There was a problem hiding this comment.
Including use_kv_events in the gate disables MM routing for --no-router-kv-events approximate KV mode even though that mode still predicts cache state and scores prefix overlap. Fix: keep mm_routing_enabled true for KV mode whenever overlap scoring is enabled, independent of event subscription.
🤖 AI Fix
In components/src/dynamo/frontend/vllm_processor.py in EngineFactory.chat_engine_factory, remove the and getattr(self.config, "use_kv_events", False) term from the mm_routing_enabled expression and leave the gate based on KV router mode and positive overlap_score_credit, or consume the same resolved uses_prefix_overlap boolean threaded from Rust.
The frontend skipped MM-aware routing whenever the KV router was not subscribed to KV events. That wrongly disabled it for approximate KV mode (`--no-router-kv-events`), where events are off but the router still scores prefix overlap — so the per-image `mm_hash` is still needed. Gate on overlap scoring directly via a new `KvRouterConfig::uses_prefix_overlap()` (`overlap_score_credit > 0`), independent of `use_kv_events`. Only `--load-aware` / `overlap_score_credit=0` or a non-KV mode now skips the work. Apply the same predicate on the vLLM chat-processor path, and update the unit test to pin the approximate-KV case. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…-when-load-balancing
2 participants
Overview:
When the KV router is enabled but not using prefix overlap —
--load-aware,overlap_score_credit=0,use_kv_events=false, or any non-KV mode — the frontend still computed per-imagemm_hash+ dim fetches for MM-aware routing. That's wasted work: the router ignores prefix overlap, so there's no KV-reuse benefit (e.g. when prefix caching is disabled on the engine) but the routing cost is still paid.Details:
Signal:
KvRouterConfig::should_subscribe_to_kv_events()(use_kv_events && overlap_score_credit > 0). The--load-awarepreset setsoverlap_score_credit=0, use_kv_events=False, so it evaluates false. The model watcher computesrouter_mode == KV && should_subscribe_to_kv_events()and threadsmm_routing_enabledinto the preprocessor.--dyn-chat-processor dynamo, covers vLLM and SGLang): when disabled,OpenAIPreprocessorskips the URL-passthroughfetch_image_dims(HTTP header fetches), themm_image_entriesbuild, andgather_mm_exact_routing_info. Media transfer (frontend decode / NIXL) is unaffected.--dyn-chat-processor vllm): skips buildingmm_routing_info;process_inputsand themm_kwargstransfer (feat(multimodal): move MM routing into vLLM frontend processor #8065) are kept — they serve the backend, not routing.Follow-up (not in this PR): in the vLLM path, when transfer is also off (
DYNAMO_DISABLE_NIXL_MM=1),process_inputsitself could be skipped — needs the backend raw-request path validated.Closes DIS-2156
Where should the reviewer start?
lib/llm/src/discovery/watcher.rs— themm_routing_enabledcomputation.lib/llm/src/preprocessor.rs— thegather_multi_modal_data/gather_mm_exact_routing_infogates.components/src/dynamo/frontend/vllm_processor.py— themm_routing_infogate.Verification
-D warningsin both default +mm-routingfeature configs;cargo fmtclean.cargo test -p dynamo-llm --lib --features mm-routing: 1296 passed. New unit testshould_subscribe_to_kv_events_gates_overlap_routingpins the gate matrix (default→on,--load-aware/either-knob-off→off).--load-aware) — running.Related Issues
🚫 This PR is NOT linked to an issue:
🤖 Generated with Claude Code
Summary by CodeRabbit
New Features
Tests