[Refactor][Diffusion] Resolve offload topology through one plan resolver (#7209)

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

docs/design/feature/offloader/README.md

@@ -53,6 +53,11 @@ 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.
+
 Pipeline component discovery prefers `SupportsComponentDiscovery` declarations:
 
 - `_dit_modules`;
@@ -61,16 +66,26 @@ Pipeline component discovery prefers `SupportsComponentDiscovery` declarations:
 - `_resident_modules`.
 
 Dotted paths are supported. Legacy pipelines may use the fallback scan of
-well-known attribute names, but new integrations should declare components
-explicitly.
+well-known attribute names; it now warns once per pipeline class, and new
+integrations must declare components explicitly.
 
-Both layerwise backends first consume pipeline `OffloadPlan` metadata. A
-plan's `block_attrs` maps each DiT path to its ordered block containers, while
-`encoder_block_attrs` declares streamable encoder stacks. DiTs absent from the
-plan fall back to `_layerwise_offload_blocks_attrs` (including the deprecated
-singular-name compatibility path). Discovery metadata describes structure
-only; the backend remains responsible for transfer, synchronization, and
-storage ownership.
+Block topology comes from the pipeline `OffloadPlan` first: `block_attrs` maps
+each DiT path to its ordered block containers, `encoder_block_attrs` declares
+streamable encoder stacks, `on_demand_component_paths` marks pipeline-managed
+residency, `resident_dit_paths` marks the DiTs that may hold resident layers,
+and `encoder_dlo_weight_replication` marks the encoders whose loader-produced
+weights are safe for AllGather. DiTs absent from the plan fall back to
+`_layerwise_offload_blocks_attrs` (including the deprecated singular-name
+compatibility path).
+
+The resolver is pure: it moves no tensor, installs no hook, writes no module
+attribute, and reads no process group — multi-rank facts come from
+`OffloadConfig`. It owns topology validation, so a configuration the model
+cannot serve fails before the first component is placed. Explicit component
+selection turns every declaration problem into an error, while the
+compatibility topology (no component selector) keeps the historical
+warn-and-skip behavior. Metadata still describes structure only; the backend
+remains responsible for transfer, synchronization, and storage ownership.
 
 ## Cross-strategy invariants
 

tests/diffusion/offloader/test_distributed_layerwise_backend.py

@@ -52,6 +52,7 @@ from vllm_omni.diffusion.offloader.offload_plan import (
     OffloadPlan,
     get_offload_plan,
 )
+from vllm_omni.diffusion.offloader.plan_resolver import resolve_offload_plan
 from vllm_omni.diffusion.offloader.startup import OffloadStartupState, attach_offload_startup_state
 from vllm_omni.host_weight_runtime import MappedHostRegion
 from vllm_omni.platforms import current_omni_platform
@@ -2813,18 +2814,21 @@ class TestDistributedComponentSelection:
 
     def test_encoder_allgather_rejects_stub_rank_before_block_discovery(self):
         """Every rank must reject an unsafe encoder group, including stub ranks."""
-        backend = DistributedLayerwiseOffloadBackend(
-            OffloadConfig(
-                strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
-                pin_cpu_memory=False,
-                dp_size=2,
-                components=frozenset({"text_encoder"}),
-                dlo_transfers={"dit": "rank-local", "text_encoder": "allgather"},
-            ),
-            torch.device("cpu"),
+
+        class StubEncoderPipeline(nn.Module):
+            _offload_plan = OffloadPlan(encoder_block_attrs={"text_encoder": ("missing.blocks",)})
+
+            def __init__(self):
+                super().__init__()
+                self.text_encoder = nn.Module()
+
+        config = OffloadConfig(
+            strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
+            pin_cpu_memory=False,
+            dp_size=2,
+            components=frozenset({"text_encoder"}),
+            dlo_transfers={"dit": "rank-local", "text_encoder": "allgather"},
         )
-        backend.dp_group = object()
-        stub_plan = OffloadPlan(encoder_block_attrs={"text_encoder": ("missing.blocks",)})
 
         with pytest.raises(ValueError, match="not declared replicated"):
-            backend._try_layerwise_offload_encoder(nn.Module(), "text_encoder", stub_plan)
+            resolve_offload_plan(StubEncoderPipeline(), config)

tests/diffusion/offloader/test_layerwise_backend.py

@@ -488,7 +488,8 @@ class TestLayerwiseComponentSelection:
         with pytest.raises(ValueError, match="Selected text encoder 'text_encoder_2' requires"):
             backend.enable(pipeline)
 
-        assert not pipeline.text_encoder._omni_layerwise_enabled
+        # The plan resolver rejects the topology before any encoder is hooked.
+        assert not getattr(pipeline.text_encoder, "_omni_layerwise_enabled", False)
 
     def test_default_selection_preserves_unplanned_auxiliaries(self, patched_offload_runtime):
         pipeline = _LegacyComponentPipeline()

tests/diffusion/offloader/test_plan_resolver.py

@@ -0,0 +1,458 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Unit tests for the shared offload plan resolver."""
+
+from typing import ClassVar
+
+import pytest
+import torch
+from torch import nn
+
+from tests.diffusion.offloader.helpers import (
+    _DummyBlock,
+    _PlainEncoder,
+    _SingleBlockModel,
+    _StagedEncoder,
+    _StagedVAE,
+    patch_offload_runtime,
+)
+from vllm_omni.diffusion.offloader import layerwise_backend as layerwise_backend_module
+from vllm_omni.diffusion.offloader import plan_resolver as plan_resolver_module
+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
+
+pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
+
+CPU = torch.device("cpu")
+
+
+@pytest.fixture
+def patched_offload_runtime(monkeypatch):
+    patch_offload_runtime(monkeypatch, layerwise_backend_module.current_omni_platform)
+
+
+def _config(
+    components: set[str] | None = None,
+    strategy: OffloadStrategy = OffloadStrategy.LAYER_WISE,
+    **kwargs,
+) -> OffloadConfig:
+    return OffloadConfig(
+        strategy=strategy,
+        pin_cpu_memory=False,
+        components=None if components is None else frozenset(components),
+        **kwargs,
+    )
+
+
+def _stack_of(component):
+    assert len(component.stacks) == 1
+    return component.stacks[0]
+
+
+class _WanStylePipeline(nn.Module):
+    """Streamable DiT plus a plan-declared, replicated text encoder."""
+
+    _dit_modules: ClassVar[list[str]] = ["transformer"]
+    _encoder_modules: ClassVar[list[str]] = ["text_encoder"]
+    _vae_modules: ClassVar[list[str]] = ["vae"]
+    _resident_modules: ClassVar[list[str]] = []
+    _offload_plan = OffloadPlan(
+        encoder_component_types={"text_encoder": "text_encoder"},
+        encoder_block_attrs={"text_encoder": ("encoder.block",)},
+        encoder_dlo_weight_replication=frozenset({"text_encoder"}),
+    )
+
+    def __init__(self):
+        super().__init__()
+        self.transformer = _SingleBlockModel(num_blocks=4)
+        self.text_encoder = _PlainEncoder()
+        self.vae = _StagedVAE()
+
+
+class _StagedComponentPipeline(nn.Module):
+    """MiniMax-H3 shape: two encoder stacks plus pipeline-staged components."""
+
+    _dit_modules: ClassVar[list[str]] = ["transformer"]
+    _encoder_modules: ClassVar[list[str]] = ["text_encoder"]
+    _vae_modules: ClassVar[list[str]] = ["vae"]
+    _resident_modules: ClassVar[list[str]] = []
+    _offload_plan = OffloadPlan(
+        encoder_component_types={"text_encoder": "text_encoder"},
+        encoder_block_attrs={"text_encoder": ("vision.blocks", "text_model.layers")},
+        on_demand_component_paths=frozenset({"text_encoder", "vae"}),
+        resident_dit_paths=frozenset({"transformer"}),
+    )
+
+    def __init__(self):
+        super().__init__()
+        self.transformer = _SingleBlockModel(num_blocks=4)
+        self.text_encoder = _StagedEncoder()
+        self.vae = _StagedVAE()
+
+
+class _LegacyScanPipeline(nn.Module):
+    """No SupportsComponentDiscovery declaration and no OffloadPlan."""
+
+    def __init__(self):
+        super().__init__()
+        self.transformer = _SingleBlockModel(num_blocks=3)
+        self.text_encoder = _PlainEncoder()
+
+
+class TestProducerPrecedence:
+    def test_plan_block_attrs_win_over_class_attributes(self):
+        class TwoStackDiT(nn.Module):
+            _layerwise_offload_blocks_attrs = ["blocks"]
+
+            def __init__(self):
+                super().__init__()
+                self.blocks = nn.ModuleList([_DummyBlock() for _ in range(2)])
+                self.alt_blocks = nn.ModuleList([_DummyBlock() for _ in range(3)])
+
+        class Pipeline(nn.Module):
+            _offload_plan = OffloadPlan(block_attrs={"transformer": ("alt_blocks",)})
+
+            def __init__(self):
+                super().__init__()
+                self.transformer = TwoStackDiT()
+
+        pipeline = Pipeline()
+        resolved = resolve_offload_plan(pipeline, _config({"dit"}))
+
+        stack = _stack_of(resolved.dits[0])
+        assert stack.attrs == ("alt_blocks",)
+        assert list(stack.blocks) == list(pipeline.transformer.alt_blocks)
+
+    def test_class_attributes_are_used_without_a_plan(self):
+        pipeline = _LegacyScanPipeline()
+        resolved = resolve_offload_plan(pipeline, _config({"dit"}))
+
+        assert list(_stack_of(resolved.dits[0]).blocks) == list(pipeline.transformer.blocks)
+
+    def test_declared_component_type_wins_over_leaf_name(self):
+        class Pipeline(nn.Module):
+            _dit_modules: ClassVar[list[str]] = ["transformer"]
+            _encoder_modules: ClassVar[list[str]] = ["mllm"]
+            _vae_modules: ClassVar[list[str]] = []
+            _resident_modules: ClassVar[list[str]] = []
+            _offload_plan = OffloadPlan(
+                encoder_component_types={"mllm": "text_encoder"},
+                encoder_block_attrs={"mllm": ("encoder.block",)},
+            )
+
+            def __init__(self):
+                super().__init__()
+                self.transformer = _SingleBlockModel(num_blocks=2)
+                self.mllm = _PlainEncoder()
+
+        resolved = resolve_offload_plan(Pipeline(), _config({"text_encoder"}))
+
+        encoder = resolved.encoders[0]
+        assert encoder.path == "mllm"
+        assert encoder.selected
+        assert len(_stack_of(encoder).blocks) == 2
+
+    def test_component_declaration_wins_over_attribute_scan(self):
+        class Pipeline(nn.Module):
+            _dit_modules: ClassVar[list[str]] = ["custom_dit"]
+            _encoder_modules: ClassVar[list[str]] = []
+            _vae_modules: ClassVar[list[str]] = []
+            _resident_modules: ClassVar[list[str]] = []
+
+            def __init__(self):
+                super().__init__()
+                self.custom_dit = _SingleBlockModel(num_blocks=2)
+                # Would be picked up by the deprecated attribute scan.
+                self.transformer = _SingleBlockModel(num_blocks=2)
+
+        resolved = resolve_offload_plan(Pipeline(), _config({"dit"}))
+
+        assert [component.path for component in resolved.dits] == ["custom_dit"]
+
+    def test_attribute_scan_warns_once_per_pipeline_class(self, monkeypatch):
+        warnings: list[str] = []
+
+        class _Recorder:
+            def warning(self, message, *args):
+                warnings.append(message % args if args else message)
+
+            def __getattr__(self, _name):
+                return lambda *args, **kwargs: None
+
+        monkeypatch.setattr(plan_resolver_module, "logger", _Recorder())
+
+        class UndeclaredPipeline(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.transformer = _SingleBlockModel(num_blocks=2)
+
+        resolve_offload_plan(UndeclaredPipeline(), _config({"dit"}))
+        resolve_offload_plan(UndeclaredPipeline(), _config({"dit"}))
+
+        scan_warnings = [message for message in warnings if "SupportsComponentDiscovery" in message]
+        assert len(scan_warnings) == 1
+        assert "deprecated" in scan_warnings[0]
+
+
+class TestResolvedTopology:
+    def test_encoder_plan_resolves_one_stack_per_declared_path(self):
+        pipeline = _WanStylePipeline()
+        resolved = resolve_offload_plan(pipeline, _config({"dit", "text_encoder"}))
+
+        encoder = resolved.encoders[0]
+        assert encoder.selected and not encoder.on_demand
+        assert list(_stack_of(encoder).blocks) == list(pipeline.text_encoder.encoder.block)
+        assert [component.path for component in resolved.vaes] == ["vae"]
+
+    def test_staged_plan_resolves_every_declared_stack(self):
+        pipeline = _StagedComponentPipeline()
+        resolved = resolve_offload_plan(pipeline, _config({"text_encoder"}))
+
+        encoder = resolved.encoders[0]
+        assert encoder.on_demand
+        assert [len(stack.blocks) for stack in encoder.stacks] == [2, 2]
+        # An explicit selection never stages a VAE through the legacy plan.
+        assert not resolved.vaes[0].on_demand
+        # DiT is unselected, so it keeps no streaming topology.
+        assert resolved.dits[0].selected is False
+        assert resolved.dits[0].stacks == ()
+
+    def test_omitted_selector_keeps_the_legacy_staged_vae(self):
+        resolved = resolve_offload_plan(_StagedComponentPipeline(), _config())
+
+        assert resolved.vaes[0].on_demand
+        assert resolved.encoders[0].on_demand
+        assert resolved.dits[0].selected
+
+    def test_on_demand_only_plan_resolves_without_stacks(self):
+        class Pipeline(nn.Module):
+            _dit_modules: ClassVar[list[str]] = ["transformer"]
+            _encoder_modules: ClassVar[list[str]] = ["text_encoder"]
+            _vae_modules: ClassVar[list[str]] = ["vae"]
+            _resident_modules: ClassVar[list[str]] = []
+            _offload_plan = OffloadPlan(on_demand_component_paths=frozenset({"text_encoder", "vae"}))
+
+            def __init__(self):
+                super().__init__()
+                self.transformer = _SingleBlockModel(num_blocks=2)
+                self.text_encoder = _StagedEncoder()
+                self.vae = _StagedVAE()
+
+        resolved = resolve_offload_plan(Pipeline(), _config({"text_encoder"}))
+
+        assert resolved.encoders[0].on_demand
+        assert resolved.encoders[0].stacks == ()
+
+    def test_module_strategy_resolves_selection_without_stacks(self):
+        resolved = resolve_offload_plan(
+            _WanStylePipeline(),
+            _config({"dit", "text_encoder"}, strategy=OffloadStrategy.MODEL_LEVEL),
+        )
+
+        assert [component.path for component in resolved.components] == ["transformer", "text_encoder", "vae"]
+        assert all(component.stacks == () for component in resolved.components)
+        assert resolved.encoders[0].selected
+
+    def test_module_strategy_leaves_staged_residency_to_the_pipeline(self):
+        resolved = resolve_offload_plan(
+            _StagedComponentPipeline(),
+            _config({"dit", "text_encoder"}, strategy=OffloadStrategy.MODEL_LEVEL),
+        )
+
+        assert not any(component.on_demand for component in resolved.components)
+
+
+class TestResidentLayers:
+    def _dlo_config(self, resident_layers: int, components: set[str] | None = None) -> OffloadConfig:
+        return _config(
+            components,
+            strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
+            dlo_transfers={"dit": "rank-local", "text_encoder": "rank-local"},
+            dlo_resident_layers=resident_layers,
+        )
+
+    def test_resident_head_splits_the_declared_dit(self):
+        pipeline = _StagedComponentPipeline()
+        resolved = resolve_offload_plan(pipeline, self._dlo_config(2, {"dit"}))
+
+        stack = _stack_of(resolved.dits[0])
+        assert stack.resident_head == 2
+        assert list(stack.resident) == list(pipeline.transformer.blocks)[:2]
+        assert list(stack.streaming) == list(pipeline.transformer.blocks)[2:]
+
+    def test_resident_count_beyond_the_stack_keeps_every_block(self):
+        resolved = resolve_offload_plan(_StagedComponentPipeline(), self._dlo_config(9, {"dit"}))
+
+        stack = _stack_of(resolved.dits[0])
+        assert stack.resident_head == 4
+        assert stack.streaming == ()
+
+    def test_resident_head_leaving_one_streaming_block_is_rejected(self):
+        with pytest.raises(ValueError, match="leaves only one streaming block"):
+            resolve_offload_plan(_StagedComponentPipeline(), self._dlo_config(3, {"dit"}))
+
+    def test_legacy_single_streaming_block_stays_resident(self):
+        resolved = resolve_offload_plan(_StagedComponentPipeline(), self._dlo_config(3))
+
+        stack = _stack_of(resolved.dits[0])
+        assert stack.resident_head == 4
+        assert stack.streaming == ()
+
+    def test_resident_layers_without_a_declared_path_is_rejected(self):
+        with pytest.raises(ValueError, match="no matching resident_dit_paths"):
+            resolve_offload_plan(_WanStylePipeline(), self._dlo_config(2, {"dit"}))
+
+
+class TestValidation:
+    def test_selected_dit_requires_a_discovered_module(self):
+        pipeline = nn.Module()
+        pipeline.text_encoder = _PlainEncoder()
+
+        with pytest.raises(ValueError, match="No DiT/transformer modules found"):
+            resolve_offload_plan(pipeline, _config({"dit"}))
+
+    def test_model_level_selection_requires_both_swap_sides(self):
+        pipeline = nn.Module()
+        pipeline.transformer = _SingleBlockModel(num_blocks=2)
+
+        with pytest.raises(ValueError, match="requires an encoder execution stage"):
+            resolve_offload_plan(pipeline, _config({"dit"}, strategy=OffloadStrategy.MODEL_LEVEL))
+
+    def test_selected_text_encoder_requires_a_matching_module(self):
+        pipeline = nn.Module()
+        pipeline.transformer = _SingleBlockModel(num_blocks=2)
+
+        with pytest.raises(ValueError, match="No text encoder modules found"):
+            resolve_offload_plan(pipeline, _config({"text_encoder"}))
+
+    def test_selected_encoder_requires_a_streamable_or_staged_plan(self):
+        pipeline = nn.Module()
+        pipeline.transformer = _SingleBlockModel(num_blocks=2)
+        pipeline.text_encoder = _PlainEncoder()
+
+        with pytest.raises(ValueError, match="Selected text encoder 'text_encoder' requires"):
+            resolve_offload_plan(pipeline, _config({"text_encoder"}))
+
+    def test_staged_component_requires_the_explicit_lifecycle(self):
+        class Pipeline(nn.Module):
+            _offload_plan = OffloadPlan(on_demand_component_paths=frozenset({"text_encoder"}))
+
+            def __init__(self):
+                super().__init__()
+                self.transformer = _SingleBlockModel(num_blocks=2)
+                self.text_encoder = _PlainEncoder()
+
+        with pytest.raises(ValueError, match="must implement load_to_device"):
+            resolve_offload_plan(Pipeline(), _config({"text_encoder"}))
+
+    def test_single_block_dit_is_rejected_for_layer_offload(self):
+        pipeline = nn.Module()
+        pipeline.transformer = _SingleBlockModel(num_blocks=1)
+
+        with pytest.raises(ValueError, match="requires at least two streamable"):
+            resolve_offload_plan(pipeline, _config({"dit"}))
+
+    def test_encoder_allgather_requires_declared_replication(self):
+        config = _config(
+            {"text_encoder"},
+            strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
+            dp_size=2,
+            dlo_transfers={"dit": "rank-local", "text_encoder": "allgather"},
+        )
+
+        with pytest.raises(ValueError, match="not declared replicated"):
+            resolve_offload_plan(_StagedComponentPipeline(), config)
+
+    def test_two_dits_cannot_own_the_same_blocks(self):
+        class Pipeline(nn.Module):
+            _dit_modules: ClassVar[list[str]] = ["transformer", "transformer_2"]
+            _encoder_modules: ClassVar[list[str]] = []
+            _vae_modules: ClassVar[list[str]] = []
+            _resident_modules: ClassVar[list[str]] = []
+
+            def __init__(self):
+                super().__init__()
+                self.transformer = _SingleBlockModel(num_blocks=2)
+                self.transformer_2 = _SingleBlockModel(num_blocks=2)
+                self.transformer_2.blocks = self.transformer.blocks
+
+        with pytest.raises(ValueError, match="claimed by both"):
+            resolve_offload_plan(Pipeline(), _config({"dit"}))
+
+    def test_legacy_selection_keeps_unstreamable_components(self):
+        pipeline = nn.Module()
+        pipeline.transformer = _SingleBlockModel(num_blocks=1)
+        pipeline.text_encoder = _PlainEncoder()
+
+        resolved = resolve_offload_plan(pipeline, _config())
+
+        assert resolved.dits[0].stacks == ()
+        assert resolved.encoders[0].stacks == ()
+
+
+class TestPurity:
+    def _snapshot(self, pipeline: nn.Module):
+        return {
+            "devices": [(name, tensor.device) for name, tensor in pipeline.named_parameters()],
+            "attrs": {name: sorted(vars(module)) for name, module in pipeline.named_modules()},
+        }
+
+    def test_resolution_does_not_touch_the_model(self):
+        pipeline = _StagedComponentPipeline()
+        before = self._snapshot(pipeline)
+
+        resolve_offload_plan(pipeline, _config({"dit", "text_encoder"}))
+
+        assert self._snapshot(pipeline) == before
+        assert pipeline.text_encoder.offload_calls == 0
+        assert pipeline.text_encoder.to_calls == 0
+        assert pipeline.vae.to_calls == 0
+
+    def test_resolution_is_repeatable(self):
+        pipeline = _WanStylePipeline()
+        config = _config({"dit", "text_encoder"})
+
+        first = resolve_offload_plan(pipeline, config)
+        second = resolve_offload_plan(pipeline, config)
+
+        assert [component.path for component in first.components] == [component.path for component in second.components]
+        for left, right in zip(first.components, second.components, strict=True):
+            assert left.selected == right.selected
+            assert left.on_demand == right.on_demand
+            assert [list(stack.blocks) for stack in left.stacks] == [list(stack.blocks) for stack in right.stacks]
+
+    def test_backend_enable_rejects_before_touching_earlier_components(self, patched_offload_runtime):
+        class Pipeline(nn.Module):
+            _dit_modules: ClassVar[list[str]] = ["transformer"]
+            _encoder_modules: ClassVar[list[str]] = ["text_encoder", "text_encoder_2"]
+            _vae_modules: ClassVar[list[str]] = ["vae"]
+            _resident_modules: ClassVar[list[str]] = []
+            _offload_plan = OffloadPlan(
+                encoder_component_types={
+                    "text_encoder": "text_encoder",
+                    "text_encoder_2": "text_encoder",
+                },
+                encoder_block_attrs={"text_encoder": ("encoder.block",)},
+            )
+
+            def __init__(self):
+                super().__init__()
+                self.transformer = _SingleBlockModel(num_blocks=2)
+                self.text_encoder = _PlainEncoder()
+                self.text_encoder_2 = _PlainEncoder()
+                self.vae = _StagedVAE()
+
+        pipeline = Pipeline()
+        backend = LayerWiseOffloadBackend(_config({"text_encoder"}), CPU)
+
+        with pytest.raises(ValueError, match="Selected text encoder 'text_encoder_2' requires"):
+            backend.enable(pipeline)
+
+        assert not backend.enabled
+        assert not getattr(pipeline.text_encoder, "_omni_layerwise_enabled", False)
+        for block in pipeline.text_encoder.encoder.block:
+            assert getattr(block, "_hook_registry", None) is None
+        assert pipeline.vae.to_calls == 0

vllm_omni/diffusion/offloader/__init__.py

@@ -20,6 +20,12 @@ from .distributed_layerwise_backend import (
 from .layerwise_backend import LayerWiseOffloadBackend
 from .module_residency import BoundedAllocatorCache, PinnedModuleStager
 from .offload_plan import OffloadPlan, get_offload_plan
+from .plan_resolver import (
+    BlockStack,
+    ResolvedComponent,
+    ResolvedOffloadPlan,
+    resolve_offload_plan,
+)
 from .sequential_backend import (
     ModelLevelOffloadBackend,
     apply_sequential_offload,
@@ -42,6 +48,10 @@ __all__ = [
     "OffloadConfig",
     "OffloadPlan",
     "OffloadStrategy",
+    "BlockStack",
+    "ResolvedComponent",
+    "ResolvedOffloadPlan",
+    "resolve_offload_plan",
     "SupportsModelCpuOffload",
     "LayerWiseOffloadBackend",
     "DistributedLayerwiseOffloadBackend",

vllm_omni/diffusion/offloader/component_utils.py

@@ -4,7 +4,7 @@
 
 from __future__ import annotations
 
-from collections.abc import Callable, Iterator
+from collections.abc import Callable, Iterator, Sequence
 from itertools import chain
 from operator import attrgetter
 from typing import TYPE_CHECKING, Any
@@ -14,14 +14,12 @@ from torch import nn
 from torch.distributed.tensor import DTensor
 from vllm.logger import init_logger
 
-from .block_discovery import get_blocks_from_dit
-from .config import DIT_COMPONENT, TEXT_ENCODER_COMPONENT
+from .config import TEXT_ENCODER_COMPONENT
 from .offload_plan import OffloadPlan
 from .tensor_utils import set_tensor_storage
 
 if TYPE_CHECKING:
-    from .base import OffloadConfig
-    from .module_collector import PipelineModules
+    from .plan_resolver import BlockStack, ResolvedComponent, ResolvedOffloadPlan
 
 logger = init_logger(__name__)
 
@@ -76,30 +74,25 @@ def get_encoder_block_groups(
 
 
 def iter_streamable_dits(
-    modules: PipelineModules,
-    config: OffloadConfig,
+    resolved: ResolvedOffloadPlan,
     device: torch.device,
-    plan: OffloadPlan | None,
-) -> Iterator[tuple[str, nn.Module, list[str], list[nn.Module]]]:
-    """Yield selected DiTs whose block metadata resolves successfully."""
-    if not config.offloads(DIT_COMPONENT):
-        return
-    for name, module in zip(modules.dit_names, modules.dits):
-        logger.info("Applying hooks on %s (%s)", name, module.__class__.__name__)
-        planned_attrs = None if plan is None else plan.block_attrs.get(name)
-        block_attrs, blocks = get_blocks_from_dit(module, planned_attrs)
-        if blocks:
-            yield name, module, block_attrs, blocks
+) -> Iterator[tuple[ResolvedComponent, BlockStack]]:
+    """Yield selected DiTs that resolved to a streamable block stack."""
+    for component in resolved.dits:
+        if not component.selected:
             continue
-        if config.components is not None:
-            raise ValueError(f"Selected DiT {name!r} has no streamable layerwise-offload blocks")
-        logger.warning("Target layers (blocks) not found. Skipping offloading on %s (%s)", name, type(module).__name__)
-        module.to(device)
+        logger.info("Applying hooks on %s (%s)", component.path, component.module.__class__.__name__)
+        if not component.stacks:
+            # The resolver already warned; a legacy selection keeps the
+            # unstreamable DiT resident instead of failing the run.
+            component.module.to(device)
+            continue
+        yield component, component.stacks[0]
 
 
 def move_non_block_state_to_device(
     module: nn.Module,
-    block_groups: list[nn.ModuleList],
+    block_groups: Sequence[Sequence[nn.Module]],
     device: torch.device,
 ) -> None:
     """Keep component state outside streamed block lists resident on device."""
@@ -120,7 +113,7 @@ def move_non_block_state_to_device(
 def set_encoder_layerwise_state(
     module: nn.Module,
     hooks: list[Any],
-    block_groups: list[nn.ModuleList],
+    block_groups: Sequence[Sequence[nn.Module]],
 ) -> None:
     """Publish the backend-neutral state used by encoder stage lifecycles."""
     module._omni_layerwise_hooks = hooks
@@ -163,59 +156,38 @@ def prepare_component(
 
 
 def prepare_pipeline_components(
-    modules: PipelineModules,
-    config: OffloadConfig,
-    plan: OffloadPlan | None,
+    resolved: ResolvedOffloadPlan,
     *,
     device: torch.device,
     staged_components: list[nn.Module],
-    enable_encoder_blocks: Callable[[nn.Module, str, OffloadPlan | None, bool], bool],
+    enable_encoder_blocks: Callable[[ResolvedComponent], bool],
 ) -> None:
     """Apply the shared encoder/VAE/resident placement policy."""
-    if config.components is not None and config.offloads(TEXT_ENCODER_COMPONENT):
-        selected_encoder_names = [name for name in modules.encoder_names if config.offloads_encoder(name, plan)]
-        if not selected_encoder_names:
-            raise ValueError("No text encoder modules found for selected text_encoder offload")
-
-    if plan is not None:
-        for encoder, name in zip(modules.encoders, modules.encoder_names):
-            if config.should_offload_encoder(name, plan) and name in plan.on_demand_component_paths:
-                validate_on_demand_component(encoder, name)
-
-    for encoder, name in zip(modules.encoders, modules.encoder_names):
-        selected = config.should_offload_encoder(name, plan)
-        stage_on_demand = bool(selected and plan is not None and name in plan.on_demand_component_paths)
-        blockwise = selected and enable_encoder_blocks(encoder, name, plan, stage_on_demand)
-        if stage_on_demand and not blockwise and config.uses_allgather(TEXT_ENCODER_COMPONENT):
-            raise ValueError(
-                f"Text encoder {name!r} cannot use AllGather without a model-declared streamable block plan"
-            )
-        if selected and config.components is not None and not (blockwise or stage_on_demand):
-            raise ValueError(f"Selected text encoder {name!r} requires a model-declared streamable or on-demand plan")
+    for component in resolved.encoders:
+        blockwise = bool(component.stacks) and enable_encoder_blocks(component)
         prepare_component(
-            encoder,
-            name,
+            component.module,
+            component.path,
             device=device,
-            stage_on_demand=stage_on_demand,
+            stage_on_demand=component.on_demand,
             blockwise=blockwise,
             staged_components=staged_components,
         )
 
-    for vae, name in zip(modules.vaes, modules.vae_names):
-        legacy_staged = config.components is None and plan is not None and name in plan.on_demand_component_paths
+    for component in resolved.vaes:
         prepare_component(
-            vae,
-            name,
+            component.module,
+            component.path,
             device=device,
-            stage_on_demand=legacy_staged,
+            stage_on_demand=component.on_demand,
             blockwise=False,
             staged_components=staged_components,
         )
 
-    for name, module in zip(modules.resident_names, modules.resident_modules):
-        module.to(device)
-        logger.debug("Moved resident module %s to %s", name, device)
+    for component in resolved.residents:
+        component.module.to(device)
+        logger.debug("Moved resident module %s to %s", component.path, device)
 
-    if not config.offloads(DIT_COMPONENT):
-        for dit in modules.dits:
-            dit.to(device)
+    for component in resolved.dits:
+        if not component.selected:
+            component.module.to(device)

vllm_omni/diffusion/offloader/distributed_layerwise_backend.py

@@ -21,6 +21,7 @@ from __future__ import annotations
 import threading
 import time
 import weakref
+from collections.abc import Sequence
 from itertools import chain
 from typing import Any
 
@@ -39,7 +40,6 @@ from vllm_omni.platforms import current_omni_platform
 from .base import OffloadBackend, OffloadConfig, run_cleanup_steps
 from .component_utils import (
     clear_encoder_layerwise_state,
-    get_encoder_block_groups,
     iter_streamable_dits,
     move_non_block_state_to_device,
     prepare_component,
@@ -53,11 +53,8 @@ from .host_registration import (
     HostRegistrationError,
     register_host_mappings,
 )
-from .module_collector import ModuleDiscovery
-from .offload_plan import (
-    OffloadPlan,
-    get_offload_plan,
-)
+from .offload_plan import OffloadPlan
+from .plan_resolver import ResolvedComponent, resolve_offload_plan
 from .tensor_utils import (
     clear_block_storage,
     clear_tensor_storage,
@@ -1437,7 +1434,7 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
 
     def _install_hook_group(
         self,
-        blocks: list[nn.Module] | nn.ModuleList,
+        blocks: Sequence[nn.Module] | nn.ModuleList,
         component: str,
         *,
         use_dit_mmap: bool = False,
@@ -1483,48 +1480,27 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
         hooks[1]._is_group_first = True
         return hooks
 
-    def _try_layerwise_offload_encoder(
-        self,
-        module: nn.Module,
-        name: str,
-        plan: OffloadPlan | None,
-        stage_on_demand: bool = False,
-    ) -> bool:
-        """Apply DLO hooks to a plan-declared text encoder.
+    def _try_layerwise_offload_encoder(self, component: ResolvedComponent) -> bool:
+        """Apply DLO hooks to a text encoder's resolved block stacks.
 
         Rank-local transfer is always safe because it preserves the tensors
         produced by the model loader. Multi-rank AllGather is enabled only
         when the model declares that those tensors are replicated across the
-        selected DLO group; encoder TP groups contain different shards and
-        are deliberately not used for this transport.
+        selected DLO group; the resolver rejects every other layout before
+        this method installs anything.
         """
-        if plan is None:
-            return False
-
-        group_size = self._component_transport(TEXT_ENCODER_COMPONENT)[1]
-        if group_size > 1 and name not in plan.encoder_dlo_weight_replication:
-            raise ValueError(
-                f"Text encoder {name!r} cannot use DLO AllGather across the DiT offload group: "
-                "its loader-produced weights are not declared replicated across that group. "
-                "Set layer_options.text_encoder.weight_transfer='rank-local' in diffusion_offload_config "
-                "for encoder-TP or rank-specific layouts."
-            )
-        block_groups = get_encoder_block_groups(
-            module,
-            name,
-            plan,
-            strict=self.config.components is not None or group_size > 1,
-        )
+        block_groups = [stack.blocks for stack in component.stacks]
         if not block_groups:
             return False
 
+        module = component.module
         encoder_hooks: list[DistributedLayerwiseOffloadHook] = []
         for blocks in block_groups:
             encoder_hooks.extend(self._install_hook_group(blocks, TEXT_ENCODER_COMPONENT))
         # Track the module before placement, which may fail, so outer rollback
         # also removes its partially-installed block hooks and marker state.
         self._encoder_modules.append(module)
-        if not stage_on_demand:
+        if not component.on_demand:
             move_non_block_state_to_device(module, block_groups, self.device)
         set_encoder_layerwise_state(
             module,
@@ -1534,10 +1510,10 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
         logger.info(
             "Enabled %s DLO transfer for text encoder %s (%d blocks across %d stacks, group_size=%d)",
             self.config.transfer_for(TEXT_ENCODER_COMPONENT).value,
-            name,
+            component.path,
             sum(len(blocks) for blocks in block_groups),
             len(block_groups),
-            group_size,
+            self._component_transport(TEXT_ENCODER_COMPONENT)[1],
         )
         return True
 
@@ -1700,31 +1676,15 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
         if self.dp_group is None and self._has_multirank_allgather():
             self._init_dp_group()
 
-        modules = ModuleDiscovery.discover(pipeline)
-        if not modules.dits and self.config.offloads(DIT_COMPONENT):
+        # One resolver call validates the whole topology; an explicit selection
+        # has already failed by the time any of the checks below run.
+        resolved = resolve_offload_plan(pipeline, self.config)
+        if not resolved.dits and self.config.offloads(DIT_COMPONENT):
             if self.host_weight_plan is not None:
                 raise RuntimeError(
                     "DLO received a loader-owned host-weight plan, but no DiT modules were discovered to consume it"
                 )
-            message = "No DiT/transformer modules found for selected DiT offload"
-            if self.config.components is not None:
-                raise ValueError(message)
-            logger.warning(message)
-
-        # Retrieve optional declarative OffloadPlan from the pipeline.
-        # When present, replaces heuristic block discovery.
-        plan = get_offload_plan(pipeline)
-
-        if self.config.dlo_resident_layers:
-            resident_paths = frozenset() if plan is None else plan.resident_dit_paths
-            if not resident_paths.intersection(modules.dit_names):
-                message = (
-                    f"resident_layers={self.config.dlo_resident_layers} was requested, but this model declares "
-                    "no matching resident_dit_paths"
-                )
-                if self.config.components is not None:
-                    raise ValueError(message)
-                logger.warning("%s; all blocks will be streamed.", message)
+            logger.warning("No DiT/transformer modules found for selected DiT offload")
 
         # Storage selection belongs to the loader.  DLO consumes the exact
         # prevalidated plan that caused the loader to skip materialization;
@@ -1755,7 +1715,7 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
             elif host_weight_plan.backing_kind == "checkpoint_mmap":
                 self._load_weights_via_mmap(
                     pipeline,
-                    modules,
+                    resolved.modules,
                     host_weight_plan,
                 )
             else:
@@ -1768,10 +1728,10 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
         else:
             remaining_meta = [
                 name
-                for dit_name, dit_module in zip(modules.dit_names, modules.dits)
+                for component in resolved.dits
                 for name, tensor in chain(
-                    dit_module.named_parameters(),
-                    dit_module.named_buffers(),
+                    component.module.named_parameters(),
+                    component.module.named_buffers(),
                 )
                 if getattr(tensor, "is_meta", False)
             ]
@@ -1784,66 +1744,46 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
         # Apply each selected encoder transfer while keeping explicit VAEs and
         # unselected components resident.
         prepare_pipeline_components(
-            modules,
-            self.config,
-            plan,
+            resolved,
             device=self.device,
             staged_components=self._staged_components,
             enable_encoder_blocks=self._try_layerwise_offload_encoder,
         )
 
         if self.config.offloads(DIT_COMPONENT):
-            logger.info("Applying distributed layer-wise offloading on %s", modules.dit_names)
+            logger.info(
+                "Applying distributed layer-wise offloading on %s",
+                [component.path for component in resolved.dits],
+            )
 
         # Collect all DiT module objects to detect submodules that are
         # already handled as a separate DiT module (avoids duplicate hooks).
-        all_dit_modules = set(id(m) for m in modules.dits)
+        all_dit_modules = set(id(component.module) for component in resolved.dits)
 
         # Apply hooks for each DiT module
-        for dit_name, dit_module, blocks_attr_names, blocks in iter_streamable_dits(
-            modules, self.config, self.device, plan
-        ):
-            resident_count = 0
-            if plan is not None and dit_name in plan.resident_dit_paths:
-                resident_count = min(self.config.dlo_resident_layers, len(blocks))
-            if resident_count:
-                resident_blocks = blocks[:resident_count]
-                self._resident_blocks.extend(resident_blocks)
-                blocks = blocks[resident_count:]
+        for component, stack in iter_streamable_dits(resolved, self.device):
+            streaming = list(stack.streaming)
+            if stack.resident_head:
+                self._resident_blocks.extend(stack.resident)
                 logger.info(
                     "Keeping %d leading blocks resident on %s; streaming %d tail blocks",
-                    resident_count,
-                    dit_name,
-                    len(blocks),
+                    stack.resident_head,
+                    component.path,
+                    len(streaming),
                 )
 
             self._prepare_dit_non_block_modules(
-                dit_module,
-                blocks_attr_names,
+                component.module,
+                list(stack.attrs),
                 all_dit_modules,
-                plan,
+                resolved.declaration,
             )
 
-            num_blocks = len(blocks)
-            if num_blocks == 0:
-                logger.info("All blocks for %s are resident; no streaming hooks required", dit_name)
-                continue
-            if num_blocks <= 1:
-                if self.config.components is not None:
-                    raise ValueError(
-                        f"Selected DiT {dit_name!r} leaves only one streaming block after "
-                        f"resident_layers={resident_count}; choose a resident count that "
-                        "leaves zero or at least two streaming blocks"
-                    )
-                logger.warning(
-                    "#Streaming target layers (blocks) <= 1. Keeping the final block resident on %s (%s)",
-                    dit_name,
-                    dit_module.__class__.__name__,
-                )
-                self._resident_blocks.extend(blocks)
+            if not streaming:
+                logger.info("All blocks for %s are resident; no streaming hooks required", component.path)
                 continue
 
-            self._install_hook_group(blocks, DIT_COMPONENT, use_dit_mmap=True)
+            self._install_hook_group(streaming, DIT_COMPONENT, use_dit_mmap=True)
         if self._resident_blocks:
             self._resident_layer_group = PinnedResidentLayerGroup(
                 self._resident_blocks,

vllm_omni/diffusion/offloader/layerwise_backend.py

@@ -1,7 +1,8 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 from __future__ import annotations
 
+from collections.abc import Sequence
 from itertools import chain
 from typing import Any
 
@@ -20,15 +21,13 @@ from .block_discovery import (
 )
 from .component_utils import (
     clear_encoder_layerwise_state,
-    get_encoder_block_groups,
     iter_streamable_dits,
     move_non_block_state_to_device,
     prepare_pipeline_components,
     set_encoder_layerwise_state,
 )
 from .config import DIT_COMPONENT
-from .module_collector import ModuleDiscovery
-from .offload_plan import OffloadPlan, get_offload_plan
+from .plan_resolver import ResolvedComponent, resolve_offload_plan
 from .tensor_utils import (
     clear_block_storage,
     clear_tensor_storage,
@@ -279,7 +278,7 @@ def remove_block_hook(module: nn.Module) -> None:
 
 
 def _install_layerwise_hook_group(
-    blocks: list[nn.Module] | nn.ModuleList,
+    blocks: Sequence[nn.Module] | nn.ModuleList,
     device: torch.device,
     stream: Any,
     pin_memory: bool,
@@ -327,36 +326,29 @@ def _install_layerwise_hook_group(
 
 
 def enable_plan_encoder_layerwise_offload(
-    module: nn.Module,
-    name: str,
-    plan: OffloadPlan | None,
+    component: ResolvedComponent,
     *,
     device: torch.device,
     stream: current_omni_platform.Stream,
     pin_memory: bool,
-    stage_on_demand: bool = False,
-    strict: bool = False,
 ) -> bool:
-    """Apply rank-local layerwise hooks to plan-declared encoder stacks."""
+    """Apply rank-local layerwise hooks to the encoder's resolved stacks."""
+    module = component.module
     if getattr(module, "_omni_layerwise_enabled", False):
         return True
 
-    hooks: list[LayerwiseOffloadHook] = []
-    hooked_blocks: list[nn.Module] = []
-    block_groups = get_encoder_block_groups(
-        module,
-        name,
-        plan,
-        strict=strict,
-    )
+    block_groups = [stack.blocks for stack in component.stacks]
     if not block_groups:
         return False
+
+    hooks: list[LayerwiseOffloadHook] = []
+    hooked_blocks: list[nn.Module] = []
     try:
         for blocks in block_groups:
             group_hooks = _install_layerwise_hook_group(blocks, device, stream, pin_memory)
             hooks.extend(group_hooks)
             hooked_blocks.extend(blocks)
-        if not stage_on_demand:
+        if not component.on_demand:
             move_non_block_state_to_device(module, block_groups, device)
     except BaseException:
         run_cleanup_steps(
@@ -376,7 +368,7 @@ def enable_plan_encoder_layerwise_offload(
     )
     logger.info(
         "Enabled rank-local layerwise offload for encoder %s (%d blocks across %d stacks)",
-        name,
+        component.path,
         sum(len(blocks) for blocks in block_groups),
         len(block_groups),
     )
@@ -440,41 +432,27 @@ class LayerWiseOffloadBackend(OffloadBackend):
             logger.warning("LayerWiseOffloadBackend already enabled")
             return
 
-        modules = ModuleDiscovery.discover(pipeline)
-        plan = get_offload_plan(pipeline)
-        if not modules.dits and self.config.offloads(DIT_COMPONENT):
-            message = "No DiT/transformer modules found for selected DiT layerwise offload"
-            if self.config.components is not None:
-                raise ValueError(message)
-            logger.warning(message)
+        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")
             return
 
-        def enable_encoder_blocks(
-            module: nn.Module,
-            name: str,
-            component_plan: OffloadPlan | None,
-            stage_on_demand: bool,
-        ) -> bool:
+        def enable_encoder_blocks(component: ResolvedComponent) -> bool:
             enabled = enable_plan_encoder_layerwise_offload(
-                module,
-                name,
-                component_plan,
+                component,
                 device=self.device,
                 stream=self.copy_stream,
                 pin_memory=self.config.pin_cpu_memory,
-                stage_on_demand=stage_on_demand,
-                strict=self.config.components is not None,
             )
             if enabled:
                 # Record each successful installation immediately so a later
                 # component failure can remove these hooks transactionally.
-                self._encoder_modules.append(module)
+                self._encoder_modules.append(component.module)
             return enabled
 
         prepare_pipeline_components(
-            modules,
-            self.config,
-            plan,
+            resolved,
             device=self.device,
             staged_components=self._staged_components,
             enable_encoder_blocks=enable_encoder_blocks,
@@ -489,30 +467,17 @@ class LayerWiseOffloadBackend(OffloadBackend):
                 )
             return
 
-        logger.info("Applying layer-wise offloading on %s", modules.dit_names)
+        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)
         # Note that there might exist multiple DiT models in specific pipelines
-        for dit_name, dit_module, blocks_attr_names, blocks in iter_streamable_dits(
-            modules, self.config, self.device, plan
-        ):
-            num_blocks = len(blocks)
-            if num_blocks <= 1:
-                if self.config.components is not None:
-                    raise ValueError(
-                        f"Selected DiT {dit_name!r} requires at least two streamable layerwise-offload blocks"
-                    )
-                logger.warning(
-                    "#Target layers (blocks) <= 1. Skipping offloading on %s (%s)",
-                    dit_name,
-                    dit_module.__class__.__name__,
-                )
-                dit_module.to(self.device)
-                continue
+        for component, stack in iter_streamable_dits(resolved, self.device):
+            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 blocks_attr_names:
+                if name not in stack.attrs:
                     m.to(self.device)
                     logger.debug(f"Moved {name} to device {self.device}")
                 else:
@@ -540,7 +505,7 @@ class LayerWiseOffloadBackend(OffloadBackend):
             # zero once; later denoising iterations prefetch it from the ring.
             block_hooks[0].prefetch_layer(non_blocking=False)
 
-            logger.info(f"Layer-wise offloading enabled on {num_blocks} layers (blocks)")
+            logger.info(f"Layer-wise offloading enabled on {len(blocks)} layers (blocks)")
 
             # Track hooked blocks for cleanup
             self._blocks.append(blocks)

vllm_omni/diffusion/offloader/offload_plan.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
 """Declarative model capabilities for component layerwise offload.
 
 Models declare an ``_offload_plan`` class attribute on the pipeline. Both the
@@ -22,9 +22,11 @@ class OffloadPlan:
     pipeline class. When present, both layerwise backends use it instead of
     model-specific backend branches, making new integrations data-driven.
 
-    If not declared, the offloader falls back to:
-    1. ``_layerwise_offload_blocks_attrs`` on each DiT module class.
-    2. Heuristic search for ``layers`` / ``blocks`` / ``h`` attributes.
+    If not declared, the offloader falls back to
+    ``_layerwise_offload_blocks_attrs`` on each DiT module class.
+
+    :func:`~vllm_omni.diffusion.offloader.plan_resolver.resolve_offload_plan`
+    is the only consumer; backends read the resolved artifact it returns.
 
     Attributes:
         block_attrs: Maps DiT path → tuple of block-list attribute names.

vllm_omni/diffusion/offloader/plan_resolver.py

@@ -0,0 +1,330 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""One topology resolver shared by every diffusion offload backend.
+
+Model topology is declared through several mechanisms that grew independently:
+:class:`~vllm_omni.diffusion.models.interface.SupportsComponentDiscovery` class
+variables, the pipeline-level :class:`OffloadPlan`, per-DiT
+``_layerwise_offload_blocks_attrs``, and a leaf-name rule for encoders. This
+module runs them in one defined order and returns one backend-neutral artifact.
+
+The resolver is pure: it moves no tensor, installs no hook, writes no module
+attribute, and touches no process group. Everything it can reject is rejected
+here, so plan-dependent failures happen before a backend mutates the model.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from functools import cache
+from typing import TYPE_CHECKING
+
+from torch import nn
+from vllm.logger import init_logger
+
+from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery
+
+from .block_discovery import get_blocks_from_dit
+from .component_utils import get_encoder_block_groups, validate_on_demand_component
+from .config import DIT_COMPONENT, TEXT_ENCODER_COMPONENT, OffloadStrategy
+from .module_collector import ModuleDiscovery, PipelineModules
+from .offload_plan import OffloadPlan, get_offload_plan
+
+if TYPE_CHECKING:
+    from .base import OffloadConfig
+
+logger = init_logger(__name__)
+
+
+@dataclass(frozen=True)
+class BlockStack:
+    """One ring of repeated blocks that a backend streams together.
+
+    ``attrs`` names the owner attributes that contributed blocks; backends use
+    it to skip those children when they place the non-streamed remainder. It is
+    filled for DiT stacks only — encoder stacks are hooked as whole stacks and
+    their remainder is placed by tensor identity.
+    """
+
+    blocks: tuple[nn.Module, ...]
+    attrs: tuple[str, ...] = ()
+    resident_head: int = 0
+
+    @property
+    def resident(self) -> tuple[nn.Module, ...]:
+        """Leading blocks that stay on the device instead of streaming."""
+        return self.blocks[: self.resident_head]
+
+    @property
+    def streaming(self) -> tuple[nn.Module, ...]:
+        """Blocks a hook ring transfers on demand."""
+        return self.blocks[self.resident_head :]
+
+
+@dataclass(frozen=True)
+class ResolvedComponent:
+    """One pipeline component with its resolved, validated offload topology.
+
+    ``selected`` means the active selector covers this component. ``on_demand``
+    means the pipeline owns its residency through ``load_to_device`` /
+    ``offload_to_cpu``; a VAE staged by a legacy model plan is ``on_demand``
+    without being selectable through the public component grammar.
+    """
+
+    path: str
+    module: nn.Module
+    selected: bool
+    stacks: tuple[BlockStack, ...] = ()
+    on_demand: bool = False
+
+
+@dataclass(frozen=True)
+class ResolvedOffloadPlan:
+    """Backend-neutral topology for one pipeline under one offload config."""
+
+    dits: tuple[ResolvedComponent, ...]
+    encoders: tuple[ResolvedComponent, ...]
+    vaes: tuple[ResolvedComponent, ...]
+    residents: tuple[ResolvedComponent, ...]
+    # Raw discovery output and the model declaration remain reachable for the
+    # loader-owned mmap path and for DLO's nested-submodule staging, which
+    # still interpret topology themselves.
+    modules: PipelineModules
+    declaration: OffloadPlan | None = None
+
+    @property
+    def components(self) -> tuple[ResolvedComponent, ...]:
+        return self.dits + self.encoders + self.vaes + self.residents
+
+
+@cache
+def _warn_legacy_discovery(pipeline_class: str) -> None:
+    """Warn once per pipeline class that component discovery guessed.
+
+    The key is the fully qualified class name so two same-named pipeline
+    classes in different modules each get their own warning.
+    """
+    logger.warning(
+        "%s does not declare SupportsComponentDiscovery; offload components were "
+        "discovered by scanning well-known attribute names. Declare _dit_modules, "
+        "_encoder_modules, and _vae_modules; the attribute scan is deprecated.",
+        pipeline_class,
+    )
+
+
+def _resolve_dit_stacks(
+    module: nn.Module,
+    path: str,
+    declaration: OffloadPlan | None,
+    config: OffloadConfig,
+    *,
+    explicit: bool,
+) -> tuple[BlockStack, ...]:
+    """Resolve one DiT's streamed blocks and its device-resident head."""
+    planned_attrs = None if declaration is None else declaration.block_attrs.get(path)
+    attrs, blocks = get_blocks_from_dit(module, planned_attrs)
+    if not blocks:
+        if explicit:
+            raise ValueError(f"Selected DiT {path!r} has no streamable layerwise-offload blocks")
+        logger.warning(
+            "Target layers (blocks) not found. Skipping offloading on %s (%s)",
+            path,
+            type(module).__name__,
+        )
+        return ()
+
+    distributed = config.strategy is OffloadStrategy.DISTRIBUTED_LAYER_WISE
+    resident_head = 0
+    if (
+        distributed
+        and config.dlo_resident_layers
+        and declaration is not None
+        and path in declaration.resident_dit_paths
+    ):
+        resident_head = min(config.dlo_resident_layers, len(blocks))
+
+    if len(blocks) - resident_head == 1:
+        if explicit and distributed:
+            raise ValueError(
+                f"Selected DiT {path!r} leaves only one streaming block after "
+                f"resident_layers={resident_head}; choose a resident count that "
+                "leaves zero or at least two streaming blocks"
+            )
+        if explicit:
+            raise ValueError(f"Selected DiT {path!r} requires at least two streamable layerwise-offload blocks")
+        if distributed:
+            # One streamed block cannot form a prefetch ring; keeping it
+            # resident preserves the model instead of skipping placement.
+            logger.warning(
+                "#Streaming target layers (blocks) <= 1. Keeping the final block resident on %s (%s)",
+                path,
+                type(module).__name__,
+            )
+            resident_head = len(blocks)
+        else:
+            logger.warning(
+                "#Target layers (blocks) <= 1. Skipping offloading on %s (%s)",
+                path,
+                type(module).__name__,
+            )
+            return ()
+
+    return (BlockStack(blocks=tuple(blocks), attrs=tuple(attrs), resident_head=resident_head),)
+
+
+def _resolve_encoder_stacks(
+    module: nn.Module,
+    path: str,
+    declaration: OffloadPlan | None,
+    config: OffloadConfig,
+    *,
+    explicit: bool,
+) -> tuple[BlockStack, ...]:
+    """Resolve one encoder's streamable stacks and check transfer safety."""
+    if declaration is None:
+        return ()
+
+    group_size = 1
+    if config.strategy is OffloadStrategy.DISTRIBUTED_LAYER_WISE and config.uses_allgather(TEXT_ENCODER_COMPONENT):
+        group_size = config.dp_size
+    if group_size > 1 and path not in declaration.encoder_dlo_weight_replication:
+        # Every rank must reject an unsafe group, including ranks whose local
+        # encoder is an unloaded stub, so the check precedes block discovery.
+        raise ValueError(
+            f"Text encoder {path!r} cannot use DLO AllGather across the DiT offload group: "
+            "its loader-produced weights are not declared replicated across that group. "
+            "Set layer_options.text_encoder.weight_transfer='rank-local' in diffusion_offload_config "
+            "for encoder-TP or rank-specific layouts."
+        )
+
+    groups = get_encoder_block_groups(module, path, declaration, strict=explicit or group_size > 1)
+    return tuple(BlockStack(blocks=tuple(blocks)) for blocks in groups)
+
+
+def resolve_offload_plan(pipeline: nn.Module, config: OffloadConfig) -> ResolvedOffloadPlan:
+    """Resolve and validate one pipeline's offload topology.
+
+    Raises ``ValueError`` for every topology that the requested configuration
+    cannot serve, before the caller places a module or installs a hook.
+    """
+    modules = ModuleDiscovery.discover(pipeline)
+    declaration = get_offload_plan(pipeline)
+    if not isinstance(pipeline, SupportsComponentDiscovery) and (modules.dits or modules.encoders or modules.vaes):
+        pipeline_class = type(pipeline)
+        _warn_legacy_discovery(f"{pipeline_class.__module__}.{pipeline_class.__qualname__}")
+
+    explicit = config.components is not None
+    layerwise = config.strategy in (OffloadStrategy.LAYER_WISE, OffloadStrategy.DISTRIBUTED_LAYER_WISE)
+    dit_selected = config.offloads(DIT_COMPONENT)
+
+    if explicit:
+        if config.strategy is OffloadStrategy.MODEL_LEVEL:
+            # Model-level offload swaps the DiT against an encoder stage, so
+            # both sides must exist regardless of which side was selected.
+            if not modules.dits:
+                raise ValueError("Component-selective model offload requires a DiT/transformer module")
+            if not modules.encoders:
+                raise ValueError("Component-selective model offload requires an encoder execution stage")
+        elif dit_selected and not modules.dits:
+            raise ValueError("No DiT/transformer modules found for selected DiT offload")
+        if config.offloads(TEXT_ENCODER_COMPONENT) and not any(
+            config.offloads_encoder(path, declaration) for path in modules.encoder_names
+        ):
+            raise ValueError("No text encoder modules found for selected text_encoder offload")
+
+    if config.strategy is OffloadStrategy.DISTRIBUTED_LAYER_WISE and config.dlo_resident_layers:
+        resident_paths = frozenset() if declaration is None else declaration.resident_dit_paths
+        if not resident_paths.intersection(modules.dit_names):
+            message = (
+                f"resident_layers={config.dlo_resident_layers} was requested, but this model declares "
+                "no matching resident_dit_paths"
+            )
+            if explicit:
+                raise ValueError(message)
+            logger.warning("%s; all blocks will be streamed.", message)
+
+    dits: list[ResolvedComponent] = []
+    encoders: list[ResolvedComponent] = []
+    vaes: list[ResolvedComponent] = []
+    residents: list[ResolvedComponent] = []
+
+    for path, module in zip(modules.dit_names, modules.dits):
+        stacks: tuple[BlockStack, ...] = ()
+        if dit_selected and layerwise:
+            stacks = _resolve_dit_stacks(module, path, declaration, config, explicit=explicit)
+        dits.append(ResolvedComponent(path=path, module=module, selected=dit_selected, stacks=stacks))
+
+    for path, module in zip(modules.encoder_names, modules.encoders):
+        selected = config.should_offload_encoder(path, declaration)
+        # Staged residency belongs to the layer backends; model-level offload
+        # swaps whole modules through its own hooks.
+        on_demand = bool(
+            selected and layerwise and declaration is not None and path in declaration.on_demand_component_paths
+        )
+        if on_demand:
+            validate_on_demand_component(module, path)
+        stacks = ()
+        if selected and layerwise:
+            stacks = _resolve_encoder_stacks(module, path, declaration, config, explicit=explicit)
+            if on_demand and not stacks and config.uses_allgather(TEXT_ENCODER_COMPONENT):
+                raise ValueError(
+                    f"Text encoder {path!r} cannot use AllGather without a model-declared streamable block plan"
+                )
+            if explicit and not (stacks or on_demand):
+                raise ValueError(
+                    f"Selected text encoder {path!r} requires a model-declared streamable or on-demand plan"
+                )
+        encoders.append(
+            ResolvedComponent(
+                path=path,
+                module=module,
+                selected=selected,
+                stacks=stacks,
+                on_demand=on_demand,
+            )
+        )
+
+    for path, module in zip(modules.vae_names, modules.vaes):
+        # VAEs are not part of the public selector. A model plan may still own
+        # their residency, which the compatibility topology preserves.
+        legacy_staged = (
+            layerwise and not explicit and declaration is not None and path in declaration.on_demand_component_paths
+        )
+        if legacy_staged:
+            validate_on_demand_component(module, path)
+        vaes.append(ResolvedComponent(path=path, module=module, selected=False, on_demand=legacy_staged))
+
+    for path, module in zip(modules.resident_names, modules.resident_modules):
+        residents.append(ResolvedComponent(path=path, module=module, selected=False))
+
+    resolved = ResolvedOffloadPlan(
+        dits=tuple(dits),
+        encoders=tuple(encoders),
+        vaes=tuple(vaes),
+        residents=tuple(residents),
+        modules=modules,
+        declaration=declaration,
+    )
+    _validate_unique_ownership(resolved)
+    return resolved
+
+
+def _validate_unique_ownership(resolved: ResolvedOffloadPlan) -> None:
+    """Reject topologies where two components would hook the same blocks."""
+    owner_by_block: dict[int, str] = {}
+    for component in resolved.components:
+        for block in (block for stack in component.stacks for block in stack.blocks):
+            owner = owner_by_block.setdefault(id(block), component.path)
+            if owner != component.path:
+                raise ValueError(
+                    f"Offload block is claimed by both {owner!r} and {component.path!r}; "
+                    "declare one owner for each block list"
+                )
+
+
+__all__ = [
+    "BlockStack",
+    "ResolvedComponent",
+    "ResolvedOffloadPlan",
+    "resolve_offload_plan",
+]

vllm_omni/diffusion/offloader/sequential_backend.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
 
 from collections.abc import Collection, Iterator
 from contextlib import contextmanager
@@ -13,9 +13,8 @@ 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, TEXT_ENCODER_COMPONENT
-from .module_collector import ModuleDiscovery
-from .offload_plan import get_offload_plan
+from .config import DIT_COMPONENT
+from .plan_resolver import resolve_offload_plan
 
 logger = init_logger(__name__)
 
@@ -329,59 +328,44 @@ class ModelLevelOffloadBackend(OffloadBackend):
             )
             return
 
-        modules = ModuleDiscovery.discover(pipeline)
-        plan = get_offload_plan(pipeline)
-        selected_encoders = [
-            encoder
-            for encoder, name in zip(modules.encoders, modules.encoder_names)
-            if self.config.should_offload_encoder(name, plan)
-        ]
-        if self.config.components is not None:
-            if not modules.dits:
-                raise ValueError("Component-selective model offload requires a DiT/transformer module")
-            if not modules.encoders:
-                raise ValueError("Component-selective model offload requires an encoder execution stage")
-            if self.config.offloads(TEXT_ENCODER_COMPONENT) and not selected_encoders:
-                raise ValueError("No text encoder modules found for selected text_encoder module offload")
+        resolved = resolve_offload_plan(pipeline, self.config)
+        dits = [component.module for component in resolved.dits]
+        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_encoders = [component.module for component in resolved.encoders if component.selected]
 
-        all_modules = [
-            *modules.dits,
-            *modules.encoders,
-            *modules.vaes,
-            *modules.resident_modules,
-        ]
+        all_modules = [*dits, *encoders, *vaes, *residents]
         initial_devices = _capture_tensor_devices(all_modules)
         try:
-            for encoder in modules.encoders:
+            for encoder in encoders:
                 encoder.to(self.device)
-            for vae in modules.vaes:
+            for vae in vaes:
                 vae.to(self.device, non_blocking=True)
-            for resident in modules.resident_modules:
+            for resident in residents:
                 resident.to(self.device)
 
-            if not modules.dits:
+            if not dits:
                 logger.warning("No DiT/transformer modules found, skipping model-level offloading")
                 return
-            if not modules.encoders:
-                for dit in modules.dits:
+            if not encoders:
+                for dit in dits:
                     dit.to(self.device)
                 logger.warning("No encoder modules found, skipping model-level offloading")
                 return
 
             apply_sequential_offload(
-                dit_modules=modules.dits,
-                encoder_modules=modules.encoders,
+                dit_modules=dits,
+                encoder_modules=encoders,
                 device=self.device,
                 pin_memory=self.config.pin_cpu_memory,
                 use_hsdp=self.config.use_hsdp,
-                offload_dit_modules=(
-                    modules.dits if self.config.components is None or self.config.offloads(DIT_COMPONENT) else ()
-                ),
+                offload_dit_modules=(dits if self.config.offloads(DIT_COMPONENT) else ()),
                 offload_encoder_modules=selected_encoders,
             )
         except BaseException:
             try:
-                remove_sequential_offload([*modules.dits, *modules.encoders])
+                remove_sequential_offload([*dits, *encoders])
             except BaseException:
                 logger.exception("Failed to remove every model-level hook during rollback")
             try:
@@ -391,15 +375,19 @@ class ModelLevelOffloadBackend(OffloadBackend):
             raise
 
         # Track modules for cleanup
-        self._offload_modules = [*modules.dits, *modules.encoders]
+        self._offload_modules = [*dits, *encoders]
 
         self.enabled = True
 
         logger.info(
             "Model-level offloading enabled: %s <-> %s (mutual exclusion)%s",
-            ", ".join(modules.dit_names),
-            ", ".join(modules.encoder_names),
-            f"; resident on GPU: {', '.join(modules.resident_names)}" if modules.resident_names else "",
+            ", ".join(component.path for component in resolved.dits),
+            ", ".join(component.path for component in resolved.encoders),
+            (
+                f"; resident on GPU: {', '.join(component.path for component in resolved.residents)}"
+                if resolved.residents
+                else ""
+            ),
         )
 
     def disable(self) -> None: