[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_model_runner.py

@@ -11,6 +11,7 @@ import vllm_omni.diffusion.worker.diffusion_model_runner as model_runner_module
 from tests.helpers.mark import hardware_test
 from vllm_omni.diffusion.data import DiffusionOutput
 from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner
+from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch, split_diffusion_output_by_request
 
 pytestmark = [pytest.mark.diffusion]
 
@@ -22,14 +23,39 @@ def _noop_forward_context(*args, **kwargs):
 
 
 class _DummyPipeline:
+    supports_request_batch = True
+
     def __init__(self, output):
         self._output = output
         self.forward_calls = 0
+        self.last_req = None
 
     def forward(self, req):
-        del req
+        self.last_req = req
         self.forward_calls += 1
-        return self._output
+        return [self._output]
+
+
+class _SingleRequestBatchPipeline:
+    supports_request_batch = False
+
+    def __init__(self):
+        self.last_req = None
+
+    def forward(self, req):
+        self.last_req = req
+        return [DiffusionOutput(output=req.prompts[0])]
+
+
+class _SingleRequestDiffusionOutputPipeline:
+    supports_request_batch = False
+
+    def __init__(self):
+        self.last_req = None
+
+    def forward(self, req):
+        self.last_req = req
+        return DiffusionOutput(output=req.prompts[0])
 
 
 class _ChunkStepPipeline:
@@ -81,9 +107,28 @@ def _make_request(skip_cache_refresh: bool = True):
     )
     return SimpleNamespace(
         request_id="req-test",
-        prompts=["a prompt"],
+        prompt="a prompt",
         sampling_params=sampling_params,
         skip_cache_refresh=skip_cache_refresh,
+        kv_sender_info=None,
+    )
+
+
+def _make_request_with_params(req_id: str, sampling_params):
+    return SimpleNamespace(
+        request_id=req_id,
+        prompt=f"prompt-{req_id}",
+        prompts=[f"prompt-{req_id}"],
+        sampling_params=sampling_params,
+        skip_cache_refresh=True,
+    )
+
+
+def _fake_platform_for_peak_memory():
+    return SimpleNamespace(
+        reset_peak_memory_stats=lambda: None,
+        max_memory_reserved=lambda: 0,
+        max_memory_allocated=lambda: 0,
     )
 
 
@@ -91,7 +136,7 @@ def _make_runner(cache_backend, cache_backend_name: str, enable_cache_dit_summar
     runner = object.__new__(DiffusionModelRunner)
     runner.vllm_config = object()
     runner.device = torch.device("cpu")
-    runner.pipeline = _DummyPipeline(output=SimpleNamespace(output="ok"))
+    runner.pipeline = _DummyPipeline(output=DiffusionOutput(output="ok"))
     runner.cache_backend = cache_backend
     runner.offload_backend = None
     runner.state_cache = {}
@@ -238,6 +283,232 @@ def test_execute_model_emits_cache_summary_with_active_cache_dit_backend(monkeyp
     assert cache_summary_calls == [(runner.pipeline, True)]
 
 
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_execute_model_passes_single_request_batch_to_non_admission_pipeline(monkeypatch):
+    runner = _make_runner(cache_backend=None, cache_backend_name="none")
+    runner.pipeline = _SingleRequestBatchPipeline()
+    req = _make_request(skip_cache_refresh=True)
+
+    monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
+
+    output = DiffusionModelRunner.execute_model(runner, req)
+
+    assert output.output == "a prompt"
+    assert isinstance(runner.pipeline.last_req, DiffusionRequestBatch)
+    assert runner.pipeline.last_req.num_reqs == 1
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_execute_model_accepts_bare_diffusion_output_from_single_request_pipeline(monkeypatch):
+    runner = _make_runner(cache_backend=None, cache_backend_name="none")
+    runner.pipeline = _SingleRequestDiffusionOutputPipeline()
+    req = _make_request(skip_cache_refresh=True)
+
+    monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
+
+    output = DiffusionModelRunner.execute_model(runner, req)
+
+    assert output.output == "a prompt"
+    assert isinstance(runner.pipeline.last_req, DiffusionRequestBatch)
+    assert runner.pipeline.last_req.num_reqs == 1
+
+
+class _BatchPipeline:
+    """Pipeline returning a configurable list of outputs from forward()."""
+
+    supports_request_batch = True
+
+    def __init__(self, outputs):
+        self._outputs = outputs
+        self.last_batch = None
+
+    def forward(self, batch):
+        self.last_batch = batch
+        return list(self._outputs)
+
+
+class _SingleRequestPipeline:
+    def forward(self, batch):
+        del batch
+        return [DiffusionOutput(output="single")]
+
+
+class _BatchSingleOutputPipeline:
+    supports_request_batch = True
+
+    def forward(self, batch):
+        return DiffusionOutput(output=batch.prompts[0])
+
+
+def _make_batch_runner(pipeline):
+    runner = object.__new__(DiffusionModelRunner)
+    runner.vllm_config = object()
+    runner.device = torch.device("cpu")
+    runner.pipeline = pipeline
+    runner.cache_backend = None
+    runner.offload_backend = None
+    runner.od_config = SimpleNamespace(
+        cache_backend="none",
+        enable_cache_dit_summary=False,
+        parallel_config=SimpleNamespace(use_hsdp=False),
+    )
+    runner.kv_transfer_manager = SimpleNamespace(
+        receive_multi_kv_cache_distributed=lambda req, cfg_kv_collect_func=None, target_device=None: None,
+    )
+    return runner
+
+
+def _make_scheduler_output(num_reqs: int):
+    reqs = [_make_request() for _ in range(num_reqs)]
+    for i, req in enumerate(reqs):
+        req.request_id = f"req-{i}"
+    return SimpleNamespace(scheduled_new_reqs=[SimpleNamespace(req=req) for req in reqs])
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_execute_model_batch_rejects_output_count_mismatch(monkeypatch):
+    """A pipeline returning the wrong number of outputs must fail loudly,
+    not silently drop requests or IndexError on the per-request mapping."""
+    monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
+    monkeypatch.setattr(model_runner_module, "current_omni_platform", _fake_platform_for_peak_memory())
+    # forward returns 1 output for 2 scheduled requests
+    runner = _make_batch_runner(_BatchPipeline(outputs=[DiffusionOutput(output="only-one")]))
+    sched = _make_scheduler_output(num_reqs=2)
+
+    with pytest.raises(RuntimeError, match="returned 1 outputs for 2 requests"):
+        DiffusionModelRunner.execute_model_batch(runner, sched, runner.od_config)
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_execute_model_batch_routes_one_output_per_request(monkeypatch):
+    monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
+    monkeypatch.setattr(model_runner_module, "current_omni_platform", _fake_platform_for_peak_memory())
+    outs = [DiffusionOutput(output="a"), DiffusionOutput(output="b")]
+    runner = _make_batch_runner(_BatchPipeline(outputs=outs))
+    sched = _make_scheduler_output(num_reqs=2)
+
+    result = DiffusionModelRunner.execute_model_batch(runner, sched, runner.od_config)
+
+    assert len(result.runner_outputs) == 2
+    assert result.runner_outputs[0].request_id == "req-0"
+    assert result.runner_outputs[0].result.output == "a"
+    assert result.runner_outputs[1].request_id == "req-1"
+    assert result.runner_outputs[1].result.output == "b"
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_execute_model_batch_preserves_per_request_sampling_and_seeds_generators(monkeypatch):
+    monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
+    monkeypatch.setattr(model_runner_module, "current_omni_platform", _fake_platform_for_peak_memory())
+    outputs = [DiffusionOutput(output="a"), DiffusionOutput(output="b")]
+    pipeline = _BatchPipeline(outputs=outputs)
+    runner = _make_batch_runner(pipeline)
+    sched = _make_scheduler_output(num_reqs=2)
+    sched.scheduled_new_reqs[0].req.sampling_params.seed = 111
+    sched.scheduled_new_reqs[1].req.sampling_params.seed = 222
+
+    DiffusionModelRunner.execute_model_batch(runner, sched, runner.od_config)
+
+    assert pipeline.last_batch is not None
+    assert [sp.seed for sp in pipeline.last_batch.sampling_params_list] == [111, 222]
+    assert [sp.generator.initial_seed() for sp in pipeline.last_batch.sampling_params_list] == [111, 222]
+    with pytest.raises(AssertionError, match="multiple requests"):
+        _ = pipeline.last_batch.sampling_params
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_execute_model_batch_rejects_pipeline_without_request_batch_support(monkeypatch):
+    monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
+    monkeypatch.setattr(model_runner_module, "current_omni_platform", _fake_platform_for_peak_memory())
+    runner = _make_batch_runner(_SingleRequestPipeline())
+    sched = _make_scheduler_output(num_reqs=2)
+
+    with pytest.raises(RuntimeError, match="does not support request-batch forward"):
+        DiffusionModelRunner.execute_model_batch(runner, sched, runner.od_config)
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_execute_model_batch_rejects_single_diffusion_output(monkeypatch):
+    monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
+    monkeypatch.setattr(model_runner_module, "current_omni_platform", _fake_platform_for_peak_memory())
+    runner = _make_batch_runner(_BatchSingleOutputPipeline())
+    sched = _make_scheduler_output(num_reqs=1)
+
+    with pytest.raises(RuntimeError, match="request-batch forward must return list\\[DiffusionOutput\\]"):
+        DiffusionModelRunner.execute_model_batch(runner, sched, runner.od_config)
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_execute_model_batch_uses_runner_output_helper(monkeypatch):
+    monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
+    monkeypatch.setattr(model_runner_module, "current_omni_platform", _fake_platform_for_peak_memory())
+    outs = [DiffusionOutput(output="a"), DiffusionOutput(output="b")]
+    runner = _make_batch_runner(_BatchPipeline(outputs=outs))
+    sched = _make_scheduler_output(num_reqs=2)
+    helper_calls = []
+    original_from_outputs = DiffusionModelRunner._runner_output_from_outputs
+
+    def _recording_from_outputs(self, reqs, outputs):
+        helper_calls.append(([req.request_id for req in reqs], [output.output for output in outputs]))
+        return original_from_outputs(self, reqs, outputs)
+
+    monkeypatch.setattr(DiffusionModelRunner, "_runner_output_from_outputs", _recording_from_outputs)
+
+    result = DiffusionModelRunner.execute_model_batch(runner, sched, runner.od_config)
+
+    assert [runner_output.result.output for runner_output in result.runner_outputs] == ["a", "b"]
+    assert helper_calls == [(["req-0", "req-1"], ["a", "b"])]
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_split_diffusion_output_by_request_slices_single_and_multi_request_outputs():
+    reqs = [_make_request(), _make_request()]
+    reqs[0].request_id = "req-0"
+    reqs[1].request_id = "req-1"
+    batch = DiffusionRequestBatch(requests=reqs)
+    result = DiffusionOutput(output=["img-0a", "img-0b", "img-1a", "img-1b"], stage_durations={"decode": 1.0})
+
+    outputs = split_diffusion_output_by_request(result, batch, num_outputs_per_prompt=2)
+
+    assert [output.output for output in outputs] == [["img-0a", "img-0b"], ["img-1a", "img-1b"]]
+    assert [output.stage_durations for output in outputs] == [{"decode": 1.0}, {"decode": 1.0}]
+
+    single = split_diffusion_output_by_request(
+        result, DiffusionRequestBatch(requests=reqs[:1]), num_outputs_per_prompt=2
+    )
+
+    assert single[0].output == ["img-0a", "img-0b"]
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_split_diffusion_output_by_request_slices_tuple_outputs():
+    reqs = [_make_request(), _make_request()]
+    batch = DiffusionRequestBatch(requests=reqs)
+    result = DiffusionOutput(
+        output=(
+            ["video-0", "video-1"],
+            torch.tensor([10, 20]),
+        )
+    )
+
+    outputs = split_diffusion_output_by_request(result, batch, num_outputs_per_prompt=1)
+
+    assert outputs[0].output[0] == ["video-0"]
+    assert torch.equal(outputs[0].output[1], torch.tensor([10]))
+    assert outputs[1].output[0] == ["video-1"]
+    assert torch.equal(outputs[1].output[1], torch.tensor([20]))
+
+
 @pytest.mark.core_model
 @pytest.mark.cpu
 def test_load_model_clears_cache_backend_for_unsupported_pipeline(monkeypatch):