[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>

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: