diff --git a/tests/core/sched/test_omni_ar_scheduler_kv_transfer.py b/tests/core/sched/test_omni_ar_scheduler_kv_transfer.py new file mode 100644 index 000000000..cffe3f289 --- /dev/null +++ b/tests/core/sched/test_omni_ar_scheduler_kv_transfer.py @@ -0,0 +1,57 @@ +from types import SimpleNamespace + +import pytest +from vllm import SamplingParams +from vllm.v1.engine import EngineCoreRequest + +from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler +from vllm_omni.engine.async_engine_utils import apply_omni_final_stage_metadata +from vllm_omni.engine.serialization import deserialize_additional_information + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _make_engine_request() -> EngineCoreRequest: + return EngineCoreRequest( + request_id="req", + prompt_token_ids=[1], + mm_features=None, + sampling_params=SamplingParams(max_tokens=1), + pooling_params=None, + arrival_time=0.0, + lora_request=None, + cache_salt=None, + data_parallel_rank=None, + ) + + +def _request_omits_kv_transfer(*, force_kv_transfer: bool) -> tuple[bool, dict]: + tagged = apply_omni_final_stage_metadata( + _make_engine_request(), + final_stage_id=0, + force_kv_transfer=force_kv_transfer, + ) + scheduler = OmniARScheduler.__new__(OmniARScheduler) + scheduler._omits_kv_transfer_cache = {} + request = SimpleNamespace( + request_id="req", + additional_information=tagged.additional_information, + ) + result = scheduler._request_omits_kv_transfer_to_next_stage(request) + metadata = deserialize_additional_information(tagged.additional_information) + return result, metadata + + +def test_stage_zero_request_omits_kv_transfer(): + omits_transfer, metadata = _request_omits_kv_transfer(force_kv_transfer=False) + + assert omits_transfer + assert "omni_force_kv_transfer" not in metadata + + +def test_cfg_companion_forces_kv_transfer_without_downstream_payload(): + omits_transfer, metadata = _request_omits_kv_transfer(force_kv_transfer=True) + + assert not omits_transfer + assert metadata["omni_final_stage_id"] == 0 + assert metadata["omni_force_kv_transfer"] is True diff --git a/tests/engine/test_async_omni_engine_input.py b/tests/engine/test_async_omni_engine_input.py index 91fdc752d..740064173 100644 --- a/tests/engine/test_async_omni_engine_input.py +++ b/tests/engine/test_async_omni_engine_input.py @@ -7,7 +7,9 @@ from vllm.v1.engine import EngineCoreRequest from vllm_omni.distributed.omni_coordinator import ReplicaInfo, ReplicaStatus from vllm_omni.engine import OmniEngineCoreRequest from vllm_omni.engine.async_omni_engine import AsyncOmniEngine, StageRuntimeInfo +from vllm_omni.engine.serialization import deserialize_additional_information from vllm_omni.engine.stage_pool import StagePool +from vllm_omni.model_executor.stage_input_processors.bagel import ExpandedPrompt pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -107,6 +109,34 @@ def test_build_add_request_message_preserves_model_intermediate_buffer(mocker: M assert torch.equal(info["hidden_states"]["tts"], hidden) +def test_cfg_companion_suppresses_payload_but_forces_kv_transfer(mocker: MockerFixture): + engine = object.__new__(AsyncOmniEngine) + params = SamplingParams(max_tokens=8) + engine.prompt_expand_func = lambda *_args: [ + ExpandedPrompt( + prompt={"prompt": "negative"}, + role="cfg_text", + request_id_suffix="__cfg_text", + ) + ] + engine.supported_tasks = ("generate",) + engine.input_processor = mocker.Mock() + engine.input_processor.process_inputs.return_value = _make_engine_core_request("req__cfg_text") + engine.request_queue = mocker.Mock() + + engine._enqueue_cfg_companions( + parent_id="req", + original_prompt={"prompt": "positive"}, + stage0_params=params, + sampling_params_list=[params], + ) + + message = engine.request_queue.sync_q.put.call_args.args[0] + metadata = deserialize_additional_information(message.prompt.additional_information) + assert metadata["omni_final_stage_id"] == 0 + assert metadata["omni_force_kv_transfer"] is True + + def test_build_add_request_message_with_resumable_streaming(mocker: MockerFixture): engine = object.__new__(AsyncOmniEngine) params = SamplingParams(max_tokens=8) diff --git a/tests/worker/test_omni_connector_mixin.py b/tests/worker/test_omni_connector_mixin.py index cf4e09cb9..7a2a0053b 100644 --- a/tests/worker/test_omni_connector_mixin.py +++ b/tests/worker/test_omni_connector_mixin.py @@ -17,6 +17,9 @@ from unittest.mock import MagicMock, patch import pytest import torch +from vllm_omni.distributed.omni_connectors.kv_transfer_manager import ( + OmniKVTransferManager, +) from vllm_omni.outputs import OmniConnectorOutput from vllm_omni.worker.omni_connector_model_runner_mixin import ( OmniConnectorModelRunnerMixin, @@ -1198,6 +1201,38 @@ class TestRankAwareKVRouting(unittest.TestCase): ) host.shutdown_omni_connectors() + def test_single_rank_hooks_use_default_manager_key(self): + host = self._make_host(from_tp=1, to_tp=1, local_rank=0) + expected = ["omni_0_to_1_kv_cache_req"] + + self.assertEqual(host.get_rank_aware_kv_send_keys("req", from_stage=0, to_stage=1), expected) + self.assertEqual(host.get_rank_aware_kv_keys("req", from_stage=0, to_stage=1), expected) + host.shutdown_omni_connectors() + + def test_init_uses_effective_manager_tp2_topology(self): + model_config = _make_model_config(stage_id=0) + model_config.omni_kv_config = { + "need_send_cache": True, + "omni_from_stage": "0", + "omni_to_stage": "1", + "rank_mapping": {"from_tp": 2, "to_tp": 2}, + } + with patch( + "vllm_omni.distributed.omni_connectors.kv_transfer_manager.get_local_tp_rank", + return_value=1, + ): + manager = OmniKVTransferManager.from_model_config(model_config) + + host = MixinHost() + host.init_omni_connectors(model_config=model_config, kv_transfer_manager=manager) + + expected = ["req_0_0_1_1"] + self.assertEqual(host.get_rank_aware_kv_send_keys("req", from_stage=0, to_stage=1), expected) + self.assertEqual(host.get_rank_aware_kv_keys("req", from_stage=0, to_stage=1), expected) + self.assertEqual(host.get_kv_rank_mapping()["from_tp"], 2) + self.assertEqual(host.get_kv_rank_mapping()["to_tp"], 2) + host.shutdown_omni_connectors() + def test_send_keys_route_from_rank_gt_to_rank(self): host = self._make_host(from_tp=4, to_tp=2, local_rank=3) self.assertEqual(host.get_rank_aware_kv_send_keys("req", from_stage=0), ["req_0_0_3_1"]) diff --git a/vllm_omni/core/sched/omni_ar_scheduler.py b/vllm_omni/core/sched/omni_ar_scheduler.py index b8a2fbd9b..514cc9b94 100644 --- a/vllm_omni/core/sched/omni_ar_scheduler.py +++ b/vllm_omni/core/sched/omni_ar_scheduler.py @@ -147,7 +147,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): return None def _request_omits_kv_transfer_to_next_stage(self, request: Request) -> bool: - """True when orchestrator will not run stage 1+ for this request (e.g. text-only). + """True when this stage-zero-final request does not need downstream KV. The result is cached per request to avoid repeated deserialization of additional_information on every scheduler tick. @@ -162,7 +162,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): result = False else: info = deserialize_additional_information(payload) - result = info.get("omni_final_stage_id") == 0 + result = info.get("omni_final_stage_id") == 0 and not bool(info.get("omni_force_kv_transfer", False)) self._omits_kv_transfer_cache[rid] = result return result diff --git a/vllm_omni/distributed/omni_connectors/kv_transfer_manager.py b/vllm_omni/distributed/omni_connectors/kv_transfer_manager.py index efed75241..8616cf9a6 100644 --- a/vllm_omni/distributed/omni_connectors/kv_transfer_manager.py +++ b/vllm_omni/distributed/omni_connectors/kv_transfer_manager.py @@ -403,6 +403,7 @@ class OmniKVTransferManager: self._prefetch_executor: ThreadPoolExecutor | None = ( ThreadPoolExecutor(max_workers=1, thread_name_prefix="kv-prefetch") if async_prefetch else None ) + self._prefetch_futures: dict[str, Any] = {} self._bg_copy_stream: current_omni_platform.Stream | None = None @@ -415,6 +416,11 @@ class OmniKVTransferManager: except Exception as e: logger.warning("Failed to eagerly initialize sender connector: %s", e) + @property + def tp_topology(self) -> KVTPTopology: + """Return the effective topology after config and runtime detection.""" + return self._tp_topo + @property def topo_config(self) -> _TransferTopoConfig: """Lazily build the parallel topology snapshot on first access.""" diff --git a/vllm_omni/engine/async_engine_utils.py b/vllm_omni/engine/async_engine_utils.py index 58f8351b0..6f544239b 100644 --- a/vllm_omni/engine/async_engine_utils.py +++ b/vllm_omni/engine/async_engine_utils.py @@ -77,12 +77,17 @@ def upgrade_to_omni_request( def apply_omni_final_stage_metadata( request: EngineCoreRequest, final_stage_id: int, + *, + force_kv_transfer: bool = False, ) -> EngineCoreRequest: - """Tag a request so the AR scheduler can skip DiT KV for stage zero.""" + """Tag a request with its final stage and optional KV-transfer override.""" merged: dict[str, Any] = {} if isinstance(request, OmniEngineCoreRequest) and request.additional_information is not None: merged = deserialize_additional_information(request.additional_information) merged["omni_final_stage_id"] = final_stage_id + merged.pop("omni_force_kv_transfer", None) + if force_kv_transfer: + merged["omni_force_kv_transfer"] = True payload = serialize_additional_information(merged) return OmniEngineCoreRequest.from_request( request, diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index 454266b7a..730018f9e 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/vllm_omni/engine/async_omni_engine.py @@ -798,11 +798,9 @@ class AsyncOmniEngine: supported_tasks=self.supported_tasks, ) request.external_req_id = cid - # Companions are stage-0-final: without this worker-visible tag - # the AR runner would still emit downstream connector payloads - # for them (the orchestrator-side final_stage_id=0 registration - # alone does not reach the worker). - request = apply_omni_final_stage_metadata(request, 0) + # Companions are stage-0-final for ordinary downstream payloads, + # but diffusion still needs their CFG KV caches. + request = apply_omni_final_stage_metadata(request, 0, force_kv_transfer=True) # Registration of this companion on stage-0's output processor is # deferred to Orchestrator._handle_add_companion, which routes diff --git a/vllm_omni/worker/omni_connector_model_runner_mixin.py b/vllm_omni/worker/omni_connector_model_runner_mixin.py index 363b850d4..ae7a0c0fd 100644 --- a/vllm_omni/worker/omni_connector_model_runner_mixin.py +++ b/vllm_omni/worker/omni_connector_model_runner_mixin.py @@ -122,6 +122,19 @@ class OmniConnectorModelRunnerMixin: # -- heterogeneous TP rank support -- rank_cfg = self._parse_rank_mapping(model_config) + if self._kv_transfer_manager is not None: + topology = getattr(self._kv_transfer_manager, "tp_topology", None) + effective_mapping = ( + getattr(topology, "source_tp_size", None), + getattr(topology, "target_tp_size", None), + getattr(topology, "local_rank", None), + ) + if all(isinstance(value, int) for value in effective_mapping): + rank_cfg = { + "from_tp": effective_mapping[0], + "to_tp": effective_mapping[1], + "local_rank": effective_mapping[2], + } self._from_tp: int = rank_cfg["from_tp"] self._to_tp: int = rank_cfg["to_tp"] self._local_rank: int = rank_cfg["local_rank"] @@ -1275,6 +1288,10 @@ class OmniConnectorModelRunnerMixin: For heterogeneous TP receive, the local rank is the target rank and must fetch one or more source-rank shards keyed as ``from_rank -> to_rank``. """ + if self._from_tp <= 1 and self._to_tp <= 1: + resolved_to_stage = self._next_stage_id if to_stage is None else to_stage + return [f"omni_{from_stage}_to_{resolved_to_stage}_kv_cache_{req_id}"] + remote_ranks = self.get_kv_remote_ranks() return [ self.get_kv_connector_key( @@ -1307,6 +1324,10 @@ class OmniConnectorModelRunnerMixin: chunk_id: int = 0, ) -> list[str]: """Build send-side connector keys for this rank's KV shard(s).""" + if self._from_tp <= 1 and self._to_tp <= 1: + resolved_to_stage = self._next_stage_id if to_stage is None else to_stage + return [f"omni_{from_stage}_to_{resolved_to_stage}_kv_cache_{req_id}"] + target_ranks = self.get_kv_target_ranks_for_send() return [ self.get_kv_connector_key(