[Diffusion] Support LoRA in step-wise execution (#3639)

Signed-off-by: samithuang <285365963@qq.com>

tests/diffusion/test_diffusion_scheduler.py

@@ -153,6 +153,54 @@ class _StubScheduler(SchedulerInterface):
         return None
 
 
+class TestGetSamplingParamsKey:
+    """Pure-function tests for the batch-compatibility key builder."""
+
+    @staticmethod
+    def _make(lora_int_id: int | None = None, lora_scale: float = 1.0) -> OmniDiffusionRequest:
+        from vllm_omni.lora.request import LoRARequest
+
+        sp = OmniDiffusionSamplingParams(num_inference_steps=2)
+        if lora_int_id is not None:
+            sp.lora_request = LoRARequest(
+                lora_name=f"adapter-{lora_int_id}",
+                lora_int_id=lora_int_id,
+                lora_path=f"/tmp/lora-{lora_int_id}",
+            )
+        sp.lora_scale = lora_scale
+        return OmniDiffusionRequest(
+            prompts=["prompt"],
+            sampling_params=sp,
+            request_ids=[f"req-{lora_int_id}-{lora_scale}"],
+        )
+
+    def test_distinguishes_lora_id(self) -> None:
+        from vllm_omni.diffusion.sched.base_scheduler import get_sampling_params_key
+
+        assert get_sampling_params_key(self._make(lora_int_id=1)) != get_sampling_params_key(self._make(lora_int_id=2))
+
+    def test_distinguishes_lora_scale(self) -> None:
+        from vllm_omni.diffusion.sched.base_scheduler import get_sampling_params_key
+
+        assert get_sampling_params_key(self._make(lora_int_id=1, lora_scale=0.5)) != get_sampling_params_key(
+            self._make(lora_int_id=1, lora_scale=1.0)
+        )
+
+    def test_treats_no_lora_as_distinct_bucket(self) -> None:
+        from vllm_omni.diffusion.sched.base_scheduler import get_sampling_params_key
+
+        assert get_sampling_params_key(self._make(lora_int_id=None)) != get_sampling_params_key(
+            self._make(lora_int_id=1)
+        )
+
+    def test_equal_for_same_lora_identity(self) -> None:
+        from vllm_omni.diffusion.sched.base_scheduler import get_sampling_params_key
+
+        a = get_sampling_params_key(self._make(lora_int_id=1, lora_scale=0.5))
+        b = get_sampling_params_key(self._make(lora_int_id=1, lora_scale=0.5))
+        assert a == b
+
+
 class TestRequestScheduler:
     def setup_method(self) -> None:
         self.scheduler: RequestScheduler = RequestScheduler()
@@ -728,6 +776,114 @@ class TestStepScheduler:
         assert second.num_running_reqs == 1
         assert second.num_waiting_reqs == 1
 
+    def test_step_batch_co_schedules_requests_sharing_lora(self) -> None:
+        """Multiple requests with the same LoRA (id + scale) co-batch."""
+        from vllm_omni.lora.request import LoRARequest
+
+        scheduler = StepScheduler()
+        scheduler.initialize(SimpleNamespace(max_num_seqs=3))
+
+        lora = LoRARequest(lora_name="adapter", lora_int_id=42, lora_path="/tmp/lora")
+
+        def _with_lora(req_id: str) -> OmniDiffusionRequest:
+            sp = OmniDiffusionSamplingParams(num_inference_steps=4)
+            sp.lora_request = lora
+            sp.lora_scale = 0.5
+            return _make_step_request(req_id, sampling_params=sp)
+
+        req_a = scheduler.add_request(_with_lora("a"))
+        req_b = scheduler.add_request(_with_lora("b"))
+        req_c = scheduler.add_request(_with_lora("c"))
+
+        sched_output = scheduler.schedule()
+
+        assert _new_ids(sched_output) == [req_a, req_b, req_c]
+        assert sched_output.num_running_reqs == 3
+        assert sched_output.num_waiting_reqs == 0
+
+    def test_step_batch_separates_requests_with_different_lora_ids(self) -> None:
+        """Different LoRA adapters → distinct batches admitted in FIFO order."""
+        from vllm_omni.lora.request import LoRARequest
+
+        scheduler = StepScheduler()
+        scheduler.initialize(SimpleNamespace(max_num_seqs=4))
+
+        lora_a = LoRARequest(lora_name="adapter-A", lora_int_id=1, lora_path="/tmp/lora-a")
+        lora_b = LoRARequest(lora_name="adapter-B", lora_int_id=2, lora_path="/tmp/lora-b")
+
+        def _build(req_id: str, lora: LoRARequest) -> OmniDiffusionRequest:
+            sp = OmniDiffusionSamplingParams(num_inference_steps=2)
+            sp.lora_request = lora
+            return _make_step_request(req_id, sampling_params=sp)
+
+        req_a1 = scheduler.add_request(_build("a1", lora_a))
+        req_b1 = scheduler.add_request(_build("b1", lora_b))
+        req_a2 = scheduler.add_request(_build("a2", lora_a))
+
+        # Strict FIFO admission: a1 starts; b1 (different LoRA) blocks the
+        # queue head, so a2 (compatible with a1) is *not* skipped ahead.
+        first = scheduler.schedule()
+        assert _new_ids(first) == [req_a1]
+        assert first.num_waiting_reqs == 2
+
+        # Drain a1 → b1 becomes head-of-line and is admitted with its LoRA.
+        scheduler.update_from_output(first, _make_step_output(req_a1, step_index=2, finished=True))
+        second = scheduler.schedule()
+        assert _new_ids(second) == [req_b1]
+        assert second.num_waiting_reqs == 1
+
+        # Drain b1 → a2 is admitted next; LoRA-A is re-activated for it.
+        scheduler.update_from_output(second, _make_step_output(req_b1, step_index=2, finished=True))
+        third = scheduler.schedule()
+        assert _new_ids(third) == [req_a2]
+        assert third.num_waiting_reqs == 0
+
+    def test_step_batch_separates_requests_with_different_lora_scale(self) -> None:
+        """Same adapter id but different scales → still separate batches."""
+        from vllm_omni.lora.request import LoRARequest
+
+        scheduler = StepScheduler()
+        scheduler.initialize(SimpleNamespace(max_num_seqs=4))
+
+        lora = LoRARequest(lora_name="adapter", lora_int_id=7, lora_path="/tmp/lora")
+
+        def _build(req_id: str, scale: float) -> OmniDiffusionRequest:
+            sp = OmniDiffusionSamplingParams(num_inference_steps=2)
+            sp.lora_request = lora
+            sp.lora_scale = scale
+            return _make_step_request(req_id, sampling_params=sp)
+
+        req_full = scheduler.add_request(_build("full", 1.0))
+        req_half = scheduler.add_request(_build("half", 0.5))
+
+        sched_output = scheduler.schedule()
+
+        admitted = _new_ids(sched_output)
+        assert admitted == [req_full]
+        assert req_half not in admitted
+        assert sched_output.num_waiting_reqs == 1
+
+    def test_step_batch_separates_lora_from_no_lora(self) -> None:
+        """A LoRA request and a no-LoRA request do not share a batch."""
+        from vllm_omni.lora.request import LoRARequest
+
+        scheduler = StepScheduler()
+        scheduler.initialize(SimpleNamespace(max_num_seqs=4))
+
+        lora = LoRARequest(lora_name="adapter", lora_int_id=3, lora_path="/tmp/lora")
+
+        sp_with = OmniDiffusionSamplingParams(num_inference_steps=2)
+        sp_with.lora_request = lora
+        req_with = scheduler.add_request(_make_step_request("with", sampling_params=sp_with))
+        req_without = scheduler.add_request(_make_step_request("without", num_inference_steps=2))
+
+        sched_output = scheduler.schedule()
+
+        admitted = _new_ids(sched_output)
+        assert admitted == [req_with]
+        assert req_without not in admitted
+        assert sched_output.num_waiting_reqs == 1
+
     def test_preempt_request_preserves_step_index(self) -> None:
         request = _make_step_request("preempt", num_inference_steps=3)
         req_id = self.scheduler.add_request(request)

tests/diffusion/test_diffusion_step_pipeline.py

@@ -285,6 +285,19 @@ def _make_scheduler_output(req, sched_req_id="req-1", step_id=0, finished_req_id
     )
 
 
+def _make_batch_scheduler_output(reqs, *, step_id=0, finished_req_ids=None):
+    """Scheduler output for a homogeneous batch (one NewRequestData per req)."""
+    new_reqs = [NewRequestData(sched_req_id=r.request_ids[0], req=r) for r in reqs]
+    return DiffusionSchedulerOutput(
+        step_id=step_id,
+        scheduled_new_reqs=new_reqs,
+        scheduled_cached_reqs=CachedRequestData.make_empty(),
+        finished_req_ids=set() if finished_req_ids is None else set(finished_req_ids),
+        num_running_reqs=len(new_reqs),
+        num_waiting_reqs=0,
+    )
+
+
 def _make_cached_scheduler_output(sched_req_id="req-1", step_id=1, finished_req_ids=None):
     return DiffusionSchedulerOutput(
         step_id=step_id,
@@ -489,71 +502,128 @@ class TestRunner:
             DiffusionModelRunner.load_model(runner)
 
 
+class _RecordingLoRAManager:
+    def __init__(self) -> None:
+        self.calls: list[tuple[object | None, float]] = []
+
+    def set_active_adapter(self, adapter, scale: float = 1.0) -> None:
+        self.calls.append((adapter, scale))
+
+
+def _make_step_worker(lora_manager=None, *, expected_output=None):
+    """Build a bare DiffusionWorker primed for execute_stepwise tests."""
+    worker = object.__new__(DiffusionWorker)
+    worker.lora_manager = lora_manager
+    worker._step_lora_state = {}
+    output = expected_output if expected_output is not None else RunnerOutput(req_id="req-1")
+    worker.model_runner = SimpleNamespace(execute_stepwise=lambda arg: output)
+    return worker
+
+
 @pytest.mark.cpu
 class TestWorker:
     """DiffusionWorker.execute_stepwise"""
 
     def test_delegates_to_model_runner(self):
-        worker = object.__new__(DiffusionWorker)
         expected = RunnerOutput(req_id="req-1", step_index=1, finished=False, result=None)
-        scheduler_output = SimpleNamespace(
-            scheduled_new_reqs=[
-                SimpleNamespace(
-                    req=SimpleNamespace(
-                        sampling_params=SimpleNamespace(lora_request=None),
-                    )
-                )
-            ]
-        )
-        worker.lora_manager = None
-        worker.model_runner = SimpleNamespace(
-            execute_stepwise=lambda arg: expected if arg is scheduler_output else None
-        )
+        worker = _make_step_worker(expected_output=expected)
+        scheduler_output = _make_scheduler_output(_make_engine_request("req-1"), sched_req_id="req-1")
 
         output = DiffusionWorker.execute_stepwise(worker, scheduler_output)
 
         assert output is expected
 
-    def test_clears_active_lora_before_stepwise_execution(self):
-        worker = object.__new__(DiffusionWorker)
-        scheduler_output = SimpleNamespace(
-            scheduled_new_reqs=[
-                SimpleNamespace(
-                    req=SimpleNamespace(
-                        sampling_params=SimpleNamespace(lora_request=None),
-                    )
-                )
-            ]
-        )
-        calls: list[object | None] = []
-
-        class _FakeLoRAManager:
-            def set_active_adapter(self, adapter):
-                calls.append(adapter)
-
-        worker.lora_manager = _FakeLoRAManager()
-        worker.model_runner = SimpleNamespace(execute_stepwise=lambda arg: RunnerOutput(req_id="req-1"))
+    def test_deactivates_lora_when_request_has_no_adapter(self):
+        manager = _RecordingLoRAManager()
+        worker = _make_step_worker(lora_manager=manager)
+        scheduler_output = _make_scheduler_output(_make_engine_request("req-1"), sched_req_id="req-1")
 
         DiffusionWorker.execute_stepwise(worker, scheduler_output)
 
-        assert calls == [None]
+        assert manager.calls == [(None, 1.0)]
 
-    def test_rejects_lora_requests_in_step_mode(self):
-        worker = object.__new__(DiffusionWorker)
-        scheduler_output = SimpleNamespace(
-            scheduled_new_reqs=[
-                SimpleNamespace(
-                    req=SimpleNamespace(
-                        sampling_params=SimpleNamespace(lora_request=object()),
-                    )
-                )
-            ]
+    def test_activates_lora_for_step_requests(self):
+        from vllm_omni.lora.request import LoRARequest
+
+        lora_request = LoRARequest(lora_name="adapter", lora_int_id=7, lora_path="/tmp/lora")
+        request = _make_engine_request("req-1")
+        request.sampling_params.lora_request = lora_request
+        request.sampling_params.lora_scale = 0.75
+
+        manager = _RecordingLoRAManager()
+        worker = _make_step_worker(lora_manager=manager)
+        scheduler_output = _make_scheduler_output(request, sched_req_id="req-1")
+
+        DiffusionWorker.execute_stepwise(worker, scheduler_output)
+
+        assert manager.calls == [(lora_request, 0.75)]
+
+    def test_recovers_lora_for_cached_step_requests(self):
+        from vllm_omni.lora.request import LoRARequest
+
+        lora_request = LoRARequest(lora_name="adapter", lora_int_id=11, lora_path="/tmp/lora")
+        request = _make_engine_request("req-1")
+        request.sampling_params.lora_request = lora_request
+        request.sampling_params.lora_scale = 0.5
+
+        manager = _RecordingLoRAManager()
+        worker = _make_step_worker(lora_manager=manager)
+        first = _make_scheduler_output(request, sched_req_id="req-1")
+        second = _make_cached_scheduler_output(sched_req_id="req-1", step_id=1)
+
+        DiffusionWorker.execute_stepwise(worker, first)
+        DiffusionWorker.execute_stepwise(worker, second)
+
+        assert manager.calls == [(lora_request, 0.5), (lora_request, 0.5)]
+
+    def test_activates_single_lora_for_homogeneous_batch(self):
+        """Multiple requests sharing the same LoRA → exactly one activation,
+        and every request id is registered in ``_step_lora_state``."""
+        from vllm_omni.lora.request import LoRARequest
+
+        lora_request = LoRARequest(lora_name="adapter", lora_int_id=9, lora_path="/tmp/lora")
+        reqs = []
+        for rid in ("req-1", "req-2", "req-3"):
+            r = _make_engine_request(rid)
+            r.sampling_params.lora_request = lora_request
+            r.sampling_params.lora_scale = 0.6
+            reqs.append(r)
+
+        manager = _RecordingLoRAManager()
+        worker = _make_step_worker(lora_manager=manager)
+        scheduler_output = _make_batch_scheduler_output(reqs)
+
+        DiffusionWorker.execute_stepwise(worker, scheduler_output)
+
+        assert manager.calls == [(lora_request, 0.6)]
+        assert set(worker._step_lora_state) == {"req-1", "req-2", "req-3"}
+        for entry in worker._step_lora_state.values():
+            assert entry == (lora_request, 0.6)
+
+    def test_evicts_step_lora_state_for_finished_requests(self):
+        from vllm_omni.lora.request import LoRARequest
+
+        lora_request = LoRARequest(lora_name="adapter", lora_int_id=3, lora_path="/tmp/lora")
+        finishing = _make_engine_request("req-1")
+        finishing.sampling_params.lora_request = lora_request
+        next_request = _make_engine_request("req-2")
+        next_request.sampling_params.lora_request = lora_request
+
+        worker = _make_step_worker(lora_manager=_RecordingLoRAManager())
+        first = _make_scheduler_output(finishing, sched_req_id="req-1")
+        next_batch = _make_scheduler_output(
+            next_request,
+            sched_req_id="req-2",
+            step_id=1,
+            finished_req_ids={"req-1"},
         )
-        worker.lora_manager = None
-        worker.model_runner = SimpleNamespace(execute_stepwise=lambda arg: RunnerOutput(req_id="req-1"))
 
-        with pytest.raises(ValueError, match="does not support LoRA"):
-            DiffusionWorker.execute_stepwise(worker, scheduler_output)
+        DiffusionWorker.execute_stepwise(worker, first)
+        assert "req-1" in worker._step_lora_state
+
+        DiffusionWorker.execute_stepwise(worker, next_batch)
+        assert "req-1" not in worker._step_lora_state
+        assert worker._step_lora_state == {"req-2": (lora_request, 1.0)}
 
 
 @pytest.mark.cpu