[diffusion][feature] Add component-selective offload policies (#5929)

Signed-off-by: lishunyang12 <lishunyang12@163.com>

docs/.nav.yml

@@ -72,7 +72,7 @@ nav:
         - Feature Compatibility: user_guide/feature_compatibility.md
       - CPU Offloading:
         - Overview: user_guide/diffusion/cpu_offload.md
-        - Model-Level: user_guide/diffusion/offloader/cpu_offload.md
+        - Model-Level: user_guide/diffusion/offloader/module_offload.md
         - Layerwise: user_guide/diffusion/offloader/layerwise_offload.md
         - Distributed Layerwise: user_guide/diffusion/offloader/distributed_layerwise_offload.md
       - Cache Acceleration:
@@ -151,7 +151,7 @@ nav:
           - Skip-Softmax: design/feature/skip_softmax.md
         - CPU Offloading:
           - Overview: design/feature/offloader/README.md
-          - Model-Level: design/feature/offloader/cpu_offload.md
+          - Model-Level: design/feature/offloader/module_offload.md
           - Layerwise: design/feature/offloader/layerwise_offload.md
           - Distributed Layerwise: design/feature/offloader/distributed_layerwise_offload.md
         - Cache-DiT: design/feature/cache_dit.md

docs/design/feature/offloader/README.md

@@ -5,23 +5,27 @@ The three strategies have separate user and design pages:
 
 | Strategy | User guide | Design contract |
 | --- | --- | --- |
-| Model-level | [Guide](../../../user_guide/diffusion/offloader/cpu_offload.md) | [Design](cpu_offload.md) |
+| Model-level | [Guide](../../../user_guide/diffusion/offloader/module_offload.md) | [Design](module_offload.md) |
 | Layerwise | [Guide](../../../user_guide/diffusion/offloader/layerwise_offload.md) | [Design](layerwise_offload.md) |
 | Distributed layerwise | [Guide](../../../user_guide/diffusion/offloader/distributed_layerwise_offload.md) | [Design](distributed_layerwise_offload.md) |
 
 ## Strategy selection
 
 `OffloadConfig.from_od_config()` converts diffusion configuration into one
-`OffloadStrategy`. If callers set multiple flags, selection is deterministic:
+`OffloadStrategy`. The public config separates component selection from
+policy:
 
-1. `DISTRIBUTED_LAYERWISE`;
-2. `LAYERWISE`;
-3. `MODEL_LEVEL`; and
-4. `NONE`.
+- `mode="module"` selects model-level offload;
+- `mode="layer"` with rank-local transfers selects ordinary layerwise
+  offload; and
+- AllGather transfer or resident layers selects the distributed layerwise
+  backend that implements those capabilities.
 
-The factory logs when a higher-priority strategy disables another. It derives
-parallel and HSDP state from `DiffusionParallelConfig`; callers do not provide
-a separate offload group size.
+The compatibility boolean flags retain their historical priority (distributed
+layerwise, layerwise, then model-level). A compact config rejects a conflicting
+legacy strategy or non-default legacy DLO tuning. The factory derives parallel
+and HSDP state from `DiffusionParallelConfig`; callers do not provide a
+separate offload group size.
 
 `get_offload_backend()` then validates platform offload support, resolves the
 device, and creates exactly one backend. Returning `None` means offloading is
@@ -39,6 +43,11 @@ Every backend implements `OffloadBackend`:
 `disable()` does not promise to restore the pipeline's original device
 placement. The caller owns any subsequent rematerialization.
 
+An enable failure must remove every partially installed hook. Rank-local
+backends restore ordinary tensors so the pipeline can be retried. A failed
+multi-rank AllGather startup must not enter unmatched recovery collectives;
+that worker is safely discarded after local resource cleanup.
+
 Hooks are registered through `HookRegistry` and `ModelHook`; offload backends
 must use distinct hook names and remove only hooks they own.
 
@@ -55,13 +64,13 @@ Dotted paths are supported. Legacy pipelines may use the fallback scan of
 well-known attribute names, but new integrations should declare components
 explicitly.
 
-Ordinary layerwise offload obtains topology from declared
-`_layerwise_offload_blocks_attrs` (with a deprecated singular-name
-compatibility path). Distributed layerwise offload may instead use a pipeline
-`OffloadPlan`, then falls back to declared block attributes and its documented
-compatibility discovery. Discovery metadata describes structure only; the
-backend remains responsible for transfer, synchronization, and storage
-ownership.
+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.
 
 ## Cross-strategy invariants
 

docs/design/feature/offloader/distributed_layerwise_offload.md

@@ -1,10 +1,11 @@
 # Distributed Layerwise Offload
 
 This document describes distributed layerwise offload (DLO) for diffusion
-models. DLO keeps only a small number of DiT blocks on the accelerator and
-streams the remaining blocks from host memory. The distributed backend can
-either shard those host-side weights across an existing parallel group or keep
-complete rank-local block sources and avoid an additional collective.
+models. DLO keeps only a small number of selected DiT or text-encoder blocks
+on the accelerator and streams the remaining blocks from host memory. The
+distributed backend can either shard those host-side weights across an
+existing parallel group or keep complete rank-local block sources and avoid
+an additional collective.
 
 For user-facing commands, see the
 [distributed layerwise offloading guide](../../../user_guide/diffusion/offloader/distributed_layerwise_offload.md)
@@ -38,10 +39,11 @@ detailed contracts and validation boundaries.
 
 ## Status
 
-DLO is implemented for multi-device diffusion execution. The default
-AllGather path is the primary path for DP and SP deployments. The
-`--dlo-no-use-allgather` path streams complete blocks independently and adds no
-DLO weight collective.
+DLO is implemented for multi-device diffusion execution. The default public
+transfer is `rank-local`; setting a selected component's
+`layer_options.NAME.weight_transfer` to `allgather` enables weight sharding across a
+compatible DP or SP group. Transfers are configured independently for `dit`
+and `text_encoder`.
 
 Host storage is selected separately from the transfer protocol. The loader can
 produce a direct-checkpoint mmap plan for a proven-compatible runtime layout;
@@ -109,7 +111,7 @@ checkpoint adapter beside the ordinary loader.
 
 ### AllGather path
 
-With the default `dlo_use_allgather=True`, each rank stores approximately
+For a component configured with `weight_transfer="allgather"`, each rank stores approximately
 `1 / group_size` of each streamable block in pinned host memory. The next
 block's shard is copied to a device buffer and reconstructed with
 `all_gather_into_tensor` on a communication stream while the current block is
@@ -132,9 +134,8 @@ source used to prepare each rank's persistent shard. They can be closed after
 shard preparation. Across the AllGather group, those private shards total
 approximately one runtime model copy.
 
-An effective DLO group size of one performs no collective, even when
-`dlo_use_allgather=True`; it follows the rank-local transfer path described
-below.
+An effective DLO group size of one performs no collective; it follows the
+rank-local transfer path described below.
 
 When DP is greater than one, the engine can process one request per DP rank in
 the same denoising wave. Because AllGather is a collective, all participating
@@ -142,9 +143,10 @@ requests must take the same execution path at every denoising step.
 
 ### Rank-local path without DLO AllGather
 
-With `--dlo-no-use-allgather`, DLO forces its internal offload shard size to
-one and streams complete blocks using H2D copies only. The host backing may be
-either a loader-approved checkpoint mapping or ordinary runtime tensors.
+With `weight_transfer="rank-local"`, DLO forces the affected component's internal
+offload shard size to one and streams complete blocks using H2D copies only.
+The host backing may be either a loader-approved checkpoint mapping or
+ordinary runtime tensors.
 
 For direct mmap, each process retains immutable safetensors views and uses two
 bounded pinned host staging slots. Processes on the same node that map the same
@@ -204,7 +206,7 @@ adopt it.
 > optimizations without changing AllGather behavior.
 
 Final-layout Host Weight Runtime (HWR) backing is opt-in and currently applies
-only to no-AllGather DLO. Enable it with
+only to a selected rank-local DiT. Enable it with
 `--host-weight-runtime-mode preferred` (or `required`) and
 `--host-weight-runtime-root <node-local-root>`.
 
@@ -422,8 +424,8 @@ failure.
   collective. When the model declares the final-layout contract, HWR keys the
   reusable artifact by TP rank/size and SP layout.
 - **HSDP + SP:** the general parallel configuration permits HSDP over SP, but
-  DLO must use `--dlo-no-use-allgather`. HSDP remains responsible for weight
-  materialization and synchronization.
+  each selected component must use `weight_transfer="rank-local"`. HSDP remains
+  responsible for weight materialization and synchronization.
 - **HSDP + DP or TP:** rejected independently by the diffusion parallel
   configuration.
 
@@ -447,8 +449,9 @@ recorded layouts (they all keep plain 1-byte dtypes over ordinary strided
 views: online INT8 has an int8 weight plus fp32 scale, online MXFP8 has an
 fp8 weight plus e8m0 block scale, each with a contiguous or transposed-view
 finalized tensor, all covered by the same physical-order packing as online
-FP8). Other online methods must use `--dlo-no-use-allgather` or disable
-online quantization until their runtime layouts are validated.
+FP8). Other online methods must set
+`layer_options.dit.weight_transfer="rank-local"` or disable online quantization
+until their runtime layouts are validated.
 
 The Host Weight Runtime representation, publication, and no-AllGather consumer
 contracts are merged; see

docs/design/feature/offloader/layerwise_offload.md

@@ -36,12 +36,14 @@ dependencies, not incidental global synchronization, define correctness.
 ## Topology contract
 
 A DiT declares ordered block-container attributes using
-`_layerwise_offload_blocks_attrs` or an `OffloadPlan`. Each resolved item must
-be an executable `nn.Module`; ordering must match forward execution. Multiple
-containers are concatenated in execution order.
+`_layerwise_offload_blocks_attrs` or an `OffloadPlan`. A text encoder declares
+its block-list paths through `OffloadPlan.encoder_block_attrs`. Each resolved
+item must be an executable `nn.Module`; ordering must match forward execution.
+Multiple containers are streamed as separate hook rings.
 
-Non-block DiT modules, encoders, VAEs, and declared resident modules are not
-streamed by this backend.
+Non-block state of a selected DiT or text encoder remains device resident.
+Unselected components, VAEs, and declared resident modules are not streamed by
+this backend.
 
 ## Invariants and limitations
 
@@ -50,7 +52,8 @@ streamed by this backend.
 - Blocks with heterogeneous shapes and dtypes use their recorded metadata;
   implementations must not assume identical blocks.
 - The ring is initialized before the first forward so block zero is available.
-- This backend is single-device; distributed weight reconstruction belongs to
+- This backend uses rank-local host weights. Multi-device compute may remain
+  active, but cross-rank weight reconstruction belongs to
   [Distributed Layerwise Offload](distributed_layerwise_offload.md).
 
 Shared selection and lifecycle behavior is defined in the

docs/design/feature/offloader/module_offload.md

@@ -2,7 +2,7 @@
 
 Model-level offload enforces mutual exclusion between pipeline component
 groups. For commands and model support, see the
-[model-level user guide](../../../user_guide/diffusion/offloader/cpu_offload.md).
+[model-level user guide](../../../user_guide/diffusion/offloader/module_offload.md).
 
 ## Generic hook path
 
@@ -10,10 +10,15 @@ groups. For commands and model support, see the
 makes encoders and VAEs device resident for initialization, then installs
 `SequentialOffloadHook` instances:
 
-- each DiT hook offloads all encoders and other DiTs before loading itself;
-- each encoder hook offloads every DiT before loading itself; and
+- each DiT hook offloads selected encoders and other selected DiTs before
+  loading itself;
+- each encoder hook offloads selected DiTs before loading itself; and
 - VAEs are not part of the mutual-exclusion hook set.
 
+All execution-stage modules still receive a hook so an unselected resident
+stage can evict a selected component before it runs. Only components listed in
+`diffusion_offload_config.components` are eligible to move to CPU.
+
 The hook moves parameters and buffers without recursively calling
 `module.to()`. This avoids recursion through modules that retain references to
 other transformer blocks. Host tensors are pinned when configured; DTensors
@@ -36,7 +41,9 @@ contract:
 
 - only the executing component is device resident;
 - enable and disable are idempotent;
-- transfers use the supplied device, pinning, and HSDP settings; and
+- transfers use the supplied device, pinning, and HSDP settings;
+- explicit component selection is either honored or rejected before partial
+  setup; and
 - disabling removes all model-local offload hooks or contexts.
 
 The delegation is currently duck typed. If another independent implementation

docs/design/index.md

@@ -66,7 +66,7 @@ The design contracts separate selection mechanics from backend algorithms:
 #### CPU offloading
 
 - [Overview and Shared Contracts](feature/offloader/README.md)
-- [Model-Level Offload](feature/offloader/cpu_offload.md)
+- [Model-Level Offload](feature/offloader/module_offload.md)
 - [Layerwise Offload](feature/offloader/layerwise_offload.md)
 - [Distributed Layerwise Offload](feature/offloader/distributed_layerwise_offload.md)
 

docs/user_guide/diffusion/cpu_offload.md

@@ -11,7 +11,7 @@ For the shared factory, discovery, and lifecycle contract, see the
 
 | Strategy | Device residency | Parallel scope | Guide |
 | --- | --- | --- | --- |
-| Model-level (sequential) | One pipeline component group at a time | Single device | [Model-Level Offloading](offloader/cpu_offload.md) |
+| Model-level (sequential) | One pipeline component group at a time | Single device | [Model-Level Offloading](offloader/module_offload.md) |
 | Layerwise (blockwise) | One transformer block, with next-block prefetch | Single device | [Layerwise Offloading](offloader/layerwise_offload.md) |
 | Distributed layerwise | Fixed two-block device buffer; optional host-weight sharding and AllGather | Multiple GPU/NPU ranks | [Distributed Layerwise Offloading](offloader/distributed_layerwise_offload.md) |
 
@@ -27,7 +27,7 @@ higher-priority strategy is selected.
 
 ## Quick selection
 
-- Use [model-level offloading](offloader/cpu_offload.md) when swapping whole
+- Use [model-level offloading](offloader/module_offload.md) when swapping whole
   encoders and DiTs is enough to fit and phase-boundary transfers are
   acceptable.
 - Use [layerwise offloading](offloader/layerwise_offload.md) for compute-heavy video
@@ -63,7 +63,7 @@ These headings preserve links to sections that moved into dedicated guides.
 
 ## Model-level (Sequential) Offloading
 
-See [Model-Level Offloading](offloader/cpu_offload.md).
+See [Model-Level Offloading](offloader/module_offload.md).
 
 ## Layerwise (Blockwise) Offloading
 

docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md

@@ -1,11 +1,13 @@
 # Distributed Layerwise Offloading
 
 Distributed layerwise offloading (DLO) extends block streaming to multi-device
-deployments. With AllGather enabled, each rank stores roughly `1 / dp_size` of
-the host weights and reconstructs each layer at runtime. Without AllGather,
-each rank streams a complete block independently. Compatible TP1 deployments
-can share checkpoint-backed host pages among processes on the same node;
-otherwise DLO streams the ordinary loader's rank-local tensors.
+deployments. Each selected component chooses its transfer independently. With
+`allgather`, each rank stores roughly one shard per member of the effective DLO
+group (DP, or SP when DP is one) and reconstructs each layer at runtime. With
+`rank-local`, each rank streams a complete loader-produced block independently.
+Compatible TP1 DiT deployments can share checkpoint-backed host pages among
+processes on the same node; otherwise DLO streams the ordinary loader's
+rank-local tensors.
 
 See the [DLO feature design](../../../design/feature/offloader/distributed_layerwise_offload.md)
 for the implementation contract and compatibility matrix.
@@ -24,52 +26,120 @@ Slots:           slot 0: Layer N    slot 1: Layer N+1
 AllGather communicates only request-independent weight shards, so data-
 parallel ranks may process different requests concurrently.
 
+The public API describes intent in one nested configuration:
+
+| Level | Setting | Question answered |
+| --- | --- | --- |
+| Granularity | `mode` | move a complete component (`module`) or stream its blocks (`layer`) |
+| Selection | `components` list | whether `dit` and/or `text_encoder` may move to CPU |
+| Per-component layer options | `layer_options` | choose rank-local or AllGather weight transfer and optionally retain leading DiT blocks |
+
+`weight_transfer` applies only to model weights. It does not change encoder
+activations or the encoder's compute parallelism. `rank-local` copies each
+rank's loader-produced block to its device. `allgather` stores one shard of
+that block per rank, copies each shard to its device, and reconstructs the
+complete block with a device-side AllGather before compute.
+
+The historical `enable_cpu_offload`, `enable_layerwise_offload`, and
+`enable_distributed_layerwise_offload` fields remain compatibility aliases
+for existing callers and model-specific stage lifecycles. Existing
+combinations retain their priority (`distributed layerwise` > `layerwise` >
+`module`). A compact config rejects conflicting compatibility strategies and
+non-default legacy DLO tuning.
+`distributed-layerwise` remains an internal backend name; users select
+`mode="layer"`, and the runtime chooses the capable backend.
+
 ## Usage
 
 ```bash
-# Four ranks with sharded host weights and AllGather
-vllm serve /path/to/model --omni \
-  --enable-distributed-layerwise-offload \
-  --data-parallel-size 4
+# Four SP ranks; Wan declares replicated text-encoder weights, so both
+# selected components can use AllGather safely.
+vllm serve Wan-AI/Wan2.2-T2V-A14B-Diffusers --omni \
+  --diffusion-offload-config \
+  '{"mode":"layer","components":["dit","text_encoder"],"layer_options":{"dit":{"weight_transfer":"allgather"},"text_encoder":{"weight_transfer":"allgather"}}}' \
+  --usp 4
 
-# Standard-loader rank-local weights, without DLO AllGather
+# Mix transfers independently: sharded DiT, loader-produced encoder layout
 vllm serve /path/to/model --omni \
-  --enable-distributed-layerwise-offload \
-  --data-parallel-size 4 \
-  --dlo-no-use-allgather
+  --diffusion-offload-config \
+  '{"mode":"layer","components":["dit","text_encoder"],"layer_options":{"dit":{"weight_transfer":"allgather"},"text_encoder":{"weight_transfer":"rank-local"}}}' \
+  --usp 4
 
-# Sequence parallel deployment
+# Full-topology rank-local DLO for an existing model-specific lifecycle
 vllm serve /path/to/model --omni \
   --enable-distributed-layerwise-offload \
+  --dlo-no-use-allgather \
   --usp 4
 ```
 
+A compact layer config whose selected components are all `rank-local` and
+whose `resident_layers` is zero resolves to the ordinary layerwise backend. It
+uses the ordinary loader and does not request DLO's direct-checkpoint mmap.
+Keep the compatibility DLO flags for a model-specific full-topology rank-local
+lifecycle; do not add resident layers solely to force backend selection.
+
 ```python
 from vllm_omni import Omni
 
 omni = Omni(
     model="/path/to/model",
-    enable_distributed_layerwise_offload=True,
-    dlo_use_allgather=True,
+    diffusion_offload_config={
+        "mode": "layer",
+        "components": ["dit", "text_encoder"],
+        "layer_options": {
+            "dit": {
+                "weight_transfer": "rank-local",
+                "resident_layers": 20,
+            },
+            "text_encoder": {
+                "weight_transfer": "allgather",
+            },
+        },
+        "pin_memory": True,
+    },
 )
 ```
 
 ## Flags
 
-| Flag | Meaning | Default |
+| Setting | Meaning | Default |
 | --- | --- | --- |
-| `--enable-distributed-layerwise-offload` | Enable DLO | `false` |
-| `--data-parallel-size N` | DP ranks and AllGather weight-sharding group | `1` |
-| `--dlo-use-allgather` | Shard host weights and reconstruct with AllGather | `true` |
-| `--dlo-no-use-allgather` | Stream complete rank-local blocks without a DLO weight collective | `false` |
-| `--dlo-resident-layers N` | Keep N leading main-DiT blocks on device; requires no-AllGather and model-declared resident paths | `0` |
+| `diffusion_offload_config.mode` | `module` or `layer` granularity | required |
+| `diffusion_offload_config.components` | Non-empty list containing `dit`, `text_encoder`, or both | required |
+| `layer_options.NAME.weight_transfer` | `rank-local` or `allgather` | `rank-local` |
+| `layer_options.dit.resident_layers` | Leading main-DiT blocks kept on device; requires `rank-local` and model-declared resident paths | `0` |
+| `diffusion_offload_config.pin_memory` | Pin streamed host memory for faster H2D copies | `true` |
+| `--data-parallel-size N` | DP ranks; DP is the DLO group when greater than one, otherwise SP is used | `1` |
 | `--host-weight-runtime-mode {disabled,preferred,required}` | HWR policy: no interaction, populate on a miss, or require an exact hit | `disabled` |
 | `--host-weight-runtime-root PATH` | Writable node-local HWR store shared by workers in one storage domain; required for `preferred` and `required` | unset |
 | `--dlo-host-registration-limit-gib N` | Optional per-worker ceiling for registering an HWR mmap; zero adds no ceiling | `0` |
 
+HWR and host-registration tuning remain on the compatibility DLO interface and
+cannot be combined with `diffusion_offload_config` in this release.
+
+## Component and weight-transfer matrix
+
+| Component | `allgather` | `rank-local` |
+| --- | --- | --- |
+| DiT | Shard each loader-visible block across the DLO DP group, or the SP group when DP is one | Stream each rank's complete loader-produced block |
+| Text encoder | Same shard + AllGather path when the model declares identical encoder weights across the DLO group | Stream each rank's complete encoder block, including encoder-TP shards |
+
+The text encoder's compute group and weight-transfer group are separate
+concepts. An encoder TP group owns different parameter shards, so it is not a
+valid AllGather offload group. Models opt in through
+`OffloadPlan.encoder_dlo_weight_replication`; otherwise multi-rank encoder
+AllGather fails at startup with guidance to use `rank-local`. This makes SP
+AllGather available to replicated encoders without silently corrupting an
+encoder-TP layout.
+
+`resident_layers` remains a DiT-only setting and currently requires the DiT
+transfer to be `rank-local`.
+
 ## Host-weight loading
 
-The diffusion loader chooses host storage before DLO is enabled. It first
+When the distributed backend is selected by AllGather, resident DiT layers, or
+the compatibility DLO flag, the diffusion loader chooses host storage before
+DLO is enabled. It first
 attempts to build a complete, validated direct-checkpoint mmap plan. If names,
 coverage, shape, dtype, topology, or loader-callback compatibility cannot be
 proven, it runs the ordinary model loader instead. DLO consumes that result and
@@ -97,20 +167,20 @@ With direct checkpoint mmap, the loader:
 6. preserves `post_load_weights()` and `validate_loaded_weights()` lifecycle
    hooks.
 
-For AllGather with a group larger than one, each process copies only its
-persistent shard and then releases the source mapping. For no-AllGather, each
-process keeps the mapping open and packs complete blocks through two bounded
-pinned staging slots. Processes mapping the same files on one node share the
-immutable pages; no-AllGather still performs a complete-block H2D copy in each
-process.
+For a DiT using `allgather` with a group larger than one, each process copies
+only its persistent shard and then releases the source mapping. For a DiT using
+`rank-local`, each process keeps the mapping open and packs complete blocks
+through two bounded pinned staging slots. Processes mapping the same files on
+one node share the immutable pages; rank-local transfer still performs a
+complete-block H2D copy in each process.
 
-When the effective DLO group size is one, `dlo_use_allgather=True` does not
+When the effective DLO group size is one, `weight_transfer="allgather"` does not
 perform a collective and uses the same rank-local transfer behavior.
 
 ### Final-layout Host Weight Runtime
 
 HWR is an opt-in startup optimization for models that declare the final-layout
-BF16 restore contract. Use it only with no-AllGather DLO:
+BF16 restore contract. Use it only when the selected DiT uses `rank-local`:
 
 The validated BF16 model contracts currently cover MiniMax H3 and
 `black-forest-labs/FLUX.2-klein-4B`. FLUX.2-klein-9B shares the same model
@@ -204,7 +274,7 @@ H2D behavior.
 #### Registered direct H2D
 
 Registration is attempted automatically only for an eligible warm HWR hit when
-`pin_cpu_memory` is enabled. A successful path is:
+the existing `pin_cpu_memory` policy is enabled. A successful path is:
 
 ```text
 shared read-only HWR mmap -> existing rotating HBM block buffer -> GPU kernel
@@ -217,10 +287,10 @@ kernels access host memory directly.
 CUDA requires read-only host-registration capability for the immutable HWR
 mapping. Unsupported capability, a positive registration limit smaller than
 the complete page-aligned mapping, or a safely rolled-back registration error
-falls back to bounded staging. Programmatic configurations can set
-`pin_cpu_memory=False` to disable registration explicitly. A successful
-registration locks the complete mapped range in host memory for that worker's
-lifetime, so use
+falls back to bounded staging. Programmatic compatibility configurations can
+set `pin_cpu_memory=False` to disable registration explicitly.
+A successful registration locks the complete mapped range in host memory for
+that worker's lifetime, so use
 `--dlo-host-registration-limit-gib` when an operator-enforced ceiling is
 required.
 
@@ -228,8 +298,8 @@ Shutdown drains H2D work, releases hook/source references, unregisters every
 range, and only then closes the HWR lease. AllGather and direct checkpoint mmap
 retain their existing transfer paths.
 
-When HWR is disabled, DLO is disabled, or AllGather is enabled, the loader does
-not resolve HWR sources or construct its store.
+When HWR is disabled, DLO is disabled, or the DiT uses `allgather`, the loader
+does not resolve HWR sources or construct its store.
 
 ## Declarative topology
 
@@ -241,38 +311,54 @@ from vllm_omni.diffusion.offloader import OffloadPlan
 
 class MyPipeline(nn.Module):
     _dit_modules = ["transformer"]
+    _encoder_modules = ["prompt_model"]
     _offload_plan = OffloadPlan(
         block_attrs={"transformer": ("blocks",)},
         offload_submodules={"context_encoder": "layers"},
+        encoder_component_types={"prompt_model": "text_encoder"},
+        encoder_block_attrs={"prompt_model": ("encoder.layers",)},
+        encoder_dlo_weight_replication=frozenset({"prompt_model"}),
+        on_demand_component_paths=frozenset({"prompt_model"}),
     )
 ```
 
-When no plan exists, discovery falls back to
-`_layerwise_offload_blocks_attrs` and then heuristic attribute lookup.
+`encoder_component_types` maps arbitrary encoder paths to the public
+`text_encoder` selector. `encoder_dlo_weight_replication` is the explicit
+safety declaration for reusing the DiT DLO group. On-demand encoder non-block
+state is loaded and offloaded by its pipeline phase.
+
+When no plan exists, DiT discovery falls back to
+`_layerwise_offload_blocks_attrs` and then heuristic attribute lookup;
+undeclared auxiliary components remain resident.
 
 ## Data-parallel concurrency
 
 With `data_parallel_size > 1` and AllGather enabled, the scheduler can process
-up to `dp_size` requests per denoising step. Every concurrent request must set
-the same explicit `num_inference_steps`; `None` is rejected because every rank
-must enter each collective.
+up to `dp_size` requests per denoising step. Concurrent requests must resolve
+to the same denoise schedule, so every request must provide the same explicit
+`num_inference_steps`. Pipeline defaults are rejected because different request
+modes can resolve the same `None` value to different schedules.
 
 ## Limitations
 
 - Direct checkpoint mmap currently requires TP1. TP greater than one is
   outside the Phase A shared-mmap support scope and falls back before model
   mutation to the ordinary TP-aware loader. DLO can stream that runtime layout,
-  and eligible no-AllGather configurations may use HWR final-layout artifacts,
+  and eligible DiT rank-local configurations may use HWR final-layout artifacts,
   but direct checkpoint mmap provides no shared-mmap host-memory guarantee for
   TP greater than one.
-- HSDP plus AllGather is rejected to avoid double sharding. HSDP without
-  AllGather has limited end-to-end validation.
+- HSDP with `allgather` on any selected component is rejected to avoid double
+  sharding. HSDP with rank-local transfers has limited end-to-end validation.
+- With data parallelism, `allgather` cannot be combined with TeaCache or
+  Cache-DiT on any selected component. Prompt-embedding cache is also
+  incompatible when the text encoder uses `allgather`. Rank-local cache hits
+  could otherwise make ranks enter different weight collectives.
 - Per-tensor online FP8 linears use the ordinary loader and can run with either
-  DLO transfer path. With AllGather, every rank temporarily materializes the
-  complete FP8 model in host memory before DLO retains only its shard. Other
-  online quantization methods require no-AllGather until their runtime layouts
-  are validated.
-- Resident leading layers require `--dlo-no-use-allgather` and a model
+  DiT transfer path. With DiT `allgather`, every rank temporarily materializes
+  the complete FP8 model in host memory before DLO retains only its shard.
+  Other online quantization methods require rank-local transfer for the
+  affected component until their runtime layouts are validated.
+- Resident leading layers require DiT `rank-local` transfer and a model
   `OffloadPlan` that declares eligible `resident_dit_paths`.
 - DP concurrency requires an explicit, identical inference-step count.
 

docs/user_guide/diffusion/offloader/layerwise_offload.md

@@ -18,8 +18,9 @@ stream.
 | ... | ... | ... | ... |
 | last block | Prefetch block 0 | Compute last block | Free last block |
 
-Encoders, VAE modules, and non-block DiT modules such as embeddings and norms
-remain device resident.
+Selected, plan-declared text-encoder blocks can use the same rank-local
+streaming mechanism. Image encoders, VAEs, and non-block DiT modules remain
+device resident.
 
 ## Usage
 
@@ -28,15 +29,44 @@ from vllm_omni import Omni
 
 omni = Omni(
     model="Wan-AI/Wan2.2-T2V-A14B-Diffusers",
-    enable_layerwise_offload=True,
+    diffusion_offload_config={
+        "mode": "layer",
+        "components": ["dit"],
+    },
 )
 ```
 
 ```bash
 vllm serve Wan-AI/Wan2.2-T2V-A14B-Diffusers \
-  --omni --enable-layerwise-offload
+  --omni \
+  --diffusion-offload-config '{"mode":"layer","components":["dit"]}'
 ```
 
+## Component selection
+
+Add `dit`, `text_encoder`, or both to the `components` list. Omitting
+`layer_options` uses safe rank-local defaults:
+
+- `["dit"]` streams only DiT blocks.
+- `["text_encoder"]` streams only declared text-encoder blocks.
+- Listing both streams both components.
+
+```bash
+# DiT-only layer offload
+vllm serve /path/to/model --omni \
+  --diffusion-offload-config \
+  '{"mode":"layer","components":["dit"]}'
+
+# Stream a model-declared text encoder while keeping the DiT resident
+vllm serve /path/to/model --omni \
+  --diffusion-offload-config \
+  '{"mode":"layer","components":["text_encoder"]}'
+```
+
+Encoder categories are resolved from `OffloadPlan.encoder_component_types`
+first. A name-based fallback is retained for pipelines that predate
+`OffloadPlan`.
+
 ## Model integration
 
 Transformer classes declare containers of executable blocks:
@@ -53,13 +83,29 @@ class Flux2Transformer2DModel(nn.Module):
     ]
 ```
 
-See the [layerwise design](../../../design/feature/offloader/layerwise_offload.md) for
-the discovery and hook invariants. `OffloadPlan` is a separate declarative
-topology path for distributed layerwise offload.
+Auxiliary components use declarative pipeline metadata:
+
+```python
+from vllm_omni.diffusion.offloader import OffloadPlan
+
+
+class MyPipeline(nn.Module):
+    _encoder_modules = ["prompt_model"]
+    _offload_plan = OffloadPlan(
+        encoder_component_types={"prompt_model": "text_encoder"},
+        encoder_block_attrs={"prompt_model": ("encoder.layers",)},
+    )
+```
+
+See the [layerwise design](../../../design/feature/offloader/layerwise_offload.md)
+for the discovery and hook invariants. Both ordinary and distributed layerwise
+offload consume the same `OffloadPlan` metadata.
 
 ## Limitations
 
-- Single device only.
+- The default weight transfer is `rank-local`. Set a selected component's
+  `weight_transfer` to `allgather` to shard its host weights across a
+  compatible multi-device group; backend selection is automatic.
 - Setup consolidates and pins block parameters, increasing cold-start time.
 - Performance depends on block compute time and host-to-device bandwidth;
   lightweight blocks may not hide transfers.

docs/user_guide/diffusion/offloader/module_offload.md

@@ -2,14 +2,14 @@
 
 Model-level, or sequential, offloading keeps only the pipeline component group
 currently executing on the accelerator. It is the simplest offload strategy
-and is selected with `--enable-cpu-offload`.
+and is selected with `diffusion_offload_config.mode="module"`.
 
 ## How it works
 
 Pre-forward hooks enforce mutual exclusion between DiT and encoder modules:
 
-- before an encoder runs, the DiT moves to CPU;
-- before a DiT runs, encoders and other DiTs move to CPU; and
+- before an encoder runs, selected DiTs move to CPU;
+- before a DiT runs, selected encoders and other selected DiTs move to CPU; and
 - VAE modules remain on the accelerator.
 
 Pinned host memory reduces transfer overhead. Transfers occur at phase
@@ -22,15 +22,28 @@ from vllm_omni import Omni
 
 omni = Omni(
     model="Wan-AI/Wan2.2-T2V-A14B-Diffusers",
-    enable_cpu_offload=True,
+    diffusion_offload_config={
+        "mode": "module",
+        "components": ["dit", "text_encoder"],
+    },
 )
 ```
 
 ```bash
 vllm serve Wan-AI/Wan2.2-T2V-A14B-Diffusers \
-  --omni --enable-cpu-offload
+  --omni \
+  --diffusion-offload-config \
+  '{"mode":"module","components":["dit","text_encoder"]}'
 ```
 
+List a component to select it for offload. Module mode rejects `layer_options`
+such as `weight_transfer` and `resident_layers`. The
+`enable_cpu_offload=True` compatibility entry point remains supported. New
+integrations should prefer the explicit config; existing model-specific stage
+lifecycles do not need to migrate until equivalent component coverage exists.
+For example, MiniMax-H3's compatibility lifecycle also stages its VAEs, while
+the compact selector intentionally covers only `dit` and `text_encoder`.
+
 ## Model integration
 
 Pipelines should implement `SupportsComponentDiscovery`:
@@ -73,9 +86,10 @@ while reusing sequential `.to()` movers.
 
 ## Limitations
 
-- Single device only.
+- Transfers are rank-local; module mode does not shard host payloads or add a
+  weight AllGather across ranks.
 - Higher cold-start latency.
 - Transfers between encoder and denoising phases add latency.
 
-See the [model-level design](../../../design/feature/offloader/cpu_offload.md)
+See the [model-level design](../../../design/feature/offloader/module_offload.md)
 for lifecycle and extension invariants.

docs/user_guide/examples/online_serving/diffusers_pipeline_adapter.md

@@ -16,6 +16,7 @@ implemented with clear Diffusers or vLLM-Omni behavior.
 - Sequence parallel execution
 - TeaCache / Cache-DiT acceleration
 - Step-wise execution (continuous batching)
+- Component-selective `diffusion_offload_config` (the legacy whole-pipeline CPU-offload options remain supported)
 
 For these features, it is recommended to use natively supported pipelines instead.
 

examples/offline_inference/image_to_video/README.md

@@ -152,9 +152,6 @@ python image_to_video.py \
 | `--vae-use-tiling` | flag | off | Enable VAE tiling for memory optimization |
 | `--enable-cpu-offload` | flag | off | Enable CPU offloading for diffusion models |
 | `--enable-layerwise-offload` | flag | off | Enable layerwise offloading on DiT modules |
-| `--enable-distributed-layerwise-offload` | flag | off | Enable distributed layerwise offload |
-| `--dlo-use-allgather` / `--dlo-no-use-allgather` | flag | on / off | Select sharded AllGather reconstruction (default) or rank-local streaming |
-| `--dlo-resident-layers` | int | `0` | Leading main-DiT blocks kept device-resident during distributed layerwise offload |
 | `--cfg-parallel-size` | int | `1` | Set to `2` to enable CFG Parallel |
 | `--tensor-parallel-size` | int | `1` | Tensor parallel size (effective for models that support TP, e.g. LTX2) |
 | `--ulysses-degree` | int | `1` | Ulysses sequence parallel degree |
@@ -347,6 +344,8 @@ Key arguments:
 - `--hsdp-shard-size`: Number of GPUs to shard model weights across within each replica group. -1 (default) auto-calculates as world_size / replicate_size.
 - `--hsdp-replicate-size`: Number of replica groups for HSDP. Each replica holds a full sharded copy. Default 1 means pure sharding (no replication).
 
+
+
 > ℹ️ If you encounter OOM errors, try using `--vae-use-slicing` and `--vae-use-tiling` to reduce memory usage.
 
 ## Wan2.1 VACE Conditional Tasks

examples/offline_inference/minimax_h3/dlo_lifecycle.py

@@ -121,7 +121,6 @@ def engine_kwargs(args: argparse.Namespace) -> dict[str, Any]:
             data_parallel_size=1,
             text_encoder_tp_size=args.tp_size,
             vae_patch_parallel_size=args.tp_size,
-            enable_distributed_layerwise_offload=False,
         )
     return common
 

examples/offline_inference/text_to_video/text_to_video.md

@@ -262,16 +262,14 @@ python text_to_video.py \
 - `--num-inference-steps`: sampling steps. Default depends on model.
 - `--fps`: frames per second for the saved MP4.
 - `--output`: path to save the generated video.
-- `--extra-body`: JSON dict of model-specific knobs (declared in `vllm_omni/model_extras/`), merged into sampling `extra_args`. See the Helios recipes above.
+- `--extra-body`: JSON object of model-specific generation knobs, filtered against the model's declared `extra_body_params` and merged into sampling `extra_args` (see the Helios and Cosmos3 examples above and [`vllm_omni/model_extras`](../../../vllm_omni/model_extras)).
 - `--vae-use-slicing`: enable VAE slicing for memory optimization.
 - `--vae-use-tiling`: enable VAE tiling for memory optimization.
 - `--cfg-parallel-size`: set it to 2 to enable CFG Parallel. See more examples in [`user_guide`](../../../docs/user_guide/diffusion/parallelism_acceleration.md#cfg-parallel).
 - `--tensor-parallel-size`: tensor parallel size (effective for models that support TP, e.g. LTX2).
 - `--enable-cpu-offload`: enable CPU offloading for diffusion models.
 - `--enable-layerwise-offload`: enable layerwise offloading on DiT modules.
-- `--enable-distributed-layerwise-offload`: enable distributed layerwise offload.
-- `--dlo-use-allgather` / `--dlo-no-use-allgather`: use sharded AllGather reconstruction (the default) or stream rank-local weights without AllGather.
-- `--dlo-resident-layers`: keep this many leading main-DiT blocks device-resident during distributed layerwise offload (default: `0`).
+- `--diffusion-offload-config`: component-selective offload JSON; do not combine it with the legacy flags above.
 - `--frame-rate`: generation FPS for pipelines that require it (e.g., LTX2).
 - `--audio-sample-rate`: fallback audio sample rate when the pipeline returns audio.
 - `--quantization`: quantization method (such as `fp8` for FP8).
@@ -279,7 +277,6 @@ python text_to_video.py \
 - `--lora-path`: path to PEFT LoRA adapter folder or checkpoint file.
 - `--lora-scale`: scale factor for LoRA weights.
 - `--lora-backend`: backend for loading LoRA adapters. Default: peft. Available options: peft, distill.
-- `--extra-body`: JSON object of model-specific generation params, filtered against the model's declared `extra_body_params` (see [`vllm_omni/model_extras`](../../../vllm_omni/model_extras)). Used by Cosmos3 (see above).
 
 ### Wan2.2-specific
 
@@ -299,4 +296,6 @@ python text_to_video.py \
 | 720p I2V | 7.0 | 6.0 | 50 |
 | CFG-distilled | (same) | 1.0 | 50 |
 
-> If you encounter OOM errors, try `--vae-use-slicing`, `--vae-use-tiling`, `--enable-cpu-offload`, or `--quantization fp8`.
+> If you encounter OOM errors, try `--vae-use-slicing`, `--vae-use-tiling`,
+> module/layer offload through `--diffusion-offload-config`, or
+> `--quantization fp8`.

examples/offline_inference/text_to_video/text_to_video.py

@@ -337,6 +337,12 @@ def parse_args() -> argparse.Namespace:
         action="store_true",
         help="Enable layerwise (blockwise) offloading on DiT modules.",
     )
+    parser.add_argument(
+        "--diffusion-offload-config",
+        type=json.loads,
+        default=None,
+        help="Component-selective diffusion offload config as JSON.",
+    )
     parser.add_argument(
         "--enable-distributed-layerwise-offload",
         action="store_true",
@@ -534,6 +540,7 @@ def main():
     omni_kwargs = dict(
         model=args.model,
         enable_layerwise_offload=args.enable_layerwise_offload,
+        diffusion_offload_config=args.diffusion_offload_config,
         vae_use_slicing=args.vae_use_slicing,
         vae_use_tiling=args.vae_use_tiling,
         enable_cpu_offload=args.enable_cpu_offload,

examples/online_serving/diffusers_pipeline_adapter/README.md

@@ -13,6 +13,7 @@ implemented with clear Diffusers or vLLM-Omni behavior.
 - Sequence parallel execution
 - TeaCache / Cache-DiT acceleration
 - Step-wise execution (continuous batching)
+- Component-selective `diffusion_offload_config` (the legacy whole-pipeline CPU-offload options remain supported)
 
 For these features, it is recommended to use natively supported pipelines instead.
 

mkdocs.yml

@@ -83,7 +83,9 @@ plugins:
         design/feature/diffusion_request_level_batching.md: design/feature/diffusion_continuous_batching.md
         design/feature/diffusion_step_execution.md: design/feature/diffusion_continuous_batching.md
         design/feature/cpu_offloading.md: design/feature/offloader/README.md
-        design/feature/model_level_offload.md: design/feature/offloader/cpu_offload.md
+        user_guide/diffusion/offloader/cpu_offload.md: user_guide/diffusion/offloader/module_offload.md
+        design/feature/model_level_offload.md: design/feature/offloader/module_offload.md
+        design/feature/offloader/cpu_offload.md: design/feature/offloader/module_offload.md
         design/feature/layerwise_offload.md: design/feature/offloader/layerwise_offload.md
         design/feature/distributed_layerwise_offload.md: design/feature/offloader/distributed_layerwise_offload.md
   - glightbox

recipes/MiniMaxAI/MiniMax-H3-Disaggregated.md

@@ -42,7 +42,8 @@ encoder's BF16 configuration and the video/audio VAEs' FP32 precision. Select
 one offload strategy per deployment:
 
 ```bash
-# Stage 1 model-level CPU offload.
+# Stage 1 full-topology model offload, including MiniMax-H3's VAEs. Keep the
+# compatibility alias because the compact selector covers only DiT/text encoder.
 --stage-overrides '{"1":{"enable_cpu_offload":true}}'
 
 # Stage 1 distributed layerwise offload. Tune resident layers for available RAM.

recipes/MiniMaxAI/MiniMax-H3.md

@@ -117,11 +117,12 @@ same time on a host sized for this minimum.
 The consumer-GPU profiles below are HBM budgets only. They still require the
 host-RAM budget above.
 
-### Single GPU: accuracy and memory first
+### Single GPU: blockwise capacity path
 
-The single-GPU configuration uses model-level CPU offload.
-This matches the accuracy-qualified reference path and prevents the Qwen3-VL
-encoder and DiT from being resident on the GPU at the same time.
+Use ordinary layerwise offload when one GPU cannot keep the Qwen3-VL encoder
+and DiT resident together. The component list below streams the active DiT,
+the Qwen vision blocks, and the first 50 Qwen text layers. Encoder blocks stay
+rank-local; the video/audio VAEs remain resident.
 
 ```bash
 export MODEL=MiniMaxAI/MiniMax-H3
@@ -129,20 +130,37 @@ export PORT=8091
 
 CUDA_VISIBLE_DEVICES=0 \
 VLLM_WORKER_MULTIPROC_METHOD=spawn \
-VLLM_OMNI_VIDEO_SYNC_TIMEOUT=1800 \
+VLLM_OMNI_VIDEO_SYNC_TIMEOUT=14400 \
 vllm serve "${MODEL}" \
   --omni \
   --host 0.0.0.0 \
   --port "${PORT}" \
   --trust-remote-code \
+  --task-type fl2va \
   --num-gpus 1 \
-  --enable-cpu-offload \
+  --diffusion-offload-config \
+  '{"mode":"layer","components":["dit","text_encoder"]}' \
+  --enforce-eager \
   --diffusion-attention-backend FLASH_ATTN
 ```
 
-Use a GPU with enough memory for the active H3 component and enough system RAM
-for both offloaded DiTs plus the shared components. Model-level offload keeps
-the two DiTs mutually exclusive on GPU, but adds PCIe/NVLink transfer latency.
+This is a capacity profile, not a latency profile: every denoising step streams
+DiT blocks over the host link, while encoder blocks are streamed only during
+the short conditioning phase. It requires host memory for the complete
+checkpoint plus pinned transfer buffers. Re-measure peak HBM on the target
+shape; this command does not claim a particular GPU model as validated.
+
+A one-B300 correctness smoke selected only `text_encoder` (keeping the DiT and
+VAEs resident) and completed a 384x672, 5-second T2VA request with two denoise
+steps. It installed 77 encoder hooks across the vision and text stacks, used a
+77,728 MiB worker peak, and measured 2.803 seconds encode, 2.706 seconds
+denoise, and 4.503 seconds decode. These reduced-step numbers validate the
+execution path; they are not a quality or production-latency benchmark.
+
+To use whole-component behavior, change the config to
+`{"mode":"module","components":["dit","text_encoder"]}`. Module
+offload swaps the complete encoder and DiT and therefore has a higher
+encode-phase peak than encoder blockwise offload.
 
 ### Two 24/32 GB GPUs: TP2 distributed layerwise offload
 
@@ -483,10 +501,9 @@ Add this option to an existing H3 server command:
 
 Use the FL2VA-only partition for this capacity test. Loading the combined
 service would also load the Ref2VA DiT and would test a different memory
-budget. A no-offload capacity check should contain none of
-`--enable-cpu-offload`, `--enable-layerwise-offload`, or
-`--enable-distributed-layerwise-offload`. VAE tiling changes decode placement
-but does not offload model weights to the CPU.
+budget. A no-offload capacity check should omit
+`--diffusion-offload-config` and all legacy `--enable-*-offload` aliases. VAE
+tiling changes decode placement but does not offload model weights to the CPU.
 
 The run passes the capacity check when the server initializes, the request
 finishes without CUDA OOM or Xid errors, `peak_used_mib` remains below the
@@ -518,10 +535,11 @@ For example, keep the first main block's attention projections in BF16 with:
 ```
 
 The structured option replaces `--quantization fp8`. Online FP8 can be used
-with H3 layerwise offload and with both DLO transfer paths. The default
-AllGather path uses the ordinary loader to finalize FP8 weights and scales
-before sharding them across ranks. `--dlo-no-use-allgather` instead retains
-complete rank-local tensors and avoids the synchronized request-wave contract.
+with H3 layerwise offload and with either DiT DLO transfer. DiT `allgather`
+uses the ordinary loader to finalize FP8 weights and scales before sharding
+them across ranks. DiT `rank-local` instead retains complete loader-produced
+tensors and avoids the synchronized request-wave contract. H3's TP-sharded
+text encoder uses `rank-local`.
 
 ## AMD ROCm (gfx942 / gfx950)
 
@@ -567,6 +585,13 @@ vllm serve "${MODEL}" \
   --diffusion-attention-backend FLASH_ATTN
 ```
 
+This validated ROCm capacity recipe intentionally retains the compatibility
+full-topology alias because MiniMax-H3 also stages its VAEs on that path. The
+compact API in this release selects only `dit` and `text_encoder`, so replacing
+the flag would change residency rather than perform a mechanical migration.
+No removal deadline is assigned until the compact API offers equivalent
+component coverage.
+
 ### ROCm four GPUs
 
 The best-practice CUDA four-GPU configuration works on ROCm with the changes above.

tests/config/test_omni_config.py

@@ -6,6 +6,7 @@ from __future__ import annotations
 
 from dataclasses import fields
 from inspect import Parameter, signature
+from multiprocessing.reduction import ForkingPickler
 from pathlib import Path
 
 import msgspec
@@ -1390,3 +1391,62 @@ def test_diffusion_quantization_mapping_reaches_terminal_config(monkeypatch):
 
     assert cfg.quantization_config is not None
     assert cfg.quantization_config.get_name() == "int8"
+
+
+def test_compact_offload_config_reaches_terminal_config(monkeypatch):
+    from vllm_omni.diffusion.data import OmniDiffusionConfig
+    from vllm_omni.diffusion.offloader.config import OffloadStrategy, resolve_offload_strategy
+
+    compact_config = {
+        "mode": "layer",
+        "components": ["dit", "text_encoder"],
+        "layer_options": {
+            "dit": {"weight_transfer": "rank-local", "resident_layers": 12},
+            "text_encoder": {"weight_transfer": "allgather"},
+        },
+    }
+    cfg = omni_config_module._DiffusionConfigProjection.from_kwargs(
+        diffusion_offload_config=compact_config,
+    )
+
+    monkeypatch.setattr(OmniDiffusionConfig, "_resolve_master_port", lambda _self: 29500)
+    monkeypatch.setattr(OmniDiffusionConfig, "enrich_config", lambda _self: None)
+    cfg.enrich_config()
+
+    assert cfg.diffusion_offload_config == compact_config
+    assert cfg.extras == {}
+    assert resolve_offload_strategy(cfg) is OffloadStrategy.DISTRIBUTED_LAYER_WISE
+    restored = ForkingPickler.loads(ForkingPickler.dumps(cfg))
+    assert resolve_offload_strategy(restored) is OffloadStrategy.DISTRIBUTED_LAYER_WISE
+
+
+def test_global_diffusion_offload_config_targets_only_diffusion_stages():
+    compact_config = {"mode": "layer", "components": ["dit"]}
+
+    config = _from_pipeline_key(
+        "minimax_h3_disaggregated",
+        cli_overrides={"diffusion_offload_config": compact_config},
+    )
+
+    assert not hasattr(config.stage_by_id(0), "diffusion_config")
+    assert config.stage_by_id(1).diffusion_config.diffusion_offload_config == compact_config
+
+
+def test_explicit_llm_stage_diffusion_offload_override_is_rejected():
+    with pytest.raises(ValueError, match="no structured config owner: diffusion_offload_config"):
+        _from_pipeline_key(
+            "minimax_h3_disaggregated",
+            cli_overrides={
+                "stage_0_diffusion_offload_config": {"mode": "layer", "components": ["dit"]},
+            },
+        )
+
+
+def test_compact_offload_config_is_validated_during_projection():
+    with pytest.raises(ValueError, match="Unknown diffusion offload mode"):
+        omni_config_module._DiffusionConfigProjection.from_kwargs(
+            diffusion_offload_config={
+                "mode": "layerwise",
+                "components": ["dit"],
+            }
+        )

tests/diffusion/diffusion_backend/test_diffusers_backend.py

@@ -224,6 +224,7 @@ class TestPipelineArgumentsHandling:
     @pytest.mark.parametrize(
         "feature_id",
         [
+            "compact_offload",
             "cfg_parallel",
             "tensor_parallel",
             "ulysses",
@@ -235,7 +236,11 @@ class TestPipelineArgumentsHandling:
         ],
     )
     def test_adapter_guard_unsupported_feature(self, feature_id):
-        if feature_id == "cfg_parallel":
+        if feature_id == "compact_offload":
+            od_config = _make_od_config(
+                diffusion_offload_config={"mode": "module", "components": ["dit"]},
+            )
+        elif feature_id == "cfg_parallel":
             od_config = _make_od_config(
                 parallel_config=DiffusionParallelConfig(cfg_parallel_size=2, sequence_parallel_size=1),
                 cache_backend="none",

tests/diffusion/model_loader/test_diffusers_loader.py

@@ -734,6 +734,153 @@ def test_dlo_transfers_loader_plan_and_skips_ordinary_weight_loading(monkeypatch
     assert loader.take_host_weight_plan() is None
 
 
+def test_compact_rank_local_layer_offload_skips_dlo_mmap_planning(monkeypatch):
+    import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
+    from vllm_omni.diffusion.offloader.config import materialize_legacy_offload_flags
+
+    od_config = SimpleNamespace(
+        dtype=torch.float32,
+        parallel_config=SimpleNamespace(use_hsdp=False, tensor_parallel_size=1),
+        quantization_config=None,
+        diffusion_offload_config={"mode": "layer", "components": ["dit"]},
+        enable_cpu_offload=False,
+        enable_layerwise_offload=False,
+        enable_distributed_layerwise_offload=False,
+        dlo_use_allgather=True,
+        dlo_resident_layers=0,
+        pin_cpu_memory=True,
+        model="unused",
+    )
+    materialize_legacy_offload_flags(od_config)
+    loader = DiffusersPipelineLoader(LoadConfig(), od_config)
+    model = nn.Module()
+    model.transformer = nn.Linear(2, 2, bias=False)
+    calls: list[str] = []
+
+    loader._init_from_load_format = lambda *_args, **_kwargs: model  # type: ignore[method-assign]
+    loader.load_weights = lambda _model: calls.append("load")  # type: ignore[method-assign]
+    loader._process_weights_after_loading = lambda *_args: calls.append("process")  # type: ignore[method-assign]
+    loader._apply_skip_softmax_calibration = lambda _model: None  # type: ignore[method-assign]
+    monkeypatch.setattr(
+        loader_mod,
+        "build_checkpoint_mmap_plan",
+        lambda *_args, **_kwargs: pytest.fail("ordinary rank-local offload must not plan DLO mmap"),
+    )
+
+    assert loader.load_model(load_device="cpu") is model
+    assert calls == ["load", "process"]
+    assert loader.take_host_weight_plan() is None
+
+
+class _UnsupportedOnlineQuantMethod:
+    uses_meta_device = True
+
+
+def _compact_layer_offload_config(
+    components: list[str],
+    layer_options: dict[str, dict[str, str]],
+) -> SimpleNamespace:
+    return SimpleNamespace(
+        dtype=torch.float32,
+        parallel_config=SimpleNamespace(
+            use_hsdp=False,
+            tensor_parallel_size=1,
+            data_parallel_size=2,
+            sequence_parallel_size=1,
+        ),
+        quantization_config=None,
+        diffusion_offload_config={
+            "mode": "layer",
+            "components": components,
+            "layer_options": layer_options,
+        },
+        enable_cpu_offload=False,
+        enable_layerwise_offload=False,
+        enable_distributed_layerwise_offload=False,
+        dlo_use_allgather=True,
+        dlo_resident_layers=0,
+        pin_cpu_memory=False,
+        host_weight_runtime_mode="disabled",
+        model="unused",
+    )
+
+
+def _stub_ordinary_loader(loader, model, monkeypatch, loader_mod) -> None:
+    loader._init_from_load_format = lambda *_args, **_kwargs: model  # type: ignore[method-assign]
+    loader.load_weights = lambda _model: None  # type: ignore[method-assign]
+    loader._process_weights_after_loading = lambda *_args: None  # type: ignore[method-assign]
+    loader._apply_skip_softmax_calibration = lambda _model: None  # type: ignore[method-assign]
+    monkeypatch.setattr(
+        loader_mod,
+        "build_checkpoint_mmap_plan",
+        lambda *_args, **_kwargs: HostWeightPlanResult(None, "test fallback"),
+    )
+
+
+def test_text_encoder_allgather_rejects_unsupported_online_quantization(monkeypatch):
+    config = _compact_layer_offload_config(
+        ["text_encoder"],
+        {"text_encoder": {"weight_transfer": "allgather"}},
+    )
+    loader = DiffusersPipelineLoader(LoadConfig(), config)
+    model = nn.Module()
+    model.text_encoder = nn.Linear(2, 2, bias=False)
+    model.text_encoder.quant_method = _UnsupportedOnlineQuantMethod()
+    loader._init_from_load_format = lambda *_args, **_kwargs: model  # type: ignore[method-assign]
+
+    with pytest.raises(ValueError, match="unsupported online methods: _UnsupportedOnlineQuantMethod"):
+        loader.load_model(load_device="cpu")
+
+
+def test_dit_allgather_ignores_unsupported_quantization_on_unselected_encoder(monkeypatch):
+    import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
+
+    config = _compact_layer_offload_config(
+        ["dit"],
+        {"dit": {"weight_transfer": "allgather"}},
+    )
+    loader = DiffusersPipelineLoader(LoadConfig(), config)
+    model = nn.Module()
+    model.transformer = nn.Linear(2, 2, bias=False)
+    model.text_encoder = nn.Linear(2, 2, bias=False)
+    model.text_encoder.quant_method = _UnsupportedOnlineQuantMethod()
+    _stub_ordinary_loader(loader, model, monkeypatch, loader_mod)
+
+    assert loader.load_model(load_device="cpu") is model
+
+
+def test_rank_local_encoder_quantization_does_not_mark_dit_mmap_plan_online(monkeypatch):
+    import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
+
+    config = _compact_layer_offload_config(
+        ["dit", "text_encoder"],
+        {
+            "dit": {"weight_transfer": "allgather"},
+            "text_encoder": {"weight_transfer": "rank-local"},
+        },
+    )
+    loader = DiffusersPipelineLoader(LoadConfig(), config)
+    model = nn.Module()
+    model.transformer = nn.Linear(2, 2, bias=False)
+    model.text_encoder = nn.Linear(2, 2, bias=False)
+    model.text_encoder.quant_method = _UnsupportedOnlineQuantMethod()
+    planned_online_quantization: list[bool] = []
+
+    loader._init_from_load_format = lambda *_args, **_kwargs: model  # type: ignore[method-assign]
+    loader.load_weights = lambda _model: None  # type: ignore[method-assign]
+    loader._process_weights_after_loading = lambda *_args: None  # type: ignore[method-assign]
+    loader._apply_skip_softmax_calibration = lambda _model: None  # type: ignore[method-assign]
+
+    def capture_plan(*_args, **kwargs):
+        planned_online_quantization.append(kwargs["online_quantization"])
+        return HostWeightPlanResult(None, "test fallback")
+
+    monkeypatch.setattr(loader_mod, "build_checkpoint_mmap_plan", capture_plan)
+
+    assert loader.load_model(load_device="cpu") is model
+    assert planned_online_quantization == [False]
+
+
 def test_dlo_plan_loads_component_sources_outside_planned_dit(monkeypatch):
     import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
 
@@ -866,7 +1013,7 @@ def test_dlo_allgather_online_fp8_uses_ordinary_loader(monkeypatch):
     )
 
     assert loader.load_model(load_device="cpu", device=torch.device("cpu")) is model
-    assert allowlist_models == [model]
+    assert allowlist_models == [model.transformer]
     assert calls == [("load", True), "process"]
     assert loader.take_host_weight_plan() is None
 
@@ -908,7 +1055,7 @@ def test_dlo_allgather_online_int8_uses_ordinary_loader(monkeypatch):
     )
 
     assert loader.load_model(load_device="cpu", device=torch.device("cpu")) is model
-    assert allowlist_models == [model]
+    assert allowlist_models == [model.transformer]
     assert calls == [("load", True), "process"]
     assert loader.take_host_weight_plan() is None
 
@@ -1024,7 +1171,7 @@ def test_dlo_allgather_online_mxfp8_uses_ordinary_loader(monkeypatch):
     )
 
     assert loader.load_model(load_device="cpu", device=torch.device("cpu")) is model
-    assert allowlist_models == [model]
+    assert allowlist_models == [model.transformer]
     assert calls == [("load", True), "process"]
     assert loader.take_host_weight_plan() is None
 

tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py

@@ -30,6 +30,36 @@ def test_pipeline_declares_layerwise_offload_components() -> None:
     assert Cosmos3OmniDiffusersPipeline._resident_modules == []
     assert hasattr(Cosmos3OmniDiffusersPipeline, "enable_omni_model_cpu_offload")
 
+    from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import (
+        Cosmos3LanguageModel,
+        Cosmos3VFMTransformer,
+    )
+
+    assert Cosmos3LanguageModel._layerwise_offload_blocks_attrs == ["layers"]
+    assert Cosmos3VFMTransformer._layerwise_offload_blocks_attrs == ["gen_layers"]
+
+
+def test_component_selective_model_offload_fails_before_component_loading(monkeypatch) -> None:
+    from vllm_omni.diffusion.models.cosmos3 import pipeline_cosmos3 as pipeline_module
+
+    monkeypatch.setattr(
+        pipeline_module,
+        "get_local_device",
+        lambda: pytest.fail("component loading must not start for an unsupported selector"),
+    )
+    config = SimpleNamespace(
+        diffusion_offload_config={"mode": "module", "components": ["dit"]},
+        enable_cpu_offload=False,
+        enable_layerwise_offload=False,
+        enable_distributed_layerwise_offload=False,
+        dlo_use_allgather=True,
+        dlo_resident_layers=0,
+        pin_cpu_memory=True,
+    )
+
+    with pytest.raises(ValueError, match="does not support the dit/text_encoder component selector"):
+        pipeline_module.Cosmos3OmniDiffusersPipeline(od_config=config)
+
 
 class StubScheduler:
     def __init__(

tests/diffusion/models/magi2/test_native_dlo_allgather.py

@@ -220,7 +220,9 @@ def _load_mmap_transform_and_reconstruct(
         tensor_transforms=backend._mmap_transforms_by_tensor_id,
     )
     assert all(not shard.requires_grad for shard in cpu_shards.values())
-    assert all(parameter.numel() == 0 for parameter in target_block.parameters())
+    # Shard preparation is non-destructive; initialize_hook clears source
+    # storage only after every shard and metadata entry has been prepared.
+    assert all(parameter.numel() > 0 for parameter in target_block.parameters())
 
     for dtype, local_shard in cpu_shards.items():
         gathered = torch.empty(

tests/diffusion/models/magi2/test_pipeline_magi2.py

@@ -578,6 +578,8 @@ class _FakeLayerwiseHook:
     def __init__(self) -> None:
         self._prev_hook = None
         self.current_slot = 0
+        self.dp_size = 1
+        self.rank_local_mmap = False
         self.prefetch_layer = Mock()
         self.get_weights = Mock()
 
@@ -625,7 +627,7 @@ def test_dlo_no_allgather_enable_preserves_staged_aux_and_streams_blocks(
     assert len(hooks) == len(pipeline.transformer.block)
     assert all(stager.offload.call_count == 1 for stager in stagers)
     assert all(stager.load.call_count == 0 for stager in stagers)
-    hooks[-1].prefetch_layer.assert_called_once_with(
-        slot=hooks[0].current_slot,
-        non_blocking=False,
-    )
+    # First prefetch stays lazy so ranks with different local groups cannot
+    # enter collectives in a different order during startup.
+    for hook in hooks:
+        hook.prefetch_layer.assert_not_called()

tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py

@@ -487,6 +487,10 @@ def test_dlo_offload_plan_includes_token_refiner():
 
     assert MiniMaxH3Pipeline._offload_plan.offload_submodules == {"token_refiner": "blocks"}
     assert MiniMaxH3Pipeline._offload_plan.resident_dit_paths == frozenset({"transformer"})
+    assert MiniMaxH3Pipeline._offload_plan.encoder_component_types == {"text_encoder": "text_encoder"}
+    assert MiniMaxH3Pipeline._offload_plan.encoder_block_attrs == {
+        "text_encoder": ("vision.blocks", "text_model.layers")
+    }
 
 
 def test_joint_postprocess_is_multiprocessing_picklable():
@@ -1185,6 +1189,12 @@ def test_text_encoder_stub_constructs_without_group_or_weights():
     from vllm_omni.diffusion.models.minimax_h3.encoder import (
         MiniMaxH3Qwen3VLEncoder,
     )
+    from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import (
+        MiniMaxH3Pipeline,
+    )
+    from vllm_omni.diffusion.offloader.component_utils import (
+        get_encoder_block_groups,
+    )
 
     encoder = MiniMaxH3Qwen3VLEncoder(
         "/nonexistent/text_encoder",
@@ -1197,6 +1207,25 @@ def test_text_encoder_stub_constructs_without_group_or_weights():
     # The stub has no parameters, so it never contributes to the runner's
     # strict missing-parameter check on non-encoder ranks.
     assert list(encoder.named_parameters()) == []
+    assert (
+        get_encoder_block_groups(
+            encoder,
+            "text_encoder",
+            MiniMaxH3Pipeline._offload_plan,
+            strict=True,
+        )
+        == []
+    )
+
+    # An arbitrary empty module is not silently treated as a distributed
+    # stub; explicit plans still validate their declared paths.
+    with pytest.raises(ValueError, match=r"text_encoder\.vision\.blocks was not found"):
+        get_encoder_block_groups(
+            nn.Module(),
+            "text_encoder",
+            MiniMaxH3Pipeline._offload_plan,
+            strict=True,
+        )
 
 
 def test_global_quant_config_is_shared_by_dit_and_encoder():
@@ -1303,10 +1332,38 @@ def test_model_offload_uses_hooked_text_encoder_call():
 
 
 @pytest.mark.parametrize(
-    "offload_flag",
-    ["enable_layerwise_offload", "enable_distributed_layerwise_offload"],
+    ("enable_layerwise", "enable_distributed"),
+    [(True, False), (False, True)],
 )
-def test_layerwise_offload_releases_text_encoder(offload_flag):
+def test_legacy_layer_offload_preserves_minimax_stage_lifecycle(enable_layerwise, enable_distributed):
+    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+
+    pipeline = object.__new__(MiniMaxH3Pipeline)
+    torch.nn.Module.__init__(pipeline)
+    pipeline.od_config = SimpleNamespace(
+        diffusion_offload_config=None,
+        enable_cpu_offload=False,
+        enable_layerwise_offload=enable_layerwise,
+        enable_distributed_layerwise_offload=enable_distributed,
+    )
+    pipeline.text_encoder = Mock()
+    expected = torch.ones(2, 3)
+    pipeline.text_encoder.encode_ids.return_value = expected
+
+    actual = pipeline._encode_text_hidden(torch.tensor([1, 2]), {})
+
+    assert actual is expected
+    pipeline.text_encoder.load_to_device.assert_called_once_with()
+    pipeline.text_encoder.offload_to_cpu.assert_called_once_with()
+
+    vae = Mock()
+    with pipeline._component_on_device(vae):
+        pass
+    vae.load_to_device.assert_called_once_with()
+    vae.offload_to_cpu.assert_called_once_with()
+
+
+def test_layerwise_encoder_selection_releases_text_encoder():
     from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
 
     pipeline = object.__new__(MiniMaxH3Pipeline)
@@ -1315,8 +1372,11 @@ def test_layerwise_offload_releases_text_encoder(offload_flag):
         enable_cpu_offload=False,
         enable_layerwise_offload=False,
         enable_distributed_layerwise_offload=False,
+        diffusion_offload_config={
+            "mode": "layer",
+            "components": ["text_encoder"],
+        },
     )
-    setattr(pipeline.od_config, offload_flag, True)
     pipeline.text_encoder = Mock()
     expected = torch.ones(2, 3)
     pipeline.text_encoder.encode_ids.return_value = expected
@@ -1328,7 +1388,7 @@ def test_layerwise_offload_releases_text_encoder(offload_flag):
     pipeline.text_encoder.offload_to_cpu.assert_called_once_with()
 
 
-def test_distributed_layerwise_offload_releases_text_encoder():
+def test_layerwise_dit_only_keeps_text_encoder_resident():
     from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
 
     pipeline = object.__new__(MiniMaxH3Pipeline)
@@ -1336,7 +1396,8 @@ def test_distributed_layerwise_offload_releases_text_encoder():
     pipeline.od_config = SimpleNamespace(
         enable_cpu_offload=False,
         enable_layerwise_offload=False,
-        enable_distributed_layerwise_offload=True,
+        enable_distributed_layerwise_offload=False,
+        diffusion_offload_config={"mode": "layer", "components": ["dit"]},
     )
     pipeline.text_encoder = Mock()
     expected = torch.ones(2, 3)
@@ -1346,25 +1407,30 @@ def test_distributed_layerwise_offload_releases_text_encoder():
 
     assert actual is expected
     pipeline.text_encoder.load_to_device.assert_called_once_with()
-    pipeline.text_encoder.offload_to_cpu.assert_called_once_with()
+    pipeline.text_encoder.offload_to_cpu.assert_not_called()
 
 
-def test_distributed_layerwise_offload_stages_vae_component():
+def test_dit_encoder_selection_keeps_vae_resident():
     from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
 
     pipeline = object.__new__(MiniMaxH3Pipeline)
     torch.nn.Module.__init__(pipeline)
     pipeline.od_config = SimpleNamespace(
+        enable_cpu_offload=False,
         enable_layerwise_offload=False,
-        enable_distributed_layerwise_offload=True,
+        enable_distributed_layerwise_offload=False,
+        diffusion_offload_config={
+            "mode": "layer",
+            "components": ["dit", "text_encoder"],
+        },
     )
     component = Mock()
 
     with pipeline._component_on_device(component):
-        component.load_to_device.assert_called_once_with()
-        component.offload_to_cpu.assert_not_called()
+        pass
 
-    component.offload_to_cpu.assert_called_once_with()
+    component.load_to_device.assert_not_called()
+    component.offload_to_cpu.assert_not_called()
 
 
 def test_distributed_layerwise_resident_blocks_are_stage_scoped():
@@ -1413,57 +1479,6 @@ def test_distributed_layerwise_resident_blocks_can_be_skipped():
     controller.offload_resident_layers.assert_not_called()
 
 
-def test_encoder_layerwise_offload_keeps_tp_blocks_rank_local(monkeypatch):
-    from vllm_omni.diffusion.models.minimax_h3.encoder import (
-        MiniMaxH3Qwen3VLEncoder,
-    )
-
-    class Stack(torch.nn.Module):
-        def __init__(self, count):
-            super().__init__()
-            self.blocks = torch.nn.ModuleList([torch.nn.Linear(2, 2) for _ in range(count)])
-
-    class TextStack(torch.nn.Module):
-        def __init__(self, count):
-            super().__init__()
-            self.layers = torch.nn.ModuleList([torch.nn.Linear(2, 2) for _ in range(count)])
-
-    hooks = []
-
-    def fake_apply(block, next_block, device, stream, pin_memory):
-        hook = SimpleNamespace(
-            block=block,
-            next_block=next_block,
-            device=device,
-            stream=stream,
-            pin_memory=pin_memory,
-            _prev_hook=None,
-            offload_layer=Mock(),
-        )
-        hooks.append(hook)
-        return hook
-
-    monkeypatch.setattr(
-        "vllm_omni.diffusion.offloader.layerwise_backend.apply_block_hook",
-        fake_apply,
-    )
-    monkeypatch.setattr(
-        "vllm_omni.platforms.current_omni_platform.Stream",
-        Mock(return_value="copy-stream"),
-    )
-    encoder = object.__new__(MiniMaxH3Qwen3VLEncoder)
-    torch.nn.Module.__init__(encoder)
-    encoder.device_target = torch.device("cpu")
-    encoder.vision = Stack(2)
-    encoder.text_model = TextStack(3)
-
-    encoder.enable_omni_layerwise_offload(pin_memory=False)
-
-    assert len(hooks) == 5
-    assert all(hook._prev_hook is not None for hook in hooks)
-    assert all(hook.device == torch.device("cpu") for hook in hooks)
-
-
 def test_video_vae_keeps_reference_fp32_weights(monkeypatch):
     from vllm_omni.diffusion.models.minimax_h3 import vae as vae_module
 

tests/diffusion/models/minimax_h3/test_minimax_h3_dlo_lifecycle.py

@@ -1,6 +1,8 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
+from types import SimpleNamespace
+
 import pytest
 import torch
 
@@ -11,6 +13,44 @@ class _OffloadAbort(BaseException):
     pass
 
 
+def test_no_allgather_example_selects_distributed_backend():
+    from examples.offline_inference.minimax_h3.dlo_lifecycle import engine_kwargs
+    from vllm_omni.diffusion.offloader.base import OffloadConfig
+    from vllm_omni.diffusion.offloader.config import OffloadStrategy
+
+    args = SimpleNamespace(
+        model="/fake/MiniMax-H3/FL2VA",
+        mode="dlo-no-allgather",
+        dp_size=2,
+        tp_size=2,
+        batch_wait_ms=500.0,
+        init_timeout=1800.0,
+    )
+    kwargs = engine_kwargs(args)
+    od_config = SimpleNamespace(
+        diffusion_offload_config=None,
+        enable_cpu_offload=False,
+        enable_layerwise_offload=False,
+        enable_distributed_layerwise_offload=kwargs["enable_distributed_layerwise_offload"],
+        dlo_use_allgather=kwargs["dlo_use_allgather"],
+        dlo_resident_layers=kwargs["dlo_resident_layers"],
+        dlo_host_registration_limit_gib=0.0,
+        host_weight_runtime_mode="disabled",
+        pin_cpu_memory=True,
+        parallel_config=SimpleNamespace(
+            data_parallel_size=kwargs["data_parallel_size"],
+            sequence_parallel_size=1,
+            use_hsdp=False,
+        ),
+    )
+
+    config = OffloadConfig.from_od_config(od_config)
+
+    assert config.strategy is OffloadStrategy.DISTRIBUTED_LAYER_WISE
+    assert config.dlo_resident_layers == 0
+    assert not config.uses_allgather("dit")
+
+
 def test_encoder_non_block_children_use_one_shared_snapshot_stager(monkeypatch, mocker):
     from vllm_omni.diffusion.models.minimax_h3 import encoder as encoder_module
 
@@ -61,6 +101,7 @@ def test_manual_component_failure_forces_retained_cache_release(mocker):
     pipeline = object.__new__(MiniMaxH3Pipeline)
     torch.nn.Module.__init__(pipeline)
     pipeline.od_config = mocker.Mock()
+    pipeline.od_config.diffusion_offload_config = None
     pipeline.od_config.enable_layerwise_offload = False
     pipeline.od_config.enable_distributed_layerwise_offload = True
     pipeline._model_cpu_offload_modules = []
@@ -83,6 +124,7 @@ def test_manual_component_offload_failure_forces_retained_cache_release(mocker):
     pipeline = object.__new__(MiniMaxH3Pipeline)
     torch.nn.Module.__init__(pipeline)
     pipeline.od_config = mocker.Mock()
+    pipeline.od_config.diffusion_offload_config = None
     pipeline.od_config.enable_layerwise_offload = False
     pipeline.od_config.enable_distributed_layerwise_offload = True
     pipeline._model_cpu_offload_modules = []

tests/diffusion/models/minimax_h3/test_minimax_h3_lora.py

@@ -247,6 +247,55 @@ def test_h3_turbo_allows_distributed_layerwise_offload(monkeypatch):
     assert captured["unsupported_offload_mode"] is None
 
 
+@pytest.mark.parametrize(
+    ("mode", "components", "expected_mode"),
+    [
+        ("module", ["text_encoder"], None),
+        ("layer", ["text_encoder"], None),
+        ("module", ["dit"], "model-level CPU offload"),
+        ("layer", ["dit"], "layerwise offload"),
+    ],
+)
+def test_h3_turbo_offload_compatibility_follows_dit_selection(
+    monkeypatch,
+    mode,
+    components,
+    expected_mode,
+):
+    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+    from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as pipeline_module
+
+    pipeline = object.__new__(MiniMaxH3Pipeline)
+    torch.nn.Module.__init__(pipeline)
+    pipeline.partition = "fl2va"
+    pipeline._turbo_lora_adapter_ids = set()
+    pipeline._native_lora_adapter_ids = set()
+    pipeline._lora_sigma_schedules = {}
+    pipeline.od_config = SimpleNamespace(
+        diffusion_offload_config={"mode": mode, "components": components},
+        enable_cpu_offload=False,
+        enable_layerwise_offload=False,
+        enable_distributed_layerwise_offload=False,
+        dlo_use_allgather=True,
+        dlo_resident_layers=0,
+        pin_cpu_memory=True,
+    )
+    captured = {}
+
+    def load_turbo(**kwargs):
+        captured.update(kwargs)
+        return object(), object()
+
+    monkeypatch.setattr(pipeline_module, "load_minimax_h3_turbo_lora", load_turbo)
+
+    assert pipeline._load_diffusion_lora_adapter(
+        lora_request=_request("turbo"),
+        lora_path="turbo",
+        dtype=torch.bfloat16,
+    )
+    assert captured["unsupported_offload_mode"] == expected_mode
+
+
 def test_h3_turbo_accepts_only_the_declared_v1_artifact(tmp_path):
     unsupported = tmp_path / "minimax_h3_fl2v_turbo_8step_v1.0_bf16.safetensors"
     _write_tiny_turbo(unsupported)

tests/diffusion/models/minimax_h3/test_minimax_h3_offload.py

@@ -1,8 +1,9 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 from contextlib import contextmanager
 from types import SimpleNamespace
+from typing import Any
 from unittest.mock import Mock
 
 import pytest
@@ -111,6 +112,66 @@ def test_h3_model_cpu_offload_registers_direct_vae_stages(monkeypatch):
     remove_offload.assert_called_once_with([*dits, *stages])
 
 
+def test_h3_model_cpu_offload_keeps_unselected_vaes_resident(monkeypatch):
+    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+    from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as module
+
+    pipeline = object.__new__(MiniMaxH3Pipeline)
+    torch.nn.Module.__init__(pipeline)
+    pipeline.transformer = torch.nn.Linear(2, 2)
+    pipeline.transformers_ref = torch.nn.Linear(2, 2)
+    pipeline.text_encoder = torch.nn.Linear(2, 2)
+    pipeline.video_vae = torch.nn.Linear(2, 2)
+    pipeline.audio_vae = torch.nn.Linear(2, 2)
+    apply_offload = Mock()
+    monkeypatch.setattr(module, "apply_sequential_offload", apply_offload)
+
+    pipeline.enable_omni_model_cpu_offload(
+        device=torch.device("cpu"),
+        pin_memory=False,
+        use_hsdp=False,
+        offload_components=frozenset({"dit", "text_encoder"}),
+    )
+
+    dits = [pipeline.transformer, pipeline.transformers_ref]
+    stages = [pipeline.text_encoder, pipeline.video_vae, pipeline.audio_vae]
+    apply_offload.assert_called_once_with(
+        dit_modules=dits,
+        encoder_modules=stages,
+        device=torch.device("cpu"),
+        pin_memory=False,
+        use_hsdp=False,
+        offload_initial_dits=True,
+        offload_dit_modules=dits,
+        offload_encoder_modules=[pipeline.text_encoder],
+    )
+
+
+def test_h3_model_cpu_offload_rejects_unloaded_selected_encoder(monkeypatch):
+    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+    from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as module
+
+    pipeline = object.__new__(MiniMaxH3Pipeline)
+    torch.nn.Module.__init__(pipeline)
+    pipeline.transformer = torch.nn.Linear(2, 2)
+    pipeline.transformers_ref = torch.nn.Linear(2, 2)
+    pipeline.text_encoder = None
+    pipeline.video_vae = torch.nn.Linear(2, 2)
+    pipeline.audio_vae = torch.nn.Linear(2, 2)
+    apply_offload = Mock()
+    monkeypatch.setattr(module, "apply_sequential_offload", apply_offload)
+
+    with pytest.raises(ValueError, match="no loaded text encoder"):
+        pipeline.enable_omni_model_cpu_offload(
+            device=torch.device("cpu"),
+            pin_memory=False,
+            use_hsdp=False,
+            offload_components=frozenset({"text_encoder"}),
+        )
+
+    apply_offload.assert_not_called()
+
+
 @pytest.mark.parametrize("decode_fails", [False, True])
 def test_h3_model_cpu_offload_scopes_direct_vae_call(monkeypatch, decode_fails):
     from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
@@ -119,7 +180,7 @@ def test_h3_model_cpu_offload_scopes_direct_vae_call(monkeypatch, decode_fails):
     pipeline = object.__new__(MiniMaxH3Pipeline)
     torch.nn.Module.__init__(pipeline)
     component = torch.nn.Linear(2, 2)
-    events = []
+    events: list[tuple[Any, ...]] = []
 
     @contextmanager
     def record_component(value):
@@ -262,7 +323,7 @@ def test_h3_model_cpu_offload_shares_embedded_and_standalone_audio_scope(monkeyp
     torch.nn.Module.__init__(pipeline)
     pipeline.audio_vae = Mock()
     pipeline._model_cpu_offload_modules = [pipeline.audio_vae]
-    events = []
+    events: list[tuple[Any, ...]] = []
 
     @contextmanager
     def record_component(value):

tests/diffusion/models/sana_video/test_cache_offload.py

@@ -16,7 +16,7 @@ from tests.diffusion.models.sana_video.test_transformer_sana_video import (  # n
     _TINY_CONFIG,
     _init_distributed,
 )
-from tests.diffusion.offloader.test_layerwise_backend import DummyEvent, DummyStream, dummy_stream
+from tests.diffusion.offloader.helpers import DummyEvent, DummyStream, dummy_stream
 from vllm_omni.diffusion.attention import selector as attention_selector
 from vllm_omni.diffusion.attention.backends.sdpa import SDPABackend
 from vllm_omni.diffusion.cache.cachedit import CacheDiTBackend
@@ -226,7 +226,7 @@ def test_component_loading_uses_loader_device_not_runtime_device(monkeypatch):
     monkeypatch.setattr(
         pipeline_module,
         "_load_json",
-        lambda _model, filename, _local: ({"vae": [None, "FakeVAE"]} if filename == "model_index.json" else {}),
+        lambda _model, filename, _local: {"vae": [None, "FakeVAE"]} if filename == "model_index.json" else {},
     )
     monkeypatch.setattr(
         pipeline_module,

tests/diffusion/models/sensenova_u1/test_sensenova_u1_paged_decode.py

@@ -173,9 +173,14 @@ def test_length_lives_in_device_tensors():
 # ---------------------------------------------------------------------------
 
 cuda_only = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+vllm_flash_attn_only = pytest.mark.skipif(
+    current_omni_platform.is_rocm(),
+    reason="vLLM's bundled FlashAttention extension is CUDA-only",
+)
 
 
 @cuda_only
+@vllm_flash_attn_only
 @pytest.mark.cuda
 @pytest.mark.L4
 def test_paged_attention_matches_sdpa_over_the_same_buffer():
@@ -217,6 +222,7 @@ def test_paged_attention_matches_sdpa_over_the_same_buffer():
 
 
 @cuda_only
+@vllm_flash_attn_only
 @pytest.mark.cuda
 @pytest.mark.L4
 def test_a_captured_graph_follows_a_later_length():
@@ -278,6 +284,7 @@ def test_a_captured_graph_follows_a_later_length():
 
 
 @cuda_only
+@vllm_flash_attn_only
 @pytest.mark.cuda
 @pytest.mark.L4
 def test_a_captured_graph_writes_each_step_to_its_own_slot():
@@ -339,6 +346,7 @@ def test_a_captured_graph_writes_each_step_to_its_own_slot():
 
 
 @cuda_only
+@vllm_flash_attn_only
 @pytest.mark.cuda
 @pytest.mark.L4
 def test_a_whole_decode_loop_matches_the_unpaged_path():

tests/diffusion/models/wan2_2/test_wan22_i2v_pipeline.py

@@ -9,6 +9,7 @@ from PIL import Image
 from torch import nn
 
 from tests.diffusion.models.wan2_2.conftest import StubScheduler, StubTransformer, StubVAE, noop_progress_bar
+from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import _WAN_TEXT_ENCODER_OFFLOAD_PLAN
 from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_i2v import (
     Wan22I2VPipeline,
     get_wan22_i2v_post_process_func,
@@ -30,6 +31,10 @@ def test_wan22_i2v_postprocess_honors_request_output_type() -> None:
     assert output is video
 
 
+def test_i2v_pipeline_declares_text_encoder_offload_blocks() -> None:
+    assert Wan22I2VPipeline._offload_plan is _WAN_TEXT_ENCODER_OFFLOAD_PLAN
+
+
 def _make_i2v_pipeline(*, expand_timesteps: bool) -> Wan22I2VPipeline:
     pipeline = object.__new__(Wan22I2VPipeline)
     nn.Module.__init__(pipeline)

tests/diffusion/models/wan2_2/test_wan22_pipeline_helpers.py

@@ -18,6 +18,15 @@ from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import (
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
 
 
+def test_t2v_pipeline_declares_text_encoder_offload_blocks() -> None:
+    plan = wan22_module.Wan22Pipeline._offload_plan
+
+    assert plan.encoder_component_types == {"text_encoder": "text_encoder"}
+    assert plan.encoder_block_attrs == {"text_encoder": ("encoder.block",)}
+    assert plan.encoder_dlo_weight_replication == frozenset({"text_encoder"})
+    assert plan.on_demand_component_paths == frozenset()
+
+
 class _LatentDist:
     def sample(self, generator):
         assert isinstance(generator, torch.Generator)

tests/diffusion/models/wan2_2/test_wan22_s2v_pipeline.py

@@ -392,52 +392,31 @@ def test_encode_audio_skips_unshard_reshard_when_not_fsdp():
     assert "audio_emb" in result
 
 
-def test_s2v_pipeline_skips_cpu_offload_when_hsdp_enabled():
-    """Test that transformer.to('cpu') is NOT called when HSDP is active."""
-    from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_s2v import Wan22S2VPipeline
-
+@pytest.mark.parametrize(
+    ("components", "use_hsdp", "expected"),
+    [
+        (None, False, True),
+        (["dit"], False, True),
+        (["text_encoder"], False, False),
+        (["dit"], True, False),
+    ],
+)
+def test_s2v_dit_release_honors_component_selection(components, use_hsdp, expected):
     pipeline = object.__new__(Wan22S2VPipeline)
     nn.Module.__init__(pipeline)
+    compact = None if components is None else {"mode": "module", "components": components}
+    pipeline.od_config = SimpleNamespace(
+        diffusion_offload_config=compact,
+        enable_cpu_offload=True,
+        enable_layerwise_offload=False,
+        enable_distributed_layerwise_offload=False,
+        dlo_use_allgather=True,
+        dlo_resident_layers=0,
+        pin_cpu_memory=True,
+        parallel_config=SimpleNamespace(use_hsdp=use_hsdp),
+    )
 
-    od_config = MagicMock()
-    od_config.enable_cpu_offload = True
-    parallel_config = MagicMock()
-    parallel_config.use_hsdp = True
-    od_config.parallel_config = parallel_config
-    pipeline.od_config = od_config
-
-    mock_transformer = MagicMock()
-    pipeline.transformer = mock_transformer
-
-    # Simulate the offload decision from the forward loop
-    if pipeline.od_config.enable_cpu_offload and not getattr(pipeline.od_config.parallel_config, "use_hsdp", False):
-        pipeline.transformer.to("cpu")
-
-    mock_transformer.to.assert_not_called()
-
-
-def test_s2v_pipeline_allows_cpu_offload_when_hsdp_disabled():
-    """Test that transformer.to('cpu') IS called when HSDP is not active."""
-    from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_s2v import Wan22S2VPipeline
-
-    pipeline = object.__new__(Wan22S2VPipeline)
-    nn.Module.__init__(pipeline)
-
-    od_config = MagicMock()
-    od_config.enable_cpu_offload = True
-    parallel_config = MagicMock()
-    parallel_config.use_hsdp = False
-    od_config.parallel_config = parallel_config
-    pipeline.od_config = od_config
-
-    mock_transformer = MagicMock()
-    pipeline.transformer = mock_transformer
-
-    # Simulate the offload decision from the forward loop
-    if pipeline.od_config.enable_cpu_offload and not getattr(pipeline.od_config.parallel_config, "use_hsdp", False):
-        pipeline.transformer.to("cpu")
-
-    mock_transformer.to.assert_called_once_with("cpu")
+    assert pipeline._should_release_dit_before_decode() is expected
 
 
 def test_s2v_pipeline_hsdp_forward_complete_process():

tests/diffusion/models/wan2_2/test_wan22_vace_pipeline.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+from dataclasses import dataclass, field
 from types import SimpleNamespace
 
 import pytest
@@ -20,6 +21,46 @@ from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
 
 
+def test_vace_default_flow_shift_preserves_compact_offload_runtime_state(monkeypatch) -> None:
+    @dataclass
+    class RuntimeConfig:
+        flow_shift: float | None = None
+        diffusion_offload_config: dict[str, object] = field(
+            default_factory=lambda: {
+                "mode": "layer",
+                "components": ["text_encoder"],
+            }
+        )
+        enable_layerwise_offload: bool = True
+        post_init_calls: int = 0
+
+        def __post_init__(self) -> None:
+            self.post_init_calls += 1
+
+    config = RuntimeConfig()
+    resolved_offload = object()
+    config._resolved_diffusion_offload = resolved_offload
+    config._diffusion_offload_flags_materialized = True
+    captured = {}
+
+    def capture_base_init(self, *, od_config, prefix="") -> None:
+        captured.update(od_config=od_config, prefix=prefix)
+
+    monkeypatch.setattr(Wan22VACEPipeline.__mro__[1], "__init__", capture_base_init)
+
+    Wan22VACEPipeline(od_config=config, prefix="stage")
+
+    forwarded = captured["od_config"]
+    assert forwarded is not config
+    assert config.flow_shift is None
+    assert forwarded.flow_shift == 3.0
+    assert forwarded.diffusion_offload_config is config.diffusion_offload_config
+    assert forwarded._resolved_diffusion_offload is resolved_offload
+    assert forwarded._diffusion_offload_flags_materialized is True
+    assert forwarded.post_init_calls == 1
+    assert captured["prefix"] == "stage"
+
+
 def _make_vace_pipeline() -> Wan22VACEPipeline:
     pipeline = object.__new__(Wan22VACEPipeline)
     nn.Module.__init__(pipeline)

tests/diffusion/offloader/conftest.py

@@ -0,0 +1,26 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Shared runtime fixtures for diffusion offloader tests."""
+
+import gc
+
+import pytest
+import torch.distributed as dist
+
+from tests.helpers.runtime import get_distributed_init_method
+from vllm_omni.platforms import current_omni_platform
+
+
+@pytest.fixture(scope="module")
+def dist_group():
+    dist.init_process_group("gloo", rank=0, world_size=1, init_method=get_distributed_init_method())
+    try:
+        yield
+    finally:
+        if dist.is_initialized():
+            dist.destroy_process_group()
+        gc.collect()
+        if current_omni_platform.is_available():
+            current_omni_platform.empty_cache()
+            current_omni_platform.synchronize()

tests/diffusion/offloader/helpers.py

@@ -0,0 +1,107 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Shared model fixtures for diffusion offloader tests."""
+
+from contextlib import contextmanager
+
+import torch
+from torch import nn
+
+
+class DummyStream:
+    def wait_stream(self, _stream) -> None:
+        return None
+
+    def wait_event(self, _event) -> None:
+        return None
+
+
+class DummyEvent:
+    def record(self, _stream) -> None:
+        return None
+
+    def synchronize(self) -> None:
+        return None
+
+
+@contextmanager
+def dummy_stream(_stream):
+    yield None
+
+
+def patch_offload_runtime(monkeypatch, platform, *, synchronize: bool = False) -> None:
+    monkeypatch.setattr(platform, "Stream", DummyStream)
+    monkeypatch.setattr(platform, "Event", DummyEvent)
+    monkeypatch.setattr(platform, "current_stream", lambda: DummyStream())
+    monkeypatch.setattr(platform, "stream", dummy_stream)
+    if synchronize:
+        monkeypatch.setattr(platform, "synchronize", lambda: None)
+        monkeypatch.setattr(platform, "empty_cache", lambda: None)
+
+
+class _DummyBlock(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.weight = nn.Parameter(torch.randn(10, 10))
+
+
+class _SingleBlockModel(nn.Module):
+    _layerwise_offload_blocks_attrs = ["blocks"]
+
+    def __init__(self, num_blocks: int = 3):
+        super().__init__()
+        self.blocks = nn.ModuleList([_DummyBlock() for _ in range(num_blocks)])
+
+
+class _StagedEncoder(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.vision = nn.Module()
+        self.vision.blocks = nn.ModuleList([_DummyBlock(), _DummyBlock()])
+        self.text_model = nn.Module()
+        self.text_model.layers = nn.ModuleList([_DummyBlock(), _DummyBlock()])
+        self.load_calls = 0
+        self.offload_calls = 0
+        self.to_calls = 0
+
+    def load_to_device(self):
+        self.load_calls += 1
+
+    def offload_to_cpu(self):
+        self.offload_calls += 1
+        for hook in getattr(self, "_omni_layerwise_hooks", []):
+            hook.offload_layer()
+
+    def to(self, *args, **kwargs):
+        self.to_calls += 1
+        return super().to(*args, **kwargs)
+
+
+class _StagedVAE(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.proj = nn.Linear(2, 2)
+        self.offload_calls = 0
+        self.to_calls = 0
+
+    def load_to_device(self):
+        return None
+
+    def offload_to_cpu(self):
+        self.offload_calls += 1
+        return self.to("cpu")
+
+    def to(self, *args, **kwargs):
+        self.to_calls += 1
+        return super().to(*args, **kwargs)
+
+
+class _PlainEncoder(nn.Module):
+    """Standard encoder with no offload-specific lifecycle methods."""
+
+    def __init__(self):
+        super().__init__()
+        self.encoder = nn.Module()
+        self.encoder.block = nn.ModuleList([_DummyBlock(), _DummyBlock()])
+        self.final_norm = nn.Linear(2, 2)

tests/diffusion/offloader/test_distributed_layerwise_backend.py

@@ -6,8 +6,8 @@
 import gc
 import json
 import weakref
-from contextlib import contextmanager
 from types import SimpleNamespace
+from unittest.mock import Mock
 
 import pytest
 import torch
@@ -17,7 +17,15 @@ from torch import nn
 from torch.distributed.tensor import DeviceMesh, DTensor, Replicate
 
 import vllm_omni.diffusion.offloader.distributed_layerwise_backend as dist_backend_module
-from tests.helpers.runtime import get_distributed_init_method
+from tests.diffusion.offloader.helpers import (
+    DummyStream,
+    _DummyBlock,
+    _PlainEncoder,
+    _SingleBlockModel,
+    _StagedEncoder,
+    _StagedVAE,
+    patch_offload_runtime,
+)
 from vllm_omni.diffusion.data import validate_dlo_host_registration_options
 from vllm_omni.diffusion.model_loader.host_weight_plan import (
     HostWeightPlan,
@@ -51,53 +59,9 @@ from vllm_omni.platforms import current_omni_platform
 pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
 
 
-class DummyStream:
-    def wait_stream(self, _stream) -> None:
-        return None
-
-    def wait_event(self, _event) -> None:
-        return None
-
-
-class DummyEvent:
-    def record(self, _stream) -> None:
-        return None
-
-    def synchronize(self) -> None:
-        return None
-
-
-@contextmanager
-def dummy_stream(_stream):
-    yield None
-
-
-def _cleanup_distributed() -> None:
-    if dist.is_initialized():
-        dist.destroy_process_group()
-
-    gc.collect()
-    if current_omni_platform.is_available():
-        current_omni_platform.empty_cache()
-        current_omni_platform.synchronize()
-
-
-@pytest.fixture(scope="module")
-def dist_group():
-    dist.init_process_group("gloo", rank=0, world_size=1, init_method=get_distributed_init_method())
-    try:
-        yield
-    finally:
-        _cleanup_distributed()
-
-
 @pytest.fixture
 def patched_offload_runtime(monkeypatch):
-    monkeypatch.setattr(dist_backend_module.current_omni_platform, "Stream", DummyStream)
-    monkeypatch.setattr(dist_backend_module.current_omni_platform, "Event", DummyEvent)
-    monkeypatch.setattr(dist_backend_module.current_omni_platform, "current_stream", lambda: DummyStream())
-    monkeypatch.setattr(dist_backend_module.current_omni_platform, "stream", dummy_stream)
-    monkeypatch.setattr(dist_backend_module.current_omni_platform, "synchronize", lambda: None)
+    patch_offload_runtime(monkeypatch, dist_backend_module.current_omni_platform, synchronize=True)
 
 
 class TinyBlock(nn.Module):
@@ -113,6 +77,93 @@ def _make_values(start: float) -> torch.Tensor:
 
 
 class TestDistributedLayerwiseOffloadHook:
+    def test_buffer_only_block_reports_offloaded_state(self, patched_offload_runtime):
+        current_block = nn.Module()
+        current_block.register_buffer("state", torch.ones(2))
+        next_block = nn.Module()
+        next_block.register_buffer("state", torch.ones(2))
+        hook = DistributedLayerwiseOffloadHook(
+            next_block=next_block,
+            device=torch.device("cpu"),
+            dp_group=None,
+            dp_size=1,
+            rank=0,
+            pin_memory=False,
+        )
+        hook.initialize_hook(current_block)
+
+        assert hook.is_materialized
+        hook.offload_layer()
+        assert not hook.is_materialized
+
+    def test_initialize_failure_keeps_next_block_materialized(self, monkeypatch):
+        current_block = nn.Linear(2, 2)
+        next_block = nn.Linear(2, 2)
+        expected = {name: tensor.detach().clone() for name, tensor in next_block.state_dict().items()}
+        hook = DistributedLayerwiseOffloadHook(
+            next_block=next_block,
+            device=torch.device("cpu"),
+            dp_group=None,
+            dp_size=1,
+            rank=0,
+            pin_memory=False,
+        )
+
+        def fail_buffer_allocation():
+            raise RuntimeError("injected buffer allocation failure")
+
+        monkeypatch.setattr(hook, "_allocate_device_buffers", fail_buffer_allocation)
+
+        with pytest.raises(RuntimeError, match="buffer allocation failure"):
+            hook.initialize_hook(current_block)
+
+        for name, tensor in next_block.state_dict().items():
+            torch.testing.assert_close(tensor, expected[name])
+
+    @pytest.mark.parametrize(
+        ("kwargs", "message"),
+        [
+            ({"dp_size": 0, "rank": 0}, "dp_size must be a positive integer"),
+            ({"dp_size": 2, "rank": 2, "dp_group": object()}, "rank must satisfy"),
+            ({"dp_size": 2, "rank": 0}, "dp_group is required"),
+            (
+                {"dp_size": 2, "rank": 0, "dp_group": object(), "rank_local_mmap": True},
+                "rank_local_mmap requires dp_size=1",
+            ),
+        ],
+    )
+    def test_constructor_rejects_invalid_transport_topology(self, kwargs, message):
+        options = {"dp_group": None, "dp_size": 1, "rank": 0, **kwargs}
+
+        with pytest.raises(ValueError, match=message):
+            DistributedLayerwiseOffloadHook(
+                next_block=nn.Linear(2, 2),
+                device=torch.device("cpu"),
+                pin_memory=False,
+                **options,
+            )
+
+    def test_ring_probes_are_captured_before_any_block_is_cleared(self, patched_offload_runtime):
+        blocks = nn.ModuleList([nn.Linear(2, 2, bias=False) for _ in range(3)])
+        expected_middle = blocks[1].weight.detach().clone()
+        backend = DistributedLayerwiseOffloadBackend(
+            OffloadConfig(
+                strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
+                pin_cpu_memory=False,
+                dlo_use_allgather=False,
+            ),
+            torch.device("cpu"),
+        )
+
+        hooks = backend._install_hook_group(blocks, "dit")
+        shared_buffers = backend._allocate_shared_buffers(hooks)
+        for hook in hooks:
+            hook.gpu_buffers = shared_buffers
+
+        assert all(not hook.is_materialized for hook in hooks)
+        hooks[2].pre_forward(blocks[1])
+        torch.testing.assert_close(blocks[1].weight, expected_middle)
+
     def test_shard_and_pin_single_rank(self, dist_group, patched_offload_runtime):
         """With dp_size=1, the shard should equal the full weights."""
         current_block = TinyBlock(_make_values(1.0))
@@ -511,7 +562,7 @@ class TestDistributedLayerwiseOffloadHook:
         hook = DistributedLayerwiseOffloadHook(
             next_block=next_block,
             device=torch.device("cpu"),
-            dp_group=None,
+            dp_group=object(),
             dp_size=4,
             rank=1,
             pin_memory=False,
@@ -537,7 +588,7 @@ class TestDistributedLayerwiseOffloadHook:
         hook0 = DistributedLayerwiseOffloadHook(
             next_block=next_block,
             device=torch.device("cpu"),
-            dp_group=None,
+            dp_group=object(),
             dp_size=2,
             rank=0,
             pin_memory=False,
@@ -559,7 +610,7 @@ class TestDistributedLayerwiseOffloadHook:
         hook1 = DistributedLayerwiseOffloadHook(
             next_block=next_block1,
             device=torch.device("cpu"),
-            dp_group=None,
+            dp_group=object(),
             dp_size=2,
             rank=1,
             pin_memory=False,
@@ -583,7 +634,7 @@ class TestDistributedLayerwiseOffloadHook:
         hook = DistributedLayerwiseOffloadHook(
             next_block=next_block,
             device=torch.device("cpu"),
-            dp_group=None,
+            dp_group=object(),
             dp_size=2,
             rank=0,
             pin_memory=False,
@@ -602,7 +653,7 @@ class TestDistributedLayerwiseOffloadHook:
         hook = DistributedLayerwiseOffloadHook(
             next_block=next_block,
             device=torch.device("cpu"),
-            dp_group=None,
+            dp_group=object(),
             dp_size=2,
             rank=0,
             pin_memory=False,
@@ -689,6 +740,32 @@ class TestDistributedLayerwiseOffloadHook:
 
 
 class TestPinnedResidentLayerGroup:
+    def test_constructor_failure_leaves_every_block_materialized(self, monkeypatch):
+        blocks = [nn.Linear(2, 2), nn.Linear(2, 2)]
+        expected = [block.weight.detach().clone() for block in blocks]
+        shard_and_pin = DistributedLayerwiseOffloadHook._shard_and_pin
+        calls = 0
+
+        def fail_second_block(*args, **kwargs):
+            nonlocal calls
+            calls += 1
+            if calls == 2:
+                raise RuntimeError("injected resident setup failure")
+            return shard_and_pin(*args, **kwargs)
+
+        monkeypatch.setattr(DistributedLayerwiseOffloadHook, "_shard_and_pin", fail_second_block)
+
+        with pytest.raises(RuntimeError, match="resident setup failure"):
+            PinnedResidentLayerGroup(
+                blocks,
+                device=torch.device("cpu"),
+                copy_stream=DummyStream(),
+                pin_memory=False,
+            )
+
+        for block, weight in zip(blocks, expected, strict=True):
+            torch.testing.assert_close(block.weight, weight)
+
     def test_load_offload_reuses_pinned_master_weights(self, patched_offload_runtime):
         blocks = [nn.Linear(2, 2), nn.Linear(2, 2)]
         expected = []
@@ -843,20 +920,6 @@ class TestPinnedModuleStager:
         stager.offload()
 
 
-class _DummyBlock(nn.Module):
-    def __init__(self):
-        super().__init__()
-        self.weight = nn.Parameter(torch.randn(10, 10))
-
-
-class _SingleBlockModel(nn.Module):
-    _layerwise_offload_blocks_attrs = ["blocks"]
-
-    def __init__(self, num_blocks: int = 3):
-        super().__init__()
-        self.blocks = nn.ModuleList([_DummyBlock() for _ in range(num_blocks)])
-
-
 class _HWRPipeline(nn.Module):
     def __init__(self):
         super().__init__()
@@ -1285,6 +1348,31 @@ class TestMmapWeightLoading:
         backend.disable()
         assert lease.closed
 
+        backend.enable(pipeline)
+        assert backend.host_weight_plan is None
+        assert backend.enabled
+        backend.disable()
+
+    def test_rank_local_disable_drains_pending_transfers(self, patched_offload_runtime, monkeypatch):
+        pipeline = _HWRPipeline()
+        backend = DistributedLayerwiseOffloadBackend(
+            OffloadConfig(
+                strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
+                pin_cpu_memory=False,
+                dp_size=1,
+                dlo_use_allgather=False,
+            ),
+            torch.device("cpu"),
+        )
+        backend.enable(pipeline)
+        backend._all_hook_groups[0][0].prefetch_layer(0)
+        synchronize = Mock()
+        monkeypatch.setattr(current_omni_platform, "synchronize", synchronize)
+
+        backend.disable()
+
+        synchronize.assert_called()
+
     def test_hwr_registration_failure_falls_back_to_bounded_staging(
         self,
         monkeypatch,
@@ -1563,6 +1651,7 @@ class TestOffloadPlan:
         assert plan.block_attrs == {}
         assert plan.offload_submodules == {}
         assert plan.resident_dit_paths == frozenset()
+        assert plan.encoder_dlo_weight_replication == frozenset()
 
     def test_offload_plan_is_frozen(self):
         """OffloadPlan should be immutable (frozen=True)."""
@@ -1595,6 +1684,7 @@ class TestOffloadPlan:
                 self.transformer = Transformer()
 
         pipeline = Pipeline()
+        expected_parameters = {name: parameter.detach().clone() for name, parameter in pipeline.named_parameters()}
         backend = DistributedLayerwiseOffloadBackend(
             OffloadConfig(
                 strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
@@ -1610,10 +1700,22 @@ class TestOffloadPlan:
         assert len(backend._resident_blocks) == 2
         assert len(backend._all_hook_groups) == 1
         assert len(backend._all_hook_groups[0]) == 2
+        assert pipeline._dlo_residency_controller is backend
         assert backend.enabled
 
         backend.disable()
 
+        assert pipeline._dlo_residency_controller is None
+
+        for name, parameter in pipeline.named_parameters():
+            torch.testing.assert_close(parameter, expected_parameters[name])
+
+        backend.enable(pipeline)
+        assert len(backend._resident_blocks) == 2
+        backend.disable()
+        for name, parameter in pipeline.named_parameters():
+            torch.testing.assert_close(parameter, expected_parameters[name])
+
 
 class TestMmapValidation:
     """Tests for loader preflight and backend plan realization."""
@@ -2024,8 +2126,8 @@ class TestConfigValidation:
             (2, 1, 1, True, 2),
             (1, 4, 1, True, 4),
             (2, 2, 2, True, 2),
-            (2, 2, 2, False, 1),
-            (1, 4, 2, False, 1),
+            (2, 2, 2, False, 2),
+            (1, 4, 2, False, 4),
         ],
     )
     def test_dlo_group_selection_covers_dp_sp_tp_matrix(
@@ -2199,7 +2301,7 @@ class TestConfigValidation:
             model = "/fake/path"
 
         config = OffloadConfig.from_od_config(FakeODConfig())
-        assert config.dp_size == 1  # forced to 1 when no AllGather
+        assert config.dp_size == 2  # topology is retained; the hook selects group size 1
         assert config.dlo_resident_layers == 20
 
     def test_resident_layers_with_allgather_rejected(self):
@@ -2218,96 +2320,9 @@ class TestConfigValidation:
             parallel_config = FakePC()
             model = "/fake/path"
 
-        with pytest.raises(ValueError, match="requires --dlo-no-use-allgather"):
+        with pytest.raises(ValueError, match="requires the DiT DLO transfer to be rank-local"):
             OffloadConfig.from_od_config(FakeODConfig())
 
-    def test_num_inference_steps_none_rejected(self):
-        """DP multi-concurrency should reject None num_inference_steps."""
-        from types import SimpleNamespace
-
-        # Mock requests with None steps
-        reqs = [
-            SimpleNamespace(
-                req=SimpleNamespace(
-                    request_id=f"req-{i}",
-                    sampling_params=SimpleNamespace(num_inference_steps=None),
-                )
-            )
-            for i in range(2)
-        ]
-
-        # We can't easily instantiate the full executor, but we can test
-        # the validation logic by checking that the code path raises.
-        # The validation is in execute_request, which needs self._ensure_open().
-        # Instead, test the validation logic directly:
-        step_counts = {
-            r.req.sampling_params.num_inference_steps
-            for r in reqs
-            if r.req.sampling_params.num_inference_steps is not None
-        }
-        has_none = any(r.req.sampling_params.num_inference_steps is None for r in reqs)
-        assert has_none, "Test setup: should have None steps"
-        assert len(step_counts) == 0, "Test setup: no explicit steps"
-
-        # The validation condition: (len(step_counts) > 1) or has_none → should reject
-        should_reject = (len(step_counts) > 1) or has_none
-        assert should_reject, "None steps should trigger rejection"
-
-    def test_num_inference_steps_same_explicit_allowed(self):
-        """DP multi-concurrency should allow same explicit num_inference_steps."""
-        from types import SimpleNamespace
-
-        reqs = [
-            SimpleNamespace(
-                req=SimpleNamespace(
-                    request_id=f"req-{i}",
-                    sampling_params=SimpleNamespace(num_inference_steps=35),
-                )
-            )
-            for i in range(4)
-        ]
-
-        step_counts = {
-            r.req.sampling_params.num_inference_steps
-            for r in reqs
-            if r.req.sampling_params.num_inference_steps is not None
-        }
-        has_none = any(r.req.sampling_params.num_inference_steps is None for r in reqs)
-
-        should_reject = (len(step_counts) > 1) or has_none
-        assert not should_reject, "Same explicit steps should be allowed"
-        assert step_counts == {35}
-
-    def test_num_inference_steps_different_explicit_rejected(self):
-        """DP multi-concurrency should reject different explicit steps."""
-        from types import SimpleNamespace
-
-        reqs = [
-            SimpleNamespace(
-                req=SimpleNamespace(
-                    request_id="req-0",
-                    sampling_params=SimpleNamespace(num_inference_steps=35),
-                )
-            ),
-            SimpleNamespace(
-                req=SimpleNamespace(
-                    request_id="req-1",
-                    sampling_params=SimpleNamespace(num_inference_steps=30),
-                )
-            ),
-        ]
-
-        step_counts = {
-            r.req.sampling_params.num_inference_steps
-            for r in reqs
-            if r.req.sampling_params.num_inference_steps is not None
-        }
-        has_none = any(r.req.sampling_params.num_inference_steps is None for r in reqs)
-
-        should_reject = (len(step_counts) > 1) or has_none
-        assert should_reject, "Different steps should trigger rejection"
-        assert len(step_counts) == 2
-
 
 class TestDynamicSlotTracking:
     """Tests for dynamic slot tracking with odd and even block counts.
@@ -2472,3 +2487,344 @@ class TestDynamicSlotTracking:
         # pre_forward should override to slot 1
         hook_b.pre_forward(block_b)
         assert hook_b.current_slot == 1, "pre_forward should read _prev_hook._prefetched_slot=1, not keep initial 0"
+
+
+class _DistributedComponentPipeline(nn.Module):
+    _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"}),
+    )
+
+    def __init__(self):
+        super().__init__()
+        self.transformer = _SingleBlockModel(num_blocks=2)
+        self.text_encoder = _StagedEncoder()
+        self.vae = _StagedVAE()
+
+
+class _GenericDistributedEncoderPipeline(nn.Module):
+    _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=2)
+        self.text_encoder = _PlainEncoder()
+
+
+class TestDistributedComponentSelection:
+    def test_on_demand_only_encoder_cannot_claim_allgather(self, patched_offload_runtime):
+        class Pipeline(nn.Module):
+            _offload_plan = OffloadPlan(
+                encoder_component_types={"text_encoder": "text_encoder"},
+                on_demand_component_paths=frozenset({"text_encoder"}),
+                encoder_dlo_weight_replication=frozenset({"text_encoder"}),
+            )
+
+            def __init__(self):
+                super().__init__()
+                self.text_encoder = _StagedEncoder()
+
+        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"),
+        )
+        backend.dp_group = object()
+
+        with pytest.raises(ValueError, match="cannot use AllGather without.*streamable block plan"):
+            backend.enable(Pipeline())
+
+    def test_explicit_single_block_dit_is_rejected(self, patched_offload_runtime):
+        pipeline = nn.Module()
+        pipeline.transformer = _SingleBlockModel(num_blocks=1)
+        backend = DistributedLayerwiseOffloadBackend(
+            OffloadConfig(
+                strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
+                pin_cpu_memory=False,
+                components=frozenset({"dit"}),
+            ),
+            torch.device("cpu"),
+        )
+
+        with pytest.raises(ValueError, match="leaves only one streaming block"):
+            backend.enable(pipeline)
+
+    def test_enable_defers_collectives_until_first_forward(self, patched_offload_runtime, monkeypatch):
+        pipeline = _GenericDistributedEncoderPipeline()
+        backend = DistributedLayerwiseOffloadBackend(
+            OffloadConfig(
+                strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
+                pin_cpu_memory=False,
+                dp_size=2,
+                components=frozenset({"dit", "text_encoder"}),
+                dlo_transfers={"dit": "allgather", "text_encoder": "rank-local"},
+            ),
+            torch.device("cpu"),
+        )
+        backend.dp_group = object()
+
+        def reject_startup_collective(*_args, **_kwargs):
+            raise AssertionError("enable() entered a rank-asymmetric collective")
+
+        monkeypatch.setattr(torch.distributed, "all_gather_into_tensor", reject_startup_collective)
+
+        backend.enable(pipeline)
+
+        assert all(hook._prefetched_slot is None for group in backend._all_hook_groups for hook in group)
+        encoder_hook = pipeline.text_encoder.encoder.block[0]._hook_registry.get_hook("distributed_layerwise_offload")
+        dit_hook = pipeline.transformer.blocks[0]._hook_registry.get_hook("distributed_layerwise_offload")
+        assert encoder_hook.gpu_shard_buffers == [None, None]
+        assert all(buffer is not None for buffer in dit_hook.gpu_shard_buffers)
+
+        monkeypatch.setattr(
+            torch.distributed,
+            "all_gather_into_tensor",
+            lambda output, local, group: output.copy_(local.repeat(2)),
+        )
+        backend.disable()
+
+    def test_partial_dit_enable_failure_restores_weights_and_hooks(
+        self,
+        patched_offload_runtime,
+        monkeypatch,
+    ):
+        pipeline = _DistributedComponentPipeline()
+        expected_weights = [block.weight.detach().clone() for block in pipeline.transformer.blocks]
+        backend = DistributedLayerwiseOffloadBackend(
+            OffloadConfig(
+                strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
+                pin_cpu_memory=False,
+                dlo_use_allgather=False,
+                components=frozenset({"dit"}),
+            ),
+            torch.device("cpu"),
+        )
+        original_apply = dist_backend_module.apply_distributed_block_hook
+        calls = 0
+
+        def fail_second_hook(*args, **kwargs):
+            nonlocal calls
+            calls += 1
+            if calls == 2:
+                raise RuntimeError("injected DLO hook failure")
+            return original_apply(*args, **kwargs)
+
+        monkeypatch.setattr(dist_backend_module, "apply_distributed_block_hook", fail_second_hook)
+
+        with pytest.raises(RuntimeError, match="injected DLO hook failure"):
+            backend.enable(pipeline)
+
+        assert not backend.enabled
+        for block, expected in zip(pipeline.transformer.blocks, expected_weights, strict=True):
+            registry = getattr(block, "_hook_registry", None)
+            assert registry is None or registry.get_hook("distributed_layerwise_offload") is None
+            torch.testing.assert_close(block.weight, expected)
+
+    def test_multirank_enable_failure_cleanup_skips_restore_collective(self, monkeypatch, mocker):
+        backend = DistributedLayerwiseOffloadBackend(
+            OffloadConfig(
+                strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
+                pin_cpu_memory=False,
+                dp_size=2,
+                dlo_use_allgather=True,
+                components=frozenset({"dit"}),
+            ),
+            torch.device("cpu"),
+        )
+        monkeypatch.setattr(backend, "_enable", mocker.Mock(side_effect=RuntimeError("injected startup failure")))
+        cleanup = mocker.Mock()
+        monkeypatch.setattr(backend, "_disable", cleanup)
+
+        with pytest.raises(RuntimeError, match="injected startup failure"):
+            backend.enable(_DistributedComponentPipeline())
+
+        cleanup.assert_called_once_with(restore_allgather_weights=False)
+
+    def test_multirank_failure_restores_only_rank_local_hooks(self, patched_offload_runtime):
+        backend = DistributedLayerwiseOffloadBackend(
+            OffloadConfig(
+                strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
+                pin_cpu_memory=False,
+                dp_size=2,
+                components=frozenset({"dit", "text_encoder"}),
+                dlo_transfers={"dit": "allgather", "text_encoder": "rank-local"},
+            ),
+            torch.device("cpu"),
+        )
+        allgather_hook = Mock(dp_size=2, next_block=nn.Linear(1, 1))
+        rank_local_hook = Mock(dp_size=1, next_block=nn.Linear(1, 1))
+        backend._all_hook_groups = [[allgather_hook], [rank_local_hook]]
+
+        backend._disable(restore_allgather_weights=False)
+
+        allgather_hook.restore_next_block_to_cpu.assert_not_called()
+        rank_local_hook.restore_next_block_to_cpu.assert_called_once_with()
+
+        with pytest.raises(RuntimeError, match="recreate the backend and reload the pipeline"):
+            backend.enable(nn.Module())
+
+    def test_disable_runs_collectives_before_best_effort_local_cleanup(
+        self,
+        patched_offload_runtime,
+        monkeypatch,
+    ):
+        events: list[str] = []
+        backend = DistributedLayerwiseOffloadBackend(
+            OffloadConfig(
+                strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
+                pin_cpu_memory=False,
+                dp_size=2,
+            ),
+            torch.device("cpu"),
+        )
+        rank_local_hook = Mock(dp_size=1, next_block=nn.Linear(1, 1))
+        allgather_hook = Mock(dp_size=2, next_block=nn.Linear(1, 1))
+        allgather_hook.restore_next_block_to_cpu.side_effect = lambda: events.append("allgather")
+
+        def fail_rank_local_restore():
+            events.append("rank-local")
+            raise RuntimeError("injected rank-local restore failure")
+
+        rank_local_hook.restore_next_block_to_cpu.side_effect = fail_rank_local_restore
+        backend._all_hook_groups = [[rank_local_hook], [allgather_hook]]
+        backend._blocks = [[nn.Module()]]
+        backend.enabled = True
+        monkeypatch.setattr(
+            dist_backend_module,
+            "remove_distributed_block_hook",
+            lambda _block: events.append("remove"),
+        )
+        monkeypatch.setattr(backend, "_release_registered_mmap", lambda: events.append("unregister"))
+        release_handles = Mock()
+        monkeypatch.setattr(backend, "_release_mmap_handles", release_handles)
+
+        with pytest.raises(RuntimeError, match="injected rank-local restore failure"):
+            backend.disable()
+
+        assert events == ["allgather", "rank-local", "remove", "unregister"]
+        release_handles.assert_not_called()
+        assert backend.enabled
+        assert backend._all_hook_groups
+
+    def test_all_streams_encoder_and_keeps_vae_resident(self, patched_offload_runtime, monkeypatch):
+        pipeline = _DistributedComponentPipeline()
+        move_non_block_state = Mock()
+        monkeypatch.setattr(dist_backend_module, "move_non_block_state_to_device", move_non_block_state)
+        backend = DistributedLayerwiseOffloadBackend(
+            OffloadConfig(
+                strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
+                pin_cpu_memory=False,
+                dlo_use_allgather=False,
+                components=frozenset({"dit", "text_encoder"}),
+            ),
+            torch.device("cpu"),
+        )
+
+        backend.enable(pipeline)
+
+        assert pipeline.text_encoder._omni_layerwise_enabled
+        assert len(pipeline.text_encoder._omni_layerwise_hooks) == 4
+        # The pipeline-owned stage lifecycle releases the encoder's non-block
+        # state after DLO has installed blockwise hooks.
+        assert pipeline.text_encoder.offload_calls == 1
+        assert pipeline.vae.offload_calls == 0
+        assert pipeline.vae.to_calls == 1
+        move_non_block_state.assert_not_called()
+
+        backend.disable()
+
+        assert not pipeline.text_encoder._omni_layerwise_enabled
+        encoder_blocks = [
+            *pipeline.text_encoder.vision.blocks,
+            *pipeline.text_encoder.text_model.layers,
+        ]
+        assert all(block._hook_registry.get_hook("distributed_layerwise_offload") is None for block in encoder_blocks)
+
+    def test_dit_and_encoder_choose_transfer_independently(self, patched_offload_runtime, monkeypatch):
+        pipeline = _GenericDistributedEncoderPipeline()
+        expected_parameters = {name: parameter.detach().clone() for name, parameter in pipeline.named_parameters()}
+        full_weights_by_rank_zero_shard = {
+            tuple(parameter.flatten()[: (parameter.numel() + 1) // 2].tolist()): parameter.flatten()
+            for parameter in expected_parameters.values()
+            if parameter.numel() > 4
+        }
+        backend = DistributedLayerwiseOffloadBackend(
+            OffloadConfig(
+                strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
+                pin_cpu_memory=False,
+                dp_size=2,
+                components=frozenset({"dit", "text_encoder"}),
+                dlo_transfers={"dit": "rank-local", "text_encoder": "allgather"},
+            ),
+            torch.device("cpu"),
+        )
+        backend.dp_group = object()
+
+        def fake_allgather(output, local, group):
+            del group
+            output.copy_(full_weights_by_rank_zero_shard[tuple(local.tolist())])
+
+        monkeypatch.setattr(
+            torch.distributed,
+            "all_gather_into_tensor",
+            fake_allgather,
+        )
+
+        backend.enable(pipeline)
+
+        encoder_hook = pipeline.text_encoder.encoder.block[0]._hook_registry.get_hook("distributed_layerwise_offload")
+        dit_hook = pipeline.transformer.blocks[0]._hook_registry.get_hook("distributed_layerwise_offload")
+        assert encoder_hook.dp_size == 2
+        assert dit_hook.dp_size == 1
+        assert all(buffer is not None for buffer in encoder_hook.gpu_shard_buffers)
+        assert dit_hook.gpu_shard_buffers == [None, None]
+
+        backend.disable()
+        actual_parameters = dict(pipeline.named_parameters())
+        for name, expected in expected_parameters.items():
+            torch.testing.assert_close(actual_parameters[name], expected)
+
+    def test_encoder_allgather_rejects_undeclared_replication(self, patched_offload_runtime):
+        pipeline = _DistributedComponentPipeline()
+        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"),
+        )
+        backend.dp_group = object()
+
+        with pytest.raises(ValueError, match="not declared replicated"):
+            backend.enable(pipeline)
+
+    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"),
+        )
+        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)

tests/diffusion/offloader/test_layerwise_backend.py

@@ -3,66 +3,39 @@
 
 """Unit tests for LayerwiseOffloadHook and LayerWiseOffloadBackend utilities."""
 
-import gc
-from contextlib import contextmanager
+from multiprocessing.reduction import ForkingPickler
+from types import SimpleNamespace
+from unittest.mock import Mock
 
 import pytest
 import torch
-import torch.distributed as dist
 from torch import nn
 from torch.distributed.tensor import DeviceMesh, DTensor, Replicate
 
 import vllm_omni.diffusion.offloader.layerwise_backend as layerwise_backend_module
-from tests.helpers.runtime import get_distributed_init_method
-from vllm_omni.diffusion.offloader.layerwise_backend import LayerWiseOffloadBackend, LayerwiseOffloadHook
-from vllm_omni.platforms import current_omni_platform
+import vllm_omni.diffusion.offloader.tensor_utils as tensor_utils_module
+from tests.diffusion.offloader.helpers import (
+    DummyStream,
+    _DummyBlock,
+    _PlainEncoder,
+    _SingleBlockModel,
+    _StagedEncoder,
+    _StagedVAE,
+    patch_offload_runtime,
+)
+from vllm_omni.diffusion.offloader.base import OffloadConfig, OffloadStrategy
+from vllm_omni.diffusion.offloader.layerwise_backend import (
+    LayerWiseOffloadBackend,
+    LayerwiseOffloadHook,
+)
+from vllm_omni.diffusion.offloader.offload_plan import OffloadPlan
 
 pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
 
 
-class DummyStream:
-    def wait_stream(self, _stream) -> None:
-        return None
-
-    def wait_event(self, _event) -> None:
-        return None
-
-
-class DummyEvent:
-    def record(self, _stream) -> None:
-        return None
-
-
-@contextmanager
-def dummy_stream(_stream):
-    yield None
-
-
-def _cleanup_distributed() -> None:
-    if dist.is_initialized():
-        dist.destroy_process_group()
-
-    gc.collect()
-    if current_omni_platform.is_available():
-        current_omni_platform.empty_cache()
-        current_omni_platform.synchronize()
-
-
-@pytest.fixture(scope="module")
-def dist_group():
-    dist.init_process_group("gloo", rank=0, world_size=1, init_method=get_distributed_init_method())
-    try:
-        yield
-    finally:
-        _cleanup_distributed()
-
-
 @pytest.fixture
 def patched_offload_runtime(monkeypatch):
-    monkeypatch.setattr(layerwise_backend_module.current_omni_platform, "Stream", DummyStream)
-    monkeypatch.setattr(layerwise_backend_module.current_omni_platform, "Event", DummyEvent)
-    monkeypatch.setattr(layerwise_backend_module.current_omni_platform, "current_stream", lambda: DummyStream())
-    monkeypatch.setattr(layerwise_backend_module.current_omni_platform, "stream", dummy_stream)
+    patch_offload_runtime(monkeypatch, layerwise_backend_module.current_omni_platform)
 
 
 class TinyBlock(nn.Module):
@@ -77,7 +50,66 @@ def _make_values(start: float) -> torch.Tensor:
     return torch.arange(start, start + 4, dtype=torch.float32)
 
 
+def test_clear_tensor_storage_rolls_back_partial_commit(monkeypatch):
+    first = nn.Parameter(torch.tensor([1.0, 2.0]))
+    second = nn.Parameter(torch.tensor([3.0, 4.0]))
+    expected_first = first.detach().clone()
+    expected_second = second.detach().clone()
+    second_data_ptr = second.data_ptr()
+    original_set_storage = tensor_utils_module.set_tensor_storage
+    calls = 0
+
+    def fail_second_commit(target, value):
+        nonlocal calls
+        calls += 1
+        if calls == 2:
+            raise RuntimeError("injected storage clear failure")
+        original_set_storage(target, value)
+
+    monkeypatch.setattr(tensor_utils_module, "set_tensor_storage", fail_second_commit)
+
+    with pytest.raises(RuntimeError, match="injected storage clear failure"):
+        tensor_utils_module.clear_tensor_storage([first, second])
+
+    assert calls == 3  # first commit, failed second commit, first rollback
+    torch.testing.assert_close(first, expected_first)
+    torch.testing.assert_close(second, expected_second)
+    assert second.data_ptr() == second_data_ptr
+
+
 class TestLayerwiseOffloadHook:
+    def test_buffer_only_block_reports_offloaded_state(self, patched_offload_runtime):
+        current_block = nn.Module()
+        current_block.register_buffer("state", torch.ones(2))
+        next_block = nn.Module()
+        next_block.register_buffer("state", torch.ones(2))
+        hook = LayerwiseOffloadHook(
+            next_block=next_block,
+            device=torch.device("cpu"),
+            stream=DummyStream(),
+            pin_memory=False,
+        )
+        hook.initialize_hook(current_block)
+
+        assert hook.is_materialized
+        hook.offload_layer()
+        assert not hook.is_materialized
+
+    def test_ring_probes_are_captured_before_any_block_is_cleared(self, patched_offload_runtime):
+        blocks = nn.ModuleList([nn.Linear(2, 2, bias=False) for _ in range(3)])
+        expected_middle = blocks[1].weight.detach().clone()
+
+        hooks = layerwise_backend_module._install_layerwise_hook_group(
+            blocks,
+            torch.device("cpu"),
+            DummyStream(),
+            pin_memory=False,
+        )
+
+        assert all(not hook.is_materialized for hook in hooks)
+        hooks[2].pre_forward(blocks[1])
+        torch.testing.assert_close(blocks[1].weight, expected_middle)
+
     def test_dtensor_wrapper_is_preserved_across_prefetch_and_offload(self, dist_group, patched_offload_runtime):
         current_block = TinyBlock(_make_values(1.0))
         next_block = TinyBlock(_make_values(10.0))
@@ -135,20 +167,6 @@ class TestLayerwiseOffloadHook:
         assert torch.equal(next_block.weight, expected)
 
 
-class _DummyBlock(nn.Module):
-    def __init__(self):
-        super().__init__()
-        self.weight = nn.Parameter(torch.randn(10, 10))
-
-
-class _SingleBlockModel(nn.Module):
-    _layerwise_offload_blocks_attrs = ["blocks"]
-
-    def __init__(self, num_blocks: int = 3):
-        super().__init__()
-        self.blocks = nn.ModuleList([_DummyBlock() for _ in range(num_blocks)])
-
-
 class _MultiBlockModel(nn.Module):
     _layerwise_offload_blocks_attrs = ["transformer_blocks", "single_transformer_blocks"]
 
@@ -246,3 +264,732 @@ class TestGetBlocksAttrNames:
         LayerWiseOffloadBackend.set_blocks_attr_names(model, ["new_blocks"])
         assert hasattr(model.__class__, "_layerwise_offload_blocks_attrs")
         assert model.__class__._layerwise_offload_blocks_attrs == ["new_blocks"]
+
+
+class _ComponentPipeline(nn.Module):
+    _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"}),
+    )
+
+    def __init__(self):
+        super().__init__()
+        self.transformer = _SingleBlockModel()
+        self.text_encoder = _StagedEncoder()
+        self.vae = _StagedVAE()
+
+
+class _LegacyComponentPipeline(nn.Module):
+    """Pipeline with DiT block metadata but no auxiliary OffloadPlan."""
+
+    def __init__(self):
+        super().__init__()
+        self.transformer = _SingleBlockModel()
+        self.text_encoder = _StagedEncoder()
+        self.vae = _StagedVAE()
+
+
+class _LegacyEncoderOnlyPipeline(nn.Module):
+    _offload_plan = _ComponentPipeline._offload_plan
+
+    def __init__(self):
+        super().__init__()
+        self.text_encoder = _StagedEncoder()
+        self.vae = _StagedVAE()
+
+
+class _PlannedDitPipeline(nn.Module):
+    _offload_plan = OffloadPlan(block_attrs={"transformer": ("blocks",)})
+
+    def __init__(self):
+        super().__init__()
+        self.transformer = _NoAttrsModel()
+
+
+class _GenericEncoderPipeline(nn.Module):
+    _offload_plan = OffloadPlan(
+        encoder_component_types={"text_encoder": "text_encoder"},
+        encoder_block_attrs={"text_encoder": ("encoder.block",)},
+    )
+
+    def __init__(self):
+        super().__init__()
+        self.transformer = _SingleBlockModel()
+        self.text_encoder = _PlainEncoder()
+
+
+def _layer_backend(components: set[str] | None = None) -> LayerWiseOffloadBackend:
+    options = {}
+    if components is not None:
+        options = {"components": frozenset(components)}
+    return LayerWiseOffloadBackend(
+        OffloadConfig(
+            strategy=OffloadStrategy.LAYER_WISE,
+            pin_cpu_memory=False,
+            **options,
+        ),
+        torch.device("cpu"),
+    )
+
+
+class TestLayerwiseComponentSelection:
+    def test_plan_block_attrs_drive_dit_discovery(self, patched_offload_runtime):
+        pipeline = _PlannedDitPipeline()
+        backend = _layer_backend({"dit"})
+
+        backend.enable(pipeline)
+
+        assert hasattr(pipeline.transformer.blocks[0], "_hook_registry")
+        backend.disable()
+
+    def test_legacy_missing_dit_preserves_noop(self, patched_offload_runtime):
+        pipeline = _LegacyEncoderOnlyPipeline()
+        backend = _layer_backend()
+
+        backend.enable(pipeline)
+
+        assert not backend.enabled
+        assert pipeline.text_encoder.offload_calls == 0
+        assert pipeline.vae.offload_calls == 0
+
+    def test_encoder_only_streams_planned_blocks(self, patched_offload_runtime, monkeypatch):
+        pipeline = _ComponentPipeline()
+        backend = _layer_backend({"text_encoder"})
+        move_non_block_state = Mock()
+        monkeypatch.setattr(layerwise_backend_module, "move_non_block_state_to_device", move_non_block_state)
+
+        backend.enable(pipeline)
+
+        assert pipeline.text_encoder._omni_layerwise_enabled
+        assert len(pipeline.text_encoder._omni_layerwise_hooks) == 4
+        assert pipeline.text_encoder.offload_calls == 1
+        assert pipeline.vae.offload_calls == 0
+        assert pipeline.vae.to_calls == 1
+        assert not hasattr(pipeline.transformer.blocks[0], "_hook_registry")
+        assert backend.enabled
+        move_non_block_state.assert_not_called()
+
+        backend.disable()
+        assert not pipeline.text_encoder._omni_layerwise_enabled
+
+    def test_single_gpu_dit_only_keeps_encoder_and_vae_resident(self, patched_offload_runtime):
+        pipeline = _ComponentPipeline()
+        expected_weights = [block.weight.detach().clone() for block in pipeline.transformer.blocks]
+        backend = _layer_backend({"dit"})
+
+        backend.enable(pipeline)
+
+        assert not hasattr(pipeline.text_encoder, "_omni_layerwise_enabled")
+        assert pipeline.text_encoder.to_calls == 1
+        assert pipeline.text_encoder.offload_calls == 0
+        assert pipeline.vae.to_calls == 1
+        assert pipeline.vae.offload_calls == 0
+        assert hasattr(pipeline.transformer.blocks[0], "_hook_registry")
+
+        backend.disable()
+
+        for block, expected in zip(pipeline.transformer.blocks, expected_weights, strict=True):
+            torch.testing.assert_close(block.weight, expected)
+
+        backend.enable(pipeline)
+        backend.disable()
+        for block, expected in zip(pipeline.transformer.blocks, expected_weights, strict=True):
+            torch.testing.assert_close(block.weight, expected)
+
+    def test_partial_dit_enable_failure_restores_weights_and_hooks(self, patched_offload_runtime, monkeypatch):
+        pipeline = _ComponentPipeline()
+        expected_weights = [block.weight.detach().clone() for block in pipeline.transformer.blocks]
+        backend = _layer_backend({"dit"})
+        original_apply = layerwise_backend_module.apply_block_hook
+        calls = 0
+
+        def fail_second_hook(*args, **kwargs):
+            nonlocal calls
+            calls += 1
+            if calls == 2:
+                raise RuntimeError("injected hook failure")
+            return original_apply(*args, **kwargs)
+
+        monkeypatch.setattr(layerwise_backend_module, "apply_block_hook", fail_second_hook)
+
+        with pytest.raises(RuntimeError, match="injected hook failure"):
+            backend.enable(pipeline)
+
+        assert not backend.enabled
+        for block, expected in zip(pipeline.transformer.blocks, expected_weights, strict=True):
+            registry = getattr(block, "_hook_registry", None)
+            assert registry is None or registry.get_hook("layerwise_offload") is None
+            torch.testing.assert_close(block.weight, expected)
+
+    def test_later_component_failure_removes_installed_encoder_hooks(
+        self,
+        patched_offload_runtime,
+        monkeypatch,
+    ):
+        pipeline = _ComponentPipeline()
+        blocks = [*pipeline.text_encoder.vision.blocks, *pipeline.text_encoder.text_model.layers]
+        expected = [block.weight.detach().clone() for block in blocks]
+        backend = _layer_backend({"text_encoder"})
+
+        def fail_vae_placement(*_args, **_kwargs):
+            raise RuntimeError("injected VAE placement failure")
+
+        monkeypatch.setattr(pipeline.vae, "to", fail_vae_placement)
+
+        with pytest.raises(RuntimeError, match="injected VAE placement failure"):
+            backend.enable(pipeline)
+
+        assert not backend.enabled
+        assert not pipeline.text_encoder._omni_layerwise_enabled
+        assert not backend._encoder_modules
+        for block, original in zip(blocks, expected, strict=True):
+            assert block._hook_registry.get_hook("layerwise_offload") is None
+            torch.testing.assert_close(block.weight, original)
+
+    def test_encoder_only_requires_streamable_offload_plan(self, patched_offload_runtime):
+        pipeline = nn.Module()
+        pipeline.transformer = _SingleBlockModel()
+        pipeline.text_encoder = _StagedEncoder()
+        backend = _layer_backend({"text_encoder"})
+
+        with pytest.raises(ValueError, match="Selected text encoder 'text_encoder' requires"):
+            backend.enable(pipeline)
+
+    def test_explicit_text_encoder_selection_requires_a_matching_module(self, patched_offload_runtime):
+        pipeline = nn.Module()
+        pipeline.transformer = _SingleBlockModel(num_blocks=2)
+        backend = _layer_backend({"dit", "text_encoder"})
+
+        with pytest.raises(ValueError, match="No text encoder modules found"):
+            backend.enable(pipeline)
+
+        assert not hasattr(pipeline.transformer.blocks[0], "_hook_registry")
+
+    def test_every_selected_encoder_requires_its_own_plan(self, patched_offload_runtime):
+        class Pipeline(nn.Module):
+            _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()
+                self.text_encoder = _PlainEncoder()
+                self.text_encoder_2 = _PlainEncoder()
+
+        pipeline = Pipeline()
+        backend = _layer_backend({"text_encoder"})
+
+        with pytest.raises(ValueError, match="Selected text encoder 'text_encoder_2' requires"):
+            backend.enable(pipeline)
+
+        assert not pipeline.text_encoder._omni_layerwise_enabled
+
+    def test_default_selection_preserves_unplanned_auxiliaries(self, patched_offload_runtime):
+        pipeline = _LegacyComponentPipeline()
+        backend = _layer_backend()
+
+        backend.enable(pipeline)
+
+        assert hasattr(pipeline.transformer.blocks[0], "_hook_registry")
+        assert not hasattr(pipeline.text_encoder, "_omni_layerwise_enabled")
+        assert pipeline.text_encoder.to_calls == 1
+        assert pipeline.text_encoder.offload_calls == 0
+        assert pipeline.vae.to_calls == 1
+        assert pipeline.vae.offload_calls == 0
+
+        backend.disable()
+
+    def test_disable_failure_keeps_host_masters_for_retry(self, patched_offload_runtime, monkeypatch):
+        pipeline = nn.Module()
+        pipeline.transformer = _SingleBlockModel(num_blocks=3)
+        expected = {name: tensor.detach().clone() for name, tensor in pipeline.state_dict().items()}
+        backend = _layer_backend()
+        backend.enable(pipeline)
+        failing_hook = backend._dit_hooks[0]
+        restore = failing_hook.restore_next_block
+        attempts = 0
+
+        def fail_once():
+            nonlocal attempts
+            attempts += 1
+            if attempts == 1:
+                raise RuntimeError("injected restore failure")
+            restore()
+
+        monkeypatch.setattr(failing_hook, "restore_next_block", fail_once)
+
+        with pytest.raises(RuntimeError, match="Failed to fully disable"):
+            backend.disable()
+
+        assert backend.enabled
+        assert backend._dit_hooks
+
+        backend.disable()
+
+        assert not backend.enabled
+        for name, tensor in pipeline.state_dict().items():
+            torch.testing.assert_close(tensor, expected[name])
+
+    def test_encoder_disable_failure_keeps_host_masters_for_retry(
+        self,
+        patched_offload_runtime,
+        monkeypatch,
+    ):
+        pipeline = _ComponentPipeline()
+        encoder = pipeline.text_encoder
+        expected = {name: tensor.detach().clone() for name, tensor in encoder.state_dict().items()}
+        backend = _layer_backend({"text_encoder"})
+        backend.enable(pipeline)
+        failing_hook = encoder._omni_layerwise_hooks[0]
+        restore = failing_hook.restore_next_block
+        attempts = 0
+
+        def fail_once():
+            nonlocal attempts
+            attempts += 1
+            if attempts == 1:
+                raise RuntimeError("injected encoder restore failure")
+            restore()
+
+        monkeypatch.setattr(failing_hook, "restore_next_block", fail_once)
+
+        with pytest.raises(RuntimeError, match="Failed to fully disable"):
+            backend.disable()
+
+        assert backend.enabled
+        assert backend._encoder_modules == [encoder]
+        assert encoder._omni_layerwise_enabled
+        assert encoder._omni_layerwise_hooks
+
+        backend.disable()
+
+        assert not backend.enabled
+        assert not encoder._omni_layerwise_enabled
+        for name, tensor in encoder.state_dict().items():
+            torch.testing.assert_close(tensor, expected[name])
+
+    def test_legacy_selector_preserves_planned_encoder_and_vae_lifecycle(self, patched_offload_runtime):
+        pipeline = _ComponentPipeline()
+        backend = _layer_backend()
+
+        backend.enable(pipeline)
+
+        assert pipeline.text_encoder._omni_layerwise_enabled
+        assert pipeline.text_encoder.offload_calls == 1
+        assert pipeline.vae.offload_calls == 1
+
+        backend.disable()
+
+    def test_standard_encoder_needs_only_declared_block_paths(self, patched_offload_runtime):
+        pipeline = _GenericEncoderPipeline()
+        backend = _layer_backend({"text_encoder"})
+
+        backend.enable(pipeline)
+
+        blocks = pipeline.text_encoder.encoder.block
+        assert pipeline.text_encoder._omni_layerwise_enabled
+        assert all(block._hook_registry.get_hook("layerwise_offload") is not None for block in blocks)
+        assert pipeline.text_encoder.final_norm.weight.numel() == 4
+
+        backend.disable()
+
+        assert not pipeline.text_encoder._omni_layerwise_enabled
+        assert all(block._hook_registry.get_hook("layerwise_offload") is None for block in blocks)
+        assert all(block.weight.numel() == 100 for block in blocks)
+
+
+def _offload_od_config(**overrides):
+    values = {
+        "diffusion_offload_config": None,
+        "enable_cpu_offload": False,
+        "enable_layerwise_offload": False,
+        "enable_distributed_layerwise_offload": False,
+        "dlo_use_allgather": True,
+        "dlo_resident_layers": 0,
+        "pin_cpu_memory": True,
+        "parallel_config": None,
+        "model": "/fake/model",
+    }
+    values.update(overrides)
+    return SimpleNamespace(**values)
+
+
+def _resolve_offload_config(public=None, **overrides) -> OffloadConfig:
+    return OffloadConfig.from_od_config(_offload_od_config(diffusion_offload_config=public, **overrides))
+
+
+class TestLayerwiseComponentConfig:
+    def test_layer_mode_selects_components_without_backend_jargon(self):
+        config = _resolve_offload_config(
+            {
+                "mode": "layer",
+                "components": ["dit", "text_encoder"],
+            }
+        )
+
+        assert config.strategy is OffloadStrategy.LAYER_WISE
+        assert config.components == frozenset({"dit", "text_encoder"})
+        assert not config.uses_allgather("dit")
+        assert not config.uses_allgather("text_encoder")
+        plan = OffloadPlan(encoder_component_types={"mllm": "text_encoder"})
+        assert config.offloads_encoder("mllm", plan)
+
+    @pytest.mark.parametrize("component", ["image_encoder", "vae", "scheduler", "text-encoder"])
+    def test_unknown_or_noncanonical_component_is_rejected(self, component):
+        with pytest.raises(ValueError, match="Unknown diffusion offload component"):
+            OffloadConfig.from_od_config(
+                _offload_od_config(diffusion_offload_config={"mode": "layer", "components": [component]})
+            )
+
+    def test_components_is_selection_only_not_an_options_mapping(self):
+        with pytest.raises(TypeError, match="components must be a non-empty list"):
+            OffloadConfig.from_od_config(
+                _offload_od_config(
+                    diffusion_offload_config={
+                        "mode": "layer",
+                        "components": {"dit": {}},
+                    }
+                )
+            )
+
+    def test_layer_options_requires_component_name_keys(self):
+        with pytest.raises(TypeError, match="layer_options keys must be strings"):
+            OffloadConfig.from_od_config(
+                _offload_od_config(
+                    diffusion_offload_config={
+                        "mode": "layer",
+                        "components": ["dit"],
+                        "layer_options": {0: {}},
+                    }
+                )
+            )
+
+    @pytest.mark.parametrize(
+        ("config", "message"),
+        [
+            ({0: "layer"}, "diffusion_offload_config keys must be strings"),
+            (
+                {
+                    "mode": "layer",
+                    "components": ["dit"],
+                    "layer_options": {"dit": {0: "allgather"}},
+                },
+                "layer_options\\['dit'\\] keys must be strings",
+            ),
+        ],
+    )
+    def test_all_public_mapping_levels_require_string_keys(self, config, message):
+        with pytest.raises(TypeError, match=message):
+            OffloadConfig.from_od_config(_offload_od_config(diffusion_offload_config=config))
+
+    @pytest.mark.parametrize(
+        "config,match",
+        [
+            ({"components": ["dit"]}, "requires 'mode'"),
+            ({"mode": "layer"}, "requires 'components'"),
+            ({"mode": "layer", "components": []}, "must not be empty"),
+            ({"mode": "layer", "components": [" dit "]}, "Unknown diffusion offload component"),
+            ({"mode": "layerwise", "components": ["dit"]}, "Unknown diffusion offload mode"),
+            (
+                {
+                    "mode": "layer",
+                    "components": ["dit"],
+                    "layer_options": {"dit": {"weight_transfer": "rank_local"}},
+                },
+                "Unknown offload transfer",
+            ),
+            (
+                {
+                    "mode": "layer",
+                    "components": ["dit"],
+                    "layer_options": {"dit": {"prefetch": 2}},
+                },
+                "Unknown diffusion offload setting",
+            ),
+            (
+                {
+                    "mode": "layer",
+                    "components": ["text_encoder"],
+                    "layer_options": {"text_encoder": {"resident_layers": 1}},
+                },
+                "supports only the 'dit'",
+            ),
+            (
+                {
+                    "mode": "layer",
+                    "components": ["dit"],
+                    "layer_options": {"text_encoder": {}},
+                },
+                "requires selecting the same component",
+            ),
+        ],
+    )
+    def test_schema_rejects_ambiguous_or_unsupported_values(self, config, match):
+        with pytest.raises(ValueError, match=match):
+            OffloadConfig.from_od_config(_offload_od_config(diffusion_offload_config=config))
+
+    def test_module_mode_accepts_component_selection(self):
+        config = _resolve_offload_config(
+            {
+                "mode": "module",
+                "components": ["dit", "text_encoder"],
+            }
+        )
+
+        assert config.strategy is OffloadStrategy.MODEL_LEVEL
+        assert config.components == frozenset({"dit", "text_encoder"})
+
+    @pytest.mark.parametrize("setting", [{"weight_transfer": "rank-local"}, {"resident_layers": 1}])
+    def test_module_mode_rejects_layer_settings(self, setting):
+        with pytest.raises(ValueError, match="layer_options requires mode='layer'"):
+            OffloadConfig.from_od_config(
+                _offload_od_config(
+                    diffusion_offload_config={
+                        "mode": "module",
+                        "components": ["dit"],
+                        "layer_options": {"dit": setting},
+                    }
+                )
+            )
+
+    def test_each_component_can_choose_its_transfer(self):
+        config = _resolve_offload_config(
+            {
+                "mode": "layer",
+                "components": ["dit", "text_encoder"],
+                "layer_options": {
+                    "dit": {"weight_transfer": "rank-local"},
+                    "text_encoder": {"weight_transfer": "allgather"},
+                },
+            }
+        )
+
+        assert config.strategy is OffloadStrategy.DISTRIBUTED_LAYER_WISE
+        assert not config.uses_allgather("dit")
+        assert config.uses_allgather("text_encoder")
+        assert config.dlo_use_allgather is False
+
+    def test_dit_residency_selects_capable_backend(self):
+        config = _resolve_offload_config(
+            {
+                "mode": "layer",
+                "components": ["dit"],
+                "layer_options": {"dit": {"resident_layers": 20}},
+            }
+        )
+
+        assert config.strategy is OffloadStrategy.DISTRIBUTED_LAYER_WISE
+        assert config.dlo_resident_layers == 20
+        assert not config.uses_allgather("dit")
+
+    def test_dit_residency_rejects_allgather(self):
+        with pytest.raises(ValueError, match="resident_layers requires dit.weight_transfer='rank-local'"):
+            OffloadConfig.from_od_config(
+                _offload_od_config(
+                    diffusion_offload_config={
+                        "mode": "layer",
+                        "components": ["dit"],
+                        "layer_options": {"dit": {"weight_transfer": "allgather", "resident_layers": 1}},
+                    }
+                )
+            )
+
+    def test_host_weight_runtime_remains_separate_from_compact_config(self):
+        with pytest.raises(ValueError, match="cannot be combined with Host Weight Runtime"):
+            OffloadConfig.from_od_config(
+                _offload_od_config(
+                    diffusion_offload_config={"mode": "layer", "components": ["dit"]},
+                    host_weight_runtime_mode="preferred",
+                )
+            )
+
+    def test_pin_memory_override_is_scoped_to_offload_config(self):
+        config = _resolve_offload_config(
+            {
+                "mode": "layer",
+                "components": ["dit"],
+                "pin_memory": False,
+            },
+            pin_cpu_memory=True,
+        )
+
+        assert config.pin_cpu_memory is False
+
+    def test_legacy_omitted_selector_preserves_dit_only_behavior(self):
+        config = _resolve_offload_config(enable_layerwise_offload=True)
+
+        assert config.strategy is OffloadStrategy.LAYER_WISE
+        assert config.components is None
+
+    def test_offload_resolution_is_cached_for_runtime_accessors(self, monkeypatch):
+        import vllm_omni.diffusion.offloader.config as config_module
+
+        od_config = _offload_od_config(
+            diffusion_offload_config={
+                "mode": "layer",
+                "components": ["dit", "text_encoder"],
+            }
+        )
+        parse_calls = 0
+        original_parse = config_module.parse_diffusion_offload_config
+
+        def track_parse(value):
+            nonlocal parse_calls
+            parse_calls += 1
+            return original_parse(value)
+
+        monkeypatch.setattr(config_module, "parse_diffusion_offload_config", track_parse)
+
+        assert config_module.resolve_offload_strategy(od_config) is OffloadStrategy.LAYER_WISE
+        assert config_module.selected_offload_components(od_config) == {"dit", "text_encoder"}
+        assert not config_module.component_uses_allgather(od_config, "dit")
+        assert not config_module.any_selected_component_uses_allgather(od_config)
+        assert parse_calls == 1
+
+    def test_cached_policy_is_deeply_immutable_and_pickle_safe(self):
+        import vllm_omni.diffusion.offloader.config as config_module
+
+        od_config = _offload_od_config(
+            diffusion_offload_config={
+                "mode": "layer",
+                "components": ["dit"],
+                "layer_options": {"dit": {"weight_transfer": "rank-local"}},
+            }
+        )
+        resolved = config_module.resolve_offload(od_config)
+
+        with pytest.raises(TypeError):
+            resolved.transfers["dit"] = config_module.DLOTransfer.ALLGATHER  # type: ignore[index]
+        assert resolved.public is not None
+        with pytest.raises(TypeError):
+            resolved.public.layer_options["dit"] = config_module.LayerOffloadOptions()  # type: ignore[index]
+
+        restored = ForkingPickler.loads(ForkingPickler.dumps(resolved))
+        assert not restored.uses_allgather("dit")
+
+    def test_dp_text_encoder_allgather_rejects_rank_local_prompt_cache(self):
+        config = _offload_od_config(
+            diffusion_offload_config={
+                "mode": "layer",
+                "components": ["text_encoder"],
+                "layer_options": {"text_encoder": {"weight_transfer": "allgather"}},
+            },
+            enable_prompt_embed_cache=True,
+            parallel_config=SimpleNamespace(
+                data_parallel_size=2,
+                sequence_parallel_size=1,
+                use_hsdp=False,
+            ),
+        )
+
+        with pytest.raises(ValueError, match="rank-local cache hits would skip different encoder collectives"):
+            OffloadConfig.from_od_config(config)
+
+    def test_dp_allgather_rejects_rank_local_dit_cache_decisions(self):
+        config = _offload_od_config(
+            diffusion_offload_config={
+                "mode": "layer",
+                "components": ["dit"],
+                "layer_options": {"dit": {"weight_transfer": "allgather"}},
+            },
+            cache_backend="tea_cache",
+            parallel_config=SimpleNamespace(
+                data_parallel_size=2,
+                sequence_parallel_size=1,
+                use_hsdp=False,
+            ),
+        )
+
+        with pytest.raises(ValueError, match="rank-local cache decisions can skip different weight collectives"):
+            OffloadConfig.from_od_config(config)
+
+    def test_internal_transfer_map_requires_every_component(self):
+        with pytest.raises(ValueError, match="missing: text_encoder"):
+            OffloadConfig(
+                strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
+                dlo_transfers={"dit": "rank-local"},
+            )
+
+    def test_legacy_distributed_settings_still_work(self):
+        config = _resolve_offload_config(
+            enable_distributed_layerwise_offload=True,
+            dlo_use_allgather=False,
+            dlo_resident_layers=3,
+        )
+
+        assert config.strategy is OffloadStrategy.DISTRIBUTED_LAYER_WISE
+        assert config.dlo_resident_layers == 3
+        assert not config.uses_allgather("dit")
+        assert not config.uses_allgather("text_encoder")
+
+    def test_legacy_allgather_remains_dit_only(self):
+        config = _resolve_offload_config(
+            enable_distributed_layerwise_offload=True,
+            dlo_use_allgather=True,
+        )
+
+        assert config.uses_allgather("dit")
+        assert not config.uses_allgather("text_encoder")
+
+    @pytest.mark.parametrize(
+        ("legacy_flags", "expected_strategy"),
+        [
+            (
+                {"enable_cpu_offload": True, "enable_layerwise_offload": True},
+                OffloadStrategy.LAYER_WISE,
+            ),
+            (
+                {
+                    "enable_cpu_offload": True,
+                    "enable_layerwise_offload": True,
+                    "enable_distributed_layerwise_offload": True,
+                },
+                OffloadStrategy.DISTRIBUTED_LAYER_WISE,
+            ),
+        ],
+    )
+    def test_legacy_aliases_preserve_existing_priority(self, legacy_flags, expected_strategy):
+        config = _resolve_offload_config(**legacy_flags)
+
+        assert config.strategy is expected_strategy
+
+    def test_compact_config_cannot_mix_with_legacy_alias(self):
+        with pytest.raises(ValueError, match="cannot be combined"):
+            OffloadConfig.from_od_config(
+                _offload_od_config(
+                    diffusion_offload_config={"mode": "layer", "components": ["dit"]},
+                    enable_cpu_offload=True,
+                )
+            )
+
+    @pytest.mark.parametrize(
+        "legacy_options",
+        [
+            {"dlo_use_allgather": False},
+            {"dlo_resident_layers": 12},
+        ],
+    )
+    def test_compact_config_cannot_mix_with_legacy_dlo_options(self, legacy_options):
+        with pytest.raises(ValueError, match="cannot be combined with legacy DLO option"):
+            OffloadConfig.from_od_config(
+                _offload_od_config(
+                    diffusion_offload_config={"mode": "layer", "components": ["dit"]},
+                    **legacy_options,
+                )
+            )
+
+    def test_legacy_allgather_resident_layers_fail_during_config_resolution(self):
+        with pytest.raises(ValueError, match="requires the DiT DLO transfer to be rank-local"):
+            OffloadConfig.from_od_config(
+                _offload_od_config(
+                    enable_distributed_layerwise_offload=True,
+                    dlo_use_allgather=True,
+                    dlo_resident_layers=12,
+                )
+            )

tests/diffusion/offloader/test_sequential_backend.py

@@ -7,7 +7,9 @@ import pytest
 import torch
 from torch import nn
 
+from vllm_omni.diffusion.hooks import HookRegistry
 from vllm_omni.diffusion.offloader.base import OffloadConfig, OffloadStrategy
+from vllm_omni.diffusion.offloader.offload_plan import OffloadPlan
 from vllm_omni.diffusion.offloader.sequential_backend import (
     ModelLevelOffloadBackend,
     SequentialOffloadHook,
@@ -48,11 +50,10 @@ def _track_pin_memory_calls():
     return tracker, mock
 
 
-def test_model_level_backend_delegates_to_custom_pipeline_offload() -> None:
-    class CustomPipeline(nn.Module):
+class _CustomPipeline(nn.Module):
     def __init__(self) -> None:
         super().__init__()
-            self.enable_args = None
+        self.enable_args: dict[str, object] | None = None
         self.disable_called = False
 
     def enable_omni_model_cpu_offload(self, **kwargs) -> None:
@@ -61,7 +62,9 @@ def test_model_level_backend_delegates_to_custom_pipeline_offload() -> None:
     def disable_omni_model_cpu_offload(self) -> None:
         self.disable_called = True
 
-    pipeline = CustomPipeline()
+
+def test_model_level_backend_delegates_to_custom_pipeline_offload() -> None:
+    pipeline = _CustomPipeline()
     backend = ModelLevelOffloadBackend(
         OffloadConfig(strategy=OffloadStrategy.MODEL_LEVEL, pin_cpu_memory=False),
         torch.device("cpu"),
@@ -82,6 +85,205 @@ def test_model_level_backend_delegates_to_custom_pipeline_offload() -> None:
     assert pipeline.disable_called is True
 
 
+def test_model_level_backend_passes_explicit_component_selection() -> None:
+    pipeline = _CustomPipeline()
+    backend = ModelLevelOffloadBackend(
+        OffloadConfig(
+            strategy=OffloadStrategy.MODEL_LEVEL,
+            components=frozenset({"dit", "text_encoder"}),
+        ),
+        torch.device("cpu"),
+    )
+
+    backend.enable(pipeline)
+
+    assert pipeline.enable_args is not None
+    assert pipeline.enable_args["offload_components"] == frozenset({"dit", "text_encoder"})
+
+
+def test_model_level_backend_rolls_back_partial_custom_enable() -> None:
+    class CustomPipeline(nn.Module):
+        def __init__(self) -> None:
+            super().__init__()
+            self.disable_called = False
+
+        def enable_omni_model_cpu_offload(self, **kwargs) -> None:
+            del kwargs
+            raise RuntimeError("injected custom enable failure")
+
+        def disable_omni_model_cpu_offload(self) -> None:
+            self.disable_called = True
+
+    pipeline = CustomPipeline()
+    backend = ModelLevelOffloadBackend(
+        OffloadConfig(strategy=OffloadStrategy.MODEL_LEVEL),
+        torch.device("cpu"),
+    )
+
+    with pytest.raises(RuntimeError, match="injected custom enable failure"):
+        backend.enable(pipeline)
+
+    assert pipeline.disable_called
+    assert not backend.enabled
+
+
+def test_sequential_offload_rolls_back_partial_hook_registration(monkeypatch: pytest.MonkeyPatch) -> None:
+    dit = _create_simple_module()
+    encoder = _create_simple_module()
+    original_register = HookRegistry.register_hook
+    calls = 0
+
+    def fail_second_registration(self, name, hook):
+        nonlocal calls
+        calls += 1
+        if calls == 2:
+            raise RuntimeError("injected registration failure")
+        return original_register(self, name, hook)
+
+    monkeypatch.setattr(HookRegistry, "register_hook", fail_second_registration)
+
+    with pytest.raises(RuntimeError, match="injected registration failure"):
+        apply_sequential_offload(
+            dit_modules=[dit],
+            encoder_modules=[encoder],
+            device=torch.device("cpu"),
+        )
+
+    for module in (dit, encoder):
+        registry = getattr(module, "_hook_registry", None)
+        assert registry is None or registry.get_hook(SequentialOffloadHook._HOOK_NAME) is None
+
+
+def test_model_level_disable_cleans_remaining_hooks_and_can_retry(monkeypatch: pytest.MonkeyPatch) -> None:
+    pipeline = nn.Module()
+    pipeline.transformer = _create_simple_module()
+    pipeline.text_encoder = _create_simple_module()
+    backend = ModelLevelOffloadBackend(
+        OffloadConfig(strategy=OffloadStrategy.MODEL_LEVEL, pin_cpu_memory=False),
+        torch.device("cpu"),
+    )
+    backend.enable(pipeline)
+
+    original_remove = HookRegistry.remove_hook
+    calls = 0
+
+    def fail_first_removal(self, name):
+        nonlocal calls
+        calls += 1
+        if calls == 1:
+            raise RuntimeError("injected hook removal failure")
+        original_remove(self, name)
+
+    monkeypatch.setattr(HookRegistry, "remove_hook", fail_first_removal)
+
+    with pytest.raises(RuntimeError, match="Failed to remove one or more sequential offload hooks"):
+        backend.disable()
+
+    assert backend.enabled
+    assert backend._offload_modules == [pipeline.transformer, pipeline.text_encoder]
+    assert pipeline.transformer._hook_registry.get_hook(SequentialOffloadHook._HOOK_NAME) is not None
+    assert pipeline.text_encoder._hook_registry.get_hook(SequentialOffloadHook._HOOK_NAME) is None
+
+    backend.disable()
+
+    assert not backend.enabled
+    assert not backend._offload_modules
+    for module in (pipeline.transformer, pipeline.text_encoder):
+        assert module._hook_registry.get_hook(SequentialOffloadHook._HOOK_NAME) is None
+
+
+def test_sequential_offload_filters_cpu_eligible_components() -> None:
+    dit = _create_simple_module()
+    encoder = _create_simple_module()
+    resident_stage = _create_simple_module()
+
+    apply_sequential_offload(
+        dit_modules=[dit],
+        encoder_modules=[encoder, resident_stage],
+        device=torch.device("cpu"),
+        offload_dit_modules=[dit],
+        offload_encoder_modules=[encoder],
+    )
+
+    dit_hook = dit._hook_registry.get_hook(SequentialOffloadHook._HOOK_NAME)
+    encoder_hook = encoder._hook_registry.get_hook(SequentialOffloadHook._HOOK_NAME)
+    resident_hook = resident_stage._hook_registry.get_hook(SequentialOffloadHook._HOOK_NAME)
+    assert dit_hook.offload_targets == [encoder]
+    assert encoder_hook.offload_targets == [dit]
+    assert encoder_hook.offload_after_context is True
+    assert resident_hook.offload_targets == [dit]
+    assert resident_hook.offload_after_context is False
+
+    remove_sequential_offload([dit, encoder, resident_stage])
+
+
+def test_move_params_handles_buffer_only_modules() -> None:
+    module = nn.Module()
+    module.register_buffer("state", torch.ones(2))
+
+    # An indexed CPU device compares differently while remaining usable on
+    # CPU-only test hosts, so the helper must inspect the buffer itself.
+    moved = SequentialOffloadHook._move_params(module, torch.device("cpu:1"))
+
+    assert moved
+    assert module.state.device.type == "cpu"
+
+
+def test_model_level_backend_keeps_declared_image_encoder_resident() -> None:
+    class MixedEncoderPipeline(nn.Module):
+        _offload_plan = OffloadPlan(
+            encoder_component_types={
+                "text_encoder": "text_encoder",
+            }
+        )
+
+        def __init__(self) -> None:
+            super().__init__()
+            self.transformer = _create_simple_module()
+            self.text_encoder = _create_simple_module()
+            self.image_encoder = _create_simple_module()
+
+    pipeline = MixedEncoderPipeline()
+    backend = ModelLevelOffloadBackend(
+        OffloadConfig(
+            strategy=OffloadStrategy.MODEL_LEVEL,
+            components=frozenset({"dit", "text_encoder"}),
+        ),
+        torch.device("cpu"),
+    )
+
+    backend.enable(pipeline)
+
+    text_hook = pipeline.text_encoder._hook_registry.get_hook(SequentialOffloadHook._HOOK_NAME)
+    image_hook = pipeline.image_encoder._hook_registry.get_hook(SequentialOffloadHook._HOOK_NAME)
+    assert text_hook.offload_after_context is True
+    assert image_hook.offload_after_context is False
+
+    backend.disable()
+
+
+@pytest.mark.parametrize(
+    ("component", "attribute", "message"),
+    [
+        ("text_encoder", "text_encoder", "requires a DiT/transformer"),
+        ("dit", "transformer", "requires an encoder execution stage"),
+    ],
+)
+def test_component_selective_model_offload_requires_swap_counterpart(component, attribute, message) -> None:
+    pipeline = nn.Module()
+    setattr(pipeline, attribute, _create_simple_module())
+    backend = ModelLevelOffloadBackend(
+        OffloadConfig(
+            strategy=OffloadStrategy.MODEL_LEVEL,
+            components=frozenset({component}),
+        ),
+        torch.device("cpu"),
+    )
+
+    with pytest.raises(ValueError, match=message):
+        backend.enable(pipeline)
+
+
 def test_sequential_offload_can_begin_with_dit_on_cpu(monkeypatch: pytest.MonkeyPatch) -> None:
     dit = _create_simple_module()
     encoder = _create_simple_module()

tests/diffusion/test_diffusion_engine.py

@@ -600,6 +600,34 @@ class TestDiffusionCompileConfig:
             ),
             ({"enable_cpu_offload": True}, "CPU offload"),
             ({"enable_layerwise_offload": True}, "layerwise offload"),
+            (
+                {
+                    "diffusion_offload_config": {
+                        "mode": "module",
+                        "components": ["dit"],
+                    }
+                },
+                "CPU offload",
+            ),
+            (
+                {
+                    "diffusion_offload_config": {
+                        "mode": "layer",
+                        "components": ["dit"],
+                    }
+                },
+                "layerwise offload",
+            ),
+            (
+                {
+                    "diffusion_offload_config": {
+                        "mode": "layer",
+                        "components": ["dit"],
+                        "layer_options": {"dit": {"weight_transfer": "allgather"}},
+                    }
+                },
+                "distributed layerwise offload",
+            ),
         ],
     )
     def test_full_compile_rejects_incompatible_features(self, kwargs, feature) -> None:
@@ -610,6 +638,34 @@ class TestDiffusionCompileConfig:
                 **kwargs,
             )
 
+    def test_raw_parallel_mapping_is_normalized_before_allgather_cache_validation(self) -> None:
+        with pytest.raises(ValueError, match="rank-local cache decisions"):
+            OmniDiffusionConfig(
+                model="test",
+                num_gpus=2,
+                parallel_config={"data_parallel_size": 2},
+                cache_backend="tea_cache",
+                diffusion_offload_config={
+                    "mode": "layer",
+                    "components": ["dit"],
+                    "layer_options": {"dit": {"weight_transfer": "allgather"}},
+                },
+            )
+
+    def test_auto_dp_is_resolved_before_encoder_prompt_cache_validation(self) -> None:
+        with pytest.raises(ValueError, match="rank-local cache hits"):
+            OmniDiffusionConfig(
+                model="test",
+                num_gpus=2,
+                parallel_config={"data_parallel_size": None},
+                enable_prompt_embed_cache=True,
+                diffusion_offload_config={
+                    "mode": "layer",
+                    "components": ["text_encoder"],
+                    "layer_options": {"text_encoder": {"weight_transfer": "allgather"}},
+                },
+            )
+
 
 class TestRequestBatchAdmission:
     pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]

tests/diffusion/test_diffusion_engine_dummy_run.py

@@ -52,6 +52,23 @@ def test_dummy_run_uses_enough_steps_for_execution_mode(
     assert captured_requests[0].sampling_params.num_inference_steps == 2
 
 
+def test_allgather_startup_runs_broadcast_dummy_request() -> None:
+    engine = object.__new__(DiffusionEngine)
+    engine.od_config = SimpleNamespace(
+        diffusion_offload_config={
+            "mode": "layer",
+            "components": ["dit"],
+            "layer_options": {"dit": {"weight_transfer": "allgather"}},
+        },
+        parallel_config=SimpleNamespace(data_parallel_size=2, sequence_parallel_size=1),
+    )
+    engine._dummy_run = Mock()
+
+    engine.run_startup_warmup()
+
+    engine._dummy_run.assert_called_once_with()
+
+
 def test_dummy_run_num_frames_uses_explicit_model_setting(monkeypatch: pytest.MonkeyPatch) -> None:
     class JointAudioVideoModel:
         dummy_run_num_frames = 2

tests/diffusion/test_diffusion_model_runner.py

@@ -484,20 +484,23 @@ def test_compile_transformer_uses_full_granularity(monkeypatch):
 
 @pytest.mark.core_model
 @pytest.mark.cpu
-def test_compile_transformer_falls_back_after_synchronous_setup_failure(monkeypatch, caplog):
+def test_compile_transformer_falls_back_after_synchronous_setup_failure(monkeypatch):
     model = _CompileTrackingModel()
     runner = _make_compile_runner(model)
+    warning = Mock()
 
     def _regionally_compile(*args, **kwargs):
         raise RuntimeError("compile setup failed")
 
     monkeypatch.setattr(model_runner_module, "regionally_compile", _regionally_compile)
+    monkeypatch.setattr(model_runner_module.logger, "warning", warning)
 
     DiffusionModelRunner._compile_transformer(runner, "transformer")
 
     assert runner.pipeline.transformer is model
-    assert "failed before activation" in caplog.text
-    assert "lazy compilation errors" in caplog.text
+    warning.assert_called_once()
+    assert "failed before activation" in warning.call_args.args[0]
+    assert "lazy compilation errors" in warning.call_args.args[0]
 
 
 @pytest.mark.core_model
@@ -1061,6 +1064,51 @@ def test_load_model_clears_cache_backend_for_unsupported_pipeline(monkeypatch):
     assert dummy_cache_backend.enabled is False
 
 
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_load_model_rejects_env_prompt_cache_before_text_encoder_allgather(monkeypatch):
+    runner = object.__new__(DiffusionModelRunner)
+    runner.vllm_config = object()
+    runner.device = torch.device("cpu")
+    runner.pipeline = None
+    runner.offload_backend = None
+    runner.od_config = SimpleNamespace(
+        diffusion_offload_config={
+            "mode": "layer",
+            "components": ["text_encoder"],
+            "layer_options": {"text_encoder": {"weight_transfer": "allgather"}},
+        },
+        enable_cpu_offload=False,
+        enable_layerwise_offload=False,
+        enable_distributed_layerwise_offload=True,
+        dlo_use_allgather=True,
+        dlo_resident_layers=0,
+        pin_cpu_memory=True,
+        parallel_config=SimpleNamespace(data_parallel_size=2),
+        enable_prompt_embed_cache=False,
+        prompt_embed_cache_size=32,
+        cache_backend="none",
+        model_class_name="Wan22Pipeline",
+        enforce_eager=True,
+        streaming_output=False,
+        step_execution=False,
+    )
+    monkeypatch.setenv("OMNI_DIFFUSION_PROMPT_EMBED_CACHE", "1")
+    runtime_init = Mock(side_effect=pytest.fail)
+    loader = Mock(side_effect=pytest.fail)
+    monkeypatch.setattr(model_runner_module.current_omni_platform, "init_diffusion_model_runner_runtime", runtime_init)
+    monkeypatch.setattr(model_runner_module, "DiffusersPipelineLoader", loader)
+    enable_offload = Mock(side_effect=pytest.fail)
+    monkeypatch.setattr(model_runner_module, "enable_offload_backend", enable_offload)
+
+    with pytest.raises(ValueError, match="Prompt embedding cache cannot be combined"):
+        DiffusionModelRunner.load_model(runner)
+
+    runtime_init.assert_not_called()
+    loader.assert_not_called()
+    enable_offload.assert_not_called()
+
+
 def test_set_forward_context_enters_vllm_config_contexts(monkeypatch):
     """Ensure `with set_forward_context(...):` enters vllm's context managers internally and calls desired vllm functions."""
     import vllm.config.vllm as vllm_config_module

tests/diffusion/test_multiproc_engine_concurrency.py

@@ -254,8 +254,11 @@ class TestRequestModeDispatch:
         executor.od_config = SimpleNamespace(
             step_execution=False,
             parallel_config=SimpleNamespace(data_parallel_size=2),
-            enable_distributed_layerwise_offload=True,
-            dlo_use_allgather=True,
+            diffusion_offload_config={
+                "mode": "layer",
+                "components": ["dit"],
+                "layer_options": {"dit": {"weight_transfer": "allgather"}},
+            },
         )
         executor.execute_request = Mock(return_value="dlo-dp")
         executor.collective_rpc = Mock()
@@ -358,7 +361,7 @@ class TestRequestModeDispatch:
         for rank, result_mq in enumerate(executor._result_mqs):
             assert result_mq.dequeue.call_count == (1 if rank in expected_primary_ranks else 0)
 
-    @pytest.mark.parametrize("empty_prompt", ["", {"prompt": ""}])
+    @pytest.mark.parametrize("empty_prompt", ["", [], {}, {"prompt": ""}])
     def test_dlo_dp_rejects_empty_prompt_before_worker_dispatch(self, empty_prompt):
         executor, _, _ = _make_executor(num_gpus=2)
         executor.od_config = SimpleNamespace(
@@ -376,6 +379,71 @@ class TestRequestModeDispatch:
 
         executor.collective_rpc.assert_not_called()
 
+    def test_dlo_dp_accepts_precomputed_embeddings_without_prompt_text(self):
+        executor, _, _ = _make_executor(num_gpus=2)
+        executor.od_config = SimpleNamespace(
+            step_execution=False,
+            parallel_config=SimpleNamespace(data_parallel_size=2),
+            diffusion_offload_config={
+                "mode": "layer",
+                "components": ["text_encoder"],
+                "layer_options": {"text_encoder": {"weight_transfer": "allgather"}},
+            },
+        )
+        executor.collective_rpc = Mock(return_value=[_tagged_output("A"), _tagged_output("B")])
+        scheduler_output = _make_sched_output("A", "B")
+        for new_req in scheduler_output.scheduled_new_reqs:
+            new_req.req.prompt = {"prompt": "", "prompt_embeds": object()}
+
+        result = executor.execute_request(scheduler_output)
+
+        assert [output.result.error for output in result.runner_outputs] == ["A", "B"]
+
+    @pytest.mark.parametrize(
+        "prompt",
+        [
+            [1, 2],
+            {"type": "token", "prompt_token_ids": [1, 2]},
+            {"prompt_ids": [1, 2], "prompt_mask": object()},
+        ],
+    )
+    def test_dlo_dp_accepts_tokenized_prompts(self, prompt):
+        executor, _, _ = _make_executor(num_gpus=2)
+        executor.od_config = SimpleNamespace(
+            step_execution=False,
+            parallel_config=SimpleNamespace(data_parallel_size=2),
+            enable_distributed_layerwise_offload=True,
+            dlo_use_allgather=True,
+        )
+        executor.collective_rpc = Mock(return_value=[_tagged_output("A"), _tagged_output("B")])
+        scheduler_output = _make_sched_output("A", "B")
+        for new_req in scheduler_output.scheduled_new_reqs:
+            new_req.req.prompt = prompt
+
+        result = executor.execute_request(scheduler_output)
+
+        assert [output.result.error for output in result.runner_outputs] == ["A", "B"]
+
+    def test_text_encoder_allgather_rejects_mismatched_embedding_paths(self):
+        executor, _, _ = _make_executor(num_gpus=2)
+        executor.od_config = SimpleNamespace(
+            step_execution=False,
+            parallel_config=SimpleNamespace(data_parallel_size=2),
+            diffusion_offload_config={
+                "mode": "layer",
+                "components": ["text_encoder"],
+                "layer_options": {"text_encoder": {"weight_transfer": "allgather"}},
+            },
+        )
+        executor.collective_rpc = Mock()
+        scheduler_output = _make_sched_output("A", "B")
+        scheduler_output.scheduled_new_reqs[1].req.prompt = {"prompt_embeds": object()}
+
+        with pytest.raises(ValueError, match="same positive/negative prompt embedding fields"):
+            executor.execute_request(scheduler_output)
+
+        executor.collective_rpc.assert_not_called()
+
     def test_dlo_dp_valid_wave_still_runs_after_rejected_wave(self):
         executor, _, _ = _make_executor(num_gpus=2)
         executor.od_config = SimpleNamespace(
@@ -396,24 +464,6 @@ class TestRequestModeDispatch:
         assert [output.result.error for output in result.runner_outputs] == ["A", "B"]
         executor.collective_rpc.assert_called_once()
 
-    def test_dlo_dp_allows_shared_default_denoise_steps(self):
-        executor, _, _ = _make_executor(num_gpus=2)
-        executor.od_config = SimpleNamespace(
-            step_execution=False,
-            parallel_config=SimpleNamespace(data_parallel_size=2),
-            enable_distributed_layerwise_offload=True,
-            dlo_use_allgather=True,
-        )
-        executor.collective_rpc = Mock(return_value=[_tagged_output("A"), _tagged_output("B")])
-        scheduler_output = _make_sched_output("A", "B")
-        for new_req in scheduler_output.scheduled_new_reqs:
-            new_req.req.sampling_params.num_inference_steps = None
-
-        result = executor.execute_request(scheduler_output)
-
-        assert [output.result.error for output in result.runner_outputs] == ["A", "B"]
-        executor.collective_rpc.assert_called_once()
-
     def test_dlo_dp_forwards_request_metadata_envelopes_as_one_wave(self):
         executor, _, _ = _make_executor(num_gpus=2)
         executor.od_config = SimpleNamespace(
@@ -496,6 +546,28 @@ class TestRequestModeDispatch:
         executor._fail_closed_on_dp_wave_timeout.assert_called_once()
         assert isinstance(executor._fail_closed_on_dp_wave_timeout.call_args.args[0], TimeoutError)
 
+    def test_single_allgather_request_times_out_and_fails_closed(self):
+        from vllm_omni.diffusion.executor import multiproc_executor as executor_module
+
+        executor, _, _ = _make_executor(num_gpus=2)
+        executor.od_config = SimpleNamespace(
+            step_execution=False,
+            parallel_config=SimpleNamespace(data_parallel_size=2),
+            diffusion_offload_config={
+                "mode": "layer",
+                "components": ["dit"],
+                "layer_options": {"dit": {"weight_transfer": "allgather"}},
+            },
+        )
+        executor.collective_rpc = Mock(side_effect=TimeoutError("timed out"))
+        executor._fail_closed_on_dp_wave_timeout = Mock()
+
+        result = executor.execute_request(_make_sched_output("A"))
+
+        assert result.runner_outputs[0].result.error == "timed out"
+        assert executor.collective_rpc.call_args.kwargs["timeout"] == executor_module._DLO_DP_WAVE_TIMEOUT_S
+        executor._fail_closed_on_dp_wave_timeout.assert_called_once()
+
 
 # ───────────────── concurrent collective RPC ─────────────────
 

tests/e2e/online_serving/test_minimax_h3_dlo_dp2_t2va.py

@@ -51,7 +51,18 @@ SERVER_ARGS = [
     "--vae-parallel-mode",
     "tile",
     "--vae-use-tiling",
-    "--enable-distributed-layerwise-offload",
+    "--diffusion-offload-config",
+    json.dumps(
+        {
+            "mode": "layer",
+            "components": ["dit", "text_encoder"],
+            "layer_options": {
+                "dit": {"weight_transfer": "allgather"},
+                "text_encoder": {"weight_transfer": "rank-local"},
+            },
+        },
+        separators=(",", ":"),
+    ),
 ]
 
 

tests/entrypoints/test_async_omni_diffusion_config.py

@@ -6,6 +6,7 @@ from types import SimpleNamespace
 import pytest
 
 from vllm_omni.diffusion.data import AttentionConfig
+from vllm_omni.engine import stage_init_utils
 from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
 from vllm_omni.entrypoints.cli.serve import OmniServeCommand
 from vllm_omni.utils.tracking_parser import TrackingArgumentParser
@@ -426,8 +427,8 @@ def test_serve_cli_forwards_distilled_lora_to_diffusion_stage():
     ]
 
 
-def test_serve_cli_forwards_distributed_offload_residency():
-    """Ensure the two-GPU DLO placement controls reach the diffusion stage."""
+def test_serve_cli_forwards_compact_diffusion_offload_config():
+    """Ensure component-specific layer settings reach the diffusion stage."""
     parser = TrackingArgumentParser()
     subparsers = parser.add_subparsers(dest="command")
     OmniServeCommand().subparser_init(subparsers)
@@ -437,10 +438,11 @@ def test_serve_cli_forwards_distributed_offload_residency():
             "serve",
             "MiniMaxAI/MiniMax-H3",
             "--omni",
-            "--enable-distributed-layerwise-offload",
-            "--dlo-no-use-allgather",
-            "--dlo-resident-layers",
-            "20",
+            "--diffusion-offload-config",
+            '{"mode":"layer","components":["dit","text_encoder"],'
+            '"layer_options":{"dit":{"weight_transfer":"rank-local","resident_layers":20},'
+            '"text_encoder":{"weight_transfer":"allgather"}},'
+            '"pin_memory":true}',
         ]
     )
 
@@ -448,12 +450,45 @@ def test_serve_cli_forwards_distributed_offload_residency():
     stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(explicit_kwargs)[0]
     engine_args = stage_cfg["engine_args"]
 
-    assert args.enable_distributed_layerwise_offload is True
-    assert args.dlo_use_allgather is False
-    assert args.dlo_resident_layers == 20
-    assert engine_args["enable_distributed_layerwise_offload"] is True
-    assert engine_args["dlo_use_allgather"] is False
-    assert engine_args["dlo_resident_layers"] == 20
+    expected = {
+        "mode": "layer",
+        "components": ["dit", "text_encoder"],
+        "layer_options": {
+            "dit": {"weight_transfer": "rank-local", "resident_layers": 20},
+            "text_encoder": {"weight_transfer": "allgather"},
+        },
+        "pin_memory": True,
+    }
+    assert args.diffusion_offload_config == expected
+    assert engine_args["diffusion_offload_config"] == expected
+
+
+def test_invalid_diffusion_offload_config_fails_before_model_loading(monkeypatch, mocker):
+    load_model = mocker.patch("vllm_omni.diffusion.model_loader.diffusers_loader.DiffusersPipelineLoader.load_model")
+    create_client = mocker.patch("vllm_omni.diffusion.stage_diffusion_client.create_diffusion_client")
+    monkeypatch.setattr(
+        stage_init_utils,
+        "build_engine_args_dict",
+        lambda *_args, **_kwargs: {
+            "model": "test",
+            "diffusion_offload_config": {
+                "mode": "layerwise",
+                "components": ["dit"],
+            },
+        },
+    )
+
+    with pytest.raises(ValueError, match="Unknown diffusion offload mode"):
+        stage_init_utils.initialize_diffusion_stage(
+            stage_id=0,
+            model="test",
+            stage_cfg=object(),
+            metadata=mocker.Mock(),
+            stage_init_timeout=30,
+        )
+
+    create_client.assert_not_called()
+    load_model.assert_not_called()
 
 
 def test_serve_cli_forwards_hwr_policy_for_no_allgather_dlo():

vllm_omni/config/omni_config.py

@@ -328,9 +328,11 @@ def _stage_cli_overrides(
 ) -> dict[str, Any]:
     runtime_overrides = build_stage_runtime_overrides(stage_id, dict(cli_overrides))
     global_stage_fields = _global_stage_cli_fields()
+    owned_fields = None if execution_type is None else _STAGE_ENGINE_FIELDS_BY_EXECUTION_TYPE[execution_type]
     result: dict[str, Any] = {}
     for key, value in runtime_overrides.items():
-        if key in global_stage_fields or f"stage_{stage_id}_{key}" in cli_overrides:
+        stage_specific = f"stage_{stage_id}_{key}" in cli_overrides
+        if stage_specific or (key in global_stage_fields and (owned_fields is None or key in owned_fields)):
             result[key] = _copy_value(value)
 
     # step_execution is a diffusion execution protocol, not an LLM engine
@@ -340,6 +342,25 @@ def _stage_cli_overrides(
     return result
 
 
+def _validate_global_stage_cli_ownership(
+    pipeline: PipelineConfig,
+    cli_overrides: Mapping[str, Any],
+) -> None:
+    """Reject global stage arguments that no stage in the pipeline owns."""
+    explicit_global_fields = {
+        key for key, value in cli_overrides.items() if value is not None and key in _global_stage_cli_fields()
+    }
+    owned_fields = {
+        field for stage in pipeline.stages for field in _STAGE_ENGINE_FIELDS_BY_EXECUTION_TYPE[stage.execution_type]
+    }
+    unowned_fields = explicit_global_fields - owned_fields
+    if unowned_fields:
+        names = ", ".join(sorted(unowned_fields))
+        raise ValueError(
+            f"Pipeline {pipeline.model_type!r} has explicit engine argument(s) with no structured config owner: {names}"
+        )
+
+
 def _resolve_deploy_path(deploy_config_path: str) -> Path:
     deploy_path = Path(deploy_config_path)
     if not deploy_path.exists() and deploy_path.parent == Path("."):
@@ -707,6 +728,9 @@ class _DiffusionConfigProjection:
     lora_backend: str = "peft"
     max_cpu_loras: int | None = None
     output_type: str = "pil"
+    diffusion_offload_config: dict[str, Any] | None = None
+    # Compatibility aliases for existing callers and model-specific stage
+    # lifecycles that are broader than the compact dit/text_encoder selector.
     enable_cpu_offload: bool = False
     enable_layerwise_offload: bool = False
     enable_distributed_layerwise_offload: bool = False
@@ -759,8 +783,12 @@ class _DiffusionConfigProjection:
     @classmethod
     def from_kwargs(cls, **kwargs: Any) -> _DiffusionConfigProjection:
         from vllm_omni.diffusion.data import normalize_omni_diffusion_kwargs
+        from vllm_omni.diffusion.offloader.config import parse_diffusion_offload_config
 
         normalized_kwargs = normalize_omni_diffusion_kwargs(kwargs)
+        # Validate before stage construction while retaining the raw mapping
+        # needed by dataclass/config serialization across process boundaries.
+        parse_diffusion_offload_config(normalized_kwargs.get("diffusion_offload_config"))
         valid_fields = {f.name for f in fields(cast(Any, cls))}
         return cls(**{k: v for k, v in normalized_kwargs.items() if k in valid_fields})
 
@@ -1865,6 +1893,7 @@ class VllmOmniConfig:
         if cli_overrides is None:
             cli_overrides = {}
         cli_overrides = normalize_pipeline_cli_overrides(pipeline_cfg, cli_overrides)
+        _validate_global_stage_cli_ownership(pipeline_cfg, cli_overrides)
 
         deploy, loaded_deploy_config_path = _get_deploy_config(
             pipeline_cfg,

vllm_omni/config/stage_config.py

@@ -472,6 +472,9 @@ class StageDeployConfig:
     # Runtime optimizations used by diffusion loading/execution.
     enable_multithread_weight_load: bool | None = None
     num_weight_load_threads: int | None = None
+    diffusion_offload_config: dict[str, Any] | None = None
+    # Compatibility aliases for existing callers and model-specific stage
+    # lifecycles that are broader than the compact dit/text_encoder selector.
     enable_cpu_offload: bool | None = None
     enable_layerwise_offload: bool | None = None
 

vllm_omni/diffusion/data.py

@@ -804,7 +804,12 @@ class OmniDiffusionConfig:
 
     output_type: str = "pil"
 
-    # CPU offload parameters
+    # CPU offload parameters. Keep the public mapping raw so stage configs can
+    # serialize it across processes; __post_init__ validates it once and caches
+    # the internal typed resolution used at runtime.
+    diffusion_offload_config: dict[str, Any] | None = None
+    # Compatibility aliases. Some model-specific legacy stage lifecycles are
+    # intentionally broader than the compact dit/text_encoder selector.
     # When enabled, DiT and encoders swap GPU access (mutual exclusion):
     # - Text encoders run on GPU while DiT is on CPU
     # - DiT runs on GPU while encoders are on CPU
@@ -1056,6 +1061,11 @@ class OmniDiffusionConfig:
         )
 
     def __post_init__(self):
+        from vllm_omni.diffusion.offloader.config import (
+            OffloadStrategy,
+            materialize_legacy_offload_flags,
+        )
+
         if self.diffusion_compile_granularity not in {"regional", "full"}:
             raise ValueError(
                 "diffusion_compile_granularity must be 'regional' or 'full', "
@@ -1128,6 +1138,9 @@ class OmniDiffusionConfig:
                 self.num_gpus = 1
 
         self.parallel_config.resolve_data_parallel_size(self.num_gpus)
+        # Resolve offload only after DP/SP normalization so cached policy
+        # validation observes the actual execution topology.
+        offload_strategy = materialize_legacy_offload_flags(self)
 
         if self.diffusion_compile_granularity == "full":
             incompatible_features = []
@@ -1136,10 +1149,14 @@ class OmniDiffusionConfig:
                 incompatible_features.append("HSDP")
             if self.parallel_config.sequence_parallel_size > 1:
                 incompatible_features.append("sequence parallelism")
-            if self.enable_cpu_offload:
-                incompatible_features.append("CPU offload")
-            if self.enable_layerwise_offload:
-                incompatible_features.append("layerwise offload")
+            if offload_strategy is not OffloadStrategy.NONE:
+                incompatible_features.append(
+                    {
+                        OffloadStrategy.MODEL_LEVEL: "CPU offload",
+                        OffloadStrategy.LAYER_WISE: "layerwise offload",
+                        OffloadStrategy.DISTRIBUTED_LAYER_WISE: "distributed layerwise offload",
+                    }[offload_strategy]
+                )
             if incompatible_features:
                 features = ", ".join(incompatible_features)
                 raise ValueError(

vllm_omni/diffusion/diffusion_engine.py

@@ -41,6 +41,7 @@ from vllm_omni.diffusion.io_support import (
     supports_audio_output,
     supports_multimodal_input,
 )
+from vllm_omni.diffusion.offloader.config import any_selected_component_uses_allgather
 from vllm_omni.diffusion.output_formatter import (
     format_diffusion_outputs,
     format_empty_diffusion_outputs,
@@ -181,11 +182,7 @@ def _max_num_seqs(od_config: OmniDiffusionConfig) -> int:
 def _uses_dlo_dp_concurrency(od_config: OmniDiffusionConfig) -> bool:
     parallel_config = getattr(od_config, "parallel_config", None)
     dp_size = getattr(parallel_config, "data_parallel_size", 1)
-    return (
-        dp_size > 1
-        and getattr(od_config, "enable_distributed_layerwise_offload", False)
-        and getattr(od_config, "dlo_use_allgather", True)
-    )
+    return dp_size > 1 and any_selected_component_uses_allgather(od_config)
 
 
 def _move_tensor_tree_to_cpu(value: object) -> object:
@@ -1012,33 +1009,12 @@ class DiffusionEngine:
                 raise RuntimeError(f"Could not {action} profiler: {e}") from e
 
     def run_startup_warmup(self) -> None:
-        dlo_use_allgather = getattr(self.od_config, "dlo_use_allgather", True)
-        # Skip dummy run when AllGather is used with more than 1 rank,
-        # because the dummy run sends only 1 request but AllGather requires
-        # all ranks to participate simultaneously.  This covers both DP > 1
-        # and SP > 1 (where dp_size is derived from sp_size in OffloadConfig).
-        pc = getattr(self.od_config, "parallel_config", None)
-        dp_size = getattr(pc, "data_parallel_size", 1) if pc else 1
-        sp_size = getattr(pc, "sequence_parallel_size", 1) if pc else 1
-        effective_shard_size = max(dp_size, sp_size)
-        skip_dummy = (
-            getattr(self.od_config, "enable_distributed_layerwise_offload", False)
-            and dlo_use_allgather
-            and effective_shard_size > 1
-        )
-        if skip_dummy:
-            logger.info(
-                "Skipping dummy run (dist_offload with AllGather, dp_size=%d, sp_size=%d)",
-                dp_size,
-                sp_size,
-            )
-            return
         try:
             self._dummy_run()
         except Exception as e:
             logger.error(f"Dummy run failed: {e}")
             self.close()
-            raise e
+            raise
 
     def _make_dummy_request(
         self,
@@ -1050,10 +1026,13 @@ class DiffusionEngine:
         num_inference_steps: int = 1,
     ) -> OmniDiffusionRequest | None:
         """Build a minimal model request for startup profiling or warmup."""
-        prompt: OmniTextPrompt = {"prompt": "dummy run"}
+        prompt = OmniTextPrompt(prompt="dummy run")
+        model_class_name = self.od_config.model_class_name
+        if model_class_name is None:
+            raise RuntimeError("Dummy request requires a resolved model_class_name")
         supports_image_input, supports_audio_input = supports_multimodal_input(self.od_config)
         if supports_image_input:
-            color_format = image_color_format(self.od_config.model_class_name)
+            color_format = image_color_format(model_class_name)
             images = [PIL.Image.new(color_format, (width, height)) for _ in range(num_image_inputs)]
             prompt.setdefault("multi_modal_data", {})["image"] = images[0] if len(images) == 1 else images
 
@@ -1061,7 +1040,7 @@ class DiffusionEngine:
             audio_sr = 16000
             prompt.setdefault("multi_modal_data", {})["audio"] = np.random.randn(audio_sr * 2).astype(np.float32)
 
-        num_frames = get_dummy_run_num_frames(self.od_config.model_class_name, supports_audio_input)
+        num_frames = get_dummy_run_num_frames(model_class_name, supports_audio_input)
         if num_frames <= 0:
             return None
         return OmniDiffusionRequest(
@@ -1101,11 +1080,14 @@ class DiffusionEngine:
             is not DiffusionKVCacheMode.PAGED_SCHEDULER
         ):
             return None
+        model_class_name = self.od_config.model_class_name
+        if model_class_name is None:
+            raise RuntimeError("Diffusion KV profiling requires a resolved model_class_name")
         request = self._make_dummy_request(
             height=1024,
             width=1024,
             guidance_scale=5.0,
-            num_image_inputs=get_dummy_run_num_image_inputs(self.od_config.model_class_name),
+            num_image_inputs=get_dummy_run_num_image_inputs(model_class_name),
         )
         if request is None:
             raise RuntimeError("paged_scheduler requires a runnable Diffusion KV memory profile request")

vllm_omni/diffusion/executor/multiproc_executor.py

@@ -23,6 +23,11 @@ from vllm.v1.executor.multiproc_executor import set_multiprocessing_worker_envs
 from vllm_omni.diffusion.data import SHUTDOWN_MESSAGE, AsyncDiffusionOutput, AsyncOutputKind, DiffusionOutput
 from vllm_omni.diffusion.executor.abstract import DiffusionExecutor
 from vllm_omni.diffusion.ipc import DIFFUSION_RPC_RESULT_ENVELOPE, unpack_diffusion_output_shm
+from vllm_omni.diffusion.offloader.config import (
+    TEXT_ENCODER_COMPONENT,
+    any_selected_component_uses_allgather,
+    resolve_offload,
+)
 from vllm_omni.diffusion.sched.request_scheduler import build_request_batch_sampling_params_key
 from vllm_omni.diffusion.utils.future_utils import try_set_exception, try_set_result
 from vllm_omni.diffusion.worker import WorkerProc
@@ -44,13 +49,30 @@ def _is_empty_dp_prompt(prompt: object) -> bool:
     """Return whether a DP request has no usable text prompt."""
     if prompt is None:
         return True
-    if isinstance(prompt, str):
+    if isinstance(prompt, (str, list, tuple)):
         return not prompt
-    if isinstance(prompt, dict) and "prompt" in prompt:
-        return not prompt["prompt"]
+    if isinstance(prompt, dict):
+        return (
+            not prompt.get("prompt")
+            and not prompt.get("prompt_token_ids")
+            and not prompt.get("prompt_ids")
+            and prompt.get("prompt_embeds") is None
+        )
     return False
 
 
+def _text_encoder_input_signature(prompt: object) -> tuple[bool, bool]:
+    """Describe precomputed embeddings that change encoder forward counts."""
+    if not isinstance(prompt, dict):
+        return False, False
+    return prompt.get("prompt_embeds") is not None, prompt.get("negative_prompt_embeds") is not None
+
+
+def _uses_text_encoder_allgather(config: object) -> bool:
+    resolved = resolve_offload(config)
+    return resolved.offloads(TEXT_ENCODER_COMPONENT) and resolved.uses_allgather(TEXT_ENCODER_COMPONENT)
+
+
 @dataclass
 class _ExecutorShutdownCleaner:
     """Finalizer that shuts down executor worker processes."""
@@ -471,16 +493,11 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
         # request and Diffusion KV metadata inseparable.
         # All ranks reply (unique_reply_rank=None) so we collect dp_size
         # responses and match by dp_rank.
-        if (
-            len(new_reqs) > 1
-            and getattr(self.od_config, "enable_distributed_layerwise_offload", False)
-            and getattr(self.od_config, "dlo_use_allgather", True)
-        ):
+        if len(new_reqs) > 1 and any_selected_component_uses_allgather(self.od_config):
             # Reuse the request scheduler's complete compatibility key. DLO
             # AllGather requires every DP rank to execute the same collective
             # schedule, including shape, CFG, denoise steps, output count,
-            # and LoRA settings. Two default (None) step counts are valid and
-            # resolve identically inside the same pipeline.
+            # and LoRA settings.
             compatibility_keys = [build_request_batch_sampling_params_key(nr.req) for nr in new_reqs]
             if any(key != compatibility_keys[0] for key in compatibility_keys[1:]):
                 raise ValueError(
@@ -498,6 +515,13 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
                     "share identical extra_args. Different extra_args can change "
                     "the forward schedule and cause AllGather deadlock."
                 )
+            if _uses_text_encoder_allgather(self.od_config):
+                encoder_signatures = {_text_encoder_input_signature(nr.req.prompt) for nr in new_reqs}
+                if len(encoder_signatures) > 1:
+                    raise ValueError(
+                        "DLO text_encoder AllGather requires every concurrent request "
+                        "to provide the same positive/negative prompt embedding fields."
+                    )
             empty_prompt_ids = [nr.request_id for nr in new_reqs if _is_empty_dp_prompt(nr.req.prompt)]
             if empty_prompt_ids:
                 raise ValueError(
@@ -547,11 +571,16 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
                 args: tuple = (req, self.od_config, scheduler_output.kv_prefetch_job)
                 if new_req.diffusion_kv_metadata is not None:
                     args += (new_req.diffusion_kv_metadata,)
+                allgather_active = any_selected_component_uses_allgather(self.od_config)
+                timeout_options: dict[str, Any] = {}
+                if allgather_active:
+                    timeout_options["timeout"] = _DLO_DP_WAVE_TIMEOUT_S
                 result = self.collective_rpc(
                     "execute_model",
                     args=args,
                     unique_reply_rank=0,
                     exec_all_ranks=True,
+                    **timeout_options,
                 )
                 if isinstance(result, AsyncDiffusionOutput) and result.kind == AsyncOutputKind.COMPUTE_DONE:
                     runner_outputs.append(
@@ -575,6 +604,8 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
                 else:
                     raise RuntimeError(f"Unexpected response type: {type(result)!r}")
             except Exception as exc:
+                if isinstance(exc, TimeoutError):
+                    self._fail_closed_on_dp_wave_timeout(exc)
                 runner_outputs.append(
                     RunnerOutput(
                         request_id=new_req.request_id,
@@ -605,11 +636,7 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
 
         parallel_config = getattr(self.od_config, "parallel_config", None)
         dp_size = getattr(parallel_config, "data_parallel_size", 1)
-        if (
-            dp_size > 1
-            and getattr(self.od_config, "enable_distributed_layerwise_offload", False)
-            and getattr(self.od_config, "dlo_use_allgather", True)
-        ):
+        if dp_size > 1 and any_selected_component_uses_allgather(self.od_config):
             # DLO DP uses one independent request per DP replica.  It is not a
             # fused pipeline request batch, so models such as MiniMax-H3 do not
             # need to advertise supports_request_batch=True.

vllm_omni/diffusion/model_loader/diffusers_loader.py

@@ -42,7 +42,15 @@ from vllm_omni.diffusion.model_loader.host_weight_plan import (
     has_online_quantization,
 )
 from vllm_omni.diffusion.models.diffusers_adapter.pipeline_diffusers_adapter import DiffusersAdapterPipeline
+from vllm_omni.diffusion.offloader.component_utils import encoder_component_type
+from vllm_omni.diffusion.offloader.config import (
+    DIT_COMPONENT,
+    TEXT_ENCODER_COMPONENT,
+    OffloadStrategy,
+    resolve_offload,
+)
 from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery
+from vllm_omni.diffusion.offloader.offload_plan import get_offload_plan
 from vllm_omni.diffusion.registry import initialize_model
 from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific
 
@@ -546,14 +554,51 @@ class DiffusersPipelineLoader(HWRLoaderMixin):
             else:
                 model = self._init_from_load_format(load_format, target_device, custom_pipeline_name, is_hsdp=False)
 
-                _dist_offload = getattr(self.od_config, "enable_distributed_layerwise_offload", False)
-                _use_ag = getattr(self.od_config, "dlo_use_allgather", True)
-                _has_online_quant = self._has_online_quant(model)
-                _tp_size = int(getattr(self.parallel_config, "tensor_parallel_size", 1))
-                _use_hsdp = bool(getattr(self.parallel_config, "use_hsdp", False))
-                _dp_size = int(getattr(self.parallel_config, "data_parallel_size", 1))
-                _sp_size = int(getattr(self.parallel_config, "sequence_parallel_size", 1))
-                _dlo_group_size = _dp_size if _dp_size > 1 else _sp_size
+                resolved_offload = resolve_offload(self.od_config)
+                distributed_offload = resolved_offload.strategy is OffloadStrategy.DISTRIBUTED_LAYER_WISE
+                dit_distributed_offload = distributed_offload and resolved_offload.offloads(DIT_COMPONENT)
+                dit_uses_allgather = dit_distributed_offload and resolved_offload.uses_allgather(DIT_COMPONENT)
+                tensor_parallel_size = int(getattr(self.parallel_config, "tensor_parallel_size", 1))
+                use_hsdp = bool(getattr(self.parallel_config, "use_hsdp", False))
+                data_parallel_size = int(getattr(self.parallel_config, "data_parallel_size", 1))
+                sequence_parallel_size = int(getattr(self.parallel_config, "sequence_parallel_size", 1))
+                dlo_group_size = data_parallel_size if data_parallel_size > 1 else sequence_parallel_size
+                modules = ModuleDiscovery.discover(model)
+                plan = get_offload_plan(model)
+                selected_encoders = [
+                    encoder
+                    for name, encoder in zip(modules.encoder_names, modules.encoders)
+                    if resolved_offload.offloads(TEXT_ENCODER_COMPONENT)
+                    and encoder_component_type(name, plan) == TEXT_ENCODER_COMPONENT
+                ]
+                allgather_modules: list[nn.Module] = []
+                if dlo_group_size > 1:
+                    if dit_uses_allgather:
+                        allgather_modules.extend(modules.dits)
+                    if (
+                        distributed_offload
+                        and resolved_offload.offloads(TEXT_ENCODER_COMPONENT)
+                        and resolved_offload.uses_allgather(TEXT_ENCODER_COMPONENT)
+                    ):
+                        allgather_modules.extend(selected_encoders)
+                allgather_online_quant = any(self._has_online_quant(module) for module in allgather_modules)
+                if allgather_online_quant:
+                    unsupported_methods = {
+                        method
+                        for module in allgather_modules
+                        for method in self._unsupported_dlo_allgather_online_quant_methods(module)
+                    }
+                    if unsupported_methods:
+                        raise ValueError(
+                            "DLO+AllGather supports online quantization only for "
+                            "per-tensor FP8, INT8, and MXFP8 linears; unsupported "
+                            f"online methods: {', '.join(sorted(unsupported_methods))}. "
+                            "Use rank-local transfer for the affected component or "
+                            "disable online quantization."
+                        )
+                    logger.info(
+                        "Validated online methods (per-tensor FP8, INT8, MXFP8) for every component using DLO+AllGather"
+                    )
 
                 plan_result = None
                 weight_sources = self._get_weight_sources(model)
@@ -562,9 +607,9 @@ class DiffusersPipelineLoader(HWRLoaderMixin):
                     try:
                         hwr_state = self._resolve_hwr(
                             model,
-                            ModuleDiscovery.discover(model),
-                            dist_offload=_dist_offload,
-                            use_allgather=_use_ag,
+                            modules,
+                            dist_offload=dit_distributed_offload,
+                            use_allgather=dit_uses_allgather,
                             load_format=load_format,
                             sources=weight_sources,
                         )
@@ -584,16 +629,15 @@ class DiffusersPipelineLoader(HWRLoaderMixin):
                 hwr_active = hwr_state is not None
                 if hwr_active and hwr_state is not None:
                     self.host_weight_plan = cast(HostWeightPlan | None, hwr_state.get("plan"))
-                if _dist_offload and not hwr_active and not self._force_canonical_load:
-                    modules = ModuleDiscovery.discover(model)
+                if dit_distributed_offload and not hwr_active and not self._force_canonical_load:
                     plan_result = build_checkpoint_mmap_plan(
                         model,
                         dit_modules=tuple(zip(modules.dit_names, modules.dits)),
                         sources=weight_sources,
                         model_path=str(getattr(self.od_config, "model", "")) or None,
-                        tensor_parallel_size=_tp_size,
-                        use_hsdp=_use_hsdp,
-                        online_quantization=_has_online_quant,
+                        tensor_parallel_size=tensor_parallel_size,
+                        use_hsdp=use_hsdp,
+                        online_quantization=any(self._has_online_quant(dit) for dit in modules.dits),
                     )
                     self.host_weight_plan = plan_result.plan
 
@@ -602,7 +646,7 @@ class DiffusersPipelineLoader(HWRLoaderMixin):
                 if host_weight_plan is not None:
                     logger.info(
                         "DLO host-weight plan active (%s, %s): skipping ordinary materialization for %s",
-                        "AllGather" if _use_ag and _dlo_group_size > 1 else "rank-local",
+                        "AllGather" if dit_uses_allgather and dlo_group_size > 1 else "rank-local",
                         host_weight_plan.backing_kind,
                         sorted(host_weight_plan.planned_source_prefixes) or "legacy DiT sources",
                     )
@@ -623,24 +667,7 @@ class DiffusersPipelineLoader(HWRLoaderMixin):
                         )
                     self._maybe_fuse_distilled_lora(model)
                 else:
-                    if _dist_offload and _use_ag and _dlo_group_size > 1 and _has_online_quant:
-                        # An effective DLO group size of one performs no weight
-                        # collective, so the AllGather layout allowlist does not
-                        # apply there even with dlo_use_allgather=True.
-                        unsupported_methods = self._unsupported_dlo_allgather_online_quant_methods(model)
-                        if unsupported_methods:
-                            raise ValueError(
-                                "DLO+AllGather supports online quantization only for "
-                                "per-tensor FP8, INT8, and MXFP8 linears; unsupported online "
-                                f"methods: {', '.join(unsupported_methods)}. Please use "
-                                "--dlo-no-use-allgather or disable online quantization."
-                            )
-                        logger.info(
-                            "Validated online methods (per-tensor FP8, INT8, MXFP8) with "
-                            "DLO+AllGather: using the ordinary loader before sharding "
-                            "finalized weights and scales"
-                        )
-                    if _dist_offload and plan_result is not None:
+                    if dit_distributed_offload and plan_result is not None:
                         logger.info(
                             "DLO direct checkpoint mmap unavailable; using ordinary loader: %s",
                             plan_result.fallback_reason,
@@ -846,14 +873,6 @@ class DiffusersPipelineLoader(HWRLoaderMixin):
                 continue
 
             if has_online_quant:
-                # finalize_layerwise_processing() and the synchronous online
-                # loader already processed these layers.  Avoid moving their
-                # quantized CPU weights back to the accelerator merely to call
-                # an idempotent no-op; doing so rebuilds a large CUDA allocator
-                # cache and defeats streaming CPU offload's startup-memory bound.
-                if getattr(module, "_already_called_process_weights_after_loading", False):
-                    continue
-
                 # Online quant may leave straggler params on the ``meta`` device.
                 # Move only real (non-meta) params onto the target device for
                 # processing and restore them afterward, mirroring upstream vLLM's

vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py

@@ -68,6 +68,7 @@ from vllm_omni.diffusion.models.schedulers.scheduling_flow_match_euler_discrete
 from vllm_omni.diffusion.models.schedulers.scheduling_flow_unipc_multistep import (
     FlowUniPCMultistepScheduler,
 )
+from vllm_omni.diffusion.offloader.config import OffloadStrategy, resolve_offload
 from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
@@ -934,6 +935,12 @@ class Cosmos3OmniDiffusersPipeline(
     ) -> None:
         super().__init__()
         self.od_config = od_config
+        resolved_offload = resolve_offload(od_config)
+        if resolved_offload.public is not None and resolved_offload.strategy is OffloadStrategy.MODEL_LEVEL:
+            raise ValueError(
+                "Cosmos3 model offload uses reasoner/generator topology and does "
+                "not support the dit/text_encoder component selector"
+            )
         self.device = get_local_device()
         self.dtype = od_config.dtype
 
@@ -1086,6 +1093,7 @@ class Cosmos3OmniDiffusersPipeline(
         device: torch.device,
         pin_memory: bool = True,
         use_hsdp: bool = False,
+        offload_components: frozenset[str] | None = None,
     ) -> None:
         """Enable Cosmos3 component-level model offload.
 
@@ -1094,6 +1102,11 @@ class Cosmos3OmniDiffusersPipeline(
         mutual-exclusion swaps.  The VAE stays resident on GPU like the generic
         model-level offloader.
         """
+        if offload_components is not None:
+            raise ValueError(
+                "Cosmos3 model offload uses reasoner/generator topology and does not support "
+                "the dit/text_encoder offload_components selector"
+            )
         self.vae.to(device, non_blocking=True)
         if isinstance(self._sound_tokenizer, nn.Module):
             self._sound_tokenizer.to(device)

vllm_omni/diffusion/models/diffusers_adapter/pipeline_diffusers_adapter.py

@@ -237,6 +237,13 @@ class DiffusersAdapterPipeline(nn.Module, DiffusionPipelineProfilerMixin):
 
     def _raise_unsupported_features(self) -> None:
         """Raise an error for incompatible feature switches."""
+        if self.od_config.diffusion_offload_config is not None:
+            raise NotImplementedError(
+                "diffusion_offload_config is not supported with the diffusers backend. "
+                "Its component selectors cannot be represented by Diffusers' pipeline-wide "
+                "CPU-offload hooks. Use the legacy enable_cpu_offload or "
+                "enable_layerwise_offload option, or use a native pipeline."
+            )
         pc = self.od_config.parallel_config
         if pc.tensor_parallel_size > 1:
             raise NotImplementedError(

vllm_omni/diffusion/models/minimax_h3/encoder.py

@@ -1275,55 +1275,6 @@ class MiniMaxH3Qwen3VLEncoder(nn.Module):
         else:
             cache.release_if_needed()
 
-    def enable_omni_layerwise_offload(self, *, pin_memory: bool = True) -> None:
-        """Stream the TP-local Qwen vision/text blocks for low-HBM serving.
-
-        The encoder has its own TP process group, so these blocks must remain
-        rank-local.  Reusing the DiT AllGather group would concatenate
-        different TP shards and corrupt the encoder weights.
-        """
-        if not self.is_loaded or getattr(self, "_omni_layerwise_enabled", False):
-            return
-
-        from vllm_omni.diffusion.offloader.layerwise_backend import apply_block_hook
-        from vllm_omni.platforms import current_omni_platform
-
-        self._omni_layerwise_hooks = []
-        self._omni_layerwise_block_groups = []
-        copy_stream = current_omni_platform.Stream()
-        for blocks in (self.vision.blocks, self.text_model.layers):
-            if len(blocks) <= 1:
-                continue
-            last_hook = apply_block_hook(
-                blocks[-1],
-                blocks[0],
-                self.device_target,
-                copy_stream,
-                pin_memory,
-            )
-            hooks = [last_hook]
-            for index, block in enumerate(blocks[:-1]):
-                hooks.append(
-                    apply_block_hook(
-                        block,
-                        blocks[index + 1],
-                        self.device_target,
-                        copy_stream,
-                        pin_memory,
-                    )
-                )
-            for index, hook in enumerate(hooks):
-                hook._prev_hook = hooks[index - 1]
-            self._omni_layerwise_hooks.extend(hooks)
-            self._omni_layerwise_block_groups.append(blocks)
-
-        self._omni_layerwise_enabled = bool(self._omni_layerwise_hooks)
-        logger.info(
-            "MiniMax H3 encoder layerwise offload enabled on %d blocks across %d stacks",
-            sum(len(blocks) for blocks in self._omni_layerwise_block_groups),
-            len(self._omni_layerwise_block_groups),
-        )
-
     def _encode(
         self,
         input_ids: torch.Tensor,

vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py

@@ -52,6 +52,13 @@ from vllm_omni.diffusion.offloader import (
     remove_sequential_offload,
     sequential_offload_component,
 )
+from vllm_omni.diffusion.offloader.config import (
+    DIT_COMPONENT,
+    TEXT_ENCODER_COMPONENT,
+    OffloadStrategy,
+    resolve_offload,
+    should_offload_component,
+)
 from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery
 from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import (
     DiffusionPipelineProfilerMixin,
@@ -668,6 +675,7 @@ class MiniMaxH3Pipeline(
     _offload_plan: ClassVar[OffloadPlan] = OffloadPlan(
         offload_submodules={"token_refiner": "blocks"},
         resident_dit_paths=frozenset({"transformer"}),
+        encoder_component_types={"text_encoder": TEXT_ENCODER_COMPONENT},
         encoder_block_attrs={"text_encoder": ("vision.blocks", "text_model.layers")},
         on_demand_component_paths=frozenset({"text_encoder", "video_vae", "audio_vae"}),
     )
@@ -703,10 +711,13 @@ class MiniMaxH3Pipeline(
         self._lora_sigma_schedules.pop(lora_request.lora_int_id, None)
         od_config = getattr(self, "od_config", None)
         offload_modes = []
-        if getattr(od_config, "enable_cpu_offload", False):
-            offload_modes.append("model-level CPU offload (--enable-cpu-offload)")
-        if getattr(od_config, "enable_layerwise_offload", False):
-            offload_modes.append("layerwise offload (--enable-layerwise-offload)")
+        if od_config is not None:
+            resolved_offload = resolve_offload(od_config)
+            if resolved_offload.offloads(DIT_COMPONENT):
+                if resolved_offload.strategy is OffloadStrategy.MODEL_LEVEL:
+                    offload_modes.append("model-level CPU offload")
+                elif resolved_offload.strategy is OffloadStrategy.LAYER_WISE:
+                    offload_modes.append("layerwise offload")
         loaded = load_minimax_h3_turbo_lora(
             partition=self.partition,
             lora_request=lora_request,
@@ -1030,10 +1041,13 @@ class MiniMaxH3Pipeline(
             self.text_encoder_group = None
             self.text_encoder = None
             self._encoder_modules = []
-        stage_components = bool(
+        legacy_manual_components = getattr(od_config, "diffusion_offload_config", None) is None and bool(
             od_config.enable_layerwise_offload or getattr(od_config, "enable_distributed_layerwise_offload", False)
         )
-        component_load_device = torch.device("cpu") if stage_components else self.device
+        # Preserve the legacy MiniMax-H3 low-residency path. The compact API
+        # deliberately limits explicit component selection to dit/text_encoder,
+        # so VAEs stay resident for new configurations.
+        component_load_device = torch.device("cpu") if legacy_manual_components else self.device
         self.video_vae = MiniMaxH3VideoVAE(
             os.path.join(model_path, "video_vae"),
             device=self.device,
@@ -1048,14 +1062,19 @@ class MiniMaxH3Pipeline(
         self.vae = self.video_vae
 
         self._dlo_component_cache = None
-        if getattr(od_config, "enable_distributed_layerwise_offload", False):
+        offloads_text_encoder = should_offload_component(od_config, TEXT_ENCODER_COMPONENT)
+        needs_component_cache = legacy_manual_components or offloads_text_encoder
+        if getattr(od_config, "enable_distributed_layerwise_offload", False) and needs_component_cache:
             self._dlo_component_cache = BoundedAllocatorCache(self.device)
+            if legacy_manual_components:
                 _register_dlo_component_cache(
                     self._dlo_component_cache,
                     self.text_encoder,
                     self.video_vae,
                     self.audio_vae,
                 )
+            elif offloads_text_encoder:
+                _register_dlo_component_cache(self._dlo_component_cache, self.text_encoder)
 
         self._quality_policy = MiniMaxH3QualityPolicy(od_config)
         self._cache_dit_runtime = RequestScopedCacheDiTRuntime(self)
@@ -1519,26 +1538,26 @@ class MiniMaxH3Pipeline(
             # swaps the resident DiT and encoder.
             return self.text_encoder(input_ids, **vision_kwargs)
 
-        if self.od_config.enable_layerwise_offload or getattr(
-            self.od_config, "enable_distributed_layerwise_offload", False
-        ):
-            # Layerwise DiT offload already provides the low-residency encoder
-            # phase used by the checkpoint reference.
+        if self._uses_manual_component_offload(self.text_encoder):
             with self._component_on_device(self.text_encoder):
                 return self.text_encoder.encode_ids(input_ids, **vision_kwargs)
 
-        # Keep both Qwen and DiT resident across requests. Moving either model
-        # here makes encoder latency include a tens-of-gigabytes PCIe transfer,
-        # which defeats the no-offload contract.
+        # Keep Qwen resident when it is not selected for layerwise offload.
         self.text_encoder.load_to_device()
         return self.text_encoder.encode_ids(input_ids, **vision_kwargs)
 
-    def _uses_manual_component_offload(self) -> bool:
+    def _uses_manual_component_offload(self, component: nn.Module) -> bool:
         od_config = getattr(self, "od_config", None)
+        if od_config is None:
+            return False
+        if getattr(od_config, "diffusion_offload_config", None) is None:
             return bool(
                 getattr(od_config, "enable_layerwise_offload", False)
                 or getattr(od_config, "enable_distributed_layerwise_offload", False)
             )
+        return component is getattr(self, "text_encoder", None) and should_offload_component(
+            od_config, TEXT_ENCODER_COMPONENT
+        )
 
     def enable_omni_model_cpu_offload(
         self,
@@ -1546,6 +1565,7 @@ class MiniMaxH3Pipeline(
         device: torch.device,
         pin_memory: bool,
         use_hsdp: bool,
+        offload_components: frozenset[str] | None = None,
     ) -> None:
         if getattr(self, "_model_cpu_offload_modules", None):
             return
@@ -1554,19 +1574,32 @@ class MiniMaxH3Pipeline(
         dits = components.dits
         stages = [*components.encoders, *components.vaes]
         modules = [*dits, *stages]
+        selection_options: dict[str, Any] = {}
+        if offload_components is not None:
+            if DIT_COMPONENT in offload_components and not dits:
+                raise ValueError("MiniMax-H3 has no loaded DiT for selected module offload")
+            if TEXT_ENCODER_COMPONENT in offload_components and not components.encoders:
+                raise ValueError("MiniMax-H3 has no loaded text encoder for selected module offload")
+            selection_options = {
+                "offload_dit_modules": dits if DIT_COMPONENT in offload_components else (),
+                "offload_encoder_modules": (
+                    components.encoders if TEXT_ENCODER_COMPONENT in offload_components else ()
+                ),
+            }
         apply_sequential_offload(
             dit_modules=dits,
             encoder_modules=stages,
             device=device,
             pin_memory=pin_memory,
             use_hsdp=use_hsdp,
-            offload_initial_dits=True,
+            offload_initial_dits=offload_components is None or DIT_COMPONENT in offload_components,
+            **selection_options,
         )
 
         self._model_cpu_offload_modules = modules
         logger.info(
-            "MiniMax-H3 model-level CPU offload enabled for %d DiT(s), text encoder, video VAE, and audio VAE",
-            len(dits),
+            "MiniMax-H3 model-level CPU offload enabled for selected components: %s",
+            sorted(offload_components) if offload_components is not None else "legacy full topology",
         )
 
     def disable_omni_model_cpu_offload(self) -> None:
@@ -1582,7 +1615,7 @@ class MiniMaxH3Pipeline(
             with sequential_offload_component(component):
                 yield
             return
-        staged = self._uses_manual_component_offload()
+        staged = self._uses_manual_component_offload(component)
         try:
             if staged:
                 component.load_to_device()

vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py

@@ -34,6 +34,7 @@ from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin, _is_rank_z
 from vllm_omni.diffusion.models.schedulers import FlowUniPCMultistepScheduler
 from vllm_omni.diffusion.models.wan2_2.scheduling_wan_euler import WanEulerScheduler
 from vllm_omni.diffusion.models.wan2_2.wan2_2_transformer import WanSelfAttention, WanTransformer3DModel
+from vllm_omni.diffusion.offloader import OffloadPlan
 from vllm_omni.diffusion.postprocess import interpolate_video_tensor
 from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
 from vllm_omni.diffusion.request import OmniDiffusionRequest
@@ -308,6 +309,13 @@ def get_wan22_pre_process_func(
     return pre_process_func
 
 
+_WAN_TEXT_ENCODER_OFFLOAD_PLAN = OffloadPlan(
+    encoder_component_types={"text_encoder": "text_encoder"},
+    encoder_block_attrs={"text_encoder": ("encoder.block",)},
+    encoder_dlo_weight_replication=frozenset({"text_encoder"}),
+)
+
+
 class Wan22Pipeline(
     nn.Module,
     PipelineParallelMixin,
@@ -322,6 +330,7 @@ class Wan22Pipeline(
     _dit_modules: ClassVar[list[str]] = ["transformer", "transformer_2"]
     _encoder_modules: ClassVar[list[str]] = ["text_encoder"]
     _vae_modules: ClassVar[list[str]] = ["vae"]
+    _offload_plan = _WAN_TEXT_ENCODER_OFFLOAD_PLAN
 
     def __init__(
         self,

vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py

@@ -32,6 +32,7 @@ from vllm_omni.diffusion.models.interface import SupportImageInput, SupportsComp
 from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin, _is_rank_zero
 from vllm_omni.diffusion.models.utils import _load_json
 from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import (
+    _WAN_TEXT_ENCODER_OFFLOAD_PLAN,
     build_wan_scheduler,
     create_transformer_from_config,
     load_transformer_config,
@@ -192,6 +193,7 @@ class Wan22I2VPipeline(
     _dit_modules: ClassVar[list[str]] = ["transformer", "transformer_2"]
     _encoder_modules: ClassVar[list[str]] = ["text_encoder", "image_encoder"]
     _vae_modules: ClassVar[list[str]] = ["vae"]
+    _offload_plan = _WAN_TEXT_ENCODER_OFFLOAD_PLAN
 
     def __init__(
         self,

vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_s2v.py

@@ -41,6 +41,7 @@ from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import (
 from vllm_omni.diffusion.models.wan2_2.wan2_2_s2v_transformer import (
     create_s2v_transformer_from_config,
 )
+from vllm_omni.diffusion.offloader.config import DIT_COMPONENT, selected_offload_components
 from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch, split_diffusion_output_by_request
@@ -1090,6 +1091,11 @@ class Wan22S2VPipeline(
     # Main forward
     # ------------------------------------------------------------------
 
+    def _should_release_dit_before_decode(self) -> bool:
+        if getattr(self.od_config.parallel_config, "use_hsdp", False):
+            return False
+        return self.od_config.enable_cpu_offload and DIT_COMPONENT in selected_offload_components(self.od_config)
+
     def forward(
         self,
         req: DiffusionRequestBatch,
@@ -1401,7 +1407,7 @@ class Wan22S2VPipeline(
             )
 
             # ---- Decode this clip ----
-            if self.od_config.enable_cpu_offload and not getattr(self.od_config.parallel_config, "use_hsdp", False):
+            if self._should_release_dit_before_decode():
                 self.transformer.to("cpu")
                 current_omni_platform.empty_cache()
 

vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_vace.py

@@ -16,7 +16,7 @@ determined by which inputs are provided (no explicit mode flag):
 from __future__ import annotations
 
 from collections.abc import Iterable
-from dataclasses import replace
+from copy import copy
 
 import PIL.Image
 import torch
@@ -198,7 +198,10 @@ class Wan22VACEPipeline(Wan22Pipeline, SupportImageInput):
     ):
         # VACE defaults to flow_shift=3.0 for 480p (base WAN T2V uses 5.0 for 720p)
         if od_config.flow_shift is None:
-            od_config = replace(od_config, flow_shift=3.0)
+            # Preserve runtime-only config provenance (including compact
+            # offload aliases) instead of rerunning dataclass __post_init__.
+            od_config = copy(od_config)
+            od_config.flow_shift = 3.0
 
         super().__init__(od_config=od_config, prefix=prefix)
 

vllm_omni/diffusion/offloader/__init__.py

@@ -8,8 +8,9 @@ from vllm_omni.diffusion.data import OmniDiffusionConfig
 from vllm_omni.diffusion.model_loader.host_weight_plan import HostWeightPlan
 from vllm_omni.platforms import current_omni_platform
 
-from .base import OffloadBackend, OffloadConfig, OffloadStrategy, SupportsModelCpuOffload
+from .base import OffloadBackend, OffloadConfig, SupportsModelCpuOffload
 from .block_discovery import get_blocks_attr_names, get_blocks_from_dit, set_blocks_attr_names
+from .config import OffloadStrategy
 from .distributed_layerwise_backend import (
     DistributedLayerwiseOffloadBackend,
     DistributedLayerwiseOffloadHook,

vllm_omni/diffusion/offloader/base.py

@@ -2,8 +2,8 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 from abc import ABC, abstractmethod
+from collections.abc import Callable, Iterable
 from dataclasses import dataclass
-from enum import Enum
 from typing import Protocol, runtime_checkable
 
 import torch
@@ -12,9 +12,33 @@ from vllm.logger import init_logger
 
 from vllm_omni.diffusion.data import OmniDiffusionConfig, validate_dlo_host_registration_options
 
+from .component_utils import encoder_component_type
+from .config import (
+    DIT_COMPONENT,
+    TEXT_ENCODER_COMPONENT,
+    DLOTransfer,
+    OffloadStrategy,
+    parse_dlo_transfer,
+    parse_offload_components,
+    resolve_offload,
+)
+from .offload_plan import OffloadPlan
+
 logger = init_logger(__name__)
 
 
+def run_cleanup_steps(steps: Iterable[tuple[str, Callable[[], None]]]) -> BaseException | None:
+    """Run every cleanup step and return the first failure, if any."""
+    first_error: BaseException | None = None
+    for description, step in steps:
+        try:
+            step()
+        except BaseException as exc:
+            logger.exception("Cleanup failed while %s", description)
+            first_error = first_error or exc
+    return first_error
+
+
 @runtime_checkable
 class SupportsModelCpuOffload(Protocol):
     """Pipeline-owned lifecycle for model-level CPU offload.
@@ -30,18 +54,12 @@ class SupportsModelCpuOffload(Protocol):
         device: torch.device,
         pin_memory: bool,
         use_hsdp: bool,
+        offload_components: frozenset[str] | None = None,
     ) -> None: ...
 
     def disable_omni_model_cpu_offload(self) -> None: ...
 
 
-class OffloadStrategy(Enum):
-    NONE = "none"
-    MODEL_LEVEL = "model_level"  # Sequential offloading between DiT and encoders
-    LAYER_WISE = "layer_wise"  # Block-level
-    DISTRIBUTED_LAYER_WISE = "distributed_layer_wise"  # Block-level with DP sharding + H2D/AllGather overlap
-
-
 @dataclass
 class OffloadConfig:
     strategy: OffloadStrategy
@@ -55,14 +73,59 @@ class OffloadConfig:
     # Optional per-worker ceiling for registering an HWR mmap. Zero means no
     # additional ceiling; pin_cpu_memory controls whether registration is tried.
     dlo_host_registration_limit_gib: float = 0.0
+    # ``None`` preserves the model's legacy plan-driven component topology;
+    # a frozenset is an explicit compact-API selection.
+    components: frozenset[str] | None = None
+    dlo_transfers: dict[str, DLOTransfer] | None = None
+
+    def __post_init__(self) -> None:
+        if self.components is not None:
+            self.components = parse_offload_components(self.components)
+        if self.dlo_transfers is None:
+            self.dlo_transfers = {
+                DIT_COMPONENT: DLOTransfer.ALLGATHER if self.dlo_use_allgather else DLOTransfer.RANK_LOCAL,
+                TEXT_ENCODER_COMPONENT: DLOTransfer.RANK_LOCAL,
+            }
+        self.dlo_transfers = parse_dlo_transfer(self.dlo_transfers)
+        # Preserve the old field as the DiT transfer compatibility view.
+        self.dlo_use_allgather = self.uses_allgather(DIT_COMPONENT)
+
+    def offloads(self, component: str) -> bool:
+        # Historically the generic selector covered DiT directly while model
+        # plans opted auxiliary components into their lifecycle.
+        return component == DIT_COMPONENT if self.components is None else component in self.components
+
+    def transfer_for(self, component: str) -> DLOTransfer:
+        if self.dlo_transfers is None:
+            raise RuntimeError("DLO transfers were not initialized")
+        try:
+            return self.dlo_transfers[component]
+        except KeyError as exc:
+            raise ValueError(f"Unknown DLO component {component!r}") from exc
+
+    def uses_allgather(self, component: str) -> bool:
+        return self.transfer_for(component) is DLOTransfer.ALLGATHER
+
+    def offloads_encoder(self, name: str, plan: OffloadPlan | None = None) -> bool:
+        """Return whether the selector covers a discovered encoder path.
+
+        Plans declare non-standard encoder names explicitly. The name-based
+        fallback preserves compatibility with pipelines that predate OffloadPlan.
+        """
+        component = encoder_component_type(name, plan)
+        return component is not None and self.offloads(component)
+
+    def should_offload_encoder(self, name: str, plan: OffloadPlan | None = None) -> bool:
+        """Apply explicit selection while preserving the legacy encoder topology."""
+        return self.components is None or self.offloads_encoder(name, plan)
 
     @classmethod
     def from_od_config(cls, od_config: OmniDiffusionConfig) -> "OffloadConfig":
         """Extract and validate offload settings from OmniDiffusionConfig.
 
-        Enforces mutual exclusion among the three offload strategies.
-        Distributed layer-wise takes the highest priority, then layer-wise,
-        then model-level.
+        ``diffusion_offload_config`` is the canonical public selector. The
+        historical ``enable_*_offload`` booleans remain compatibility aliases;
+        ambiguous combinations fail instead of using silent precedence.
 
         The ``dp_size`` is automatically derived from ``parallel_config`` —
         it is NOT a user-configurable parameter. The distributed layerwise
@@ -74,10 +137,11 @@ class OffloadConfig:
         Returns:
             OffloadConfig with validated settings
         """
-        enable_cpu_offload = getattr(od_config, "enable_cpu_offload", False)
-        enable_layerwise_offload = getattr(od_config, "enable_layerwise_offload", False)
-        enable_distributed_layerwise_offload = getattr(od_config, "enable_distributed_layerwise_offload", False)
-        pin_cpu_memory = getattr(od_config, "pin_cpu_memory", True)
+        resolved = resolve_offload(od_config)
+        strategy = resolved.strategy
+        public_config = resolved.public
+        enable_distributed_layerwise_offload = strategy is OffloadStrategy.DISTRIBUTED_LAYER_WISE
+        pin_cpu_memory = resolved.pin_memory
 
         parallel_config = getattr(od_config, "parallel_config", None)
         use_hsdp = getattr(parallel_config, "use_hsdp", False) if parallel_config else False
@@ -102,61 +166,40 @@ class OffloadConfig:
                 if sp_size and sp_size > 1:
                     dp_size = sp_size
 
-        # Determine strategy (mutual exclusion, distributed layer-wise takes priority)
-        if enable_distributed_layerwise_offload:
-            strategy = OffloadStrategy.DISTRIBUTED_LAYER_WISE
-            if enable_layerwise_offload or enable_cpu_offload:
-                logger.info("Distributed layer-wise offloading takes priority, disabling other offloading strategies.")
-        elif enable_layerwise_offload:
-            strategy = OffloadStrategy.LAYER_WISE
-            if enable_cpu_offload:
-                logger.info(
-                    "Both model-level and layer-wise offloading enabled. "
-                    "Layer-wise takes priority, disabling model-level offloading."
-                )
-        elif enable_cpu_offload:
-            strategy = OffloadStrategy.MODEL_LEVEL
+        if public_config is not None:
+            components: frozenset[str] | None = resolved.components
         else:
-            strategy = OffloadStrategy.NONE
+            components = None
 
-        # With dlo_use_allgather=False, do not add another DP shard. Each rank
-        # streams the tensors produced by the standard loader, which may
-        # already be TP-local shards. This avoids AllGather synchronization
-        # requirements (concurrent requests, dummy run skip).
-        dlo_use_allgather = getattr(od_config, "dlo_use_allgather", True)
-        dlo_resident_layers = int(getattr(od_config, "dlo_resident_layers", 0))
+        dlo_transfers = dict(resolved.transfers)
+        dlo_resident_layers = resolved.resident_layers
+        dit_uses_allgather = resolved.uses_allgather(DIT_COMPONENT)
         dlo_host_registration_limit_gib = validate_dlo_host_registration_options(
             limit_gib=getattr(od_config, "dlo_host_registration_limit_gib", 0.0),
             enable_dlo=enable_distributed_layerwise_offload,
-            use_allgather=dlo_use_allgather,
+            use_allgather=dit_uses_allgather,
             hwr_mode=getattr(od_config, "host_weight_runtime_mode", "disabled"),
         )
-        if dlo_resident_layers < 0:
-            raise ValueError(f"dlo_resident_layers must be >= 0, got {dlo_resident_layers}")
-        if dlo_resident_layers and dlo_use_allgather:
-            raise ValueError(
-                "dlo_resident_layers currently requires --dlo-no-use-allgather so "
-                "resident blocks use weights prepared by the standard TP-aware loader"
-            )
 
-        # If dlo_use_allgather=False, force dp_size=1 (each rank independent)
-        if enable_distributed_layerwise_offload and not dlo_use_allgather:
-            dp_size = 1
+        if enable_distributed_layerwise_offload and all(
+            transfer is DLOTransfer.RANK_LOCAL
+            for component, transfer in dlo_transfers.items()
+            if component in resolved.components
+        ):
             logger.info(
-                "Distributed layerwise offload: dlo_use_allgather=False, "
-                "streaming complete rank-local blocks (no DLO shard or AllGather); "
-                "the backend will select mmap or standard-loader host storage"
+                "Distributed layerwise offload: all selected components use "
+                "rank-local transfer (no DLO shard or AllGather)"
             )
 
         # HSDP already shards parameters into DTensors.  Running distributed
         # layerwise offload on top would shard each to_local() again, producing
         # incorrect reconstruction after AllGather.  Reject this combination.
-        if enable_distributed_layerwise_offload and use_hsdp and dlo_use_allgather:
+        if enable_distributed_layerwise_offload and use_hsdp and resolved.any_allgather:
             raise ValueError(
                 "Distributed layerwise offload with AllGather is incompatible with "
                 "HSDP: HSDP parameters are already sharded DTensors, and the offloader "
-                "would double-shard them. Use --dlo-no-use-allgather (standard-loader "
-                "rank-local weights) or disable HSDP."
+                "would double-shard them. Set weight_transfer='rank-local' for the affected "
+                "component in diffusion_offload_config, or disable HSDP."
             )
 
         return cls(
@@ -164,9 +207,11 @@ class OffloadConfig:
             pin_cpu_memory=pin_cpu_memory,
             use_hsdp=use_hsdp,
             dp_size=dp_size,
-            dlo_use_allgather=dlo_use_allgather,
+            dlo_use_allgather=dit_uses_allgather,
             dlo_resident_layers=dlo_resident_layers,
             dlo_host_registration_limit_gib=dlo_host_registration_limit_gib,
+            components=components,
+            dlo_transfers=dlo_transfers,
         )
 
 

vllm_omni/diffusion/offloader/block_discovery.py

@@ -35,9 +35,13 @@ def set_blocks_attr_names(model: nn.Module, names: list[str]) -> None:
         setattr(model.__class__, "_layerwise_offload_blocks_attrs", names)
 
 
-def get_blocks_from_dit(model: nn.Module) -> tuple[list[str], list[nn.Module]]:
-    """Retrieve blocks and attribute names from provided DiT model."""
-    blocks_attr_names = get_blocks_attr_names(model)
+def get_blocks_from_dit(
+    model: nn.Module,
+    block_attrs: tuple[str, ...] | None = None,
+) -> tuple[list[str], list[nn.Module]]:
+    """Retrieve blocks from an explicit plan or the DiT class metadata."""
+    explicit_attrs = block_attrs is not None
+    blocks_attr_names = list(block_attrs) if explicit_attrs else get_blocks_attr_names(model)
     if not blocks_attr_names:
         logger.warning(
             f"No _layerwise_offload_blocks_attrs defined for {model.__class__.__name__}, "
@@ -49,9 +53,9 @@ def get_blocks_from_dit(model: nn.Module) -> tuple[list[str], list[nn.Module]]:
     for name in blocks_attr_names:
         attr = getattr(model, name, None)
         if attr is None:
+            source = "OffloadPlan.block_attrs" if explicit_attrs else "_layerwise_offload_blocks_attrs"
             raise AttributeError(
-                f"Attribute '{name}' declared in _layerwise_offload_blocks_attrs "
-                f"does not exist on model {model.__class__.__name__}"
+                f"Attribute '{name}' declared in {source} does not exist on model {model.__class__.__name__}"
             )
         try:
             attr_iter = iter(attr)

vllm_omni/diffusion/offloader/component_utils.py

@@ -0,0 +1,221 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Shared component-plan helpers for diffusion offload backends."""
+
+from __future__ import annotations
+
+from collections.abc import Callable, Iterator
+from itertools import chain
+from operator import attrgetter
+from typing import TYPE_CHECKING, Any
+
+import torch
+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 .offload_plan import OffloadPlan
+from .tensor_utils import set_tensor_storage
+
+if TYPE_CHECKING:
+    from .base import OffloadConfig
+    from .module_collector import PipelineModules
+
+logger = init_logger(__name__)
+
+
+def encoder_component_type(name: str, plan: OffloadPlan | None) -> str | None:
+    """Map a discovered encoder path to its public offload component."""
+    declared = None if plan is None else plan.encoder_component_types.get(name)
+    if declared is not None:
+        if declared != TEXT_ENCODER_COMPONENT:
+            raise ValueError(f"OffloadPlan maps encoder {name!r} to unknown component {declared!r}")
+        return declared
+
+    leaf_name = name.rsplit(".", 1)[-1]
+    if leaf_name.startswith(TEXT_ENCODER_COMPONENT) or leaf_name.endswith(TEXT_ENCODER_COMPONENT):
+        return TEXT_ENCODER_COMPONENT
+    return None
+
+
+def get_encoder_block_groups(
+    module: nn.Module,
+    name: str,
+    plan: OffloadPlan | None,
+    *,
+    strict: bool = False,
+) -> list[nn.ModuleList]:
+    """Resolve the streamable block lists declared for one encoder."""
+    if plan is None:
+        return []
+
+    # Some distributed models expose an unloaded stub on ranks where the
+    # component never executes.  That is a valid local no-op even when the
+    # component was selected explicitly.
+    if getattr(module, "is_loaded", True) is False:
+        return []
+
+    groups: list[nn.ModuleList] = []
+    for block_path in plan.encoder_block_attrs.get(name, ()):
+        try:
+            blocks = attrgetter(block_path)(module)
+        except AttributeError:
+            if strict:
+                raise ValueError(f"Encoder offload path {name}.{block_path} was not found") from None
+            logger.warning("Encoder offload path %s.%s was not found", name, block_path)
+            continue
+        if not isinstance(blocks, nn.ModuleList) or len(blocks) <= 1:
+            if strict:
+                raise ValueError(f"Encoder offload path {name}.{block_path} is not a streamable block list")
+            logger.warning("Encoder offload path %s.%s is not a streamable block list", name, block_path)
+            continue
+        groups.append(blocks)
+    return groups
+
+
+def iter_streamable_dits(
+    modules: PipelineModules,
+    config: OffloadConfig,
+    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
+            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)
+
+
+def move_non_block_state_to_device(
+    module: nn.Module,
+    block_groups: list[nn.ModuleList],
+    device: torch.device,
+) -> None:
+    """Keep component state outside streamed block lists resident on device."""
+    block_tensors = {
+        id(tensor)
+        for blocks in block_groups
+        for block in blocks
+        for tensor in chain(block.parameters(), block.buffers())
+    }
+    for tensor in chain(module.parameters(), module.buffers()):
+        if id(tensor) in block_tensors:
+            continue
+        local = tensor.to_local() if isinstance(tensor, DTensor) else tensor
+        if local.device != device:
+            set_tensor_storage(tensor, local.to(device, non_blocking=True))
+
+
+def set_encoder_layerwise_state(
+    module: nn.Module,
+    hooks: list[Any],
+    block_groups: list[nn.ModuleList],
+) -> None:
+    """Publish the backend-neutral state used by encoder stage lifecycles."""
+    module._omni_layerwise_hooks = hooks
+    module._omni_layerwise_block_groups = block_groups
+    module._omni_layerwise_enabled = True
+
+
+def clear_encoder_layerwise_state(module: nn.Module) -> None:
+    """Clear encoder layerwise state after its backend hooks are removed."""
+    module._omni_layerwise_hooks = []
+    module._omni_layerwise_block_groups = []
+    module._omni_layerwise_enabled = False
+
+
+def validate_on_demand_component(module: nn.Module, name: str) -> None:
+    """Require the explicit lifecycle used by pipeline-managed components."""
+    if not callable(getattr(module, "load_to_device", None)) or not callable(getattr(module, "offload_to_cpu", None)):
+        raise ValueError(
+            f"Component {name!r} declares on-demand offload but must implement load_to_device() and offload_to_cpu()"
+        )
+
+
+def prepare_component(
+    module: nn.Module,
+    name: str,
+    *,
+    device: torch.device,
+    stage_on_demand: bool,
+    blockwise: bool,
+    staged_components: list[nn.Module],
+) -> None:
+    """Stage a selected component or keep its non-streamed form resident."""
+    if stage_on_demand:
+        validate_on_demand_component(module, name)
+        getattr(module, "offload_to_cpu")()
+        staged_components.append(module)
+        logger.info("Prepared %s for pipeline-managed staged offload", name)
+    elif not blockwise:
+        module.to(device)
+
+
+def prepare_pipeline_components(
+    modules: PipelineModules,
+    config: OffloadConfig,
+    plan: OffloadPlan | None,
+    *,
+    device: torch.device,
+    staged_components: list[nn.Module],
+    enable_encoder_blocks: Callable[[nn.Module, str, OffloadPlan | None, bool], 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")
+        prepare_component(
+            encoder,
+            name,
+            device=device,
+            stage_on_demand=stage_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
+        prepare_component(
+            vae,
+            name,
+            device=device,
+            stage_on_demand=legacy_staged,
+            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)
+
+    if not config.offloads(DIT_COMPONENT):
+        for dit in modules.dits:
+            dit.to(device)

vllm_omni/diffusion/offloader/config.py

@@ -0,0 +1,429 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Public diffusion CPU-offload configuration helpers."""
+
+from __future__ import annotations
+
+from collections.abc import Collection, Iterator, Mapping
+from dataclasses import dataclass
+from enum import Enum
+from typing import Any, TypeVar
+
+DIT_COMPONENT = "dit"
+TEXT_ENCODER_COMPONENT = "text_encoder"
+OFFLOAD_COMPONENTS = frozenset({DIT_COMPONENT, TEXT_ENCODER_COMPONENT})
+DEFAULT_OFFLOAD_COMPONENTS = frozenset({DIT_COMPONENT})
+
+_KeyT = TypeVar("_KeyT")
+_ValueT = TypeVar("_ValueT")
+
+
+class _FrozenMapping(Mapping[_KeyT, _ValueT]):
+    """Small immutable mapping that remains safe to pickle between workers."""
+
+    def __init__(self, value: Mapping[_KeyT, _ValueT]) -> None:
+        self._items = tuple(value.items())
+
+    def __getitem__(self, key: _KeyT) -> _ValueT:
+        for candidate, value in self._items:
+            if candidate == key:
+                return value
+        raise KeyError(key)
+
+    def __iter__(self) -> Iterator[_KeyT]:
+        return (key for key, _ in self._items)
+
+    def __len__(self) -> int:
+        return len(self._items)
+
+    def __reduce__(self):
+        return type(self), (dict(self._items),)
+
+
+class OffloadMode(str, Enum):
+    """User-facing offload granularity."""
+
+    MODULE = "module"
+    LAYER = "layer"
+
+
+class OffloadStrategy(str, Enum):
+    """Resolved internal backend strategy."""
+
+    NONE = "none"
+    MODEL_LEVEL = "model_level"
+    LAYER_WISE = "layer_wise"
+    DISTRIBUTED_LAYER_WISE = "distributed_layer_wise"
+
+
+class DLOTransfer(str, Enum):
+    """How one component's next block reaches the device."""
+
+    ALLGATHER = "allgather"
+    RANK_LOCAL = "rank-local"
+
+
+@dataclass(frozen=True)
+class LayerOffloadOptions:
+    """Layer-mode settings for one selected diffusion component."""
+
+    weight_transfer: DLOTransfer | None = None
+    resident_layers: int = 0
+
+
+@dataclass(frozen=True)
+class ParsedDiffusionOffloadConfig:
+    """Validated internal form of the raw public offload mapping."""
+
+    mode: OffloadMode
+    components: frozenset[str]
+    layer_options: Mapping[str, LayerOffloadOptions]
+    pin_memory: bool | None = None
+
+    def __post_init__(self) -> None:
+        object.__setattr__(self, "layer_options", _FrozenMapping(self.layer_options))
+
+
+@dataclass(frozen=True)
+class ResolvedOffload:
+    """Canonical offload policy derived once from public or legacy inputs."""
+
+    strategy: OffloadStrategy
+    public: ParsedDiffusionOffloadConfig | None
+    components: frozenset[str]
+    transfers: Mapping[str, DLOTransfer]
+    pin_memory: bool
+    resident_layers: int
+
+    def __post_init__(self) -> None:
+        object.__setattr__(self, "transfers", _FrozenMapping(self.transfers))
+
+    def offloads(self, component: str) -> bool:
+        return component in self.components
+
+    def uses_allgather(self, component: str = DIT_COMPONENT) -> bool:
+        try:
+            return self.transfers[component] is DLOTransfer.ALLGATHER
+        except KeyError as exc:
+            raise ValueError(f"Unknown offload component {component!r}") from exc
+
+    @property
+    def any_allgather(self) -> bool:
+        return self.strategy is OffloadStrategy.DISTRIBUTED_LAYER_WISE and any(
+            self.uses_allgather(component) for component in self.components
+        )
+
+
+_LEGACY_STRATEGY_FIELDS = {
+    "enable_cpu_offload": OffloadStrategy.MODEL_LEVEL,
+    "enable_layerwise_offload": OffloadStrategy.LAYER_WISE,
+    "enable_distributed_layerwise_offload": OffloadStrategy.DISTRIBUTED_LAYER_WISE,
+}
+_LEGACY_STRATEGY_PRIORITY = (
+    OffloadStrategy.DISTRIBUTED_LAYER_WISE,
+    OffloadStrategy.LAYER_WISE,
+    OffloadStrategy.MODEL_LEVEL,
+)
+
+
+_EnumT = TypeVar("_EnumT", bound=Enum)
+
+
+def _parse_enum(value: Any, enum_type: type[_EnumT], label: str) -> _EnumT:
+    if isinstance(value, enum_type):
+        return value
+    if not isinstance(value, str):
+        raise TypeError(f"{label} must be a string, got {type(value).__name__}")
+    try:
+        return enum_type(value)
+    except ValueError as exc:
+        choices = ", ".join(item.value for item in enum_type)
+        raise ValueError(f"Unknown {label} {value!r}; choose from: {choices}") from exc
+
+
+def _validate_string_keys(value: Mapping[Any, Any], label: str) -> None:
+    if any(not isinstance(key, str) for key in value):
+        raise TypeError(f"{label} keys must be strings")
+
+
+def _parse_layer_options(component: str, value: Any) -> LayerOffloadOptions:
+    if not isinstance(value, Mapping):
+        raise TypeError(
+            f"diffusion_offload_config.layer_options[{component!r}] must be a mapping, got {type(value).__name__}"
+        )
+    _validate_string_keys(value, f"diffusion_offload_config.layer_options[{component!r}]")
+    unknown = sorted(set(value) - {"weight_transfer", "resident_layers"})
+    if unknown:
+        raise ValueError(
+            f"Unknown diffusion offload setting(s) for {component}: {', '.join(unknown)}; "
+            "choose from: weight_transfer, resident_layers"
+        )
+    weight_transfer = (
+        _parse_enum(value["weight_transfer"], DLOTransfer, "offload transfer") if "weight_transfer" in value else None
+    )
+    resident_layers = value.get("resident_layers", 0)
+
+    if type(resident_layers) is not int or resident_layers < 0:
+        raise ValueError(f"resident_layers for {component} must be a non-negative integer")
+    parsed = LayerOffloadOptions(
+        weight_transfer=weight_transfer,
+        resident_layers=resident_layers,
+    )
+    if component != DIT_COMPONENT and parsed.resident_layers:
+        raise ValueError("resident_layers currently supports only the 'dit' component")
+    return parsed
+
+
+def parse_diffusion_offload_config(value: Any) -> ParsedDiffusionOffloadConfig | None:
+    """Validate the compact ``diffusion_offload_config`` public API."""
+    if value is None:
+        return None
+    if not isinstance(value, Mapping):
+        raise TypeError(f"diffusion_offload_config must be a mapping, got {type(value).__name__}")
+
+    _validate_string_keys(value, "diffusion_offload_config")
+    unknown = sorted(set(value) - {"mode", "components", "layer_options", "pin_memory"})
+    if unknown:
+        raise ValueError(
+            f"Unknown diffusion_offload_config field(s): {', '.join(unknown)}; "
+            "choose from: mode, components, layer_options, pin_memory"
+        )
+    if "mode" not in value:
+        raise ValueError("diffusion_offload_config requires 'mode'")
+    if "components" not in value:
+        raise ValueError("diffusion_offload_config requires 'components'")
+
+    mode = _parse_enum(value["mode"], OffloadMode, "diffusion offload mode")
+    components = parse_offload_components(value["components"])
+
+    raw_layer_options = value.get("layer_options", {})
+    if not isinstance(raw_layer_options, Mapping):
+        raise TypeError("diffusion_offload_config.layer_options must be a mapping")
+    _validate_string_keys(raw_layer_options, "diffusion_offload_config.layer_options")
+    unselected_options = sorted(set(raw_layer_options) - components)
+    if unselected_options:
+        raise ValueError(
+            "diffusion_offload_config.layer_options requires selecting the same component(s): "
+            + ", ".join(unselected_options)
+        )
+    layer_options = {
+        component: _parse_layer_options(component, raw_layer_options.get(component, {})) for component in components
+    }
+
+    pin_memory = value.get("pin_memory")
+    if pin_memory is not None and type(pin_memory) is not bool:
+        raise TypeError("diffusion_offload_config.pin_memory must be a bool")
+
+    if mode is OffloadMode.MODULE and raw_layer_options:
+        raise ValueError(
+            "diffusion_offload_config.layer_options requires mode='layer'; "
+            f"configured for: {', '.join(sorted(raw_layer_options))}"
+        )
+    dit = layer_options.get(DIT_COMPONENT)
+    if dit is not None and dit.resident_layers and dit.weight_transfer is DLOTransfer.ALLGATHER:
+        raise ValueError("resident_layers requires dit.weight_transfer='rank-local'")
+
+    return ParsedDiffusionOffloadConfig(
+        mode=mode,
+        components=components,
+        layer_options=layer_options,
+        pin_memory=pin_memory,
+    )
+
+
+def _legacy_strategy(config: Any) -> OffloadStrategy:
+    enabled = {strategy for field, strategy in _LEGACY_STRATEGY_FIELDS.items() if bool(getattr(config, field, False))}
+    return next((strategy for strategy in _LEGACY_STRATEGY_PRIORITY if strategy in enabled), OffloadStrategy.NONE)
+
+
+def _public_strategy(public: ParsedDiffusionOffloadConfig) -> OffloadStrategy:
+    if public.mode is OffloadMode.MODULE:
+        return OffloadStrategy.MODEL_LEVEL
+
+    needs_distributed_backend = any(
+        settings.weight_transfer is DLOTransfer.ALLGATHER or settings.resident_layers
+        for settings in public.layer_options.values()
+    )
+    return OffloadStrategy.DISTRIBUTED_LAYER_WISE if needs_distributed_backend else OffloadStrategy.LAYER_WISE
+
+
+def _validate_legacy_layer_options(config: Any, public: ParsedDiffusionOffloadConfig | None) -> None:
+    """Reject ambiguous or invalid compatibility DLO tuning before loading."""
+    resident_layers = getattr(config, "dlo_resident_layers", 0)
+    if type(resident_layers) is not int or resident_layers < 0:
+        raise ValueError(f"dlo_resident_layers must be a non-negative integer, got {resident_layers!r}")
+
+    if public is not None:
+        if getattr(config, "host_weight_runtime_mode", "disabled") != "disabled":
+            raise ValueError("diffusion_offload_config cannot be combined with Host Weight Runtime")
+        conflicting_fields = []
+        if getattr(config, "dlo_use_allgather", True) is not True:
+            conflicting_fields.append("dlo_use_allgather")
+        if resident_layers:
+            conflicting_fields.append("dlo_resident_layers")
+        if float(getattr(config, "dlo_host_registration_limit_gib", 0.0)) > 0:
+            conflicting_fields.append("dlo_host_registration_limit_gib")
+        if conflicting_fields:
+            raise ValueError(
+                "diffusion_offload_config cannot be combined with legacy DLO option(s): "
+                + ", ".join(conflicting_fields)
+            )
+        return
+
+    if resident_layers and bool(getattr(config, "dlo_use_allgather", True)):
+        raise ValueError(
+            "dlo_resident_layers requires the DiT DLO transfer to be rank-local; set dlo_use_allgather=False"
+        )
+
+
+def resolve_offload(config: Any) -> ResolvedOffload:
+    """Resolve and cache the complete offload policy on its owning config."""
+    cached = getattr(config, "_resolved_diffusion_offload", None)
+    if isinstance(cached, ResolvedOffload):
+        return cached
+
+    public = parse_diffusion_offload_config(getattr(config, "diffusion_offload_config", None))
+    _validate_legacy_layer_options(config, public)
+    legacy = _legacy_strategy(config)
+    if public is None:
+        transfers = {
+            DIT_COMPONENT: (
+                DLOTransfer.ALLGATHER if bool(getattr(config, "dlo_use_allgather", True)) else DLOTransfer.RANK_LOCAL
+            ),
+            TEXT_ENCODER_COMPONENT: DLOTransfer.RANK_LOCAL,
+        }
+        resolved = ResolvedOffload(
+            strategy=legacy,
+            public=None,
+            components=DEFAULT_OFFLOAD_COMPONENTS,
+            transfers=transfers,
+            pin_memory=bool(getattr(config, "pin_cpu_memory", True)),
+            resident_layers=int(getattr(config, "dlo_resident_layers", 0)),
+        )
+    else:
+        strategy = _public_strategy(public)
+        if legacy is not OffloadStrategy.NONE and legacy is not strategy:
+            raise ValueError("diffusion_offload_config cannot be combined with legacy enable_*_offload flags")
+        transfers = {component: DLOTransfer.RANK_LOCAL for component in OFFLOAD_COMPONENTS}
+        for component, options in public.layer_options.items():
+            transfers[component] = options.weight_transfer or DLOTransfer.RANK_LOCAL
+        dit_options = public.layer_options.get(DIT_COMPONENT)
+        resolved = ResolvedOffload(
+            strategy=strategy,
+            public=public,
+            components=public.components,
+            transfers=transfers,
+            pin_memory=(
+                public.pin_memory if public.pin_memory is not None else bool(getattr(config, "pin_cpu_memory", True))
+            ),
+            resident_layers=0 if dit_options is None else dit_options.resident_layers,
+        )
+
+    parallel_config = getattr(config, "parallel_config", None)
+    data_parallel_size = int(getattr(parallel_config, "data_parallel_size", 1))
+    cache_backend = getattr(config, "cache_backend", None)
+    if data_parallel_size > 1 and resolved.any_allgather and cache_backend not in (None, "none"):
+        raise ValueError(
+            "Diffusion cache acceleration cannot be combined with DLO AllGather "
+            "across data-parallel ranks because rank-local cache decisions can skip "
+            "different weight collectives. Disable cache_backend or use rank-local transfer."
+        )
+    if (
+        data_parallel_size > 1
+        and bool(getattr(config, "enable_prompt_embed_cache", False))
+        and resolved.offloads(TEXT_ENCODER_COMPONENT)
+        and resolved.uses_allgather(TEXT_ENCODER_COMPONENT)
+    ):
+        raise ValueError(
+            "Prompt embedding cache cannot be combined with text_encoder AllGather "
+            "across data-parallel ranks: rank-local cache hits would skip different "
+            "encoder collectives. Use text_encoder weight_transfer='rank-local' or "
+            "disable enable_prompt_embed_cache."
+        )
+
+    # Actual runtime configs are immutable after construction. Keeping the
+    # derived object off the public dataclass fields also keeps projections and
+    # serialization transport-safe while making hot-path accessors constant-time.
+    try:
+        setattr(config, "_resolved_diffusion_offload", resolved)
+    except (AttributeError, TypeError):
+        pass
+    return resolved
+
+
+def resolve_offload_strategy(config: Any) -> OffloadStrategy:
+    """Resolve the compact config or compatibility boolean entry points."""
+    return resolve_offload(config).strategy
+
+
+def materialize_legacy_offload_flags(config: Any) -> OffloadStrategy:
+    """Keep existing strategy readers working after resolving the compact API."""
+    resolved = resolve_offload(config)
+    for field, field_strategy in _LEGACY_STRATEGY_FIELDS.items():
+        setattr(config, field, resolved.strategy is field_strategy)
+    setattr(config, "pin_cpu_memory", resolved.pin_memory)
+    return resolved.strategy
+
+
+def parse_offload_components(value: Collection[str]) -> frozenset[str]:
+    """Validate an internal component collection."""
+    if isinstance(value, (str, Mapping)) or not isinstance(value, Collection):
+        raise TypeError("diffusion_offload_config.components must be a non-empty list of component names")
+    if any(not isinstance(item, str) for item in value):
+        raise TypeError("offload component entries must be strings")
+    components = list(value)
+    if not components:
+        raise ValueError("offload components must not be empty")
+    unknown = sorted(set(components) - OFFLOAD_COMPONENTS)
+    if unknown:
+        choices = ", ".join(sorted(OFFLOAD_COMPONENTS))
+        raise ValueError(f"Unknown diffusion offload component(s): {', '.join(unknown)}; choose from: {choices}")
+    return frozenset(components)
+
+
+def parse_dlo_transfer(value: Mapping[str, str | DLOTransfer]) -> dict[str, DLOTransfer]:
+    """Resolve an internal per-component transfer mapping."""
+    if not isinstance(value, Mapping):
+        raise TypeError(f"offload transfers must be a mapping, got {type(value).__name__}")
+    _validate_string_keys(value, "offload transfers")
+    missing = sorted(OFFLOAD_COMPONENTS - set(value))
+    if missing:
+        raise ValueError("offload transfers require every component; missing: " + ", ".join(missing))
+    resolved: dict[str, DLOTransfer] = {}
+    for component, raw_transfer in value.items():
+        if component not in OFFLOAD_COMPONENTS:
+            choices = ", ".join(sorted(OFFLOAD_COMPONENTS))
+            raise ValueError(f"Unknown offload transfer component {component!r}; choose from: {choices}")
+        resolved[component] = _parse_enum(raw_transfer, DLOTransfer, "offload transfer")
+    return resolved
+
+
+def component_uses_allgather(config: Any, component: str = DIT_COMPONENT) -> bool:
+    """Return whether one selected component uses AllGather transport."""
+    resolved = resolve_offload(config)
+    if not resolved.offloads(component):
+        raise ValueError(f"Offload component {component!r} is not selected")
+    return resolved.uses_allgather(component)
+
+
+def selected_offload_components(config: Any) -> frozenset[str]:
+    """Resolve selected components while preserving legacy topology defaults."""
+    return resolve_offload(config).components
+
+
+def should_offload_component(config: Any, component: str) -> bool:
+    """Return whether an active layer policy selects ``component``."""
+    if component not in OFFLOAD_COMPONENTS:
+        raise ValueError(f"Unknown offload component: {component}")
+    resolved = resolve_offload(config)
+    if resolved.strategy not in {
+        OffloadStrategy.LAYER_WISE,
+        OffloadStrategy.DISTRIBUTED_LAYER_WISE,
+    }:
+        return False
+    return resolved.offloads(component)
+
+
+def any_selected_component_uses_allgather(config: Any) -> bool:
+    """Return whether an enabled layer backend requires weight collectives."""
+    return resolve_offload(config).any_allgather

vllm_omni/diffusion/offloader/layerwise_backend.py

@@ -7,14 +7,41 @@ from typing import Any
 
 import torch
 from torch import nn
-from torch.distributed.tensor import DTensor
 from vllm.logger import init_logger
 
 from vllm_omni.diffusion.hooks import HookRegistry, ModelHook
 from vllm_omni.platforms import current_omni_platform
 
-from .base import OffloadBackend, OffloadConfig
+from .base import OffloadBackend, OffloadConfig, run_cleanup_steps
+from .block_discovery import (
+    get_blocks_attr_names,
+    get_blocks_from_dit,
+    set_blocks_attr_names,
+)
+from .component_utils import (
+    clear_encoder_layerwise_state,
+    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 .tensor_utils import (
+    clear_block_storage,
+    clear_tensor_storage,
+    describe_tensor_storage,
+    flatten_physical_storage,
+    group_named_tensors_by_dtype,
+    is_materialized_tensor,
+    materialization_probe,
+    module_materialization_probe,
+    restore_tensor_storage,
+    set_tensor_storage,
+    tensor_storage_metadata,
+)
 
 logger = init_logger(__name__)
 
@@ -41,6 +68,7 @@ class LayerwiseOffloadHook(ModelHook):
         device: torch.device,
         stream: current_omni_platform.Stream | None = None,
         pin_memory: bool = True,
+        materialization_probe_tensor: torch.Tensor | None = None,
     ):
         assert isinstance(next_block, nn.Module), "transformer block must be type `torch.nn.Module`"
 
@@ -59,31 +87,7 @@ class LayerwiseOffloadHook(ModelHook):
         self.next_block_buffers: dict[str, torch.Tensor] = {}
         self.dtype_cpu_flattened_weights: dict[torch.dtype, torch.Tensor] = {}
         self.dtype_metadata: dict[torch.dtype, list[dict[str, Any]]] = {}
-
-    @staticmethod
-    def _is_dtensor(t: torch.Tensor) -> bool:
-        return isinstance(t, DTensor)
-
-    @staticmethod
-    def _set_tensor_storage(target: torch.Tensor, value: torch.Tensor) -> None:
-        if LayerwiseOffloadHook._is_dtensor(target):
-            target._local_tensor = value
-        else:
-            target.data = value
-
-    @staticmethod
-    def _make_offload_placeholder(tensor: torch.Tensor) -> torch.Tensor:
-        if LayerwiseOffloadHook._is_dtensor(tensor):
-            local_shape = tuple(tensor.to_local().shape)
-            return torch.empty(local_shape, device="meta", dtype=tensor.dtype)
-        return torch.empty((0,), device=tensor.device, dtype=tensor.dtype)
-
-    @staticmethod
-    def _is_materialized_tensor(t: torch.Tensor) -> bool:
-        if LayerwiseOffloadHook._is_dtensor(t):
-            local_t = t.to_local()
-            return not local_t.is_meta
-        return not t.is_meta and t.data.numel() > 0
+        self._materialization_probe = materialization_probe_tensor
 
     def initialize_hook(self, module: nn.Module) -> nn.Module:
         # This all happen during the hook instance being registered to hook registry;
@@ -92,6 +96,8 @@ class LayerwiseOffloadHook(ModelHook):
 
         self.block_parameters: dict[str, nn.Parameter] = dict(module.named_parameters())
         self.block_buffers: dict[str, torch.Tensor] = dict(module.named_buffers())
+        if self._materialization_probe is None:
+            self._materialization_probe = materialization_probe(self.block_parameters, self.block_buffers)
 
         self.next_block_parameters: dict[str, nn.Parameter] = dict(self.next_block.named_parameters())
         self.next_block_buffers: dict[str, torch.Tensor] = dict(self.next_block.named_buffers())
@@ -100,7 +106,6 @@ class LayerwiseOffloadHook(ModelHook):
         self.dtype_cpu_flattened_weights, self.dtype_metadata = LayerwiseOffloadHook._to_cpu(
             self.next_block_parameters,
             self.next_block_buffers,
-            self.device,
             self.pin_memory,
         )
 
@@ -110,7 +115,6 @@ class LayerwiseOffloadHook(ModelHook):
     def _to_cpu(
         params: dict[str, nn.Parameter],
         bufs: dict[str, torch.Tensor],
-        device: torch.device,
         pin_memory: bool = True,
     ) -> tuple[dict[torch.dtype, torch.Tensor], dict[torch.dtype, list[dict[str, Any]]]]:
         """Helper method to move block parameters and buffers to CPU, flattening by dtype.
@@ -123,85 +127,39 @@ class LayerwiseOffloadHook(ModelHook):
                 flattened CPU tensors by dtype,
                 metadata for reconstruction by dtype
         """
-        dtype_grouped_weights: dict[torch.dtype, dict[str, torch.Tensor]] = {}
         dtype_cpu_flattened_weights: dict[torch.dtype, torch.Tensor] = {}
         # NOTE: order does matter
         dtype_metadata: dict[torch.dtype, list[dict[str, Any]]] = {}
+        targets_to_offload: list[torch.Tensor] = []
 
-        for name, param_or_buf in chain(params.items(), bufs.items()):
-            dtype = param_or_buf.dtype
-            if dtype not in dtype_grouped_weights:
-                dtype_grouped_weights[dtype] = {}
-            dtype_grouped_weights[dtype][name] = param_or_buf
-
-        for dtype, name2weights in dtype_grouped_weights.items():
+        for dtype, named_weights in group_named_tensors_by_dtype(params, bufs).items():
             # total # of parameters + buffers
-            weights_with_local = []
-            for name, t in name2weights.items():
-                local_t = t.to_local() if hasattr(t, "to_local") else t
-                stride = local_t.stride()
-                storage_numel = (
-                    0
-                    if local_t.numel() == 0
-                    else 1 + sum((size - 1) * axis_stride for size, axis_stride in zip(local_t.shape, stride))
-                )
-                weights_with_local.append((name, t, local_t, storage_numel, stride))
-            total_numel = sum(storage_numel for _, _, _, storage_numel, _ in weights_with_local)
+            specs = describe_tensor_storage(named_weights)
+            total_numel = sum(spec.storage_numel for spec in specs)
             cpu_tensor = torch.empty(total_numel, dtype=dtype, device="cpu", pin_memory=pin_memory)
 
             current_offset = 0
-            for (
-                name,
-                original_tensor,
-                local_tensor,
-                storage_numel,
-                stride,
-            ) in weights_with_local:
-                if local_tensor.is_contiguous():
-                    flat_storage = local_tensor.flatten()
-                else:
-                    # Cutlass FP8 weights use a transposed physical layout.
-                    # Preserve it across the flattened CPU staging buffer.
-                    flat_storage = torch.zeros(
-                        storage_numel,
-                        dtype=dtype,
-                        device=local_tensor.device,
+            for spec in specs:
+                flat_storage = flatten_physical_storage(spec.value, spec.storage_numel)
+                cpu_tensor[current_offset : current_offset + spec.storage_numel].copy_(flat_storage)
+                dtype_metadata.setdefault(dtype, []).append(
+                    tensor_storage_metadata(spec, current_offset, include_device=True)
                 )
-                    physical_view = torch.as_strided(
-                        flat_storage,
-                        size=local_tensor.shape,
-                        stride=stride,
-                    )
-                    physical_view.copy_(local_tensor)
-                cpu_tensor[current_offset : current_offset + storage_numel].copy_(flat_storage)
-                if dtype not in dtype_metadata:
-                    dtype_metadata[dtype] = []
-                dtype_metadata[dtype].append(
-                    {
-                        "name": name,
-                        "offset": current_offset,
-                        "numel": storage_numel,
-                        "shape": local_tensor.shape,
-                        "stride": stride,
-                    }
-                )
-
-                LayerwiseOffloadHook._set_tensor_storage(
-                    original_tensor, LayerwiseOffloadHook._make_offload_placeholder(original_tensor)
-                )
-                current_offset += storage_numel
+                targets_to_offload.append(spec.target)
+                current_offset += spec.storage_numel
 
             dtype_cpu_flattened_weights[dtype] = cpu_tensor
 
+        # Do not mutate the module until every host master has been built.
+        # Allocation/copy failures therefore leave the input module intact.
+        clear_tensor_storage(targets_to_offload)
+
         return dtype_cpu_flattened_weights, dtype_metadata
 
     @property
     def is_materialized(self) -> bool:
         """Check whether this block's parameters hold real data on device."""
-        for param in self.block_parameters.values():
-            return LayerwiseOffloadHook._is_materialized_tensor(param)
-
-        return True
+        return self._materialization_probe is None or is_materialized_tensor(self._materialization_probe)
 
     @torch.compiler.disable
     def prefetch_layer(self, non_blocking: bool = True) -> None:
@@ -236,7 +194,7 @@ class LayerwiseOffloadHook(ModelHook):
                     layer_params[target_name] if target_name in layer_params else layer_bufs[target_name]
                 )
 
-                LayerwiseOffloadHook._set_tensor_storage(
+                set_tensor_storage(
                     target_param_or_buf,
                     torch.as_strided(
                         gpu_weight[metadata["offset"] : metadata["offset"] + metadata["numel"]],
@@ -252,17 +210,27 @@ class LayerwiseOffloadHook(ModelHook):
         """Free GPU memory for layer by replacing tensors with empty placeholders.
         This function does not actually offload weights from GPU back to CPU.
         """
-        evt = self._prefetch_done
-        if evt is not None:
-            current_omni_platform.current_stream().wait_event(evt)
-
+        clear_block_storage(self.block_parameters, self.block_buffers, self._prefetch_done)
         self._prefetch_done = None
 
-        # free GPU residency
-        for _, param in self.block_parameters.items():
-            LayerwiseOffloadHook._set_tensor_storage(param, LayerwiseOffloadHook._make_offload_placeholder(param))
-        for _, buf in self.block_buffers.items():
-            LayerwiseOffloadHook._set_tensor_storage(buf, LayerwiseOffloadHook._make_offload_placeholder(buf))
+    @torch.compiler.disable
+    def restore_next_block(self) -> None:
+        """Detach the next block from this hook's host backing store."""
+        for dtype, ordered_metadata in self.dtype_metadata.items():
+            flat = self.dtype_cpu_flattened_weights[dtype]
+            for metadata in ordered_metadata:
+                value = torch.as_strided(
+                    flat[metadata["offset"] : metadata["offset"] + metadata["numel"]],
+                    size=metadata["shape"],
+                    stride=metadata["stride"],
+                )
+                target_name = metadata["name"]
+                target = (
+                    self.next_block_parameters[target_name]
+                    if target_name in self.next_block_parameters
+                    else self.next_block_buffers[target_name]
+                )
+                restore_tensor_storage(target, value, device=metadata["device"])
 
     def pre_forward(self, module: nn.Module, *args: Any, **kwargs: Any) -> tuple[tuple, dict]:
         # if the previous hook was skipped and the weights are not on device,
@@ -287,9 +255,17 @@ def apply_block_hook(
     device: torch.device,
     stream: current_omni_platform.Stream | None = None,
     pin_memory: bool = True,
+    *,
+    materialization_probe_tensor: torch.Tensor | None = None,
 ) -> LayerwiseOffloadHook:
     registry = HookRegistry.get_or_create(module)
-    hook = LayerwiseOffloadHook(next_block, device, stream, pin_memory)
+    hook = LayerwiseOffloadHook(
+        next_block,
+        device,
+        stream,
+        pin_memory,
+        materialization_probe_tensor,
+    )
     registry.register_hook(LayerwiseOffloadHook._HOOK_NAME, hook)
 
     return hook
@@ -302,6 +278,135 @@ def remove_block_hook(module: nn.Module) -> None:
         logger.debug("Removed offload hook from %s", module.__class__.__name__)
 
 
+def _install_layerwise_hook_group(
+    blocks: list[nn.Module] | nn.ModuleList,
+    device: torch.device,
+    stream: Any,
+    pin_memory: bool,
+) -> list[LayerwiseOffloadHook]:
+    """Install one circular hook ring and roll it back transactionally."""
+    block_list = list(blocks)
+    if len(block_list) <= 1:
+        raise ValueError("A layerwise hook group requires at least two blocks")
+
+    hooks: list[LayerwiseOffloadHook] = []
+    hooked_blocks: list[nn.Module] = []
+    probes = {id(block): module_materialization_probe(block) for block in block_list}
+    try:
+        for block, next_block in zip(
+            chain((block_list[-1],), block_list[:-1]),
+            block_list,
+            strict=True,
+        ):
+            hooks.append(
+                apply_block_hook(
+                    block,
+                    next_block,
+                    device,
+                    stream,
+                    pin_memory,
+                    materialization_probe_tensor=probes[id(block)],
+                )
+            )
+            hooked_blocks.append(block)
+    except BaseException:
+        run_cleanup_steps(
+            [
+                *(("restoring a partially installed layerwise block", hook.restore_next_block) for hook in hooks),
+                *(
+                    ("removing a partially installed layerwise hook", lambda block=block: remove_block_hook(block))
+                    for block in hooked_blocks
+                ),
+            ]
+        )
+        raise
+
+    for index, hook in enumerate(hooks):
+        hook._prev_hook = hooks[index - 1]
+    return hooks
+
+
+def enable_plan_encoder_layerwise_offload(
+    module: nn.Module,
+    name: str,
+    plan: OffloadPlan | None,
+    *,
+    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."""
+    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,
+    )
+    if not block_groups:
+        return False
+    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:
+            move_non_block_state_to_device(module, block_groups, device)
+    except BaseException:
+        run_cleanup_steps(
+            [
+                *(("restoring a partially installed encoder block", hook.restore_next_block) for hook in hooks),
+                *(
+                    ("removing a partially installed encoder hook", lambda block=block: remove_block_hook(block))
+                    for block in hooked_blocks
+                ),
+            ]
+        )
+        raise
+    set_encoder_layerwise_state(
+        module,
+        hooks,
+        block_groups,
+    )
+    logger.info(
+        "Enabled rank-local layerwise offload for encoder %s (%d blocks across %d stacks)",
+        name,
+        sum(len(blocks) for blocks in block_groups),
+        len(block_groups),
+    )
+    return True
+
+
+def disable_plan_encoder_layerwise_offload(
+    module: nn.Module,
+    *,
+    restore_weights: bool = True,
+) -> None:
+    """Remove hooks installed by :func:`enable_plan_encoder_layerwise_offload`."""
+    if not getattr(module, "_omni_layerwise_enabled", False):
+        return
+    hooks = getattr(module, "_omni_layerwise_hooks", [])
+    block_groups = getattr(module, "_omni_layerwise_block_groups", [])
+    steps = []
+    if restore_weights:
+        steps.extend(("restoring an encoder block", hook.restore_next_block) for hook in hooks)
+    steps.extend(
+        ("removing an encoder hook", lambda block=block: remove_block_hook(block))
+        for blocks in block_groups
+        for block in blocks
+    )
+    cleanup_error = run_cleanup_steps(steps)
+    if cleanup_error is not None:
+        raise RuntimeError("Failed to fully disable encoder layerwise offload") from cleanup_error
+    clear_encoder_layerwise_state(module)
+
+
 class LayerWiseOffloadBackend(OffloadBackend):
     """Layer-wise (block-level) offloading backend.
 
@@ -315,56 +420,88 @@ class LayerWiseOffloadBackend(OffloadBackend):
 
         self.copy_stream = current_omni_platform.Stream()
         self._blocks: list[list[nn.Module]] = []
+        self._dit_hooks: list[LayerwiseOffloadHook] = []
+        self._hooked_dit_blocks: list[nn.Module] = []
+        self._encoder_modules: list[nn.Module] = []
+        self._staged_components: list[nn.Module] = []
 
     def enable(self, pipeline: nn.Module) -> None:
+        try:
+            self._enable(pipeline)
+        except BaseException:
+            try:
+                self.disable()
+            except BaseException:
+                logger.exception("LayerWiseOffloadBackend cleanup failed while handling an enable failure")
+            raise
+
+    def _enable(self, pipeline: nn.Module) -> None:
         if self.enabled:
             logger.warning("LayerWiseOffloadBackend already enabled")
             return
 
         modules = ModuleDiscovery.discover(pipeline)
-        if not modules.dits:
-            logger.warning("No DiT/transformer modules found, skipping layer-wise offloading")
+        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)
             return
 
-        # Move encoders to GPU (they stay resident)
-        for enc in modules.encoders:
-            enc.to(self.device)
+        def enable_encoder_blocks(
+            module: nn.Module,
+            name: str,
+            component_plan: OffloadPlan | None,
+            stage_on_demand: bool,
+        ) -> bool:
+            enabled = enable_plan_encoder_layerwise_offload(
+                module,
+                name,
+                component_plan,
+                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)
+            return enabled
 
-        # Move VAE(s) to GPU if available
-        for vae in modules.vaes:
-            try:
-                vae.to(self.device, non_blocking=True)
-            except Exception as exc:
-                logger.debug("Failed to move VAE to GPU: %s", exc)
+        prepare_pipeline_components(
+            modules,
+            self.config,
+            plan,
+            device=self.device,
+            staged_components=self._staged_components,
+            enable_encoder_blocks=enable_encoder_blocks,
+        )
 
-        # Move resident modules to GPU (small modules needed every forward)
-        for name, module in zip(modules.resident_names, modules.resident_modules):
-            try:
-                module.to(self.device)
-            except Exception as exc:
-                logger.debug("Failed to move resident module %s to GPU: %s", name, exc)
+        if not self.config.offloads(DIT_COMPONENT):
+            self.enabled = bool(self._encoder_modules or self._staged_components)
+            if not self.enabled:
+                raise ValueError(
+                    "None of the selected layerwise offload components have "
+                    "a model-declared streamable or on-demand plan"
+                )
+            return
 
         logger.info("Applying layer-wise offloading on %s", modules.dit_names)
 
         # 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 i, dit_module in enumerate(modules.dits):
-            dit_name = modules.dit_names[i]
-            logger.info(f"Applying hooks on {dit_name} ({dit_module.__class__.__name__})")
-
-            blocks_attr_names, blocks = LayerWiseOffloadBackend.get_blocks_from_dit(dit_module)
-
-            if not blocks:
-                logger.warning(
-                    "Target layers (blocks) not found. Skipping offloading on %s (%s)",
-                    dit_name,
-                    dit_module.__class__.__name__,
-                )
-                dit_module.to(self.device)
-                continue
-
+        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,
@@ -390,137 +527,66 @@ class LayerWiseOffloadBackend(OffloadBackend):
                 if buffer is not None:
                     buffer.data = buffer.data.to(self.device, non_blocking=True)
 
-            # Pre-fetch the first layer by manually calling the hook function on the last layer;
-            # For subsequent requests, the first layer/block will be pre-fetched
-            # during the last layer compute of the previous request.
-            last_block, first_block = blocks[-1], blocks[0]
-            last_hook = apply_block_hook(
-                last_block,
-                first_block,
+            block_hooks = _install_layerwise_hook_group(
+                blocks,
                 self.device,
                 self.copy_stream,
                 self.config.pin_cpu_memory,
             )
-            last_hook.prefetch_layer(non_blocking=False)
+            self._dit_hooks.extend(block_hooks)
+            self._hooked_dit_blocks.extend(blocks)
 
-            block_hooks: list[LayerwiseOffloadHook] = [last_hook]
-            # Register hook for each of blocks
-            for i, block in enumerate(blocks[:-1]):
-                next_block = blocks[(i + 1) % num_blocks]
-                hook = apply_block_hook(
-                    block,
-                    next_block,
-                    self.device,
-                    self.copy_stream,
-                    self.config.pin_cpu_memory,
-                )
-                block_hooks.append(hook)
-
-            # NOTE(yuanheng-zhao): We make each hook gets a backward reference to the hook
-            # that is responsible for prefetching its block's weights. This is specifically a
-            # workaround for that arbitrary blocks are skipped by caching systems (e.g., cache-dit)
-            for i in range(len(block_hooks)):
-                block_hooks[i]._prev_hook = block_hooks[i - 1]
+            # The last block owns block zero's host backing. Materialize block
+            # 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)")
 
             # Track hooked blocks for cleanup
             self._blocks.append(blocks)
 
-        if len(self._blocks) > 0 and len(self._blocks[0]) > 0:
-            self.enabled = True
+        self.enabled = bool(self._blocks or self._encoder_modules or self._staged_components)
 
-    def disable(self) -> None:
-        if not self.enabled:
+    def _disable(self, *, restore_weights: bool) -> None:
+        if not self.enabled and not (
+            self._dit_hooks or self._hooked_dit_blocks or self._encoder_modules or self._staged_components
+        ):
             return
 
-        for blocks in self._blocks:
-            for block in blocks:
-                remove_block_hook(block)
-
+        steps = []
+        if restore_weights:
+            steps.extend(("restoring a DiT block", hook.restore_next_block) for hook in self._dit_hooks)
+        steps.extend(
+            ("removing a DiT block hook", lambda block=block: remove_block_hook(block))
+            for block in self._hooked_dit_blocks
+        )
+        steps.extend(
+            (
+                "disabling encoder layerwise offload",
+                lambda module=module: disable_plan_encoder_layerwise_offload(
+                    module,
+                    restore_weights=restore_weights,
+                ),
+            )
+            for module in self._encoder_modules
+        )
+        cleanup_error = run_cleanup_steps(steps)
+        if cleanup_error is not None:
+            # Keep every host-master reference so a transient teardown error
+            # can be retried without losing weights behind placeholders.
+            raise RuntimeError("Failed to fully disable layerwise offload") from cleanup_error
         self._blocks.clear()
+        self._dit_hooks.clear()
+        self._hooked_dit_blocks.clear()
+        self._encoder_modules.clear()
+        self._staged_components.clear()
         self.enabled = False
         logger.info("Layer-wise offloading disabled")
 
-    @staticmethod
-    def get_blocks_attr_names(model: nn.Module) -> list[str]:
-        """Get block attribute names from model class."""
-        attrs: list[str] = getattr(model.__class__, "_layerwise_offload_blocks_attrs", [])
+    def disable(self) -> None:
+        self._disable(restore_weights=True)
 
-        if not attrs:
-            old_attr = getattr(model.__class__, "_layerwise_offload_blocks_attr", None)
-            if old_attr is not None:
-                logger.warning(
-                    "'_layerwise_offload_blocks_attr' is deprecated, "
-                    "please use '_layerwise_offload_blocks_attrs' instead. "
-                    "Example: _layerwise_offload_blocks_attrs = ['blocks']"
-                )
-                attrs = [old_attr] if isinstance(old_attr, str) else list(old_attr)
-
-        return attrs
-
-    @staticmethod
-    def set_blocks_attr_names(model: nn.Module, names: list[str]) -> None:
-        if not hasattr(model.__class__, "_layerwise_offload_blocks_attrs"):
-            setattr(model.__class__, "_layerwise_offload_blocks_attrs", names)
-
-    @staticmethod
-    def get_blocks_from_dit(model: nn.Module) -> tuple[list[str], list[nn.Module]]:
-        """
-        Retrieve blocks and attribute names from provided DiT model. Blocks attribute names
-        are found by `_layerwise_offload_blocks_attrs` set to DiT models. For example,
-
-        ```
-        class WanTransformer3DModel(nn.Module):
-            _layerwise_offload_blocks_attrs = ["blocks"]
-        ```
-
-        Returns:
-            Tuple of (blocks_attr_names, blocks)
-        """
-        blocks_attr_names = LayerWiseOffloadBackend.get_blocks_attr_names(model)
-        if not blocks_attr_names:
-            logger.warning(
-                f"No _layerwise_offload_blocks_attrs defined for {model.__class__.__name__}, "
-                "skipping layerwise offloading"
-            )
-            return [], []
-
-        blocks = []
-        for name in blocks_attr_names:
-            attr = getattr(model, name, None)
-            if attr is None:
-                raise AttributeError(
-                    f"Attribute '{name}' declared in _layerwise_offload_blocks_attrs "
-                    f"does not exist on model {model.__class__.__name__}"
-                )
-            try:
-                attr_iter = iter(attr)
-            except TypeError:
-                if isinstance(attr, nn.Module):
-                    logger.warning(
-                        "Attribute '%s' on %s is not iterable; treating it as one block.",
-                        name,
-                        model.__class__.__name__,
-                    )
-                    blocks.append(attr)
-                    continue
-
-                logger.warning(
-                    "Attribute '%s' on %s is not iterable (got %s); skipping it.",
-                    name,
-                    model.__class__.__name__,
-                    type(attr).__name__,
-                )
-            else:
-                blocks.extend(attr_iter)
-
-        if not blocks:
-            logger.warning(
-                "No blocks found in %s for %s, skipping layerwise offloading",
-                blocks_attr_names,
-                model.__class__.__name__,
-            )
-            return [], []
-
-        return blocks_attr_names, blocks
+    # Compatibility aliases for existing model integrations.
+    get_blocks_attr_names = staticmethod(get_blocks_attr_names)
+    set_blocks_attr_names = staticmethod(set_blocks_attr_names)
+    get_blocks_from_dit = staticmethod(get_blocks_from_dit)

vllm_omni/diffusion/offloader/offload_plan.py

@@ -1,10 +1,10 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Declarative OffloadPlan for distributed layerwise offload.
+"""Declarative model capabilities for component layerwise offload.
 
-Models declare this as a class attribute ``_offload_plan`` on the
-pipeline class. When present, the offloader uses it instead of
-heuristic block discovery.
+Models declare an ``_offload_plan`` class attribute on the pipeline. Both the
+ordinary and distributed layerwise backends consume the same plan instead of
+embedding model-specific block paths in backend code.
 """
 
 from __future__ import annotations
@@ -16,12 +16,11 @@ from torch import nn
 
 @dataclass(frozen=True)
 class OffloadPlan:
-    """Optional declarative metadata for distributed layerwise offload.
+    """Optional declarative metadata for component layerwise offload.
 
     Models declare this as a class attribute ``_offload_plan`` on the
-    pipeline class.  When present, the offloader uses it instead of
-    heuristic block discovery, making it easier to support new models
-    without offload-specific logic in constructor or forward code.
+    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.
@@ -39,9 +38,14 @@ class OffloadPlan:
             device when ``dlo_resident_layers`` is nonzero. Keeping this
             model-declared avoids applying a consumer-GPU tuning knob to
             auxiliary or dual DiTs unintentionally.
-        encoder_block_attrs: Maps encoder paths to rank-local block-list paths.
-            These blocks are streamed with ordinary layerwise hooks, never
-            with the DiT AllGather group.
+        encoder_component_types: Maps encoder paths to public selector types
+            (currently text_encoder). This declaration is used before the
+            compatibility name heuristic.
+        encoder_block_attrs: Maps encoder paths to streamable block-list paths.
+        encoder_dlo_weight_replication: Encoder paths whose loader-produced
+            block tensors are identical across the DiT DLO group. Only these
+            encoders may use multi-rank AllGather transfer; this must not be
+            declared for encoder-TP shards.
     """
 
     on_demand_component_paths: frozenset[str] = field(default_factory=frozenset)
@@ -49,7 +53,9 @@ class OffloadPlan:
     block_attrs: dict[str, tuple[str, ...]] = field(default_factory=dict)
     offload_submodules: dict[str, str] = field(default_factory=dict)
     resident_dit_paths: frozenset[str] = field(default_factory=frozenset)
+    encoder_component_types: dict[str, str] = field(default_factory=dict)
     encoder_block_attrs: dict[str, tuple[str, ...]] = field(default_factory=dict)
+    encoder_dlo_weight_replication: frozenset[str] = field(default_factory=frozenset)
 
 
 def get_offload_plan(pipeline: nn.Module) -> OffloadPlan | None:

vllm_omni/diffusion/offloader/sequential_backend.py

@@ -1,7 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-from collections.abc import Iterator
+from collections.abc import Collection, Iterator
 from contextlib import contextmanager
 
 import torch
@@ -13,11 +13,38 @@ 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
 
 logger = init_logger(__name__)
 
 
+def _capture_tensor_devices(modules: Collection[nn.Module]) -> list[tuple[torch.Tensor, torch.device]]:
+    seen: set[int] = set()
+    devices: list[tuple[torch.Tensor, torch.device]] = []
+    for module in modules:
+        for tensor in (*module.parameters(), *module.buffers()):
+            if id(tensor) in seen:
+                continue
+            seen.add(id(tensor))
+            devices.append((tensor, tensor.device))
+    return devices
+
+
+def _restore_tensor_devices(devices: list[tuple[torch.Tensor, torch.device]]) -> None:
+    first_error: BaseException | None = None
+    for tensor, device in devices:
+        try:
+            if tensor.device != device:
+                tensor.data = tensor.data.to(device)
+        except BaseException as exc:
+            logger.exception("Failed to restore tensor placement to %s", device)
+            first_error = first_error or exc
+    if first_error is not None:
+        raise RuntimeError("Failed to restore one or more tensor placements") from first_error
+
+
 class SequentialOffloadHook(ModelHook):
     """Hook for sequential offloading with mutual exclusion on encoder and DiT modules.
 
@@ -34,12 +61,14 @@ class SequentialOffloadHook(ModelHook):
         device: torch.device,
         pin_memory: bool = True,
         use_hsdp: bool = False,
+        offload_after_context: bool = True,
     ):
         # Modules to offload to CPU before this module runs
         self.offload_targets = offload_targets
         self.device = device
         self.pin_memory = pin_memory
         self.use_hsdp = use_hsdp
+        self.offload_after_context = offload_after_context
 
     @staticmethod
     def _move_params(
@@ -48,7 +77,7 @@ class SequentialOffloadHook(ModelHook):
         *,
         non_blocking: bool = False,
         pin_memory: bool = False,
-    ) -> None:
+    ) -> bool:
         """Move module parameters and buffers to device.
 
         This cls method specifically prevents recursion device movement,
@@ -57,47 +86,38 @@ class SequentialOffloadHook(ModelHook):
         refer to
         https://github.com/vipshop/cache-dit/blob/v1.2.3/src/cache_dit/caching/cache_blocks/__init__.py#L83
         """
+        moved = False
         for p in module.parameters():
             if p.data.device != target_device:
                 data = p.data.to(target_device, non_blocking=non_blocking)
                 if pin_memory and target_device.type == "cpu" and not isinstance(data, DTensor):
                     data = data.pin_memory()
                 p.data = data
+                moved = True
         for b in module.buffers():
             if b.device != target_device:
                 data = b.data.to(target_device, non_blocking=non_blocking)
                 if pin_memory and target_device.type == "cpu" and not isinstance(data, DTensor):
                     data = data.pin_memory()
                 b.data = data
+                moved = True
+        return moved
 
     def _to_cpu(self, module: nn.Module) -> None:
-        try:
-            param = next(module.parameters())
-        except StopIteration:
-            return
-
-        if param.device.type == "cpu":
-            return
-
         # XPU's allocator doesn't respect stream dependencies in empty_cache,
         # so non-blocking copies can race with cache eviction. Use blocking
         # copies on XPU to avoid NULL pointer errors during DMA.
         non_blocking = not self.use_hsdp and not current_omni_platform.is_xpu()
-        self._move_params(
+        moved = self._move_params(
             module,
             torch.device("cpu"),
             non_blocking=non_blocking,
             pin_memory=self.pin_memory,
         )
+        if moved:
             current_omni_platform.empty_cache()
 
     def _to_gpu(self, module: nn.Module) -> None:
-        try:
-            if next(module.parameters()).device == self.device:
-                return
-        except StopIteration:
-            return
-
         self._move_params(module, self.device, non_blocking=False)
 
     def pre_forward(self, module: nn.Module, *args, **kwargs) -> tuple[tuple, dict]:
@@ -127,6 +147,8 @@ def apply_sequential_offload(
     pin_memory: bool = True,
     use_hsdp: bool = False,
     offload_initial_dits: bool = False,
+    offload_dit_modules: Collection[nn.Module] | None = None,
+    offload_encoder_modules: Collection[nn.Module] | None = None,
 ) -> None:
     """Apply sequential offloading hooks to DiT and encoder modules.
 
@@ -141,6 +163,10 @@ def apply_sequential_offload(
         pin_memory: Whether to pin CPU memory for faster transfers
         use_hsdp: Whether HSDP is enabled (affects non_blocking behavior)
         offload_initial_dits: Whether to begin with all DiT modules on CPU.
+        offload_dit_modules: DiT modules allowed to move to CPU. None selects
+            every DiT for backward compatibility.
+        offload_encoder_modules: Encoder/stage modules allowed to move to CPU.
+            None selects every supplied module for backward compatibility.
 
     Example:
         >>> apply_sequential_offload(
@@ -150,37 +176,58 @@ def apply_sequential_offload(
         ... )
         >>> # Modules of pipeline now automatically swap between CPU and GPU
     """
-    # Register hooks on DiT modules (offload encoders AND other DiTs when a DiT runs)
+    selected_dit_ids = {id(module) for module in (dit_modules if offload_dit_modules is None else offload_dit_modules)}
+    selected_encoder_ids = {
+        id(module) for module in (encoder_modules if offload_encoder_modules is None else offload_encoder_modules)
+    }
+
+    all_modules = [*dit_modules, *encoder_modules]
+    initial_devices = _capture_tensor_devices(dit_modules) if offload_initial_dits else []
+    try:
+        selected_encoders = [encoder for encoder in encoder_modules if id(encoder) in selected_encoder_ids]
+        # Register hooks on DiT modules (offload selected encoders and other selected DiTs).
         for i, dit_mod in enumerate(dit_modules):
-        other_dits = [d for j, d in enumerate(dit_modules) if j != i]
+            other_dits = [d for j, d in enumerate(dit_modules) if j != i and id(d) in selected_dit_ids]
             registry = HookRegistry.get_or_create(dit_mod)
             hook = SequentialOffloadHook(
-            offload_targets=encoder_modules + other_dits,
+                offload_targets=selected_encoders + other_dits,
                 device=device,
                 pin_memory=pin_memory,
                 use_hsdp=use_hsdp,
+                offload_after_context=id(dit_mod) in selected_dit_ids,
             )
             registry.register_hook(SequentialOffloadHook._HOOK_NAME, hook)
             logger.debug("Registered offload hook for %s", dit_mod.__class__.__name__)
 
-    # Register hooks on encoders (offload DiTs when encoder runs)
+        # Register hooks on all execution stages so unselected resident modules can
+        # still evict selected DiTs before they run.
+        selected_dits = [dit for dit in dit_modules if id(dit) in selected_dit_ids]
         for enc in encoder_modules:
             registry = HookRegistry.get_or_create(enc)
             hook = SequentialOffloadHook(
-            offload_targets=dit_modules,
+                offload_targets=selected_dits,
                 device=device,
                 pin_memory=pin_memory,
                 use_hsdp=use_hsdp,
+                offload_after_context=id(enc) in selected_encoder_ids,
             )
             registry.register_hook(SequentialOffloadHook._HOOK_NAME, hook)
             logger.debug("Registered offload hook for %s", enc.__class__.__name__)
 
         if offload_initial_dits:
-        try:
             for dit_mod in dit_modules:
+                if id(dit_mod) not in selected_dit_ids:
+                    continue
                 _get_sequential_offload_hook(dit_mod)._to_cpu(dit_mod)
-        except Exception:
-            remove_sequential_offload([*dit_modules, *encoder_modules])
+    except BaseException:
+        try:
+            remove_sequential_offload(all_modules)
+        except BaseException:
+            logger.exception("Failed to remove every sequential hook during rollback")
+        try:
+            _restore_tensor_devices(initial_devices)
+        except BaseException:
+            logger.exception("Failed to restore initial DiT placement during rollback")
         raise
 
 
@@ -194,11 +241,18 @@ def remove_sequential_offload(modules: list[nn.Module]) -> None:
         >>> all_modules = [*dit_modules, *encoder_modules]
         >>> remove_sequential_offload(all_modules)
     """
+    first_error: BaseException | None = None
     for module in modules:
+        try:
             registry: HookRegistry | None = getattr(module, "_hook_registry", None)
             if registry is not None:
                 registry.remove_hook(SequentialOffloadHook._HOOK_NAME)
                 logger.debug("Removed offload hook from %s", module.__class__.__name__)
+        except BaseException as exc:
+            logger.exception("Failed to remove offload hook from %s", module.__class__.__name__)
+            first_error = first_error or exc
+    if first_error is not None:
+        raise RuntimeError("Failed to remove one or more sequential offload hooks") from first_error
 
 
 def _get_sequential_offload_hook(module: nn.Module) -> SequentialOffloadHook:
@@ -218,11 +272,13 @@ def sequential_offload_component(module: nn.Module) -> Iterator[None]:
         yield
     except BaseException:
         try:
+            if hook.offload_after_context:
                 hook._to_cpu(module)
         except Exception:
             logger.exception("Failed to release %s after component failure", module.__class__.__name__)
         raise
     else:
+        if hook.offload_after_context:
             hook._to_cpu(module)
 
 
@@ -245,11 +301,26 @@ class ModelLevelOffloadBackend(OffloadBackend):
         # Pipelines with non-forward component entry points own their complete
         # mutual-exclusion lifecycle. Delegate through the explicit protocol.
         if isinstance(pipeline, SupportsModelCpuOffload):
+            try:
+                if self.config.components is not None:
+                    pipeline.enable_omni_model_cpu_offload(
+                        device=self.device,
+                        pin_memory=self.config.pin_cpu_memory,
+                        use_hsdp=self.config.use_hsdp,
+                        offload_components=self.config.components,
+                    )
+                else:
                     pipeline.enable_omni_model_cpu_offload(
                         device=self.device,
                         pin_memory=self.config.pin_cpu_memory,
                         use_hsdp=self.config.use_hsdp,
                     )
+            except BaseException:
+                try:
+                    pipeline.disable_omni_model_cpu_offload()
+                except BaseException:
+                    logger.exception("Model-level cleanup failed while handling an enable failure")
+                raise
             self._custom_pipeline = pipeline
             self.enabled = True
             logger.info(
@@ -259,44 +330,65 @@ 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")
 
-        # Move encoders to GPU
-        for enc in modules.encoders:
-            enc.to(self.device)
-
-        # Move VAE(s) to GPU if available
+        all_modules = [
+            *modules.dits,
+            *modules.encoders,
+            *modules.vaes,
+            *modules.resident_modules,
+        ]
+        initial_devices = _capture_tensor_devices(all_modules)
+        try:
+            for encoder in modules.encoders:
+                encoder.to(self.device)
             for vae in modules.vaes:
-            try:
                 vae.to(self.device, non_blocking=True)
-            except Exception as exc:
-                logger.debug("Failed to move VAE to GPU: %s", exc)
-
-        # Pin resident modules on GPU (small hot submodules called inside the DiT loop).
-        for res, name in zip(modules.resident_modules, modules.resident_names):
-            try:
-                res.to(self.device)
-            except Exception as exc:
-                logger.warning("Failed to move resident module '%s' to GPU: %s", name, exc)
+            for resident in modules.resident_modules:
+                resident.to(self.device)
 
             if not modules.dits:
                 logger.warning("No DiT/transformer modules found, skipping model-level offloading")
                 return
-
             if not modules.encoders:
-            # Nothing to swap against — move DiTs to GPU and skip hooks.
                 for dit in modules.dits:
                     dit.to(self.device)
                 logger.warning("No encoder modules found, skipping model-level offloading")
                 return
 
-        # Apply sequential offloading hooks
             apply_sequential_offload(
                 dit_modules=modules.dits,
                 encoder_modules=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_encoder_modules=selected_encoders,
             )
+        except BaseException:
+            try:
+                remove_sequential_offload([*modules.dits, *modules.encoders])
+            except BaseException:
+                logger.exception("Failed to remove every model-level hook during rollback")
+            try:
+                _restore_tensor_devices(initial_devices)
+            except BaseException:
+                logger.exception("Failed to restore model placement during rollback")
+            raise
 
         # Track modules for cleanup
         self._offload_modules = [*modules.dits, *modules.encoders]

vllm_omni/diffusion/offloader/tensor_utils.py

@@ -9,15 +9,100 @@ offload backends.
 
 from __future__ import annotations
 
+from collections.abc import Callable, Iterable, Mapping
+from dataclasses import dataclass
+from itertools import chain
+from typing import Any
+
 import torch
+from torch import nn
 from torch.distributed.tensor import DTensor
 
 
+@dataclass(frozen=True)
+class TensorStorageSpec:
+    """One tensor's target and physical-storage layout."""
+
+    name: str
+    target: torch.Tensor
+    value: torch.Tensor
+    storage_numel: int
+    stride: tuple[int, ...]
+
+
 def dtype_size(dtype: torch.dtype) -> int:
     """Return element size in bytes for a torch.dtype."""
     return torch.empty(1, dtype=dtype).element_size()
 
 
+def group_named_tensors_by_dtype(
+    params: Mapping[str, torch.Tensor],
+    buffers: Mapping[str, torch.Tensor],
+) -> dict[torch.dtype, list[tuple[str, torch.Tensor]]]:
+    """Group parameters and buffers by dtype while preserving their order."""
+    grouped: dict[torch.dtype, list[tuple[str, torch.Tensor]]] = {}
+    for name, tensor in chain(params.items(), buffers.items()):
+        grouped.setdefault(tensor.dtype, []).append((name, tensor))
+    return grouped
+
+
+def physical_storage_numel(tensor: torch.Tensor) -> int:
+    """Return the storage span needed to preserve ``tensor``'s stride."""
+    if tensor.numel() == 0:
+        return 0
+    return 1 + sum((size - 1) * stride for size, stride in zip(tensor.shape, tensor.stride()))
+
+
+def flatten_physical_storage(
+    tensor: torch.Tensor,
+    storage_numel: int | None = None,
+) -> torch.Tensor:
+    """Pack a possibly strided tensor in physical storage order."""
+    if tensor.is_contiguous():
+        return tensor.flatten()
+    storage_numel = physical_storage_numel(tensor) if storage_numel is None else storage_numel
+    storage = torch.zeros(storage_numel, dtype=tensor.dtype, device=tensor.device)
+    torch.as_strided(storage, size=tensor.shape, stride=tensor.stride()).copy_(tensor)
+    return storage
+
+
+def describe_tensor_storage(
+    named_tensors: Iterable[tuple[str, torch.Tensor]],
+    transforms: Mapping[int, Callable[[torch.Tensor], torch.Tensor]] | None = None,
+) -> list[TensorStorageSpec]:
+    """Resolve local tensors and their transport-preserving layouts."""
+    specs: list[TensorStorageSpec] = []
+    for name, target in named_tensors:
+        value = target.to_local() if hasattr(target, "to_local") else target
+        transform = (transforms or {}).get(id(target))
+        if callable(transform):
+            value = transform(value)
+        specs.append(
+            TensorStorageSpec(
+                name=name,
+                target=target,
+                value=value,
+                storage_numel=physical_storage_numel(value),
+                stride=value.stride(),
+            )
+        )
+    return specs
+
+
+def tensor_storage_metadata(spec: TensorStorageSpec, offset: int, *, include_device: bool = False) -> dict[str, Any]:
+    """Build reconstruction metadata shared by offload transports."""
+    metadata: dict[str, Any] = {
+        "name": spec.name,
+        "offset": offset,
+        "numel": spec.storage_numel,
+        "shape": spec.value.shape,
+        "stride": spec.stride,
+    }
+    if include_device:
+        metadata["device"] = spec.value.device
+    return metadata
+
+
 def is_dtensor(t: torch.Tensor) -> bool:
     """Check if tensor is a DTensor."""
     return isinstance(t, DTensor)
@@ -39,9 +124,74 @@ def make_offload_placeholder(tensor: torch.Tensor) -> torch.Tensor:
     return torch.empty((0,), device=tensor.device, dtype=tensor.dtype)
 
 
+def clear_tensor_storage(tensors: Iterable[torch.Tensor]) -> None:
+    """Atomically replace tensor storage with offload placeholders."""
+    targets = list(tensors)
+    originals = [tensor.to_local() if is_dtensor(tensor) else tensor.data for tensor in targets]
+    placeholders = [make_offload_placeholder(tensor) for tensor in targets]
+    committed = 0
+    try:
+        for target, placeholder in zip(targets, placeholders, strict=True):
+            set_tensor_storage(target, placeholder)
+            committed += 1
+    except BaseException:
+        rollback_error: BaseException | None = None
+        for target, original in zip(targets[:committed], originals[:committed], strict=True):
+            try:
+                set_tensor_storage(target, original)
+            except BaseException as exc:
+                rollback_error = rollback_error or exc
+        if rollback_error is not None:
+            raise RuntimeError("Tensor storage clear failed and rollback was incomplete") from rollback_error
+        raise
+
+
+def clear_block_storage(
+    params: Mapping[str, torch.Tensor],
+    buffers: Mapping[str, torch.Tensor],
+    ready_event: Any | None,
+) -> None:
+    """Wait for an outstanding transfer, then release one block's residency."""
+    if ready_event is not None:
+        from vllm_omni.platforms import current_omni_platform
+
+        current_omni_platform.current_stream().wait_event(ready_event)
+    clear_tensor_storage(chain(params.values(), buffers.values()))
+
+
 def is_materialized_tensor(t: torch.Tensor) -> bool:
     """Check if tensor holds real data (not meta or empty placeholder)."""
     if is_dtensor(t):
         local_t = t.to_local()
         return not local_t.is_meta
     return not t.is_meta and t.data.numel() > 0
+
+
+def materialization_probe(
+    params: Mapping[str, torch.Tensor],
+    buffers: Mapping[str, torch.Tensor],
+) -> torch.Tensor | None:
+    """Pick one non-empty tensor that tracks a block's atomic storage state."""
+    return next((tensor for tensor in chain(params.values(), buffers.values()) if tensor.numel()), None)
+
+
+def module_materialization_probe(module: nn.Module) -> torch.Tensor | None:
+    """Capture a probe before another hook releases the module's storage."""
+    return materialization_probe(dict(module.named_parameters()), dict(module.named_buffers()))
+
+
+def restore_tensor_storage(
+    target: torch.Tensor,
+    value: torch.Tensor,
+    *,
+    device: torch.device | str,
+) -> None:
+    """Detach ``target`` from an offloader-owned backing with a real copy."""
+    restored = torch.empty_strided(
+        value.shape,
+        value.stride(),
+        dtype=value.dtype,
+        device=device,
+    )
+    restored.copy_(value)
+    set_tensor_storage(target, restored)

vllm_omni/diffusion/worker/diffusion_model_runner.py

@@ -51,6 +51,7 @@ from vllm_omni.diffusion.models.interface import (
     supports_step_execution,
 )
 from vllm_omni.diffusion.offloader import enable_offload_backend
+from vllm_omni.diffusion.offloader.config import TEXT_ENCODER_COMPONENT, resolve_offload
 from vllm_omni.diffusion.registry import _NO_CACHE_ACCELERATION
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.sched.interface import (
@@ -279,6 +280,26 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
         if load_format == "dummy":
             return
 
+        # Resolve environment overrides before model loading. Rank-local
+        # prompt-cache hits cannot safely skip text-encoder collectives.
+        enable_pec, pec_size = resolve_prompt_embed_cache_config(
+            enable=getattr(self.od_config, "enable_prompt_embed_cache", False),
+            max_size=getattr(self.od_config, "prompt_embed_cache_size", 32),
+        )
+        resolved_offload = resolve_offload(self.od_config)
+        dp_size = int(getattr(getattr(self.od_config, "parallel_config", None), "data_parallel_size", 1))
+        if (
+            enable_pec
+            and dp_size > 1
+            and resolved_offload.offloads(TEXT_ENCODER_COMPONENT)
+            and resolved_offload.uses_allgather(TEXT_ENCODER_COMPONENT)
+        ):
+            raise ValueError(
+                "Prompt embedding cache cannot be combined with text_encoder "
+                "AllGather across data-parallel ranks; disable the cache or use "
+                "rank-local text_encoder transfer."
+            )
+
         current_omni_platform.init_diffusion_model_runner_runtime(
             vllm_config=self.vllm_config,
             od_config=self.od_config,
@@ -398,10 +419,6 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
         # Install prompt-embedding cache (transparent wrapper around
         # ``pipeline.encode_prompt``). Enabled via config or env var; a no-op
         # when the pipeline does not expose ``encode_prompt``.
-        enable_pec, pec_size = resolve_prompt_embed_cache_config(
-            enable=getattr(self.od_config, "enable_prompt_embed_cache", False),
-            max_size=getattr(self.od_config, "prompt_embed_cache_size", 32),
-        )
         if enable_pec:
             self.prompt_embed_cache = install_prompt_embed_cache(
                 self.pipeline,

vllm_omni/engine/arg_utils.py

@@ -572,6 +572,9 @@ class OrchestratorArgs:
     vae_use_tiling: bool = False
     enable_multithread_weight_load: bool = True
     num_weight_load_threads: int = 4
+    diffusion_offload_config: dict[str, Any] | None = None
+    # Compatibility aliases for existing callers and model-specific stage
+    # lifecycles that are broader than the compact dit/text_encoder selector.
     enable_cpu_offload: bool = False
     enable_layerwise_offload: bool = False
     enable_distributed_layerwise_offload: bool = False

vllm_omni/engine/async_omni_engine.py

@@ -1137,6 +1137,7 @@ class AsyncOmniEngine:
             "cache_backend": cache_backend,
             "cache_config": cache_config,
             "enable_cache_dit_summary": kwargs.get("enable_cache_dit_summary", False),
+            "diffusion_offload_config": kwargs.get("diffusion_offload_config", None),
             "enable_cpu_offload": kwargs.get("enable_cpu_offload", False),
             "enable_layerwise_offload": kwargs.get("enable_layerwise_offload", False),
             "enable_distributed_layerwise_offload": kwargs.get("enable_distributed_layerwise_offload", False),

vllm_omni/entrypoints/cli/serve.py

@@ -703,30 +703,41 @@ class OmniServeCommand(CLISubcommand):
         )
 
         # diffusion model offload parameters
+        omni_config_group.add_argument(
+            "--diffusion-offload-config",
+            type=json.loads,
+            default=None,
+            help="Diffusion CPU-offload config as JSON. "
+            "Set mode to module or layer, list dit and/or text_encoder in "
+            "components, and put layer-only tuning under layer_options. "
+            "Layer settings are weight_transfer (rank-local or allgather) and "
+            "resident_layers (DiT only).",
+        )
         omni_config_group.add_argument(
             "--enable-cpu-offload",
             action="store_true",
-            help="Enable CPU offloading for diffusion models.",
+            help="Compatibility alias for model-level CPU offload. New integrations should use "
+            "--diffusion-offload-config with mode=module and explicit components.",
         )
         omni_config_group.add_argument(
             "--enable-layerwise-offload",
             action="store_true",
-            help="Enable layerwise (blockwise) offloading on DiT modules.",
+            help="Compatibility alias for layerwise CPU offload. New integrations should use "
+            "--diffusion-offload-config with mode=layer and explicit components.",
         )
         omni_config_group.add_argument(
             "--enable-distributed-layerwise-offload",
             action="store_true",
-            help="Enable distributed layerwise offloading with H2D + AllGather overlap. "
-            "Shards weights across DP ranks, stores only 1/DP_size on each host, "
-            "and overlaps H2D transfers and AllGather with computation. "
-            "DP size is automatically derived from the parallel configuration.",
+            help="Compatibility alias for distributed layerwise CPU offload. "
+            "New integrations should use mode=layer and configure weight transfer per component.",
         )
         omni_config_group.add_argument(
             "--dlo-use-allgather",
             dest="dlo_use_allgather",
             action="store_true",
             default=True,
-            help="Use shard + AllGather for weight reconstruction (default: True). "
+            help="Compatibility option; use component weight_transfer=allgather in new configurations. "
+            "Use shard + AllGather for weight reconstruction (default: True). "
             "When disabled (--dlo-no-use-allgather), each rank streams the "
             "standard loader's rank-local tensors via H2D only — no additional "
             "DP sharding, no AllGather, and no concurrent-request requirement.",
@@ -736,6 +747,7 @@ class OmniServeCommand(CLISubcommand):
             dest="dlo_use_allgather",
             action="store_false",
             help=(
+                "Compatibility option; use component weight_transfer=rank-local in new configurations. "
                 "Disable AllGather and stream standard-loader rank-local weights "
                 "independently (including existing TP shards)."
             ),
@@ -744,8 +756,7 @@ class OmniServeCommand(CLISubcommand):
             "--dlo-resident-layers",
             type=int,
             default=0,
-            help="Keep this many leading main-DiT blocks resident on the device "
-            "while distributed layerwise offload streams the remaining blocks.",
+            help="Compatibility option; use layer_options.dit.resident_layers in new configurations.",
         )
         omni_config_group.add_argument(
             "--host-weight-runtime-mode",