[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),
         )