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

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

tests/diffusion/diffusion_backend/test_diffusers_backend.py

@@ -224,6 +224,7 @@ class TestPipelineArgumentsHandling:
     @pytest.mark.parametrize(
         "feature_id",
         [
+            "compact_offload",
             "cfg_parallel",
             "tensor_parallel",
             "ulysses",
@@ -235,7 +236,11 @@ class TestPipelineArgumentsHandling:
         ],
     )
     def test_adapter_guard_unsupported_feature(self, feature_id):
-        if feature_id == "cfg_parallel":
+        if feature_id == "compact_offload":
+            od_config = _make_od_config(
+                diffusion_offload_config={"mode": "module", "components": ["dit"]},
+            )
+        elif feature_id == "cfg_parallel":
             od_config = _make_od_config(
                 parallel_config=DiffusionParallelConfig(cfg_parallel_size=2, sequence_parallel_size=1),
                 cache_backend="none",

tests/diffusion/model_loader/test_diffusers_loader.py

@@ -734,6 +734,153 @@ def test_dlo_transfers_loader_plan_and_skips_ordinary_weight_loading(monkeypatch
     assert loader.take_host_weight_plan() is None
 
 
+def test_compact_rank_local_layer_offload_skips_dlo_mmap_planning(monkeypatch):
+    import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
+    from vllm_omni.diffusion.offloader.config import materialize_legacy_offload_flags
+
+    od_config = SimpleNamespace(
+        dtype=torch.float32,
+        parallel_config=SimpleNamespace(use_hsdp=False, tensor_parallel_size=1),
+        quantization_config=None,
+        diffusion_offload_config={"mode": "layer", "components": ["dit"]},
+        enable_cpu_offload=False,
+        enable_layerwise_offload=False,
+        enable_distributed_layerwise_offload=False,
+        dlo_use_allgather=True,
+        dlo_resident_layers=0,
+        pin_cpu_memory=True,
+        model="unused",
+    )
+    materialize_legacy_offload_flags(od_config)
+    loader = DiffusersPipelineLoader(LoadConfig(), od_config)
+    model = nn.Module()
+    model.transformer = nn.Linear(2, 2, bias=False)
+    calls: list[str] = []
+
+    loader._init_from_load_format = lambda *_args, **_kwargs: model  # type: ignore[method-assign]
+    loader.load_weights = lambda _model: calls.append("load")  # type: ignore[method-assign]
+    loader._process_weights_after_loading = lambda *_args: calls.append("process")  # type: ignore[method-assign]
+    loader._apply_skip_softmax_calibration = lambda _model: None  # type: ignore[method-assign]
+    monkeypatch.setattr(
+        loader_mod,
+        "build_checkpoint_mmap_plan",
+        lambda *_args, **_kwargs: pytest.fail("ordinary rank-local offload must not plan DLO mmap"),
+    )
+
+    assert loader.load_model(load_device="cpu") is model
+    assert calls == ["load", "process"]
+    assert loader.take_host_weight_plan() is None
+
+
+class _UnsupportedOnlineQuantMethod:
+    uses_meta_device = True
+
+
+def _compact_layer_offload_config(
+    components: list[str],
+    layer_options: dict[str, dict[str, str]],
+) -> SimpleNamespace:
+    return SimpleNamespace(
+        dtype=torch.float32,
+        parallel_config=SimpleNamespace(
+            use_hsdp=False,
+            tensor_parallel_size=1,
+            data_parallel_size=2,
+            sequence_parallel_size=1,
+        ),
+        quantization_config=None,
+        diffusion_offload_config={
+            "mode": "layer",
+            "components": components,
+            "layer_options": layer_options,
+        },
+        enable_cpu_offload=False,
+        enable_layerwise_offload=False,
+        enable_distributed_layerwise_offload=False,
+        dlo_use_allgather=True,
+        dlo_resident_layers=0,
+        pin_cpu_memory=False,
+        host_weight_runtime_mode="disabled",
+        model="unused",
+    )
+
+
+def _stub_ordinary_loader(loader, model, monkeypatch, loader_mod) -> None:
+    loader._init_from_load_format = lambda *_args, **_kwargs: model  # type: ignore[method-assign]
+    loader.load_weights = lambda _model: None  # type: ignore[method-assign]
+    loader._process_weights_after_loading = lambda *_args: None  # type: ignore[method-assign]
+    loader._apply_skip_softmax_calibration = lambda _model: None  # type: ignore[method-assign]
+    monkeypatch.setattr(
+        loader_mod,
+        "build_checkpoint_mmap_plan",
+        lambda *_args, **_kwargs: HostWeightPlanResult(None, "test fallback"),
+    )
+
+
+def test_text_encoder_allgather_rejects_unsupported_online_quantization(monkeypatch):
+    config = _compact_layer_offload_config(
+        ["text_encoder"],
+        {"text_encoder": {"weight_transfer": "allgather"}},
+    )
+    loader = DiffusersPipelineLoader(LoadConfig(), config)
+    model = nn.Module()
+    model.text_encoder = nn.Linear(2, 2, bias=False)
+    model.text_encoder.quant_method = _UnsupportedOnlineQuantMethod()
+    loader._init_from_load_format = lambda *_args, **_kwargs: model  # type: ignore[method-assign]
+
+    with pytest.raises(ValueError, match="unsupported online methods: _UnsupportedOnlineQuantMethod"):
+        loader.load_model(load_device="cpu")
+
+
+def test_dit_allgather_ignores_unsupported_quantization_on_unselected_encoder(monkeypatch):
+    import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
+
+    config = _compact_layer_offload_config(
+        ["dit"],
+        {"dit": {"weight_transfer": "allgather"}},
+    )
+    loader = DiffusersPipelineLoader(LoadConfig(), config)
+    model = nn.Module()
+    model.transformer = nn.Linear(2, 2, bias=False)
+    model.text_encoder = nn.Linear(2, 2, bias=False)
+    model.text_encoder.quant_method = _UnsupportedOnlineQuantMethod()
+    _stub_ordinary_loader(loader, model, monkeypatch, loader_mod)
+
+    assert loader.load_model(load_device="cpu") is model
+
+
+def test_rank_local_encoder_quantization_does_not_mark_dit_mmap_plan_online(monkeypatch):
+    import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
+
+    config = _compact_layer_offload_config(
+        ["dit", "text_encoder"],
+        {
+            "dit": {"weight_transfer": "allgather"},
+            "text_encoder": {"weight_transfer": "rank-local"},
+        },
+    )
+    loader = DiffusersPipelineLoader(LoadConfig(), config)
+    model = nn.Module()
+    model.transformer = nn.Linear(2, 2, bias=False)
+    model.text_encoder = nn.Linear(2, 2, bias=False)
+    model.text_encoder.quant_method = _UnsupportedOnlineQuantMethod()
+    planned_online_quantization: list[bool] = []
+
+    loader._init_from_load_format = lambda *_args, **_kwargs: model  # type: ignore[method-assign]
+    loader.load_weights = lambda _model: None  # type: ignore[method-assign]
+    loader._process_weights_after_loading = lambda *_args: None  # type: ignore[method-assign]
+    loader._apply_skip_softmax_calibration = lambda _model: None  # type: ignore[method-assign]
+
+    def capture_plan(*_args, **kwargs):
+        planned_online_quantization.append(kwargs["online_quantization"])
+        return HostWeightPlanResult(None, "test fallback")
+
+    monkeypatch.setattr(loader_mod, "build_checkpoint_mmap_plan", capture_plan)
+
+    assert loader.load_model(load_device="cpu") is model
+    assert planned_online_quantization == [False]
+
+
 def test_dlo_plan_loads_component_sources_outside_planned_dit(monkeypatch):
     import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
 
@@ -866,7 +1013,7 @@ def test_dlo_allgather_online_fp8_uses_ordinary_loader(monkeypatch):
     )
 
     assert loader.load_model(load_device="cpu", device=torch.device("cpu")) is model
-    assert allowlist_models == [model]
+    assert allowlist_models == [model.transformer]
     assert calls == [("load", True), "process"]
     assert loader.take_host_weight_plan() is None
 
@@ -908,7 +1055,7 @@ def test_dlo_allgather_online_int8_uses_ordinary_loader(monkeypatch):
     )
 
     assert loader.load_model(load_device="cpu", device=torch.device("cpu")) is model
-    assert allowlist_models == [model]
+    assert allowlist_models == [model.transformer]
     assert calls == [("load", True), "process"]
     assert loader.take_host_weight_plan() is None
 
@@ -1024,7 +1171,7 @@ def test_dlo_allgather_online_mxfp8_uses_ordinary_loader(monkeypatch):
     )
 
     assert loader.load_model(load_device="cpu", device=torch.device("cpu")) is model
-    assert allowlist_models == [model]
+    assert allowlist_models == [model.transformer]
     assert calls == [("load", True), "process"]
     assert loader.take_host_weight_plan() is None
 

tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py

@@ -30,6 +30,36 @@ def test_pipeline_declares_layerwise_offload_components() -> None:
     assert Cosmos3OmniDiffusersPipeline._resident_modules == []
     assert hasattr(Cosmos3OmniDiffusersPipeline, "enable_omni_model_cpu_offload")
 
+    from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import (
+        Cosmos3LanguageModel,
+        Cosmos3VFMTransformer,
+    )
+
+    assert Cosmos3LanguageModel._layerwise_offload_blocks_attrs == ["layers"]
+    assert Cosmos3VFMTransformer._layerwise_offload_blocks_attrs == ["gen_layers"]
+
+
+def test_component_selective_model_offload_fails_before_component_loading(monkeypatch) -> None:
+    from vllm_omni.diffusion.models.cosmos3 import pipeline_cosmos3 as pipeline_module
+
+    monkeypatch.setattr(
+        pipeline_module,
+        "get_local_device",
+        lambda: pytest.fail("component loading must not start for an unsupported selector"),
+    )
+    config = SimpleNamespace(
+        diffusion_offload_config={"mode": "module", "components": ["dit"]},
+        enable_cpu_offload=False,
+        enable_layerwise_offload=False,
+        enable_distributed_layerwise_offload=False,
+        dlo_use_allgather=True,
+        dlo_resident_layers=0,
+        pin_cpu_memory=True,
+    )
+
+    with pytest.raises(ValueError, match="does not support the dit/text_encoder component selector"):
+        pipeline_module.Cosmos3OmniDiffusersPipeline(od_config=config)
+
 
 class StubScheduler:
     def __init__(

tests/diffusion/models/magi2/test_native_dlo_allgather.py

@@ -220,7 +220,9 @@ def _load_mmap_transform_and_reconstruct(
         tensor_transforms=backend._mmap_transforms_by_tensor_id,
     )
     assert all(not shard.requires_grad for shard in cpu_shards.values())
-    assert all(parameter.numel() == 0 for parameter in target_block.parameters())
+    # Shard preparation is non-destructive; initialize_hook clears source
+    # storage only after every shard and metadata entry has been prepared.
+    assert all(parameter.numel() > 0 for parameter in target_block.parameters())
 
     for dtype, local_shard in cpu_shards.items():
         gathered = torch.empty(

tests/diffusion/models/magi2/test_pipeline_magi2.py

@@ -578,6 +578,8 @@ class _FakeLayerwiseHook:
     def __init__(self) -> None:
         self._prev_hook = None
         self.current_slot = 0
+        self.dp_size = 1
+        self.rank_local_mmap = False
         self.prefetch_layer = Mock()
         self.get_weights = Mock()
 
@@ -625,7 +627,7 @@ def test_dlo_no_allgather_enable_preserves_staged_aux_and_streams_blocks(
     assert len(hooks) == len(pipeline.transformer.block)
     assert all(stager.offload.call_count == 1 for stager in stagers)
     assert all(stager.load.call_count == 0 for stager in stagers)
-    hooks[-1].prefetch_layer.assert_called_once_with(
-        slot=hooks[0].current_slot,
-        non_blocking=False,
-    )
+    # First prefetch stays lazy so ranks with different local groups cannot
+    # enter collectives in a different order during startup.
+    for hook in hooks:
+        hook.prefetch_layer.assert_not_called()

tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py

@@ -487,6 +487,10 @@ def test_dlo_offload_plan_includes_token_refiner():
 
     assert MiniMaxH3Pipeline._offload_plan.offload_submodules == {"token_refiner": "blocks"}
     assert MiniMaxH3Pipeline._offload_plan.resident_dit_paths == frozenset({"transformer"})
+    assert MiniMaxH3Pipeline._offload_plan.encoder_component_types == {"text_encoder": "text_encoder"}
+    assert MiniMaxH3Pipeline._offload_plan.encoder_block_attrs == {
+        "text_encoder": ("vision.blocks", "text_model.layers")
+    }
 
 
 def test_joint_postprocess_is_multiprocessing_picklable():
@@ -1185,6 +1189,12 @@ def test_text_encoder_stub_constructs_without_group_or_weights():
     from vllm_omni.diffusion.models.minimax_h3.encoder import (
         MiniMaxH3Qwen3VLEncoder,
     )
+    from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import (
+        MiniMaxH3Pipeline,
+    )
+    from vllm_omni.diffusion.offloader.component_utils import (
+        get_encoder_block_groups,
+    )
 
     encoder = MiniMaxH3Qwen3VLEncoder(
         "/nonexistent/text_encoder",
@@ -1197,6 +1207,25 @@ def test_text_encoder_stub_constructs_without_group_or_weights():
     # The stub has no parameters, so it never contributes to the runner's
     # strict missing-parameter check on non-encoder ranks.
     assert list(encoder.named_parameters()) == []
+    assert (
+        get_encoder_block_groups(
+            encoder,
+            "text_encoder",
+            MiniMaxH3Pipeline._offload_plan,
+            strict=True,
+        )
+        == []
+    )
+
+    # An arbitrary empty module is not silently treated as a distributed
+    # stub; explicit plans still validate their declared paths.
+    with pytest.raises(ValueError, match=r"text_encoder\.vision\.blocks was not found"):
+        get_encoder_block_groups(
+            nn.Module(),
+            "text_encoder",
+            MiniMaxH3Pipeline._offload_plan,
+            strict=True,
+        )
 
 
 def test_global_quant_config_is_shared_by_dit_and_encoder():
@@ -1303,10 +1332,38 @@ def test_model_offload_uses_hooked_text_encoder_call():
 
 
 @pytest.mark.parametrize(
-    "offload_flag",
-    ["enable_layerwise_offload", "enable_distributed_layerwise_offload"],
+    ("enable_layerwise", "enable_distributed"),
+    [(True, False), (False, True)],
 )
-def test_layerwise_offload_releases_text_encoder(offload_flag):
+def test_legacy_layer_offload_preserves_minimax_stage_lifecycle(enable_layerwise, enable_distributed):
+    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+
+    pipeline = object.__new__(MiniMaxH3Pipeline)
+    torch.nn.Module.__init__(pipeline)
+    pipeline.od_config = SimpleNamespace(
+        diffusion_offload_config=None,
+        enable_cpu_offload=False,
+        enable_layerwise_offload=enable_layerwise,
+        enable_distributed_layerwise_offload=enable_distributed,
+    )
+    pipeline.text_encoder = Mock()
+    expected = torch.ones(2, 3)
+    pipeline.text_encoder.encode_ids.return_value = expected
+
+    actual = pipeline._encode_text_hidden(torch.tensor([1, 2]), {})
+
+    assert actual is expected
+    pipeline.text_encoder.load_to_device.assert_called_once_with()
+    pipeline.text_encoder.offload_to_cpu.assert_called_once_with()
+
+    vae = Mock()
+    with pipeline._component_on_device(vae):
+        pass
+    vae.load_to_device.assert_called_once_with()
+    vae.offload_to_cpu.assert_called_once_with()
+
+
+def test_layerwise_encoder_selection_releases_text_encoder():
     from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
 
     pipeline = object.__new__(MiniMaxH3Pipeline)
@@ -1315,8 +1372,11 @@ def test_layerwise_offload_releases_text_encoder(offload_flag):
         enable_cpu_offload=False,
         enable_layerwise_offload=False,
         enable_distributed_layerwise_offload=False,
+        diffusion_offload_config={
+            "mode": "layer",
+            "components": ["text_encoder"],
+        },
     )
-    setattr(pipeline.od_config, offload_flag, True)
     pipeline.text_encoder = Mock()
     expected = torch.ones(2, 3)
     pipeline.text_encoder.encode_ids.return_value = expected
@@ -1328,7 +1388,7 @@ def test_layerwise_offload_releases_text_encoder(offload_flag):
     pipeline.text_encoder.offload_to_cpu.assert_called_once_with()
 
 
-def test_distributed_layerwise_offload_releases_text_encoder():
+def test_layerwise_dit_only_keeps_text_encoder_resident():
     from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
 
     pipeline = object.__new__(MiniMaxH3Pipeline)
@@ -1336,7 +1396,8 @@ def test_distributed_layerwise_offload_releases_text_encoder():
     pipeline.od_config = SimpleNamespace(
         enable_cpu_offload=False,
         enable_layerwise_offload=False,
-        enable_distributed_layerwise_offload=True,
+        enable_distributed_layerwise_offload=False,
+        diffusion_offload_config={"mode": "layer", "components": ["dit"]},
     )
     pipeline.text_encoder = Mock()
     expected = torch.ones(2, 3)
@@ -1346,25 +1407,30 @@ def test_distributed_layerwise_offload_releases_text_encoder():
 
     assert actual is expected
     pipeline.text_encoder.load_to_device.assert_called_once_with()
-    pipeline.text_encoder.offload_to_cpu.assert_called_once_with()
+    pipeline.text_encoder.offload_to_cpu.assert_not_called()
 
 
-def test_distributed_layerwise_offload_stages_vae_component():
+def test_dit_encoder_selection_keeps_vae_resident():
     from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
 
     pipeline = object.__new__(MiniMaxH3Pipeline)
     torch.nn.Module.__init__(pipeline)
     pipeline.od_config = SimpleNamespace(
+        enable_cpu_offload=False,
         enable_layerwise_offload=False,
-        enable_distributed_layerwise_offload=True,
+        enable_distributed_layerwise_offload=False,
+        diffusion_offload_config={
+            "mode": "layer",
+            "components": ["dit", "text_encoder"],
+        },
     )
     component = Mock()
 
     with pipeline._component_on_device(component):
-        component.load_to_device.assert_called_once_with()
-        component.offload_to_cpu.assert_not_called()
+        pass
 
-    component.offload_to_cpu.assert_called_once_with()
+    component.load_to_device.assert_not_called()
+    component.offload_to_cpu.assert_not_called()
 
 
 def test_distributed_layerwise_resident_blocks_are_stage_scoped():
@@ -1413,57 +1479,6 @@ def test_distributed_layerwise_resident_blocks_can_be_skipped():
     controller.offload_resident_layers.assert_not_called()
 
 
-def test_encoder_layerwise_offload_keeps_tp_blocks_rank_local(monkeypatch):
-    from vllm_omni.diffusion.models.minimax_h3.encoder import (
-        MiniMaxH3Qwen3VLEncoder,
-    )
-
-    class Stack(torch.nn.Module):
-        def __init__(self, count):
-            super().__init__()
-            self.blocks = torch.nn.ModuleList([torch.nn.Linear(2, 2) for _ in range(count)])
-
-    class TextStack(torch.nn.Module):
-        def __init__(self, count):
-            super().__init__()
-            self.layers = torch.nn.ModuleList([torch.nn.Linear(2, 2) for _ in range(count)])
-
-    hooks = []
-
-    def fake_apply(block, next_block, device, stream, pin_memory):
-        hook = SimpleNamespace(
-            block=block,
-            next_block=next_block,
-            device=device,
-            stream=stream,
-            pin_memory=pin_memory,
-            _prev_hook=None,
-            offload_layer=Mock(),
-        )
-        hooks.append(hook)
-        return hook
-
-    monkeypatch.setattr(
-        "vllm_omni.diffusion.offloader.layerwise_backend.apply_block_hook",
-        fake_apply,
-    )
-    monkeypatch.setattr(
-        "vllm_omni.platforms.current_omni_platform.Stream",
-        Mock(return_value="copy-stream"),
-    )
-    encoder = object.__new__(MiniMaxH3Qwen3VLEncoder)
-    torch.nn.Module.__init__(encoder)
-    encoder.device_target = torch.device("cpu")
-    encoder.vision = Stack(2)
-    encoder.text_model = TextStack(3)
-
-    encoder.enable_omni_layerwise_offload(pin_memory=False)
-
-    assert len(hooks) == 5
-    assert all(hook._prev_hook is not None for hook in hooks)
-    assert all(hook.device == torch.device("cpu") for hook in hooks)
-
-
 def test_video_vae_keeps_reference_fp32_weights(monkeypatch):
     from vllm_omni.diffusion.models.minimax_h3 import vae as vae_module
 

tests/diffusion/models/minimax_h3/test_minimax_h3_dlo_lifecycle.py

@@ -1,6 +1,8 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
+from types import SimpleNamespace
+
 import pytest
 import torch
 
@@ -11,6 +13,44 @@ class _OffloadAbort(BaseException):
     pass
 
 
+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
+    from vllm_omni.diffusion.offloader.config import OffloadStrategy
+
+    args = SimpleNamespace(
+        model="/fake/MiniMax-H3/FL2VA",
+        mode="dlo-no-allgather",
+        dp_size=2,
+        tp_size=2,
+        batch_wait_ms=500.0,
+        init_timeout=1800.0,
+    )
+    kwargs = engine_kwargs(args)
+    od_config = SimpleNamespace(
+        diffusion_offload_config=None,
+        enable_cpu_offload=False,
+        enable_layerwise_offload=False,
+        enable_distributed_layerwise_offload=kwargs["enable_distributed_layerwise_offload"],
+        dlo_use_allgather=kwargs["dlo_use_allgather"],
+        dlo_resident_layers=kwargs["dlo_resident_layers"],
+        dlo_host_registration_limit_gib=0.0,
+        host_weight_runtime_mode="disabled",
+        pin_cpu_memory=True,
+        parallel_config=SimpleNamespace(
+            data_parallel_size=kwargs["data_parallel_size"],
+            sequence_parallel_size=1,
+            use_hsdp=False,
+        ),
+    )
+
+    config = OffloadConfig.from_od_config(od_config)
+
+    assert config.strategy is OffloadStrategy.DISTRIBUTED_LAYER_WISE
+    assert config.dlo_resident_layers == 0
+    assert not config.uses_allgather("dit")
+
+
 def test_encoder_non_block_children_use_one_shared_snapshot_stager(monkeypatch, mocker):
     from vllm_omni.diffusion.models.minimax_h3 import encoder as encoder_module
 
@@ -61,6 +101,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._model_cpu_offload_modules = []
@@ -83,6 +124,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._model_cpu_offload_modules = []

tests/diffusion/models/minimax_h3/test_minimax_h3_lora.py

@@ -247,6 +247,55 @@ def test_h3_turbo_allows_distributed_layerwise_offload(monkeypatch):
     assert captured["unsupported_offload_mode"] is None
 
 
+@pytest.mark.parametrize(
+    ("mode", "components", "expected_mode"),
+    [
+        ("module", ["text_encoder"], None),
+        ("layer", ["text_encoder"], None),
+        ("module", ["dit"], "model-level CPU offload"),
+        ("layer", ["dit"], "layerwise offload"),
+    ],
+)
+def test_h3_turbo_offload_compatibility_follows_dit_selection(
+    monkeypatch,
+    mode,
+    components,
+    expected_mode,
+):
+    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+    from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as pipeline_module
+
+    pipeline = object.__new__(MiniMaxH3Pipeline)
+    torch.nn.Module.__init__(pipeline)
+    pipeline.partition = "fl2va"
+    pipeline._turbo_lora_adapter_ids = set()
+    pipeline._native_lora_adapter_ids = set()
+    pipeline._lora_sigma_schedules = {}
+    pipeline.od_config = SimpleNamespace(
+        diffusion_offload_config={"mode": mode, "components": components},
+        enable_cpu_offload=False,
+        enable_layerwise_offload=False,
+        enable_distributed_layerwise_offload=False,
+        dlo_use_allgather=True,
+        dlo_resident_layers=0,
+        pin_cpu_memory=True,
+    )
+    captured = {}
+
+    def load_turbo(**kwargs):
+        captured.update(kwargs)
+        return object(), object()
+
+    monkeypatch.setattr(pipeline_module, "load_minimax_h3_turbo_lora", load_turbo)
+
+    assert pipeline._load_diffusion_lora_adapter(
+        lora_request=_request("turbo"),
+        lora_path="turbo",
+        dtype=torch.bfloat16,
+    )
+    assert captured["unsupported_offload_mode"] == expected_mode
+
+
 def test_h3_turbo_accepts_only_the_declared_v1_artifact(tmp_path):
     unsupported = tmp_path / "minimax_h3_fl2v_turbo_8step_v1.0_bf16.safetensors"
     _write_tiny_turbo(unsupported)

tests/diffusion/models/minimax_h3/test_minimax_h3_offload.py

@@ -1,8 +1,9 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 from contextlib import contextmanager
 from types import SimpleNamespace
+from typing import Any
 from unittest.mock import Mock
 
 import pytest
@@ -111,6 +112,66 @@ def test_h3_model_cpu_offload_registers_direct_vae_stages(monkeypatch):
     remove_offload.assert_called_once_with([*dits, *stages])
 
 
+def test_h3_model_cpu_offload_keeps_unselected_vaes_resident(monkeypatch):
+    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+    from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as module
+
+    pipeline = object.__new__(MiniMaxH3Pipeline)
+    torch.nn.Module.__init__(pipeline)
+    pipeline.transformer = torch.nn.Linear(2, 2)
+    pipeline.transformers_ref = torch.nn.Linear(2, 2)
+    pipeline.text_encoder = torch.nn.Linear(2, 2)
+    pipeline.video_vae = torch.nn.Linear(2, 2)
+    pipeline.audio_vae = torch.nn.Linear(2, 2)
+    apply_offload = Mock()
+    monkeypatch.setattr(module, "apply_sequential_offload", apply_offload)
+
+    pipeline.enable_omni_model_cpu_offload(
+        device=torch.device("cpu"),
+        pin_memory=False,
+        use_hsdp=False,
+        offload_components=frozenset({"dit", "text_encoder"}),
+    )
+
+    dits = [pipeline.transformer, pipeline.transformers_ref]
+    stages = [pipeline.text_encoder, pipeline.video_vae, pipeline.audio_vae]
+    apply_offload.assert_called_once_with(
+        dit_modules=dits,
+        encoder_modules=stages,
+        device=torch.device("cpu"),
+        pin_memory=False,
+        use_hsdp=False,
+        offload_initial_dits=True,
+        offload_dit_modules=dits,
+        offload_encoder_modules=[pipeline.text_encoder],
+    )
+
+
+def test_h3_model_cpu_offload_rejects_unloaded_selected_encoder(monkeypatch):
+    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+    from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as module
+
+    pipeline = object.__new__(MiniMaxH3Pipeline)
+    torch.nn.Module.__init__(pipeline)
+    pipeline.transformer = torch.nn.Linear(2, 2)
+    pipeline.transformers_ref = torch.nn.Linear(2, 2)
+    pipeline.text_encoder = None
+    pipeline.video_vae = torch.nn.Linear(2, 2)
+    pipeline.audio_vae = torch.nn.Linear(2, 2)
+    apply_offload = Mock()
+    monkeypatch.setattr(module, "apply_sequential_offload", apply_offload)
+
+    with pytest.raises(ValueError, match="no loaded text encoder"):
+        pipeline.enable_omni_model_cpu_offload(
+            device=torch.device("cpu"),
+            pin_memory=False,
+            use_hsdp=False,
+            offload_components=frozenset({"text_encoder"}),
+        )
+
+    apply_offload.assert_not_called()
+
+
 @pytest.mark.parametrize("decode_fails", [False, True])
 def test_h3_model_cpu_offload_scopes_direct_vae_call(monkeypatch, decode_fails):
     from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
@@ -119,7 +180,7 @@ def test_h3_model_cpu_offload_scopes_direct_vae_call(monkeypatch, decode_fails):
     pipeline = object.__new__(MiniMaxH3Pipeline)
     torch.nn.Module.__init__(pipeline)
     component = torch.nn.Linear(2, 2)
-    events = []
+    events: list[tuple[Any, ...]] = []
 
     @contextmanager
     def record_component(value):
@@ -262,7 +323,7 @@ def test_h3_model_cpu_offload_shares_embedded_and_standalone_audio_scope(monkeyp
     torch.nn.Module.__init__(pipeline)
     pipeline.audio_vae = Mock()
     pipeline._model_cpu_offload_modules = [pipeline.audio_vae]
-    events = []
+    events: list[tuple[Any, ...]] = []
 
     @contextmanager
     def record_component(value):

tests/diffusion/models/sana_video/test_cache_offload.py

@@ -16,7 +16,7 @@ from tests.diffusion.models.sana_video.test_transformer_sana_video import (  # n
     _TINY_CONFIG,
     _init_distributed,
 )
-from tests.diffusion.offloader.test_layerwise_backend import DummyEvent, DummyStream, dummy_stream
+from tests.diffusion.offloader.helpers import DummyEvent, DummyStream, dummy_stream
 from vllm_omni.diffusion.attention import selector as attention_selector
 from vllm_omni.diffusion.attention.backends.sdpa import SDPABackend
 from vllm_omni.diffusion.cache.cachedit import CacheDiTBackend
@@ -226,7 +226,7 @@ def test_component_loading_uses_loader_device_not_runtime_device(monkeypatch):
     monkeypatch.setattr(
         pipeline_module,
         "_load_json",
-        lambda _model, filename, _local: ({"vae": [None, "FakeVAE"]} if filename == "model_index.json" else {}),
+        lambda _model, filename, _local: {"vae": [None, "FakeVAE"]} if filename == "model_index.json" else {},
     )
     monkeypatch.setattr(
         pipeline_module,

tests/diffusion/models/sensenova_u1/test_sensenova_u1_paged_decode.py

@@ -173,9 +173,14 @@ def test_length_lives_in_device_tensors():
 # ---------------------------------------------------------------------------
 
 cuda_only = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+vllm_flash_attn_only = pytest.mark.skipif(
+    current_omni_platform.is_rocm(),
+    reason="vLLM's bundled FlashAttention extension is CUDA-only",
+)
 
 
 @cuda_only
+@vllm_flash_attn_only
 @pytest.mark.cuda
 @pytest.mark.L4
 def test_paged_attention_matches_sdpa_over_the_same_buffer():
@@ -217,6 +222,7 @@ def test_paged_attention_matches_sdpa_over_the_same_buffer():
 
 
 @cuda_only
+@vllm_flash_attn_only
 @pytest.mark.cuda
 @pytest.mark.L4
 def test_a_captured_graph_follows_a_later_length():
@@ -278,6 +284,7 @@ def test_a_captured_graph_follows_a_later_length():
 
 
 @cuda_only
+@vllm_flash_attn_only
 @pytest.mark.cuda
 @pytest.mark.L4
 def test_a_captured_graph_writes_each_step_to_its_own_slot():
@@ -339,6 +346,7 @@ def test_a_captured_graph_writes_each_step_to_its_own_slot():
 
 
 @cuda_only
+@vllm_flash_attn_only
 @pytest.mark.cuda
 @pytest.mark.L4
 def test_a_whole_decode_loop_matches_the_unpaged_path():

tests/diffusion/models/wan2_2/test_wan22_i2v_pipeline.py

@@ -9,6 +9,7 @@ from PIL import Image
 from torch import nn
 
 from tests.diffusion.models.wan2_2.conftest import StubScheduler, StubTransformer, StubVAE, noop_progress_bar
+from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import _WAN_TEXT_ENCODER_OFFLOAD_PLAN
 from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_i2v import (
     Wan22I2VPipeline,
     get_wan22_i2v_post_process_func,
@@ -30,6 +31,10 @@ def test_wan22_i2v_postprocess_honors_request_output_type() -> None:
     assert output is video
 
 
+def test_i2v_pipeline_declares_text_encoder_offload_blocks() -> None:
+    assert Wan22I2VPipeline._offload_plan is _WAN_TEXT_ENCODER_OFFLOAD_PLAN
+
+
 def _make_i2v_pipeline(*, expand_timesteps: bool) -> Wan22I2VPipeline:
     pipeline = object.__new__(Wan22I2VPipeline)
     nn.Module.__init__(pipeline)

tests/diffusion/models/wan2_2/test_wan22_pipeline_helpers.py

@@ -18,6 +18,15 @@ from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import (
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
 
 
+def test_t2v_pipeline_declares_text_encoder_offload_blocks() -> None:
+    plan = wan22_module.Wan22Pipeline._offload_plan
+
+    assert plan.encoder_component_types == {"text_encoder": "text_encoder"}
+    assert plan.encoder_block_attrs == {"text_encoder": ("encoder.block",)}
+    assert plan.encoder_dlo_weight_replication == frozenset({"text_encoder"})
+    assert plan.on_demand_component_paths == frozenset()
+
+
 class _LatentDist:
     def sample(self, generator):
         assert isinstance(generator, torch.Generator)

tests/diffusion/models/wan2_2/test_wan22_s2v_pipeline.py

@@ -392,52 +392,31 @@ def test_encode_audio_skips_unshard_reshard_when_not_fsdp():
     assert "audio_emb" in result
 
 
-def test_s2v_pipeline_skips_cpu_offload_when_hsdp_enabled():
-    """Test that transformer.to('cpu') is NOT called when HSDP is active."""
-    from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_s2v import Wan22S2VPipeline
-
+@pytest.mark.parametrize(
+    ("components", "use_hsdp", "expected"),
+    [
+        (None, False, True),
+        (["dit"], False, True),
+        (["text_encoder"], False, False),
+        (["dit"], True, False),
+    ],
+)
+def test_s2v_dit_release_honors_component_selection(components, use_hsdp, expected):
     pipeline = object.__new__(Wan22S2VPipeline)
     nn.Module.__init__(pipeline)
+    compact = None if components is None else {"mode": "module", "components": components}
+    pipeline.od_config = SimpleNamespace(
+        diffusion_offload_config=compact,
+        enable_cpu_offload=True,
+        enable_layerwise_offload=False,
+        enable_distributed_layerwise_offload=False,
+        dlo_use_allgather=True,
+        dlo_resident_layers=0,
+        pin_cpu_memory=True,
+        parallel_config=SimpleNamespace(use_hsdp=use_hsdp),
+    )
 
-    od_config = MagicMock()
-    od_config.enable_cpu_offload = True
-    parallel_config = MagicMock()
-    parallel_config.use_hsdp = True
-    od_config.parallel_config = parallel_config
-    pipeline.od_config = od_config
-
-    mock_transformer = MagicMock()
-    pipeline.transformer = mock_transformer
-
-    # Simulate the offload decision from the forward loop
-    if pipeline.od_config.enable_cpu_offload and not getattr(pipeline.od_config.parallel_config, "use_hsdp", False):
-        pipeline.transformer.to("cpu")
-
-    mock_transformer.to.assert_not_called()
-
-
-def test_s2v_pipeline_allows_cpu_offload_when_hsdp_disabled():
-    """Test that transformer.to('cpu') IS called when HSDP is not active."""
-    from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_s2v import Wan22S2VPipeline
-
-    pipeline = object.__new__(Wan22S2VPipeline)
-    nn.Module.__init__(pipeline)
-
-    od_config = MagicMock()
-    od_config.enable_cpu_offload = True
-    parallel_config = MagicMock()
-    parallel_config.use_hsdp = False
-    od_config.parallel_config = parallel_config
-    pipeline.od_config = od_config
-
-    mock_transformer = MagicMock()
-    pipeline.transformer = mock_transformer
-
-    # Simulate the offload decision from the forward loop
-    if pipeline.od_config.enable_cpu_offload and not getattr(pipeline.od_config.parallel_config, "use_hsdp", False):
-        pipeline.transformer.to("cpu")
-
-    mock_transformer.to.assert_called_once_with("cpu")
+    assert pipeline._should_release_dit_before_decode() is expected
 
 
 def test_s2v_pipeline_hsdp_forward_complete_process():

tests/diffusion/models/wan2_2/test_wan22_vace_pipeline.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+from dataclasses import dataclass, field
 from types import SimpleNamespace
 
 import pytest
@@ -20,6 +21,46 @@ from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
 
 
+def test_vace_default_flow_shift_preserves_compact_offload_runtime_state(monkeypatch) -> None:
+    @dataclass
+    class RuntimeConfig:
+        flow_shift: float | None = None
+        diffusion_offload_config: dict[str, object] = field(
+            default_factory=lambda: {
+                "mode": "layer",
+                "components": ["text_encoder"],
+            }
+        )
+        enable_layerwise_offload: bool = True
+        post_init_calls: int = 0
+
+        def __post_init__(self) -> None:
+            self.post_init_calls += 1
+
+    config = RuntimeConfig()
+    resolved_offload = object()
+    config._resolved_diffusion_offload = resolved_offload
+    config._diffusion_offload_flags_materialized = True
+    captured = {}
+
+    def capture_base_init(self, *, od_config, prefix="") -> None:
+        captured.update(od_config=od_config, prefix=prefix)
+
+    monkeypatch.setattr(Wan22VACEPipeline.__mro__[1], "__init__", capture_base_init)
+
+    Wan22VACEPipeline(od_config=config, prefix="stage")
+
+    forwarded = captured["od_config"]
+    assert forwarded is not config
+    assert config.flow_shift is None
+    assert forwarded.flow_shift == 3.0
+    assert forwarded.diffusion_offload_config is config.diffusion_offload_config
+    assert forwarded._resolved_diffusion_offload is resolved_offload
+    assert forwarded._diffusion_offload_flags_materialized is True
+    assert forwarded.post_init_calls == 1
+    assert captured["prefix"] == "stage"
+
+
 def _make_vace_pipeline() -> Wan22VACEPipeline:
     pipeline = object.__new__(Wan22VACEPipeline)
     nn.Module.__init__(pipeline)

tests/diffusion/offloader/conftest.py

@@ -0,0 +1,26 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Shared runtime fixtures for diffusion offloader tests."""
+
+import gc
+
+import pytest
+import torch.distributed as dist
+
+from tests.helpers.runtime import get_distributed_init_method
+from vllm_omni.platforms import current_omni_platform
+
+
+@pytest.fixture(scope="module")
+def dist_group():
+    dist.init_process_group("gloo", rank=0, world_size=1, init_method=get_distributed_init_method())
+    try:
+        yield
+    finally:
+        if dist.is_initialized():
+            dist.destroy_process_group()
+        gc.collect()
+        if current_omni_platform.is_available():
+            current_omni_platform.empty_cache()
+            current_omni_platform.synchronize()

tests/diffusion/offloader/helpers.py

@@ -0,0 +1,107 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Shared model fixtures for diffusion offloader tests."""
+
+from contextlib import contextmanager
+
+import torch
+from torch import nn
+
+
+class DummyStream:
+    def wait_stream(self, _stream) -> None:
+        return None
+
+    def wait_event(self, _event) -> None:
+        return None
+
+
+class DummyEvent:
+    def record(self, _stream) -> None:
+        return None
+
+    def synchronize(self) -> None:
+        return None
+
+
+@contextmanager
+def dummy_stream(_stream):
+    yield None
+
+
+def patch_offload_runtime(monkeypatch, platform, *, synchronize: bool = False) -> None:
+    monkeypatch.setattr(platform, "Stream", DummyStream)
+    monkeypatch.setattr(platform, "Event", DummyEvent)
+    monkeypatch.setattr(platform, "current_stream", lambda: DummyStream())
+    monkeypatch.setattr(platform, "stream", dummy_stream)
+    if synchronize:
+        monkeypatch.setattr(platform, "synchronize", lambda: None)
+        monkeypatch.setattr(platform, "empty_cache", lambda: None)
+
+
+class _DummyBlock(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.weight = nn.Parameter(torch.randn(10, 10))
+
+
+class _SingleBlockModel(nn.Module):
+    _layerwise_offload_blocks_attrs = ["blocks"]
+
+    def __init__(self, num_blocks: int = 3):
+        super().__init__()
+        self.blocks = nn.ModuleList([_DummyBlock() for _ in range(num_blocks)])
+
+
+class _StagedEncoder(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.vision = nn.Module()
+        self.vision.blocks = nn.ModuleList([_DummyBlock(), _DummyBlock()])
+        self.text_model = nn.Module()
+        self.text_model.layers = nn.ModuleList([_DummyBlock(), _DummyBlock()])
+        self.load_calls = 0
+        self.offload_calls = 0
+        self.to_calls = 0
+
+    def load_to_device(self):
+        self.load_calls += 1
+
+    def offload_to_cpu(self):
+        self.offload_calls += 1
+        for hook in getattr(self, "_omni_layerwise_hooks", []):
+            hook.offload_layer()
+
+    def to(self, *args, **kwargs):
+        self.to_calls += 1
+        return super().to(*args, **kwargs)
+
+
+class _StagedVAE(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.proj = nn.Linear(2, 2)
+        self.offload_calls = 0
+        self.to_calls = 0
+
+    def load_to_device(self):
+        return None
+
+    def offload_to_cpu(self):
+        self.offload_calls += 1
+        return self.to("cpu")
+
+    def to(self, *args, **kwargs):
+        self.to_calls += 1
+        return super().to(*args, **kwargs)
+
+
+class _PlainEncoder(nn.Module):
+    """Standard encoder with no offload-specific lifecycle methods."""
+
+    def __init__(self):
+        super().__init__()
+        self.encoder = nn.Module()
+        self.encoder.block = nn.ModuleList([_DummyBlock(), _DummyBlock()])
+        self.final_norm = nn.Linear(2, 2)

tests/diffusion/offloader/test_distributed_layerwise_backend.py

@@ -6,8 +6,8 @@
 import gc
 import json
 import weakref
-from contextlib import contextmanager
 from types import SimpleNamespace
+from unittest.mock import Mock
 
 import pytest
 import torch
@@ -17,7 +17,15 @@ from torch import nn
 from torch.distributed.tensor import DeviceMesh, DTensor, Replicate
 
 import vllm_omni.diffusion.offloader.distributed_layerwise_backend as dist_backend_module
-from tests.helpers.runtime import get_distributed_init_method
+from tests.diffusion.offloader.helpers import (
+    DummyStream,
+    _DummyBlock,
+    _PlainEncoder,
+    _SingleBlockModel,
+    _StagedEncoder,
+    _StagedVAE,
+    patch_offload_runtime,
+)
 from vllm_omni.diffusion.data import validate_dlo_host_registration_options
 from vllm_omni.diffusion.model_loader.host_weight_plan import (
     HostWeightPlan,
@@ -51,53 +59,9 @@ from vllm_omni.platforms import current_omni_platform
 pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
 
 
-class DummyStream:
-    def wait_stream(self, _stream) -> None:
-        return None
-
-    def wait_event(self, _event) -> None:
-        return None
-
-
-class DummyEvent:
-    def record(self, _stream) -> None:
-        return None
-
-    def synchronize(self) -> None:
-        return None
-
-
-@contextmanager
-def dummy_stream(_stream):
-    yield None
-
-
-def _cleanup_distributed() -> None:
-    if dist.is_initialized():
-        dist.destroy_process_group()
-
-    gc.collect()
-    if current_omni_platform.is_available():
-        current_omni_platform.empty_cache()
-        current_omni_platform.synchronize()
-
-
-@pytest.fixture(scope="module")
-def dist_group():
-    dist.init_process_group("gloo", rank=0, world_size=1, init_method=get_distributed_init_method())
-    try:
-        yield
-    finally:
-        _cleanup_distributed()
-
-
 @pytest.fixture
 def patched_offload_runtime(monkeypatch):
-    monkeypatch.setattr(dist_backend_module.current_omni_platform, "Stream", DummyStream)
-    monkeypatch.setattr(dist_backend_module.current_omni_platform, "Event", DummyEvent)
-    monkeypatch.setattr(dist_backend_module.current_omni_platform, "current_stream", lambda: DummyStream())
-    monkeypatch.setattr(dist_backend_module.current_omni_platform, "stream", dummy_stream)
-    monkeypatch.setattr(dist_backend_module.current_omni_platform, "synchronize", lambda: None)
+    patch_offload_runtime(monkeypatch, dist_backend_module.current_omni_platform, synchronize=True)
 
 
 class TinyBlock(nn.Module):
@@ -113,6 +77,93 @@ def _make_values(start: float) -> torch.Tensor:
 
 
 class TestDistributedLayerwiseOffloadHook:
+    def test_buffer_only_block_reports_offloaded_state(self, patched_offload_runtime):
+        current_block = nn.Module()
+        current_block.register_buffer("state", torch.ones(2))
+        next_block = nn.Module()
+        next_block.register_buffer("state", torch.ones(2))
+        hook = DistributedLayerwiseOffloadHook(
+            next_block=next_block,
+            device=torch.device("cpu"),
+            dp_group=None,
+            dp_size=1,
+            rank=0,
+            pin_memory=False,
+        )
+        hook.initialize_hook(current_block)
+
+        assert hook.is_materialized
+        hook.offload_layer()
+        assert not hook.is_materialized
+
+    def test_initialize_failure_keeps_next_block_materialized(self, monkeypatch):
+        current_block = nn.Linear(2, 2)
+        next_block = nn.Linear(2, 2)
+        expected = {name: tensor.detach().clone() for name, tensor in next_block.state_dict().items()}
+        hook = DistributedLayerwiseOffloadHook(
+            next_block=next_block,
+            device=torch.device("cpu"),
+            dp_group=None,
+            dp_size=1,
+            rank=0,
+            pin_memory=False,
+        )
+
+        def fail_buffer_allocation():
+            raise RuntimeError("injected buffer allocation failure")
+
+        monkeypatch.setattr(hook, "_allocate_device_buffers", fail_buffer_allocation)
+
+        with pytest.raises(RuntimeError, match="buffer allocation failure"):
+            hook.initialize_hook(current_block)
+
+        for name, tensor in next_block.state_dict().items():
+            torch.testing.assert_close(tensor, expected[name])
+
+    @pytest.mark.parametrize(
+        ("kwargs", "message"),
+        [
+            ({"dp_size": 0, "rank": 0}, "dp_size must be a positive integer"),
+            ({"dp_size": 2, "rank": 2, "dp_group": object()}, "rank must satisfy"),
+            ({"dp_size": 2, "rank": 0}, "dp_group is required"),
+            (
+                {"dp_size": 2, "rank": 0, "dp_group": object(), "rank_local_mmap": True},
+                "rank_local_mmap requires dp_size=1",
+            ),
+        ],
+    )
+    def test_constructor_rejects_invalid_transport_topology(self, kwargs, message):
+        options = {"dp_group": None, "dp_size": 1, "rank": 0, **kwargs}
+
+        with pytest.raises(ValueError, match=message):
+            DistributedLayerwiseOffloadHook(
+                next_block=nn.Linear(2, 2),
+                device=torch.device("cpu"),
+                pin_memory=False,
+                **options,
+            )
+
+    def test_ring_probes_are_captured_before_any_block_is_cleared(self, patched_offload_runtime):
+        blocks = nn.ModuleList([nn.Linear(2, 2, bias=False) for _ in range(3)])
+        expected_middle = blocks[1].weight.detach().clone()
+        backend = DistributedLayerwiseOffloadBackend(
+            OffloadConfig(
+                strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
+                pin_cpu_memory=False,
+                dlo_use_allgather=False,
+            ),
+            torch.device("cpu"),
+        )
+
+        hooks = backend._install_hook_group(blocks, "dit")
+        shared_buffers = backend._allocate_shared_buffers(hooks)
+        for hook in hooks:
+            hook.gpu_buffers = shared_buffers
+
+        assert all(not hook.is_materialized for hook in hooks)
+        hooks[2].pre_forward(blocks[1])
+        torch.testing.assert_close(blocks[1].weight, expected_middle)
+
     def test_shard_and_pin_single_rank(self, dist_group, patched_offload_runtime):
         """With dp_size=1, the shard should equal the full weights."""
         current_block = TinyBlock(_make_values(1.0))
@@ -511,7 +562,7 @@ class TestDistributedLayerwiseOffloadHook:
         hook = DistributedLayerwiseOffloadHook(
             next_block=next_block,
             device=torch.device("cpu"),
-            dp_group=None,
+            dp_group=object(),
             dp_size=4,
             rank=1,
             pin_memory=False,
@@ -537,7 +588,7 @@ class TestDistributedLayerwiseOffloadHook:
         hook0 = DistributedLayerwiseOffloadHook(
             next_block=next_block,
             device=torch.device("cpu"),
-            dp_group=None,
+            dp_group=object(),
             dp_size=2,
             rank=0,
             pin_memory=False,
@@ -559,7 +610,7 @@ class TestDistributedLayerwiseOffloadHook:
         hook1 = DistributedLayerwiseOffloadHook(
             next_block=next_block1,
             device=torch.device("cpu"),
-            dp_group=None,
+            dp_group=object(),
             dp_size=2,
             rank=1,
             pin_memory=False,
@@ -583,7 +634,7 @@ class TestDistributedLayerwiseOffloadHook:
         hook = DistributedLayerwiseOffloadHook(
             next_block=next_block,
             device=torch.device("cpu"),
-            dp_group=None,
+            dp_group=object(),
             dp_size=2,
             rank=0,
             pin_memory=False,
@@ -602,7 +653,7 @@ class TestDistributedLayerwiseOffloadHook:
         hook = DistributedLayerwiseOffloadHook(
             next_block=next_block,
             device=torch.device("cpu"),
-            dp_group=None,
+            dp_group=object(),
             dp_size=2,
             rank=0,
             pin_memory=False,
@@ -689,6 +740,32 @@ class TestDistributedLayerwiseOffloadHook:
 
 
 class TestPinnedResidentLayerGroup:
+    def test_constructor_failure_leaves_every_block_materialized(self, monkeypatch):
+        blocks = [nn.Linear(2, 2), nn.Linear(2, 2)]
+        expected = [block.weight.detach().clone() for block in blocks]
+        shard_and_pin = DistributedLayerwiseOffloadHook._shard_and_pin
+        calls = 0
+
+        def fail_second_block(*args, **kwargs):
+            nonlocal calls
+            calls += 1
+            if calls == 2:
+                raise RuntimeError("injected resident setup failure")
+            return shard_and_pin(*args, **kwargs)
+
+        monkeypatch.setattr(DistributedLayerwiseOffloadHook, "_shard_and_pin", fail_second_block)
+
+        with pytest.raises(RuntimeError, match="resident setup failure"):
+            PinnedResidentLayerGroup(
+                blocks,
+                device=torch.device("cpu"),
+                copy_stream=DummyStream(),
+                pin_memory=False,
+            )
+
+        for block, weight in zip(blocks, expected, strict=True):
+            torch.testing.assert_close(block.weight, weight)
+
     def test_load_offload_reuses_pinned_master_weights(self, patched_offload_runtime):
         blocks = [nn.Linear(2, 2), nn.Linear(2, 2)]
         expected = []
@@ -843,20 +920,6 @@ class TestPinnedModuleStager:
         stager.offload()
 
 
-class _DummyBlock(nn.Module):
-    def __init__(self):
-        super().__init__()
-        self.weight = nn.Parameter(torch.randn(10, 10))
-
-
-class _SingleBlockModel(nn.Module):
-    _layerwise_offload_blocks_attrs = ["blocks"]
-
-    def __init__(self, num_blocks: int = 3):
-        super().__init__()
-        self.blocks = nn.ModuleList([_DummyBlock() for _ in range(num_blocks)])
-
-
 class _HWRPipeline(nn.Module):
     def __init__(self):
         super().__init__()
@@ -1285,6 +1348,31 @@ class TestMmapWeightLoading:
         backend.disable()
         assert lease.closed
 
+        backend.enable(pipeline)
+        assert backend.host_weight_plan is None
+        assert backend.enabled
+        backend.disable()
+
+    def test_rank_local_disable_drains_pending_transfers(self, patched_offload_runtime, monkeypatch):
+        pipeline = _HWRPipeline()
+        backend = DistributedLayerwiseOffloadBackend(
+            OffloadConfig(
+                strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
+                pin_cpu_memory=False,
+                dp_size=1,
+                dlo_use_allgather=False,
+            ),
+            torch.device("cpu"),
+        )
+        backend.enable(pipeline)
+        backend._all_hook_groups[0][0].prefetch_layer(0)
+        synchronize = Mock()
+        monkeypatch.setattr(current_omni_platform, "synchronize", synchronize)
+
+        backend.disable()
+
+        synchronize.assert_called()
+
     def test_hwr_registration_failure_falls_back_to_bounded_staging(
         self,
         monkeypatch,
@@ -1563,6 +1651,7 @@ class TestOffloadPlan:
         assert plan.block_attrs == {}
         assert plan.offload_submodules == {}
         assert plan.resident_dit_paths == frozenset()
+        assert plan.encoder_dlo_weight_replication == frozenset()
 
     def test_offload_plan_is_frozen(self):
         """OffloadPlan should be immutable (frozen=True)."""
@@ -1595,6 +1684,7 @@ class TestOffloadPlan:
                 self.transformer = Transformer()
 
         pipeline = Pipeline()
+        expected_parameters = {name: parameter.detach().clone() for name, parameter in pipeline.named_parameters()}
         backend = DistributedLayerwiseOffloadBackend(
             OffloadConfig(
                 strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
@@ -1610,10 +1700,22 @@ class TestOffloadPlan:
         assert len(backend._resident_blocks) == 2
         assert len(backend._all_hook_groups) == 1
         assert len(backend._all_hook_groups[0]) == 2
+        assert pipeline._dlo_residency_controller is backend
         assert backend.enabled
 
         backend.disable()
 
+        assert pipeline._dlo_residency_controller is None
+
+        for name, parameter in pipeline.named_parameters():
+            torch.testing.assert_close(parameter, expected_parameters[name])
+
+        backend.enable(pipeline)
+        assert len(backend._resident_blocks) == 2
+        backend.disable()
+        for name, parameter in pipeline.named_parameters():
+            torch.testing.assert_close(parameter, expected_parameters[name])
+
 
 class TestMmapValidation:
     """Tests for loader preflight and backend plan realization."""
@@ -2024,8 +2126,8 @@ class TestConfigValidation:
             (2, 1, 1, True, 2),
             (1, 4, 1, True, 4),
             (2, 2, 2, True, 2),
-            (2, 2, 2, False, 1),
-            (1, 4, 2, False, 1),
+            (2, 2, 2, False, 2),
+            (1, 4, 2, False, 4),
         ],
     )
     def test_dlo_group_selection_covers_dp_sp_tp_matrix(
@@ -2199,7 +2301,7 @@ class TestConfigValidation:
             model = "/fake/path"
 
         config = OffloadConfig.from_od_config(FakeODConfig())
-        assert config.dp_size == 1  # forced to 1 when no AllGather
+        assert config.dp_size == 2  # topology is retained; the hook selects group size 1
         assert config.dlo_resident_layers == 20
 
     def test_resident_layers_with_allgather_rejected(self):
@@ -2218,96 +2320,9 @@ class TestConfigValidation:
             parallel_config = FakePC()
             model = "/fake/path"
 
-        with pytest.raises(ValueError, match="requires --dlo-no-use-allgather"):
+        with pytest.raises(ValueError, match="requires the DiT DLO transfer to be rank-local"):
             OffloadConfig.from_od_config(FakeODConfig())
 
-    def test_num_inference_steps_none_rejected(self):
-        """DP multi-concurrency should reject None num_inference_steps."""
-        from types import SimpleNamespace
-
-        # Mock requests with None steps
-        reqs = [
-            SimpleNamespace(
-                req=SimpleNamespace(
-                    request_id=f"req-{i}",
-                    sampling_params=SimpleNamespace(num_inference_steps=None),
-                )
-            )
-            for i in range(2)
-        ]
-
-        # We can't easily instantiate the full executor, but we can test
-        # the validation logic by checking that the code path raises.
-        # The validation is in execute_request, which needs self._ensure_open().
-        # Instead, test the validation logic directly:
-        step_counts = {
-            r.req.sampling_params.num_inference_steps
-            for r in reqs
-            if r.req.sampling_params.num_inference_steps is not None
-        }
-        has_none = any(r.req.sampling_params.num_inference_steps is None for r in reqs)
-        assert has_none, "Test setup: should have None steps"
-        assert len(step_counts) == 0, "Test setup: no explicit steps"
-
-        # The validation condition: (len(step_counts) > 1) or has_none → should reject
-        should_reject = (len(step_counts) > 1) or has_none
-        assert should_reject, "None steps should trigger rejection"
-
-    def test_num_inference_steps_same_explicit_allowed(self):
-        """DP multi-concurrency should allow same explicit num_inference_steps."""
-        from types import SimpleNamespace
-
-        reqs = [
-            SimpleNamespace(
-                req=SimpleNamespace(
-                    request_id=f"req-{i}",
-                    sampling_params=SimpleNamespace(num_inference_steps=35),
-                )
-            )
-            for i in range(4)
-        ]
-
-        step_counts = {
-            r.req.sampling_params.num_inference_steps
-            for r in reqs
-            if r.req.sampling_params.num_inference_steps is not None
-        }
-        has_none = any(r.req.sampling_params.num_inference_steps is None for r in reqs)
-
-        should_reject = (len(step_counts) > 1) or has_none
-        assert not should_reject, "Same explicit steps should be allowed"
-        assert step_counts == {35}
-
-    def test_num_inference_steps_different_explicit_rejected(self):
-        """DP multi-concurrency should reject different explicit steps."""
-        from types import SimpleNamespace
-
-        reqs = [
-            SimpleNamespace(
-                req=SimpleNamespace(
-                    request_id="req-0",
-                    sampling_params=SimpleNamespace(num_inference_steps=35),
-                )
-            ),
-            SimpleNamespace(
-                req=SimpleNamespace(
-                    request_id="req-1",
-                    sampling_params=SimpleNamespace(num_inference_steps=30),
-                )
-            ),
-        ]
-
-        step_counts = {
-            r.req.sampling_params.num_inference_steps
-            for r in reqs
-            if r.req.sampling_params.num_inference_steps is not None
-        }
-        has_none = any(r.req.sampling_params.num_inference_steps is None for r in reqs)
-
-        should_reject = (len(step_counts) > 1) or has_none
-        assert should_reject, "Different steps should trigger rejection"
-        assert len(step_counts) == 2
-
 
 class TestDynamicSlotTracking:
     """Tests for dynamic slot tracking with odd and even block counts.
@@ -2472,3 +2487,344 @@ class TestDynamicSlotTracking:
         # pre_forward should override to slot 1
         hook_b.pre_forward(block_b)
         assert hook_b.current_slot == 1, "pre_forward should read _prev_hook._prefetched_slot=1, not keep initial 0"
+
+
+class _DistributedComponentPipeline(nn.Module):
+    _offload_plan = OffloadPlan(
+        encoder_component_types={"text_encoder": "text_encoder"},
+        encoder_block_attrs={"text_encoder": ("vision.blocks", "text_model.layers")},
+        on_demand_component_paths=frozenset({"text_encoder", "vae"}),
+    )
+
+    def __init__(self):
+        super().__init__()
+        self.transformer = _SingleBlockModel(num_blocks=2)
+        self.text_encoder = _StagedEncoder()
+        self.vae = _StagedVAE()
+
+
+class _GenericDistributedEncoderPipeline(nn.Module):
+    _offload_plan = OffloadPlan(
+        encoder_component_types={"text_encoder": "text_encoder"},
+        encoder_block_attrs={"text_encoder": ("encoder.block",)},
+        encoder_dlo_weight_replication=frozenset({"text_encoder"}),
+    )
+
+    def __init__(self):
+        super().__init__()
+        self.transformer = _SingleBlockModel(num_blocks=2)
+        self.text_encoder = _PlainEncoder()
+
+
+class TestDistributedComponentSelection:
+    def test_on_demand_only_encoder_cannot_claim_allgather(self, patched_offload_runtime):
+        class Pipeline(nn.Module):
+            _offload_plan = OffloadPlan(
+                encoder_component_types={"text_encoder": "text_encoder"},
+                on_demand_component_paths=frozenset({"text_encoder"}),
+                encoder_dlo_weight_replication=frozenset({"text_encoder"}),
+            )
+
+            def __init__(self):
+                super().__init__()
+                self.text_encoder = _StagedEncoder()
+
+        backend = DistributedLayerwiseOffloadBackend(
+            OffloadConfig(
+                strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
+                pin_cpu_memory=False,
+                dp_size=2,
+                components=frozenset({"text_encoder"}),
+                dlo_transfers={"dit": "rank-local", "text_encoder": "allgather"},
+            ),
+            torch.device("cpu"),
+        )
+        backend.dp_group = object()
+
+        with pytest.raises(ValueError, match="cannot use AllGather without.*streamable block plan"):
+            backend.enable(Pipeline())
+
+    def test_explicit_single_block_dit_is_rejected(self, patched_offload_runtime):
+        pipeline = nn.Module()
+        pipeline.transformer = _SingleBlockModel(num_blocks=1)
+        backend = DistributedLayerwiseOffloadBackend(
+            OffloadConfig(
+                strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
+                pin_cpu_memory=False,
+                components=frozenset({"dit"}),
+            ),
+            torch.device("cpu"),
+        )
+
+        with pytest.raises(ValueError, match="leaves only one streaming block"):
+            backend.enable(pipeline)
+
+    def test_enable_defers_collectives_until_first_forward(self, patched_offload_runtime, monkeypatch):
+        pipeline = _GenericDistributedEncoderPipeline()
+        backend = DistributedLayerwiseOffloadBackend(
+            OffloadConfig(
+                strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
+                pin_cpu_memory=False,
+                dp_size=2,
+                components=frozenset({"dit", "text_encoder"}),
+                dlo_transfers={"dit": "allgather", "text_encoder": "rank-local"},
+            ),
+            torch.device("cpu"),
+        )
+        backend.dp_group = object()
+
+        def reject_startup_collective(*_args, **_kwargs):
+            raise AssertionError("enable() entered a rank-asymmetric collective")
+
+        monkeypatch.setattr(torch.distributed, "all_gather_into_tensor", reject_startup_collective)
+
+        backend.enable(pipeline)
+
+        assert all(hook._prefetched_slot is None for group in backend._all_hook_groups for hook in group)
+        encoder_hook = pipeline.text_encoder.encoder.block[0]._hook_registry.get_hook("distributed_layerwise_offload")
+        dit_hook = pipeline.transformer.blocks[0]._hook_registry.get_hook("distributed_layerwise_offload")
+        assert encoder_hook.gpu_shard_buffers == [None, None]
+        assert all(buffer is not None for buffer in dit_hook.gpu_shard_buffers)
+
+        monkeypatch.setattr(
+            torch.distributed,
+            "all_gather_into_tensor",
+            lambda output, local, group: output.copy_(local.repeat(2)),
+        )
+        backend.disable()
+
+    def test_partial_dit_enable_failure_restores_weights_and_hooks(
+        self,
+        patched_offload_runtime,
+        monkeypatch,
+    ):
+        pipeline = _DistributedComponentPipeline()
+        expected_weights = [block.weight.detach().clone() for block in pipeline.transformer.blocks]
+        backend = DistributedLayerwiseOffloadBackend(
+            OffloadConfig(
+                strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
+                pin_cpu_memory=False,
+                dlo_use_allgather=False,
+                components=frozenset({"dit"}),
+            ),
+            torch.device("cpu"),
+        )
+        original_apply = dist_backend_module.apply_distributed_block_hook
+        calls = 0
+
+        def fail_second_hook(*args, **kwargs):
+            nonlocal calls
+            calls += 1
+            if calls == 2:
+                raise RuntimeError("injected DLO hook failure")
+            return original_apply(*args, **kwargs)
+
+        monkeypatch.setattr(dist_backend_module, "apply_distributed_block_hook", fail_second_hook)
+
+        with pytest.raises(RuntimeError, match="injected DLO hook failure"):
+            backend.enable(pipeline)
+
+        assert not backend.enabled
+        for block, expected in zip(pipeline.transformer.blocks, expected_weights, strict=True):
+            registry = getattr(block, "_hook_registry", None)
+            assert registry is None or registry.get_hook("distributed_layerwise_offload") is None
+            torch.testing.assert_close(block.weight, expected)
+
+    def test_multirank_enable_failure_cleanup_skips_restore_collective(self, monkeypatch, mocker):
+        backend = DistributedLayerwiseOffloadBackend(
+            OffloadConfig(
+                strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
+                pin_cpu_memory=False,
+                dp_size=2,
+                dlo_use_allgather=True,
+                components=frozenset({"dit"}),
+            ),
+            torch.device("cpu"),
+        )
+        monkeypatch.setattr(backend, "_enable", mocker.Mock(side_effect=RuntimeError("injected startup failure")))
+        cleanup = mocker.Mock()
+        monkeypatch.setattr(backend, "_disable", cleanup)
+
+        with pytest.raises(RuntimeError, match="injected startup failure"):
+            backend.enable(_DistributedComponentPipeline())
+
+        cleanup.assert_called_once_with(restore_allgather_weights=False)
+
+    def test_multirank_failure_restores_only_rank_local_hooks(self, patched_offload_runtime):
+        backend = DistributedLayerwiseOffloadBackend(
+            OffloadConfig(
+                strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
+                pin_cpu_memory=False,
+                dp_size=2,
+                components=frozenset({"dit", "text_encoder"}),
+                dlo_transfers={"dit": "allgather", "text_encoder": "rank-local"},
+            ),
+            torch.device("cpu"),
+        )
+        allgather_hook = Mock(dp_size=2, next_block=nn.Linear(1, 1))
+        rank_local_hook = Mock(dp_size=1, next_block=nn.Linear(1, 1))
+        backend._all_hook_groups = [[allgather_hook], [rank_local_hook]]
+
+        backend._disable(restore_allgather_weights=False)
+
+        allgather_hook.restore_next_block_to_cpu.assert_not_called()
+        rank_local_hook.restore_next_block_to_cpu.assert_called_once_with()
+
+        with pytest.raises(RuntimeError, match="recreate the backend and reload the pipeline"):
+            backend.enable(nn.Module())
+
+    def test_disable_runs_collectives_before_best_effort_local_cleanup(
+        self,
+        patched_offload_runtime,
+        monkeypatch,
+    ):
+        events: list[str] = []
+        backend = DistributedLayerwiseOffloadBackend(
+            OffloadConfig(
+                strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
+                pin_cpu_memory=False,
+                dp_size=2,
+            ),
+            torch.device("cpu"),
+        )
+        rank_local_hook = Mock(dp_size=1, next_block=nn.Linear(1, 1))
+        allgather_hook = Mock(dp_size=2, next_block=nn.Linear(1, 1))
+        allgather_hook.restore_next_block_to_cpu.side_effect = lambda: events.append("allgather")
+
+        def fail_rank_local_restore():
+            events.append("rank-local")
+            raise RuntimeError("injected rank-local restore failure")
+
+        rank_local_hook.restore_next_block_to_cpu.side_effect = fail_rank_local_restore
+        backend._all_hook_groups = [[rank_local_hook], [allgather_hook]]
+        backend._blocks = [[nn.Module()]]
+        backend.enabled = True
+        monkeypatch.setattr(
+            dist_backend_module,
+            "remove_distributed_block_hook",
+            lambda _block: events.append("remove"),
+        )
+        monkeypatch.setattr(backend, "_release_registered_mmap", lambda: events.append("unregister"))
+        release_handles = Mock()
+        monkeypatch.setattr(backend, "_release_mmap_handles", release_handles)
+
+        with pytest.raises(RuntimeError, match="injected rank-local restore failure"):
+            backend.disable()
+
+        assert events == ["allgather", "rank-local", "remove", "unregister"]
+        release_handles.assert_not_called()
+        assert backend.enabled
+        assert backend._all_hook_groups
+
+    def test_all_streams_encoder_and_keeps_vae_resident(self, patched_offload_runtime, monkeypatch):
+        pipeline = _DistributedComponentPipeline()
+        move_non_block_state = Mock()
+        monkeypatch.setattr(dist_backend_module, "move_non_block_state_to_device", move_non_block_state)
+        backend = DistributedLayerwiseOffloadBackend(
+            OffloadConfig(
+                strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
+                pin_cpu_memory=False,
+                dlo_use_allgather=False,
+                components=frozenset({"dit", "text_encoder"}),
+            ),
+            torch.device("cpu"),
+        )
+
+        backend.enable(pipeline)
+
+        assert pipeline.text_encoder._omni_layerwise_enabled
+        assert len(pipeline.text_encoder._omni_layerwise_hooks) == 4
+        # The pipeline-owned stage lifecycle releases the encoder's non-block
+        # state after DLO has installed blockwise hooks.
+        assert pipeline.text_encoder.offload_calls == 1
+        assert pipeline.vae.offload_calls == 0
+        assert pipeline.vae.to_calls == 1
+        move_non_block_state.assert_not_called()
+
+        backend.disable()
+
+        assert not pipeline.text_encoder._omni_layerwise_enabled
+        encoder_blocks = [
+            *pipeline.text_encoder.vision.blocks,
+            *pipeline.text_encoder.text_model.layers,
+        ]
+        assert all(block._hook_registry.get_hook("distributed_layerwise_offload") is None for block in encoder_blocks)
+
+    def test_dit_and_encoder_choose_transfer_independently(self, patched_offload_runtime, monkeypatch):
+        pipeline = _GenericDistributedEncoderPipeline()
+        expected_parameters = {name: parameter.detach().clone() for name, parameter in pipeline.named_parameters()}
+        full_weights_by_rank_zero_shard = {
+            tuple(parameter.flatten()[: (parameter.numel() + 1) // 2].tolist()): parameter.flatten()
+            for parameter in expected_parameters.values()
+            if parameter.numel() > 4
+        }
+        backend = DistributedLayerwiseOffloadBackend(
+            OffloadConfig(
+                strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
+                pin_cpu_memory=False,
+                dp_size=2,
+                components=frozenset({"dit", "text_encoder"}),
+                dlo_transfers={"dit": "rank-local", "text_encoder": "allgather"},
+            ),
+            torch.device("cpu"),
+        )
+        backend.dp_group = object()
+
+        def fake_allgather(output, local, group):
+            del group
+            output.copy_(full_weights_by_rank_zero_shard[tuple(local.tolist())])
+
+        monkeypatch.setattr(
+            torch.distributed,
+            "all_gather_into_tensor",
+            fake_allgather,
+        )
+
+        backend.enable(pipeline)
+
+        encoder_hook = pipeline.text_encoder.encoder.block[0]._hook_registry.get_hook("distributed_layerwise_offload")
+        dit_hook = pipeline.transformer.blocks[0]._hook_registry.get_hook("distributed_layerwise_offload")
+        assert encoder_hook.dp_size == 2
+        assert dit_hook.dp_size == 1
+        assert all(buffer is not None for buffer in encoder_hook.gpu_shard_buffers)
+        assert dit_hook.gpu_shard_buffers == [None, None]
+
+        backend.disable()
+        actual_parameters = dict(pipeline.named_parameters())
+        for name, expected in expected_parameters.items():
+            torch.testing.assert_close(actual_parameters[name], expected)
+
+    def test_encoder_allgather_rejects_undeclared_replication(self, patched_offload_runtime):
+        pipeline = _DistributedComponentPipeline()
+        backend = DistributedLayerwiseOffloadBackend(
+            OffloadConfig(
+                strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
+                pin_cpu_memory=False,
+                dp_size=2,
+                components=frozenset({"text_encoder"}),
+                dlo_transfers={"dit": "rank-local", "text_encoder": "allgather"},
+            ),
+            torch.device("cpu"),
+        )
+        backend.dp_group = object()
+
+        with pytest.raises(ValueError, match="not declared replicated"):
+            backend.enable(pipeline)
+
+    def test_encoder_allgather_rejects_stub_rank_before_block_discovery(self):
+        """Every rank must reject an unsafe encoder group, including stub ranks."""
+        backend = DistributedLayerwiseOffloadBackend(
+            OffloadConfig(
+                strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
+                pin_cpu_memory=False,
+                dp_size=2,
+                components=frozenset({"text_encoder"}),
+                dlo_transfers={"dit": "rank-local", "text_encoder": "allgather"},
+            ),
+            torch.device("cpu"),
+        )
+        backend.dp_group = object()
+        stub_plan = OffloadPlan(encoder_block_attrs={"text_encoder": ("missing.blocks",)})
+
+        with pytest.raises(ValueError, match="not declared replicated"):
+            backend._try_layerwise_offload_encoder(nn.Module(), "text_encoder", stub_plan)

tests/diffusion/offloader/test_layerwise_backend.py

@@ -3,66 +3,39 @@
 
 """Unit tests for LayerwiseOffloadHook and LayerWiseOffloadBackend utilities."""
 
-import gc
-from contextlib import contextmanager
+from multiprocessing.reduction import ForkingPickler
+from types import SimpleNamespace
+from unittest.mock import Mock
 
 import pytest
 import torch
-import torch.distributed as dist
 from torch import nn
 from torch.distributed.tensor import DeviceMesh, DTensor, Replicate
 
 import vllm_omni.diffusion.offloader.layerwise_backend as layerwise_backend_module
-from tests.helpers.runtime import get_distributed_init_method
-from vllm_omni.diffusion.offloader.layerwise_backend import LayerWiseOffloadBackend, LayerwiseOffloadHook
-from vllm_omni.platforms import current_omni_platform
+import vllm_omni.diffusion.offloader.tensor_utils as tensor_utils_module
+from tests.diffusion.offloader.helpers import (
+    DummyStream,
+    _DummyBlock,
+    _PlainEncoder,
+    _SingleBlockModel,
+    _StagedEncoder,
+    _StagedVAE,
+    patch_offload_runtime,
+)
+from vllm_omni.diffusion.offloader.base import OffloadConfig, OffloadStrategy
+from vllm_omni.diffusion.offloader.layerwise_backend import (
+    LayerWiseOffloadBackend,
+    LayerwiseOffloadHook,
+)
+from vllm_omni.diffusion.offloader.offload_plan import OffloadPlan
 
 pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
 
 
-class DummyStream:
-    def wait_stream(self, _stream) -> None:
-        return None
-
-    def wait_event(self, _event) -> None:
-        return None
-
-
-class DummyEvent:
-    def record(self, _stream) -> None:
-        return None
-
-
-@contextmanager
-def dummy_stream(_stream):
-    yield None
-
-
-def _cleanup_distributed() -> None:
-    if dist.is_initialized():
-        dist.destroy_process_group()
-
-    gc.collect()
-    if current_omni_platform.is_available():
-        current_omni_platform.empty_cache()
-        current_omni_platform.synchronize()
-
-
-@pytest.fixture(scope="module")
-def dist_group():
-    dist.init_process_group("gloo", rank=0, world_size=1, init_method=get_distributed_init_method())
-    try:
-        yield
-    finally:
-        _cleanup_distributed()
-
-
 @pytest.fixture
 def patched_offload_runtime(monkeypatch):
-    monkeypatch.setattr(layerwise_backend_module.current_omni_platform, "Stream", DummyStream)
-    monkeypatch.setattr(layerwise_backend_module.current_omni_platform, "Event", DummyEvent)
-    monkeypatch.setattr(layerwise_backend_module.current_omni_platform, "current_stream", lambda: DummyStream())
-    monkeypatch.setattr(layerwise_backend_module.current_omni_platform, "stream", dummy_stream)
+    patch_offload_runtime(monkeypatch, layerwise_backend_module.current_omni_platform)
 
 
 class TinyBlock(nn.Module):
@@ -77,7 +50,66 @@ def _make_values(start: float) -> torch.Tensor:
     return torch.arange(start, start + 4, dtype=torch.float32)
 
 
+def test_clear_tensor_storage_rolls_back_partial_commit(monkeypatch):
+    first = nn.Parameter(torch.tensor([1.0, 2.0]))
+    second = nn.Parameter(torch.tensor([3.0, 4.0]))
+    expected_first = first.detach().clone()
+    expected_second = second.detach().clone()
+    second_data_ptr = second.data_ptr()
+    original_set_storage = tensor_utils_module.set_tensor_storage
+    calls = 0
+
+    def fail_second_commit(target, value):
+        nonlocal calls
+        calls += 1
+        if calls == 2:
+            raise RuntimeError("injected storage clear failure")
+        original_set_storage(target, value)
+
+    monkeypatch.setattr(tensor_utils_module, "set_tensor_storage", fail_second_commit)
+
+    with pytest.raises(RuntimeError, match="injected storage clear failure"):
+        tensor_utils_module.clear_tensor_storage([first, second])
+
+    assert calls == 3  # first commit, failed second commit, first rollback
+    torch.testing.assert_close(first, expected_first)
+    torch.testing.assert_close(second, expected_second)
+    assert second.data_ptr() == second_data_ptr
+
+
 class TestLayerwiseOffloadHook:
+    def test_buffer_only_block_reports_offloaded_state(self, patched_offload_runtime):
+        current_block = nn.Module()
+        current_block.register_buffer("state", torch.ones(2))
+        next_block = nn.Module()
+        next_block.register_buffer("state", torch.ones(2))
+        hook = LayerwiseOffloadHook(
+            next_block=next_block,
+            device=torch.device("cpu"),
+            stream=DummyStream(),
+            pin_memory=False,
+        )
+        hook.initialize_hook(current_block)
+
+        assert hook.is_materialized
+        hook.offload_layer()
+        assert not hook.is_materialized
+
+    def test_ring_probes_are_captured_before_any_block_is_cleared(self, patched_offload_runtime):
+        blocks = nn.ModuleList([nn.Linear(2, 2, bias=False) for _ in range(3)])
+        expected_middle = blocks[1].weight.detach().clone()
+
+        hooks = layerwise_backend_module._install_layerwise_hook_group(
+            blocks,
+            torch.device("cpu"),
+            DummyStream(),
+            pin_memory=False,
+        )
+
+        assert all(not hook.is_materialized for hook in hooks)
+        hooks[2].pre_forward(blocks[1])
+        torch.testing.assert_close(blocks[1].weight, expected_middle)
+
     def test_dtensor_wrapper_is_preserved_across_prefetch_and_offload(self, dist_group, patched_offload_runtime):
         current_block = TinyBlock(_make_values(1.0))
         next_block = TinyBlock(_make_values(10.0))
@@ -135,20 +167,6 @@ class TestLayerwiseOffloadHook:
         assert torch.equal(next_block.weight, expected)
 
 
-class _DummyBlock(nn.Module):
-    def __init__(self):
-        super().__init__()
-        self.weight = nn.Parameter(torch.randn(10, 10))
-
-
-class _SingleBlockModel(nn.Module):
-    _layerwise_offload_blocks_attrs = ["blocks"]
-
-    def __init__(self, num_blocks: int = 3):
-        super().__init__()
-        self.blocks = nn.ModuleList([_DummyBlock() for _ in range(num_blocks)])
-
-
 class _MultiBlockModel(nn.Module):
     _layerwise_offload_blocks_attrs = ["transformer_blocks", "single_transformer_blocks"]
 
@@ -246,3 +264,732 @@ class TestGetBlocksAttrNames:
         LayerWiseOffloadBackend.set_blocks_attr_names(model, ["new_blocks"])
         assert hasattr(model.__class__, "_layerwise_offload_blocks_attrs")
         assert model.__class__._layerwise_offload_blocks_attrs == ["new_blocks"]
+
+
+class _ComponentPipeline(nn.Module):
+    _offload_plan = OffloadPlan(
+        encoder_component_types={"text_encoder": "text_encoder"},
+        encoder_block_attrs={"text_encoder": ("vision.blocks", "text_model.layers")},
+        on_demand_component_paths=frozenset({"text_encoder", "vae"}),
+    )
+
+    def __init__(self):
+        super().__init__()
+        self.transformer = _SingleBlockModel()
+        self.text_encoder = _StagedEncoder()
+        self.vae = _StagedVAE()
+
+
+class _LegacyComponentPipeline(nn.Module):
+    """Pipeline with DiT block metadata but no auxiliary OffloadPlan."""
+
+    def __init__(self):
+        super().__init__()
+        self.transformer = _SingleBlockModel()
+        self.text_encoder = _StagedEncoder()
+        self.vae = _StagedVAE()
+
+
+class _LegacyEncoderOnlyPipeline(nn.Module):
+    _offload_plan = _ComponentPipeline._offload_plan
+
+    def __init__(self):
+        super().__init__()
+        self.text_encoder = _StagedEncoder()
+        self.vae = _StagedVAE()
+
+
+class _PlannedDitPipeline(nn.Module):
+    _offload_plan = OffloadPlan(block_attrs={"transformer": ("blocks",)})
+
+    def __init__(self):
+        super().__init__()
+        self.transformer = _NoAttrsModel()
+
+
+class _GenericEncoderPipeline(nn.Module):
+    _offload_plan = OffloadPlan(
+        encoder_component_types={"text_encoder": "text_encoder"},
+        encoder_block_attrs={"text_encoder": ("encoder.block",)},
+    )
+
+    def __init__(self):
+        super().__init__()
+        self.transformer = _SingleBlockModel()
+        self.text_encoder = _PlainEncoder()
+
+
+def _layer_backend(components: set[str] | None = None) -> LayerWiseOffloadBackend:
+    options = {}
+    if components is not None:
+        options = {"components": frozenset(components)}
+    return LayerWiseOffloadBackend(
+        OffloadConfig(
+            strategy=OffloadStrategy.LAYER_WISE,
+            pin_cpu_memory=False,
+            **options,
+        ),
+        torch.device("cpu"),
+    )
+
+
+class TestLayerwiseComponentSelection:
+    def test_plan_block_attrs_drive_dit_discovery(self, patched_offload_runtime):
+        pipeline = _PlannedDitPipeline()
+        backend = _layer_backend({"dit"})
+
+        backend.enable(pipeline)
+
+        assert hasattr(pipeline.transformer.blocks[0], "_hook_registry")
+        backend.disable()
+
+    def test_legacy_missing_dit_preserves_noop(self, patched_offload_runtime):
+        pipeline = _LegacyEncoderOnlyPipeline()
+        backend = _layer_backend()
+
+        backend.enable(pipeline)
+
+        assert not backend.enabled
+        assert pipeline.text_encoder.offload_calls == 0
+        assert pipeline.vae.offload_calls == 0
+
+    def test_encoder_only_streams_planned_blocks(self, patched_offload_runtime, monkeypatch):
+        pipeline = _ComponentPipeline()
+        backend = _layer_backend({"text_encoder"})
+        move_non_block_state = Mock()
+        monkeypatch.setattr(layerwise_backend_module, "move_non_block_state_to_device", move_non_block_state)
+
+        backend.enable(pipeline)
+
+        assert pipeline.text_encoder._omni_layerwise_enabled
+        assert len(pipeline.text_encoder._omni_layerwise_hooks) == 4
+        assert pipeline.text_encoder.offload_calls == 1
+        assert pipeline.vae.offload_calls == 0
+        assert pipeline.vae.to_calls == 1
+        assert not hasattr(pipeline.transformer.blocks[0], "_hook_registry")
+        assert backend.enabled
+        move_non_block_state.assert_not_called()
+
+        backend.disable()
+        assert not pipeline.text_encoder._omni_layerwise_enabled
+
+    def test_single_gpu_dit_only_keeps_encoder_and_vae_resident(self, patched_offload_runtime):
+        pipeline = _ComponentPipeline()
+        expected_weights = [block.weight.detach().clone() for block in pipeline.transformer.blocks]
+        backend = _layer_backend({"dit"})
+
+        backend.enable(pipeline)
+
+        assert not hasattr(pipeline.text_encoder, "_omni_layerwise_enabled")
+        assert pipeline.text_encoder.to_calls == 1
+        assert pipeline.text_encoder.offload_calls == 0
+        assert pipeline.vae.to_calls == 1
+        assert pipeline.vae.offload_calls == 0
+        assert hasattr(pipeline.transformer.blocks[0], "_hook_registry")
+
+        backend.disable()
+
+        for block, expected in zip(pipeline.transformer.blocks, expected_weights, strict=True):
+            torch.testing.assert_close(block.weight, expected)
+
+        backend.enable(pipeline)
+        backend.disable()
+        for block, expected in zip(pipeline.transformer.blocks, expected_weights, strict=True):
+            torch.testing.assert_close(block.weight, expected)
+
+    def test_partial_dit_enable_failure_restores_weights_and_hooks(self, patched_offload_runtime, monkeypatch):
+        pipeline = _ComponentPipeline()
+        expected_weights = [block.weight.detach().clone() for block in pipeline.transformer.blocks]
+        backend = _layer_backend({"dit"})
+        original_apply = layerwise_backend_module.apply_block_hook
+        calls = 0
+
+        def fail_second_hook(*args, **kwargs):
+            nonlocal calls
+            calls += 1
+            if calls == 2:
+                raise RuntimeError("injected hook failure")
+            return original_apply(*args, **kwargs)
+
+        monkeypatch.setattr(layerwise_backend_module, "apply_block_hook", fail_second_hook)
+
+        with pytest.raises(RuntimeError, match="injected hook failure"):
+            backend.enable(pipeline)
+
+        assert not backend.enabled
+        for block, expected in zip(pipeline.transformer.blocks, expected_weights, strict=True):
+            registry = getattr(block, "_hook_registry", None)
+            assert registry is None or registry.get_hook("layerwise_offload") is None
+            torch.testing.assert_close(block.weight, expected)
+
+    def test_later_component_failure_removes_installed_encoder_hooks(
+        self,
+        patched_offload_runtime,
+        monkeypatch,
+    ):
+        pipeline = _ComponentPipeline()
+        blocks = [*pipeline.text_encoder.vision.blocks, *pipeline.text_encoder.text_model.layers]
+        expected = [block.weight.detach().clone() for block in blocks]
+        backend = _layer_backend({"text_encoder"})
+
+        def fail_vae_placement(*_args, **_kwargs):
+            raise RuntimeError("injected VAE placement failure")
+
+        monkeypatch.setattr(pipeline.vae, "to", fail_vae_placement)
+
+        with pytest.raises(RuntimeError, match="injected VAE placement failure"):
+            backend.enable(pipeline)
+
+        assert not backend.enabled
+        assert not pipeline.text_encoder._omni_layerwise_enabled
+        assert not backend._encoder_modules
+        for block, original in zip(blocks, expected, strict=True):
+            assert block._hook_registry.get_hook("layerwise_offload") is None
+            torch.testing.assert_close(block.weight, original)
+
+    def test_encoder_only_requires_streamable_offload_plan(self, patched_offload_runtime):
+        pipeline = nn.Module()
+        pipeline.transformer = _SingleBlockModel()
+        pipeline.text_encoder = _StagedEncoder()
+        backend = _layer_backend({"text_encoder"})
+
+        with pytest.raises(ValueError, match="Selected text encoder 'text_encoder' requires"):
+            backend.enable(pipeline)
+
+    def test_explicit_text_encoder_selection_requires_a_matching_module(self, patched_offload_runtime):
+        pipeline = nn.Module()
+        pipeline.transformer = _SingleBlockModel(num_blocks=2)
+        backend = _layer_backend({"dit", "text_encoder"})
+
+        with pytest.raises(ValueError, match="No text encoder modules found"):
+            backend.enable(pipeline)
+
+        assert not hasattr(pipeline.transformer.blocks[0], "_hook_registry")
+
+    def test_every_selected_encoder_requires_its_own_plan(self, patched_offload_runtime):
+        class Pipeline(nn.Module):
+            _offload_plan = OffloadPlan(
+                encoder_component_types={
+                    "text_encoder": "text_encoder",
+                    "text_encoder_2": "text_encoder",
+                },
+                encoder_block_attrs={"text_encoder": ("encoder.block",)},
+            )
+
+            def __init__(self):
+                super().__init__()
+                self.transformer = _SingleBlockModel()
+                self.text_encoder = _PlainEncoder()
+                self.text_encoder_2 = _PlainEncoder()
+
+        pipeline = Pipeline()
+        backend = _layer_backend({"text_encoder"})
+
+        with pytest.raises(ValueError, match="Selected text encoder 'text_encoder_2' requires"):
+            backend.enable(pipeline)
+
+        assert not pipeline.text_encoder._omni_layerwise_enabled
+
+    def test_default_selection_preserves_unplanned_auxiliaries(self, patched_offload_runtime):
+        pipeline = _LegacyComponentPipeline()
+        backend = _layer_backend()
+
+        backend.enable(pipeline)
+
+        assert hasattr(pipeline.transformer.blocks[0], "_hook_registry")
+        assert not hasattr(pipeline.text_encoder, "_omni_layerwise_enabled")
+        assert pipeline.text_encoder.to_calls == 1
+        assert pipeline.text_encoder.offload_calls == 0
+        assert pipeline.vae.to_calls == 1
+        assert pipeline.vae.offload_calls == 0
+
+        backend.disable()
+
+    def test_disable_failure_keeps_host_masters_for_retry(self, patched_offload_runtime, monkeypatch):
+        pipeline = nn.Module()
+        pipeline.transformer = _SingleBlockModel(num_blocks=3)
+        expected = {name: tensor.detach().clone() for name, tensor in pipeline.state_dict().items()}
+        backend = _layer_backend()
+        backend.enable(pipeline)
+        failing_hook = backend._dit_hooks[0]
+        restore = failing_hook.restore_next_block
+        attempts = 0
+
+        def fail_once():
+            nonlocal attempts
+            attempts += 1
+            if attempts == 1:
+                raise RuntimeError("injected restore failure")
+            restore()
+
+        monkeypatch.setattr(failing_hook, "restore_next_block", fail_once)
+
+        with pytest.raises(RuntimeError, match="Failed to fully disable"):
+            backend.disable()
+
+        assert backend.enabled
+        assert backend._dit_hooks
+
+        backend.disable()
+
+        assert not backend.enabled
+        for name, tensor in pipeline.state_dict().items():
+            torch.testing.assert_close(tensor, expected[name])
+
+    def test_encoder_disable_failure_keeps_host_masters_for_retry(
+        self,
+        patched_offload_runtime,
+        monkeypatch,
+    ):
+        pipeline = _ComponentPipeline()
+        encoder = pipeline.text_encoder
+        expected = {name: tensor.detach().clone() for name, tensor in encoder.state_dict().items()}
+        backend = _layer_backend({"text_encoder"})
+        backend.enable(pipeline)
+        failing_hook = encoder._omni_layerwise_hooks[0]
+        restore = failing_hook.restore_next_block
+        attempts = 0
+
+        def fail_once():
+            nonlocal attempts
+            attempts += 1
+            if attempts == 1:
+                raise RuntimeError("injected encoder restore failure")
+            restore()
+
+        monkeypatch.setattr(failing_hook, "restore_next_block", fail_once)
+
+        with pytest.raises(RuntimeError, match="Failed to fully disable"):
+            backend.disable()
+
+        assert backend.enabled
+        assert backend._encoder_modules == [encoder]
+        assert encoder._omni_layerwise_enabled
+        assert encoder._omni_layerwise_hooks
+
+        backend.disable()
+
+        assert not backend.enabled
+        assert not encoder._omni_layerwise_enabled
+        for name, tensor in encoder.state_dict().items():
+            torch.testing.assert_close(tensor, expected[name])
+
+    def test_legacy_selector_preserves_planned_encoder_and_vae_lifecycle(self, patched_offload_runtime):
+        pipeline = _ComponentPipeline()
+        backend = _layer_backend()
+
+        backend.enable(pipeline)
+
+        assert pipeline.text_encoder._omni_layerwise_enabled
+        assert pipeline.text_encoder.offload_calls == 1
+        assert pipeline.vae.offload_calls == 1
+
+        backend.disable()
+
+    def test_standard_encoder_needs_only_declared_block_paths(self, patched_offload_runtime):
+        pipeline = _GenericEncoderPipeline()
+        backend = _layer_backend({"text_encoder"})
+
+        backend.enable(pipeline)
+
+        blocks = pipeline.text_encoder.encoder.block
+        assert pipeline.text_encoder._omni_layerwise_enabled
+        assert all(block._hook_registry.get_hook("layerwise_offload") is not None for block in blocks)
+        assert pipeline.text_encoder.final_norm.weight.numel() == 4
+
+        backend.disable()
+
+        assert not pipeline.text_encoder._omni_layerwise_enabled
+        assert all(block._hook_registry.get_hook("layerwise_offload") is None for block in blocks)
+        assert all(block.weight.numel() == 100 for block in blocks)
+
+
+def _offload_od_config(**overrides):
+    values = {
+        "diffusion_offload_config": None,
+        "enable_cpu_offload": False,
+        "enable_layerwise_offload": False,
+        "enable_distributed_layerwise_offload": False,
+        "dlo_use_allgather": True,
+        "dlo_resident_layers": 0,
+        "pin_cpu_memory": True,
+        "parallel_config": None,
+        "model": "/fake/model",
+    }
+    values.update(overrides)
+    return SimpleNamespace(**values)
+
+
+def _resolve_offload_config(public=None, **overrides) -> OffloadConfig:
+    return OffloadConfig.from_od_config(_offload_od_config(diffusion_offload_config=public, **overrides))
+
+
+class TestLayerwiseComponentConfig:
+    def test_layer_mode_selects_components_without_backend_jargon(self):
+        config = _resolve_offload_config(
+            {
+                "mode": "layer",
+                "components": ["dit", "text_encoder"],
+            }
+        )
+
+        assert config.strategy is OffloadStrategy.LAYER_WISE
+        assert config.components == frozenset({"dit", "text_encoder"})
+        assert not config.uses_allgather("dit")
+        assert not config.uses_allgather("text_encoder")
+        plan = OffloadPlan(encoder_component_types={"mllm": "text_encoder"})
+        assert config.offloads_encoder("mllm", plan)
+
+    @pytest.mark.parametrize("component", ["image_encoder", "vae", "scheduler", "text-encoder"])
+    def test_unknown_or_noncanonical_component_is_rejected(self, component):
+        with pytest.raises(ValueError, match="Unknown diffusion offload component"):
+            OffloadConfig.from_od_config(
+                _offload_od_config(diffusion_offload_config={"mode": "layer", "components": [component]})
+            )
+
+    def test_components_is_selection_only_not_an_options_mapping(self):
+        with pytest.raises(TypeError, match="components must be a non-empty list"):
+            OffloadConfig.from_od_config(
+                _offload_od_config(
+                    diffusion_offload_config={
+                        "mode": "layer",
+                        "components": {"dit": {}},
+                    }
+                )
+            )
+
+    def test_layer_options_requires_component_name_keys(self):
+        with pytest.raises(TypeError, match="layer_options keys must be strings"):
+            OffloadConfig.from_od_config(
+                _offload_od_config(
+                    diffusion_offload_config={
+                        "mode": "layer",
+                        "components": ["dit"],
+                        "layer_options": {0: {}},
+                    }
+                )
+            )
+
+    @pytest.mark.parametrize(
+        ("config", "message"),
+        [
+            ({0: "layer"}, "diffusion_offload_config keys must be strings"),
+            (
+                {
+                    "mode": "layer",
+                    "components": ["dit"],
+                    "layer_options": {"dit": {0: "allgather"}},
+                },
+                "layer_options\\['dit'\\] keys must be strings",
+            ),
+        ],
+    )
+    def test_all_public_mapping_levels_require_string_keys(self, config, message):
+        with pytest.raises(TypeError, match=message):
+            OffloadConfig.from_od_config(_offload_od_config(diffusion_offload_config=config))
+
+    @pytest.mark.parametrize(
+        "config,match",
+        [
+            ({"components": ["dit"]}, "requires 'mode'"),
+            ({"mode": "layer"}, "requires 'components'"),
+            ({"mode": "layer", "components": []}, "must not be empty"),
+            ({"mode": "layer", "components": [" dit "]}, "Unknown diffusion offload component"),
+            ({"mode": "layerwise", "components": ["dit"]}, "Unknown diffusion offload mode"),
+            (
+                {
+                    "mode": "layer",
+                    "components": ["dit"],
+                    "layer_options": {"dit": {"weight_transfer": "rank_local"}},
+                },
+                "Unknown offload transfer",
+            ),
+            (
+                {
+                    "mode": "layer",
+                    "components": ["dit"],
+                    "layer_options": {"dit": {"prefetch": 2}},
+                },
+                "Unknown diffusion offload setting",
+            ),
+            (
+                {
+                    "mode": "layer",
+                    "components": ["text_encoder"],
+                    "layer_options": {"text_encoder": {"resident_layers": 1}},
+                },
+                "supports only the 'dit'",
+            ),
+            (
+                {
+                    "mode": "layer",
+                    "components": ["dit"],
+                    "layer_options": {"text_encoder": {}},
+                },
+                "requires selecting the same component",
+            ),
+        ],
+    )
+    def test_schema_rejects_ambiguous_or_unsupported_values(self, config, match):
+        with pytest.raises(ValueError, match=match):
+            OffloadConfig.from_od_config(_offload_od_config(diffusion_offload_config=config))
+
+    def test_module_mode_accepts_component_selection(self):
+        config = _resolve_offload_config(
+            {
+                "mode": "module",
+                "components": ["dit", "text_encoder"],
+            }
+        )
+
+        assert config.strategy is OffloadStrategy.MODEL_LEVEL
+        assert config.components == frozenset({"dit", "text_encoder"})
+
+    @pytest.mark.parametrize("setting", [{"weight_transfer": "rank-local"}, {"resident_layers": 1}])
+    def test_module_mode_rejects_layer_settings(self, setting):
+        with pytest.raises(ValueError, match="layer_options requires mode='layer'"):
+            OffloadConfig.from_od_config(
+                _offload_od_config(
+                    diffusion_offload_config={
+                        "mode": "module",
+                        "components": ["dit"],
+                        "layer_options": {"dit": setting},
+                    }
+                )
+            )
+
+    def test_each_component_can_choose_its_transfer(self):
+        config = _resolve_offload_config(
+            {
+                "mode": "layer",
+                "components": ["dit", "text_encoder"],
+                "layer_options": {
+                    "dit": {"weight_transfer": "rank-local"},
+                    "text_encoder": {"weight_transfer": "allgather"},
+                },
+            }
+        )
+
+        assert config.strategy is OffloadStrategy.DISTRIBUTED_LAYER_WISE
+        assert not config.uses_allgather("dit")
+        assert config.uses_allgather("text_encoder")
+        assert config.dlo_use_allgather is False
+
+    def test_dit_residency_selects_capable_backend(self):
+        config = _resolve_offload_config(
+            {
+                "mode": "layer",
+                "components": ["dit"],
+                "layer_options": {"dit": {"resident_layers": 20}},
+            }
+        )
+
+        assert config.strategy is OffloadStrategy.DISTRIBUTED_LAYER_WISE
+        assert config.dlo_resident_layers == 20
+        assert not config.uses_allgather("dit")
+
+    def test_dit_residency_rejects_allgather(self):
+        with pytest.raises(ValueError, match="resident_layers requires dit.weight_transfer='rank-local'"):
+            OffloadConfig.from_od_config(
+                _offload_od_config(
+                    diffusion_offload_config={
+                        "mode": "layer",
+                        "components": ["dit"],
+                        "layer_options": {"dit": {"weight_transfer": "allgather", "resident_layers": 1}},
+                    }
+                )
+            )
+
+    def test_host_weight_runtime_remains_separate_from_compact_config(self):
+        with pytest.raises(ValueError, match="cannot be combined with Host Weight Runtime"):
+            OffloadConfig.from_od_config(
+                _offload_od_config(
+                    diffusion_offload_config={"mode": "layer", "components": ["dit"]},
+                    host_weight_runtime_mode="preferred",
+                )
+            )
+
+    def test_pin_memory_override_is_scoped_to_offload_config(self):
+        config = _resolve_offload_config(
+            {
+                "mode": "layer",
+                "components": ["dit"],
+                "pin_memory": False,
+            },
+            pin_cpu_memory=True,
+        )
+
+        assert config.pin_cpu_memory is False
+
+    def test_legacy_omitted_selector_preserves_dit_only_behavior(self):
+        config = _resolve_offload_config(enable_layerwise_offload=True)
+
+        assert config.strategy is OffloadStrategy.LAYER_WISE
+        assert config.components is None
+
+    def test_offload_resolution_is_cached_for_runtime_accessors(self, monkeypatch):
+        import vllm_omni.diffusion.offloader.config as config_module
+
+        od_config = _offload_od_config(
+            diffusion_offload_config={
+                "mode": "layer",
+                "components": ["dit", "text_encoder"],
+            }
+        )
+        parse_calls = 0
+        original_parse = config_module.parse_diffusion_offload_config
+
+        def track_parse(value):
+            nonlocal parse_calls
+            parse_calls += 1
+            return original_parse(value)
+
+        monkeypatch.setattr(config_module, "parse_diffusion_offload_config", track_parse)
+
+        assert config_module.resolve_offload_strategy(od_config) is OffloadStrategy.LAYER_WISE
+        assert config_module.selected_offload_components(od_config) == {"dit", "text_encoder"}
+        assert not config_module.component_uses_allgather(od_config, "dit")
+        assert not config_module.any_selected_component_uses_allgather(od_config)
+        assert parse_calls == 1
+
+    def test_cached_policy_is_deeply_immutable_and_pickle_safe(self):
+        import vllm_omni.diffusion.offloader.config as config_module
+
+        od_config = _offload_od_config(
+            diffusion_offload_config={
+                "mode": "layer",
+                "components": ["dit"],
+                "layer_options": {"dit": {"weight_transfer": "rank-local"}},
+            }
+        )
+        resolved = config_module.resolve_offload(od_config)
+
+        with pytest.raises(TypeError):
+            resolved.transfers["dit"] = config_module.DLOTransfer.ALLGATHER  # type: ignore[index]
+        assert resolved.public is not None
+        with pytest.raises(TypeError):
+            resolved.public.layer_options["dit"] = config_module.LayerOffloadOptions()  # type: ignore[index]
+
+        restored = ForkingPickler.loads(ForkingPickler.dumps(resolved))
+        assert not restored.uses_allgather("dit")
+
+    def test_dp_text_encoder_allgather_rejects_rank_local_prompt_cache(self):
+        config = _offload_od_config(
+            diffusion_offload_config={
+                "mode": "layer",
+                "components": ["text_encoder"],
+                "layer_options": {"text_encoder": {"weight_transfer": "allgather"}},
+            },
+            enable_prompt_embed_cache=True,
+            parallel_config=SimpleNamespace(
+                data_parallel_size=2,
+                sequence_parallel_size=1,
+                use_hsdp=False,
+            ),
+        )
+
+        with pytest.raises(ValueError, match="rank-local cache hits would skip different encoder collectives"):
+            OffloadConfig.from_od_config(config)
+
+    def test_dp_allgather_rejects_rank_local_dit_cache_decisions(self):
+        config = _offload_od_config(
+            diffusion_offload_config={
+                "mode": "layer",
+                "components": ["dit"],
+                "layer_options": {"dit": {"weight_transfer": "allgather"}},
+            },
+            cache_backend="tea_cache",
+            parallel_config=SimpleNamespace(
+                data_parallel_size=2,
+                sequence_parallel_size=1,
+                use_hsdp=False,
+            ),
+        )
+
+        with pytest.raises(ValueError, match="rank-local cache decisions can skip different weight collectives"):
+            OffloadConfig.from_od_config(config)
+
+    def test_internal_transfer_map_requires_every_component(self):
+        with pytest.raises(ValueError, match="missing: text_encoder"):
+            OffloadConfig(
+                strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
+                dlo_transfers={"dit": "rank-local"},
+            )
+
+    def test_legacy_distributed_settings_still_work(self):
+        config = _resolve_offload_config(
+            enable_distributed_layerwise_offload=True,
+            dlo_use_allgather=False,
+            dlo_resident_layers=3,
+        )
+
+        assert config.strategy is OffloadStrategy.DISTRIBUTED_LAYER_WISE
+        assert config.dlo_resident_layers == 3
+        assert not config.uses_allgather("dit")
+        assert not config.uses_allgather("text_encoder")
+
+    def test_legacy_allgather_remains_dit_only(self):
+        config = _resolve_offload_config(
+            enable_distributed_layerwise_offload=True,
+            dlo_use_allgather=True,
+        )
+
+        assert config.uses_allgather("dit")
+        assert not config.uses_allgather("text_encoder")
+
+    @pytest.mark.parametrize(
+        ("legacy_flags", "expected_strategy"),
+        [
+            (
+                {"enable_cpu_offload": True, "enable_layerwise_offload": True},
+                OffloadStrategy.LAYER_WISE,
+            ),
+            (
+                {
+                    "enable_cpu_offload": True,
+                    "enable_layerwise_offload": True,
+                    "enable_distributed_layerwise_offload": True,
+                },
+                OffloadStrategy.DISTRIBUTED_LAYER_WISE,
+            ),
+        ],
+    )
+    def test_legacy_aliases_preserve_existing_priority(self, legacy_flags, expected_strategy):
+        config = _resolve_offload_config(**legacy_flags)
+
+        assert config.strategy is expected_strategy
+
+    def test_compact_config_cannot_mix_with_legacy_alias(self):
+        with pytest.raises(ValueError, match="cannot be combined"):
+            OffloadConfig.from_od_config(
+                _offload_od_config(
+                    diffusion_offload_config={"mode": "layer", "components": ["dit"]},
+                    enable_cpu_offload=True,
+                )
+            )
+
+    @pytest.mark.parametrize(
+        "legacy_options",
+        [
+            {"dlo_use_allgather": False},
+            {"dlo_resident_layers": 12},
+        ],
+    )
+    def test_compact_config_cannot_mix_with_legacy_dlo_options(self, legacy_options):
+        with pytest.raises(ValueError, match="cannot be combined with legacy DLO option"):
+            OffloadConfig.from_od_config(
+                _offload_od_config(
+                    diffusion_offload_config={"mode": "layer", "components": ["dit"]},
+                    **legacy_options,
+                )
+            )
+
+    def test_legacy_allgather_resident_layers_fail_during_config_resolution(self):
+        with pytest.raises(ValueError, match="requires the DiT DLO transfer to be rank-local"):
+            OffloadConfig.from_od_config(
+                _offload_od_config(
+                    enable_distributed_layerwise_offload=True,
+                    dlo_use_allgather=True,
+                    dlo_resident_layers=12,
+                )
+            )

tests/diffusion/offloader/test_sequential_backend.py

@@ -7,7 +7,9 @@ import pytest
 import torch
 from torch import nn
 
+from vllm_omni.diffusion.hooks import HookRegistry
 from vllm_omni.diffusion.offloader.base import OffloadConfig, OffloadStrategy
+from vllm_omni.diffusion.offloader.offload_plan import OffloadPlan
 from vllm_omni.diffusion.offloader.sequential_backend import (
     ModelLevelOffloadBackend,
     SequentialOffloadHook,
@@ -48,20 +50,21 @@ def _track_pin_memory_calls():
     return tracker, mock
 
 
+class _CustomPipeline(nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+        self.enable_args: dict[str, object] | None = None
+        self.disable_called = False
+
+    def enable_omni_model_cpu_offload(self, **kwargs) -> None:
+        self.enable_args = kwargs
+
+    def disable_omni_model_cpu_offload(self) -> None:
+        self.disable_called = True
+
+
 def test_model_level_backend_delegates_to_custom_pipeline_offload() -> None:
-    class CustomPipeline(nn.Module):
-        def __init__(self) -> None:
-            super().__init__()
-            self.enable_args = None
-            self.disable_called = False
-
-        def enable_omni_model_cpu_offload(self, **kwargs) -> None:
-            self.enable_args = kwargs
-
-        def disable_omni_model_cpu_offload(self) -> None:
-            self.disable_called = True
-
-    pipeline = CustomPipeline()
+    pipeline = _CustomPipeline()
     backend = ModelLevelOffloadBackend(
         OffloadConfig(strategy=OffloadStrategy.MODEL_LEVEL, pin_cpu_memory=False),
         torch.device("cpu"),
@@ -82,6 +85,205 @@ def test_model_level_backend_delegates_to_custom_pipeline_offload() -> None:
     assert pipeline.disable_called is True
 
 
+def test_model_level_backend_passes_explicit_component_selection() -> None:
+    pipeline = _CustomPipeline()
+    backend = ModelLevelOffloadBackend(
+        OffloadConfig(
+            strategy=OffloadStrategy.MODEL_LEVEL,
+            components=frozenset({"dit", "text_encoder"}),
+        ),
+        torch.device("cpu"),
+    )
+
+    backend.enable(pipeline)
+
+    assert pipeline.enable_args is not None
+    assert pipeline.enable_args["offload_components"] == frozenset({"dit", "text_encoder"})
+
+
+def test_model_level_backend_rolls_back_partial_custom_enable() -> None:
+    class CustomPipeline(nn.Module):
+        def __init__(self) -> None:
+            super().__init__()
+            self.disable_called = False
+
+        def enable_omni_model_cpu_offload(self, **kwargs) -> None:
+            del kwargs
+            raise RuntimeError("injected custom enable failure")
+
+        def disable_omni_model_cpu_offload(self) -> None:
+            self.disable_called = True
+
+    pipeline = CustomPipeline()
+    backend = ModelLevelOffloadBackend(
+        OffloadConfig(strategy=OffloadStrategy.MODEL_LEVEL),
+        torch.device("cpu"),
+    )
+
+    with pytest.raises(RuntimeError, match="injected custom enable failure"):
+        backend.enable(pipeline)
+
+    assert pipeline.disable_called
+    assert not backend.enabled
+
+
+def test_sequential_offload_rolls_back_partial_hook_registration(monkeypatch: pytest.MonkeyPatch) -> None:
+    dit = _create_simple_module()
+    encoder = _create_simple_module()
+    original_register = HookRegistry.register_hook
+    calls = 0
+
+    def fail_second_registration(self, name, hook):
+        nonlocal calls
+        calls += 1
+        if calls == 2:
+            raise RuntimeError("injected registration failure")
+        return original_register(self, name, hook)
+
+    monkeypatch.setattr(HookRegistry, "register_hook", fail_second_registration)
+
+    with pytest.raises(RuntimeError, match="injected registration failure"):
+        apply_sequential_offload(
+            dit_modules=[dit],
+            encoder_modules=[encoder],
+            device=torch.device("cpu"),
+        )
+
+    for module in (dit, encoder):
+        registry = getattr(module, "_hook_registry", None)
+        assert registry is None or registry.get_hook(SequentialOffloadHook._HOOK_NAME) is None
+
+
+def test_model_level_disable_cleans_remaining_hooks_and_can_retry(monkeypatch: pytest.MonkeyPatch) -> None:
+    pipeline = nn.Module()
+    pipeline.transformer = _create_simple_module()
+    pipeline.text_encoder = _create_simple_module()
+    backend = ModelLevelOffloadBackend(
+        OffloadConfig(strategy=OffloadStrategy.MODEL_LEVEL, pin_cpu_memory=False),
+        torch.device("cpu"),
+    )
+    backend.enable(pipeline)
+
+    original_remove = HookRegistry.remove_hook
+    calls = 0
+
+    def fail_first_removal(self, name):
+        nonlocal calls
+        calls += 1
+        if calls == 1:
+            raise RuntimeError("injected hook removal failure")
+        original_remove(self, name)
+
+    monkeypatch.setattr(HookRegistry, "remove_hook", fail_first_removal)
+
+    with pytest.raises(RuntimeError, match="Failed to remove one or more sequential offload hooks"):
+        backend.disable()
+
+    assert backend.enabled
+    assert backend._offload_modules == [pipeline.transformer, pipeline.text_encoder]
+    assert pipeline.transformer._hook_registry.get_hook(SequentialOffloadHook._HOOK_NAME) is not None
+    assert pipeline.text_encoder._hook_registry.get_hook(SequentialOffloadHook._HOOK_NAME) is None
+
+    backend.disable()
+
+    assert not backend.enabled
+    assert not backend._offload_modules
+    for module in (pipeline.transformer, pipeline.text_encoder):
+        assert module._hook_registry.get_hook(SequentialOffloadHook._HOOK_NAME) is None
+
+
+def test_sequential_offload_filters_cpu_eligible_components() -> None:
+    dit = _create_simple_module()
+    encoder = _create_simple_module()
+    resident_stage = _create_simple_module()
+
+    apply_sequential_offload(
+        dit_modules=[dit],
+        encoder_modules=[encoder, resident_stage],
+        device=torch.device("cpu"),
+        offload_dit_modules=[dit],
+        offload_encoder_modules=[encoder],
+    )
+
+    dit_hook = dit._hook_registry.get_hook(SequentialOffloadHook._HOOK_NAME)
+    encoder_hook = encoder._hook_registry.get_hook(SequentialOffloadHook._HOOK_NAME)
+    resident_hook = resident_stage._hook_registry.get_hook(SequentialOffloadHook._HOOK_NAME)
+    assert dit_hook.offload_targets == [encoder]
+    assert encoder_hook.offload_targets == [dit]
+    assert encoder_hook.offload_after_context is True
+    assert resident_hook.offload_targets == [dit]
+    assert resident_hook.offload_after_context is False
+
+    remove_sequential_offload([dit, encoder, resident_stage])
+
+
+def test_move_params_handles_buffer_only_modules() -> None:
+    module = nn.Module()
+    module.register_buffer("state", torch.ones(2))
+
+    # An indexed CPU device compares differently while remaining usable on
+    # CPU-only test hosts, so the helper must inspect the buffer itself.
+    moved = SequentialOffloadHook._move_params(module, torch.device("cpu:1"))
+
+    assert moved
+    assert module.state.device.type == "cpu"
+
+
+def test_model_level_backend_keeps_declared_image_encoder_resident() -> None:
+    class MixedEncoderPipeline(nn.Module):
+        _offload_plan = OffloadPlan(
+            encoder_component_types={
+                "text_encoder": "text_encoder",
+            }
+        )
+
+        def __init__(self) -> None:
+            super().__init__()
+            self.transformer = _create_simple_module()
+            self.text_encoder = _create_simple_module()
+            self.image_encoder = _create_simple_module()
+
+    pipeline = MixedEncoderPipeline()
+    backend = ModelLevelOffloadBackend(
+        OffloadConfig(
+            strategy=OffloadStrategy.MODEL_LEVEL,
+            components=frozenset({"dit", "text_encoder"}),
+        ),
+        torch.device("cpu"),
+    )
+
+    backend.enable(pipeline)
+
+    text_hook = pipeline.text_encoder._hook_registry.get_hook(SequentialOffloadHook._HOOK_NAME)
+    image_hook = pipeline.image_encoder._hook_registry.get_hook(SequentialOffloadHook._HOOK_NAME)
+    assert text_hook.offload_after_context is True
+    assert image_hook.offload_after_context is False
+
+    backend.disable()
+
+
+@pytest.mark.parametrize(
+    ("component", "attribute", "message"),
+    [
+        ("text_encoder", "text_encoder", "requires a DiT/transformer"),
+        ("dit", "transformer", "requires an encoder execution stage"),
+    ],
+)
+def test_component_selective_model_offload_requires_swap_counterpart(component, attribute, message) -> None:
+    pipeline = nn.Module()
+    setattr(pipeline, attribute, _create_simple_module())
+    backend = ModelLevelOffloadBackend(
+        OffloadConfig(
+            strategy=OffloadStrategy.MODEL_LEVEL,
+            components=frozenset({component}),
+        ),
+        torch.device("cpu"),
+    )
+
+    with pytest.raises(ValueError, match=message):
+        backend.enable(pipeline)
+
+
 def test_sequential_offload_can_begin_with_dit_on_cpu(monkeypatch: pytest.MonkeyPatch) -> None:
     dit = _create_simple_module()
     encoder = _create_simple_module()

tests/diffusion/test_diffusion_engine.py

@@ -600,6 +600,34 @@ class TestDiffusionCompileConfig:
             ),
             ({"enable_cpu_offload": True}, "CPU offload"),
             ({"enable_layerwise_offload": True}, "layerwise offload"),
+            (
+                {
+                    "diffusion_offload_config": {
+                        "mode": "module",
+                        "components": ["dit"],
+                    }
+                },
+                "CPU offload",
+            ),
+            (
+                {
+                    "diffusion_offload_config": {
+                        "mode": "layer",
+                        "components": ["dit"],
+                    }
+                },
+                "layerwise offload",
+            ),
+            (
+                {
+                    "diffusion_offload_config": {
+                        "mode": "layer",
+                        "components": ["dit"],
+                        "layer_options": {"dit": {"weight_transfer": "allgather"}},
+                    }
+                },
+                "distributed layerwise offload",
+            ),
         ],
     )
     def test_full_compile_rejects_incompatible_features(self, kwargs, feature) -> None:
@@ -610,6 +638,34 @@ class TestDiffusionCompileConfig:
                 **kwargs,
             )
 
+    def test_raw_parallel_mapping_is_normalized_before_allgather_cache_validation(self) -> None:
+        with pytest.raises(ValueError, match="rank-local cache decisions"):
+            OmniDiffusionConfig(
+                model="test",
+                num_gpus=2,
+                parallel_config={"data_parallel_size": 2},
+                cache_backend="tea_cache",
+                diffusion_offload_config={
+                    "mode": "layer",
+                    "components": ["dit"],
+                    "layer_options": {"dit": {"weight_transfer": "allgather"}},
+                },
+            )
+
+    def test_auto_dp_is_resolved_before_encoder_prompt_cache_validation(self) -> None:
+        with pytest.raises(ValueError, match="rank-local cache hits"):
+            OmniDiffusionConfig(
+                model="test",
+                num_gpus=2,
+                parallel_config={"data_parallel_size": None},
+                enable_prompt_embed_cache=True,
+                diffusion_offload_config={
+                    "mode": "layer",
+                    "components": ["text_encoder"],
+                    "layer_options": {"text_encoder": {"weight_transfer": "allgather"}},
+                },
+            )
+
 
 class TestRequestBatchAdmission:
     pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]

tests/diffusion/test_diffusion_engine_dummy_run.py

@@ -52,6 +52,23 @@ def test_dummy_run_uses_enough_steps_for_execution_mode(
     assert captured_requests[0].sampling_params.num_inference_steps == 2
 
 
+def test_allgather_startup_runs_broadcast_dummy_request() -> None:
+    engine = object.__new__(DiffusionEngine)
+    engine.od_config = SimpleNamespace(
+        diffusion_offload_config={
+            "mode": "layer",
+            "components": ["dit"],
+            "layer_options": {"dit": {"weight_transfer": "allgather"}},
+        },
+        parallel_config=SimpleNamespace(data_parallel_size=2, sequence_parallel_size=1),
+    )
+    engine._dummy_run = Mock()
+
+    engine.run_startup_warmup()
+
+    engine._dummy_run.assert_called_once_with()
+
+
 def test_dummy_run_num_frames_uses_explicit_model_setting(monkeypatch: pytest.MonkeyPatch) -> None:
     class JointAudioVideoModel:
         dummy_run_num_frames = 2

tests/diffusion/test_diffusion_model_runner.py

@@ -484,20 +484,23 @@ def test_compile_transformer_uses_full_granularity(monkeypatch):
 
 @pytest.mark.core_model
 @pytest.mark.cpu
-def test_compile_transformer_falls_back_after_synchronous_setup_failure(monkeypatch, caplog):
+def test_compile_transformer_falls_back_after_synchronous_setup_failure(monkeypatch):
     model = _CompileTrackingModel()
     runner = _make_compile_runner(model)
+    warning = Mock()
 
     def _regionally_compile(*args, **kwargs):
         raise RuntimeError("compile setup failed")
 
     monkeypatch.setattr(model_runner_module, "regionally_compile", _regionally_compile)
+    monkeypatch.setattr(model_runner_module.logger, "warning", warning)
 
     DiffusionModelRunner._compile_transformer(runner, "transformer")
 
     assert runner.pipeline.transformer is model
-    assert "failed before activation" in caplog.text
-    assert "lazy compilation errors" in caplog.text
+    warning.assert_called_once()
+    assert "failed before activation" in warning.call_args.args[0]
+    assert "lazy compilation errors" in warning.call_args.args[0]
 
 
 @pytest.mark.core_model
@@ -1061,6 +1064,51 @@ def test_load_model_clears_cache_backend_for_unsupported_pipeline(monkeypatch):
     assert dummy_cache_backend.enabled is False
 
 
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_load_model_rejects_env_prompt_cache_before_text_encoder_allgather(monkeypatch):
+    runner = object.__new__(DiffusionModelRunner)
+    runner.vllm_config = object()
+    runner.device = torch.device("cpu")
+    runner.pipeline = None
+    runner.offload_backend = None
+    runner.od_config = SimpleNamespace(
+        diffusion_offload_config={
+            "mode": "layer",
+            "components": ["text_encoder"],
+            "layer_options": {"text_encoder": {"weight_transfer": "allgather"}},
+        },
+        enable_cpu_offload=False,
+        enable_layerwise_offload=False,
+        enable_distributed_layerwise_offload=True,
+        dlo_use_allgather=True,
+        dlo_resident_layers=0,
+        pin_cpu_memory=True,
+        parallel_config=SimpleNamespace(data_parallel_size=2),
+        enable_prompt_embed_cache=False,
+        prompt_embed_cache_size=32,
+        cache_backend="none",
+        model_class_name="Wan22Pipeline",
+        enforce_eager=True,
+        streaming_output=False,
+        step_execution=False,
+    )
+    monkeypatch.setenv("OMNI_DIFFUSION_PROMPT_EMBED_CACHE", "1")
+    runtime_init = Mock(side_effect=pytest.fail)
+    loader = Mock(side_effect=pytest.fail)
+    monkeypatch.setattr(model_runner_module.current_omni_platform, "init_diffusion_model_runner_runtime", runtime_init)
+    monkeypatch.setattr(model_runner_module, "DiffusersPipelineLoader", loader)
+    enable_offload = Mock(side_effect=pytest.fail)
+    monkeypatch.setattr(model_runner_module, "enable_offload_backend", enable_offload)
+
+    with pytest.raises(ValueError, match="Prompt embedding cache cannot be combined"):
+        DiffusionModelRunner.load_model(runner)
+
+    runtime_init.assert_not_called()
+    loader.assert_not_called()
+    enable_offload.assert_not_called()
+
+
 def test_set_forward_context_enters_vllm_config_contexts(monkeypatch):
     """Ensure `with set_forward_context(...):` enters vllm's context managers internally and calls desired vllm functions."""
     import vllm.config.vllm as vllm_config_module

tests/diffusion/test_multiproc_engine_concurrency.py

@@ -254,8 +254,11 @@ class TestRequestModeDispatch:
         executor.od_config = SimpleNamespace(
             step_execution=False,
             parallel_config=SimpleNamespace(data_parallel_size=2),
-            enable_distributed_layerwise_offload=True,
-            dlo_use_allgather=True,
+            diffusion_offload_config={
+                "mode": "layer",
+                "components": ["dit"],
+                "layer_options": {"dit": {"weight_transfer": "allgather"}},
+            },
         )
         executor.execute_request = Mock(return_value="dlo-dp")
         executor.collective_rpc = Mock()
@@ -358,7 +361,7 @@ class TestRequestModeDispatch:
         for rank, result_mq in enumerate(executor._result_mqs):
             assert result_mq.dequeue.call_count == (1 if rank in expected_primary_ranks else 0)
 
-    @pytest.mark.parametrize("empty_prompt", ["", {"prompt": ""}])
+    @pytest.mark.parametrize("empty_prompt", ["", [], {}, {"prompt": ""}])
     def test_dlo_dp_rejects_empty_prompt_before_worker_dispatch(self, empty_prompt):
         executor, _, _ = _make_executor(num_gpus=2)
         executor.od_config = SimpleNamespace(
@@ -376,6 +379,71 @@ class TestRequestModeDispatch:
 
         executor.collective_rpc.assert_not_called()
 
+    def test_dlo_dp_accepts_precomputed_embeddings_without_prompt_text(self):
+        executor, _, _ = _make_executor(num_gpus=2)
+        executor.od_config = SimpleNamespace(
+            step_execution=False,
+            parallel_config=SimpleNamespace(data_parallel_size=2),
+            diffusion_offload_config={
+                "mode": "layer",
+                "components": ["text_encoder"],
+                "layer_options": {"text_encoder": {"weight_transfer": "allgather"}},
+            },
+        )
+        executor.collective_rpc = Mock(return_value=[_tagged_output("A"), _tagged_output("B")])
+        scheduler_output = _make_sched_output("A", "B")
+        for new_req in scheduler_output.scheduled_new_reqs:
+            new_req.req.prompt = {"prompt": "", "prompt_embeds": object()}
+
+        result = executor.execute_request(scheduler_output)
+
+        assert [output.result.error for output in result.runner_outputs] == ["A", "B"]
+
+    @pytest.mark.parametrize(
+        "prompt",
+        [
+            [1, 2],
+            {"type": "token", "prompt_token_ids": [1, 2]},
+            {"prompt_ids": [1, 2], "prompt_mask": object()},
+        ],
+    )
+    def test_dlo_dp_accepts_tokenized_prompts(self, prompt):
+        executor, _, _ = _make_executor(num_gpus=2)
+        executor.od_config = SimpleNamespace(
+            step_execution=False,
+            parallel_config=SimpleNamespace(data_parallel_size=2),
+            enable_distributed_layerwise_offload=True,
+            dlo_use_allgather=True,
+        )
+        executor.collective_rpc = Mock(return_value=[_tagged_output("A"), _tagged_output("B")])
+        scheduler_output = _make_sched_output("A", "B")
+        for new_req in scheduler_output.scheduled_new_reqs:
+            new_req.req.prompt = prompt
+
+        result = executor.execute_request(scheduler_output)
+
+        assert [output.result.error for output in result.runner_outputs] == ["A", "B"]
+
+    def test_text_encoder_allgather_rejects_mismatched_embedding_paths(self):
+        executor, _, _ = _make_executor(num_gpus=2)
+        executor.od_config = SimpleNamespace(
+            step_execution=False,
+            parallel_config=SimpleNamespace(data_parallel_size=2),
+            diffusion_offload_config={
+                "mode": "layer",
+                "components": ["text_encoder"],
+                "layer_options": {"text_encoder": {"weight_transfer": "allgather"}},
+            },
+        )
+        executor.collective_rpc = Mock()
+        scheduler_output = _make_sched_output("A", "B")
+        scheduler_output.scheduled_new_reqs[1].req.prompt = {"prompt_embeds": object()}
+
+        with pytest.raises(ValueError, match="same positive/negative prompt embedding fields"):
+            executor.execute_request(scheduler_output)
+
+        executor.collective_rpc.assert_not_called()
+
     def test_dlo_dp_valid_wave_still_runs_after_rejected_wave(self):
         executor, _, _ = _make_executor(num_gpus=2)
         executor.od_config = SimpleNamespace(
@@ -396,24 +464,6 @@ class TestRequestModeDispatch:
         assert [output.result.error for output in result.runner_outputs] == ["A", "B"]
         executor.collective_rpc.assert_called_once()
 
-    def test_dlo_dp_allows_shared_default_denoise_steps(self):
-        executor, _, _ = _make_executor(num_gpus=2)
-        executor.od_config = SimpleNamespace(
-            step_execution=False,
-            parallel_config=SimpleNamespace(data_parallel_size=2),
-            enable_distributed_layerwise_offload=True,
-            dlo_use_allgather=True,
-        )
-        executor.collective_rpc = Mock(return_value=[_tagged_output("A"), _tagged_output("B")])
-        scheduler_output = _make_sched_output("A", "B")
-        for new_req in scheduler_output.scheduled_new_reqs:
-            new_req.req.sampling_params.num_inference_steps = None
-
-        result = executor.execute_request(scheduler_output)
-
-        assert [output.result.error for output in result.runner_outputs] == ["A", "B"]
-        executor.collective_rpc.assert_called_once()
-
     def test_dlo_dp_forwards_request_metadata_envelopes_as_one_wave(self):
         executor, _, _ = _make_executor(num_gpus=2)
         executor.od_config = SimpleNamespace(
@@ -496,6 +546,28 @@ class TestRequestModeDispatch:
         executor._fail_closed_on_dp_wave_timeout.assert_called_once()
         assert isinstance(executor._fail_closed_on_dp_wave_timeout.call_args.args[0], TimeoutError)
 
+    def test_single_allgather_request_times_out_and_fails_closed(self):
+        from vllm_omni.diffusion.executor import multiproc_executor as executor_module
+
+        executor, _, _ = _make_executor(num_gpus=2)
+        executor.od_config = SimpleNamespace(
+            step_execution=False,
+            parallel_config=SimpleNamespace(data_parallel_size=2),
+            diffusion_offload_config={
+                "mode": "layer",
+                "components": ["dit"],
+                "layer_options": {"dit": {"weight_transfer": "allgather"}},
+            },
+        )
+        executor.collective_rpc = Mock(side_effect=TimeoutError("timed out"))
+        executor._fail_closed_on_dp_wave_timeout = Mock()
+
+        result = executor.execute_request(_make_sched_output("A"))
+
+        assert result.runner_outputs[0].result.error == "timed out"
+        assert executor.collective_rpc.call_args.kwargs["timeout"] == executor_module._DLO_DP_WAVE_TIMEOUT_S
+        executor._fail_closed_on_dp_wave_timeout.assert_called_once()
+
 
 # ───────────────── concurrent collective RPC ─────────────────