[Refactor][Diffusion] Complete generic backend plan consumption (#7313)

Signed-off-by: specture724 <specture724@gmail.com>
Co-authored-by: Claude Opus 5 <noreply@anthropic.com>

docs/design/feature/offloader/README.md

@@ -41,7 +41,8 @@ Every backend implements `OffloadBackend`:
 - `is_enabled()` reports lifecycle state.
 
 `disable()` does not promise to restore the pipeline's original device
-placement. The caller owns any subsequent rematerialization.
+placement. It must restore usable parameter and buffer storage before releasing
+backend-owned host weights. The caller owns subsequent device placement.
 
 An enable failure must remove every partially installed hook. Rank-local
 backends restore ordinary tensors so the pipeline can be retried. A failed
@@ -54,9 +55,19 @@ must use distinct hook names and remove only hooks they own.
 ## Discovery and topology
 
 `resolve_offload_plan(pipeline, config)` is the single topology entry point.
-Every backend calls it once when `enable()` starts and then consumes the
-resolved components, block stacks, resident heads, and staged-residency flags
-it returns. Backends do not re-read model declarations.
+The generic model-level and ordinary layerwise paths call it once when
+`enable()` starts and consume its component selection and block stacks.
+Ordinary layerwise offload also respects a resolved `skip_reason`, preserving
+the legacy missing-DiT no-op without placing auxiliary components. Non-block
+DiT and encoder state is placed by resolved block tensor identity, so the
+backend does not rediscover block containers by attribute name.
+
+Two compatibility paths remain outside this completed generic cutover:
+model-level pipelines implementing `SupportsModelCpuOffload` still own their
+custom phase lifecycle (RFC #6648 J4), and distributed layerwise offload still
+reads declarations for nested-submodule staging and loader-owned mmap paths
+(J3). The custom model-level protocol is delegated to before generic plan
+resolution; it must not be subjected to the generic encoder/DiT swap checks.
 
 Pipeline component discovery prefers `SupportsComponentDiscovery` declarations:
 
@@ -102,3 +113,17 @@ remains responsible for transfer, synchronization, and storage ownership.
 
 The diffusion [Offloader module design](../../module/diffusion/offloader.md)
 describes how these feature contracts fit into the larger diffusion runtime.
+
+## Backend acceptance tests
+
+`tests/diffusion/offloader/test_backend_plan_contract.py` covers the two
+behaviors the generic cutover changed. A DiT attribute aliasing a streamed
+block must keep the ring's host residency, so placement cannot follow attribute
+names; the CUDA case fails against the previous name-based placement. Both
+backends must then run with their selector helpers and the model declaration
+made unreadable, proving execution consumes the resolved plan alone.
+
+Selection errors, rollback, residency, and enable/disable cycles stay in
+`test_plan_resolver.py`, `test_layerwise_backend.py` and
+`test_sequential_backend.py`. Model-specific phases and executed tied aliases
+remain the J4 contract suite.

tests/diffusion/models/flux2/test_flux2_transformer_tp.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 import pytest
 import torch
 from pytest_mock import MockerFixture
@@ -265,7 +268,7 @@ class TestFlux2TransformerLayerwiseOffload:
         ]
 
     def test_get_blocks_from_dit(self):
-        from vllm_omni.diffusion.offloader.layerwise_backend import LayerWiseOffloadBackend
+        from vllm_omni.diffusion.offloader.block_discovery import get_blocks_from_dit
 
         # Block discovery only needs the ModuleList structure, not production dims.
         model = Flux2Transformer2DModel(
@@ -275,6 +278,6 @@ class TestFlux2TransformerLayerwiseOffload:
             attention_head_dim=4,
             joint_attention_dim=16,
         )
-        attr_names, blocks = LayerWiseOffloadBackend.get_blocks_from_dit(model)
+        attr_names, blocks = get_blocks_from_dit(model)
         assert attr_names == ["transformer_blocks", "single_transformer_blocks"]
         assert len(blocks) == 4

tests/diffusion/offloader/test_backend_plan_contract.py

@@ -0,0 +1,163 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Executable contracts for the two behaviors the generic cutover changed.
+
+Placement now follows resolved block identity, and the backends run on the
+resolved plan alone. Declaration parsing, selection errors, rollback,
+residency and enable/disable cycles are covered by ``test_plan_resolver.py``,
+``test_layerwise_backend.py`` and ``test_sequential_backend.py``.
+"""
+
+from typing import ClassVar
+
+import pytest
+import torch
+from torch import nn
+
+from tests.diffusion.offloader.helpers import patch_offload_runtime
+from vllm_omni.diffusion.offloader import layerwise_backend, sequential_backend
+from vllm_omni.diffusion.offloader.base import OffloadConfig, OffloadStrategy
+from vllm_omni.diffusion.offloader.layerwise_backend import LayerWiseOffloadBackend
+from vllm_omni.diffusion.offloader.offload_plan import OffloadPlan
+from vllm_omni.diffusion.offloader.plan_resolver import resolve_offload_plan
+from vllm_omni.diffusion.offloader.sequential_backend import ModelLevelOffloadBackend
+from vllm_omni.platforms import current_omni_platform
+
+pytestmark = [pytest.mark.diffusion, pytest.mark.core_model]
+
+STRATEGIES = [OffloadStrategy.MODEL_LEVEL, OffloadStrategy.LAYER_WISE]
+CPU = torch.device("cpu")
+
+
+class _Block(nn.Linear):
+    def __init__(self):
+        super().__init__(4, 4)
+        self.register_buffer("scale", torch.tensor(0.5))
+
+    def forward(self, x):
+        return super().forward(x).tanh() * self.scale
+
+
+class _Stack(nn.Module):
+    """Two block containers plus non-block state that must stay resident."""
+
+    def __init__(self):
+        super().__init__()
+        self.blocks = nn.ModuleList([_Block() for _ in range(3)])
+        self.tail = nn.ModuleList([_Block() for _ in range(2)])
+        self.bias = nn.Parameter(torch.ones(4))
+        self.register_buffer("scale", torch.tensor(0.25))
+        self.proj = nn.Linear(4, 4)
+
+    def forward(self, x):
+        for block in (*self.blocks, *self.tail):
+            x = block(x)
+        return self.proj(x) * self.scale + self.bias
+
+
+class _Pipeline(nn.Module):
+    _dit_modules: ClassVar[list[str]] = ["transformer"]
+    _encoder_modules: ClassVar[list[str]] = ["text_encoder", "image_encoder"]
+    _vae_modules: ClassVar[list[str]] = ["vae"]
+    _resident_modules: ClassVar[list[str]] = ["resident"]
+    _offload_plan = OffloadPlan(
+        block_attrs={"transformer": ("blocks", "tail")},
+        encoder_block_attrs={"text_encoder": ("blocks", "tail")},
+    )
+
+    def __init__(self):
+        super().__init__()
+        self.transformer = _Stack()
+        self.text_encoder = _Stack()
+        self.image_encoder = nn.Linear(4, 4)
+        self.vae = nn.Linear(4, 4)
+        self.resident = nn.Linear(4, 4)
+
+    def forward(self, x):
+        return self.resident(self.vae(self.transformer(self.text_encoder(self.image_encoder(x)))))
+
+
+@pytest.fixture(params=[pytest.param("cpu", marks=pytest.mark.cpu), pytest.param("cuda", marks=pytest.mark.cuda)])
+def execution_device(request, monkeypatch):
+    if request.param == "cuda":
+        if not torch.cuda.is_available():
+            pytest.skip("CUDA required for real transfer/residency coverage")
+        return torch.device("cuda:0")
+    patch_offload_runtime(monkeypatch, current_omni_platform, synchronize=True)
+    monkeypatch.setattr(current_omni_platform, "get_free_memory", lambda: 0)
+    return CPU
+
+
+def _config(strategy, components):
+    return OffloadConfig(strategy=strategy, components=components, pin_cpu_memory=False)
+
+
+def _backend(config, device):
+    kind = ModelLevelOffloadBackend if config.strategy is OffloadStrategy.MODEL_LEVEL else LayerWiseOffloadBackend
+    return kind(config, device)
+
+
+def _assert_no_offload_hooks(pipeline):
+    for module in pipeline.modules():
+        registry = getattr(module, "_hook_registry", None)
+        if registry is not None:
+            assert registry.get_hook("sequential_offload") is None
+            assert registry.get_hook("layerwise_offload") is None
+        assert not getattr(module, "_omni_layerwise_enabled", False)
+
+
+@torch.inference_mode()
+def test_alias_of_a_streamed_block_is_not_placed(execution_device):
+    """Placement moves DiT state by resolved block identity, not by attribute.
+
+    An attribute aliasing a streamed block must stay with its ring: copying it
+    to the device makes the ring record the device as that block's home and
+    strands those weights there after teardown.
+    """
+    pipeline = _Pipeline()
+    pipeline.transformer.proj.weight = pipeline.transformer.blocks[1].weight
+    original = pipeline.transformer.proj.weight.clone()
+    backend = _backend(_config(OffloadStrategy.LAYER_WISE, frozenset({"dit"})), execution_device)
+
+    backend.enable(pipeline)
+    try:
+        assert pipeline.transformer.proj.weight is pipeline.transformer.blocks[1].weight
+        assert pipeline.transformer.proj.weight.device == CPU
+        # Non-block state around the alias is still placed.
+        assert pipeline.transformer.proj.bias.device == execution_device
+        assert pipeline.transformer.bias.device == execution_device
+    finally:
+        backend.disable()
+
+    assert pipeline.transformer.proj.weight.device == CPU
+    torch.testing.assert_close(pipeline.transformer.proj.weight, original)
+
+
+@pytest.mark.parametrize("strategy", STRATEGIES)
+@torch.inference_mode()
+def test_backends_execute_the_resolved_plan_alone(execution_device, strategy, monkeypatch):
+    pipeline = _Pipeline().to(execution_device)
+    x = torch.randn(2, 4, device=execution_device)
+    expected = pipeline(x)
+    config = _config(strategy, frozenset({"dit", "text_encoder"}))
+    resolved = resolve_offload_plan(pipeline, config)
+
+    def forbidden_read(*args, **kwargs):
+        pytest.fail("Backend reinterpreted topology after plan resolution")
+
+    # Once resolved, neither declarations nor selector helpers are an input to
+    # backend execution. Transport options stay readable.
+    monkeypatch.setattr(config, "offloads", forbidden_read)
+    monkeypatch.setattr(config, "should_offload_encoder", forbidden_read)
+    monkeypatch.setattr(pipeline, "_offload_plan", None)
+    backend_module = sequential_backend if strategy is OffloadStrategy.MODEL_LEVEL else layerwise_backend
+    monkeypatch.setattr(backend_module, "resolve_offload_plan", lambda *_: resolved)
+
+    backend = _backend(config, execution_device)
+    try:
+        backend.enable(pipeline)
+        torch.testing.assert_close(pipeline(x), expected)
+    finally:
+        backend.disable()
+    _assert_no_offload_hooks(pipeline)

tests/diffusion/offloader/test_layerwise_backend.py

@@ -24,6 +24,11 @@ from tests.diffusion.offloader.helpers import (
     patch_offload_runtime,
 )
 from vllm_omni.diffusion.offloader.base import OffloadConfig, OffloadStrategy
+from vllm_omni.diffusion.offloader.block_discovery import (
+    get_blocks_attr_names,
+    get_blocks_from_dit,
+    set_blocks_attr_names,
+)
 from vllm_omni.diffusion.offloader.layerwise_backend import (
     LayerWiseOffloadBackend,
     LayerwiseOffloadHook,
@@ -209,21 +214,21 @@ class _NoAttrsModel(nn.Module):
 class TestGetBlocksFromDit:
     def test_get_blocks_from_dit_single_block_attr(self):
         model = _SingleBlockModel(num_blocks=3)
-        attr_names, blocks = LayerWiseOffloadBackend.get_blocks_from_dit(model)
+        attr_names, blocks = get_blocks_from_dit(model)
         assert attr_names == ["blocks"]
         assert len(blocks) == 3
         assert all(isinstance(b, _DummyBlock) for b in blocks)
 
     def test_get_blocks_from_dit_multi_block_attrs(self):
         model = _MultiBlockModel(num_transformer=2, num_single=3)
-        attr_names, blocks = LayerWiseOffloadBackend.get_blocks_from_dit(model)
+        attr_names, blocks = get_blocks_from_dit(model)
         assert set(attr_names) == {"transformer_blocks", "single_transformer_blocks"}
         assert len(blocks) == 5
         assert all(isinstance(b, _DummyBlock) for b in blocks)
 
     def test_get_blocks_from_dit_empty_blocks(self):
         model = _EmptyBlocksModel()
-        attr_names, blocks = LayerWiseOffloadBackend.get_blocks_from_dit(model)
+        attr_names, blocks = get_blocks_from_dit(model)
         assert attr_names == []
         assert blocks == []
 
@@ -233,17 +238,17 @@ class TestGetBlocksFromDit:
             AttributeError,
             match="Attribute 'nonexistent_blocks' declared in _layerwise_offload_blocks_attrs does not exist",
         ):
-            LayerWiseOffloadBackend.get_blocks_from_dit(model)
+            get_blocks_from_dit(model)
 
     def test_get_blocks_from_dit_no_attrs_defined(self):
         model = _NoAttrsModel(num_blocks=3)
-        attr_names, blocks = LayerWiseOffloadBackend.get_blocks_from_dit(model)
+        attr_names, blocks = get_blocks_from_dit(model)
         assert attr_names == []
         assert blocks == []
 
     def test_get_blocks_from_dit_deprecated_single_attr(self):
         model = _DeprecatedSingleAttrModel(num_blocks=2)
-        attr_names, blocks = LayerWiseOffloadBackend.get_blocks_from_dit(model)
+        attr_names, blocks = get_blocks_from_dit(model)
         assert attr_names == ["blocks"]
         assert len(blocks) == 2
 
@@ -251,17 +256,17 @@ class TestGetBlocksFromDit:
 class TestGetBlocksAttrNames:
     def test_get_blocks_attr_names_new_format(self):
         model = _MultiBlockModel()
-        attrs = LayerWiseOffloadBackend.get_blocks_attr_names(model)
+        attrs = get_blocks_attr_names(model)
         assert attrs == ["transformer_blocks", "single_transformer_blocks"]
 
     def test_get_blocks_attr_names_no_attrs(self):
         model = _NoAttrsModel()
-        attrs = LayerWiseOffloadBackend.get_blocks_attr_names(model)
+        attrs = get_blocks_attr_names(model)
         assert attrs == []
 
     def test_set_blocks_attr_names(self):
         model = _NoAttrsModel()
-        LayerWiseOffloadBackend.set_blocks_attr_names(model, ["new_blocks"])
+        set_blocks_attr_names(model, ["new_blocks"])
         assert hasattr(model.__class__, "_layerwise_offload_blocks_attrs")
         assert model.__class__._layerwise_offload_blocks_attrs == ["new_blocks"]
 

vllm_omni/diffusion/offloader/layerwise_backend.py

@@ -14,11 +14,6 @@ from vllm_omni.diffusion.hooks import HookRegistry, ModelHook
 from vllm_omni.platforms import current_omni_platform
 
 from .base import OffloadBackend, OffloadConfig, run_cleanup_steps
-from .block_discovery import (
-    get_blocks_attr_names,
-    get_blocks_from_dit,
-    set_blocks_attr_names,
-)
 from .component_utils import (
     clear_encoder_layerwise_state,
     iter_streamable_dits,
@@ -26,7 +21,6 @@ from .component_utils import (
     prepare_pipeline_components,
     set_encoder_layerwise_state,
 )
-from .config import DIT_COMPONENT
 from .plan_resolver import ResolvedComponent, resolve_offload_plan
 from .tensor_utils import (
     clear_block_storage,
@@ -433,9 +427,8 @@ class LayerWiseOffloadBackend(OffloadBackend):
             return
 
         resolved = resolve_offload_plan(pipeline, self.config)
-        if not resolved.dits and self.config.offloads(DIT_COMPONENT):
-            # An explicit selection already failed while resolving.
-            logger.warning("No DiT/transformer modules found for selected DiT layerwise offload")
+        if resolved.skip_reason is not None:
+            logger.warning("%s", resolved.skip_reason)
             return
 
         def enable_encoder_blocks(component: ResolvedComponent) -> bool:
@@ -458,15 +451,6 @@ class LayerWiseOffloadBackend(OffloadBackend):
             enable_encoder_blocks=enable_encoder_blocks,
         )
 
-        if not self.config.offloads(DIT_COMPONENT):
-            self.enabled = bool(self._encoder_modules or self._staged_components)
-            if not self.enabled:
-                raise ValueError(
-                    "None of the selected layerwise offload components have "
-                    "a model-declared streamable or on-demand plan"
-                )
-            return
-
         logger.info("Applying layer-wise offloading on %s", [component.path for component in resolved.dits])
 
         # Apply block-wise offloading hook for each of the blocks in DiT model(s)
@@ -475,22 +459,9 @@ class LayerWiseOffloadBackend(OffloadBackend):
             dit_module = component.module
             blocks = list(stack.blocks)
 
-            # Move non-block modules to GPU (they stay resident)
-            for name, m in dit_module.named_children():
-                if name not in stack.attrs:
-                    m.to(self.device)
-                    logger.debug(f"Moved {name} to device {self.device}")
-                else:
-                    logger.debug(f"Skipped blocks module {name}")
-
-            # Move top-level params/buffers to GPU (dit_module's own, not sub-modules)
-            for param in dit_module._parameters.values():
-                if param is not None:
-                    param.data = param.data.to(self.device, non_blocking=True)
-
-            for buffer in dit_module._buffers.values():
-                if buffer is not None:
-                    buffer.data = buffer.data.to(self.device, non_blocking=True)
+            # Place the remainder by resolved block tensor identity, just as
+            # for encoders. Attribute aliases must not move streamed weights.
+            move_non_block_state_to_device(dit_module, (stack.blocks,), self.device)
 
             block_hooks = _install_layerwise_hook_group(
                 blocks,
@@ -550,8 +521,3 @@ class LayerWiseOffloadBackend(OffloadBackend):
 
     def disable(self) -> None:
         self._disable(restore_weights=True)
-
-    # Compatibility aliases for existing model integrations.
-    get_blocks_attr_names = staticmethod(get_blocks_attr_names)
-    set_blocks_attr_names = staticmethod(set_blocks_attr_names)
-    get_blocks_from_dit = staticmethod(get_blocks_from_dit)

vllm_omni/diffusion/offloader/plan_resolver.py

@@ -91,6 +91,11 @@ class ResolvedOffloadPlan:
     # still interpret topology themselves.
     modules: PipelineModules
     declaration: OffloadPlan | None = None
+    # Ordinary layerwise offload turns a legacy selection it cannot serve into
+    # a complete no-op, placement included. Distributed layerwise offload warns
+    # and keeps preparing its other components instead, so this stays scoped to
+    # the strategy that honors it.
+    skip_reason: str | None = None
 
     @property
     def components(self) -> tuple[ResolvedComponent, ...]:
@@ -304,6 +309,11 @@ def resolve_offload_plan(pipeline: nn.Module, config: OffloadConfig) -> Resolved
         residents=tuple(residents),
         modules=modules,
         declaration=declaration,
+        skip_reason=(
+            "No DiT/transformer modules found for selected DiT layerwise offload"
+            if config.strategy is OffloadStrategy.LAYER_WISE and dit_selected and not dits
+            else None
+        ),
     )
     _validate_unique_ownership(resolved)
     return resolved

vllm_omni/diffusion/offloader/sequential_backend.py

@@ -13,7 +13,6 @@ from vllm_omni.diffusion.hooks import HookRegistry, ModelHook
 from vllm_omni.platforms import current_omni_platform
 
 from .base import OffloadBackend, OffloadConfig, SupportsModelCpuOffload
-from .config import DIT_COMPONENT
 from .plan_resolver import resolve_offload_plan
 
 logger = init_logger(__name__)
@@ -333,6 +332,7 @@ class ModelLevelOffloadBackend(OffloadBackend):
         encoders = [component.module for component in resolved.encoders]
         vaes = [component.module for component in resolved.vaes]
         residents = [component.module for component in resolved.residents]
+        selected_dits = [component.module for component in resolved.dits if component.selected]
         selected_encoders = [component.module for component in resolved.encoders if component.selected]
 
         all_modules = [*dits, *encoders, *vaes, *residents]
@@ -360,7 +360,7 @@ class ModelLevelOffloadBackend(OffloadBackend):
                 device=self.device,
                 pin_memory=self.config.pin_cpu_memory,
                 use_hsdp=self.config.use_hsdp,
-                offload_dit_modules=(dits if self.config.offloads(DIT_COMPONENT) else ()),
+                offload_dit_modules=selected_dits,
                 offload_encoder_modules=selected_encoders,
             )
         except BaseException: