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