async_schedule

Signed-off-by: CHEN <116010019@link.cuhk.edu.cn>

Author: Shirley125

vllm_omni/core/sched/omni_ar_scheduler.py

@@ -20,9 +20,10 @@
 from vllm.v1.spec_decode.metrics import SpecDecodingStats
 
 from vllm_omni.core.sched.output import OmniSchedulerOutput
-from vllm_omni.distributed.omni_connectors.factory import OmniConnectorFactory
-from vllm_omni.distributed.omni_connectors.transfer_manager.chunk_transfer_manager import OmniChunkTransferManager
-from vllm_omni.distributed.omni_connectors.utils.config import ConnectorSpec
+from vllm_omni.distributed.omni_connectors.transfer_manager.base import OmniModelMode
+from vllm_omni.distributed.omni_connectors.transfer_manager.chunk_transfer_manager import (
+    OmniChunkTransferManager,
+)
 
 logger = init_logger(__name__)
 
@@ -65,17 +66,12 @@ def __init__(self, *args, **kwargs):
         # 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
         self.chunk_manager = 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.chunk_manager = OmniChunkTransferManager(self.omni_connector)
+        if getattr(model_config, "async_chunk", False):
+            self.chunk_manager = OmniChunkTransferManager(
+                model_config, OmniModelMode.MODE_AR)
 
+        if self.chunk_manager:
             custom_process_next_stage_input_func = getattr(
                 self.vllm_config.model_config, "custom_process_next_stage_input_func", None
             )
@@ -192,15 +188,8 @@ def schedule(self) -> SchedulerOutput:  # type: ignore[override]
                 new_list.append(omni_nr)
 
             scheduler_output.scheduled_new_reqs = new_list  # type: ignore[assignment]
-            cached_reqs = scheduler_output.scheduled_cached_reqs
-            if not hasattr(cached_reqs, "additional_information"):
-                cached_reqs.additional_information = {}
-            for req_id in cached_reqs.req_ids:
-                request = self.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
             if self.chunk_manager:
-                self.chunk_manager.filter_scheduler_output(scheduler_output)
+                self.chunk_manager.filter_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:

vllm_omni/core/sched/omni_generation_scheduler.py

@@ -16,26 +16,21 @@
 from vllm.v1.spec_decode.metrics import SpecDecodingStats
 
 from vllm_omni.core.sched.output import OmniCachedRequestData, OmniNewRequestData
-from vllm_omni.distributed.omni_connectors.factory import OmniConnectorFactory
-from vllm_omni.distributed.omni_connectors.transfer_manager.chunk_transfer_manager import OmniChunkTransferManager
-from vllm_omni.distributed.omni_connectors.utils.config import ConnectorSpec
+from vllm_omni.distributed.omni_connectors.transfer_manager.base import OmniModelMode
+from vllm_omni.distributed.omni_connectors.transfer_manager.chunk_transfer_manager import (
+    OmniChunkTransferManager,
+)
 from vllm_omni.outputs import OmniModelRunnerOutput
 
 
 class OmniGenerationScheduler(VLLMScheduler):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         model_config = self.vllm_config.model_config
-        self.omni_connector = None
         self.chunk_manager = 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.chunk_manager = OmniChunkTransferManager(self.omni_connector)
+        if getattr(model_config, "async_chunk", False):
+            self.chunk_manager = OmniChunkTransferManager(
+                model_config, OmniModelMode.MODE_GENERATION)
 
         self.stage_id = getattr(self.vllm_config.model_config, "stage_id", None)
 
@@ -76,7 +71,7 @@ def schedule(self) -> SchedulerOutput:
             # 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_manager is None and request.status == RequestStatus.FINISHED_STOPPED
             ):
                 already_finished_reqs.add(request)
                 req_index += 1
@@ -115,7 +110,7 @@ def schedule(self) -> SchedulerOutput:
             request = self.waiting.peek_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_manager is None and request.status == RequestStatus.FINISHED_STOPPED
             ):
                 # Pop the finished request from waiting queue and don't schedule it
                 self.waiting.pop_request()
@@ -367,7 +362,7 @@ def update_from_output(
 
             # 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_manager 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

vllm_omni/distributed/omni_connectors/transfer_manager/base.py

@@ -1,23 +1,35 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+import enum
 import threading
 import time
+from typing import Any
 
 from ..utils.logging import get_connector_logger
 
 logger = get_connector_logger(__name__)
 
 
+class OmniModelMode(enum.Enum):
+    # Omni AR Model
+    MODE_AR = "ar"
+
+    # Omni Generation Model
+    MODE_GENERATION = "generate"
+
+
 class OmniTransferManagerBase:
     """Base class for managing asynchronous 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, connector):
-        self.connector = connector
+    def __init__(self, config: Any, mode: 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
@@ -37,6 +49,10 @@ def __init__(self, connector):
         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():

vllm_omni/distributed/omni_connectors/transfer_manager/chunk_transfer_manager.py

@@ -7,18 +7,21 @@
 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 OmniTransferManagerBase
+from .base import OmniModelMode, OmniTransferManagerBase
 
 logger = get_connector_logger(__name__)
 
 
 class OmniChunkTransferManager(OmniTransferManagerBase):
     """Manages asynchronous retrieval and storage of data chunks via OmniConnector."""
 
-    def __init__(self, connector):
-        super().__init__(connector)
-
+    def __init__(self, model_config: Any, mode: OmniModelMode):
+        self.connector = self.create_connector(model_config)
+        self.model_mode = mode
+        super().__init__(model_config)
         # State specific to Chunk management
         self.put_requests: dict[str, int] = defaultdict(int)
         self.get_requests: dict[str, int] = defaultdict(int)
@@ -32,6 +35,23 @@ def __init__(self, connector):
         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(self, request):
         """Request to retrieve a chunk of data for a specific request.
 
@@ -115,7 +135,7 @@ def _process_single_recv(self, req_id: str):
             self.get_requests[req_id] += 1
             req = self._pending_load_reqs[req_id]
 
-            if stage_id != 2:
+            if self.mode == OmniModelMode.MODE_AR:
                 self._update_request_payload(external_req_id, payload_data)
                 req.additional_information = payload_data
                 if payload_data.get("finished"):
@@ -211,12 +231,33 @@ def restore_queues(self, waiting_queue: Any, running_queue: list[Request]) -> No
             running_queue.extend(self.waiting_for_chunk_running_requests)
         self.waiting_for_chunk_running_requests = deque()
 
-    def filter_scheduler_output(self, scheduler_output: Any) -> None:
+    def filter_scheduler_output(
+        self,
+        scheduler_output: Any,
+        requests: dict[str, Request] | None = None,
+    ) -> None:
         """
-        Clean up ready chunks from scheduler output.
+        Add addtitional 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,
@@ -232,6 +273,7 @@ def _process_chunk_queue(
                 # Access finished_requests from self instead of connector
                 if request.request_id in self.finished_requests:
                     request.additional_information = {}
+                    self.finished_requests.remove(request.request_id)
                     continue
                 self.load(request)
                 request.status = RequestStatus.WAITING_FOR_CHUNK