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

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

docs/design/feature/diffusion_continuous_batching.md

@@ -83,7 +83,10 @@ There are two important details:
 
 Today that compatibility rule is still shape-sensitive. `height`, `width`,
 `num_frames`, and CFG-related fields remain part of the key, so different
-resolutions or incompatible guidance settings do **not** co-batch yet.
+resolutions or incompatible guidance settings do **not** co-batch yet. The
+key also covers LoRA identity (`lora_int_id`, `lora_scale`), so requests
+targeting different adapters or scales run in separate batches and the
+worker can activate exactly one adapter per step.
 
 The current batching unit is one `OmniDiffusionRequest`. Requests with
 multiple prompts do not participate in batching today.

docs/user_guide/diffusion/step_execution.md

@@ -60,6 +60,9 @@ batching only becomes relevant when `--max-num-seqs > 1`.
 - `cache_backend` is not supported together with step execution.
 - Unsupported pipelines fail early during model loading.
 - Request-mode extras such as KV transfer are not wired into step mode yet.
+- LoRA is supported in step mode, but each batch must use a single adapter:
+  requests with different `lora_request` or `lora_scale` are scheduled into
+  separate batches.
 
 ## When To Use It
 

docs/user_guide/diffusion_features.md

@@ -184,7 +184,7 @@ The following tables show which models support each feature:
 | **💾VAE Patch Parallel** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | | | | |
 | **💾FP8 Quant** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❓ | ❓ | ✅ | ✅ | ✅ | | | |
 | **🔧LoRA Inference** | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | | |
-| **🔄Step Execution** | ❌ | ❌ | ✅ | ✅ | ✅ | ✅ | ❓ | ❓ | ✅ | ❓ | ✅ | ✅ | ❌ | |
+| **🔄Step Execution** | ❌ | ❌ | ✅ | ✅ | ✅ | ✅ | ❓ | ❓ | ✅ | ❓ | ✅ | ✅ | ✅ | |
 
 !!! info
 
@@ -193,7 +193,7 @@ The following tables show which models support each feature:
     3. CPU Offloading (Layerwise) and CPU Offloading (Module-wise) are not compatible.
     4. CPU Offloading (Layerwise) supports single-card for now.
     5. Using FP8-Quant as an example of qunatization methods.
-    6. Step Execution is not compatible with cache backends (TeaCache, Cache-DiT) or LoRA.
+    6. Step Execution is not compatible with cache backends (TeaCache, Cache-DiT). LoRA is supported, but each scheduled batch must use a single adapter (requests with different `lora_request` or `lora_scale` are kept in separate batches).
 
 
 ## Multi-Thread Weight Loading

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

vllm_omni/diffusion/sched/base_scheduler.py

@@ -22,7 +22,9 @@ from vllm_omni.diffusion.sched.interface import (
 
 logger = init_logger(__name__)
 
-_KEY_FIELD_NAMES = frozenset(f.name for f in fields(SamplingParamsKey))
+# LoRA identity is derived from `sampling.lora_request`, not a same-named field
+# on sampling params, so it must be resolved separately from the bulk lookup.
+_KEY_FIELD_NAMES = frozenset(f.name for f in fields(SamplingParamsKey)) - {"lora_int_id"}
 
 
 def get_sampling_params_key(request: OmniDiffusionRequest) -> SamplingParamsKey | None:
@@ -31,7 +33,11 @@ def get_sampling_params_key(request: OmniDiffusionRequest) -> SamplingParamsKey
         return None
 
     sampling = request.sampling_params
-    return SamplingParamsKey(**{name: getattr(sampling, name) for name in _KEY_FIELD_NAMES})
+    lora_request = getattr(sampling, "lora_request", None)
+    return SamplingParamsKey(
+        lora_int_id=lora_request.lora_int_id if lora_request is not None else None,
+        **{name: getattr(sampling, name) for name in _KEY_FIELD_NAMES},
+    )
 
 
 class _BaseScheduler(SchedulerInterface):

vllm_omni/diffusion/sched/interface.py

@@ -65,6 +65,11 @@ class SamplingParamsKey:
     true_cfg_scale: float | None = None
     cfg_normalize: bool = False
 
+    # LoRA identity. Requests with different adapters or scales must run in
+    # separate batches so the worker can activate exactly one adapter per step.
+    lora_int_id: int | None = None
+    lora_scale: float = 1.0
+
 
 @dataclass
 class DiffusionRequestState:

vllm_omni/diffusion/worker/diffusion_worker.py

@@ -185,6 +185,10 @@ class DiffusionWorker:
         self.model_runner: DiffusionModelRunner | None = None
         self._sleep_saved_buffers: dict[str, torch.Tensor] = {}
         self.lora_manager: DiffusionLoRAManager | None = None
+        # Worker-side cache of (lora_request, lora_scale) per scheduled
+        # request id. Used by step mode to recover LoRA identity for cached
+        # requests, which only carry their sched_req_id in subsequent ticks.
+        self._step_lora_state: dict[str, tuple[LoRARequest | None, float]] = {}
         self.stage_id = getattr(od_config, "stage_id", 0)
         self.init_device()
         # Create model runner using the platform-specified class
@@ -367,13 +371,7 @@ class DiffusionWorker:
     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:
-            # Step mode does not support LoRA yet. Clear any previously active
-            # adapter first so worker-local LoRA state cannot leak in.
-            self.lora_manager.set_active_adapter(None)
-
-        if any(new_req.req.sampling_params.lora_request is not None for new_req in scheduler_output.scheduled_new_reqs):
-            raise ValueError("Step mode does not support LoRA yet.")
+        self._activate_step_lora(scheduler_output)
         profiler = self._get_profiler()
         ctx = profiler.annotate_context_manager("diffusion_step") if profiler else nullcontext()
         with ctx:
@@ -382,6 +380,43 @@ class DiffusionWorker:
             profiler.step()
         return output
 
+    def _activate_step_lora(self, scheduler_output: DiffusionSchedulerOutput) -> None:
+        """Activate the LoRA adapter for the scheduled step batch.
+
+        Newly scheduled requests register their (lora_request, lora_scale)
+        in ``_step_lora_state`` so cached requests can resolve to the same
+        adapter on later ticks. Finished requests are evicted. Batch
+        homogeneity is enforced by the scheduler via ``SamplingParamsKey``,
+        so any scheduled request id resolves to the active LoRA identity.
+        """
+        for sched_req_id in scheduler_output.finished_req_ids:
+            self._step_lora_state.pop(sched_req_id, None)
+
+        for new_req in scheduler_output.scheduled_new_reqs:
+            sampling = new_req.req.sampling_params
+            self._step_lora_state[new_req.sched_req_id] = (
+                sampling.lora_request,
+                sampling.lora_scale,
+            )
+
+        if self.lora_manager is None:
+            return
+
+        lora_request: LoRARequest | None = None
+        lora_scale = 1.0
+        for sched_req_id in scheduler_output.scheduled_req_ids:
+            entry = self._step_lora_state.get(sched_req_id)
+            if entry is not None:
+                lora_request, lora_scale = entry
+                break
+
+        try:
+            self.lora_manager.set_active_adapter(lora_request, lora_scale)
+        except Exception as exc:
+            if lora_request is not None:
+                raise
+            logger.warning("LoRA activation skipped: %s", exc)
+
     def load_weights(self, weights) -> set[str]:
         """Load weights by delegating to the model runner."""
         assert self.model_runner is not None, "Model runner not initialized"