[Refactor][Diffusion] Read the resolved offload policy outside the config layer (#7327)

Signed-off-by: specture724 <specture724@gmail.com>
Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

tests/diffusion/models/magi2/test_pipeline_magi2.py

@@ -78,10 +78,15 @@ class _ParallelStub:
 @dataclass
 class _TopologyStub:
     parallel_config: _ParallelStub = field(default_factory=_ParallelStub)
+    diffusion_offload_config: object = None
     enable_cpu_offload: bool = False
     enable_layerwise_offload: bool = False
     enable_distributed_layerwise_offload: bool = False
     dlo_use_allgather: bool = True
+    dlo_resident_layers: int = 0
+    dlo_host_registration_limit_gib: float = 0.0
+    host_weight_runtime_mode: str = "disabled"
+    pin_cpu_memory: bool = True
     quantization_config: object = None
     cache_backend: str = "none"
     custom_pipeline_args: dict[str, object] = field(default_factory=dict)
@@ -383,26 +388,33 @@ def test_native_topology_rejects_nondivisible_tensor_parallelism():
 
 
 def test_native_topology_accepts_single_device_layerwise_with_cpu_staging():
-    config = _topology_config(
-        enable_layerwise_offload=True,
-        additional_config={},
+    _validate_native_topology(
+        _topology_config(
+            enable_layerwise_offload=True,
+            additional_config={},
+        )
+    )
+    _validate_native_topology(
+        _topology_config(
+            enable_cpu_offload=True,
+            enable_layerwise_offload=True,
+            additional_config={},
+        )
     )
-    _validate_native_topology(config)
 
-    config.enable_cpu_offload = True
-    _validate_native_topology(config)
-
-    config.enable_layerwise_offload = False
-    with pytest.raises(ValueError, match="Combine --enable-cpu-offload"):
-        _validate_native_topology(config)
+    with pytest.raises(ValueError, match="already stages its auxiliary components"):
+        _validate_native_topology(
+            _topology_config(
+                enable_cpu_offload=True,
+                additional_config={},
+            )
+        )
 
 
 def test_native_topology_requires_dlo_rank_local_mode():
-    config = _topology_config(enable_distributed_layerwise_offload=True)
     with pytest.raises(ValueError, match="dlo-no-use-allgather"):
-        _validate_native_topology(config)
-    config.dlo_use_allgather = False
-    _validate_native_topology(config)
+        _validate_native_topology(_topology_config(enable_distributed_layerwise_offload=True))
+    _validate_native_topology(_topology_config(enable_distributed_layerwise_offload=True, dlo_use_allgather=False))
 
 
 def test_native_topology_accepts_four_device_hsdp_cfg_and_vae_layouts():

tests/diffusion/models/minimax_h3/test_minimax_h3_dlo_lifecycle.py

@@ -13,6 +13,21 @@ class _OffloadAbort(BaseException):
     pass
 
 
+def _legacy_distributed_offload_config() -> SimpleNamespace:
+    return SimpleNamespace(
+        diffusion_offload_config=None,
+        enable_cpu_offload=False,
+        enable_layerwise_offload=False,
+        enable_distributed_layerwise_offload=True,
+        dlo_use_allgather=True,
+        dlo_resident_layers=0,
+        dlo_host_registration_limit_gib=0.0,
+        host_weight_runtime_mode="disabled",
+        pin_cpu_memory=True,
+        parallel_config=SimpleNamespace(data_parallel_size=1),
+    )
+
+
 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
@@ -100,10 +115,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.od_config = _legacy_distributed_offload_config()
     pipeline._model_cpu_offload_modules = []
     pipeline._dlo_component_cache = mocker.Mock()
     component = mocker.Mock()
@@ -123,10 +135,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.od_config = _legacy_distributed_offload_config()
     pipeline._model_cpu_offload_modules = []
     pipeline._dlo_component_cache = mocker.Mock()
     component = mocker.Mock()

tests/diffusion/offloader/test_legacy_flag_readers.py

@@ -0,0 +1,51 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Runtime code reads the resolved offload policy, not the legacy booleans."""
+
+import re
+from pathlib import Path
+
+import pytest
+
+import vllm_omni
+
+pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
+
+# `enable_cpu_offload`, `enable_layerwise_offload` and
+# `enable_distributed_layerwise_offload` are compatibility aliases that
+# `resolve_offload` materializes from the resolved policy. Only the layer that
+# declares and resolves them may read them; everything else asks
+# `vllm_omni.diffusion.offloader.config` for the strategy or the selected
+# components, so the aliases can be removed without sweeping the runtime again.
+COMPATIBILITY_LAYER = {
+    "config/omni_config.py",
+    "config/stage_config.py",
+    "engine/arg_utils.py",
+    "diffusion/data.py",
+    "diffusion/offloader/base.py",
+    "diffusion/offloader/config.py",
+    "quantization/tools/compare_diffusion_trajectory_similarity.py",
+}
+
+READER = re.compile(
+    r"""\.enable_(?:cpu|layerwise|distributed_layerwise)_offload\b"""
+    r"""|getattr\(\s*[^,]+,\s*["']enable_(?:cpu|layerwise|distributed_layerwise)_offload["']"""
+)
+
+
+def test_no_runtime_module_reads_the_legacy_offload_flags():
+    root = Path(vllm_omni.__file__).parent
+    offenders = [
+        f"{path.relative_to(root)}:{number}"
+        for path in sorted(root.rglob("*.py"))
+        if str(path.relative_to(root)) not in COMPATIBILITY_LAYER
+        for number, line in enumerate(path.read_text().splitlines(), start=1)
+        if READER.search(line)
+    ]
+
+    assert offenders == [], (
+        "Read the resolved policy instead: resolve_offload_strategy(od_config), "
+        "offload_enabled(od_config), offload_streams_blocks(od_config), or "
+        f"should_offload_component(od_config, component). Offending reads: {offenders}"
+    )

tests/diffusion/quantization/test_cosmos3_mixed_precision.py

@@ -406,28 +406,38 @@ def test_runtime_allows_standard_offload_and_rejects_distributed(monkeypatch) ->
         "get_tensor_model_parallel_world_size",
         lambda: 1,
     )
-    od_config = SimpleNamespace(
-        enable_cpu_offload=True,
-        enable_layerwise_offload=True,
-        enable_distributed_layerwise_offload=False,
-        parallel_config=SimpleNamespace(use_hsdp=False),
-    )
+
+    def offload_config(**overrides: object) -> SimpleNamespace:
+        config = SimpleNamespace(
+            diffusion_offload_config=None,
+            enable_cpu_offload=False,
+            enable_layerwise_offload=True,
+            enable_distributed_layerwise_offload=False,
+            dlo_use_allgather=True,
+            dlo_resident_layers=0,
+            dlo_host_registration_limit_gib=0.0,
+            host_weight_runtime_mode="disabled",
+            pin_cpu_memory=True,
+            max_num_seqs=1,
+            parallel_config=SimpleNamespace(data_parallel_size=1, use_hsdp=False),
+        )
+        for name, value in overrides.items():
+            setattr(config, name, value)
+        return config
+
     transformer_cosmos3._validate_mixed_precision_runtime(
         Cosmos3MixedPrecisionConfig(),
-        od_config,
+        offload_config(),
     )
-    od_config.enable_distributed_layerwise_offload = True
     with pytest.raises(ValueError, match="distributed layer-wise offload"):
         transformer_cosmos3._validate_mixed_precision_runtime(
             Cosmos3MixedPrecisionConfig(),
-            od_config,
+            offload_config(enable_distributed_layerwise_offload=True),
         )
-    od_config.enable_distributed_layerwise_offload = False
-    od_config.max_num_seqs = 2
     with pytest.raises(ValueError, match="one active request"):
         transformer_cosmos3._validate_mixed_precision_runtime(
             Cosmos3MixedPrecisionConfig(),
-            od_config,
+            offload_config(max_num_seqs=2),
         )
 
 

vllm_omni/diffusion/diffusion_engine.py

@@ -44,7 +44,11 @@ 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.offloader.config import (
+    OffloadStrategy,
+    any_selected_component_uses_allgather,
+    resolve_offload_strategy,
+)
 from vllm_omni.diffusion.output_formatter import (
     format_diffusion_outputs,
     format_empty_diffusion_outputs,
@@ -573,7 +577,8 @@ class DiffusionEngine:
         if output.media is not None:
             if output.output is not None:
                 raise ValueError("DiffusionOutput cannot contain both media and legacy output")
-            media = output.media.to_cpu() if self.od_config.enable_cpu_offload else output.media
+            model_level = resolve_offload_strategy(self.od_config) is OffloadStrategy.MODEL_LEVEL
+            media = output.media.to_cpu() if model_level else output.media
             output_data = media.video.tensor
             outputs = finalize_diffusion_media(media, sampling_params=request.sampling_params)
         else:
@@ -585,7 +590,7 @@ class DiffusionEngine:
             # post-processing to avoid device OOM — model weights may still
             # reside on the device and leave no headroom for intermediates.
             output_data = output.output
-            if self.od_config.enable_cpu_offload:
+            if resolve_offload_strategy(self.od_config) is OffloadStrategy.MODEL_LEVEL:
                 output_data = _move_tensor_tree_to_cpu(output_data)
 
             if self.post_process_func is not None:

vllm_omni/diffusion/lora/manager.py

@@ -73,11 +73,17 @@ class DiffusionLoRAManager:
         self.pipeline = pipeline
         self.device = device
         self.dtype = dtype
+        # Imported here: vllm_omni.diffusion.data imports this package while it
+        # is still initializing, and the offloader package imports that module.
+        from vllm_omni.diffusion.offloader.config import OffloadStrategy, resolve_offload_strategy
+
         od_config = getattr(pipeline, "od_config", None)
         # DLO owns the base-weight lifecycle. Keep request-switchable LoRA
         # sidecars resident instead of rebuilding DLO host shards per request.
         self._resident_lora_device = (
-            device if getattr(od_config, "enable_distributed_layerwise_offload", False) is True else None
+            device
+            if od_config is not None and resolve_offload_strategy(od_config) is OffloadStrategy.DISTRIBUTED_LAYER_WISE
+            else None
         )
 
         # Cache supported/expected module suffixes once, before any layer

vllm_omni/diffusion/models/bagel/pipeline_bagel.py

@@ -32,6 +32,10 @@ from vllm_omni.diffusion.distributed.parallel_state import (
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
 from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery
+from vllm_omni.diffusion.offloader.config import (
+    OffloadStrategy,
+    resolve_offload_strategy,
+)
 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
@@ -500,7 +504,9 @@ class BagelPipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipelineProf
         #    selectively materialized/moved by the offloader.
         # 3. HSDP: weights should be loaded on CPU first and sharded afterwards,
         #    rather than eagerly placing the full model on one GPU.
-        if quant_config is None and not (od_config.enable_layerwise_offload or od_config.parallel_config.use_hsdp):
+        if quant_config is None and not (
+            resolve_offload_strategy(od_config) is OffloadStrategy.LAYER_WISE or od_config.parallel_config.use_hsdp
+        ):
             self.to(self.device)
         self.setup_diffusion_pipeline_profiler(
             enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler

vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py

@@ -38,6 +38,10 @@ from vllm_omni.diffusion.data import OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.sp_plan import SequenceParallelInput, SequenceParallelOutput
 from vllm_omni.diffusion.forward_context import get_forward_context, is_forward_context_available
 from vllm_omni.diffusion.layers.norm import RMSNorm as _VllmRMSNorm
+from vllm_omni.diffusion.offloader.config import (
+    OffloadStrategy,
+    resolve_offload_strategy,
+)
 from vllm_omni.platforms import current_omni_platform
 from vllm_omni.quantization.component_config import (
     ComponentQuantizationConfig,
@@ -203,7 +207,7 @@ def _validate_mixed_precision_runtime(
         raise ValueError("Cosmos3 mixed precision currently supports tensor parallel size 1 only")
     if int(getattr(od_config, "max_num_seqs", 1)) != 1:
         raise ValueError("Cosmos3 mixed precision currently supports one active request per worker")
-    if bool(getattr(od_config, "enable_distributed_layerwise_offload", False)):
+    if resolve_offload_strategy(od_config) is OffloadStrategy.DISTRIBUTED_LAYER_WISE:
         raise ValueError(
             "Cosmos3 mixed precision does not support distributed layer-wise offload "
             "because its direct loader bypasses ModelOpt post-load transformations"

vllm_omni/diffusion/models/diffusers_adapter/pipeline_diffusers_adapter.py

@@ -33,6 +33,10 @@ from vllm_omni.diffusion.models.diffusers_adapter.quantization_utils import (
     convert_diffusers_quantization_config,
     ensure_supported_diffusers_quantization,
 )
+from vllm_omni.diffusion.offloader.config import (
+    OffloadStrategy,
+    resolve_offload_strategy,
+)
 from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
 from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
 from vllm_omni.inputs.data import OmniPromptType, OmniTextPrompt
@@ -148,9 +152,10 @@ class DiffusersAdapterPipeline(nn.Module, DiffusionPipelineProfilerMixin):
         self._accept_call_kwargs = set(inspect.signature(self._pipeline.__call__).parameters.keys())
 
         # CPU offloading
-        if self.od_config.enable_layerwise_offload:
+        strategy = resolve_offload_strategy(self.od_config)
+        if strategy is OffloadStrategy.LAYER_WISE:
             self._pipeline.enable_sequential_cpu_offload()
-        elif self.od_config.enable_cpu_offload:
+        elif strategy is OffloadStrategy.MODEL_LEVEL:
             self._pipeline.enable_model_cpu_offload()
 
         # VAE slicing and tiling: try-catch because not all models have VAE

vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py

@@ -54,6 +54,10 @@ from vllm_omni.diffusion.models.dreamzero.utils import (
     DEFAULT_SIGMA_SHIFT,
 )
 from vllm_omni.diffusion.models.schedulers.scheduling_flow_unipc_multistep import FlowUniPCMultistepScheduler
+from vllm_omni.diffusion.offloader.config import (
+    OffloadStrategy,
+    resolve_offload_strategy,
+)
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
 from vllm_omni.experimental.ar_diffusion.capability import (
@@ -439,9 +443,7 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
         else:
             self.vae = DistributedAutoencoderKLWan()
             self.vae.init_distributed()
-        if not (
-            getattr(od_config, "enable_cpu_offload", False) or getattr(od_config, "enable_layerwise_offload", False)
-        ):
+        if resolve_offload_strategy(od_config) not in (OffloadStrategy.MODEL_LEVEL, OffloadStrategy.LAYER_WISE):
             self.vae = self.vae.to(device=get_local_device(), dtype=od_config.dtype)
         self.register_buffer(
             "vae_latents_mean",

vllm_omni/diffusion/models/flux2/pipeline_flux2.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
 import inspect
 import json
 import logging
@@ -34,6 +34,10 @@ from vllm_omni.diffusion.models.flux2 import Flux2Transformer2DModel
 from vllm_omni.diffusion.models.interface import SupportImageInput, SupportsComponentDiscovery
 from vllm_omni.diffusion.models.mistral_encoder import MistralEncoderModel
 from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin
+from vllm_omni.diffusion.offloader.config import (
+    OffloadStrategy,
+    resolve_offload_strategy,
+)
 from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
 from vllm_omni.diffusion.utils.tf_utils import get_transformer_config_kwargs
 from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
@@ -1160,7 +1164,7 @@ class Flux2Pipeline(
         text_encoder_quant_config = _resolve_component_quant_config(self.od_config.quantization_config, "text_encoder")
         transformer_quant_config = _resolve_component_quant_config(self.od_config.quantization_config, "transformer")
         if (
-            self.od_config.enable_cpu_offload
+            resolve_offload_strategy(self.od_config) is OffloadStrategy.MODEL_LEVEL
             and text_encoder_quant_config is not None
             and transformer_quant_config is None
         ):

vllm_omni/diffusion/models/lance/pipeline_lance.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
 """LancePipeline — Lance (ByteDance) packaged for the vLLM-Omni diffusion engine.
 
 Lance is BAGEL-lineage (Qwen2-MoT unified AR+diffusion), so the transformer
@@ -59,6 +59,10 @@ from vllm_omni.diffusion.models.bagel.pipeline_bagel import (
     BagelPipeline,
     add_special_tokens,
 )
+from vllm_omni.diffusion.offloader.config import (
+    OffloadStrategy,
+    resolve_offload_strategy,
+)
 from vllm_omni.model_executor.model_loader.weight_utils import (
     download_weights_from_hf_specific,
 )
@@ -401,7 +405,9 @@ class LancePipeline(BagelPipeline):
                 )
             )
 
-        if quant_config is None and not (od_config.enable_layerwise_offload or od_config.parallel_config.use_hsdp):
+        if quant_config is None and not (
+            resolve_offload_strategy(od_config) is OffloadStrategy.LAYER_WISE or od_config.parallel_config.use_hsdp
+        ):
             self.to(self.device)
         self.setup_diffusion_pipeline_profiler(
             enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler

vllm_omni/diffusion/models/lingbot_world/pipeline.py

@@ -67,6 +67,10 @@ from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin
 from vllm_omni.diffusion.models.schedulers import FlowUniPCMultistepScheduler
 from vllm_omni.diffusion.models.utils import _load_json
 from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import load_transformer_config, retrieve_latents
+from vllm_omni.diffusion.offloader.config import (
+    OffloadStrategy,
+    resolve_offload_strategy,
+)
 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
@@ -611,8 +615,9 @@ class LingBotWorldCausalDMDPipeline(
         dtype = getattr(od_config, "dtype", torch.bfloat16)
         model = od_config.model
         local_files_only = os.path.exists(model)
-        managed_component_placement = bool(
-            getattr(od_config, "enable_cpu_offload", False) or getattr(od_config, "enable_layerwise_offload", False)
+        managed_component_placement = resolve_offload_strategy(od_config) in (
+            OffloadStrategy.MODEL_LEVEL,
+            OffloadStrategy.LAYER_WISE,
         )
 
         # Standard components use from_pretrained; the custom transformer uses the loader.

vllm_omni/diffusion/models/ltx2/ltx2_components.py

@@ -29,6 +29,10 @@ from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_ltx2 import Dis
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
 from vllm_omni.diffusion.model_loader.hub_prefetch import from_pretrained_with_prefetch, prefetch_subfolders
+from vllm_omni.diffusion.offloader.config import (
+    OffloadStrategy,
+    resolve_offload_strategy,
+)
 from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery
 from vllm_omni.diffusion.offloader.offload_plan import OffloadPlan
 from vllm_omni.transformers_utils.repo_utils import hf_api
@@ -621,11 +625,10 @@ def _load_ltx25_native_diffusion_decoder(
 
 def _place_aux_components(pipeline: Any) -> None:
     parallel_config = getattr(pipeline.od_config, "parallel_config", None)
-    use_managed_placement = bool(
-        getattr(pipeline.od_config, "enable_cpu_offload", False)
-        or getattr(pipeline.od_config, "enable_layerwise_offload", False)
-        or getattr(parallel_config, "use_hsdp", False)
-    )
+    use_managed_placement = resolve_offload_strategy(pipeline.od_config) in (
+        OffloadStrategy.MODEL_LEVEL,
+        OffloadStrategy.LAYER_WISE,
+    ) or bool(getattr(parallel_config, "use_hsdp", False))
     if use_managed_placement:
         return
 

vllm_omni/diffusion/models/magi2/pipeline_magi2.py

@@ -44,6 +44,12 @@ from vllm_omni.diffusion.models.interface import (
 )
 from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin
 from vllm_omni.diffusion.offloader import OffloadPlan, PinnedModuleStager
+from vllm_omni.diffusion.offloader.config import (
+    DIT_COMPONENT,
+    OffloadStrategy,
+    offload_streams_blocks,
+    resolve_offload,
+)
 from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import (
     DiffusionPipelineProfilerMixin,
 )
@@ -298,18 +304,17 @@ def _validate_native_topology(od_config: OmniDiffusionConfig) -> None:
         raise ValueError(f"MAGI-2 tensor_parallel_size={tp_size} does not divide: " + ", ".join(invalid_tp_dimensions))
 
     configured_world_size = dp_size * cfg_size * tp_size * sp_size
-    cpu_offload = bool(od_config.enable_cpu_offload)
-    layerwise_offload = bool(od_config.enable_layerwise_offload)
-    distributed_offload = bool(od_config.enable_distributed_layerwise_offload)
-    if cpu_offload and not layerwise_offload:
+    resolved_offload = resolve_offload(od_config)
+    strategy = resolved_offload.strategy
+    layerwise_offload = strategy is OffloadStrategy.LAYER_WISE
+    distributed_offload = strategy is OffloadStrategy.DISTRIBUTED_LAYER_WISE
+    if strategy is OffloadStrategy.MODEL_LEVEL:
         raise ValueError(
             "MAGI-2 already stages its auxiliary components from CPU, while "
             "the complete Preview transformer cannot fit on one qualified GPU. "
-            "Combine --enable-cpu-offload with --enable-layerwise-offload, or "
-            "use --enable-layerwise-offload alone."
+            "Use layer offload instead: diffusion_offload_config={'mode': 'layer', "
+            "'components': ['dit']}, or the legacy --enable-layerwise-offload."
         )
-    if layerwise_offload and distributed_offload:
-        raise ValueError("MAGI-2 ordinary and distributed layerwise offload are mutually exclusive")
     if layerwise_offload and configured_world_size != 1:
         raise ValueError(
             "MAGI-2 ordinary layerwise offload is a single-worker path; "
@@ -330,7 +335,6 @@ def _validate_native_topology(od_config: OmniDiffusionConfig) -> None:
             f"expected vae_patch_parallel_size=1 or {configured_world_size}, got {vae_pp_size}."
         )
 
-    dlo_allgather = bool(getattr(od_config, "dlo_use_allgather", True))
     if cfg_size > 1 and dp_size > 1:
         raise ValueError("MAGI-2 CFG parallelism is not yet combined with DLO data parallelism")
     if distributed_offload and getattr(parallel, "use_hsdp", False):
@@ -339,7 +343,7 @@ def _validate_native_topology(od_config: OmniDiffusionConfig) -> None:
         raise ValueError("MAGI-2 data parallelism currently requires distributed layerwise offload")
     if dp_size > 1 and tp_size > 1:
         raise ValueError("MAGI-2 DLO data-parallel replicas currently require tensor_parallel_size=1")
-    if distributed_offload and dlo_allgather:
+    if distributed_offload and resolved_offload.uses_allgather(DIT_COMPONENT):
         if dp_size <= 1:
             raise ValueError(
                 "MAGI-2 DLO AllGather requires data_parallel_size > 1. SP ranks "
@@ -574,9 +578,7 @@ class Magi2Pipeline(
         )
         self._is_output_rank = self._parallel_group.rank == 0
         self._offload_aux_after_use = True
-        self._transformer_is_layerwise_offloaded = bool(
-            od_config.enable_layerwise_offload or od_config.enable_distributed_layerwise_offload
-        )
+        self._transformer_is_layerwise_offloaded = offload_streams_blocks(od_config)
         self._transformer_is_hsdp = bool(getattr(od_config.parallel_config, "use_hsdp", False))
         self._distributed_video_decode = int(od_config.parallel_config.vae_patch_parallel_size) > 1
 
@@ -585,8 +587,10 @@ class Magi2Pipeline(
         from .modeling_magi2 import Magi2PreviewTransformer
 
         MAGI2_PREVIEW_CONFIG.validate()
-        mmap_dlo = bool(
-            od_config.enable_distributed_layerwise_offload and getattr(od_config, "dlo_use_allgather", True)
+        resolved_offload = resolve_offload(od_config)
+        mmap_dlo = (
+            resolved_offload.strategy is OffloadStrategy.DISTRIBUTED_LAYER_WISE
+            and resolved_offload.uses_allgather(DIT_COMPONENT)
         )
         if mmap_dlo:
             # AllGather DLO binds checkpoint tensors as mmap views and copies

vllm_omni/diffusion/models/minimax_h3/fasth3.py

@@ -42,6 +42,7 @@ import torch
 from safetensors import safe_open
 from vllm.logger import init_logger
 
+from vllm_omni.diffusion.offloader.config import offload_enabled, resolve_offload_strategy
 from vllm_omni.diffusion.sched.sigma_schedule import DMD2SigmaSchedule
 from vllm_omni.errors import OmniClientError
 from vllm_omni.platforms import current_omni_platform
@@ -608,19 +609,14 @@ class FastH3WeightFusion:
         """Hold a starting server to the ladder this student was trained on."""
         if partition == "ref2va":
             raise ValueError("FastH3 preview v1 distills T2VA only, so it cannot serve a Ref2VA partition")
-        offloads = [
-            flag
-            for flag in ("enable_cpu_offload", "enable_layerwise_offload", "enable_distributed_layerwise_offload")
-            if getattr(od_config, flag, False)
-        ]
-        if offloads:
+        if offload_enabled(od_config):
             # A host-weight plan installs the transformer without going through
             # load_weights(), which is where the fusion and its completeness
             # check live. Serving base H3 weights under a four-step schedule
             # would otherwise degrade output with nothing to signal it.
             raise ValueError(
-                f"FastH3 is fused while the checkpoint streams in, so it cannot be combined with "
-                f"{sorted(offloads)}. Serve it without offload."
+                "FastH3 is fused while the checkpoint streams in, so it cannot be combined with "
+                f"{resolve_offload_strategy(od_config).value} offload. Serve it without offload."
             )
         if self.requires_vsa:
             backend = _resolve_dit_attention_backend(od_config)

vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py

@@ -53,7 +53,9 @@ from vllm_omni.diffusion.offloader.config import (
     DIT_COMPONENT,
     TEXT_ENCODER_COMPONENT,
     OffloadStrategy,
+    offload_streams_blocks,
     resolve_offload,
+    resolve_offload_strategy,
     should_offload_component,
 )
 from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery
@@ -1051,8 +1053,8 @@ class MiniMaxH3Pipeline(
             self.text_encoder_group = None
             self.text_encoder = None
             self._encoder_modules = []
-        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)
+        legacy_manual_components = getattr(od_config, "diffusion_offload_config", None) is None and (
+            offload_streams_blocks(od_config)
         )
         # Preserve the legacy MiniMax-H3 low-residency path. The compact API
         # deliberately limits explicit component selection to dit/text_encoder,
@@ -1078,7 +1080,7 @@ class MiniMaxH3Pipeline(
         self._dlo_component_cache = None
         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:
+        if resolve_offload_strategy(od_config) is OffloadStrategy.DISTRIBUTED_LAYER_WISE and needs_component_cache:
             self._dlo_component_cache = BoundedAllocatorCache(self.device)
             if legacy_manual_components:
                 _register_dlo_component_cache(
@@ -1437,10 +1439,8 @@ class MiniMaxH3Pipeline(
         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)
-            )
+            # The compatibility topology stages every component it can.
+            return offload_streams_blocks(od_config)
         return component is getattr(self, "text_encoder", None) and should_offload_component(
             od_config, TEXT_ENCODER_COMPONENT
         )

vllm_omni/diffusion/models/sana_video/pipeline_sana_video.py

@@ -48,6 +48,11 @@ from vllm_omni.diffusion.model_loader.hub_prefetch import from_pretrained_with_p
 from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery
 from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin
 from vllm_omni.diffusion.models.utils import _load_json
+from vllm_omni.diffusion.offloader.config import (
+    OffloadStrategy,
+    offload_enabled,
+    resolve_offload_strategy,
+)
 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 resolve_video_num_frames
@@ -280,18 +285,13 @@ def _validate_cache_offload_parallelism(od_config: OmniDiffusionConfig) -> None:
             f"Cache backend {cache_backend!r} is not supported by the native SANA-Video pipeline; "
             "use 'cache_dit' or 'none'."
         )
-    if cache_backend == "cache_dit" and od_config.enable_distributed_layerwise_offload:
+    if cache_backend == "cache_dit" and resolve_offload_strategy(od_config) is OffloadStrategy.DISTRIBUTED_LAYER_WISE:
         raise NotImplementedError(
             "SANA-Video does not support Cache-DiT with distributed layerwise offload; "
             "per-rank cache skips desynchronize the weight AllGather."
         )
 
-    offload_enabled = (
-        od_config.enable_cpu_offload
-        or od_config.enable_layerwise_offload
-        or od_config.enable_distributed_layerwise_offload
-    )
-    if cache_backend == "none" and not offload_enabled:
+    if cache_backend == "none" and not offload_enabled(od_config):
         return
 
     parallel_config = od_config.parallel_config

vllm_omni/diffusion/models/sana_wm/pipeline_sana_wm.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
 """Sana-WM pipeline integration.
 
 This module wires the registry-visible surface, release-layout validation, and
@@ -50,6 +50,10 @@ from vllm_omni.diffusion.models.sana_wm.sana_wm_transformer import (
     SanaWmTransformer3DModel,
 )
 from vllm_omni.diffusion.models.schedulers import FlowMatchEulerDiscreteScheduler
+from vllm_omni.diffusion.offloader.config import (
+    OffloadStrategy,
+    resolve_offload_strategy,
+)
 from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery
 from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import (
     DiffusionPipelineProfilerMixin,
@@ -389,11 +393,10 @@ class SanaWmPipeline(
         components to still be on CPU when they take over.
         """
         parallel_config = getattr(self.od_config, "parallel_config", None)
-        if (
-            getattr(self.od_config, "enable_cpu_offload", False)
-            or getattr(self.od_config, "enable_layerwise_offload", False)
-            or getattr(parallel_config, "use_hsdp", False)
-        ):
+        if resolve_offload_strategy(self.od_config) in (
+            OffloadStrategy.MODEL_LEVEL,
+            OffloadStrategy.LAYER_WISE,
+        ) or getattr(parallel_config, "use_hsdp", False):
             return
         modules = ModuleDiscovery.discover(self)
         for module in (*modules.encoders, *modules.vaes, *modules.resident_modules):

vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_s2v.py

@@ -46,7 +46,12 @@ 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.offloader.config import (
+    DIT_COMPONENT,
+    OffloadStrategy,
+    resolve_offload_strategy,
+    selected_offload_components,
+)
 from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.utils.chunked_video import ChunkedVideoMP4Session
@@ -659,7 +664,10 @@ class Wan22S2VPipeline(
 
         t5_checkpoint = os.path.join(model_path, "models_t5_umt5-xxl-enc-bf16.pth")
         self.text_encoder = UMT5EncoderModel(_WAN_UMT5_CONFIG)
-        _cpu_offload = self.od_config.enable_cpu_offload or self.od_config.enable_layerwise_offload
+        _cpu_offload = resolve_offload_strategy(self.od_config) in (
+            OffloadStrategy.MODEL_LEVEL,
+            OffloadStrategy.LAYER_WISE,
+        )
         self.text_encoder = _load_wan_t5_as_umt5(self.text_encoder, t5_checkpoint, dtype=dtype)
         if not _cpu_offload:
             self.text_encoder = self.text_encoder.to(self.device)
@@ -1108,7 +1116,9 @@ class Wan22S2VPipeline(
     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)
+        return resolve_offload_strategy(self.od_config) is OffloadStrategy.MODEL_LEVEL and (
+            DIT_COMPONENT in selected_offload_components(self.od_config)
+        )
 
     def forward(
         self,

vllm_omni/diffusion/offloader/config.py

@@ -356,6 +356,19 @@ def resolve_offload_strategy(config: Any) -> OffloadStrategy:
     return resolve_offload(config).strategy
 
 
+def offload_enabled(config: Any) -> bool:
+    """Return whether any CPU-offload policy is active."""
+    return resolve_offload(config).strategy is not OffloadStrategy.NONE
+
+
+def offload_streams_blocks(config: Any) -> bool:
+    """Return whether an active policy streams weights block by block."""
+    return resolve_offload(config).strategy in (
+        OffloadStrategy.LAYER_WISE,
+        OffloadStrategy.DISTRIBUTED_LAYER_WISE,
+    )
+
+
 def materialize_legacy_offload_flags(config: Any) -> OffloadStrategy:
     """Keep existing strategy readers working after resolving the compact API."""
     resolved = resolve_offload(config)

vllm_omni/diffusion/worker/diffusion_model_runner.py

@@ -54,7 +54,13 @@ 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.offloader.config import (
+    TEXT_ENCODER_COMPONENT,
+    OffloadStrategy,
+    offload_enabled,
+    resolve_offload,
+    resolve_offload_strategy,
+)
 from vllm_omni.diffusion.postprocess.device_reduction import prepare_diffusion_media_for_transport
 from vllm_omni.diffusion.registry import _NO_CACHE_ACCELERATION
 from vllm_omni.diffusion.request import OmniDiffusionRequest
@@ -321,13 +327,7 @@ class DiffusionModelRunner(DiffusionStagePayloadMixin):
             device=self.device,
         )
 
-        load_device = (
-            "cpu"
-            if self.od_config.enable_cpu_offload
-            or self.od_config.enable_layerwise_offload
-            or getattr(self.od_config, "enable_distributed_layerwise_offload", False)
-            else str(self.device)
-        )
+        load_device = "cpu" if offload_enabled(self.od_config) else str(self.device)
 
         def get_memory_context() -> AbstractContextManager[Any]:
             if memory_pool_context_fn is not None:
@@ -783,7 +783,7 @@ class DiffusionModelRunner(DiffusionStagePayloadMixin):
         # better perf. HSDP2's fully_shard pre-forward hooks need tensor version
         # counters, which inference tensors do not track.
         use_hsdp = od_config.parallel_config.use_hsdp
-        use_distributed_offload = getattr(self.od_config, "enable_distributed_layerwise_offload", False)
+        use_distributed_offload = resolve_offload_strategy(self.od_config) is OffloadStrategy.DISTRIBUTED_LAYER_WISE
         grad_context = torch.no_grad() if (use_hsdp or use_distributed_offload) else torch.inference_mode()
         with grad_context:
             for req in reqs: