[Bugfix]Fix multimodal cache routing for AR replicas (#3605)

Signed-off-by: bjf-frz <frz123db@gmail.com>

tests/engine/test_async_omni_engine_input.py

@@ -5,13 +5,14 @@ from vllm.v1.engine import EngineCoreRequest
 
 from vllm_omni.engine import OmniEngineCoreRequest
 from vllm_omni.engine.async_omni_engine import AsyncOmniEngine, StageRuntimeInfo
+from vllm_omni.engine.stage_pool import StagePool
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
-def _make_engine_core_request() -> EngineCoreRequest:
+def _make_engine_core_request(request_id: str = "req-1") -> EngineCoreRequest:
     return EngineCoreRequest(
-        request_id="req-1",
+        request_id=request_id,
         prompt_token_ids=[1, 1, 1],
         mm_features=None,
         sampling_params=SamplingParams(max_tokens=8),
@@ -88,3 +89,86 @@ def test_build_add_request_message_with_resumable_streaming(mocker: MockerFixtur
     assert msg.type == "streaming_update"
     input_processor.process_inputs.assert_called_once()
     assert input_processor.process_inputs.call_args.kwargs["resumable"] is True
+
+
+class _FakeStageClient:
+    stage_type = "llm"
+    final_output = False
+
+
+def test_build_add_request_message_scopes_mm_uuids_to_selected_stage0_replica(mocker: MockerFixture):
+    engine = object.__new__(AsyncOmniEngine)
+    params = SamplingParams(max_tokens=8)
+    engine.model = "test-model"
+    engine.default_sampling_params_list = [params]
+    engine.stage_metadata = [StageRuntimeInfo(final_output=False, final_output_type=None, stage_type="llm")]
+    engine.supported_tasks = ("generate",)
+    engine.stage_pools = [StagePool(0, [_FakeStageClient(), _FakeStageClient()])]
+
+    seen_uuids: list[str] = []
+
+    def process_inputs(**kwargs):
+        prompt = kwargs["prompt"]
+        seen_uuids.append(prompt["multi_modal_uuids"]["image"][0])
+        return _make_engine_core_request(kwargs["request_id"])
+
+    input_processor = mocker.Mock()
+    input_processor.process_inputs.side_effect = process_inputs
+    engine.input_processor = input_processor
+
+    for request_id in ("req-1", "req-2"):
+        engine._build_add_request_message(
+            request_id=request_id,
+            prompt={
+                "prompt": "describe",
+                "multi_modal_data": {"image": "same-image"},
+            },
+            sampling_params_list=[params],
+            final_stage_id=0,
+        )
+
+    assert seen_uuids[0].startswith("stage0:rep0:")
+    assert seen_uuids[1].startswith("stage0:rep1:")
+    assert seen_uuids[0].removeprefix("stage0:rep0:") == seen_uuids[1].removeprefix("stage0:rep1:")
+
+
+def test_stage_pool_replica_count_falls_back_to_clients():
+    class PoolWithoutLiveNumReplicas:
+        clients = [object(), None, object()]
+
+    assert AsyncOmniEngine._stage_pool_replica_count(PoolWithoutLiveNumReplicas()) == 2
+
+
+def test_stage_pool_is_distributed_falls_back_to_hub():
+    class PoolWithoutIsDistributed:
+        _hub = object()
+
+    assert AsyncOmniEngine._stage_pool_is_distributed(PoolWithoutIsDistributed()) is True
+
+
+def test_build_add_request_message_releases_preselected_replica_on_preprocess_error(mocker: MockerFixture):
+    engine = object.__new__(AsyncOmniEngine)
+    params = SamplingParams(max_tokens=8)
+    engine.model = "test-model"
+    engine.default_sampling_params_list = [params]
+    engine.stage_metadata = [StageRuntimeInfo(final_output=False, final_output_type=None, stage_type="llm")]
+    engine.supported_tasks = ("generate",)
+    stage_pool = StagePool(0, [_FakeStageClient(), _FakeStageClient()])
+    engine.stage_pools = [stage_pool]
+
+    input_processor = mocker.Mock()
+    input_processor.process_inputs.side_effect = RuntimeError("boom")
+    engine.input_processor = input_processor
+
+    with pytest.raises(RuntimeError, match="boom"):
+        engine._build_add_request_message(
+            request_id="req-error",
+            prompt={
+                "prompt": "describe",
+                "multi_modal_data": {"image": "same-image"},
+            },
+            sampling_params_list=[params],
+            final_stage_id=0,
+        )
+
+    assert stage_pool.get_bound_replica_id("req-error") is None

vllm_omni/engine/async_omni_engine.py

@@ -1199,6 +1199,158 @@ class AsyncOmniEngine:
 
     # ---- request helpers ----
 
+    @staticmethod
+    def _iter_multimodal_items(value: Any) -> list[Any]:
+        if value is None:
+            return []
+        if isinstance(value, list):
+            return value
+        return [value]
+
+    def _ensure_stage_replica_mm_uuids(
+        self,
+        prompt: Any,
+        *,
+        stage_id: int,
+        replica_id: int,
+    ) -> None:
+        """Make multimodal processor-cache keys local to a stage replica.
+
+        vLLM's frontend multimodal sender cache is process-global, while each
+        vllm-omni stage replica owns a separate EngineCore receiver cache. If
+        two requests with the same image are routed to different stage-0
+        replicas, a plain content hash can make the sender omit the tensor for
+        a replica that has never received it. Prefixing user/content UUIDs with
+        the selected replica keeps cache reuse within the receiver that owns it.
+        """
+
+        if not isinstance(prompt, dict):
+            return
+
+        mm_data = prompt.get("multi_modal_data")
+        if not isinstance(mm_data, dict) or not mm_data:
+            return
+
+        from vllm.multimodal.hasher import MultiModalHasher
+
+        existing_uuids = prompt.get("multi_modal_uuids")
+        if not isinstance(existing_uuids, dict):
+            existing_uuids = {}
+
+        model_id = str(getattr(self, "model", ""))
+        scoped_uuids: dict[str, list[str | None]] = dict(existing_uuids)
+        for modality, raw_items in mm_data.items():
+            items = self._iter_multimodal_items(raw_items)
+            if not items:
+                continue
+
+            modality_existing = existing_uuids.get(modality)
+            if not isinstance(modality_existing, list):
+                modality_existing = [modality_existing] if modality_existing is not None else []
+
+            modality_uuids: list[str | None] = []
+            for idx, item in enumerate(items):
+                user_uuid = modality_existing[idx] if idx < len(modality_existing) else None
+                if user_uuid is not None:
+                    base_uuid = str(user_uuid)
+                elif item is None:
+                    base_uuid = None
+                else:
+                    base_uuid = MultiModalHasher.hash_kwargs(
+                        model_id=model_id,
+                        **{modality: item},
+                    )
+
+                if base_uuid is None:
+                    modality_uuids.append(None)
+                else:
+                    modality_uuids.append(f"stage{stage_id}:rep{replica_id}:{base_uuid}")
+
+            scoped_uuids[modality] = modality_uuids
+
+        if scoped_uuids:
+            prompt["multi_modal_uuids"] = scoped_uuids
+
+    @staticmethod
+    def _stage_pool_replica_count(stage_pool: Any) -> int:
+        try:
+            live_num_replicas = getattr(stage_pool, "live_num_replicas", None)
+            if live_num_replicas is not None:
+                return int(live_num_replicas)
+        except Exception:
+            pass
+
+        try:
+            live_replica_ids = getattr(stage_pool, "live_replica_ids", None)
+            if callable(live_replica_ids):
+                return len(live_replica_ids())
+        except Exception:
+            pass
+
+        try:
+            clients = getattr(stage_pool, "clients", None)
+            if clients is not None:
+                return sum(1 for client in clients if client is not None)
+        except Exception:
+            pass
+
+        return int(getattr(stage_pool, "num_replicas", 1) or 1)
+
+    @staticmethod
+    def _stage_pool_is_distributed(stage_pool: Any) -> bool:
+        try:
+            is_distributed = getattr(stage_pool, "is_distributed", None)
+            if is_distributed is not None:
+                return bool(is_distributed() if callable(is_distributed) else is_distributed)
+        except Exception:
+            pass
+
+        return getattr(stage_pool, "_hub", None) is not None
+
+    def _scope_stage0_multimodal_cache_to_replica(
+        self,
+        request_id: str,
+        prompt: Any,
+    ) -> int | None:
+        stage_pools = getattr(self, "stage_pools", None)
+        if isinstance(prompt, EngineCoreRequest) or not stage_pools:
+            return None
+
+        stage0_pool = stage_pools[0]
+        # TODO: Currently only supports the ar -> dit process.
+        # Future scenarios (e.g., dit -> ar) need to be added, which will require modifications here.
+        if stage0_pool.stage_type == "diffusion" or self._stage_pool_replica_count(stage0_pool) <= 1:
+            return None
+
+        # This synchronous request path can safely pre-bind local in-process
+        # replicas. Distributed head mode still uses the async picker inside
+        # StagePool.submit_initial().
+        if self._stage_pool_is_distributed(stage0_pool):
+            logger.debug(
+                "[AsyncOmniEngine] Skipping stage-0 multimodal cache scoping for distributed routing req=%s",
+                request_id,
+            )
+            return None
+
+        prompts = prompt if isinstance(prompt, list) else [prompt]
+        if not any(isinstance(p, dict) and p.get("multi_modal_data") for p in prompts):
+            return None
+
+        replica_id = stage0_pool.select_replica_id(request_id)
+        for p in prompts:
+            self._ensure_stage_replica_mm_uuids(
+                p,
+                stage_id=0,
+                replica_id=replica_id,
+            )
+
+        logger.debug(
+            "[AsyncOmniEngine] Scoped multimodal cache keys to stage-0 replica-%s for req=%s",
+            replica_id,
+            request_id,
+        )
+        return replica_id
+
     def _build_add_request_message(
         self,
         request_id: str,
@@ -1232,6 +1384,7 @@ class AsyncOmniEngine:
         # Keep the original prompt for downstream stages (they need the raw
         # dict, e.g. for multi_modal_data).
         original_prompt = prompt
+        preselected_stage0_replica: int | None = None
 
         stage_type = self.stage_metadata[0].stage_type
         output_prompt_text: Any = None
@@ -1244,21 +1397,31 @@ class AsyncOmniEngine:
                 for item in prompt:
                     _inject_global_id(item, request_id)
 
+            preselected_stage0_replica = self._scope_stage0_multimodal_cache_to_replica(
+                request_id,
+                prompt,
+            )
+
             # Full input processing (tokenization, multimodal, etc.)
             _t_preprocess = time.perf_counter()
-            request = self.input_processor.process_inputs(
-                request_id=request_id,
-                prompt=prompt,
-                params=params,
-                supported_tasks=self.supported_tasks,
-                arrival_time=arrival_time,
-                lora_request=lora_request,
-                tokenization_kwargs=tokenization_kwargs,
-                trace_headers=trace_headers,
-                priority=priority,
-                data_parallel_rank=data_parallel_rank,
-                resumable=resumable,
-            )
+            try:
+                request = self.input_processor.process_inputs(
+                    request_id=request_id,
+                    prompt=prompt,
+                    params=params,
+                    supported_tasks=self.supported_tasks,
+                    arrival_time=arrival_time,
+                    lora_request=lora_request,
+                    tokenization_kwargs=tokenization_kwargs,
+                    trace_headers=trace_headers,
+                    priority=priority,
+                    data_parallel_rank=data_parallel_rank,
+                    resumable=resumable,
+                )
+            except Exception:
+                if preselected_stage0_replica is not None and self.stage_pools:
+                    self.stage_pools[0].release_binding(request_id)
+                raise
             _preprocess_ms = (time.perf_counter() - _t_preprocess) * 1000.0
             # TODO (Peiqi): add this for Qwen3-TTS only. Other models don't have
             # additional_information field in the prompt.