From 9b6a308914cf649134ccaf43879fd19a2f43cc51 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=B1=AA=E5=BF=97=E9=B9=8F?= Date: Mon, 7 Sep 2026 22:27:24 +0800 Subject: [PATCH] [Perf][Diffusion] Keep LoRA weights resident across activation cycles (#7195) Signed-off-by: princepride --- tests/diffusion/lora/helpers.py | 26 +++- tests/diffusion/lora/test_lora_manager.py | 125 ++++++++++++++++++ .../diffusion/lora/layers/base_linear.py | 30 +++++ vllm_omni/diffusion/lora/manager.py | 34 ++++- 4 files changed, 211 insertions(+), 4 deletions(-) diff --git a/tests/diffusion/lora/helpers.py b/tests/diffusion/lora/helpers.py index 8b9b1ef4d..4f1dfe176 100644 --- a/tests/diffusion/lora/helpers.py +++ b/tests/diffusion/lora/helpers.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Shared test helpers for diffusion LoRA tests.""" from __future__ import annotations @@ -26,15 +26,39 @@ class DummyBaseLayerWithLoRA(torch.nn.Module): tuple[list[torch.Tensor | None] | torch.Tensor, list[torch.Tensor | None] | torch.Tensor] ] = [] self.reset_calls: int = 0 + self.suspend_calls: int = 0 + self.resume_calls: int = 0 + self.active_slices: tuple[bool, ...] = () + self.suspended_slices: tuple[bool, ...] | None = None self.create_calls: int = 0 def set_lora(self, index: int, lora_a, lora_b): assert index == 0 self.set_calls.append((lora_a, lora_b)) + if isinstance(lora_b, list): + self.active_slices = tuple(b is not None for b in lora_b) + else: + self.active_slices = (True,) + self.suspended_slices = None def reset_lora(self, index: int): assert index == 0 self.reset_calls += 1 + self.active_slices = () + self.suspended_slices = None + + def suspend_lora(self) -> None: + if self.suspended_slices is not None: + return + self.suspended_slices = self.active_slices + self.active_slices = (False,) * len(self.active_slices) + self.suspend_calls += 1 + + def resume_lora(self) -> None: + if self.suspended_slices is not None: + self.active_slices = self.suspended_slices + self.suspended_slices = None + self.resume_calls += 1 def create_lora_weights(self, max_loras, lora_config, model_config): self.create_calls += 1 diff --git a/tests/diffusion/lora/test_lora_manager.py b/tests/diffusion/lora/test_lora_manager.py index 8edc6ce4b..0814bf15f 100644 --- a/tests/diffusion/lora/test_lora_manager.py +++ b/tests/diffusion/lora/test_lora_manager.py @@ -30,14 +30,38 @@ class _DummyLoRALayer: tuple[list[torch.Tensor | None] | torch.Tensor, list[torch.Tensor | None] | torch.Tensor] ] = [] self.reset_calls: int = 0 + self.suspend_calls: int = 0 + self.resume_calls: int = 0 + self.active_slices: tuple[bool, ...] = () + self.suspended_slices: tuple[bool, ...] | None = None def set_lora(self, index: int, lora_a, lora_b): assert index == 0 self.set_calls.append((lora_a, lora_b)) + if isinstance(lora_b, list): + self.active_slices = tuple(b is not None for b in lora_b) + else: + self.active_slices = (True,) + self.suspended_slices = None def reset_lora(self, index: int): assert index == 0 self.reset_calls += 1 + self.active_slices = () + self.suspended_slices = None + + def suspend_lora(self) -> None: + if self.suspended_slices is not None: + return + self.suspended_slices = self.active_slices + self.active_slices = (False,) * len(self.active_slices) + self.suspend_calls += 1 + + def resume_lora(self) -> None: + if self.suspended_slices is not None: + self.active_slices = self.suspended_slices + self.suspended_slices = None + self.resume_calls += 1 # Aliases for backward compatibility within this file @@ -723,3 +747,104 @@ def test_lora_manager_discovers_unet_component(monkeypatch): assert "unet.down_block.proj" in manager._lora_modules # Verify the module was actually replaced in the tree (not just recorded) assert isinstance(pipeline.unet.down_block.proj, _DummyBaseLayerWithLoRA) + + +def _suspend_harness(monkeypatch, adapter_id: int = 7, rank: int = 2): + """Manager wired to a single dummy layer, ready to activate `adapter_id`.""" + import vllm_omni.diffusion.lora.manager as manager_mod + + monkeypatch.setattr(manager_mod, "BaseLayerWithLoRA", _DummyBaseLayerWithLoRA) + + lora_model = _DummyLM(rank=rank) + manager = DiffusionLoRAManager( + pipeline=_DummyPipeline(), + device=torch.device("cpu"), + dtype=torch.bfloat16, + ) + monkeypatch.setattr(manager, "_load_adapter", lambda _req: (lora_model, type("PH", (), {"r": rank})())) + manager._registered_adapters = {adapter_id: lora_model} + manager._lora_modules = {"transformer.foo": lora_model.transformer.foo} + return manager, lora_model.transformer.foo + + +def test_reactivating_same_adapter_resumes_without_rebinding(monkeypatch): + """activate -> deactivate -> activate at the same scale must re-arm the + saved mask instead of re-uploading every layer.""" + adapter_id = 7 + req = _dummy_lora_request(adapter_id) + manager, layer = _suspend_harness(monkeypatch, adapter_id) + + manager.set_active_adapter(req, lora_scale=0.5) + assert len(layer.set_calls) == 1 + assert layer.active_slices == (True,) + + manager.set_active_adapter(None) + assert manager._suspended_adapter_id == adapter_id + assert layer.suspend_calls == 1 + assert layer.active_slices == (False,) + # Suspending must not tear the upload down. + assert layer.reset_calls == 0 + + manager.set_active_adapter(req, lora_scale=0.5) + assert layer.resume_calls == 1 + assert len(layer.set_calls) == 1, "resume must not rebind" + assert layer.active_slices == (True,) + assert manager._active_adapter_id == adapter_id + assert manager._suspended_adapter_id is None + + +@pytest.mark.parametrize("second_scale", [0.5, 0.25]) +def test_scale_or_adapter_change_forces_rebind(monkeypatch, second_scale): + """A different scale, or a different adapter, must take the full bind + path rather than re-arming a stale mask.""" + adapter_id = 7 + manager, layer = _suspend_harness(monkeypatch, adapter_id) + + manager.set_active_adapter(_dummy_lora_request(adapter_id), lora_scale=0.5) + manager.set_active_adapter(None) + assert manager._suspended_adapter_id == adapter_id + + # Same id at a new scale, or a different id entirely. + other_id = adapter_id if second_scale != 0.5 else adapter_id + 1 + if other_id != adapter_id: + manager._registered_adapters[other_id] = manager._registered_adapters[adapter_id] + + manager.set_active_adapter(_dummy_lora_request(other_id), lora_scale=second_scale) + assert len(layer.set_calls) == 2, "must rebind" + assert manager._suspended_adapter_id is None + assert manager._active_adapter_id == other_id + + +def test_partial_packed_mask_survives_suspend_resume(monkeypatch): + """A packed layer where only some slices carry LoRA must come back with + the same mask, not an all-True one.""" + adapter_id = 7 + manager, layer = _suspend_harness(monkeypatch, adapter_id) + + # Emulate a packed bind that left slice 1 empty. + layer.set_lora(0, [torch.ones(2, 2), None], [torch.ones(2, 2), None]) + manager._active_adapter_id = adapter_id + manager._update_adapter_scale(adapter_id, 0.5) + assert layer.active_slices == (True, False) + + manager.set_active_adapter(None) + assert layer.active_slices == (False, False) + + manager.set_active_adapter(_dummy_lora_request(adapter_id), lora_scale=0.5) + assert layer.active_slices == (True, False), "mask must survive unchanged" + + +def test_removing_suspended_adapter_drops_the_upload(monkeypatch): + """A removed adapter can never be resumed, so its weights must not stay + in the stacked buffers.""" + adapter_id = 7 + manager, layer = _suspend_harness(monkeypatch, adapter_id) + + manager.set_active_adapter(_dummy_lora_request(adapter_id), lora_scale=0.5) + manager.set_active_adapter(None) + assert manager._suspended_adapter_id == adapter_id + + manager.remove_adapter(adapter_id) + assert manager._suspended_adapter_id is None + assert layer.reset_calls == 1, "the upload must be torn down" + assert layer.suspended_slices is None diff --git a/vllm_omni/diffusion/lora/layers/base_linear.py b/vllm_omni/diffusion/lora/layers/base_linear.py index c1f941a10..1a78c6ad4 100644 --- a/vllm_omni/diffusion/lora/layers/base_linear.py +++ b/vllm_omni/diffusion/lora/layers/base_linear.py @@ -25,6 +25,9 @@ class DiffusionBaseLinearLayerWithLoRA(BaseLinearLayerWithLoRA): lora_a_stacked: tuple[torch.Tensor, ...] lora_b_stacked: tuple[torch.Tensor, ...] + # Mask saved while the adapter is suspended; None means "not suspended". + _diffusion_lora_suspended_slices: tuple[bool, ...] | None = None + def _move_lora_buffers(self, device: torch.device) -> None: self.lora_a_stacked = tuple(tensor.to(device=device) for tensor in self.lora_a_stacked) self.lora_b_stacked = tuple(tensor.to(device=device) for tensor in self.lora_b_stacked) @@ -54,11 +57,14 @@ class DiffusionBaseLinearLayerWithLoRA(BaseLinearLayerWithLoRA): object.__setattr__(self, "_diffusion_base_layer_ref", base_layer) n_slices = getattr(self, "n_slices", 1) self._diffusion_lora_active_slices = (False,) * int(n_slices) + self._diffusion_lora_suspended_slices = None def reset_lora(self, index: int): super().reset_lora(index) n_slices = getattr(self, "n_slices", 1) self._diffusion_lora_active_slices = (False,) * int(n_slices) + # The upload is gone, so a saved mask no longer describes the buffers. + self._diffusion_lora_suspended_slices = None def set_lora( self, @@ -82,6 +88,30 @@ class DiffusionBaseLinearLayerWithLoRA(BaseLinearLayerWithLoRA): # Single-slice layer. self._diffusion_lora_active_slices = (True,) + def suspend_lora(self) -> None: + """Gate every slice off without dropping the uploaded weights. + + `apply()` reads the stacked buffers only for slices this mask marks + active, so a suspended adapter costs nothing while staying resident. + A saved mask means the layer is already suspended; keep the original + rather than overwriting it with the all-False mask now in place. + """ + if self._diffusion_lora_suspended_slices is not None: + return + self._diffusion_lora_suspended_slices = self._diffusion_lora_active_slices + self._diffusion_lora_active_slices = (False,) * len(self._diffusion_lora_active_slices) + + def resume_lora(self) -> None: + """Restore the mask captured by `suspend_lora()`. + + A layer registered after the suspend, or one whose buffers were torn + down meanwhile, has no mask to restore and keeps its own inactive one. + """ + saved = self._diffusion_lora_suspended_slices + if saved is not None: + self._diffusion_lora_active_slices = saved + self._diffusion_lora_suspended_slices = None + def apply(self, x: torch.Tensor, bias: torch.Tensor | None = None) -> torch.Tensor: """ override: Use simple matmul instead of punica_wrapper.add_lora_linear(). diff --git a/vllm_omni/diffusion/lora/manager.py b/vllm_omni/diffusion/lora/manager.py index 1b63692fc..2f6d47780 100644 --- a/vllm_omni/diffusion/lora/manager.py +++ b/vllm_omni/diffusion/lora/manager.py @@ -48,6 +48,11 @@ class DiffusionLoRAManager: # Valid max allowed ranks for LoRA in vLLM _VALID_MAX_RANKS: list[int] = sorted(get_args(MaxLoRARanks)) + # Adapter whose weights are still uploaded while gated off. Class-level so + # a manager built without __init__ (some tests use object.__new__) reads as + # "nothing suspended" instead of raising. + _suspended_adapter_id: int | None = None + def __init__( self, pipeline: nn.Module, @@ -493,9 +498,11 @@ class DiffusionLoRAManager: fully_sharded_loras=False, ) - # Recreate per-layer buffers with the new maximum rank. + # Recreate per-layer buffers with the new maximum rank. The previous + # upload is gone, so nothing may be re-armed afterwards. for lora_layer in self._lora_modules.values(): lora_layer.create_lora_weights(max_loras=1, lora_config=lora_config, model_config=None) + self._suspended_adapter_id = None # Re-apply active adapter if needed (buffers were reset). if self._active_adapter_id is not None: @@ -679,18 +686,32 @@ class DiffusionLoRAManager: def _reset_lora_layers(self) -> None: for lora_layer in self._lora_modules.values(): lora_layer.reset_lora(0) + self._suspended_adapter_id = None def _activate_adapter(self, adapter_id: int, scale: float) -> None: if self._is_active_at_scale(adapter_id, scale): logger.debug("Adapter %d already active at scale %.3f skipping", adapter_id, scale) return + if self._suspended_adapter_id == adapter_id and self._adapter_scales.get( + adapter_id + ) == DiffusionLoRAManager._get_rounded_scale(scale): + # Weights are still uploaded from an earlier activation; re-arming + # the masks avoids rebuilding and re-uploading every layer. + for lora_layer in self._lora_modules.values(): + lora_layer.resume_lora() + self._suspended_adapter_id = None + self._active_adapter_id = adapter_id + logger.debug("Re-armed suspended adapter %d", adapter_id) + return + logger.info("Activating adapter: id=%d", adapter_id) lora_model = self._registered_adapters[adapter_id] # Binding overwrites slot 0 incrementally. Invalidate the fast-path # state before the first mutation and leave every wrapper inactive if # any set_lora() call or model validator fails. self._active_adapter_id = None + self._suspended_adapter_id = None try: self._bind_adapter_weights(lora_model, scale) except Exception: @@ -704,8 +725,10 @@ class DiffusionLoRAManager: if self._active_adapter_id is None: logger.debug("All adapters already inactive") return - logger.info("Deactivating all adapters: %d layers", len(self._lora_modules)) - self._reset_lora_layers() + logger.info("Suspending all adapters: %d layers", len(self._lora_modules)) + for lora_layer in self._lora_modules.values(): + lora_layer.suspend_lora() + self._suspended_adapter_id = self._active_adapter_id self._active_adapter_id = None logger.debug("All adapters deactivated") @@ -772,6 +795,11 @@ class DiffusionLoRAManager: if self._active_adapter_id == adapter_id: self._deactivate_all_adapters() + if self._suspended_adapter_id == adapter_id: + # The adapter is going away, so the upload can never be resumed. + # Tear it down instead of leaving it in the stacked buffers. + self._reset_lora_layers() + del self._registered_adapters[adapter_id] self._adapter_scales.pop(adapter_id, None) self._adapter_access_order.pop(adapter_id, None)