[Performance] Add per-component DLO offload policy

Signed-off-by: hsliuustc0106 <liuhongsheng4@huawei.com>

docs/design/feature/offloader/distributed_layerwise_offload.md

@@ -29,6 +29,7 @@ Legend: ✅ supported, ⚠️ compatibility path or limited validation, ❌ unsu
 | **Other online quantization methods** | ❌ Rejected until runtime packing and scale layouts are validated. | ⚠️ Allowed through the ordinary loader; validation is method-specific. |
 | **Model-level or standard layerwise CPU offload** | ❌ Disabled because DLO takes priority. | ❌ Disabled because DLO takes priority. |
 | **Resident leading layers** | ❌ Rejected. | ✅ Requires eligible resident paths in the model's `OffloadPlan`. |
+| **Auxiliary component policy** | ✅ Encoders/VAEs may remain resident. | ✅ Encoders/VAEs may remain resident. |
 
 See [Parallelism compatibility](#parallelism-compatibility) and
 [Request and loading constraints](#request-and-loading-constraints) for the
@@ -419,6 +420,22 @@ failure.
 - **HSDP + DP or TP:** rejected independently by the diffusion parallel
   configuration.
 
+### Auxiliary component policy
+
+`dlo_offload_components` maps discovered encoder/VAE paths to booleans. A
+component set to `False` is materialized on the target device once and is not
+entered through the model's manual offload lifecycle. A `True` value permits
+the existing `OffloadPlan.on_demand_component_paths` and encoder block-hook
+behavior; it does not add support where the model declares none. The optional
+`default` key handles unmatched auxiliary components, and an empty map
+preserves the current policy.
+
+The policy is resolved after component discovery and before DLO mutates DiT
+storage. Unknown names, including DiT paths, are rejected. DLO continues to
+own DiT streaming regardless of this map, which avoids an invalid state where
+the loader has skipped ordinary DiT materialization but the backend is asked
+to keep that DiT resident.
+
 ## Request and loading constraints
 
 AllGather DP multi-concurrency requires:
@@ -430,6 +447,15 @@ AllGather DP multi-concurrency requires:
 The no-AllGather path does not impose these DLO-specific synchronized-wave
 requirements.
 
+MiniMax-H3 currently builds one global encoder TP group and broadcasts the
+resulting conditioning from global rank 0. Therefore an encoder TP group that
+spans multiple DP replicas is not valid for independent-prompt AllGather
+waves. Exact-recipe DP2/TP2 and DP2/SP2 experiments with encoder TP4 completed
+the collective schedule, but two distinct prompts with the same seed produced
+byte-identical video and audio in every measured wave. Enabling or disabling
+encoder offload produced paired-identical results, so this is a pre-existing
+model topology limitation rather than component-policy behavior.
+
 Direct checkpoint mmap can back either transfer path. It is currently limited
 to proven TP1, non-HSDP, non-online-quantized layouts. Other layouts use the
 ordinary loader. Per-tensor online FP8 linears can use DLO AllGather after the
@@ -494,6 +520,9 @@ AllGather, because DP replicas did not share checkpoint-backed runtime
 weights. This is a functional and memory smoke test, not a production-quality
 performance or output-quality benchmark.
 
+These B300 DP waves intentionally reused the same prompt and seed, so they do
+not establish independent-prompt isolation across DP replicas.
+
 ### Host-memory measurement
 
 A two-worker MiniMax-H3 FL2VA measurement on one L20X node compared the

docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md

@@ -36,7 +36,8 @@ vllm serve /path/to/model --omni \
 vllm serve /path/to/model --omni \
   --enable-distributed-layerwise-offload \
   --data-parallel-size 4 \
-  --dlo-no-use-allgather
+  --dlo-no-use-allgather \
+  --dlo-offload-components '{"text_encoder": false}'
 
 # Sequence parallel deployment
 vllm serve /path/to/model --omni \
@@ -51,6 +52,7 @@ omni = Omni(
     model="/path/to/model",
     enable_distributed_layerwise_offload=True,
     dlo_use_allgather=True,
+    dlo_offload_components={"text_encoder": False},
 )
 ```
 
@@ -63,6 +65,7 @@ omni = Omni(
 | `--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` |
+| `--dlo-offload-components JSON` | Permit or disable offload for discovered encoder/VAE components | `{}` |
 | `--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` |
@@ -244,6 +247,25 @@ class MyPipeline(nn.Module):
 When no plan exists, discovery falls back to
 `_layerwise_offload_blocks_attrs` and then heuristic attribute lookup.
 
+## Auxiliary component policy
+
+`dlo_offload_components` follows the same component-map style as diffusion
+quantization. Keys are discovered encoder or VAE paths, values are booleans,
+and `default` is the optional fallback:
+
+```python
+dlo_offload_components = {
+    "text_encoder": False,
+    "default": True,
+}
+```
+
+`True` permits the model's existing DLO behavior for that component. `False`
+keeps the whole component resident, including its child blocks and non-block
+modules. An empty map preserves current behavior. Unknown keys are rejected;
+the map does not control DiT blocks because enabling DLO gives the backend
+ownership of DiT streaming.
+
 ## Data-parallel concurrency
 
 With `data_parallel_size > 1` and AllGather enabled, the scheduler can process
@@ -268,6 +290,13 @@ must enter each collective.
   are validated.
 - Resident leading layers require `--dlo-no-use-allgather` and a model
   `OffloadPlan` that declares eligible `resident_dit_paths`.
+- Component policy applies only to discovered auxiliary encoders and VAEs;
+  it cannot disable DLO for a DiT.
+- MiniMax-H3's encoder TP group is global rather than DP-replica-local. An
+  encoder TP size that spans multiple DP replicas (for example TP4 with
+  DP2/TP2 or DP2/SP2 on four ranks) is not supported for independent-prompt
+  AllGather waves: the current pipeline broadcasts rank 0's conditioning to
+  both replicas. This limitation is independent of the component policy.
 - DP concurrency requires an explicit, identical inference-step count.
 
 Sharing quantized or otherwise unvalidated transformed layouts through a

examples/offline_inference/minimax_h3/dlo_lifecycle.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Validate MiniMax-H3 request-mode and DLO lifecycle behavior.
 
 The DLO modes accept any positive data-parallel size.  With AllGather enabled,
@@ -85,12 +85,20 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--height", type=int, default=768)
     parser.add_argument("--batch-wait-ms", type=float, default=500.0)
     parser.add_argument("--init-timeout", type=float, default=1800.0)
+    parser.add_argument(
+        "--dlo-offload-components",
+        type=json.loads,
+        default=None,
+        help="Component policy JSON, e.g. '{\"text_encoder\": false}'",
+    )
     parser.add_argument("--output", type=Path)
     return parser.parse_args()
 
 
 def engine_kwargs(args: argparse.Namespace) -> dict[str, Any]:
     is_dlo = args.mode in {"dlo", "dlo-no-allgather"}
+    if args.dlo_offload_components and not is_dlo:
+        raise ValueError("--dlo-offload-components requires a DLO mode")
     common: dict[str, Any] = {
         "model": args.model,
         "trust_remote_code": True,
@@ -114,6 +122,7 @@ def engine_kwargs(args: argparse.Namespace) -> dict[str, Any]:
             enable_distributed_layerwise_offload=True,
             dlo_use_allgather=args.mode == "dlo",
             dlo_resident_layers=0,
+            dlo_offload_components=({} if args.dlo_offload_components is None else args.dlo_offload_components),
         )
     else:
         common.update(

tests/diffusion/models/minimax_h3/test_minimax_h3_dlo_lifecycle.py

@@ -99,3 +99,22 @@ def test_manual_component_offload_failure_forces_retained_cache_release(mocker):
     assert component.load_to_device.call_count == 2
     assert component.offload_to_cpu.call_count == 2
     pipeline._dlo_component_cache.release_if_needed.assert_called_once_with(force=True)
+
+
+def test_resident_component_skips_manual_dlo_lifecycle(mocker):
+    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+
+    pipeline = object.__new__(MiniMaxH3Pipeline)
+    torch.nn.Module.__init__(pipeline)
+    pipeline.od_config = mocker.Mock()
+    pipeline.od_config.enable_layerwise_offload = False
+    pipeline.od_config.enable_distributed_layerwise_offload = True
+    pipeline._model_cpu_offload_modules = []
+    component = mocker.Mock()
+    component._omni_dlo_offload_enabled = False
+
+    with pipeline._component_on_device(component):
+        pass
+
+    component.load_to_device.assert_not_called()
+    component.offload_to_cpu.assert_not_called()

tests/diffusion/offloader/test_component_policy.py

@@ -0,0 +1,142 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+from torch import nn
+
+from vllm_omni.diffusion.offloader.base import OffloadConfig, OffloadStrategy
+from vllm_omni.diffusion.offloader.distributed_layerwise_backend import (
+    DistributedLayerwiseOffloadBackend,
+)
+from vllm_omni.diffusion.offloader.module_collector import PipelineModules
+from vllm_omni.diffusion.offloader.offload_plan import OffloadPlan
+
+pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
+
+
+def _od_config(**overrides):
+    values = {
+        "enable_cpu_offload": False,
+        "enable_layerwise_offload": False,
+        "enable_distributed_layerwise_offload": True,
+        "dlo_use_allgather": False,
+        "dlo_resident_layers": 0,
+        "dlo_offload_components": {},
+        "pin_cpu_memory": False,
+        "parallel_config": SimpleNamespace(
+            use_hsdp=False,
+            data_parallel_size=1,
+            sequence_parallel_size=1,
+        ),
+    }
+    values.update(overrides)
+    return SimpleNamespace(**values)
+
+
+def _backend(policy: dict[str, bool]) -> DistributedLayerwiseOffloadBackend:
+    backend = object.__new__(DistributedLayerwiseOffloadBackend)
+    backend.config = OffloadConfig(
+        strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
+        pin_cpu_memory=False,
+        dlo_use_allgather=False,
+        dlo_offload_components=policy,
+    )
+    backend.device = torch.device("cpu")
+    return backend
+
+
+def _modules() -> PipelineModules:
+    return PipelineModules(
+        dits=[],
+        dit_names=[],
+        encoders=[nn.Module()],
+        encoder_names=["text_encoder"],
+        vaes=[nn.Module(), nn.Module()],
+        vae_names=["video_vae", "audio_vae"],
+    )
+
+
+def test_component_policy_resolves_exact_then_default():
+    backend = _backend({"text_encoder": False, "default": True})
+
+    assert not backend._component_offload_enabled("text_encoder")
+    assert backend._component_offload_enabled("video_vae")
+
+
+@pytest.mark.parametrize(
+    ("value", "error", "match"),
+    [
+        ([], TypeError, "must be a mapping"),
+        (["text_encoder"], TypeError, "must be a mapping"),
+        (False, TypeError, "must be a mapping"),
+        ({1: False}, TypeError, "keys must be non-empty strings"),
+        ({"text_encoder": 0}, TypeError, "must be a boolean"),
+    ],
+)
+def test_component_policy_rejects_invalid_shapes(value, error, match):
+    with pytest.raises(error, match=match):
+        OffloadConfig.from_od_config(_od_config(dlo_offload_components=value))
+
+
+def test_component_policy_accepts_none_as_default():
+    config = OffloadConfig.from_od_config(_od_config(dlo_offload_components=None))
+
+    assert config.dlo_offload_components == {}
+
+
+def test_component_policy_requires_dlo():
+    with pytest.raises(ValueError, match="requires distributed layerwise offload"):
+        OffloadConfig.from_od_config(
+            _od_config(
+                enable_distributed_layerwise_offload=False,
+                dlo_offload_components={"text_encoder": False},
+            )
+        )
+
+
+def test_component_policy_rejects_unknown_or_dit_names():
+    backend = _backend({"transformer": False})
+
+    with pytest.raises(ValueError, match="unknown auxiliary components.*transformer"):
+        backend._validate_component_offload_policy(_modules())
+
+
+@pytest.mark.parametrize(
+    ("policy", "encoder_stage", "video_stage", "audio_stage", "try_encoder"),
+    [
+        ({}, True, True, True, True),
+        ({"text_encoder": False}, False, True, True, False),
+        ({"default": False, "video_vae": True}, False, True, False, False),
+    ],
+)
+def test_prepare_auxiliary_components_applies_policy(
+    mocker,
+    policy,
+    encoder_stage,
+    video_stage,
+    audio_stage,
+    try_encoder,
+):
+    backend = _backend(policy)
+    backend._try_layerwise_offload_encoder = mocker.Mock()
+    backend._register_on_demand_hook = mocker.Mock()
+    modules = _modules()
+    plan = OffloadPlan(
+        encoder_block_attrs={"text_encoder": ("layers",)},
+        on_demand_component_paths=frozenset({"text_encoder", "video_vae", "audio_vae"}),
+    )
+
+    backend._prepare_auxiliary_components(modules, plan)
+
+    assert backend._try_layerwise_offload_encoder.called is try_encoder
+    assert modules.encoders[0]._omni_dlo_offload_enabled is encoder_stage
+    assert modules.vaes[0]._omni_dlo_offload_enabled is video_stage
+    assert modules.vaes[1]._omni_dlo_offload_enabled is audio_stage
+    assert backend._register_on_demand_hook.call_args_list == [
+        mocker.call(modules.encoders[0], "text_encoder", stage_on_demand=encoder_stage),
+        mocker.call(modules.vaes[0], "video_vae", stage_on_demand=video_stage),
+        mocker.call(modules.vaes[1], "audio_vae", stage_on_demand=audio_stage),
+    ]

tests/entrypoints/test_async_omni_diffusion_config.py

@@ -332,7 +332,7 @@ def test_serve_cli_forwards_distilled_lora_to_diffusion_stage():
     ]
 
 
-def test_serve_cli_forwards_distributed_offload_residency():
+def test_serve_cli_forwards_distributed_offload_options():
     """Ensure the two-GPU DLO placement controls reach the diffusion stage."""
     parser = TrackingArgumentParser()
     subparsers = parser.add_subparsers(dest="command")
@@ -347,6 +347,8 @@ def test_serve_cli_forwards_distributed_offload_residency():
             "--dlo-no-use-allgather",
             "--dlo-resident-layers",
             "20",
+            "--dlo-offload-components",
+            '{"text_encoder": false, "default": true}',
         ]
     )
 
@@ -357,9 +359,14 @@ def test_serve_cli_forwards_distributed_offload_residency():
     assert args.enable_distributed_layerwise_offload is True
     assert args.dlo_use_allgather is False
     assert args.dlo_resident_layers == 20
+    assert args.dlo_offload_components == {"text_encoder": False, "default": True}
     assert engine_args["enable_distributed_layerwise_offload"] is True
     assert engine_args["dlo_use_allgather"] is False
     assert engine_args["dlo_resident_layers"] == 20
+    assert engine_args["dlo_offload_components"] == {
+        "text_encoder": False,
+        "default": True,
+    }
 
 
 def test_serve_cli_forwards_hwr_policy_for_no_allgather_dlo():

vllm_omni/config/omni_config.py

@@ -697,6 +697,7 @@ class _DiffusionConfigProjection:
     enable_distributed_layerwise_offload: bool = False
     dlo_use_allgather: bool = True
     dlo_resident_layers: int = Field(default=0, ge=0)
+    dlo_offload_components: dict[str, bool] = Field(default_factory=dict)
     host_weight_runtime_mode: Literal["disabled", "preferred", "required"] = "disabled"
     host_weight_runtime_root: str | None = None
     dlo_host_registration_limit_gib: float = Field(default=0.0, ge=0)

vllm_omni/config/stage_config.py

@@ -389,6 +389,7 @@ class StageDeployConfig:
     enable_distributed_layerwise_offload: bool | None = None
     dlo_use_allgather: bool | None = None
     dlo_resident_layers: int | None = None
+    dlo_offload_components: dict[str, bool] | None = None
     host_weight_runtime_mode: str | None = None
     host_weight_runtime_root: str | None = None
     dlo_host_registration_limit_gib: float | None = None

vllm_omni/diffusion/data.py

@@ -795,6 +795,9 @@ class OmniDiffusionConfig:
     dlo_use_allgather: bool = True
     # Leading main-DiT blocks kept resident by distributed layerwise offload.
     dlo_resident_layers: int = 0
+    # Auxiliary encoder/VAE policy. False keeps the whole component resident;
+    # DiT offload remains owned by DLO.
+    dlo_offload_components: dict[str, bool] = field(default_factory=dict)
     # Final-layout Host Weight Runtime policy. The loader only activates this
     # for eligible no-AllGather DLO; all other configurations preserve their
     # existing loader/storage path.

vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py

@@ -1436,7 +1436,7 @@ class MiniMaxH3Pipeline(
             with sequential_offload_component(component):
                 yield
             return
-        staged = self._uses_manual_component_offload()
+        staged = self._uses_manual_component_offload() and getattr(component, "_omni_dlo_offload_enabled", True)
         try:
             if staged:
                 component.load_to_device()

vllm_omni/diffusion/offloader/base.py

@@ -2,7 +2,8 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 from abc import ABC, abstractmethod
-from dataclasses import dataclass
+from collections.abc import Mapping
+from dataclasses import dataclass, field
 from enum import Enum
 from typing import Protocol, runtime_checkable
 
@@ -52,6 +53,9 @@ class OffloadConfig:
     # blocks from the loader-selected host backing with H2D only.
     dlo_use_allgather: bool = True
     dlo_resident_layers: int = 0  # leading DiT layers kept on device
+    # Per auxiliary component: True permits model-declared DLO offload;
+    # False keeps the whole component resident. "default" is the fallback.
+    dlo_offload_components: dict[str, bool] = field(default_factory=dict)
     # 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
@@ -125,6 +129,25 @@ class OffloadConfig:
         # 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))
+        raw_component_policy = getattr(od_config, "dlo_offload_components", {})
+        if raw_component_policy is None:
+            raw_component_policy = {}
+        if not isinstance(raw_component_policy, Mapping):
+            raise TypeError(
+                "dlo_offload_components must be a mapping of component names to booleans, "
+                f"got {type(raw_component_policy).__name__}"
+            )
+        dlo_offload_components: dict[str, bool] = {}
+        for component, enabled in raw_component_policy.items():
+            if not isinstance(component, str) or not component:
+                raise TypeError("dlo_offload_components keys must be non-empty strings")
+            if not isinstance(enabled, bool):
+                raise TypeError(
+                    f"dlo_offload_components[{component!r}] must be a boolean, got {type(enabled).__name__}"
+                )
+            dlo_offload_components[component] = enabled
+        if dlo_offload_components and not enable_distributed_layerwise_offload:
+            raise ValueError("dlo_offload_components requires distributed layerwise offload")
         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,
@@ -166,6 +189,7 @@ class OffloadConfig:
             dp_size=dp_size,
             dlo_use_allgather=dlo_use_allgather,
             dlo_resident_layers=dlo_resident_layers,
+            dlo_offload_components=dlo_offload_components,
             dlo_host_registration_limit_gib=dlo_host_registration_limit_gib,
         )
 

vllm_omni/diffusion/offloader/distributed_layerwise_backend.py

@@ -43,7 +43,7 @@ from .host_registration import (
     HostRegistrationError,
     register_host_mappings,
 )
-from .module_collector import ModuleDiscovery
+from .module_collector import ModuleDiscovery, PipelineModules
 from .offload_plan import (
     OffloadPlan,
     get_offload_plan,
@@ -1400,6 +1400,39 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
         module.to(self.device)
         logger.info("Moved %s (%s) to GPU (resident)", label, module.__class__.__name__)
 
+    def _component_offload_enabled(self, name: str) -> bool:
+        policy = self.config.dlo_offload_components
+        return policy.get(name, policy.get("default", True))
+
+    def _validate_component_offload_policy(self, modules: PipelineModules) -> None:
+        known_components = {"default", *modules.encoder_names, *modules.vae_names}
+        unknown = sorted(self.config.dlo_offload_components.keys() - known_components)
+        if unknown:
+            raise ValueError(
+                f"dlo_offload_components contains unknown auxiliary components: {unknown}. DLO always owns DiT offload."
+            )
+
+    def _prepare_auxiliary_components(self, modules: PipelineModules, plan: OffloadPlan | None) -> None:
+        for encoder, name in zip(modules.encoders, modules.encoder_names, strict=True):
+            offload = self._component_offload_enabled(name)
+            encoder._omni_dlo_offload_enabled = offload
+            if offload:
+                self._try_layerwise_offload_encoder(encoder, name, plan)
+            self._register_on_demand_hook(
+                encoder,
+                name,
+                stage_on_demand=(offload and plan is not None and name in plan.on_demand_component_paths),
+            )
+
+        for vae, name in zip(modules.vaes, modules.vae_names, strict=True):
+            offload = self._component_offload_enabled(name)
+            vae._omni_dlo_offload_enabled = offload
+            self._register_on_demand_hook(
+                vae,
+                name,
+                stage_on_demand=(offload and plan is not None and name in plan.on_demand_component_paths),
+            )
+
     def _try_layerwise_offload_encoder(self, module: nn.Module, name: str, plan: OffloadPlan | None) -> bool:
         """Stream plan-declared encoder blocks on each rank without AllGather."""
         if plan is None or name not in plan.encoder_block_attrs:
@@ -1668,6 +1701,7 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
         # Retrieve optional declarative OffloadPlan from the pipeline.
         # When present, replaces heuristic block discovery.
         plan = get_offload_plan(pipeline)
+        self._validate_component_offload_policy(modules)
 
         if self.config.dlo_resident_layers and (plan is None or not plan.resident_dit_paths):
             logger.warning(
@@ -1728,21 +1762,10 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
                 )
             logger.info("DLO is using host tensors materialized by the ordinary loader")
 
-        # Keep VAE/encoders on CPU; move to GPU on-demand via hooks.
-        # This saves ~4.3 GB HBM per card (VAE 1.3 + encoder 1.1 + sound 1.9)
-        # during the DiT forward pass.  They are only needed briefly for
-        # text-encoding (before DiT) and VAE-decode (after DiT).
-        for enc, enc_name in zip(modules.encoders, modules.encoder_names):
-            self._try_layerwise_offload_encoder(enc, enc_name, plan)
-            self._register_on_demand_hook(
-                enc, "encoder", stage_on_demand=plan is not None and enc_name in plan.on_demand_component_paths
-            )
-        for vae, vae_name in zip(modules.vaes, modules.vae_names):
-            self._register_on_demand_hook(
-                vae,
-                "vae",
-                stage_on_demand=(plan is not None and vae_name in plan.on_demand_component_paths),
-            )
+        # Apply the user policy to discovered auxiliary components. The default
+        # keeps the model-declared staged lifecycle; False leaves a component
+        # resident for repeated requests.
+        self._prepare_auxiliary_components(modules, plan)
 
         # Move resident modules to GPU (small modules needed every forward)
         for name, module in zip(modules.resident_names, modules.resident_modules):

vllm_omni/engine/arg_utils.py

@@ -569,6 +569,7 @@ class OrchestratorArgs:
     enable_distributed_layerwise_offload: bool = False
     dlo_use_allgather: bool = True
     dlo_resident_layers: int = 0
+    dlo_offload_components: dict[str, bool] = field(default_factory=dict)
     host_weight_runtime_mode: str = "disabled"
     host_weight_runtime_root: str | None = None
     dlo_host_registration_limit_gib: float = 0.0

vllm_omni/engine/async_omni_engine.py

@@ -1059,6 +1059,9 @@ class AsyncOmniEngine:
             "enable_distributed_layerwise_offload": kwargs.get("enable_distributed_layerwise_offload", False),
             "dlo_use_allgather": kwargs.get("dlo_use_allgather", True),
             "dlo_resident_layers": kwargs.get("dlo_resident_layers", 0),
+            "dlo_offload_components": (
+                {} if kwargs.get("dlo_offload_components") is None else kwargs["dlo_offload_components"]
+            ),
             "host_weight_runtime_mode": kwargs.get("host_weight_runtime_mode", "disabled"),
             "host_weight_runtime_root": kwargs.get("host_weight_runtime_root"),
             "dlo_host_registration_limit_gib": kwargs.get("dlo_host_registration_limit_gib", 0.0),

vllm_omni/entrypoints/cli/serve.py

@@ -770,6 +770,18 @@ class OmniServeCommand(CLISubcommand):
                 "under the existing pinned-memory policy and fall back to bounded staging when unavailable."
             ),
         )
+        omni_config_group.add_argument(
+            "--dlo-offload-components",
+            type=json.loads,
+            default=None,
+            metavar="JSON",
+            help=(
+                "Per auxiliary encoder/VAE DLO policy as a JSON object, for example "
+                "'{\"text_encoder\": false}'. False keeps the whole component "
+                "resident; true permits the model's existing offload behavior. "
+                "DiT offload is not controlled by this map."
+            ),
+        )
         # Video model parameters (e.g., Wan2.2) - engine-level
         omni_config_group.add_argument(
             "--boundary-ratio",