feat(mm-routing): lightseek-mm Rust-frontend MM-aware KV routing

[krishung5]

Overview:

Adds a Rust-frontend MM-aware KV routing path for multimodal workloads, using the new lightseek-mm cargo feature for image token expansion.

Closes DIS-1930

Details:

How MM-aware routing works.

The frontend now extracts a structured mm_routing_info from each multimodal request before it picks a worker:

request {text + N images}
   │
   ▼
for each image:
   ├─ compute mm_hash             ── xxh3_64 over normalized URL (default path)
   │                                 OR over decoded bytes (frontend-decoding path)
   ├─ fetch image dims (W, H)     ── HTTP Range read, cached
   └─ count tokens for (W, H)     ── per-VLM-family math (Qwen3-VL etc.) via llm-multimodal from LightSeek
   │
   ▼
mm_routing_info {
   image_multiset: [(mm_hash, token_count), ...],    # sorted, dedup-aware
   total_image_tokens: sum,
}
   │
   ▼
KV router selects worker

URL-passthrough vs frontend-decoding.

Two paths share the same routing logic but differ on where image bytes live:

• URL-passthrough (default). Request body carries image_url. Frontend hashes the normalized URL string (cache-buster query params stripped), routes the request, and forwards the raw URL to the chosen worker. Worker fetches and decodes the image. mm_processor_cache is keyed on decoded-bytes hash, so re-requests of the same URL hit; cross-URL same-content requests don't.
• Frontend-decoding (--frontend-decoding). Frontend's MediaLoader (LRU-cached) fetches and decodes the image once, ships the decoded RGB tensor to the chosen worker via NIXL/RDMA, and routes by content hash (xxh3_64 over the registered SystemStorage bytes — see RdmaMediaDataDescriptor::content_hash()). Worker skips both fetch and decode entirely. Same image content reached through different signed URLs collides on the same routing key.

The dim-cache primitive.

Image-token counts depend on the dimensions of each image — but the routing decision happens before the worker has decoded anything. So the frontend has to know (W, H) for each image at request-dispatch time. To avoid an HTTP Range fetch per image per request, dimensions are cached:

• Sharded bounded LRU: moka::future::Cache<u64, (u32, u32)> with 100 k-entry capacity and 24 h TTL. Lock-free reads, sharded write locks — handles concurrent dispatch with no contention bottleneck.
• Atomic singleflight: try_get_with collapses concurrent in-flight fetches for the same mm_hash into a single HTTP request — multiple workers can dispatch the same request in parallel without 4 threads racing for the same 4 KB Range read.

The cache is feature-gated and bounded so even unbounded URL pools (e.g. signed-URL refresh, image proxies) don't blow up frontend memory.

Backend behavior unchanged.

This PR only touches the Rust frontend's routing decision. vLLM workers are unmodified — they still receive image_url (or the decoded RDMA tensor in FD mode) and run their normal preprocessing pipeline. The mm_processor_cache keying that was already content-based is what makes the routing decision deliver actual cache hits.

Where should the reviewer start?

1. lib/llm/src/preprocessor.rs (660-line diff) — the routing-decision code path, all under cfg(feature = "lightseek-mm"). Key call sites:
   • fetch_image_dims — the moka cache + singleflight pattern; this is what avoids per-request HTTP overhead.
   • mm_routing_info construction in the URL-passthrough branch and the FD branch; the only difference is the mm_hash source (URL string vs descriptor.content_hash()).
2. lib/llm/src/preprocessor/media/rdma.rs — small (~23 lines added) but load-bearing: content_hash() is what makes FD-path routing collide across (signed) URL changes.
3. lib/llm/src/preprocessor/lightseek_mm.rs (new file) — model_id → per-VLM-family token-count function dispatch (Qwen3-VL etc.). Thin wrapper over llm-multimodal.
4. lib/llm/Cargo.toml + workspace Cargo.toml — confirm the lightseek-mm feature flag is wired the way the team prefers, and that default = ["block-manager", "lightseek-mm"] (default-on) is acceptable. Flip to default = ["block-manager"] for opt-in.
5. components/src/dynamo/vllm/handlers.py — 11-line FD-path error-yield fix; not part of routing logic but blocked aiperf parsing on FD failures, so it's in this PR.
6. docs/features/multimodal/multimodal-kv-routing.md — feature description for users.

Upstream PR: smg-project/smg#1459 to avoid pulling the openssl/native-tls toolchain.

Related Issues:

• Relates to GitHub issue: #xxx

---

Summary by CodeRabbit

Release Notes

• New Features

  • Added in-process multimodal (image) aware KV-cache routing support enabled by default
  • Added launch examples for multiple router topologies: standard KV router, chat-processor variant, and frontend-decoding variant with configurable worker and frontend replicas

• Documentation

  • Updated multimodal KV routing documentation with detailed routing flows, launch commands, and expanded configuration guidance

• Bug Fixes

  • Improved error handling for Qwen-VL multimodal image metadata failures

• Tests

  • Added end-to-end tests for multimodal routing covering content-hashing, URL normalization, and cache warm/cold behavior

[github-actions]

🌿 Fern Docs Preview: https://nvidia-preview-e78b8326-317e-4fa5-a3bb-cd24536f4d09.docs.buildwithfern.com/dynamo/dev

[coderabbitai]

Walkthrough

This PR implements Lightseek-based multimodal-aware KV cache routing via a Rust frontend. It adds image token counting, content-addressed hashing (mm_hash), MM routing metadata building, and per-worker ZMQ KV event orchestration. Includes launch scripts, E2E tests validating overlap/caching behavior, and comprehensive documentation.

Changes

Lightseek MM-aware KV Routing via Rust Frontend

Layer / File(s)	Summary
Build Configuration & Dependencies Cargo.toml, lib/bindings/python/Cargo.toml, lib/llm/Cargo.toml	Workspace adds moka 0.12 dependency. Python binding and LLM crates enable lightseek-mm feature by default. LLM crate adds llm-multimodal=1.5.0 and vendored openssl optional dependencies.
MM Data Structures & Image Token Resolution lib/llm/src/preprocessor.rs, lib/llm/src/preprocessor/image_token.rs, lib/llm/src/preprocessor/lightseek_mm.rs	Introduces MmImageEntry struct (mm_hash, width, height). Implements three-tier image placeholder token ID resolver (numeric config, tokenizer-mapped string, vocab probing). Adds LightseekMmCounter wrapper around llm-multimodal's ImageProcessorRegistry.
Preprocessor MM Routing Logic lib/llm/src/preprocessor.rs	Extends OpenAIPreprocessor with Lightseek/MM initialization. Refactors gather_tokens to return (token_ids, annotations). Implements mm_hash derivation (decoded-byte hash or URL fallback), per-image Lightseek token counting, and MM routing info building with image placeholder expansion and KV block padding. Adds URL-passthrough dimension fetching with bounded async cache and singleflight.
Media Storage Content Hashing lib/llm/src/preprocessor/media/rdma.rs	Adds content_hash() method to RdmaMediaDataDescriptor for xxh3-64 hashing of local registered bytes to support content-addressed routing.
Backend Error Response Formatting components/src/dynamo/vllm/handlers.py	Updates Qwen-VL decode worker error path to yield backend-output-shaped stream objects (finish_reason, index, empty token_ids) instead of status/message dicts when required prefill MM metadata is missing.
Launch Scripts & Orchestration examples/backends/vllm/launch/agg_multimodal_router.sh, examples/backends/vllm/launch/agg_multimodal_router_chat_processor.sh	Rust frontend variant launches single frontend with lightseek-mm KV routing, per-worker ZMQ KV event ports, and backend health checks. Chat-processor variant launches multiple frontend replicas running Python vLLM processor with KV router mode, per-worker multimodal config, and readiness polling.
Documentation & Example Binary lib/llm/examples/lightseek_count.rs, docs/features/multimodal/multimodal-kv-routing.md	Adds example binary demonstrating image token counting via Lightseek. Updates multimodal KV routing guide with Rust frontend routing flow (image (W,H) derivation, placeholder expansion, mm_hash forwarding), launch commands, environment variables, and frontend-decoding behavior.
End-to-End Tests & Configuration tests/mm_router/test_router_rust_mm_router_e2e.py, tests/mm_router/test_router_rust_mm_frontend_decode_e2e.py, tests/serve/multimodal_profiles/vllm.py	Implements E2E tests validating warm-vs-cold routing overlap, data-URI/HTTP URL equivalence, content-addressed hash collisions, query parameter normalization, and lightseek initialization logs. Adds frontend-decode variant testing decoded bytes source tag. Extends Qwen/Qwen3-VL-2B-Instruct profile with three new TopologyConfig entries (agg_router, agg_router_chat_processor, agg_router_frontend_decode).

🎯 4 (Complex) | ⏱️ ~45 minutes

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly and specifically describes the main feature being added: Rust-frontend MM-aware KV routing using the lightseek-mm cargo feature, which aligns with the substantial changes throughout the codebase.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Description check	✅ Passed	The PR description is comprehensive and well-structured, covering all required template sections with detailed technical context.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

Tip

💬 Introducing Slack Agent: The best way for teams to turn conversations into code.

Slack Agent is built on CodeRabbit's deep understanding of your code, so your team can collaborate across the entire SDLC without losing context.

• Generate code and open pull requests
• Plan features and break down work
• Investigate incidents and troubleshoot customer tickets together
• Automate recurring tasks and respond to alerts with triggers
• Summarize progress and report instantly

Built for teams:

• Shared memory across your entire org—no repeating context
• Per-thread sandboxes to safely plan and execute work
• Governance built-in—scoped access, auditability, and budget controls

One agent for your entire SDLC. Right inside Slack.

👉 Get started

---

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.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

Actionable comments posted: 11

🧹 Nitpick comments (2)

lib/llm/Cargo.toml (1)

87-87: ⚡ Quick win

Make moka an optional dependency and add it to the lightseek-mm feature.

moka is only used inside #[cfg(feature = "lightseek-mm")] blocks in preprocessor.rs (line 983), yet it's compiled unconditionally. The feature documentation promises a slimmer build with --no-default-features --features block-manager, but this build still incurs the cost of compiling moka and its tokio integration despite never using them.

♻️ Suggested fix

-moka = { workspace = true }
+moka = { workspace = true, optional = true }

-lightseek-mm = ["dep:llm-multimodal", "dep:openssl"]
+lightseek-mm = ["dep:llm-multimodal", "dep:openssl", "dep:moka"]

🤖 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/Cargo.toml` at line 87, The crate currently compiles moka
unconditionally; mark the moka dependency as optional and include it under the
lightseek-mm feature so it only builds when that feature is enabled. Update
Cargo.toml to make the moka dependency optional (optional = true) and add "moka"
to the feature list for "lightseek-mm" (so the feature enables moka), matching
the cfg usage in preprocessor.rs (the blocks gated by #[cfg(feature =
"lightseek-mm")]). Ensure the feature name exactly matches "lightseek-mm" so the
conditional compilation and feature flag behave consistently.

lib/llm/src/preprocessor/lightseek_mm.rs (1)

31-44: ⚡ Quick win

try_new performs blocking filesystem I/O and is called from async contexts.

std::fs::read_to_string blocks the calling thread. The function is invoked from OpenAIPreprocessor::new_with_parts, which is called during model initialization from async fn do_worker_set_registration in lib/llm/src/discovery/watcher.rs. This can stall a tokio runtime worker thread. Use tokio::fs::read_to_string for async I/O, or wrap the call with tokio::task::spawn_blocking.

🤖 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/lightseek_mm.rs` around lines 31 - 44, The try_new
function in lightseek_mm.rs does blocking disk I/O via std::fs::read_to_string
and is called from async contexts (e.g., OpenAIPreprocessor::new_with_parts ->
do_worker_set_registration); replace the blocking call with an async-aware
version: either call tokio::fs::read_to_string awaiting the result (and make
try_new async) or keep try_new sync but offload the read_to_string into
tokio::task::spawn_blocking and await that JoinHandle; also ensure the
subsequent PreProcessorConfig::from_json parsing is performed on the async path
or inside the spawn_blocking closure to avoid blocking the runtime thread.

🤖 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 `@examples/backends/vllm/launch/agg_multimodal_router_chat_processor.sh`:
- Around line 22-27: Replace the hand-rolled GPU/lifecycle logic in
agg_multimodal_router_chat_processor.sh (the manual handling of
_PROFILE_OVERRIDE_VLLM_KV_CACHE_BYTES, readiness polling, cleanup, and final
wait path) by sourcing the shared helpers gpu_utils.sh and launch_utils.sh, call
build_vllm_gpu_mem_args to construct vLLM memory args, use print_launch_banner
for consistent startup logs, and terminate the script with wait_any_exit so the
test/serve harness gets VRAM-safe parallel runs and proper fail-fast shutdown;
update all similar blocks (lines around 46-52, 114-149, 297) to the same pattern
and remove the duplicated readiness/cleanup/wait logic.

In `@examples/backends/vllm/launch/agg_multimodal_router.sh`:
- Around line 25-30: Replace the script's custom cleanup/wait logic by sourcing
the repo-standard launch_utils.sh near the top (using SCRIPT_DIR to locate it,
just like gpu_utils.sh), retain set -euo pipefail and the EXIT trap placement,
remove the hand-rolled trap/wait handling, and replace any bare wait calls with
wait_any_exit to ensure immediate failure detection and cleanup; apply the same
change to the other occurrences referenced (around the other ranges).
- Around line 127-148: The script is passing a manual --gpu-memory-utilization
flag that conflicts with the helper build_vllm_gpu_mem_args; remove the explicit
--gpu-memory-utilization "${GPU_MEMORY_UTILIZATION}" token from the python -m
dynamo.vllm invocation so the GPU_MEM_ARGS (returned by build_vllm_gpu_mem_args)
are the single source of truth for VRAM sizing; ensure GPU_MEM_ARGS remains
included in the command alongside VLLM_EXTRA_ARGS and PASSTHRU_ARGS.

In `@lib/llm/src/preprocessor.rs`:
- Around line 189-203: The struct OpenAIPreprocessor currently declares
kv_cache_block_size twice (once as an unconditional usize and again as a
feature-gated u32), causing a duplicate field error when lightseek-mm is
enabled; fix by consolidating into a single field named kv_cache_block_size with
a single type (prefer u32 to match the feature-gated declaration) and remove the
duplicate feature-gated declaration (or alternatively rename the feature-gated
field to kv_cache_block_size_mm if you need both), and update any usages of
OpenAIPreprocessor::kv_cache_block_size to the chosen type (and adjust casts)
or, if renaming, update all call sites that reference the old feature-gated
name.
- Around line 154-162: The type MmImageEntry is feature-gated but
gather_multi_modal_data() returns Vec<MmImageEntry> unconditionally, causing
compile failures when lightseek-mm is disabled; either remove the #[cfg(feature
= "lightseek-mm")] from the MmImageEntry definition so the struct is always
available, or gate the gather_multi_modal_data() signature and every call site
(e.g., calls at gather_mm_exact_routing_info, and the callers around lines
noted) with the same #[cfg(feature = "lightseek-mm")] so the function and its
uses are only compiled when the feature is enabled; update the signature and all
call sites (or the struct visibility) consistently to eliminate the
unconditional dependency.

In `@lib/llm/src/preprocessor/image_token.rs`:
- Around line 72-77: The current conditional in image_token.rs uses chained
.or(...) on ConfigJson fields (image_token_id, image_token_index,
media_placeholder_token_id) before validating >= 0, so a sentinel -1 on an
earlier field blocks valid later values; change the logic after
load_json::<ConfigJson> to explicitly inspect each field in order (inspect
image_token_id, if Some and >= 0 use it; otherwise inspect image_token_index, if
Some and >= 0 use it; otherwise inspect media_placeholder_token_id and accept it
only if >= 0) and only then proceed when you have a non-negative id; refer to
load_json, ConfigJson, image_token_id, image_token_index, and
media_placeholder_token_id to locate and update the check.

In `@tests/mm_router/test_router_rust_mm_frontend_decode_e2e.py`:
- Around line 55-63: The module-level pytestmark lacks a profiled_vram_gib
marker; update the pytestmark list (where VLLM_MM_MODEL and
pytest.mark.requested_vllm_kv_cache_bytes(...) are declared) to include
pytest.mark.profiled_vram_gib(<appropriate_gib_value>) so the parallel scheduler
can size runs correctly; place the new marker alongside the existing markers in
the pytestmark list for tests in
tests/mm_router/test_router_rust_mm_frontend_decode_e2e.py.
- Around line 91-94: _prepare_log_dir currently writes into a repo-relative path
and mm_runtime_services mutates os.environ directly; change _prepare_log_dir to
create and return a temp directory (use pytest's tmp_path or tempfile.mkdtemp
and create a subdir named with suffix) instead of using shutil.rmtree on a repo
path, and update callers to pass tmp_path (or make _prepare_log_dir accept
tmp_path) so logs are written under the temp dir; in mm_runtime_services replace
direct os.environ[...] = ... with pytest.MonkeyPatch (or the monkeypatch
fixture) calls like monkeypatch.setenv("NATS_SERVER", value) and
monkeypatch.setenv("ETCD_ENDPOINTS", value) so env changes are scoped and
cleaned up automatically (also apply the same tmp_path/monkeypatch pattern for
the other occurrence mentioned at lines ~175-182).

In `@tests/mm_router/test_router_rust_mm_router_e2e.py`:
- Around line 94-97: The _prepare_log_dir function and mm_runtime_services
fixture must avoid writing into the repo and mutating os.environ directly:
change _prepare_log_dir to create/return a temporary directory using pytest's
tmp_path or tempfile (e.g., use request and tmp_path to build a unique dir and
rmtree only within that temp), and update mm_runtime_services to set
NATS/ETCD-related env vars via pytest.MonkeyPatch (monkeypatch.setenv) instead
of assigning os.environ[...] directly; ensure teardown uses tmp_path cleanup and
that log paths point into the provided temp dir so tests are hermetic and leave
no repo state behind.
- Around line 53-61: The file-level pytestmark list is missing a profiled VRAM
marker; add a pytest.mark.profiled_vram_gib(<N>) entry to the pytestmark array
alongside pytest.mark.model(VLLM_MM_MODEL) and
pytest.mark.requested_vllm_kv_cache_bytes(...) so the scheduler can size the
test correctly; choose the <N> value equal to the profiled VRAM (in GiB) for the
VLLM_MM_MODEL used in this test and place the new
pytest.mark.profiled_vram_gib(...) element in the existing pytestmark list.

In `@tests/serve/multimodal_profiles/vllm.py`:
- Around line 108-141: The three TopologyConfig entries "agg_router",
"agg_router_chat_processor", and "agg_router_frontend_decode" are missing the
required profiled_vram_gib marker; run tests/utils/profile_pytest.py to measure
the VRAM (GiB) used by Qwen/Qwen3-VL-2B-Instruct for each topology and add
profiled_vram_gib=<measured_value> to each TopologyConfig (alongside the
existing requested_vllm_kv_cache_bytes and env where present) so the parallel
scheduler and --max-vram-gib filtering can size and select these tests
correctly.

---

Nitpick comments:
In `@lib/llm/Cargo.toml`:
- Line 87: The crate currently compiles moka unconditionally; mark the moka
dependency as optional and include it under the lightseek-mm feature so it only
builds when that feature is enabled. Update Cargo.toml to make the moka
dependency optional (optional = true) and add "moka" to the feature list for
"lightseek-mm" (so the feature enables moka), matching the cfg usage in
preprocessor.rs (the blocks gated by #[cfg(feature = "lightseek-mm")]). Ensure
the feature name exactly matches "lightseek-mm" so the conditional compilation
and feature flag behave consistently.

In `@lib/llm/src/preprocessor/lightseek_mm.rs`:
- Around line 31-44: The try_new function in lightseek_mm.rs does blocking disk
I/O via std::fs::read_to_string and is called from async contexts (e.g.,
OpenAIPreprocessor::new_with_parts -> do_worker_set_registration); replace the
blocking call with an async-aware version: either call tokio::fs::read_to_string
awaiting the result (and make try_new async) or keep try_new sync but offload
the read_to_string into tokio::task::spawn_blocking and await that JoinHandle;
also ensure the subsequent PreProcessorConfig::from_json parsing is performed on
the async path or inside the spawn_blocking closure to avoid blocking the
runtime thread.

🪄 Autofix (Beta)

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: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: 0f39e4ac-e4d5-411e-9d87-26b7b54ab460

📥 Commits

Reviewing files that changed from the base of the PR and between 4cbec37 and 0b0efd6.

⛔ Files ignored due to path filters (2)

• Cargo.lock is excluded by !**/*.lock
• lib/bindings/python/Cargo.lock is excluded by !**/*.lock

📒 Files selected for processing (15)

• Cargo.toml
• components/src/dynamo/vllm/handlers.py
• docs/features/multimodal/multimodal-kv-routing.md
• examples/backends/vllm/launch/agg_multimodal_router.sh
• examples/backends/vllm/launch/agg_multimodal_router_chat_processor.sh
• lib/bindings/python/Cargo.toml
• lib/llm/Cargo.toml
• lib/llm/examples/lightseek_count.rs
• lib/llm/src/preprocessor.rs
• lib/llm/src/preprocessor/image_token.rs
• lib/llm/src/preprocessor/lightseek_mm.rs
• lib/llm/src/preprocessor/media/rdma.rs
• tests/mm_router/test_router_rust_mm_frontend_decode_e2e.py
• tests/mm_router/test_router_rust_mm_router_e2e.py
• tests/serve/multimodal_profiles/vllm.py