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

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

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.