diff --git a/requirements/common.txt b/requirements/common.txt index 2331e0f5d..3954b4587 100644 --- a/requirements/common.txt +++ b/requirements/common.txt @@ -13,3 +13,4 @@ openai-whisper>=20250625 imageio[ffmpeg]>=2.37.2 sox>=1.5.0 prettytable>=3.8.0 +aenum==3.1.16 diff --git a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py new file mode 100644 index 000000000..a5ebc2227 --- /dev/null +++ b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py @@ -0,0 +1,182 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import threading +from collections import deque +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest +import torch +from vllm.v1.request import RequestStatus + +from vllm_omni.distributed.omni_connectors.transfer_adapter.base import OmniTransferAdapterBase +from vllm_omni.distributed.omni_connectors.transfer_adapter.chunk_transfer_adapter import ( + OmniChunkTransferAdapter, +) +from vllm_omni.distributed.omni_connectors.utils.config import ConnectorSpec + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class DummyWaitingQueue(list): + def prepend_requests(self, requests): + self[:0] = list(requests) + + def add_request(self, request): + self.append(request) + + +def _req(req_id: str, status: RequestStatus, external_req_id: str | None = None): + return SimpleNamespace( + request_id=req_id, + external_req_id=external_req_id or req_id, + status=status, + prompt_token_ids=[], + num_computed_tokens=0, + additional_information=None, + ) + + +@pytest.fixture +def build_adapter(monkeypatch): + def _build(*, stage_id: int = 1, model_mode: str = "ar", max_num_seqs: int = 2): + connector = MagicMock() + connector.stage_id = stage_id + connector.get.return_value = None + connector.put.return_value = (True, 1, {}) + + def _fake_base_init(self, config): + self.config = config + self._pending_load_reqs = {} + self._finished_load_reqs = set() + self._pending_save_reqs = {} + self._finished_save_reqs = set() + self.stop_event = threading.Event() + self.lock = threading.Lock() + + monkeypatch.setattr(OmniTransferAdapterBase, "__init__", _fake_base_init) + monkeypatch.setattr( + OmniChunkTransferAdapter, + "create_connector", + classmethod(lambda cls, _model_config: connector), + ) + + model_config = SimpleNamespace(worker_type=model_mode) + scheduler_config = SimpleNamespace(max_num_seqs=max_num_seqs) + adapter = OmniChunkTransferAdapter( + SimpleNamespace(model_config=model_config, scheduler_config=scheduler_config) + ) + return adapter, connector + + return _build + + +@pytest.mark.parametrize( + ("raw_cfg", "expected_name", "expected_extra"), + [ + (None, "SharedMemoryConnector", {}), + (SimpleNamespace(name="YuanrongConnector", extra={"k": "v"}), "YuanrongConnector", {"k": "v"}), + ], +) +def test_create_connector_config_parsing(monkeypatch, raw_cfg, expected_name, expected_extra): + captured = {} + + def _fake_create(spec): + captured["spec"] = spec + return "ok" + + monkeypatch.setattr( + "vllm_omni.distributed.omni_connectors.transfer_adapter.chunk_transfer_adapter" + ".OmniConnectorFactory.create_connector", + _fake_create, + ) + + model_config = SimpleNamespace(stage_connector_config=raw_cfg) if raw_cfg is not None else SimpleNamespace() + connector = OmniChunkTransferAdapter.create_connector(model_config) + + assert connector == "ok" + assert isinstance(captured["spec"], ConnectorSpec) + assert captured["spec"].name == expected_name + assert captured["spec"].extra == expected_extra + + +def test_load_poll(build_adapter): + adapter, connector = build_adapter(stage_id=2, model_mode="ar") + request = _req("req-1", RequestStatus.WAITING, external_req_id="external-1") + + adapter.load_async(request) + payload = {"code_predictor_codes": [[1]], "hidden_states": torch.tensor([[2.0]]), "finished": True} + connector.get.return_value = (payload, 16) + adapter._poll_single_request("req-1") + + connector.get.assert_called_once_with("1", "2", "external-1_1_0") + assert request.additional_information == payload + assert adapter.get_req_chunk["req-1"] == 1 + assert "req-1" in adapter._finished_load_reqs + assert "req-1" in adapter.finished_requests + assert "req-1" not in adapter._pending_load_reqs + + +def test_save_async(build_adapter): + adapter, _ = build_adapter(stage_id=1) + request = SimpleNamespace(external_req_id="external-1") + + adapter.custom_process_next_stage_input_func = lambda **kwargs: {"x": [1], "finished": False} + adapter.save_async(pooling_output=None, request=request) + adapter.custom_process_next_stage_input_func = lambda **kwargs: {} + adapter.save_async(pooling_output=None, request=request) + + assert adapter.put_req_chunk["external-1"] == 1 + queued = adapter._pending_save_reqs["external-1"] + assert len(queued) == 1 + assert queued[0]["put_key"] == "external-1_1_0" + + +def test_update_request_payload(build_adapter): + adapter, _ = build_adapter() + + adapter._update_request_payload("ext", {"h": torch.tensor([[1.0]]), "codes": [1], "finished": False}) + merged = adapter._update_request_payload("ext", {"h": torch.tensor([[2.0]]), "codes": [2], "finished": True}) + + assert torch.equal(merged["h"], torch.tensor([[1.0], [2.0]])) + assert merged["codes"] == [1, 2] + assert merged["finished"] is True + + +def test_process_and_restore_queues(build_adapter): + adapter, _ = build_adapter(stage_id=1, max_num_seqs=8) + waiting_req = _req("w1", RequestStatus.WAITING) + running_req = _req("r1", RequestStatus.RUNNING) + waiting_queue = DummyWaitingQueue([waiting_req]) + running_queue = [running_req] + + adapter.process_pending_chunks(waiting_queue, running_queue) + assert waiting_req.status == RequestStatus.WAITING_FOR_CHUNK + assert running_req.status == RequestStatus.WAITING_FOR_CHUNK + assert waiting_queue == [] + assert running_queue == [] + + adapter.restore_queues(waiting_queue, running_queue) + assert waiting_queue == [waiting_req] + assert running_queue == [running_req] + assert adapter.waiting_for_chunk_waiting_requests == deque() + assert adapter.waiting_for_chunk_running_requests == deque() + + +def test_postprocess_scheduler_output(build_adapter): + adapter, _ = build_adapter() + adapter.requests_with_ready_chunks = {"new-ready", "cached-ready", "leftover"} + + scheduler_output = SimpleNamespace( + scheduled_new_reqs=[SimpleNamespace(req_id="new-ready")], + scheduled_cached_reqs=SimpleNamespace(req_ids=["cached-ready", "missing"]), + ) + requests = {"cached-ready": SimpleNamespace(additional_information={"k": "v"})} + + adapter.postprocess_scheduler_output(scheduler_output, requests) + + cached_info = scheduler_output.scheduled_cached_reqs.additional_information + assert cached_info["cached-ready"] == {"k": "v"} + assert cached_info["missing"] is None + assert adapter.requests_with_ready_chunks == {"leftover"} diff --git a/vllm_omni/config/model.py b/vllm_omni/config/model.py index dad07dc22..b099c19a0 100644 --- a/vllm_omni/config/model.py +++ b/vllm_omni/config/model.py @@ -45,6 +45,7 @@ class OmniModelConfig(ModelConfig): (default: "thinker") model_arch: Model architecture name (default: "Qwen2_5OmniForConditionalGeneration") + worker_type: Model Type, e.g., "ar" or "generation" engine_output_type: Optional output type specification for the engine. Used to route outputs to appropriate processors (e.g., "image", "audio", "latents"). If None, output type is inferred. @@ -63,6 +64,7 @@ class OmniModelConfig(ModelConfig): async_chunk: bool = False model_stage: str = "thinker" model_arch: str = "Qwen2_5OmniForConditionalGeneration" + worker_type: str | None = None engine_output_type: str | None = None hf_config_name: str | None = None custom_process_next_stage_input_func: str | None = None diff --git a/vllm_omni/core/sched/omni_ar_scheduler.py b/vllm_omni/core/sched/omni_ar_scheduler.py index 2af96e718..dc33f9cbc 100644 --- a/vllm_omni/core/sched/omni_ar_scheduler.py +++ b/vllm_omni/core/sched/omni_ar_scheduler.py @@ -1,6 +1,5 @@ from __future__ import annotations -import importlib from collections import defaultdict from dataclasses import asdict, dataclass from time import time @@ -20,9 +19,9 @@ from vllm.v1.request import Request, RequestStatus from vllm.v1.spec_decode.metrics import SpecDecodingStats from vllm_omni.core.sched.output import OmniSchedulerOutput -from vllm_omni.distributed.omni_connectors.adapter import get_chunk, put_chunk -from vllm_omni.distributed.omni_connectors.factory import OmniConnectorFactory -from vllm_omni.distributed.omni_connectors.utils.config import ConnectorSpec +from vllm_omni.distributed.omni_connectors.transfer_adapter.chunk_transfer_adapter import ( + OmniChunkTransferAdapter, +) logger = init_logger(__name__) @@ -65,24 +64,9 @@ class OmniARScheduler(VLLMScheduler): # Track requests that have already triggered prefill transfer to avoid duplicates self.transfer_triggered_requests: set[str] = set() model_config = self.vllm_config.model_config - self.omni_connector = None - if model_config.async_chunk: - connector_config = model_config.stage_connector_config - connector_specs = ConnectorSpec( - name=connector_config.get("name", "SharedMemoryConnector"), - extra=connector_config.get("extra", {}), - ) - self.omni_connector = OmniConnectorFactory.create_connector(connector_specs) - - custom_process_next_stage_input_func = getattr( - self.vllm_config.model_config, "custom_process_next_stage_input_func", None - ) - if custom_process_next_stage_input_func: - module_path, func_name = custom_process_next_stage_input_func.rsplit(".", 1) - module = importlib.import_module(module_path) - self.custom_process_next_stage_input_func = getattr(module, func_name) - - self.stage_id = getattr(self.vllm_config.model_config, "stage_id", None) + self.chunk_transfer_adapter = None + if getattr(model_config, "async_chunk", False): + self.chunk_transfer_adapter = OmniChunkTransferAdapter(self.vllm_config) def _get_kv_transfer_criteria(self) -> dict | None: # Note: vllm_config is available in Scheduler after super().__init__ @@ -152,7 +136,15 @@ class OmniARScheduler(VLLMScheduler): return False def schedule(self) -> SchedulerOutput: # type: ignore[override] - scheduler_output = super().schedule() + if self.chunk_transfer_adapter: + self.chunk_transfer_adapter.process_pending_chunks(self.waiting, self.running) + + try: + scheduler_output = super().schedule() + finally: + if self.chunk_transfer_adapter: + # Add request waiting for chunk to the waiting and running queue + self.chunk_transfer_adapter.restore_queues(self.waiting, self.running) try: # Late import to avoid circulars in some launch modes from .output import OmniNewRequestData @@ -181,9 +173,8 @@ class OmniARScheduler(VLLMScheduler): new_list.append(omni_nr) scheduler_output.scheduled_new_reqs = new_list # type: ignore[assignment] - if self.omni_connector is not None: - get_chunk(self.omni_connector, scheduler_output) - + if self.chunk_transfer_adapter: + self.chunk_transfer_adapter.postprocess_scheduler_output(scheduler_output, self.requests) # Add information about requests needing KV cache transfer finished_reqs = self.get_finished_requests_needing_kv_transfer() except Exception: @@ -312,6 +303,11 @@ class OmniARScheduler(VLLMScheduler): kv_transfer_params = self._free_request(request) if status_before_stop == RequestStatus.RUNNING: stopped_running_reqs.add(request) + elif status_before_stop == RequestStatus.WAITING_FOR_CHUNK: + # In async chunk mode, request may be in either queue. + # Remove from both to avoid stale queue entries. + stopped_running_reqs.add(request) + stopped_preempted_reqs.add(request) else: stopped_preempted_reqs.add(request) @@ -355,9 +351,8 @@ class OmniARScheduler(VLLMScheduler): num_nans_in_logits=request.num_nans_in_logits, ) ) - if self.omni_connector is not None: - custom_process_next_stage_input_func = self.custom_process_next_stage_input_func - put_chunk(self.omni_connector, pooler_output, request, custom_process_next_stage_input_func) + if self.chunk_transfer_adapter is not None: + self.chunk_transfer_adapter.save_async(pooler_output, request) else: # Invariant: EngineCore returns no partial prefill outputs. assert not prompt_logprobs_tensors diff --git a/vllm_omni/core/sched/omni_generation_scheduler.py b/vllm_omni/core/sched/omni_generation_scheduler.py index 9abcae94b..b86d49ed5 100644 --- a/vllm_omni/core/sched/omni_generation_scheduler.py +++ b/vllm_omni/core/sched/omni_generation_scheduler.py @@ -16,9 +16,9 @@ from vllm.v1.request import Request, RequestStatus from vllm.v1.spec_decode.metrics import SpecDecodingStats from vllm_omni.core.sched.output import OmniCachedRequestData, OmniNewRequestData -from vllm_omni.distributed.omni_connectors.adapter import get_chunk_for_generation -from vllm_omni.distributed.omni_connectors.factory import OmniConnectorFactory -from vllm_omni.distributed.omni_connectors.utils.config import ConnectorSpec +from vllm_omni.distributed.omni_connectors.transfer_adapter.chunk_transfer_adapter import ( + OmniChunkTransferAdapter, +) from vllm_omni.outputs import OmniModelRunnerOutput @@ -26,15 +26,9 @@ class OmniGenerationScheduler(VLLMScheduler): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) model_config = self.vllm_config.model_config - self.omni_connector = None - if model_config.async_chunk: - connector_config = model_config.stage_connector_config - connector_specs = ConnectorSpec( - name=connector_config.get("name", "SharedMemoryConnector"), - extra=connector_config.get("extra", {}), - ) - self.omni_connector = OmniConnectorFactory.create_connector(connector_specs) - self.stage_id = getattr(self.vllm_config.model_config, "stage_id", None) + self.chunk_transfer_adapter = None + if getattr(model_config, "async_chunk", False): + self.chunk_transfer_adapter = OmniChunkTransferAdapter(self.vllm_config) def schedule(self) -> SchedulerOutput: """Diffusion fast path: @@ -59,18 +53,18 @@ class OmniGenerationScheduler(VLLMScheduler): # Temporary queue: preserve waiting order, do not disturb non-diffusion requests skipped_waiting_requests = create_request_queue(self.policy) req_index = 0 + if self.chunk_transfer_adapter: + self.chunk_transfer_adapter.process_pending_chunks(self.waiting, self.running) + # OMNI: Track requests that are already finished (e.g., marked by connector) # These should be removed from running and not scheduled already_finished_reqs: set[Request] = set() while req_index < len(self.running) and token_budget > 0: request = self.running[req_index] - if self.omni_connector is not None: - get_chunk_for_generation(self.omni_connector, request) - # OMNI: Skip requests that are not in self.requests # This can happen when connector marks request as finished and it's removed from requests if request.request_id not in self.requests or ( - self.omni_connector is None and request.status == RequestStatus.FINISHED_STOPPED + self.chunk_transfer_adapter is None and request.status == RequestStatus.FINISHED_STOPPED ): already_finished_reqs.add(request) req_index += 1 @@ -107,12 +101,9 @@ class OmniGenerationScheduler(VLLMScheduler): # independent of pooling_params) while self.waiting and token_budget > 0 and len(self.running) < self.max_num_running_reqs: request = self.waiting.peek_request() - if self.omni_connector is not None: - get_chunk_for_generation(self.omni_connector, request) - # OMNI: Skip requests that are not in self.requests if request.request_id not in self.requests or ( - self.omni_connector is None and request.status == RequestStatus.FINISHED_STOPPED + self.chunk_transfer_adapter is None and request.status == RequestStatus.FINISHED_STOPPED ): # Pop the finished request from waiting queue and don't schedule it self.waiting.pop_request() @@ -154,7 +145,11 @@ class OmniGenerationScheduler(VLLMScheduler): # If fast path scheduled none, fall back to the original scheduling if not num_scheduled_tokens: - return super().schedule() + res = super().schedule() + if self.chunk_transfer_adapter: + self.chunk_transfer_adapter.restore_queues(self.waiting, self.running) + self.chunk_transfer_adapter.postprocess_scheduler_output(res) + return res # Compute common prefix blocks (aligned with v1) num_common_prefix_blocks = [0] * len(self.kv_cache_config.kv_cache_groups) @@ -254,6 +249,11 @@ class OmniGenerationScheduler(VLLMScheduler): new_list.append(omni_nr) scheduler_output.scheduled_new_reqs = new_list # type: ignore[assignment] + + if self.chunk_transfer_adapter: + self.chunk_transfer_adapter.restore_queues(self.waiting, self.running) + self.chunk_transfer_adapter.postprocess_scheduler_output(scheduler_output) + except Exception: # If anything goes wrong, leave the original output unchanged init_logger(__name__).exception("Failed to wrap scheduled_new_reqs with OmniNewRequestData") @@ -355,7 +355,7 @@ class OmniGenerationScheduler(VLLMScheduler): # Diffusion request: completes in one step; mark finished and free resources if request.status == RequestStatus.FINISHED_STOPPED or ( - self.omni_connector is None and request.num_computed_tokens >= request.num_prompt_tokens + self.chunk_transfer_adapter is None and request.num_computed_tokens >= request.num_prompt_tokens ): request.status = RequestStatus.FINISHED_STOPPED # Optional: set a stop_reason for front-end clarity @@ -371,6 +371,11 @@ class OmniGenerationScheduler(VLLMScheduler): kv_transfer_params = self._free_request(request) if status_before_stop == RequestStatus.RUNNING: stopped_running_reqs.add(request) + elif status_before_stop == RequestStatus.WAITING_FOR_CHUNK: + # In async chunk mode, request may be in either queue. + # Remove from both to avoid stale queue entries. + stopped_running_reqs.add(request) + stopped_preempted_reqs.add(request) else: stopped_preempted_reqs.add(request) diff --git a/vllm_omni/distributed/omni_connectors/adapter.py b/vllm_omni/distributed/omni_connectors/adapter.py index 66aebddc1..1d1dd6f0f 100644 --- a/vllm_omni/distributed/omni_connectors/adapter.py +++ b/vllm_omni/distributed/omni_connectors/adapter.py @@ -5,14 +5,7 @@ import time from collections.abc import Callable -from typing import TYPE_CHECKING, Any - -import torch -from vllm.v1.core.sched.output import SchedulerOutput -from vllm.v1.request import Request, RequestStatus - -if TYPE_CHECKING: - from .connectors.base import OmniConnectorBase +from typing import Any from vllm_omni.entrypoints.stage_utils import OmniStageTaskType from vllm_omni.metrics import OrchestratorAggregator @@ -180,178 +173,6 @@ def try_recv_via_connector( return None, None -def update_request_payload(connector: "OmniConnectorBase", req_id: str, payload_data: dict[str, Any]) -> dict[str, Any]: - """Update the payload data for a request in the connector. - - Args: - connector: OmniConnectorBase instance - req_id: Request ID to update - payload_data: New payload data to store - """ - origin_payload = connector.request_payload[req_id] - for key, value in payload_data.items(): - if key == "finished": - continue - elif isinstance(value, torch.Tensor) and key in origin_payload: - payload_data[key] = torch.cat([origin_payload[key], value], dim=0) - elif isinstance(value, list) and key in origin_payload: - payload_data[key] = origin_payload[key] + value - - connector.request_payload[req_id] = payload_data - return payload_data - - -def get_chunk( - connector: "OmniConnectorBase", - scheduler_output: SchedulerOutput, -) -> None: - """Retrieve a chunk of pooling output. - - Args: - connector: OmniConnectorBase instance - scheduler_output: Partial scheduler output dictionary - - Returns: - None: This function modifies scheduler_output in place - """ - stage_id = connector.stage_id - if stage_id == 0: - return - - target_stage_id = stage_id - 1 - # Handle new requests - for new_req_data in scheduler_output.scheduled_new_reqs: - connector.request_ids_mapping[new_req_data.req_id] = new_req_data.external_req_id - req_id = new_req_data.external_req_id - chunk_id = connector.get_requests[req_id] - connector_get_key = f"{req_id}_{target_stage_id}_{chunk_id}" - payload_data = get_through_connector(connector, target_stage_id, stage_id, req_id, connector_get_key) - if payload_data: - new_req_data.additional_information = payload_data - connector.request_payload[req_id] = payload_data - if payload_data.get("finished"): - connector.finished_requests.add(req_id) - - # Handle cached/running requests - cached_reqs = scheduler_output.scheduled_cached_reqs - if not hasattr(cached_reqs, "additional_information"): - cached_reqs.additional_information = {} - - for i, cached_req_id in enumerate(cached_reqs.req_ids): - req_id = connector.request_ids_mapping.get(cached_req_id, cached_req_id) - if req_id in connector.finished_requests: - continue - chunk_id = connector.get_requests[req_id] - connector_get_key = f"{req_id}_{target_stage_id}_{chunk_id}" - payload_data = get_through_connector(connector, target_stage_id, stage_id, req_id, connector_get_key) - if payload_data: - payload_data = update_request_payload(connector, req_id, payload_data) - cached_reqs.additional_information[cached_req_id] = payload_data - if payload_data.get("finished"): - connector.finished_requests.add(req_id) - - -def get_through_connector(connector, target_stage_id, stage_id, req_id, connector_get_key): - # TODO: add correct check mechanism for the payload_data - max_wait = 300 - for _ in range(max_wait): - result = connector.get( - from_stage=str(target_stage_id), - to_stage=str(stage_id), - get_key=connector_get_key, - ) - payload_data = None - if result: - payload_data, size = result - if payload_data: - connector.request_prompt_token_ids[req_id] = payload_data.get("thinker_input_ids", []) - connector.get_requests[req_id] += 1 - logger.debug("[Stage-%d] Received one chunk for request %s", stage_id, connector_get_key) - break - time.sleep(0.01) - return payload_data - - -def get_chunk_for_generation( - connector: "OmniConnectorBase", - request: Request, -) -> None: - """Retrieve a chunk of pooling output. - - Args: - connector: OmniConnectorBase instance - request: Request object - - Returns: - None: This function modifies request in place - """ - stage_id = connector.stage_id - target_stage_id = stage_id - 1 - request_id = request.external_req_id - - if request_id in connector.finished_requests: - return - - chunk_id = connector.get_requests[request_id] - connector_get_key = f"{request_id}_{target_stage_id}_{chunk_id}" - payload_data = get_through_connector(connector, target_stage_id, stage_id, request_id, connector_get_key) - if not payload_data: - return - - if payload_data.get("finished"): - connector.finished_requests.add(request_id) - request.status = RequestStatus.FINISHED_STOPPED - request.prompt_token_ids = payload_data.get("code_predictor_codes", []) - request.num_computed_tokens = 0 - - -def put_chunk( - connector: "OmniConnectorBase", - pooling_output: dict[str, Any], - request: Request, - custom_process_input_func: Callable[[dict[str, Any], Request], dict[str, Any] | None] | None = None, -) -> None: - """Store a chunk of pooling output. - - Args: - connector: OmniConnectorBase instance - pooling_output: Partial pooling output dictionary - request: Request object - custom_process_input_func: Optional custom function to process input - - Returns: - None: This function sends data via connector - """ - stage_id = connector.stage_id - next_stage_id = stage_id + 1 - request_id = request.external_req_id - chunk_id = connector.put_requests[request_id] - connector_put_key = f"{request_id}_{stage_id}_{chunk_id}" - payload_data = None - - # TODO: add default process_input_func to handle the payload_data ? - if custom_process_input_func: - try: - payload_data = custom_process_input_func( - connector=connector, - pooling_output=pooling_output, - request=request, - ) - except Exception as e: - logger.error(f"Failed to use custom_process_input_func for payload extraction: {e}") - - if not payload_data: - return - - success, size, metadata = connector.put( - from_stage=str(stage_id), to_stage=str(next_stage_id), put_key=connector_put_key, data=payload_data - ) - - if success: - connector.put_requests[request_id] += 1 - logger.debug("[Stage-%d] Sent %s", stage_id, connector_put_key) - - def compute_talker_prompt_ids_length(prompt_ids: list[int]) -> int: """Compute the length of the talker prompt ids. diff --git a/vllm_omni/distributed/omni_connectors/connectors/shm_connector.py b/vllm_omni/distributed/omni_connectors/connectors/shm_connector.py index 9ad32e1c7..719dbeae7 100644 --- a/vllm_omni/distributed/omni_connectors/connectors/shm_connector.py +++ b/vllm_omni/distributed/omni_connectors/connectors/shm_connector.py @@ -4,7 +4,6 @@ import fcntl import os import time -from collections import defaultdict from typing import Any from vllm_omni.entrypoints.stage_utils import shm_read_bytes, shm_write_bytes @@ -25,14 +24,6 @@ class SharedMemoryConnector(OmniConnectorBase): self.config = config self.stage_id = config.get("stage_id", -1) self.device = config.get("device", "cuda:0") - self.put_requests: dict[str, int] = defaultdict(int) - self.get_requests: dict[str, int] = defaultdict(int) - self.finished_requests: set[str] = set() - self.request_payload = {} - self.request_prompt_token_ids: dict[str, list[int]] = defaultdict(list) - self.code_prompt_token_ids: dict[str, list[list[int]]] = defaultdict(list) - self.request_ids_mapping: dict[str, str] = {} - # Default threshold matches legacy behavior (64KB) self.threshold = int(config.get("shm_threshold_bytes", 65536)) self._metrics = { "puts": 0, diff --git a/vllm_omni/distributed/omni_connectors/transfer_adapter/__init__.py b/vllm_omni/distributed/omni_connectors/transfer_adapter/__init__.py new file mode 100644 index 000000000..208f01a7c --- /dev/null +++ b/vllm_omni/distributed/omni_connectors/transfer_adapter/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm_omni/distributed/omni_connectors/transfer_adapter/base.py b/vllm_omni/distributed/omni_connectors/transfer_adapter/base.py new file mode 100644 index 000000000..0962808ed --- /dev/null +++ b/vllm_omni/distributed/omni_connectors/transfer_adapter/base.py @@ -0,0 +1,115 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import threading +import time +from typing import Any + +from ..utils.logging import get_connector_logger + +logger = get_connector_logger(__name__) + + +class OmniTransferAdapterBase: + """Base class for managing data transfer via OmniConnector. + + This class handles the core loop logic and connector interactions, but + leaves the specific data processing (chunks, KV cache, etc.) to subclasses. + """ + + def __init__(self, config: Any): + self.config = config + if not hasattr(self, "connector"): + self.connector = None + # Requests that are waiting to be polled + self._pending_load_reqs = {} + # Requests that have successfully retrieved data + self._finished_load_reqs = set() + + # Requests that are waiting to be saved + self._pending_save_reqs = {} + # Requests that have successfully saved data + self._finished_save_reqs = set() + + self.stop_event = threading.Event() + self.lock = threading.Lock() + + self.recv_thread = threading.Thread(target=self.recv_loop, daemon=True) + self.recv_thread.start() + + self.save_thread = threading.Thread(target=self.save_loop, daemon=True) + self.save_thread.start() + + @classmethod + def create_connector(cls, model_config: Any): + raise NotImplementedError + + def recv_loop(self): + """Loop to poll for incoming data.""" + while not self.stop_event.is_set(): + # Iterate over a snapshot of pending requests + with self.lock: + pending_reqs_ids = list(self._pending_load_reqs.keys()) + + for req_id in pending_reqs_ids: + try: + self._poll_single_request(req_id) + except Exception as e: + logger.warning(f"Error receiving data for {req_id}: {e}") + + time.sleep(0.001) + + def save_loop(self): + """Loop to send outgoing data.""" + while not self.stop_event.is_set(): + task = None + with self.lock: + pending_save_reqs_ids = list(self._pending_save_reqs.keys()) + for req_id in pending_save_reqs_ids: + if self._pending_save_reqs[req_id]: + task = self._pending_save_reqs[req_id].popleft() + if not self._pending_save_reqs[req_id]: + del self._pending_save_reqs[req_id] + break + + if task: + try: + self._send_single_request(task) + except Exception as e: + logger.error(f"Error saving data for {task.get('request_id')}: {e}") + else: + time.sleep(0.001) + + def _poll_single_request(self, *args, **kwargs): + """Poll connector for a single request task. + Subclasses should implement request-specific receive behavior.""" + raise NotImplementedError + + def _send_single_request(self, *args, **kwargs): + """Send one pending save request task to the connector. + Subclasses should implement task-specific handling logic.""" + raise NotImplementedError + + def load_async(self, *args, **kwargs): + """Register a request to load data. To be implemented by subclasses.""" + raise NotImplementedError + + def save_async(self, *args, **kwargs): + """Submit data to be saved. To be implemented by subclasses.""" + raise NotImplementedError + + def load(self, *args, **kwargs): + """Load request data from connector synchronously. To be implemented by subclasses.""" + raise NotImplementedError + + def save(self, *args, **kwargs): + """Save data to connector synchronously. To be implemented by subclasses.""" + raise NotImplementedError + + def get_finished_requests(self): + """Get finished loaded or saved requests""" + raise NotImplementedError + + def shutdown(self): + """Stop background loops and close the connector.""" + raise NotImplementedError diff --git a/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py b/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py new file mode 100644 index 000000000..70d38a9d6 --- /dev/null +++ b/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py @@ -0,0 +1,339 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import importlib +from collections import defaultdict, deque +from typing import Any + +import torch +from vllm.v1.request import Request, RequestStatus + +from ..factory import OmniConnectorFactory +from ..utils.config import ConnectorSpec +from ..utils.logging import get_connector_logger +from .base import OmniTransferAdapterBase + +logger = get_connector_logger(__name__) + + +class OmniChunkTransferAdapter(OmniTransferAdapterBase): + """Chunk-level transfer adapter for Omni connector pipelines. + + This class coordinates per-request chunk exchange between adjacent stages, + and implements asynchronous get/put of chunks via background threads. + It tracks per-request chunk indices for put/get, and accumulates + payloads across chunks (concatenating tensors/lists in AR mode). It also + caches prompt token ids and additional information for scheduler use. + + Scheduler integration is handled via WAITING_FOR_CHUNK transitions: + requests are moved to waiting for chunk deque while polling, then restored + to waiting/running queues once a chunk arrives. The requests will finish + loading chunk util detecting the payload "finished" flag. + + The base class owns background recv/save loops; load/save only enqueue + work and return immediately. + """ + + def __init__(self, vllm_config: Any): + model_config = vllm_config.model_config + self.scheduler_max_num_seqs = vllm_config.scheduler_config.max_num_seqs + self.connector = self.create_connector(model_config) + super().__init__(model_config) + self.model_mode = getattr(model_config, "worker_type", "ar") + # State specific to Chunk management + self.custom_process_next_stage_input_func = None + custom_process_next_stage_input_func = getattr(model_config, "custom_process_next_stage_input_func", None) + if custom_process_next_stage_input_func: + module_path, func_name = custom_process_next_stage_input_func.rsplit(".", 1) + module = importlib.import_module(module_path) + self.custom_process_next_stage_input_func = getattr(module, func_name) + # mapping for request id and chunk id + self.put_req_chunk: dict[str, int] = defaultdict(int) + self.get_req_chunk: dict[str, int] = defaultdict(int) + self.finished_requests: set[str] = set() + self.request_payload = {} + self.code_prompt_token_ids: dict[str, list[list[int]]] = defaultdict(list) + self.request_ids_mapping: dict[str, str] = {} + + self.waiting_for_chunk_waiting_requests: deque[Any] = deque() + self.waiting_for_chunk_running_requests: deque[Any] = deque() + self.requests_with_ready_chunks = set() + + @classmethod + def create_connector(cls, model_config: Any): + connector_config = getattr(model_config, "stage_connector_config", None) + if connector_config is None: + connector_config = {} + elif not isinstance(connector_config, dict): + connector_config = { + "name": getattr(connector_config, "name", None), + "extra": getattr(connector_config, "extra", {}), + } + + connector_specs = ConnectorSpec( + name=connector_config.get("name", "SharedMemoryConnector"), + extra=connector_config.get("extra", {}), + ) + return OmniConnectorFactory.create_connector(connector_specs) + + def load_async(self, request: Request): + """Register a request for asynchronous chunk retrieval. + + This method does not read from the connector directly. It records + request metadata and enqueues the request id for the background + receive loop to poll. + + Stage-0 has no upstream producer, so this call is a no-op there. + + Args: + request: The request object needing data. + """ + stage_id = self.connector.stage_id + request_id = request.request_id + self.request_ids_mapping[request_id] = request.external_req_id + + if stage_id == 0: + return + if not hasattr(request, "additional_information"): + request.additional_information = None + with self.lock: + self._pending_load_reqs[request_id] = request + + def save_async(self, pooling_output: torch.Tensor | None = None, request: Request | None = None): + """Build and enqueue one chunk for asynchronous sending. + + Payload extraction is executed in the caller thread via + ``custom_process_next_stage_input_func`` + + Args: + pooling_output: Partial pooling output dictionary + request: Request object + """ + stage_id = self.connector.stage_id + next_stage_id = stage_id + 1 + request_id = request.external_req_id + chunk_id = self.put_req_chunk[request_id] + + # Process payload in main thread to avoid race conditions on request state + payload_data = None + if self.custom_process_next_stage_input_func: + try: + payload_data = self.custom_process_next_stage_input_func( + transfer_manager=self, + pooling_output=pooling_output, + request=request, + ) + + except Exception as e: + logger.error(f"Failed to use custom_process_input_func for payload extraction: {e}") + + if not payload_data: + return + + # Increment chunk_id + self.put_req_chunk[request_id] += 1 + connector_put_key = f"{request_id}_{stage_id}_{chunk_id}" + + task = { + "stage_id": stage_id, + "next_stage_id": next_stage_id, + "put_key": connector_put_key, + "data": payload_data, + "request_id": request_id, + } + + with self.lock: + if request_id not in self._pending_save_reqs: + self._pending_save_reqs[request_id] = deque() + self._pending_save_reqs[request_id].append(task) + + def _poll_single_request(self, req_id: str): + stage_id = self.connector.stage_id + target_stage_id = stage_id - 1 + chunk_id = self.get_req_chunk[req_id] + external_req_id = self.request_ids_mapping.get(req_id, req_id) + connector_get_key = f"{external_req_id}_{target_stage_id}_{chunk_id}" + + # Use timeout=0 for non-blocking poll + result = self.connector.get( + str(target_stage_id), + str(stage_id), + connector_get_key, + ) + + if result is None: + return + + payload_data, size = result + + if payload_data: + # Update connector state + self.get_req_chunk[req_id] += 1 + req = self._pending_load_reqs[req_id] + + if self.model_mode == "ar": + self._update_request_payload(external_req_id, payload_data) + req.additional_information = payload_data + if payload_data.get("finished"): + self.finished_requests.add(req_id) + else: + if payload_data.get("finished"): + self.finished_requests.add(req_id) + req.status = RequestStatus.FINISHED_STOPPED + + req.prompt_token_ids = payload_data.get("code_predictor_codes", []) + req.num_computed_tokens = 0 + + # Mark as finished for consumption + with self.lock: + self._finished_load_reqs.add(req_id) + if req_id in self._pending_load_reqs: + del self._pending_load_reqs[req_id] + logger.debug(f"[Stage-{stage_id}] Received one chunk for key {connector_get_key}") + + def _update_request_payload(self, req_id: str, payload_data: dict[str, Any]) -> dict[str, Any]: + """Update the payload data for a request in the connector. + + Args: + connector: OmniConnectorBase instance + req_id: Request ID to update + payload_data: New payload data to store + """ + if req_id not in self.request_payload: + self.request_payload[req_id] = payload_data + return payload_data + origin_payload = self.request_payload[req_id] + for key, value in payload_data.items(): + if key == "finished": + continue + elif isinstance(value, torch.Tensor) and key in origin_payload: + payload_data[key] = torch.cat([origin_payload[key], value], dim=0) + elif isinstance(value, list) and key in origin_payload: + payload_data[key] = origin_payload[key] + value + + self.request_payload[req_id] = payload_data + return payload_data + + def _send_single_request(self, task: dict): + connector_put_key = task["put_key"] + stage_id = task["stage_id"] + next_stage_id = task["next_stage_id"] + payload_data = task["data"] + + success, size, metadata = self.connector.put( + from_stage=str(stage_id), + to_stage=str(next_stage_id), + put_key=connector_put_key, + data=payload_data, + ) + + if success: + logger.debug(f"[Stage-{stage_id}] Sent {connector_put_key}") + + ######################################################################## + # Schedule Helper + ######################################################################## + + def process_pending_chunks( + self, + waiting_queue: Any, + running_queue: list[Request], + ) -> None: + """ + Process pending chunks for waiting and running queues. + """ + if self.connector.stage_id == 0: + return + finished_load_reqs = self.get_finished_requests() + self._process_chunk_queue( + waiting_queue, self.waiting_for_chunk_waiting_requests, RequestStatus.WAITING, finished_load_reqs + ) + self._process_chunk_queue( + running_queue, self.waiting_for_chunk_running_requests, RequestStatus.RUNNING, finished_load_reqs + ) + while len(running_queue) > self.scheduler_max_num_seqs: + request = running_queue.pop() + waiting_queue.prepend_requests([request]) + + def restore_queues(self, waiting_queue: Any, running_queue: list[Request]) -> None: + """ + Restore requests waiting for chunk to the waiting and running queues. + """ + # Add request waiting for chunk to the waiting and running queue + for request in self.waiting_for_chunk_waiting_requests: + waiting_queue.add_request(request) + self.waiting_for_chunk_waiting_requests = deque() + + if self.waiting_for_chunk_running_requests: + running_queue.extend(self.waiting_for_chunk_running_requests) + self.waiting_for_chunk_running_requests = deque() + + def postprocess_scheduler_output( + self, + scheduler_output: Any, + requests: dict[str, Request] | None = None, + ) -> None: + """ + Add additional info for cached requests and + clean up ready chunks from scheduler output. + """ + if requests is not None: + self.attach_cached_additional_information(scheduler_output, requests) + self._clear_chunk_ready(scheduler_output) + + @staticmethod + def attach_cached_additional_information(scheduler_output: Any, requests: dict[str, Request]) -> None: + cached_reqs = getattr(scheduler_output, "scheduled_cached_reqs", None) + if not cached_reqs: + return + if not hasattr(cached_reqs, "additional_information"): + cached_reqs.additional_information = {} + for req_id in cached_reqs.req_ids: + request = requests.get(req_id) if req_id else None + additional_info = getattr(request, "additional_information", None) if request else None + cached_reqs.additional_information[req_id] = additional_info + + def _process_chunk_queue( + self, + queue: Any, + waiting_for_chunk_list: deque[Any], + target_status: RequestStatus, + finished_load_reqs: set[str], + ) -> None: + queue_snapshot = list(queue) + for request in queue_snapshot: + if request.status != RequestStatus.WAITING_FOR_CHUNK: + if request.request_id in self.requests_with_ready_chunks: + # Requests that have loaded chunk from last round + # of schedule, but have not scheduled + continue + if request.request_id in self.finished_requests: + request.additional_information = {} + continue + # Requests that waiting for chunk + self.load_async(request) + request.status = RequestStatus.WAITING_FOR_CHUNK + else: + if request.request_id in finished_load_reqs: + request.status = target_status + self.requests_with_ready_chunks.add(request.request_id) + continue + queue.remove(request) + waiting_for_chunk_list.append(request) + + def _clear_chunk_ready(self, scheduler_output: Any) -> None: + if scheduler_output.scheduled_new_reqs: + for req_data in scheduler_output.scheduled_new_reqs: + if req_data.req_id in self.requests_with_ready_chunks: + self.requests_with_ready_chunks.remove(req_data.req_id) + + if scheduler_output.scheduled_cached_reqs: + for req_id in scheduler_output.scheduled_cached_reqs.req_ids: + if req_id in self.requests_with_ready_chunks: + self.requests_with_ready_chunks.remove(req_id) + + def get_finished_requests(self): + with self.lock: + finished_load = set(self._finished_load_reqs) + self._finished_load_reqs = set() + return finished_load diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index d550cf555..6e6ee465f 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/vllm_omni/engine/arg_utils.py @@ -63,6 +63,7 @@ class OmniEngineArgs(EngineArgs): If None, default processing is used. stage_connector_spec: Extra configuration for stage connector async_chunk: If set to True, perform async chunk + worker_type: Model Type, e.g., "ar" or "generation" """ stage_id: int = 0 @@ -74,6 +75,7 @@ class OmniEngineArgs(EngineArgs): stage_connector_spec: dict[str, Any] = field(default_factory=dict) async_chunk: bool = False omni_kv_config: dict | None = None + worker_type: str | None = None def draw_hf_text_config(self, config_dict: dict) -> Qwen3OmniMoeTextConfig: # transformers' get_text_config method is used to get the text config from thinker_config. @@ -129,6 +131,7 @@ class OmniEngineArgs(EngineArgs): config_dict["async_chunk"] = self.async_chunk config_dict["model_stage"] = self.model_stage config_dict["model_arch"] = self.model_arch + config_dict["worker_type"] = self.worker_type config_dict["engine_output_type"] = self.engine_output_type # Build stage_connector_config from stage_connector_spec stage_connector_config = { @@ -165,6 +168,7 @@ class AsyncOmniEngineArgs(AsyncEngineArgs): Used to route outputs to appropriate processors (e.g., "image", "audio", "latents"). If None, output type is inferred. stage_connector_spec: Extra configuration for stage connector + worker_type: Model Type, e.g., "ar" or "generation" """ stage_id: int = 0 @@ -176,6 +180,7 @@ class AsyncOmniEngineArgs(AsyncEngineArgs): stage_connector_spec: dict[str, Any] = field(default_factory=dict) async_chunk: bool = False omni_kv_config: dict | None = None + worker_type: str | None = None def draw_hf_text_config(self, config_dict: dict) -> Qwen3OmniMoeTextConfig: # transformers' get_text_config method is used to get the text config from thinker_config. @@ -221,6 +226,7 @@ class AsyncOmniEngineArgs(AsyncEngineArgs): config_dict["async_chunk"] = self.async_chunk config_dict["model_stage"] = self.model_stage config_dict["model_arch"] = self.model_arch + config_dict["worker_type"] = self.worker_type config_dict["engine_output_type"] = self.engine_output_type stage_connector_config = { "name": self.stage_connector_spec.get("name", "SharedMemoryConnector"), diff --git a/vllm_omni/entrypoints/omni_stage.py b/vllm_omni/entrypoints/omni_stage.py index e32adf737..92ca363ff 100644 --- a/vllm_omni/entrypoints/omni_stage.py +++ b/vllm_omni/entrypoints/omni_stage.py @@ -201,7 +201,7 @@ def _sequential_init_lock(engine_args: dict[str, Any], stage_init_timeout: int = def _resolve_worker_cls(engine_args: dict[str, Any]) -> None: - worker_type = engine_args.pop("worker_type", None) + worker_type = engine_args.get("worker_type", None) if not worker_type: return if engine_args.get("worker_cls"): diff --git a/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py b/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py index 701923309..11aea4f04 100644 --- a/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py +++ b/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py @@ -84,7 +84,7 @@ def _validate_stage_inputs(stage_list, engine_input_source): def thinker2talker_async_chunk( - connector: Any, + transfer_manager: Any, pooling_output: dict[str, Any], request: OmniEngineCoreRequest, ) -> list[dict[str, Any]]: @@ -96,7 +96,7 @@ def thinker2talker_async_chunk( """ request_id = request.external_req_id - chunk_id = connector.put_requests[request_id] + chunk_id = transfer_manager.put_req_chunk[request_id] if chunk_id == 0: all_token_ids = request.all_token_ids # prefill + decode prompt_token_ids = request.prompt_token_ids @@ -114,12 +114,12 @@ def thinker2talker_async_chunk( "tts_pad_embed": pooling_output.get("tts_pad_embed").detach().cpu(), "finished": torch.tensor(request.is_finished(), dtype=torch.bool), } - if connector.request_payload.get(request_id) is None: + if transfer_manager.request_payload.get(request_id) is None: if not request.is_finished(): - connector.request_payload[request_id] = talker_additional_info + transfer_manager.request_payload[request_id] = talker_additional_info return None else: - save_payload = connector.request_payload.pop(request_id) + save_payload = transfer_manager.request_payload.pop(request_id) talker_additional_info["thinker_embeddings"] = torch.cat( (save_payload.get("thinker_embeddings"), talker_additional_info.get("thinker_embeddings")), dim=0 ) @@ -206,7 +206,7 @@ def thinker2talker( def talker2code2wav_async_chunk( - connector: Any, + transfer_manager: Any, pooling_output: dict[str, Any], request: OmniEngineCoreRequest, ): @@ -241,8 +241,8 @@ def talker2code2wav_async_chunk( request_id = request.external_req_id chunk_size = left_context_size = 25 - connector.code_prompt_token_ids[request_id].append(codec_codes) - length = len(connector.code_prompt_token_ids[request_id]) + transfer_manager.code_prompt_token_ids[request_id].append(codec_codes) + length = len(transfer_manager.code_prompt_token_ids[request_id]) chunk_length = length % chunk_size if chunk_length != 0 and not request.is_finished(): return None @@ -252,7 +252,10 @@ def talker2code2wav_async_chunk( info = { "code_predictor_codes": ( - torch.tensor(connector.code_prompt_token_ids[request_id][-end_index:]).transpose(0, 1).reshape(-1).tolist() + torch.tensor(transfer_manager.code_prompt_token_ids[request_id][-end_index:]) + .transpose(0, 1) + .reshape(-1) + .tolist() ), "finished": torch.tensor(request.is_finished(), dtype=torch.bool), } diff --git a/vllm_omni/patch.py b/vllm_omni/patch.py index 90c718d40..e2f069d8f 100644 --- a/vllm_omni/patch.py +++ b/vllm_omni/patch.py @@ -1,5 +1,6 @@ import sys +from aenum import extend_enum from vllm.inputs.data import TokensPrompt as _OriginalTokensPrompt from vllm.model_executor.layers.rotary_embedding import ( MRotaryEmbedding as _OriginalMRotaryEmbedding, @@ -8,6 +9,7 @@ from vllm.v1.engine import EngineCoreOutput as _OriginalEngineCoreOutput from vllm.v1.engine import EngineCoreOutputs as _OriginalEngineCoreOutputs from vllm.v1.engine import EngineCoreRequest as _OriginalEngineCoreRequest from vllm.v1.request import Request as _OriginalRequest +from vllm.v1.request import RequestStatus import vllm_omni.logger # noqa: F401 from vllm_omni.engine import OmniEngineCoreOutput, OmniEngineCoreOutputs, OmniEngineCoreRequest @@ -15,6 +17,12 @@ from vllm_omni.inputs.data import OmniTokensPrompt from vllm_omni.model_executor.layers.rotary_embedding import OmniMRotaryEmbedding from vllm_omni.request import OmniRequest +# Extend RequestStatus enum with omni-specific statuses +if not hasattr(RequestStatus, "WAITING_FOR_CHUNK"): + # The value - 1 is intentionally chosen to ensure it is treated + # as a non-finished state and remains compatible with existing comparisons. + extend_enum(RequestStatus, "WAITING_FOR_CHUNK", -1) + for module_name, module in sys.modules.items(): # only do patch on module of vllm, pass others if "vllm" not in module_name: