Skip to content

fix: support approx routing in mm router #8135

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
zhongdaor-nv wants to merge 4 commits into
base: main
Choose a base branch
from
Open

fix: support approx routing in mm router #8135

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
8e4f2b7
support approx routing in mm router
zhongdaor-nv Apr 13, 2026
d9bfffc
pre-commit
zhongdaor-nv Apr 13, 2026
e96984b
Merge branch 'main' into zhongdaor/support_approx_routing_mm_router
zhongdaor-nv Apr 13, 2026
946d54d
Merge branch 'main' into zhongdaor/support_approx_routing_mm_router
zhongdaor-nv Apr 22, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
25 changes: 20 additions & 5 deletions examples/backends/vllm/mm_router_worker/mm_router_worker.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,7 @@
Usage:
python -m examples.backends.vllm.mm_router_worker \
--model Qwen/Qwen3-VL-8B-Instruct \
--namespace default \
--namespace dynamo \
--component mm_router \
--endpoint generate \
--downstream-component backend \
Expand All @@ -21,6 +21,7 @@
import argparse
import asyncio
import logging
import os
import signal

import uvloop
Expand Down Expand Up @@ -58,8 +59,8 @@ def parse_args() -> argparse.Namespace:
parser.add_argument(
"--namespace",
type=str,
default="default",
help="Dynamo namespace",
default=os.environ.get("DYN_NAMESPACE", "dynamo"),
help="Dynamo namespace (default: DYN_NAMESPACE env or 'dynamo')",
)
parser.add_argument(
"--component",
Expand Down Expand Up @@ -88,6 +89,15 @@ def parse_args() -> argparse.Namespace:
help="Downstream vLLM workers' endpoint name",
)

# Router configuration
parser.add_argument(
"--no-router-kv-events",
action="store_true",
default=False,
help="Use approximate KV routing (no KV events from workers). "
"Required for hybrid models like Qwen3.5 that cannot emit KV events.",
)

return parser.parse_args()


Expand Down Expand Up @@ -135,12 +145,17 @@ def signal_handler():
logger.info(f"Found {len(instance_ids)} workers: {list(instance_ids)}")

# Create KvRouter to select workers based on KV overlap
kv_router_config = KvRouterConfig(
use_kv_events=not args.no_router_kv_events,
)
kv_router = KvRouter(
endpoint=downstream_endpoint,
block_size=args.block_size,
kv_router_config=KvRouterConfig(),
kv_router_config=kv_router_config,
)
logger.info(
f"KvRouter created successfully (use_kv_events={not args.no_router_kv_events})"
)
logger.info("KvRouter created successfully")

# Initialize tokenizer and processor for MM processing
logger.info(f"Loading tokenizer from {args.model}...")
Expand Down
124 changes: 124 additions & 0 deletions examples/backends/vllm/qwen35/launch.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,124 @@
#!/bin/bash
# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Aggregated multimodal serving for Qwen3.5 hybrid models with MM-aware approximate KV routing.
#
# Qwen3.5 is a multimodal hybrid model (GatedDeltaNet + Gated Attention + Vision Encoder).
# It supports images and video but has hybrid architecture constraints:
#
# 1. DO NOT pass --kv-events-config or --enable-kv-cache-events:
# vLLM disables the Hybrid KV Cache Manager when kv_events_config is set,
# but Qwen3.5's mixed KV cache specs (GDN + FullAttention) cannot be unified
# into one type.
#
# 2. Use --mamba-cache-mode align (not "all"):
# Qwen3.5 raises NotImplementedError with mamba_cache_mode="all".
#
# 3. Approximate KV routing (--no-router-kv-events):
# Hybrid models cannot emit KV events to the router. The MM Router Worker
# predicts cache state from its own routing decisions using prefix hashing.
#
# 4. Disaggregated P/D is NOT supported for hybrid models in vLLM:
# HybridKVCacheCoordinator asserts dcp_world_size == 1.
#
# 5. Use TCP transport for multimodal payloads (NATS has 1MB limit).
#
# Architecture:
# Frontend (--router-mode round-robin)
# -> MM Router Worker (approximate KV routing + multimodal hash)
# -> vLLM Worker (--enable-multimodal --mamba-cache-mode align)

set -e
trap 'echo Cleaning up...; kill 0' EXIT

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
source "$SCRIPT_DIR/../../../common/gpu_utils.sh"
source "$SCRIPT_DIR/../../../common/launch_utils.sh"

MODEL="${MODEL:-Qwen/Qwen3.5-0.8B}"

EXTRA_ARGS=()
while [[ $# -gt 0 ]]; do
case $1 in
--model) MODEL="$2"; shift 2 ;;
*) EXTRA_ARGS+=("$1"); shift ;;
esac
done

MAX_MODEL_LEN="${MAX_MODEL_LEN:-4096}"
MAX_CONCURRENT_SEQS="${MAX_CONCURRENT_SEQS:-2}"
HTTP_PORT="${DYN_HTTP_PORT:-8000}"

@coderabbitai coderabbitai Bot Apr 13, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major

Pass the resolved frontend port into dynamo.frontend.

HTTP_PORT drives the banner and curl examples, but the frontend command never consumes it. That makes the documented port override path rely on implicit frontend behavior and can leave the printed curl command pointing at the wrong socket.

🔧 Proposed fix 
 python -m dynamo.frontend \
+    --http-port "$HTTP_PORT" \
     --router-mode round-robin &

Also applies to: 121-122

🤖 Prompt for AI Agents 
Verify each finding against the current code and only fix it if needed.

In `@examples/backends/vllm/qwen35/launch.sh` at line 51, The script computes
HTTP_PORT but never passes it into the frontend invocation, so the banner/curl
examples can point at the wrong socket; update each call to dynamo.frontend (the
frontend command invocation seen near the HTTP_PORT assignment around line 51
and the other invocation around lines 121-122) to accept the resolved port by
supplying "$HTTP_PORT" (or the appropriate --port "$HTTP_PORT" flag) as an
argument or environment variable so the actual frontend listens on the same port
printed in banners and curl examples.

BLOCK_SIZE="${BLOCK_SIZE:-16}"

# TCP transport: avoids NATS 1MB payload limit for base64-encoded images
export DYN_REQUEST_PLANE=tcp

print_launch_banner --no-curl "Launching Qwen3.5 Multimodal + MM Router + Approx KV (1 GPU)" "$MODEL" "$HTTP_PORT" \
"Backend: dynamo.vllm --enable-multimodal --mamba-cache-mode align" \
"MM Router: MM-aware approximate KV routing (--no-router-kv-events)" \
"Frontend: round-robin to MM Router" \
"Transport: TCP (multimodal payloads)"

print_curl_footer <<CURL
# Text-only request
curl http://localhost:${HTTP_PORT}/v1/chat/completions \\
-H 'Content-Type: application/json' \\
-d '{
"model": "${MODEL}",
"messages": [{"role": "user", "content": "What is 2+2?"}],
"max_tokens": 32
}'

# Multimodal request (image)
curl http://localhost:${HTTP_PORT}/v1/chat/completions \\
-H 'Content-Type: application/json' \\
-d '{
"model": "${MODEL}",
"messages": [{"role": "user", "content": [
{"type": "text", "text": "Describe the image"},
{"type": "image_url", "image_url": {"url": "https://upload.wikimedia.org/wikipedia/commons/thumb/4/47/PNG_transparency_demonstration_1.png/300px-PNG_transparency_demonstration_1.png"}}
]}],
"max_tokens": 50
}'
CURL

GPU_MEM_ARGS=$(build_vllm_gpu_mem_args)

# vLLM worker: hybrid multimodal model
# --served-model-name __internal: hides from frontend so traffic goes through MM Router
# --mamba-cache-mode align: required for GDN+Attention hybrid architecture
# --enable-multimodal: enables vision encoder / multimodal data handling
# NOTE: do NOT use --is-decode-worker here. It causes the handler to enter
# decode-only mode which silently drops image data for models not in
# QWEN_VL_MODELS (Qwen3.5 is not listed), leading to incorrect prefix
# cache hits across different images.
CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES:-0} \
DYN_SYSTEM_PORT=${DYN_SYSTEM_PORT:-8081} \
python -m dynamo.vllm \
--model "$MODEL" \
--served-model-name "${MODEL}__internal" \
--enable-multimodal \
--mamba-cache-mode align \
--block-size "$BLOCK_SIZE" \
--enforce-eager \

@furionw furionw Apr 14, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

should we leave this out?
for perf reasons, we should not enforce eager by default.
people can pass this as extra_args

--max-model-len "$MAX_MODEL_LEN" \
--max-num-seqs "$MAX_CONCURRENT_SEQS" \
$GPU_MEM_ARGS \
"${EXTRA_ARGS[@]}" &

# MM Router Worker: multimodal-aware approximate KV routing
# --namespace dynamo: must match frontend/vllm default (DYN_NAMESPACE defaults to "dynamo")
# --no-router-kv-events: Qwen3.5 hybrid model cannot emit KV events
DYN_SYSTEM_PORT=${MM_ROUTER_SYSTEM_PORT:-8082} \
python -m examples.backends.vllm.mm_router_worker \
--model "$MODEL" \
--namespace dynamo \
--block-size "$BLOCK_SIZE" \
--no-router-kv-events &

# Frontend: round-robin dispatch to MM Router (KV routing happens inside MM Router)
python -m dynamo.frontend \
--router-mode round-robin &

wait_any_exit
107 changes: 62 additions & 45 deletions tests/mm_router/test_vllm_mm_router_e2e.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -99,30 +99,34 @@ def _prepare_log_dir(request, suffix: str) -> str:


class VLLMWorkerProcess(ManagedProcess):
"""vLLM backend worker that emits KV events."""

def __init__(self, request, *, system_port: int, kv_event_port: int):
super().__init__(
command=[
"python3",
"-m",
"dynamo.vllm",
"--model",
VLLM_MM_MODEL,
"--enable-multimodal",
"--gpu-memory-utilization",
"0.85",
"--max-model-len",
"8192",
"--served-model-name",
f"{VLLM_MM_MODEL}__internal",
"""vLLM backend worker that optionally emits KV events."""

def __init__(self, request, *, system_port: int, kv_event_port: int | None = None):
cmd = [
"python3",
"-m",
"dynamo.vllm",
"--model",
VLLM_MM_MODEL,
"--enable-multimodal",
"--gpu-memory-utilization",
"0.85",
"--max-model-len",
"8192",
"--served-model-name",
f"{VLLM_MM_MODEL}__internal",
]
if kv_event_port is not None:
cmd += [
"--kv-events-config",
(
f'{{"publisher":"zmq","topic":"kv-events",'
f'"endpoint":"tcp://*:{kv_event_port}",'
f'"enable_kv_cache_events": true}}'
),
],
]
super().__init__(
command=cmd,
env=_make_process_env(DYN_SYSTEM_PORT=str(system_port)),
health_check_urls=[
(f"http://localhost:{system_port}/health", _check_ready)
Expand All @@ -135,29 +139,32 @@ def __init__(self, request, *, system_port: int, kv_event_port: int):


class VLLMMMRouterWorkerProcess(ManagedProcess):
"""vLLM MM router worker."""

def __init__(self, request, *, system_port: int):
"""vLLM MM router worker (exact or approximate KV routing)."""

def __init__(self, request, *, system_port: int, approx_routing: bool = False):
cmd = [
"python3",
"-m",
"examples.backends.vllm.mm_router_worker",
"--model",
VLLM_MM_MODEL,
"--namespace",
NAMESPACE,
"--component",
"mm_router",
"--endpoint",
"generate",
"--downstream-component",
"backend",
"--downstream-endpoint",
"generate",
"--block-size",
str(BLOCK_SIZE),
]
if approx_routing:
cmd.append("--no-router-kv-events")
super().__init__(
command=[
"python3",
"-m",
"examples.backends.vllm.mm_router_worker",
"--model",
VLLM_MM_MODEL,
"--namespace",
NAMESPACE,
"--component",
"mm_router",
"--endpoint",
"generate",
"--downstream-component",
"backend",
"--downstream-endpoint",
"generate",
"--block-size",
str(BLOCK_SIZE),
],
command=cmd,
env=_make_process_env(
DYN_SYSTEM_USE_ENDPOINT_HEALTH_STATUS='["generate"]',
DYN_SYSTEM_PORT=str(system_port),
Expand Down Expand Up @@ -207,17 +214,27 @@ def mm_runtime_services(request):
os.environ.pop("ETCD_ENDPOINTS", None)


@pytest.fixture(scope="module")
@pytest.fixture(scope="module", params=[False, True], ids=["exact_kv", "approx_kv"])
def start_vllm_mm_services(
request, mm_runtime_services
) -> Generator[tuple[int, ManagedProcess], None, None]:
frontend_port, vllm_port, router_port, kv_event_port = allocate_ports(
count=4, start_port=10000
)
approx_routing = request.param

if approx_routing:
frontend_port, vllm_port, router_port = allocate_ports(
count=3, start_port=10000
)
kv_event_port = None
else:
frontend_port, vllm_port, router_port, kv_event_port = allocate_ports(
count=4, start_port=10000
)
Comment on lines +217 to +231

@coderabbitai coderabbitai Bot Apr 13, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major

Use the repo’s dynamic service/port fixtures here instead of extending the hand-rolled setup.

Adding the exact/approx split on top of manual allocate_ports(..., start_port=10000) keeps this module outside the xdist-safe path the rest of the suite uses. Please request the needed system ports via num_system_ports and consume runtime_services_dynamic_ports + dynamo_dynamic_ports instead of branching port allocation manually here.

As per coding guidelines, "Use runtime_services_dynamic_ports and dynamo_dynamic_ports fixtures together for xdist/parallel safety" and "Use num_system_ports parametrize ... to request multiple system ports."

🤖 Prompt for AI Agents 
Verify each finding against the current code and only fix it if needed.

In `@tests/mm_router/test_vllm_mm_router_e2e.py` around lines 216 - 230, The
start_vllm_mm_services fixture currently calls allocate_ports(...) manually and
branches on request.param; replace that manual port allocation by adding
num_system_ports and the runtime fixtures (runtime_services_dynamic_ports and
dynamo_dynamic_ports) to the fixture signature and consume the dynamic port
lists they provide; when approx_routing is True pull three ports from
runtime_services_dynamic_ports, otherwise pull four ports (including an extra
kv_event port) by combining runtime_services_dynamic_ports and
dynamo_dynamic_ports as needed, removing allocate_ports and start_port usage and
keeping the rest of the fixture logic unchanged so xdist-safe dynamic port
allocation is used.


with VLLMWorkerProcess(request, system_port=vllm_port, kv_event_port=kv_event_port):
time.sleep(10)
with VLLMMMRouterWorkerProcess(request, system_port=router_port) as router_proc:
with VLLMMMRouterWorkerProcess(
request, system_port=router_port, approx_routing=approx_routing
) as router_proc:
time.sleep(3)
with FrontendProcess(request, frontend_port=frontend_port):
yield frontend_port, router_proc
Expand Down
Loading