[Feat] DiffusionEngine Support async batch infer (#2729)

Signed-off-by: Semmer <semmer@live.cn>
Signed-off-by: jader <yjader@foxmail.com>
Signed-off-by: asukaqaq-s <1311722138@qq.com>
Co-authored-by: asukaqaq-s <1311722138@qq.com>
Co-authored-by: jader <yjader@foxmail.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Co-authored-by: Gao Han <hgaoaf@connect.ust.hk>

vllm_omni/diffusion/sched/base_scheduler.py

@@ -4,6 +4,7 @@
 from __future__ import annotations
 
 from collections import deque
+from dataclasses import fields
 
 from vllm.logger import init_logger
 
@@ -15,11 +16,23 @@ from vllm_omni.diffusion.sched.interface import (
     DiffusionRequestStatus,
     DiffusionSchedulerOutput,
     NewRequestData,
+    SamplingParamsKey,
     SchedulerInterface,
 )
 
 logger = init_logger(__name__)
 
+_KEY_FIELD_NAMES = frozenset(f.name for f in fields(SamplingParamsKey))
+
+
+def get_sampling_params_key(request: OmniDiffusionRequest) -> SamplingParamsKey | None:
+    """Build a batch-compatibility key from the request's sampling params."""
+    if len(request.prompts) != 1:
+        return None
+
+    sampling = request.sampling_params
+    return SamplingParamsKey(**{name: getattr(sampling, name) for name in _KEY_FIELD_NAMES})
+
 
 class _BaseScheduler(SchedulerInterface):
     """Shared queue/state bookkeeping for diffusion schedulers."""
@@ -31,8 +44,9 @@ class _BaseScheduler(SchedulerInterface):
         self._step_id: int = 0
         self._waiting: deque[str] = deque()
         self._running: list[str] = []
+        self._running_sampling_params_key: SamplingParamsKey | None = None
         self._finished_req_ids: set[str] = set()
-        self._max_batch_size: int = 1
+        self.max_num_running_reqs: int = 1
 
     def initialize(self, od_config: OmniDiffusionConfig) -> None:
         self.od_config = od_config
@@ -41,11 +55,13 @@ class _BaseScheduler(SchedulerInterface):
         self._step_id = 0
         self._waiting.clear()
         self._running.clear()
+        self._running_sampling_params_key = None
         self._finished_req_ids.clear()
-        # The current DiffusionEngine execution mode does not support real
-        # request batching well, so we keep this fixed at 1 for now.
-        # TODO: Add support for multiple concurrent requests
-        self.max_num_running_reqs = 1
+        max_num_seqs = getattr(od_config, "max_num_seqs", 1)
+        try:
+            self.max_num_running_reqs = max(1, int(max_num_seqs))
+        except (TypeError, ValueError):
+            self.max_num_running_reqs = 1
         self._reset_scheduler_state()
 
     def add_request(self, request: OmniDiffusionRequest) -> str:
@@ -53,7 +69,7 @@ class _BaseScheduler(SchedulerInterface):
         return self._add_request_with_sched_req_id(sched_req_id, request)
 
     def _add_request_with_sched_req_id(self, sched_req_id: str, request: OmniDiffusionRequest) -> str:
-        state = DiffusionRequestState(sched_req_id=sched_req_id, req=request)
+        state = self._make_request_state(sched_req_id, request)
         self._request_states[sched_req_id] = state
         self._register_request_ids(request.request_ids, sched_req_id)
         self._waiting.append(sched_req_id)
@@ -82,6 +98,8 @@ class _BaseScheduler(SchedulerInterface):
 
             self._waiting.popleft()
             was_new_request = state.status == DiffusionRequestStatus.WAITING
+            if not self._running:
+                self._running_sampling_params_key = state.sampling_params_key
             state.status = DiffusionRequestStatus.RUNNING
             self._running.append(sched_req_id)
             if was_new_request:
@@ -124,6 +142,8 @@ class _BaseScheduler(SchedulerInterface):
             return False
         if sched_req_id in self._running:
             self._running.remove(sched_req_id)
+            if not self._running:
+                self._running_sampling_params_key = None
             self._waiting.appendleft(sched_req_id)
             self._request_states[sched_req_id].status = DiffusionRequestStatus.PREEMPTED
             return True
@@ -140,6 +160,7 @@ class _BaseScheduler(SchedulerInterface):
         self._request_id_to_sched_req_id.clear()
         self._waiting.clear()
         self._running.clear()
+        self._running_sampling_params_key = None
         self._finished_req_ids.clear()
         self._reset_scheduler_state()
 
@@ -169,6 +190,8 @@ class _BaseScheduler(SchedulerInterface):
 
         if running_to_remove:
             self._running = [sched_req_id for sched_req_id in self._running if sched_req_id not in running_to_remove]
+            if not self._running:
+                self._running_sampling_params_key = None
         if waiting_to_remove:
             self._waiting = deque(
                 sched_req_id for sched_req_id in self._waiting if sched_req_id not in waiting_to_remove
@@ -208,9 +231,26 @@ class _BaseScheduler(SchedulerInterface):
     def _pop_extra_request_state(self, sched_req_id: str) -> None:
         """Remove subclass-owned per-request state before popping request state."""
 
+    def _make_request_state(self, sched_req_id: str, request: OmniDiffusionRequest) -> DiffusionRequestState:
+        return DiffusionRequestState(
+            sched_req_id=sched_req_id,
+            req=request,
+            sampling_params_key=get_sampling_params_key(request),
+        )
+
     def _can_schedule_waiting(self, state: DiffusionRequestState) -> bool:
-        del state
-        return True
+        if not self._running:
+            return True
+
+        current_key = self._current_sampling_params_key()
+        return current_key is not None and current_key == state.sampling_params_key
+
+    def _current_sampling_params_key(self) -> SamplingParamsKey | None:
+        if self._running_sampling_params_key is not None or not self._running:
+            return self._running_sampling_params_key
+        state = self._request_states.get(self._running[0])
+        self._running_sampling_params_key = None if state is None else state.sampling_params_key
+        return self._running_sampling_params_key
 
     def _register_request_ids(self, request_ids: list[str], sched_req_id: str) -> None:
         for request_id in request_ids:

vllm_omni/diffusion/sched/interface.py

@@ -38,6 +38,34 @@ class DiffusionRequestStatus(enum.IntEnum):
         return status >= DiffusionRequestStatus.FINISHED_COMPLETED
 
 
+@dataclass(frozen=True, eq=True)
+class SamplingParamsKey:
+    """Batch-compatibility key derived from ``OmniDiffusionSamplingParams``.
+
+    Only requests with the same key can be batched together.
+    Fields not included here are treated as request-local and do not
+    participate in the current homogeneous batching policy.
+    """
+
+    # Spatial / temporal shape.
+    height: int | None = None
+    width: int | None = None
+    num_frames: int = 1
+    resolution: int | str | None = None
+    fps: int | None = None
+    frame_rate: float | None = None
+    boundary_ratio: float | None = None
+
+    # CFG / guidance.
+    do_classifier_free_guidance: bool = False
+    guidance_scale: float = 0.0
+    guidance_scale_provided: bool = False
+    guidance_scale_2: float | None = None
+    guidance_rescale: float = 0.0
+    true_cfg_scale: float | None = None
+    cfg_normalize: bool = False
+
+
 @dataclass
 class DiffusionRequestState:
     """Scheduler-owned state for one queued OmniDiffusionRequest."""
@@ -47,6 +75,7 @@ class DiffusionRequestState:
     # TODO: Align this with OmniDiffusionRequest.request_ids once scheduler batching is supported.
     sched_req_id: str
     req: OmniDiffusionRequest
+    sampling_params_key: SamplingParamsKey | None = None
     status: DiffusionRequestStatus = DiffusionRequestStatus.WAITING
     error: str | None = None
 
@@ -81,7 +110,7 @@ class CachedRequestData:
 class DiffusionSchedulerOutput:
     """Output of a single scheduling cycle."""
 
-    step_id: int
+    step_id: int  # global step index
     scheduled_new_reqs: list[NewRequestData]
     scheduled_cached_reqs: CachedRequestData
     finished_req_ids: set[str]

vllm_omni/diffusion/sched/request_scheduler.py

@@ -32,11 +32,12 @@ class RequestScheduler(_BaseScheduler):
 
         terminal_statuses: dict[str, DiffusionRequestStatus] = {}
         terminal_errors: dict[str, str | None] = {}
-        result = output.result
         for sched_req_id in scheduled_req_ids:
             state = self._request_states.get(sched_req_id)
             if state is None or state.is_finished():
                 continue
+            req_output = output.get_req_output(sched_req_id)
+            result = req_output.result if req_output is not None else None
             if result is None:
                 terminal_statuses[sched_req_id] = DiffusionRequestStatus.FINISHED_ERROR
                 terminal_errors[sched_req_id] = "No output result"

vllm_omni/diffusion/sched/step_scheduler.py

@@ -72,19 +72,29 @@ class StepScheduler(_BaseScheduler):
 
         terminal_statuses: dict[str, DiffusionRequestStatus] = {}
         terminal_errors: dict[str, str | None] = {}
-        output_error = output.result.error if output.result is not None else None
         for sched_req_id in scheduled_req_ids:
             state = self._request_states.get(sched_req_id)
             progress = self._request_progress.get(sched_req_id)
             if state is None or progress is None or state.is_finished():
                 continue
+            req_output = output.get_req_output(sched_req_id)
+            if req_output is None:
+                logger.warning(
+                    "No RunnerOutput for request %s, treating as error",
+                    sched_req_id,
+                )
+                terminal_statuses[sched_req_id] = DiffusionRequestStatus.FINISHED_ERROR
+                terminal_errors[sched_req_id] = "No output for request"
+                continue
 
+            req_result = req_output.result
+            output_error = req_result.error if req_result is not None else None
             if output_error is not None:
                 terminal_statuses[sched_req_id] = DiffusionRequestStatus.FINISHED_ERROR
                 terminal_errors[sched_req_id] = output_error
                 continue
 
-            if output.step_index is None:
+            if req_output.step_index is None:
                 logger.warning(
                     "Received RunnerOutput with no step_index for request %s, treating as error",
                     sched_req_id,
@@ -94,9 +104,9 @@ class StepScheduler(_BaseScheduler):
                 continue
 
             # We assume that the decoding stage is executed immediately after the denoising stage completes.
-            progress.current_step = output.step_index
-            state.req.sampling_params.step_index = output.step_index
-            if output.finished:
+            progress.current_step = req_output.step_index
+            state.req.sampling_params.step_index = req_output.step_index
+            if req_output.finished:
                 terminal_statuses[sched_req_id] = DiffusionRequestStatus.FINISHED_COMPLETED
                 terminal_errors[sched_req_id] = None
             else: