[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/diffusion_engine.py

@@ -3,6 +3,7 @@
 
 from __future__ import annotations
 
+import asyncio
 import inspect
 import queue
 import threading
@@ -30,7 +31,7 @@ from vllm_omni.diffusion.registry import (
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.sched import RequestScheduler, SchedulerInterface, StepScheduler
 from vllm_omni.diffusion.sched.interface import DiffusionRequestStatus
-from vllm_omni.diffusion.worker.utils import RunnerOutput
+from vllm_omni.diffusion.worker.utils import BatchRunnerOutput, RunnerOutput
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniTextPrompt
 from vllm_omni.outputs import OmniRequestOutput
 
@@ -99,7 +100,17 @@ class DiffusionEngine:
             StepScheduler() if self.step_execution else RequestScheduler()
         )
         self.scheduler.initialize(od_config)
+        if self.scheduler.max_num_running_reqs > 1 and not self.step_execution:
+            max_num_seqs = self.scheduler.max_num_running_reqs
+            self.scheduler.max_num_running_reqs = 1
+            logger.warning(f"Non-stepwise-execution does not support max-num-seqs={max_num_seqs}, set it to 1.")
+        self.main_loop: asyncio.AbstractEventLoop | None = None
+        self.stop_event: threading.Event | None = None
+        self._loop_started = False
+        self._init_lock = asyncio.Lock()
         self._rpc_lock = threading.RLock()
+        self._cv = threading.Condition(self._rpc_lock)
+        self._out_queue: dict[str, asyncio.Future] = {}
         self.abort_queue: queue.Queue[str] = queue.Queue()
         self.execute_fn = self.executor.execute_step if self.step_execution else self.executor.execute_request
 
@@ -110,7 +121,24 @@ class DiffusionEngine:
             self.close()
             raise e
 
-    def step(self, request: OmniDiffusionRequest) -> list[OmniRequestOutput]:
+    async def _check_and_start_background_loop(self):
+        if self._loop_started:
+            return
+
+        async with self._init_lock:
+            # double check, in case of lock queue issue
+            if self._loop_started:
+                return
+
+            self.main_loop = asyncio.get_running_loop()
+            self.stop_event = threading.Event()
+            self.worker_thread = threading.Thread(target=self._busy_loop)
+            self.worker_thread.start()
+            self._loop_started = True
+
+    async def step(self, request: OmniDiffusionRequest) -> list[OmniRequestOutput]:
+        await self._check_and_start_background_loop()
+
         diffusion_engine_start_time = time.perf_counter()
 
         # Apply pre-processing if available
@@ -122,7 +150,7 @@ class DiffusionEngine:
             logger.info(f"Pre-processing completed in {preprocess_time:.4f} seconds")
 
         exec_start_time = time.perf_counter()
-        output = self.add_req_and_wait_for_response(request)
+        output = await self.async_add_req_and_wait_for_response(request)
         exec_total_time = time.perf_counter() - exec_start_time
 
         if output.aborted:
@@ -336,6 +364,62 @@ class DiffusionEngine:
 
             return results
 
+    def _busy_loop(self):
+        while not self.stop_event.is_set():
+            self._process_aborts_queue()
+
+            with self._cv:
+                while not self.scheduler.has_requests() and not self.stop_event.is_set():
+                    self._cv.wait(timeout=1.0)
+
+                if self.stop_event.is_set():
+                    break
+
+                sched_output = self.scheduler.schedule()
+
+            if sched_output.is_empty:
+                self._handle_finished_requests(sched_output.finished_req_ids, None)
+                continue
+
+            try:
+                runner_output = self.execute_fn(sched_output)
+            except Exception as exc:
+                logger.error(
+                    "Execution failed for diffusion requests %s", sched_output.scheduled_req_ids, exc_info=True
+                )
+                runner_output = BatchRunnerOutput.from_list(
+                    [
+                        RunnerOutput(
+                            req_id=req_id,
+                            step_index=None,
+                            finished=True,
+                            result=DiffusionOutput(error=str(exc)),
+                        )
+                        for req_id in sched_output.scheduled_req_ids
+                    ]
+                )
+
+            self._process_aborts_queue()
+            finished_req_ids = self.scheduler.update_from_output(sched_output, runner_output)
+            self._handle_finished_requests(finished_req_ids, runner_output)
+
+    def _handle_finished_requests(
+        self,
+        finished_ids: set[str],
+        runner_output: RunnerOutput | None = None,
+        missing_result_error: str = "Diffusion execution finished without a final output",
+    ):
+        for rid in finished_ids:
+            if rid in self._out_queue:
+                fut = self._out_queue.pop(rid)
+                if runner_output is not None:
+                    _output = runner_output.get_req_output(rid)
+                else:
+                    _output = None
+                out = self._finalize_finished_request(rid, _output, missing_result_error)
+                if not fut.done():
+                    self.main_loop.call_soon_threadsafe(fut.set_result, out)
+
     @staticmethod
     def make_engine(
         config: OmniDiffusionConfig,
@@ -351,6 +435,30 @@ class DiffusionEngine:
         """
         return DiffusionEngine(config, scheduler=scheduler)
 
+    def add_request(self, request: OmniDiffusionRequest) -> str:
+        fut = self.main_loop.create_future()
+        with self._cv:
+            sched_req_id = self.scheduler.add_request(request)
+            self._out_queue[sched_req_id] = fut
+            self._cv.notify_all()
+
+        return sched_req_id
+
+    async def get_result(self, request_id: str) -> DiffusionOutput:
+        fut = self._out_queue.get(request_id)
+
+        if fut is None:
+            raise RuntimeError(f"Request {request_id} not found in output queue.")
+        try:
+            return await fut
+        except Exception as e:
+            logger.error(f"Wait for response failed: {e}")
+            raise
+
+    async def async_add_req_and_wait_for_response(self, request: OmniDiffusionRequest) -> DiffusionOutput:
+        sched_req_id = self.add_request(request)
+        return await self.get_result(sched_req_id)
+
     def add_req_and_wait_for_response(self, request: OmniDiffusionRequest) -> DiffusionOutput:
         with self._rpc_lock:
             target_sched_req_id = self.scheduler.add_request(request)
@@ -366,10 +474,8 @@ class DiffusionEngine:
                         raise RuntimeError("Diffusion scheduler has no runnable requests.")
                     continue
 
-                # NOTE: add_req_and_wait_for_response() is synchronous, and
-                # the scheduler currently enforces _max_batch_size = 1 (see
-                # vllm_omni/diffusion/sched/base_scheduler.py), so we directly
-                # take the single scheduled request here.
+                # NOTE: add_req_and_wait_for_response() is synchronous, will be only called
+                # within _dummy_run, only one request will be scheduled
                 sched_req_id = sched_output.scheduled_req_ids[0]
                 try:
                     runner_output = self.execute_fn(sched_output)
@@ -387,7 +493,12 @@ class DiffusionEngine:
                 self._process_aborts_queue()
 
                 finished_req_ids = self.scheduler.update_from_output(sched_output, runner_output)
+
+                # sync func should receive one result
+                if not isinstance(runner_output, RunnerOutput) and not len(runner_output) == 1:
+                    raise ValueError("Sync func should receive one result at one time")
                 if target_sched_req_id in finished_req_ids:
+                    runner_output = runner_output.get_req_output(target_sched_req_id)
                     return self._finalize_finished_request(
                         target_sched_req_id,
                         runner_output=runner_output,
@@ -510,13 +621,26 @@ class DiffusionEngine:
     def close(self) -> None:
         if hasattr(self, "scheduler"):
             self.scheduler.close()
+        if hasattr(self, "stop_event") and self.stop_event is not None:
+            with self._cv:
+                self.stop_event.set()
+                self._cv.notify_all()
+        if hasattr(self, "worker_thread") and self.worker_thread.is_alive():
+            self.worker_thread.join(timeout=10)
+            if self.worker_thread.is_alive():
+                logger.warning("Worker thread did not terminate within 10s, resources may not be released.")
+            else:
+                self._loop_started = False
         if hasattr(self, "executor"):
             self.executor.shutdown()
 
     def abort(self, request_id: str | Iterable[str]) -> None:
         request_ids = [request_id] if isinstance(request_id, str) else list(request_id)
-        for req_id in request_ids:
-            self.abort_queue.put(req_id)
+
+        with self._cv:
+            for req_id in request_ids:
+                self.abort_queue.put(req_id)
+            self._cv.notify_all()
 
     def _process_aborts_queue(self) -> None:
         if self.abort_queue.empty():

vllm_omni/diffusion/executor/abstract.py

@@ -10,7 +10,7 @@ from vllm_omni.diffusion.request import OmniDiffusionRequest
 
 if TYPE_CHECKING:
     from vllm_omni.diffusion.sched.interface import DiffusionSchedulerOutput
-    from vllm_omni.diffusion.worker.utils import RunnerOutput
+    from vllm_omni.diffusion.worker.utils import BaseRunnerOutput
 
 
 class DiffusionExecutor(ABC):
@@ -74,12 +74,12 @@ class DiffusionExecutor(ABC):
         pass
 
     @abstractmethod
-    def execute_request(self, scheduler_output: DiffusionSchedulerOutput) -> RunnerOutput:
+    def execute_request(self, scheduler_output: DiffusionSchedulerOutput) -> BaseRunnerOutput:
         """Execute request-mode work from a scheduler output."""
         pass
 
     @abstractmethod
-    def execute_step(self, scheduler_output: DiffusionSchedulerOutput) -> RunnerOutput:
+    def execute_step(self, scheduler_output: DiffusionSchedulerOutput) -> BaseRunnerOutput:
         """Execute step-mode work from a scheduler output."""
         pass
 

vllm_omni/diffusion/executor/multiproc_executor.py

@@ -22,7 +22,7 @@ from vllm_omni.diffusion.worker import WorkerProc
 
 if TYPE_CHECKING:
     from vllm_omni.diffusion.sched.interface import DiffusionSchedulerOutput
-    from vllm_omni.diffusion.worker.utils import RunnerOutput
+    from vllm_omni.diffusion.worker.utils import BaseRunnerOutput
 
 logger = init_logger(__name__)
 
@@ -273,7 +273,7 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
             logger.error(f"Generate call failed: {e}")
             raise
 
-    def execute_request(self, scheduler_output: DiffusionSchedulerOutput) -> RunnerOutput:
+    def execute_request(self, scheduler_output: DiffusionSchedulerOutput) -> BaseRunnerOutput:
         """Adapt request-mode scheduler output to worker execute_model RPC."""
         from vllm_omni.diffusion.worker.utils import RunnerOutput
 
@@ -301,9 +301,9 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
             result=result,
         )
 
-    def execute_step(self, scheduler_output: DiffusionSchedulerOutput) -> RunnerOutput:
+    def execute_step(self, scheduler_output: DiffusionSchedulerOutput) -> BaseRunnerOutput:
         """Forward step-mode scheduler output to worker execute_stepwise RPC."""
-        from vllm_omni.diffusion.worker.utils import RunnerOutput
+        from vllm_omni.diffusion.worker.utils import BaseRunnerOutput, RunnerOutput
 
         self._ensure_open()
         result = self.collective_rpc(
@@ -313,7 +313,7 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
             exec_all_ranks=True,
         )
 
-        if isinstance(result, RunnerOutput):
+        if isinstance(result, BaseRunnerOutput):
             return result
         # TODO: Remove this fallback; DiffusionOutput cannot faithfully represent
         # failed multi-request step batches.

vllm_omni/diffusion/inline_stage_diffusion_client.py

@@ -108,8 +108,7 @@ class InlineStageDiffusionClient:
                 kv_sender_info=kv_sender_info,
             )
 
-            loop = asyncio.get_running_loop()
-            results = await loop.run_in_executor(self._executor, self._engine.step, request)
+            results = await self._engine.step(request)
             result = results[0]
             if not result.request_id:
                 result.request_id = request_id
@@ -161,8 +160,7 @@ class InlineStageDiffusionClient:
                 kv_sender_info=kv_sender_info,
             )
 
-            loop = asyncio.get_running_loop()
-            results = await loop.run_in_executor(self._executor, self._engine.step, request)
+            results = await self._engine.step(request)
 
             all_images: list = []
             merged_mm: dict[str, Any] = {}

vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py

@@ -44,6 +44,7 @@ from vllm_omni.diffusion.utils.size_utils import (
 from vllm_omni.diffusion.utils.tf_utils import get_transformer_config_kwargs
 
 if TYPE_CHECKING:
+    from vllm_omni.diffusion.worker.input_batch import InputBatch
     from vllm_omni.diffusion.worker.utils import DiffusionRequestState
 
 from vllm_omni.model_executor.model_loader.weight_utils import (
@@ -864,54 +865,48 @@ class QwenImagePipeline(nn.Module, QwenImageCFGParallelMixin, DiffusionPipelineP
 
     def denoise_step(
         self,
-        state: "DiffusionRequestState",
+        input_batch: "InputBatch",
         **kwargs: Any,
     ) -> torch.Tensor | None:
-        """One denoise step: read from *state*, delegate to CFGParallelMixin.
+        """One denoise step: read from *input_batch*, delegate to CFGParallelMixin.
 
         Reuses ``predict_noise_maybe_with_cfg`` so that CFG-parallel,
         sequential-CFG, and no-CFG paths are handled identically to
         ``diffuse()``.
         """
+        del kwargs
         if self.interrupt:
             return None
 
-        t = state.current_timestep
+        t = input_batch.timesteps
         self._current_timestep = t
-        self.transformer.do_true_cfg = state.do_true_cfg
-
-        # Normalize timestep to [batch_size] tensor
-        if not torch.is_tensor(t):
-            t = torch.tensor(t, device=state.latents.device, dtype=state.latents.dtype)
+        self.transformer.do_true_cfg = input_batch.do_true_cfg
 
         positive_kwargs, negative_kwargs, output_slice = self._build_denoise_kwargs(
-            latents=state.latents,
+            latents=input_batch.latents,
             timestep=t,
-            guidance=state.guidance,
-            prompt_embeds=state.prompt_embeds,
-            prompt_embeds_mask=state.prompt_embeds_mask,
-            img_shapes=state.img_shapes,
-            txt_seq_lens=state.txt_seq_lens,
-            do_true_cfg=state.do_true_cfg,
-            negative_prompt_embeds=state.negative_prompt_embeds,
-            negative_prompt_embeds_mask=state.negative_prompt_embeds_mask,
-            negative_txt_seq_lens=state.negative_txt_seq_lens,
-            image_latents=state.sampling.image_latent,
+            guidance=input_batch.guidance,
+            prompt_embeds=input_batch.prompt_embeds,
+            prompt_embeds_mask=input_batch.prompt_embeds_mask,
+            img_shapes=input_batch.img_shapes,
+            txt_seq_lens=input_batch.txt_seq_lens,
+            do_true_cfg=input_batch.do_true_cfg,
+            negative_prompt_embeds=input_batch.negative_prompt_embeds,
+            negative_prompt_embeds_mask=input_batch.negative_prompt_embeds_mask,
+            negative_txt_seq_lens=input_batch.negative_txt_seq_lens,
+            image_latents=input_batch.image_latents,
             extra_transformer_kwargs={
                 "attention_kwargs": self.attention_kwargs,
                 "return_dict": False,
             },
         )
 
-        true_cfg_scale = state.sampling.true_cfg_scale or 4.0
-        cfg_normalize = state.sampling.cfg_normalize
-
         return self.predict_noise_maybe_with_cfg(
-            state.do_true_cfg,
-            true_cfg_scale,
+            input_batch.do_true_cfg,
+            input_batch.true_cfg_scale,
             positive_kwargs,
             negative_kwargs,
-            cfg_normalize,
+            input_batch.cfg_normalize,
             output_slice,
         )
 

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:

vllm_omni/diffusion/stage_diffusion_proc.py

@@ -103,8 +103,7 @@ class StageDiffusionProc:
             kv_sender_info=kv_sender_info,
         )
 
-        loop = asyncio.get_running_loop()
-        results = await loop.run_in_executor(self._executor, self._engine.step, request)
+        results = await self._engine.step(request)
         result = results[0]
         if not result.request_id:
             result.request_id = request_id
@@ -134,8 +133,7 @@ class StageDiffusionProc:
             kv_sender_info=kv_sender_info,
         )
 
-        loop = asyncio.get_running_loop()
-        results = await loop.run_in_executor(self._executor, self._engine.step, request)
+        results = await self._engine.step(request)
 
         # Merge per-prompt results into a single combined output.
         all_images: list = []

vllm_omni/diffusion/worker/diffusion_model_runner.py

@@ -14,6 +14,7 @@ import copy
 import time
 from collections.abc import Iterable
 from contextlib import nullcontext
+from typing import Any
 
 import torch
 from torch.profiler import record_function
@@ -32,7 +33,8 @@ from vllm_omni.diffusion.offloader import get_offload_backend
 from vllm_omni.diffusion.registry import _NO_CACHE_ACCELERATION
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.sched.interface import DiffusionSchedulerOutput
-from vllm_omni.diffusion.worker.utils import DiffusionRequestState, RunnerOutput
+from vllm_omni.diffusion.worker.input_batch import InputBatch, scatter_latents
+from vllm_omni.diffusion.worker.utils import BatchRunnerOutput, DiffusionRequestState, RunnerOutput
 from vllm_omni.distributed.omni_connectors.kv_transfer_manager import OmniKVTransferManager
 from vllm_omni.platforms import current_omni_platform
 from vllm_omni.worker.omni_connector_model_runner_mixin import OmniConnectorModelRunnerMixin
@@ -297,50 +299,107 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
         """Return whether current pipeline supports step execution."""
         return self.pipeline is not None and supports_step_execution(self.pipeline)
 
-    def _update_states(self, scheduler_output: DiffusionSchedulerOutput) -> tuple[DiffusionRequestState, bool]:
+    def _update_states(
+        self, scheduler_output: DiffusionSchedulerOutput
+    ) -> tuple[list[DiffusionRequestState], list[str]]:
         """Step-before update: cleanup finished requests and get/create one running state."""
         for req_id in scheduler_output.finished_req_ids:
             self.state_cache.pop(req_id, None)
 
-        if scheduler_output.num_scheduled_reqs != 1:
-            raise ValueError(
-                "Step mode currently supports batch_size=1, "
-                f"but got {scheduler_output.num_scheduled_reqs} scheduled requests."
-            )
+        resolved: list[DiffusionRequestState] = []
+        new_req_id: list[str] = []
+        try:
+            # process new requests
+            for sched_new_req in scheduler_output.scheduled_new_reqs:
+                # new_req_data = scheduler_output.scheduled_new_reqs[0]
+                req_id = sched_new_req.sched_req_id
+                req = sched_new_req.req
+                new_req_id.append(req_id)
+                if req_id in self.state_cache:
+                    raise ValueError(f"Received duplicate new-request payload for cached request {req_id}.")
+                request_ids = req.request_ids or [req_id]
+                if len(request_ids) != len(req.prompts):
+                    raise ValueError(
+                        f"request_ids length ({len(request_ids)}) does not match prompts length ({len(req.prompts)})"
+                    )
+                new_state = DiffusionRequestState(
+                    req_id=req_id,
+                    sampling=copy.deepcopy(req.sampling_params),
+                    prompts=req.prompts,
+                )
+                self.state_cache[req_id] = new_state
+                resolved.append(new_state)
 
-        if scheduler_output.scheduled_new_reqs:
-            new_req_data = scheduler_output.scheduled_new_reqs[0]
-            req_id = new_req_data.sched_req_id
-            req = new_req_data.req
-            if req_id in self.state_cache:
-                raise ValueError(f"Received duplicate new-request payload for cached request {req_id}.")
-        else:
-            req_id = scheduler_output.scheduled_cached_reqs.sched_req_ids[0]
-            state = self.state_cache.get(req_id)
-            if state is None:
-                raise ValueError(f"Missing cached state for request {req_id}.")
-            return state, False
+            # process cached requests
+            for req_id in scheduler_output.scheduled_cached_reqs.sched_req_ids:
+                state = self.state_cache.get(req_id)
+                if state is None:
+                    raise ValueError(f"Missing cached state for request {req_id}.")
+                resolved.append(state)
+        except Exception:
+            for req_id in new_req_id:
+                self.state_cache.pop(req_id, None)
+            raise
 
-        request_ids = req.request_ids or [req_id]
-        if len(request_ids) != len(req.prompts):
-            raise ValueError(
-                f"request_ids length ({len(request_ids)}) does not match prompts length ({len(req.prompts)})"
-            )
+        return resolved, new_req_id
 
-        state = DiffusionRequestState(
-            req_id=req_id,
-            sampling=copy.deepcopy(req.sampling_params),
-            prompts=req.prompts,
+    def _prepare_batch_inputs(self, states: list[DiffusionRequestState], new_request_ids: list[str]) -> InputBatch:
+        # process new reqs
+        for state in states:
+            if state.req_id in new_request_ids:
+                # set generator
+                if state.sampling.generator is None and state.sampling.seed is not None:
+                    if state.sampling.generator_device is not None:
+                        gen_device = state.sampling.generator_device
+                    elif self.device.type == "cpu":
+                        gen_device = "cpu"
+                    else:
+                        gen_device = self.device
+                    state.sampling.generator = torch.Generator(device=gen_device).manual_seed(state.sampling.seed)
+                # encode
+                self.pipeline.prepare_encode(state)
+
+        input_batch = InputBatch.make_batch(
+            states,
+            cached_batch=getattr(self, "input_batch", None),
         )
-        self.state_cache[req_id] = state
-        return state, True
+        self.input_batch = input_batch
+        return input_batch
 
-    def _update_states_after(self, state: DiffusionRequestState, finished: bool) -> None:
+    def _update_states_after(
+        self,
+        states: list[DiffusionRequestState],
+        input_batch: InputBatch,
+        interrupted: bool = False,
+    ):
         """Step-after update: clear cached state for completed request."""
-        if finished:
-            self.state_cache.pop(state.req_id, None)
+        gathered_latents = torch.cat([state.latents for state in states], dim=0)
+        if (
+            input_batch.latents.size() == gathered_latents.size()
+            and input_batch.latents.dtype == gathered_latents.dtype
+            and input_batch.latents.device == gathered_latents.device
+        ):
+            input_batch.latents.copy_(gathered_latents)
+        else:
+            input_batch.latents = gathered_latents.clone()
 
-    def execute_stepwise(self, scheduler_output: DiffusionSchedulerOutput) -> RunnerOutput:
+        self.input_batch = input_batch
+        scatter_latents(states, input_batch)
+
+        for state in states:
+            if interrupted or state.denoise_completed:
+                self.state_cache.pop(state.req_id, None)
+
+    def _prepare_attn_metadata(self, input_batch: InputBatch) -> Any:
+        model_state = getattr(self, "model_state", None)
+        if model_state is None:
+            return {}
+        prepare_attn = getattr(model_state, "prepare_attn", None)
+        if not callable(prepare_attn):
+            return {}
+        return prepare_attn(input_batch)
+
+    def execute_stepwise(self, scheduler_output: DiffusionSchedulerOutput) -> BatchRunnerOutput:
         """Execute one step for one scheduled request and return runner output."""
         assert self.pipeline is not None, "Model not loaded. Call load_model() first."
         if not self.supports_step_mode():
@@ -354,46 +413,57 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
         use_hsdp = self.od_config.parallel_config.use_hsdp
         grad_context = torch.no_grad() if use_hsdp else torch.inference_mode()
         with grad_context:
-            state, is_new_request = self._update_states(scheduler_output)
+            states, new_request_ids = self._update_states(scheduler_output)
+            input_batch = self._prepare_batch_inputs(states, new_request_ids)
+            attn_metadata = self._prepare_attn_metadata(input_batch)
 
-            if is_new_request:
-                # TODO: support kv manager recv
-                # TODO: support cache backend
-                if state.sampling.generator is None and state.sampling.seed is not None:
-                    if state.sampling.generator_device is not None:
-                        gen_device = state.sampling.generator_device
-                    elif self.device.type == "cpu":
-                        gen_device = "cpu"
-                    else:
-                        gen_device = self.device
-                    state.sampling.generator = torch.Generator(device=gen_device).manual_seed(state.sampling.seed)
+            with set_forward_context(
+                vllm_config=self.vllm_config,
+                omni_diffusion_config=self.od_config,
+                attn_metadata=attn_metadata,
+            ):
+                noise_pred = self.pipeline.denoise_step(input_batch)
 
-            with set_forward_context(vllm_config=self.vllm_config, omni_diffusion_config=self.od_config):
-                # step0/new request: encode
-                if is_new_request:
-                    self.pipeline.prepare_encode(state)
+                runner_output_list = []
+                pipeline_interrupted = getattr(self.pipeline, "interrupt", False)
+                if noise_pred is None and pipeline_interrupted:
+                    for state in states:
+                        runner_output_list.append(
+                            RunnerOutput(
+                                req_id=state.req_id,
+                                step_index=state.step_index,
+                                finished=True,
+                                result=DiffusionOutput(error="stepwise denoise interrupted"),
+                            )
+                        )
 
-                noise_pred = self.pipeline.denoise_step(state)
-                finished = False
-
-                # In CFG parallel mode, only rank 0 gets the actual noise_pred; non-rank-0 workers receive None.
-                # A true interrupt (all ranks return None) is detected by checking self.pipeline.interrupt.
-                if noise_pred is None and getattr(self.pipeline, "interrupt", False):
-                    finished = True
-                    result = DiffusionOutput(error="stepwise denoise interrupted")
                 else:
-                    self.pipeline.step_scheduler(state, noise_pred)
-                    finished = state.denoise_completed
-                    if finished:
-                        result = self.pipeline.post_decode(state)
-                    else:
-                        result = None
+                    offset = 0
+                    for req in states:
+                        row_num = req.latents.shape[0]
+                        self.pipeline.step_scheduler(
+                            req, noise_pred[offset : offset + row_num] if noise_pred is not None else None
+                        )
+                        offset = offset + row_num
+                        if req.denoise_completed:
+                            result = self.pipeline.post_decode(req)
+                        else:
+                            result = None
+                        runner_output_list.append(
+                            RunnerOutput(
+                                req_id=req.req_id,
+                                step_index=req.step_index,
+                                finished=req.denoise_completed,
+                                result=result,
+                            )
+                        )
 
-                self._update_states_after(state, finished)
+                    if noise_pred is not None and offset != noise_pred.shape[0]:
+                        raise ValueError(
+                            f"Stepwise noise_pred consumed {offset} rows, "
+                            f"but batched noise_pred has {noise_pred.shape[0]} rows."
+                        )
 
-                return RunnerOutput(
-                    req_id=state.req_id,
-                    step_index=state.step_index,
-                    finished=finished,
-                    result=result,
-                )
+                self._update_states_after(states, input_batch, pipeline_interrupted)
+
+                return BatchRunnerOutput.from_list(runner_output_list)

vllm_omni/diffusion/worker/diffusion_worker.py

@@ -45,7 +45,7 @@ from vllm_omni.diffusion.lora.manager import DiffusionLoRAManager
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.sched.interface import DiffusionSchedulerOutput
 from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner
-from vllm_omni.diffusion.worker.utils import RunnerOutput
+from vllm_omni.diffusion.worker.utils import BaseRunnerOutput
 from vllm_omni.lora.request import LoRARequest
 from vllm_omni.platforms import current_omni_platform
 from vllm_omni.profiler import OmniTorchProfilerWrapper, create_omni_profiler
@@ -262,7 +262,7 @@ class DiffusionWorker:
             profiler.step()
         return output
 
-    def execute_stepwise(self, scheduler_output: DiffusionSchedulerOutput) -> RunnerOutput:
+    def execute_stepwise(self, scheduler_output: DiffusionSchedulerOutput) -> BaseRunnerOutput:
         """Execute one diffusion step by delegating to the model runner."""
         assert self.model_runner is not None, "Model runner not initialized"
         if self.lora_manager is not None:
@@ -834,7 +834,7 @@ class WorkerWrapperBase:
         """
         return self.worker.execute_model(reqs, od_config)
 
-    def execute_stepwise(self, scheduler_output: DiffusionSchedulerOutput) -> RunnerOutput:
+    def execute_stepwise(self, scheduler_output: DiffusionSchedulerOutput) -> BaseRunnerOutput:
         """Execute one diffusion step."""
         return self.worker.execute_stepwise(scheduler_output)
 

vllm_omni/diffusion/worker/input_batch.py

@@ -0,0 +1,740 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Diffusion input-batch structures following the MRV2-style vLLM layout.
+
+Request states remain the only persistent source of truth. Static tensors are
+normalized/padded onto the request state once, while :class:`InputBatch`
+assembles an ephemeral step-local view. Dynamic tensors are re-gathered every
+step, and step outputs are scattered back into request states by
+``scatter_latents()`` using ``idx_mapping``.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from dataclasses import dataclass
+
+import numpy as np
+import torch
+
+from vllm_omni.diffusion.worker.utils import DiffusionRequestState
+
+
+def _normalize_prompt_embeds(x: torch.Tensor) -> torch.Tensor:
+    if x.ndim == 2:
+        return x.unsqueeze(0)
+    if x.ndim == 3:
+        return x
+    raise ValueError(f"prompt_embeds must be 2D or 3D, got shape={tuple(x.shape)}")
+
+
+def _normalize_mask(x: torch.Tensor) -> torch.Tensor:
+    if x.ndim == 1:
+        x = x.unsqueeze(0)
+    elif x.ndim != 2:
+        raise ValueError(f"prompt mask must be 1D or 2D, got shape={tuple(x.shape)}")
+    if x.dtype != torch.bool:
+        x = x != 0
+    return x
+
+
+def _pad_prompt_embeds(x: torch.Tensor, target_seq_len: int) -> torch.Tensor:
+    x = _normalize_prompt_embeds(x)
+    bsz, seq_len, hidden = x.shape
+    if seq_len == target_seq_len:
+        return x
+    out = x.new_zeros((bsz, target_seq_len, hidden))
+    out[:, :seq_len] = x
+    return out
+
+
+def _pad_mask(x: torch.Tensor, target_seq_len: int) -> torch.Tensor:
+    x = _normalize_mask(x)
+    bsz, seq_len = x.shape
+    if seq_len == target_seq_len:
+        return x
+    out = torch.zeros((bsz, target_seq_len), dtype=torch.bool, device=x.device)
+    out[:, :seq_len] = x
+    return out
+
+
+def _select_states(
+    states: Sequence[DiffusionRequestState],
+    idx_mapping: torch.Tensor | None,
+) -> tuple[list[DiffusionRequestState], torch.Tensor, np.ndarray]:
+    if not states:
+        raise ValueError("Cannot build InputBatch from empty states.")
+
+    if idx_mapping is None:
+        device = states[0].latents.device if states[0].latents is not None else None
+        idx_mapping = torch.arange(len(states), dtype=torch.int32, device=device)
+    else:
+        if idx_mapping.ndim != 1:
+            raise ValueError("idx_mapping must be a 1D tensor.")
+        idx_mapping = idx_mapping.to(dtype=torch.int32)
+
+    selected_states: list[DiffusionRequestState] = []
+    for batch_idx, state_idx in enumerate(idx_mapping.tolist()):
+        if state_idx < 0 or state_idx >= len(states):
+            raise ValueError(f"idx_mapping[{batch_idx}]={state_idx} is out of range for states.")
+        selected_states.append(states[state_idx])
+    return selected_states, idx_mapping, idx_mapping.detach().cpu().numpy()
+
+
+def _prepare_req_ids(states: Sequence[DiffusionRequestState]) -> list[str]:
+    return [state.req_id for state in states]
+
+
+def _prepare_prompt_field_on_state(
+    state: DiffusionRequestState,
+    *,
+    embeds_attr: str,
+    mask_attr: str,
+) -> tuple[torch.Tensor | None, torch.Tensor | None]:
+    embeds = getattr(state, embeds_attr)
+    if embeds is None:
+        return None, None
+
+    normalized_embeds = _normalize_prompt_embeds(embeds)
+    if normalized_embeds is not embeds:
+        setattr(state, embeds_attr, normalized_embeds)
+    embeds = normalized_embeds
+
+    mask = getattr(state, mask_attr)
+    if mask is None:
+        return embeds, None
+
+    normalized_mask = _normalize_mask(mask)
+    if normalized_mask is not mask:
+        setattr(state, mask_attr, normalized_mask)
+    return embeds, normalized_mask
+
+
+def _prepare_reused_buffer(
+    current: torch.Tensor | None,
+    *,
+    shape: tuple[int, ...],
+    dtype: torch.dtype,
+    device: torch.device,
+) -> torch.Tensor:
+    if current is not None and tuple(current.shape) == shape and current.dtype == dtype and current.device == device:
+        return current
+    return torch.empty(shape, dtype=dtype, device=device)
+
+
+def _validate_gather_tensors(
+    values: Sequence[torch.Tensor],
+    *,
+    field_name: str,
+) -> tuple[torch.dtype, torch.device, tuple[int, ...], int]:
+    if not values:
+        raise ValueError(f"Cannot gather empty tensor list for {field_name}.")
+
+    first = values[0]
+    dtype = first.dtype
+    device = first.device
+    suffix_shape = tuple(first.shape[1:])
+    total_rows = 0
+
+    for value in values:
+        if value.dtype != dtype:
+            raise ValueError(f"Mixed dtypes in {field_name} batch.")
+        if value.device != device:
+            raise ValueError(f"Mixed devices in {field_name} batch.")
+        if tuple(value.shape[1:]) != suffix_shape:
+            raise ValueError(
+                f"Mixed trailing shapes in {field_name} batch: expected {suffix_shape}, got {tuple(value.shape[1:])}."
+            )
+        total_rows += int(value.shape[0])
+
+    return dtype, device, suffix_shape, total_rows
+
+
+def _gather_tensor_rows(
+    values: Sequence[torch.Tensor],
+    *,
+    field_name: str,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    dtype, device, suffix_shape, total_rows = _validate_gather_tensors(
+        values,
+        field_name=field_name,
+    )
+    gathered = _prepare_reused_buffer(
+        out,
+        shape=(total_rows, *suffix_shape),
+        dtype=dtype,
+        device=device,
+    )
+
+    row_offset = 0
+    for value in values:
+        next_row_offset = row_offset + int(value.shape[0])
+        gathered[row_offset:next_row_offset].copy_(value)
+        row_offset = next_row_offset
+    return gathered
+
+
+def _get_seq_lens_from_mask(mask: torch.Tensor) -> list[int]:
+    return mask.sum(dim=1, dtype=torch.int32).tolist()
+
+
+def _get_request_prompt_seq_lens(
+    state: DiffusionRequestState,
+    *,
+    embeds_attr: str,
+    mask_attr: str,
+    seq_lens_attr: str,
+) -> list[int]:
+    embeds, mask = _prepare_prompt_field_on_state(
+        state,
+        embeds_attr=embeds_attr,
+        mask_attr=mask_attr,
+    )
+    if embeds is None:
+        raise ValueError(f"{embeds_attr} is not initialized on request {state.req_id}.")
+
+    if mask is not None:
+        if mask.shape[0] != embeds.shape[0]:
+            raise ValueError(f"{mask_attr} batch dimension does not match {embeds_attr} for request {state.req_id}.")
+        return _get_seq_lens_from_mask(mask)
+
+    seq_lens = getattr(state, seq_lens_attr)
+    if seq_lens is not None:
+        return [int(value) for value in seq_lens]
+
+    return [int(embeds.shape[1])] * int(embeds.shape[0])
+
+
+def _prepare_request_prompt_field(
+    state: DiffusionRequestState,
+    *,
+    embeds_attr: str,
+    mask_attr: str,
+    seq_lens_attr: str,
+    target_seq_len: int,
+) -> tuple[torch.Tensor, torch.Tensor | None]:
+    embeds, mask = _prepare_prompt_field_on_state(
+        state,
+        embeds_attr=embeds_attr,
+        mask_attr=mask_attr,
+    )
+    if embeds is None:
+        raise ValueError(f"{embeds_attr} is not initialized on request {state.req_id}.")
+
+    actual_seq_lens = _get_request_prompt_seq_lens(
+        state,
+        embeds_attr=embeds_attr,
+        mask_attr=mask_attr,
+        seq_lens_attr=seq_lens_attr,
+    )
+    max_actual_seq_len = max(actual_seq_lens) if actual_seq_lens else 0
+
+    if mask is None and max_actual_seq_len != target_seq_len:
+        raise ValueError(
+            f"Variable-length {embeds_attr} in batch but {mask_attr} is None. "
+            f"Provide masks or ensure {embeds_attr} have the same seq_len."
+        )
+
+    current_seq_len = int(embeds.shape[1])
+    if current_seq_len < target_seq_len:
+        embeds = _pad_prompt_embeds(embeds, target_seq_len)
+        setattr(state, embeds_attr, embeds)
+        if mask is not None:
+            mask = _pad_mask(mask, target_seq_len)
+            setattr(state, mask_attr, mask)
+        current_seq_len = target_seq_len
+
+    if current_seq_len > target_seq_len:
+        if max_actual_seq_len > target_seq_len:
+            raise ValueError(
+                f"{embeds_attr} for request {state.req_id} requires seq_len "
+                f"{max_actual_seq_len}, got target {target_seq_len}."
+            )
+        return embeds[:, :target_seq_len], None if mask is None else mask[:, :target_seq_len]
+
+    return embeds, mask
+
+
+def _prepare_padded_prompt_fields(
+    states: Sequence[DiffusionRequestState],
+    *,
+    embeds_attr: str,
+    mask_attr: str,
+    seq_lens_attr: str,
+    embeds_out: torch.Tensor | None = None,
+    mask_out: torch.Tensor | None = None,
+) -> tuple[torch.Tensor | None, torch.Tensor | None]:
+    prepared_fields = [
+        _prepare_prompt_field_on_state(
+            state,
+            embeds_attr=embeds_attr,
+            mask_attr=mask_attr,
+        )
+        for state in states
+    ]
+    embeds_values = [embeds for embeds, _ in prepared_fields]
+    if not any(embeds is not None for embeds in embeds_values):
+        return None, None
+    if not all(embeds is not None for embeds in embeds_values):
+        raise ValueError(f"Mixed {embeds_attr} in batch.")
+
+    mask_values = [mask for _, mask in prepared_fields]
+    if any(mask is None for mask in mask_values):
+        if any(mask is not None for mask in mask_values):
+            raise ValueError(f"Mixed {mask_attr} in batch.")
+
+    target_seq_len = max(
+        max(
+            _get_request_prompt_seq_lens(
+                state,
+                embeds_attr=embeds_attr,
+                mask_attr=mask_attr,
+                seq_lens_attr=seq_lens_attr,
+            )
+        )
+        for state in states
+    )
+
+    request_embeds: list[torch.Tensor] = []
+    request_masks: list[torch.Tensor] = []
+    for state in states:
+        prepared_embeds, prepared_mask = _prepare_request_prompt_field(
+            state,
+            embeds_attr=embeds_attr,
+            mask_attr=mask_attr,
+            seq_lens_attr=seq_lens_attr,
+            target_seq_len=target_seq_len,
+        )
+        request_embeds.append(prepared_embeds)
+        if prepared_mask is not None:
+            request_masks.append(prepared_mask)
+
+    gathered_embeds = _gather_tensor_rows(
+        request_embeds,
+        field_name=embeds_attr,
+        out=embeds_out,
+    )
+
+    if not request_masks:
+        return gathered_embeds, None
+
+    gathered_masks = _gather_tensor_rows(
+        request_masks,
+        field_name=mask_attr,
+        out=mask_out,
+    )
+    return gathered_embeds, gathered_masks
+
+
+def _prepare_latents(
+    states: Sequence[DiffusionRequestState],
+    *,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    latents_values = [state.latents for state in states]
+    if any(latents is None for latents in latents_values):
+        raise ValueError("All requests must have `latents` initialized.")
+    return _gather_tensor_rows(
+        [latents for latents in latents_values if latents is not None],
+        field_name="latents",
+        out=out,
+    )
+
+
+def _require_state_latents(
+    state: DiffusionRequestState,
+    *,
+    for_field: str,
+) -> torch.Tensor:
+    latents = state.latents
+    if latents is None:
+        raise ValueError(f"Request {state.req_id} has no latents while preparing {for_field}.")
+    return latents
+
+
+def _expand_scalar_or_vector(
+    value: torch.Tensor,
+    *,
+    num_rows: int,
+    field_name: str,
+) -> torch.Tensor:
+    if value.ndim == 0:
+        return value.reshape(1).expand(num_rows)
+    if value.ndim != 1:
+        raise ValueError(f"{field_name} must be scalar or 1D, got ndim={value.ndim}.")
+    if value.shape[0] == num_rows:
+        return value
+    if value.shape[0] == 1:
+        return value.expand(num_rows)
+    raise ValueError(
+        f"Per-request {field_name} must have either 1 element or {num_rows} elements; got {value.shape[0]}."
+    )
+
+
+def _prepare_timesteps(
+    states: Sequence[DiffusionRequestState],
+    *,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    timestep_values: list[torch.Tensor] = []
+    for state in states:
+        timestep = state.current_timestep
+        if timestep is None:
+            raise ValueError("All requests must have a current timestep initialized.")
+        if not torch.is_tensor(timestep):
+            raise ValueError("InputBatch expects tensor timesteps; normalize them before batching.")
+        latents = _require_state_latents(state, for_field="timesteps")
+        timestep_values.append(
+            _expand_scalar_or_vector(
+                timestep,
+                num_rows=int(latents.shape[0]),
+                field_name="timestep tensor",
+            )
+        )
+
+    return _gather_tensor_rows(
+        timestep_values,
+        field_name="timesteps",
+        out=out,
+    )
+
+
+def _prepare_cfg_scalars(
+    states: Sequence[DiffusionRequestState],
+) -> tuple[bool, float, bool]:
+    def _cfg_scalars(state: DiffusionRequestState) -> tuple[bool, float, bool]:
+        true_cfg_scale = getattr(state.sampling, "true_cfg_scale", None) or 4.0
+        cfg_normalize = bool(getattr(state.sampling, "cfg_normalize", False))
+        return state.do_true_cfg, true_cfg_scale, cfg_normalize
+
+    scalars = _cfg_scalars(states[0])
+    for state in states[1:]:
+        if _cfg_scalars(state) != scalars:
+            raise ValueError("Mixed CFG settings in one diffusion batch are not supported.")
+    return scalars
+
+
+def _prepare_guidance(
+    states: Sequence[DiffusionRequestState],
+    *,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor | None:
+    guidance_values = [state.guidance for state in states]
+    if all(guidance is None for guidance in guidance_values):
+        return None
+    if any(guidance is None for guidance in guidance_values):
+        raise ValueError("Mixed guidance in one diffusion batch are not supported.")
+
+    gathered_guidance: list[torch.Tensor] = []
+    for state in states:
+        latents = _require_state_latents(state, for_field="guidance")
+        guidance = state.guidance
+        assert guidance is not None
+        guidance_tensor = torch.as_tensor(
+            guidance,
+            device=latents.device,
+            dtype=latents.dtype,
+        )
+        gathered_guidance.append(
+            _expand_scalar_or_vector(
+                guidance_tensor,
+                num_rows=int(latents.shape[0]),
+                field_name="guidance tensor",
+            )
+        )
+
+    return _gather_tensor_rows(
+        gathered_guidance,
+        field_name="guidance",
+        out=out,
+    )
+
+
+def _prepare_image_latents(
+    states: Sequence[DiffusionRequestState],
+    *,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor | None:
+    image_latents = [getattr(state.sampling, "image_latent", None) for state in states]
+    if all(image_latent is None for image_latent in image_latents):
+        return None
+    if any(image_latent is None for image_latent in image_latents):
+        raise ValueError("Mixed image_latent presence in one diffusion batch is not supported.")
+    return _gather_tensor_rows(
+        [image_latent for image_latent in image_latents if image_latent is not None],
+        field_name="image_latents",
+        out=out,
+    )
+
+
+def _prepare_seq_lens(
+    states: Sequence[DiffusionRequestState],
+    attr_name: str,
+) -> list[int] | None:
+    values = [getattr(state, attr_name) for state in states]
+    if all(value is None for value in values):
+        return None
+    if any(value is None for value in values):
+        raise ValueError(f"Mixed {attr_name} in batch.")
+    return [int(value[0]) for value in values if value is not None]
+
+
+def _prepare_img_shapes(states: Sequence[DiffusionRequestState]) -> list | None:
+    values = [state.img_shapes for state in states]
+    if all(value is None for value in values):
+        return None
+    if any(value is None for value in values):
+        raise ValueError("Mixed img_shapes in batch.")
+    return [value[0] if value else [] for value in values if value is not None]
+
+
+def _prepare_prompt_embeds(
+    states: Sequence[DiffusionRequestState],
+    *,
+    embeds_out: torch.Tensor | None = None,
+    mask_out: torch.Tensor | None = None,
+) -> tuple[torch.Tensor, torch.Tensor | None]:
+    prompt_embeds, prompt_embeds_mask = _prepare_padded_prompt_fields(
+        states,
+        embeds_attr="prompt_embeds",
+        mask_attr="prompt_embeds_mask",
+        seq_lens_attr="txt_seq_lens",
+        embeds_out=embeds_out,
+        mask_out=mask_out,
+    )
+    if prompt_embeds is None:
+        raise ValueError("All requests must have `prompt_embeds` initialized.")
+    return prompt_embeds, prompt_embeds_mask
+
+
+def _prepare_negative_prompt_embeds(
+    states: Sequence[DiffusionRequestState],
+    *,
+    embeds_out: torch.Tensor | None = None,
+    mask_out: torch.Tensor | None = None,
+) -> tuple[torch.Tensor | None, torch.Tensor | None]:
+    return _prepare_padded_prompt_fields(
+        states,
+        embeds_attr="negative_prompt_embeds",
+        mask_attr="negative_prompt_embeds_mask",
+        seq_lens_attr="negative_txt_seq_lens",
+        embeds_out=embeds_out,
+        mask_out=mask_out,
+    )
+
+
+def _same_composition(
+    cached_batch: InputBatch | None,
+    req_ids: list[str],
+    idx_mapping_np: np.ndarray,
+) -> bool:
+    if cached_batch is None:
+        return False
+    if cached_batch.req_ids != req_ids:
+        return False
+    return np.array_equal(cached_batch.idx_mapping_np, idx_mapping_np)
+
+
+def _scatter_batch_tensor_by_mapping(
+    states: Sequence[DiffusionRequestState],
+    idx_mapping_np: np.ndarray,
+    *,
+    attr_name: str,
+    value: torch.Tensor,
+) -> None:
+    row_offset = 0
+    for batch_idx, state_idx in enumerate(idx_mapping_np.tolist()):
+        if state_idx < 0 or state_idx >= len(states):
+            raise ValueError(f"idx_mapping[{batch_idx}]={state_idx} is out of range for states.")
+        state = states[state_idx]
+        state_value = getattr(state, attr_name)
+        num_rows = 1 if state_value is None else int(state_value.shape[0])
+        next_row_offset = row_offset + num_rows
+        value_slice = value[row_offset:next_row_offset]
+        if state_value is None:
+            setattr(state, attr_name, value_slice.clone())
+        elif (
+            tuple(state_value.shape) != tuple(value_slice.shape)
+            or state_value.dtype != value_slice.dtype
+            or state_value.device != value_slice.device
+        ):
+            setattr(state, attr_name, value_slice.clone())
+        else:
+            state_value.copy_(value_slice)
+        row_offset = next_row_offset
+
+    if row_offset != int(value.shape[0]):
+        raise ValueError(
+            f"Scatter for {attr_name} consumed {row_offset} rows, but batch has {int(value.shape[0])} rows."
+        )
+
+
+@dataclass
+class InputBatch:
+    """Ephemeral step-level batch view.
+
+    Static request-local tensors are normalized and padded onto
+    ``DiffusionRequestState`` itself, making the request state the persistent
+    source of truth. ``InputBatch`` only assembles a contiguous view for the
+    current step and refreshes dynamic fields in-place when composition is
+    unchanged.
+    """
+
+    req_ids: list[str]
+    num_reqs: int
+    num_reqs_after_padding: int
+    idx_mapping: torch.Tensor
+    idx_mapping_np: np.ndarray
+
+    latents: torch.Tensor
+    timesteps: torch.Tensor
+    prompt_embeds: torch.Tensor
+    prompt_embeds_mask: torch.Tensor | None
+    negative_prompt_embeds: torch.Tensor | None
+    negative_prompt_embeds_mask: torch.Tensor | None
+    guidance: torch.Tensor | None = None
+    do_true_cfg: bool = False
+    true_cfg_scale: float = 4.0
+    cfg_normalize: bool = False
+    image_latents: torch.Tensor | None = None
+
+    img_shapes: list | None = None
+    txt_seq_lens: list[int] | None = None
+    negative_txt_seq_lens: list[int] | None = None
+
+    def __post_init__(self) -> None:
+        if len(self.req_ids) != int(self.idx_mapping.numel()):
+            raise ValueError("`req_ids` and `idx_mapping` must have the same length.")
+        if self.num_reqs != len(self.req_ids):
+            raise ValueError("`num_reqs` must match the number of request ids.")
+        if self.num_reqs_after_padding < self.num_reqs:
+            raise ValueError("`num_reqs_after_padding` must be >= `num_reqs`.")
+
+    def _refresh_dynamic_fields(
+        self,
+        selected_states: Sequence[DiffusionRequestState],
+    ) -> None:
+        self.latents = _prepare_latents(selected_states, out=self.latents)
+        self.timesteps = _prepare_timesteps(selected_states, out=self.timesteps)
+
+    def _refresh_static_fields(
+        self,
+        states: Sequence[DiffusionRequestState],
+    ) -> None:
+        self.do_true_cfg, self.true_cfg_scale, self.cfg_normalize = _prepare_cfg_scalars(states)
+        self.guidance = _prepare_guidance(states, out=self.guidance)
+        self.image_latents = _prepare_image_latents(states, out=self.image_latents)
+        self.prompt_embeds, self.prompt_embeds_mask = _prepare_prompt_embeds(
+            states,
+            embeds_out=self.prompt_embeds,
+            mask_out=self.prompt_embeds_mask,
+        )
+        (
+            self.negative_prompt_embeds,
+            self.negative_prompt_embeds_mask,
+        ) = _prepare_negative_prompt_embeds(
+            states,
+            embeds_out=self.negative_prompt_embeds,
+            mask_out=self.negative_prompt_embeds_mask,
+        )
+        self.img_shapes = _prepare_img_shapes(states)
+        self.txt_seq_lens = _prepare_seq_lens(states, "txt_seq_lens")
+        self.negative_txt_seq_lens = _prepare_seq_lens(states, "negative_txt_seq_lens")
+
+    def _repack_dynamic_fields(
+        self,
+        selected_states: Sequence[DiffusionRequestState],
+    ) -> None:
+        self._refresh_dynamic_fields(selected_states)
+
+    def _rebuild(
+        self,
+        selected_states: Sequence[DiffusionRequestState],
+        idx_mapping: torch.Tensor,
+        idx_mapping_np: np.ndarray,
+        req_ids: list[str],
+    ) -> InputBatch:
+        self.req_ids = req_ids
+        self.num_reqs = len(req_ids)
+        self.num_reqs_after_padding = len(req_ids)
+        self.idx_mapping = idx_mapping
+        self.idx_mapping_np = idx_mapping_np
+        self.latents = _prepare_latents(selected_states, out=self.latents)
+        self.timesteps = _prepare_timesteps(selected_states, out=self.timesteps)
+        self._refresh_static_fields(selected_states)
+        self.__post_init__()
+        return self
+
+    @classmethod
+    def make_batch(
+        cls,
+        states: Sequence[DiffusionRequestState],
+        idx_mapping: torch.Tensor | None = None,
+        cached_batch: InputBatch | None = None,
+    ) -> InputBatch:
+        """Build a temporary step-local batch view from request states."""
+        selected_states, idx_mapping, idx_mapping_np = _select_states(states, idx_mapping)
+        req_ids = _prepare_req_ids(selected_states)
+
+        if _same_composition(cached_batch, req_ids, idx_mapping_np):
+            assert cached_batch is not None
+            cached_batch._repack_dynamic_fields(selected_states)
+            return cached_batch
+
+        if cached_batch is not None:
+            return cached_batch._rebuild(
+                selected_states,
+                idx_mapping,
+                idx_mapping_np,
+                req_ids,
+            )
+
+        prompt_embeds, prompt_embeds_mask = _prepare_prompt_embeds(selected_states)
+        negative_prompt_embeds, negative_prompt_embeds_mask = _prepare_negative_prompt_embeds(selected_states)
+        do_true_cfg, true_cfg_scale, cfg_normalize = _prepare_cfg_scalars(selected_states)
+        return cls(
+            req_ids=req_ids,
+            num_reqs=len(selected_states),
+            num_reqs_after_padding=len(selected_states),
+            idx_mapping=idx_mapping,
+            idx_mapping_np=idx_mapping_np,
+            latents=_prepare_latents(selected_states),
+            timesteps=_prepare_timesteps(selected_states),
+            guidance=_prepare_guidance(selected_states),
+            do_true_cfg=do_true_cfg,
+            true_cfg_scale=true_cfg_scale,
+            cfg_normalize=cfg_normalize,
+            image_latents=_prepare_image_latents(selected_states),
+            prompt_embeds=prompt_embeds,
+            prompt_embeds_mask=prompt_embeds_mask,
+            negative_prompt_embeds=negative_prompt_embeds,
+            negative_prompt_embeds_mask=negative_prompt_embeds_mask,
+            img_shapes=_prepare_img_shapes(selected_states),
+            txt_seq_lens=_prepare_seq_lens(selected_states, "txt_seq_lens"),
+            negative_txt_seq_lens=_prepare_seq_lens(
+                selected_states,
+                "negative_txt_seq_lens",
+            ),
+        )
+
+
+def scatter_latents(
+    states: Sequence[DiffusionRequestState],
+    input_batch: InputBatch,
+) -> None:
+    """Scatter the step-updated latents back into persistent request states.
+
+    This is the CPU fallback of the vLLM-style post-update path. The mapping is
+    driven entirely by ``input_batch.idx_mapping_np`` so the runner remains free
+    to keep request states in its own persistent storage layout.
+    """
+    _scatter_batch_tensor_by_mapping(
+        states,
+        input_batch.idx_mapping_np,
+        attr_name="latents",
+        value=input_batch.latents,
+    )
+
+
+DiffusionInputBatch = InputBatch

vllm_omni/diffusion/worker/utils.py

@@ -4,6 +4,7 @@
 
 from __future__ import annotations
 
+from abc import ABC, abstractmethod
 from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, Any
 
@@ -109,8 +110,14 @@ class DiffusionRequestState:
         return self.step_index == 0 or self.timesteps is None
 
 
+class BaseRunnerOutput(ABC):
+    @abstractmethod
+    def get_req_output(self, sched_req_id: str) -> RunnerOutput | None:
+        pass
+
+
 @dataclass
-class RunnerOutput:
+class RunnerOutput(BaseRunnerOutput):
     """Output of a single denoising step for a request.
 
     NOTE: `latents` may be None when returned through IPC to avoid
@@ -122,3 +129,34 @@ class RunnerOutput:
     step_index: int | None = None
     finished: bool = False
     result: DiffusionOutput | None = None
+
+    def get_req_output(self, sched_req_id: str) -> RunnerOutput | None:
+        return self if self.req_id == sched_req_id else None
+
+
+@dataclass
+class BatchRunnerOutput(BaseRunnerOutput):
+    runner_outputs: list[RunnerOutput]
+    _id_to_idx: dict[str, int] = field(init=False, repr=False)
+
+    def __post_init__(self) -> None:
+        self._id_to_idx = {out.req_id: i for i, out in enumerate(self.runner_outputs)}
+
+    def __getitem__(self, req_id: str) -> RunnerOutput | None:
+        """access single RunnerOutput by req_id"""
+        idx = self._id_to_idx.get(req_id)
+        return self.runner_outputs[idx] if idx is not None else None
+
+    def get_req_output(self, sched_req_id: str) -> RunnerOutput | None:
+        return self[sched_req_id]
+
+    @property
+    def req_ids(self) -> list[str]:
+        return list(self._id_to_idx.keys())
+
+    def __len__(self) -> int:
+        return len(self.runner_outputs)
+
+    @classmethod
+    def from_list(cls, runner_output_list: list[RunnerOutput]) -> BatchRunnerOutput:
+        return cls(runner_outputs=runner_output_list)