[Perf][Diffusion] Keep LoRA weights resident across activation cycles (#7195)

Signed-off-by: princepride <wangzhipeng628@gmail.com>

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

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

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().

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)