diff --git a/docs/design/feature/diffusion_continuous_batching.md b/docs/design/feature/diffusion_continuous_batching.md index ab8ff69fe..c443dea78 100644 --- a/docs/design/feature/diffusion_continuous_batching.md +++ b/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. diff --git a/docs/user_guide/diffusion/step_execution.md b/docs/user_guide/diffusion/step_execution.md index 189510916..07184a787 100644 --- a/docs/user_guide/diffusion/step_execution.md +++ b/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 diff --git a/docs/user_guide/diffusion_features.md b/docs/user_guide/diffusion_features.md index 26b4d9bf7..61694ac9d 100644 --- a/docs/user_guide/diffusion_features.md +++ b/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 diff --git a/tests/diffusion/test_diffusion_scheduler.py b/tests/diffusion/test_diffusion_scheduler.py index f0bc1da35..ca28b2629 100644 --- a/tests/diffusion/test_diffusion_scheduler.py +++ b/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) diff --git a/tests/diffusion/test_diffusion_step_pipeline.py b/tests/diffusion/test_diffusion_step_pipeline.py index 1954f4a05..2722de7e3 100644 --- a/tests/diffusion/test_diffusion_step_pipeline.py +++ b/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 diff --git a/vllm_omni/diffusion/sched/base_scheduler.py b/vllm_omni/diffusion/sched/base_scheduler.py index f21c03e98..0bd16597f 100644 --- a/vllm_omni/diffusion/sched/base_scheduler.py +++ b/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): diff --git a/vllm_omni/diffusion/sched/interface.py b/vllm_omni/diffusion/sched/interface.py index 8eef211ff..1bfb69458 100644 --- a/vllm_omni/diffusion/sched/interface.py +++ b/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: diff --git a/vllm_omni/diffusion/worker/diffusion_worker.py b/vllm_omni/diffusion/worker/diffusion_worker.py index e2c4b97c1..cd5fe246d 100644 --- a/vllm_omni/diffusion/worker/diffusion_worker.py +++ b/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"