[Feature]async scheduling to overlap chunk IO and compute (#951)

Signed-off-by: CHEN <116010019@link.cuhk.edu.cn>
Co-authored-by: Bhanu068 <voutharoja.bhanu06@gmail.com>
Co-authored-by: Gao Han <gaohan19@huawei.com>

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

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"}

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

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

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)
 

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.
 

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,

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

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

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

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"),

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"):

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),
     }

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: