[Core][Frontend] Support request-level batching for diffusion pipelines (#4079)

Signed-off-by: jader <yjader@foxmail.com>
Co-authored-by: Samit <285365963@qq.com>

tests/diffusion/test_diffusion_engine_rpc_routing.py

@@ -32,12 +32,12 @@ import pytest
 from vllm_omni.diffusion.data import DiffusionOutput
 from vllm_omni.diffusion.diffusion_engine import DiffusionEngine, _RpcTask
 from vllm_omni.diffusion.sched import RequestScheduler
-from vllm_omni.diffusion.sched.interface import SamplingParamsKey
+from vllm_omni.diffusion.sched.interface import RequestBatchSamplingParamsKey
 from vllm_omni.diffusion.worker.utils import RunnerOutput
 
 # Default values for every batch-key field, so SimpleNamespace-based
 # sampling_params satisfy ``get_sampling_params_key``'s attribute lookups.
-_SAMPLING_KEY_DEFAULTS = {f.name: f.default for f in _dc_fields(SamplingParamsKey)}
+_SAMPLING_KEY_DEFAULTS = {f.name: f.default for f in _dc_fields(RequestBatchSamplingParamsKey)}
 
 pytestmark = [pytest.mark.diffusion, pytest.mark.cpu]
 
@@ -59,6 +59,7 @@ class _ConcurrencyTrackingExecutor:
         self.rpc_delay = rpc_delay
         self.is_failed = False
         self._closed = False
+        self.od_config = SimpleNamespace()
 
     def collective_rpc(
         self,
@@ -101,14 +102,15 @@ class _ConcurrencyTrackingExecutor:
         # Mimic the real MultiprocDiffusionExecutor.execute_request: it
         # forwards a single request through collective_rpc.
         new_req = scheduler_output.scheduled_new_reqs[0]
+        req = new_req.req
         result = self.collective_rpc(
             "execute_model",
-            args=(new_req.req,),
+            args=(req, self.od_config),
             unique_reply_rank=0,
             exec_all_ranks=True,
         )
         return RunnerOutput(
-            request_id=new_req.request_id,
+            request_id=req.request_id,
             step_index=None,
             finished=True,
             result=result,
@@ -119,10 +121,12 @@ class _ConcurrencyTrackingExecutor:
 
 
 def _make_request(tag: str):
+    sampling_params = dict(_SAMPLING_KEY_DEFAULTS)
+    sampling_params["num_inference_steps"] = 1
     return SimpleNamespace(
         request_id=tag,
-        prompts=[f"prompt_{tag}"],
-        sampling_params=SimpleNamespace(num_inference_steps=1, **_SAMPLING_KEY_DEFAULTS),
+        prompt=f"prompt_{tag}",
+        sampling_params=SimpleNamespace(**sampling_params),
     )
 
 
@@ -137,17 +141,20 @@ def _make_engine_with_loop(
     """
     engine = DiffusionEngine.__new__(DiffusionEngine)
     engine._closed = False
+    engine.od_config = SimpleNamespace(streaming_output=False)
     engine.executor = _ConcurrencyTrackingExecutor(rpc_delay=rpc_delay)
 
     sched = RequestScheduler()
     sched.initialize(SimpleNamespace(max_num_seqs=1))
     engine.scheduler = sched
     engine.step_execution = False
+    engine.supports_request_batch = False
     engine.execute_fn = engine.executor.execute_request
 
     engine._rpc_lock = threading.RLock()
     engine._cv = threading.Condition(engine._rpc_lock)
     engine._out_queue = {}
+    engine._out_queue_streaming = {}
     engine._closed = False
     engine.abort_queue = queue.Queue()
     engine._rpc_queue = queue.Queue()